From aeafc603edbada229a2545339253965320f29dda Mon Sep 17 00:00:00 2001 From: Francisco Ramirez Date: Wed, 30 Aug 2023 20:22:21 +0200 Subject: [PATCH] (PR #19) Add cycling sequence workchain This PR adds a workflow allowing users to couple multiple protocols per battery sample in an experiment, with separate settings and monitors defined for each protocol. The main motivation of the PR is to ensure protocols run in the requested sequence. The workflow includes cross validation ensuring that settings/monitors match selected protocols. In addition, a check is applied between protocols ensuring previous protocols finished successfully, terminating the sequence otherwise. --- aiida_aurora/calculations/cycler.py | 4 +- aiida_aurora/schemas/cycling.py | 16 +- aiida_aurora/workflows/__init__.py | 11 ++ aiida_aurora/workflows/cycling_sequence.py | 154 ++++++++++++++++ examples/config_tests/code_ketchup-0.2rc2.yml | 10 ++ .../computer_localhost_tomato.yml | 12 ++ .../computer_localhost_tomato_config.yml | 2 + examples/cycling_sequence/README.md | 167 ++++++++++++++++++ .../cycling_sequence/test_cycling_sequence.py | 110 ++++++++++++ pyproject.toml | 3 + 10 files changed, 484 insertions(+), 5 deletions(-) create mode 100644 aiida_aurora/workflows/__init__.py create mode 100644 aiida_aurora/workflows/cycling_sequence.py create mode 100644 examples/config_tests/code_ketchup-0.2rc2.yml create mode 100644 examples/config_tests/computer_localhost_tomato.yml create mode 100644 examples/config_tests/computer_localhost_tomato_config.yml create mode 100644 examples/cycling_sequence/README.md create mode 100644 examples/cycling_sequence/test_cycling_sequence.py diff --git a/aiida_aurora/calculations/cycler.py b/aiida_aurora/calculations/cycler.py index 5cce42e..ff0b410 100644 --- a/aiida_aurora/calculations/cycler.py +++ b/aiida_aurora/calculations/cycler.py @@ -49,7 +49,7 @@ def define(cls, spec): default=cls._OUTPUT_FILE_PREFIX, ) spec.input("battery_sample", valid_type=BatterySampleData, help="Battery sample used.") - spec.input("technique", valid_type=CyclingSpecsData, help="Experiment specifications.") + spec.input("protocol", valid_type=CyclingSpecsData, help="Experiment specifications.") spec.input("control_settings", valid_type=TomatoSettingsData, 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.") @@ -115,7 +115,7 @@ def prepare_for_submission(self, folder): payload = TomatoPayload( version=self._INPUT_PAYLOAD_VERSION, sample=conversion_map[self._INPUT_PAYLOAD_VERSION]["sample"](self.inputs.battery_sample.get_dict()), - method=conversion_map[self._INPUT_PAYLOAD_VERSION]["method"](self.inputs.technique.get_dict()), + method=conversion_map[self._INPUT_PAYLOAD_VERSION]["method"](self.inputs.protocol.get_dict()), tomato=conversion_map[self._INPUT_PAYLOAD_VERSION]["tomato"](**tomato_dict), ) diff --git a/aiida_aurora/schemas/cycling.py b/aiida_aurora/schemas/cycling.py index 945e63b..22d527e 100644 --- a/aiida_aurora/schemas/cycling.py +++ b/aiida_aurora/schemas/cycling.py @@ -551,13 +551,22 @@ def __getitem__(self, item): class ElectroChemSequence(BaseModel): + name: str = "" method: Sequence[ElectroChemPayloads] + class Config: + validate_assignment = True + extra = Extra.forbid + @property def n_steps(self): "Number of steps of the method" return len(self.method) + def set_name(self, name: str) -> None: + """docstring""" + self.name = name + def add_step(self, elem): if not isinstance(elem, get_args(ElectroChemPayloads)): raise ValueError("Invalid technique") @@ -576,6 +585,7 @@ def move_step_forward(self, i): j = i + 1 self.method[i], self.method[j] = self.method[j], self.method[i] - class Config: - validate_assignment = True - extra = Extra.forbid + def __eq__(self, other: object) -> bool: + if not isinstance(other, ElectroChemSequence): + return NotImplemented + return self.name == other.name diff --git a/aiida_aurora/workflows/__init__.py b/aiida_aurora/workflows/__init__.py new file mode 100644 index 0000000..e6b4095 --- /dev/null +++ b/aiida_aurora/workflows/__init__.py @@ -0,0 +1,11 @@ +""" +Calculations provided by aiida_aurora. + +Register calculations via the "aiida.calculations" entry point in setup.json. +""" + +from .cycling_sequence import CyclingSequenceWorkChain + +__all__ = [ + 'CyclingSequenceWorkChain', +] diff --git a/aiida_aurora/workflows/cycling_sequence.py b/aiida_aurora/workflows/cycling_sequence.py new file mode 100644 index 0000000..73eec6e --- /dev/null +++ b/aiida_aurora/workflows/cycling_sequence.py @@ -0,0 +1,154 @@ +from aiida import orm +from aiida.engine import ToContext, WorkChain, append_, while_ +from aiida.plugins import CalculationFactory, DataFactory + +CyclerCalcjob = CalculationFactory('aurora.cycler') +CyclingSpecsData = DataFactory('aurora.cyclingspecs') +BatterySampleData = DataFactory('aurora.batterysample') +TomatoSettingsData = DataFactory('aurora.tomatosettings') + + +def validate_inputs(inputs, ctx=None): + """Validate the inputs of the entire input namespace.""" + error_message = '' + + reference_keys = set(inputs['protocols'].keys()) + for namekey in inputs['control_settings'].keys(): + + if namekey not in reference_keys: + error_message += f'namekey {namekey} missing in protocols\n' + continue + + reference_keys.remove(namekey) + + for remaining_key in reference_keys: + error_message += f'protocol {remaining_key} has no settings\n' + + if len(error_message) > 0: + return error_message + + +class CyclingSequenceWorkChain(WorkChain): + """This workflow represents a process containing a variable number of steps.""" + + @classmethod + def define(cls, spec): + """Define the process specification.""" + + super().define(spec) + + spec.input( + "battery_sample", + valid_type=BatterySampleData, + help="Battery sample to be used.", + ) + + spec.input( + "tomato_code", + valid_type=orm.Code, + help="Tomato code to use.", + ) + + spec.input_namespace( + "protocols", + dynamic=True, + valid_type=CyclingSpecsData, + help="List of experiment specifications.", + ) + + spec.input_namespace( + "control_settings", + dynamic=True, + valid_type=TomatoSettingsData, + help="List of experiment control settings.", + ) + + spec.input_namespace( + "monitor_settings", + dynamic=True, + valid_type=orm.Dict, + help="Dictionary of battery experiment monitor settings.", + ) + + spec.output_namespace( + "results", + dynamic=True, + valid_type=orm.ArrayData, + help="Results of each step by key.", + ) + + spec.inputs.validator = validate_inputs + + spec.outline( + cls.setup_workload, + while_(cls.has_steps_remaining)( + cls.run_cycling_step, + cls.inspect_cycling_step, + ), + cls.gather_results, + ) + + spec.exit_code( + 401, + 'ERROR_IN_CYCLING_STEP', + message='One of the steps of CyclingSequenceWorkChain failed', + ) + + def setup_workload(self): + """Take the inputs and wrap them together.""" + self.worksteps_keynames = list(self.inputs['protocols'].keys()) + + def has_steps_remaining(self): + """Check if there is any remaining step.""" + return len(self.worksteps_keynames) > 0 + + def inspect_cycling_step(self): + """Verify that the last cycling step finished successfully.""" + last_subprocess = self.ctx.subprocesses[-1] + + if not last_subprocess.is_finished_ok: + pkid = last_subprocess.pk + stat = last_subprocess.exit_status + self.report(f'Cycling substep failed with exit status {stat}') + return self.exit_codes.ERROR_IN_CYCLING_STEP + + def run_cycling_step(self): + """Run the next cycling step.""" + current_keyname = self.worksteps_keynames.pop(0) + + inputs = { + 'code': self.inputs.tomato_code, + 'battery_sample': self.inputs.battery_sample, + 'protocol': self.inputs.protocols[current_keyname], + 'control_settings': self.inputs.control_settings[current_keyname], + } + + has_monitors = current_keyname in self.inputs.monitor_settings + + if has_monitors: + inputs['monitors'] = self.inputs.monitor_settings[current_keyname] + + running = self.submit(CyclerCalcjob, **inputs) + + if has_monitors: + running.set_extra('monitored', True) + else: + running.set_extra('monitored', False) + + self.report(f'launching CyclerCalcjob<{running.pk}>') + return ToContext(subprocesses=append_(running)) + + def gather_results(self): + """Gather the results from all cycling steps.""" + keynames = list(self.inputs['protocols'].keys()) + if len(self.ctx.subprocesses) != len(keynames): + raise RuntimeError('Problem with subprocess!') + + multiple_results = {} + for keyname in keynames: + current_subprocess = self.ctx.subprocesses.pop(0) + if 'results' not in current_subprocess.outputs: + continue + multiple_results[keyname] = current_subprocess.outputs.results + + self.out('results', dict(multiple_results)) diff --git a/examples/config_tests/code_ketchup-0.2rc2.yml b/examples/config_tests/code_ketchup-0.2rc2.yml new file mode 100644 index 0000000..651d147 --- /dev/null +++ b/examples/config_tests/code_ketchup-0.2rc2.yml @@ -0,0 +1,10 @@ +--- +label: "ketchup-0.2rc2" +computer: "localhost-tomato" +description: "ketchup submit - tomato 0.2rc2" +input_plugin: "aurora.cycler" +on_computer: true +remote_abs_path: "/home/aiida/.conda/envs/tomato/bin/ketchup" +use_double_quotes: False +prepend_text: "" +append_text: "" diff --git a/examples/config_tests/computer_localhost_tomato.yml b/examples/config_tests/computer_localhost_tomato.yml new file mode 100644 index 0000000..b9ce410 --- /dev/null +++ b/examples/config_tests/computer_localhost_tomato.yml @@ -0,0 +1,12 @@ +--- +description: "Localhost running the tomato scheduler" +label: "localhost-tomato" +hostname: "localhost" +transport: "local" +scheduler: "tomato" +work_dir: "/home/aiida/aiida_run/" +mpirun_command: "mpirun -np {tot_num_mpiprocs}" +mpiprocs_per_machine: 1 +shebang: "#!/bin/bash" +prepend_text: "" +append_text: "" diff --git a/examples/config_tests/computer_localhost_tomato_config.yml b/examples/config_tests/computer_localhost_tomato_config.yml new file mode 100644 index 0000000..7e922d7 --- /dev/null +++ b/examples/config_tests/computer_localhost_tomato_config.yml @@ -0,0 +1,2 @@ +safe_interval: 0.0 +use_login_shell: true diff --git a/examples/cycling_sequence/README.md b/examples/cycling_sequence/README.md new file mode 100644 index 0000000..586d32f --- /dev/null +++ b/examples/cycling_sequence/README.md @@ -0,0 +1,167 @@ +# Cycling sequence workflow example + +## 1. Setup Container + +First you need to set up a conteinerized environment capable of running AiiDAlab +(or at least aiica-core, for the purposes of this test). +There are many ways to do this, from using +[aiidalab-launch](https://github.com/aiidalab/aiidalab-launch) +to setting up your own Docker image. +A compromise of ease and versatility I found useful is to make a `docker-compose.yml` file with the already +existing AiiDAlab images: + +``` +version: "3.9" + +services: + + aurora: + image: "aiidalab/aiidalab-docker-stack:22.8.1" + ports: + - "8888:8888" + volumes: + - "data:/home/aiida/data" + +volumes: + data: +``` + +Then you can start it up by running (in the same folder) `docker-compose up` +(note that the terminal will then be blocked due to the running docker instance). +Then you can access the container via de jupyter lab interface (as the aiida user) +or running `docker exec -it ${CONTAINER_ID} /bin/bash` (as root). + +Note that the data volume will be created automatically, and will serve the purpose +of data persistence. +Also, depending on how you will be working, you will need to remember to set up +your internal configuration to work with github (for example, some tools like vscode +plugins may allow you to open the files inside the container while using an external +github connection / credentials). +For me, this is just copying the ssh key inside the .ssh folder in the container and +then setting up the config file: + +``` +Host github.com + User git + HostName github.com + IdentityFile ~/.ssh/github_key + IdentitiesOnly yes +``` + +## 2. Install Tomato + +It is better to install tomato in its own separate environment. +If you are using a container with conda (suche as the one used before), just run: + +``` +conda create --name tomato +conda activate tomato +``` + +You may also specify python version at the end of the first line (`python=3.9`). +Then you will have to install tomato, which can normally be done via PyPi package +(`pip install tomato`), see the +[docs](https://dgbowl.github.io/tomato/master/installation.html) +for more information. +However, if the PyPi package is broken, you may need to install it manually by +cloning the [git repository](https://github.com/dgbowl/tomato) and running + +``` +pip install -e . +``` + +You can start the tomato server by running `tomato -vv` (this will block the window). +You can check the available pipelines using: + +``` +$ ketchup status + +pipeline ready jobid (PID) sampleid +=================================================================== +dummy-10 no None None +dummy-5 no None None +``` + +You can set up a sample with the following commands: + +``` +$ ketchup load commercial-10 dummy-10 +$ ketchup ready dummy-10 +$ ketchup status + +pipeline ready jobid (PID) sampleid +=================================================================== +dummy-10 yes None commercial-10 +dummy-5 no None None +``` + +## 3. Install AiiDA + +Although AiiDA can normally be pip installed by itself or as a dependency of the +aurora-plugin and the aurora-app, the code is currently relying on a custom +adaptation of the code that is not included in the core release. +Although this should only affect the connection to windows clusters, it may be +more convenient to just set up that version anyways: + +``` +git clone https://github.com/lorisercole/aiida-core.git +cd aiida-core +git checkout windows +pip install -e . +``` + +NOTE: in some computers python still uses the `aiida-core` installed in the root +directory, so you may need to log in as root and install the package system-wide. + +## 4. Install Aurora + +The aurora plugin can also be installed via PyPi package +(and if you are installing the +[aiidalab-auror app](https://github.com/epfl-theos/aiidalab-aurora), +the plugin will be installed as a dependency). +However, if you are going to develop the package, you will probably prefer to clone +it from the [github repo](https://github.com/epfl-theos/aiida-aurora), or your own +fork. +In this case you may also want to install with `pip install -e .[pre-commit]` and +then do `pre-commit install` too. + +Note that there is at least one line in the code that you may need to adapt in +order to run: the location of the ketchup executable. + +``` +diff --git a/aiida_aurora/scheduler.py b/aiida_aurora/scheduler.py +index dd42421..476a393 100644 +--- a/aiida_aurora/scheduler.py ++++ b/aiida_aurora/scheduler.py +@@ -102,7 +102,7 @@ class TomatoScheduler(Scheduler): +- KETCHUP = "ketchup" ++ KETCHUP = "/home/aiida/.conda/envs/tomato/bin/ketchup" +``` + +Once the package is installed, you can run: + +``` +reentry scan +verdi daemon restart --reset +``` + +And afterwards you should be able to check the new plugins by running: + +``` +verdi plugin list aiida.schedulers +verdi plugin list aiida.calculations +``` + +## 5. Database Setup + +Finally, you need to set up the base computer and code to run. +The settings for these devices can be found in the `aiida-aurora/examples/config_tests` +folder, and can be installed by runnig: + +``` +verdi computer setup --config computer_localhost_tomato.yml + +verdi computer configure local --config computer_localhost_tomato_config.yml localhost-tomato + +verdi code setup --config code_ketchup-0.2rc2.yml +``` diff --git a/examples/cycling_sequence/test_cycling_sequence.py b/examples/cycling_sequence/test_cycling_sequence.py new file mode 100644 index 0000000..d48a098 --- /dev/null +++ b/examples/cycling_sequence/test_cycling_sequence.py @@ -0,0 +1,110 @@ +#################################################################################################### +"""Script to use as a template for the submission of the cycling sequence workchain.""" +#################################################################################################### +import datetime + +from aiida import orm +from aiida.engine import calcfunction, submit +from aiida.plugins import DataFactory, WorkflowFactory + +CyclingSpecsData = DataFactory('aurora.cyclingspecs') +BatterySampleData = DataFactory('aurora.batterysample') +TomatoSettingsData = DataFactory('aurora.tomatosettings') + + +@calcfunction +def generate_test_inputs(): + """Generate the inputs.""" + nodes_dict = {} + + nodes_dict['sample'] = BatterySampleData({ + 'manufacturer': + 'fake_maufacturer', + 'composition': + dict(description='C|E|A'), + 'form_factor': + 'fake_form', + 'capacity': + dict(nominal=1.0, units='Ah'), + 'battery_id': + 666, + 'metadata': + dict( + name='commercial-10', + creation_datetime=datetime.datetime.now(tz=datetime.timezone.utc), + creation_process='This is a fake battery for testing purposes.' + ) + }) + + BASELINE_SPECS = { + 'method': [{ + 'name': 'DUMMY_SEQUENTIAL_1', + 'device': 'worker', + 'technique': 'sequential', + 'parameters': { + 'time': { + 'label': 'Time:', + 'units': 's', + 'value': 60.0, + 'required': True, + 'description': '', + 'default_value': 100.0 + }, + 'delay': { + 'label': 'Delay:', + 'units': 's', + 'value': 1.0, + 'required': True, + 'description': '', + 'default_value': 1.0 + } + }, + 'short_name': 'DUMMY_SEQUENTIAL', + 'description': 'Dummy worker - sequential numbers' + }] + } + + step1_tech = dict(BASELINE_SPECS) + nodes_dict['step1_tech'] = CyclingSpecsData(step1_tech) + step2_tech = dict(BASELINE_SPECS) + step2_tech['method'][0]['parameters']['time']['value'] = 30.0 + nodes_dict['step2_tech'] = CyclingSpecsData(step2_tech) + + BASELINE_SETTINGS = { + "output": { + "path": None, + "prefix": None + }, + "snapshot": None, + "verbosity": "INFO", + "unlock_when_done": True, + } + + nodes_dict['step1_setting'] = TomatoSettingsData(BASELINE_SETTINGS) + nodes_dict['step2_setting'] = TomatoSettingsData(BASELINE_SETTINGS) + + return nodes_dict + + +#################################################################################################### +# ACTUAL SUBMISSION + +WorkflowClass = WorkflowFactory('aurora.cycling_sequence') +workflow_builder = WorkflowClass.get_builder() +workflow_inputs = generate_test_inputs() + +workflow_builder.tomato_code = orm.load_code('ketchup-0.2rc2@localhost-tomato') + +workflow_builder.battery_sample = workflow_inputs['sample'] +workflow_builder.protocols = { + 'step1': workflow_inputs['step1_tech'], + 'step2': workflow_inputs['step2_tech'], +} +workflow_builder.control_settings = { + 'step1': workflow_inputs['step1_setting'], + 'step2': workflow_inputs['step2_setting'], +} + +workflow_node = submit(workflow_builder) + +#################################################################################################### diff --git a/pyproject.toml b/pyproject.toml index d95b243..552513c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -66,6 +66,9 @@ docs = [ "aurora.fake" = "aiida_aurora.calculations.fake:BatteryFakeExperiment" "aurora.cycler" = "aiida_aurora.calculations.cycler:BatteryCyclerExperiment" +[project.entry-points.'aiida.workflows'] +'aurora.cycling_sequence' = 'aiida_aurora.workflows.cycling_sequence:CyclingSequenceWorkChain' + [project.entry-points."aiida.calculations.monitors"] "aurora.monitors.capacity_threshold" = "aiida_aurora.monitors:monitor_capacity_threshold"