First working implementation of plugin (#6)

- Data nodes: `BatterySample`, `BatteryState`, `CyclingSpecs`, `TomatoSettings`.
The data schemas of `aiidalab-aurora` are used.
- Calculation nodes: `BatteryCyclerExperiment`, `BatteryFakeExperiment` (obsolete).
- Parser: `TomatoParser`
- Scheduler: `TomatoScheduler` (works with `tomato<=0.2a1`)

Note: tests are not yet implemented (#1).

.github/workflows/ci.yml

@@ -9,7 +9,7 @@ jobs:
     timeout-minutes: 30
     strategy:
       matrix:
-        python-version: [3.8]
+        python-version: ["3.8"]
         backend: ['django']
 
     services:

.github/workflows/publish-on-pypi.yml

@@ -20,21 +20,14 @@ jobs:
       with:
         python-version: 3.8
 
-    - name: Upgrade setuptools and install package
+    - name: Install flit
       run: |
-        python -m pip install --upgrade pip setuptools
-        python -m pip install -e .
+        python -m pip install --upgrade pip
+        python -m pip install flit~=3.4
 
-    - name: Assert package version
+    - name: Build and publish
+      run: |
+        flit publish
       env:
-        TAG_VERSION: ${{ github.ref }}
-      run: python ./.github/check_version.py
-
-    - name: Build source distribution
-      run: python ./setup.py sdist
-
-    - name: Publish package to PyPI
-      uses: pypa/gh-action-pypi-publish@master
-      with:
-        user: __token__
-        password: ${{ secrets.pypi_token }}
+        FLIT_USERNAME: __token__
+        FLIT_PASSWORD:  ${{ secrets.pypi_token }}

.gitignore

@@ -12,3 +12,4 @@
 build/
 dist/
 pip-wheel-metadata/
+fake_aurora_server/exp.out

.pre-commit-config.yaml

@@ -2,22 +2,28 @@
 # pre-commit install
 repos:
 - repo: https://github.com/pre-commit/pre-commit-hooks
-  rev: v3.4.0
+  rev: v4.1.0
   hooks:
-    - id: double-quote-string-fixer
     - id: end-of-file-fixer
-    - id: fix-encoding-pragma
     - id: mixed-line-ending
     - id: trailing-whitespace
     - id: check-json
 
-# yapf = yet another python formatter
-- repo: https://github.com/pre-commit/mirrors-yapf
-  rev: v0.30.0
+- repo: https://github.com/asottile/pyupgrade
+  rev: v2.31.0
   hooks:
-  - id: yapf
-    name: yapf
-    args: ["-i"]
+  - id: pyupgrade
+    args: ["--py37-plus"]
+
+- repo: https://github.com/PyCQA/isort
+  rev: 5.10.1
+  hooks:
+    - id: isort
+
+- repo: https://github.com/psf/black
+  rev: 22.3.0
+  hooks:
+    - id: black
 
 - repo: local
   hooks:
@@ -31,9 +37,3 @@ repos:
           docs/.*|
       )$
     entry: pylint
-
-  - id: version-number
-    name: Check version numbers
-    entry: python ./.github/check_version.py
-    language: system
-    files: '^(setup.json)|(aiida_aurora/__init__.py)'

aiida_aurora/__init__.py

@@ -1,8 +1,7 @@
-# -*- coding: utf-8 -*-
 """
 aiida_aurora
 
 AiiDA plugin for the Aurora platform.
 """
 
-__version__ = '0.1.0a0'
+__version__ = "0.2.0"

aiida_aurora/calculations/__init__.py

@@ -0,0 +1,2 @@
+from .cycler import BatteryCyclerExperiment
+from .fake import BatteryFakeExperiment

aiida_aurora/calculations/cycler.py

@@ -0,0 +1,117 @@
+"""
+Calculations provided by aiida_aurora.
+
+Register calculations via the "aiida.calculations" entry point in setup.json.
+"""
+from importlib import import_module
+
+import yaml
+
+from aiida.common import datastructures
+from aiida.engine import CalcJob
+from aiida.orm import ArrayData, SinglefileData
+
+from aiida_aurora.data.battery import BatterySample, BatteryState
+from aiida_aurora.data.control import TomatoSettings
+from aiida_aurora.data.experiment import CyclingSpecs
+
+
+class BatteryCyclerExperiment(CalcJob):
+    """
+    AiiDA calculation plugin for the tomato instrument automation package.
+    https://github.com/dgbowl/tomato
+    """
+    _INPUT_PAYLOAD_YAML_FILE = 'payload.yaml'
+    _INPUT_PAYLOAD_VERSION = '0.2'
+    _OUTPUT_FILE_PREFIX = 'results'
+
+    @classmethod
+    def define(cls, spec):
+        """Define inputs and outputs of the calculation."""
+        # yapf: disable
+        super().define(spec)
+
+        # set default values for AiiDA options
+        spec.inputs['metadata']['options']['resources'].default = {  # REQUIRED?
+            'num_machines': 1,
+        }
+        spec.inputs['metadata']['options']['parser_name'].default = 'aurora'
+        spec.inputs['metadata']['options']['withmpi'].default = False
+        spec.inputs['metadata']['options']['input_filename'].default = cls._INPUT_PAYLOAD_YAML_FILE
+        # spec.inputs['metadata']['options']['scheduler_stderr'].default = ''
+        # spec.inputs['metadata']['options']['scheduler_stdout'].default = ''
+
+        # new ports
+        spec.input('metadata.options.output_filename', valid_type=str, default=cls._OUTPUT_FILE_PREFIX)
+        spec.input('battery_sample', valid_type=BatterySample, help='Battery sample used.')
+        spec.input('technique', valid_type=CyclingSpecs, help='Experiment specifications.')
+        spec.input('control_settings', valid_type=TomatoSettings, help='Experiment control settings.')
+        spec.output('results', valid_type=ArrayData, help='Results of the experiment.')
+        spec.output('raw_data', valid_type=SinglefileData, help='Raw data retrieved.')
+        # spec.output('battery_state', valid_type=BatteryState, help='State of the battery after the experiment.')
+
+        spec.exit_code(300, 'ERROR_MISSING_OUTPUT_FILES', message='Experiment did not produce any kind of output file.')
+        spec.exit_code(301, 'ERROR_MISSING_JSON_FILE', message='Experiment did not produce an output json file.')
+        spec.exit_code(302, 'ERROR_MISSING_ZIP_FILE', message='Experiment did not produce a zip file with raw data.')
+        spec.exit_code(311, 'ERROR_COMPLETED_ERROR', message='The tomato job was marked as completed with an error.')
+        spec.exit_code(312, 'ERROR_COMPLETED_CANCELLED', message='The tomato job was marked as cancelled.')
+        spec.exit_code(400, 'ERROR_PARSING_JSON_FILE', message='Parsing of json file failed.')
+
+    def prepare_for_submission(self, folder):
+        """
+        Create input files.
+
+        :param folder: an `aiida.common.folders.Folder` where the plugin should temporarily place all files
+            needed by the calculation.
+        :return: `aiida.common.datastructures.CalcInfo` instance
+        """
+
+        # if connecting to a Windows PowerShell computer, change the extension of the submit script to '.ps1'
+        if self.inputs.code.computer.transport_type == 'sshtowin':
+            submit_script_filename = self.node.get_option('submit_script_filename')
+            if not submit_script_filename.endswith('.ps1'):
+                if submit_script_filename.endswith('.sh'):
+                    submit_script_filename = submit_script_filename[:-3] + '.ps1'
+                else:
+                    submit_script_filename(submit_script_filename + '.ps1')
+                self.node.backend_entity.set_attribute('submit_script_filename', submit_script_filename)
+
+        # prepare the payload
+        # TODO: use dgbowl_schemas module
+        ## from dgbowl_schemas.tomato.payload_0_1.tomato import Tomato
+        ## - should 'version: "0.1"' be written in the submit script?
+        ## - name of the file containing the sample and method
+        tomato_schema_module = import_module(f"aurora.schemas.tomato_{self._INPUT_PAYLOAD_VERSION.replace('.', 'p')}")
+        TomatoSchema = tomato_schema_module.tomato.Tomato
+
+        tomato_dict = self.inputs.control_settings.get_dict()
+        if tomato_dict['output']['prefix'] is None:
+            tomato_dict['output']['prefix'] = self._OUTPUT_FILE_PREFIX
+        else:
+            self._OUTPUT_FILE_PREFIX = tomato_dict['output']['prefix']
+        payload = tomato_schema_module.TomatoPayload(
+            version = self._INPUT_PAYLOAD_VERSION,
+            sample = tomato_schema_module.sample.convert_batterysample_to_sample(self.inputs.battery_sample.get_dict()),
+            method = tomato_schema_module.method.convert_electrochemsequence_to_method_list(
+                self.inputs.technique.get_dict()),
+            tomato = TomatoSchema(**tomato_dict),
+        )
+        with folder.open(self.options.input_filename, 'w', encoding='utf8') as handle:
+            handle.write(yaml.dump(payload.dict()))
+
+        codeinfo = datastructures.CodeInfo()
+
+        # the calculation code should be 'ketchup', so we add the following `cmdline_params`
+        # in order to submit the payload to tomato
+        codeinfo.cmdline_params = ['submit', '--jobname', '$JOB_TITLE', self._INPUT_PAYLOAD_YAML_FILE]
+        codeinfo.code_uuid = self.inputs.code.uuid
+        codeinfo.withmpi = self.inputs.metadata.options.withmpi
+
+        # Prepare a `CalcInfo` to be returned to the engine
+        calcinfo = datastructures.CalcInfo()
+        calcinfo.codes_info = [codeinfo]
+        calcinfo.local_copy_list = []
+        calcinfo.retrieve_list = [f'{self._OUTPUT_FILE_PREFIX}.json']
+        calcinfo.retrieve_temporary_list = [f'{self._OUTPUT_FILE_PREFIX}.zip']
+
+        return calcinfo