diff --git a/implementors/cookiecutter.json b/implementors/cookiecutter.json new file mode 100644 index 0000000..10508c9 --- /dev/null +++ b/implementors/cookiecutter.json @@ -0,0 +1,8 @@ +{ + "name": "site", + "root_name": "mimas_{{ cookiecutter.name.lower() }}", + "project_name": "{{ cookiecutter.name }} zocalo-mimas project", + "project_slug": "mimas_{{ cookiecutter.name.lower() }}", + "project_short_description": "{{ cookiecutter.name }} zocalo mimas local implementors", + "version": "0.1.0" +} diff --git a/implementors/{{cookiecutter.root_name}}/setup.cfg b/implementors/{{cookiecutter.root_name}}/setup.cfg new file mode 100644 index 0000000..6ed9d7a --- /dev/null +++ b/implementors/{{cookiecutter.root_name}}/setup.cfg @@ -0,0 +1,6 @@ +[flake8] +ignore = E501,W503,E203,W504,E251,E262,E265,E266,W291,W293 +exclude = docs + +[aliases] +test = pytest diff --git a/implementors/{{cookiecutter.root_name}}/setup.py b/implementors/{{cookiecutter.root_name}}/setup.py new file mode 100644 index 0000000..f61eae5 --- /dev/null +++ b/implementors/{{cookiecutter.root_name}}/setup.py @@ -0,0 +1,27 @@ +#!/usr/bin/env python + +from setuptools import setup, find_packages + +setup( + classifiers=[ + "Development Status :: 3 - Alpha", + "Intended Audience :: Developers", + "License :: OSI Approved :: BSD License", + "Natural Language :: English", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.7", + "Programming Language :: Python :: 3.8", + "Programming Language :: Python :: 3.9", + ], + description="{{cookiecutter.project_slug}}", + install_requires=["zocalo"], + license="BSD license", + include_package_data=True, + keywords="zocalo mimas", + name="{{cookiecutter.project_slug}}", + packages=find_packages(exclude=("tests",)), + python_requires=">=3.7", + setup_requires=[], + version="0.1.0", + zip_safe=False, +) diff --git a/implementors/{{cookiecutter.root_name}}/{{cookiecutter.project_slug}}/__init__.py b/implementors/{{cookiecutter.root_name}}/{{cookiecutter.project_slug}}/__init__.py new file mode 100644 index 0000000..3dc1f76 --- /dev/null +++ b/implementors/{{cookiecutter.root_name}}/{{cookiecutter.project_slug}}/__init__.py @@ -0,0 +1 @@ +__version__ = "0.1.0" diff --git a/implementors/{{cookiecutter.root_name}}/{{cookiecutter.project_slug}}/implementors/__init__.py b/implementors/{{cookiecutter.root_name}}/{{cookiecutter.project_slug}}/implementors/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/implementors/{{cookiecutter.root_name}}/{{cookiecutter.project_slug}}/implementors/tasks/__init__.py b/implementors/{{cookiecutter.root_name}}/{{cookiecutter.project_slug}}/implementors/tasks/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/implementors/{{cookiecutter.root_name}}/{{cookiecutter.project_slug}}/implementors/triggers/__init__.py b/implementors/{{cookiecutter.root_name}}/{{cookiecutter.project_slug}}/implementors/triggers/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/requirements_dev.txt b/requirements_dev.txt index e50d125..2a47ace 100644 --- a/requirements_dev.txt +++ b/requirements_dev.txt @@ -5,3 +5,5 @@ pytest-cov==2.12.1 pytest==6.2.4 setuptools==58.0.4 workflows==2.12 +gemmi==0.4.9 +dataclasses==0.8; python_version < '3.7' diff --git a/setup.cfg b/setup.cfg index 7bf12c7..94337a8 100644 --- a/setup.cfg +++ b/setup.cfg @@ -31,6 +31,8 @@ install_requires = marshmallow setuptools workflows>=2.7 + dataclasses;python_version<'3.7' + gemmi packages = find: package_dir = =src @@ -42,6 +44,7 @@ console_scripts = zocalo.go = zocalo.cli.go:run zocalo.service = zocalo.service:start_service zocalo.wrap = zocalo.cli.wrap:run + zocalo.mimas = zocalo.cli.mimas:run libtbx.dispatcher.script = zocalo.go = zocalo.go zocalo.service = zocalo.service @@ -50,13 +53,18 @@ libtbx.precommit = zocalo = zocalo workflows.services = Schlockmeister = zocalo.service.schlockmeister:Schlockmeister + Mimas = zocalo.service.mimas:Mimas + Trigger = zocalo.service.trigger:Trigger zocalo.configuration.plugins = graylog = zocalo.configuration.plugin_graylog:Graylog rabbitmqapi = zocalo.configuration.plugin_rabbitmqapi:RabbitAPI storage = zocalo.configuration.plugin_storage:Storage jmx = zocalo.configuration.plugin_jmx:JMX + mimas = zocalo.configuration.plugin_mimas:Mimas zocalo.wrappers = dummy = zocalo.wrapper:DummyWrapper +zocalo.system_tests = + mimas = zocalo.system_test.tests.mimas:MimasService [options.packages.find] where = src diff --git a/src/zocalo/cli/mimas.py b/src/zocalo/cli/mimas.py new file mode 100644 index 0000000..aef8a74 --- /dev/null +++ b/src/zocalo/cli/mimas.py @@ -0,0 +1,219 @@ +""" Interrogates Mimas. Shows what would happen for a given datacollection ID""" +import os +import argparse +import errno +import logging + +import workflows.recipe +import zocalo.configuration + +try: + from zocalo.ispyb.filter import ispyb_filter, ispybtbx +except ImportError: + print("Error: zocalo.mimas requires ispyb to be installed") + exit(1) + +from zocalo.mimas.core import run as mimas_run +from zocalo.mimas.classes import ( + MimasEvent, + MimasScenario, + MimasExperimentType, + MimasRunStatus, + MimasDCClass, + MimasISPyBSweep, + MimasISPyBUnitCell, + MimasISPyBSpaceGroup, + MimasDetectorClass, + MimasRecipeInvocation, + MimasISPyBJobInvocation, + validate, + zocalo_command_line, +) + +logging.basicConfig(level=logging.WARNING) + + +_readable = { + MimasEvent.START: "start of data collection", + MimasEvent.END: "end of data collection", +} + + +def get_recipe_triggers(recipe_path, recipefile): + try: + with open(os.path.join(recipe_path, recipefile + ".json"), "r") as rcp: + named_recipe = workflows.recipe.Recipe(recipe=rcp.read()) + except ValueError as e: + raise ValueError("Error reading recipe '%s': %s", recipefile, str(e)) + except IOError as e: + if e.errno == errno.ENOENT: + raise ValueError( + f"Message references non-existing recipe {recipefile}. Recipe path is {recipe_path}", + ) + raise + + triggers = [] + for step_id, step in named_recipe.recipe.items(): + if isinstance(step, dict): + if step["queue"] == "trigger": + triggers.append( + { + "target": step["parameters"]["target"], + "comment": step["parameters"]["comment"], + "automatic": step["parameters"]["automatic"], + } + ) + + return triggers + + +def get_scenarios(dcid): + _, ispyb_info = ispyb_filter({}, {"ispyb_dcid": dcid}) + + if len(ispyb_info.keys()) == 1: + exit() + + cell = ispyb_info.get("ispyb_unit_cell") + if cell: + cell = MimasISPyBUnitCell(*cell) + else: + cell = None + + spacegroup = ispyb_info.get("ispyb_space_group") + if spacegroup: + spacegroup = MimasISPyBSpaceGroup(spacegroup) + else: + spacegroup = None + + experimenttype = ispyb_info["ispyb_experimenttype"] + if not experimenttype or not isinstance(experimenttype, str): + print(f"Invalid Mimas request rejected (experimenttype = {experimenttype!r})") + exit() + + try: + experimenttype_safe = experimenttype.replace(" ", "_") + experimenttype_mimas = MimasExperimentType[experimenttype_safe.upper()] + except KeyError: + print(f"Invalid Mimas request (Experiment type = {experimenttype!r})") + experimenttype_mimas = MimasExperimentType.UNDEFINED + + dc_class = ispyb_info.get("ispyb_dc_class") + if dc_class and dc_class["grid"]: + dc_class_mimas = MimasDCClass.GRIDSCAN + elif dc_class and dc_class["screen"]: + dc_class_mimas = MimasDCClass.SCREENING + elif dc_class and dc_class["rotation"]: + dc_class_mimas = MimasDCClass.ROTATION + else: + dc_class_mimas = MimasDCClass.UNDEFINED + + run_status = ispyb_info["ispyb_dc_info"]["runStatus"].lower() + if "success" in run_status: + run_status_mimas = MimasRunStatus.SUCCESS + elif "fail" in run_status: + run_status_mimas = MimasRunStatus.FAILURE + else: + run_status_mimas = MimasRunStatus.UNKNOWN + + detectorclass = ( + MimasDetectorClass.EIGER + if ispyb_info["ispyb_detectorclass"] == "eiger" + else MimasDetectorClass.PILATUS + ) + scenarios = [] + for event in (MimasEvent.START, MimasEvent.END): + scenario = MimasScenario( + dcid=dcid, + dcclass=dc_class_mimas, + experimenttype=experimenttype_mimas, + event=event, + beamline=ispyb_info["ispyb_beamline"], + proposalcode=ispyb_info["ispyb_dc_info"].get("proposalCode"), + runstatus=run_status_mimas, + spacegroup=spacegroup, + unitcell=cell, + getsweepslistfromsamedcg=tuple( + MimasISPyBSweep(*sweep) + for sweep in ispyb_info.get("ispyb_related_sweeps", []) + ), + preferred_processing=ispyb_info.get("ispyb_preferred_processing"), + detectorclass=detectorclass, + ) + try: + validate(scenario) + except ValueError: + print( + f"Can not generate a valid Mimas scenario for {_readable.get(scenario.event)} {dcid}" + ) + raise + scenarios.append(scenario) + return scenarios + + +def run(args=None): + zc = zocalo.configuration.from_file() + zc.activate() + + parser = argparse.ArgumentParser(usage="zocalo.mimas [options] dcid") + parser.add_argument("dcids", type=int, nargs="+", help="Data collection ids") + parser.add_argument("-?", action="help", help=argparse.SUPPRESS) + parser.add_argument( + "--commands", + "-c", + action="store_true", + dest="show_commands", + default=False, + help="Show commands that would trigger the individual processing steps", + ) + zc.add_command_line_options(parser) + + args = parser.parse_args(args) + + i = ispybtbx() + + for dcid in args.dcids: + dc_info = i.get_dc_info(dcid) + + for scenario in get_scenarios(dcid): + actions = mimas_run(scenario, zc.mimas.get("implementors")) + print( + f"At the {_readable.get(scenario.event)} {dcid} ({dc_info['visit']} on {dc_info['beamLineName']}):" + ) + for a in sorted(actions, key=lambda a: str(type(a)) + " " + a.recipe): + try: + validate(a) + except ValueError: + print( + f"Mimas scenario for dcid {dcid}, {scenario.event} returned invalid action {a!r}" + ) + raise + + if isinstance(a, MimasRecipeInvocation): + if args.show_commands: + print(" - " + zocalo_command_line(a)) + else: + print(f" - for dcid {a.dcid} call recipe {a.recipe}") + elif isinstance(a, MimasISPyBJobInvocation): + full_recipe = f"ispyb-{a.recipe}" + if args.show_commands: + print(" - " + zocalo_command_line(a)) + else: + print( + f" - create ISPyB job for dcid {a.dcid} named {a.displayname!r} with recipe '{full_recipe}' (autostart={a.autostart})" + ) + + triggers = get_recipe_triggers( + zc.storage.get("zocalo.recipe_directory"), full_recipe + ) + if triggers: + print(" Then trigger: ") + for trigger in triggers: + print( + f" - {trigger['target']}: {trigger['comment']} (autostart={a.autostart})" + ) + + else: + raise RuntimeError(f"Encountered unknown action {a!r}") + if not actions: + print(" - do nothing") + print() diff --git a/src/zocalo/configuration/plugin_mimas.py b/src/zocalo/configuration/plugin_mimas.py new file mode 100644 index 0000000..5bcba8c --- /dev/null +++ b/src/zocalo/configuration/plugin_mimas.py @@ -0,0 +1,12 @@ +from marshmallow import fields + +from zocalo.configuration import PluginSchema + + +class Mimas: + class Schema(PluginSchema): + implementors = fields.Str() + + @staticmethod + def activate(configuration): + return configuration diff --git a/src/zocalo/mimas/README.md b/src/zocalo/mimas/README.md new file mode 100644 index 0000000..0344672 --- /dev/null +++ b/src/zocalo/mimas/README.md @@ -0,0 +1,119 @@ +# Mimas: a zocalo decision making service + +Mimas includes the following services: + +* Mimas - the decision maker +* Trigger - a downstream trigger service (for starting jobs after another has finished) + +CLI tools: + +* zocalo.mimas - interrogate a specific dcid for the expected execution plan + +## Installation + +Cookiecutter the local implementation project: + +```python +pip install cookiecutter +cookiecutter https://github.com/DiamondLightSource/python-zocalo/implementors +cd mimas_ +pip install -e . +``` + +Then set the new package as the mimas implementors: + +```yaml +mimas: + plugin: mimas + implementors: mimas_ +``` + +## Tasks + +Multiple task files can be created in the implementors.tasks module. These tell mimas what to do, an example is provided. All classes should inherit from `Tasks`. If `beamline` is specified this class will only be run if the scenario specifies the beamline. `event` can also be specified to run this class only on a specific event. Classes should be named the same as the file. The `run` method should return a list of `MimasRecipeInvocation` or `MimasISPyBJobInvocation`. + +```python +file: mimas_/implementors/tasks/mybeamline.py + +from zocalo.mimas.tasks import Tasks +from zocalo.mimas.classes import ( + MimasExperimentType, + MimasRecipeInvocation, + MimasISPyBJobInvocation, + MimasEvent, +) + + +class mybeamline(Tasks): + beamline = "bl" + event = MimasEvent.START + + def run(self, scenario): + tasks = [] + + if scenario.experimenttype == MimasExperimentType.ENERGY_SCAN: + tasks.append( + MimasRecipeInvocation(dcid=scenario.dcid, recipe="myrecipe") + ) + + return tasks +``` + +## Triggers + +Downstream triggers are defined in a similar way to tasks. + +A recipe as follows: +```yaml +{ + "1": { + "service": "Trigger", + "queue": "trigger", + "parameters": { + "target": "test", + ... + } + }, + "start": [[1, []]] +} +``` + +will cause the trigger service to search for a module called `test` in the `implementors.triggers` package. The class should be named the same as the file. The `run` function should return an instance of `TriggerResponse`. + +```python +file: mimas_/implementors/triggers/test.py + +from zocalo.trigger import Trigger, TriggerResponse + + +class test(Trigger): + name = "TestTrigger" + + def run(self): + self._jobid = 12 + self._trigger_job({"testid": 42}) + + return TriggerResponse(success=True, return_value=self._jobid) +``` + +## System Test + +Required configuration for associated system test + +```yaml +storage: + plugin: storage + + system_tests: + mimas: + event: start + beamline: bl + experimenttype: Energy scan + proposalcode: bl + dc_class: + grid: None + screen: None + rotation: None + run_status: Successful + expected_recipe: exafs-qa +``` diff --git a/src/zocalo/mimas/__init__.py b/src/zocalo/mimas/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/zocalo/mimas/classes.py b/src/zocalo/mimas/classes.py new file mode 100644 index 0000000..d3d4f12 --- /dev/null +++ b/src/zocalo/mimas/classes.py @@ -0,0 +1,395 @@ +import dataclasses +import enum +import functools +import numbers +from typing import Tuple + +import gemmi + + +MimasExperimentType = enum.Enum( + "MimasExperimentType", + "OSC MESH SAD ENERGY_SCAN XRF_MAP XRF_MAP_XAS XRF_SPECTRUM UNDEFINED", +) + +# enum.Enum( +# "SAD", +# "SAD - Inverse Beam", +# "OSC", +# "Collect - Multiwedge", +# "MAD", +# "Helical", +# "Multi-positional", +# "Mesh", +# "Burn", +# "MAD - Inverse Beam", +# "Characterization", +# "Dehydration", +# "tomo", +# "experiment", +# "EM", +# "PDF", +# "PDF+Bragg", +# "Bragg", +# "single particle", +# "Serial Fixed", +# "Serial Jet", +# "Standard", +# "Time Resolved", +# "Diamond Anvil High Pressure", +# "Custom", +# "XRF map", +# "Energy scan", +# "XRF spectrum", +# "XRF map xas", +# ) + +MimasDCClass = enum.Enum("MimasDCClass", "GRIDSCAN ROTATION SCREENING UNDEFINED") + +MimasDetectorClass = enum.Enum("MimasDetectorClass", "PILATUS EIGER") + +MimasEvent = enum.Enum("MimasEvent", "START END") + +MimasRunStatus = enum.Enum("MimasRunStatus", "SUCCESS FAILURE UNKNOWN") + + +@dataclasses.dataclass(frozen=True) +class MimasISPyBUnitCell: + a: float + b: float + c: float + alpha: float + beta: float + gamma: float + + @property + def string(self): + return f"{self.a},{self.b},{self.c},{self.alpha},{self.beta},{self.gamma}" + + +@dataclasses.dataclass(frozen=True) +class MimasISPyBSpaceGroup: + symbol: str + + @property + def string(self): + return gemmi.SpaceGroup(self.symbol).hm.replace(" ", "") + + +@dataclasses.dataclass(frozen=True) +class MimasISPyBSweep: + dcid: int + start: int + end: int + + +@dataclasses.dataclass(frozen=True) +class MimasScenario: + dcid: int + experimenttype: MimasExperimentType + event: MimasEvent + beamline: str + runstatus: MimasRunStatus + dcclass: MimasDCClass = None + spacegroup: MimasISPyBSpaceGroup = None + unitcell: MimasISPyBUnitCell = None + getsweepslistfromsamedcg: Tuple[MimasISPyBSweep] = () + preferred_processing: str = None + detectorclass: MimasDetectorClass = None + proposalcode: str = None + + +@dataclasses.dataclass(frozen=True) +class MimasISPyBParameter: + key: str + value: str + + +@dataclasses.dataclass(frozen=True) +class MimasISPyBTriggerVariable: + key: str + value: str + + +@dataclasses.dataclass(frozen=True) +class MimasISPyBJobInvocation: + dcid: int + autostart: bool + recipe: str + source: str + comment: str = "" + displayname: str = "" + parameters: Tuple[MimasISPyBParameter] = () + sweeps: Tuple[MimasISPyBSweep] = () + triggervariables: Tuple[MimasISPyBTriggerVariable] = () + + +@dataclasses.dataclass(frozen=True) +class MimasRecipeInvocation: + dcid: int + recipe: str + + +@functools.singledispatch +def validate(mimasobject, expectedtype=None): + """ + A generic validation function that can (recursively) validate any Mimas* + object for consistency and semantic correctness. + If any issues are found a ValueError is raised, returns None otherwise. + """ + raise ValueError(f"{mimasobject!r} is not a known Mimas object") + + +@validate.register(MimasScenario) +def _(mimasobject: MimasScenario, expectedtype=None): + if expectedtype and not isinstance(mimasobject, expectedtype): + raise ValueError(f"{mimasobject!r} is not a {expectedtype}") + if type(mimasobject.dcid) != int: + raise ValueError(f"{mimasobject!r} has non-integer dcid") + validate(mimasobject.experimenttype, expectedtype=MimasExperimentType) + validate(mimasobject.runstatus, expectedtype=MimasRunStatus) + if mimasobject.dcclass is not None: + validate(mimasobject.dcclass, expectedtype=MimasDCClass) + validate(mimasobject.event, expectedtype=MimasEvent) + if type(mimasobject.getsweepslistfromsamedcg) not in (list, tuple): + raise ValueError( + f"{mimasobject!r} getsweepslistfromsamedcg must be a tuple, not {type(mimasobject.getsweepslistfromsamedcg)}" + ) + for sweep in mimasobject.getsweepslistfromsamedcg: + validate(sweep, expectedtype=MimasISPyBSweep) + if mimasobject.unitcell is not None: + validate(mimasobject.unitcell, expectedtype=MimasISPyBUnitCell) + if mimasobject.spacegroup is not None: + validate(mimasobject.spacegroup, expectedtype=MimasISPyBSpaceGroup) + if mimasobject.detectorclass is not None: + validate(mimasobject.detectorclass, expectedtype=MimasDetectorClass) + + +@validate.register(MimasExperimentType) +def _(mimasobject: MimasExperimentType, expectedtype=None): + if expectedtype and not isinstance(mimasobject, expectedtype): + raise ValueError(f"{mimasobject!r} is not a {expectedtype}") + + +@validate.register(MimasRunStatus) +def _(mimasobject: MimasRunStatus, expectedtype=None): + if expectedtype and not isinstance(mimasobject, expectedtype): + raise ValueError(f"{mimasobject!r} is not a {expectedtype}") + + +@validate.register(MimasEvent) +def _(mimasobject: MimasEvent, expectedtype=None): + if expectedtype and not isinstance(mimasobject, expectedtype): + raise ValueError(f"{mimasobject!r} is not a {expectedtype}") + + +@validate.register(MimasDCClass) +def _(mimasobject: MimasDCClass, expectedtype=None): + if expectedtype and not isinstance(mimasobject, expectedtype): + raise ValueError(f"{mimasobject!r} is not a {expectedtype}") + + +@validate.register(MimasDetectorClass) +def _(mimasobject: MimasDetectorClass, expectedtype=None): + if expectedtype and not isinstance(mimasobject, expectedtype): + raise ValueError(f"{mimasobject!r} is not a {expectedtype}") + + +@validate.register(MimasRecipeInvocation) +def _(mimasobject: MimasRecipeInvocation, expectedtype=None): + if expectedtype and not isinstance(mimasobject, expectedtype): + raise ValueError(f"{mimasobject!r} is not a {expectedtype}") + if type(mimasobject.dcid) != int: + raise ValueError(f"{mimasobject!r} has non-integer dcid") + if type(mimasobject.recipe) != str: + raise ValueError(f"{mimasobject!r} has non-string recipe") + if not mimasobject.recipe: + raise ValueError(f"{mimasobject!r} has empty recipe string") + + +@validate.register(MimasISPyBJobInvocation) +def _(mimasobject: MimasISPyBJobInvocation, expectedtype=None): + if expectedtype and not isinstance(mimasobject, expectedtype): + raise ValueError(f"{mimasobject!r} is not a {expectedtype}") + if type(mimasobject.dcid) != int: + raise ValueError(f"{mimasobject!r} has non-integer dcid") + if mimasobject.autostart not in (True, False): + raise ValueError(f"{mimasobject!r} has invalid autostart property") + if type(mimasobject.parameters) not in (list, tuple): + raise ValueError( + f"{mimasobject!r} parameters must be a tuple, not {type(mimasobject.parameters)}" + ) + for parameter in mimasobject.parameters: + validate(parameter, expectedtype=MimasISPyBParameter) + if type(mimasobject.recipe) != str: + raise ValueError(f"{mimasobject!r} has non-string recipe") + if not mimasobject.recipe: + raise ValueError(f"{mimasobject!r} has empty recipe string") + if type(mimasobject.sweeps) not in (list, tuple): + raise ValueError( + f"{mimasobject!r} sweeps must be a tuple, not {type(mimasobject.sweeps)}" + ) + for sweep in mimasobject.sweeps: + validate(sweep, expectedtype=MimasISPyBSweep) + + +@validate.register(MimasISPyBParameter) +def _(mimasobject: MimasISPyBParameter, expectedtype=None): + if expectedtype and not isinstance(mimasobject, expectedtype): + raise ValueError(f"{mimasobject!r} is not a {expectedtype}") + if type(mimasobject.key) != str: + raise ValueError(f"{mimasobject!r} has non-string key") + if not mimasobject.key: + raise ValueError(f"{mimasobject!r} has an empty key") + if type(mimasobject.value) != str: + raise ValueError( + f"{mimasobject!r} value must be a string, not {type(mimasobject.value)}" + ) + + +@validate.register(MimasISPyBSweep) +def _(mimasobject: MimasISPyBSweep, expectedtype=None): + if expectedtype and not isinstance(mimasobject, expectedtype): + raise ValueError(f"{mimasobject!r} is not a {expectedtype}") + if type(mimasobject.dcid) != int: + raise ValueError(f"{mimasobject!r} has non-integer dcid") + if mimasobject.dcid <= 0: + raise ValueError(f"{mimasobject!r} has an invalid dcid") + if type(mimasobject.start) != int: + raise ValueError(f"{mimasobject!r} has non-integer start image") + if mimasobject.start <= 0: + raise ValueError(f"{mimasobject!r} has an invalid start image") + if type(mimasobject.end) != int: + raise ValueError(f"{mimasobject!r} has non-integer end image") + if mimasobject.end < mimasobject.start: + raise ValueError(f"{mimasobject!r} has an invalid end image") + + +@validate.register(MimasISPyBUnitCell) +def _(mimasobject: MimasISPyBUnitCell, expectedtype=None): + if expectedtype and not isinstance(mimasobject, expectedtype): + raise ValueError(f"{mimasobject!r} is not a {expectedtype}") + if not isinstance(mimasobject.a, numbers.Real) or mimasobject.a <= 0: + raise ValueError(f"{mimasobject!r} has invalid length a") + if not isinstance(mimasobject.b, numbers.Real) or mimasobject.b <= 0: + raise ValueError(f"{mimasobject!r} has invalid length b") + if not isinstance(mimasobject.c, numbers.Real) or mimasobject.c <= 0: + raise ValueError(f"{mimasobject!r} has invalid length c") + if ( + not isinstance(mimasobject.alpha, numbers.Real) + or not 0 < mimasobject.alpha < 180 + ): + raise ValueError(f"{mimasobject!r} has invalid angle alpha") + if not isinstance(mimasobject.beta, numbers.Real) or not 0 < mimasobject.beta < 180: + raise ValueError(f"{mimasobject!r} has invalid angle beta") + if ( + not isinstance(mimasobject.gamma, numbers.Real) + or not 0 < mimasobject.gamma < 180 + ): + raise ValueError(f"{mimasobject!r} has invalid angle gamma") + + +@validate.register(MimasISPyBSpaceGroup) +def _(mimasobject: MimasISPyBSpaceGroup, expectedtype=None): + if expectedtype and not isinstance(mimasobject, expectedtype): + raise ValueError(f"{mimasobject!r} is not a {expectedtype}") + gemmi.SpaceGroup(mimasobject.symbol) + + +@functools.singledispatch +def zocalo_message(mimasobject): + """ + A generic function that (recursively) transforms any Mimas* object + into serializable objects that can be sent via zocalo. + If any issues are found a ValueError is raised. + """ + if isinstance(mimasobject, (bool, int, float, str, type(None))): + # trivial base types + return mimasobject + raise ValueError(f"{mimasobject!r} is not a known Mimas object") + + +@zocalo_message.register(MimasRecipeInvocation) +def _(mimasobject: MimasRecipeInvocation): + return { + "recipes": [mimasobject.recipe], + "parameters": {"ispyb_dcid": mimasobject.dcid}, + } + + +@zocalo_message.register(MimasISPyBJobInvocation) +def _(mimasobject: MimasISPyBJobInvocation): + return dataclasses.asdict(mimasobject) + + +@zocalo_message.register(MimasISPyBSweep) +def _(mimasobject: MimasISPyBSweep): + return dataclasses.asdict(mimasobject) + + +@zocalo_message.register(MimasISPyBParameter) +def _(mimasobject: MimasISPyBParameter): + return dataclasses.asdict(mimasobject) + + +@zocalo_message.register(MimasISPyBUnitCell) +def _(mimasobject: MimasISPyBUnitCell): + return dataclasses.astuple(mimasobject) + + +@zocalo_message.register(MimasISPyBSpaceGroup) +def _(mimasobject: MimasISPyBSpaceGroup): + return mimasobject.string + + +@zocalo_message.register(list) +def _(list_: list): + return [zocalo_message(element) for element in list_] + + +@zocalo_message.register(tuple) +def _(tuple_: tuple): + return tuple(zocalo_message(element) for element in tuple_) + + +@functools.singledispatch +def zocalo_command_line(mimasobject): + """ + Return the command line equivalent to execute a given Mimas* object + """ + raise ValueError(f"{mimasobject!r} is not a known Mimas object") + + +@zocalo_command_line.register(MimasRecipeInvocation) +def _(mimasobject: MimasRecipeInvocation): + return f"zocalo.go -r {mimasobject.recipe} {mimasobject.dcid}" + + +@zocalo_command_line.register(MimasISPyBJobInvocation) +def _(mimasobject: MimasISPyBJobInvocation): + if mimasobject.comment: + comment = f"--comment={mimasobject.comment!r} " + else: + comment = "" + if mimasobject.displayname: + displayname = f"--display={mimasobject.displayname!r} " + else: + displayname = "" + parameters = " ".join( + f"--add-param={p.key}:{p.value}" for p in mimasobject.parameters + ) + sweeps = " ".join( + f"--add-sweep={s.dcid}:{s.start}:{s.end}" for s in mimasobject.sweeps + ) + if mimasobject.autostart: + trigger = "--trigger " + else: + trigger = "" + triggervars = " ".join( + f"--trigger-variable={tv.key}:{tv.value}" for tv in mimasobject.triggervariables + ) + + return ( + f"ispyb.job --new --dcid={mimasobject.dcid} --source={mimasobject.source} --recipe={mimasobject.recipe} " + f"{sweeps} {parameters} {displayname}{comment}{trigger}{triggervars}" + ) diff --git a/src/zocalo/mimas/core.py b/src/zocalo/mimas/core.py new file mode 100644 index 0000000..3331147 --- /dev/null +++ b/src/zocalo/mimas/core.py @@ -0,0 +1,58 @@ +import importlib +import pkgutil +import logging +from typing import List, Union + +from zocalo.mimas.classes import ( + MimasScenario, + MimasRecipeInvocation, + MimasISPyBJobInvocation, +) +from zocalo.mimas.tasks import Tasks + +logger = logging.getLogger(__name__) + + +def load(mod_file: str, cls_name: str) -> Tasks: + mod = importlib.import_module(mod_file) + mod = importlib.reload(mod) + return getattr(mod, cls_name) + + +def run( + scenario: MimasScenario, + implementors=None, +) -> List[Union[MimasRecipeInvocation, MimasISPyBJobInvocation]]: + tasks = [] + + if implementors is None: + implementors = "zocalo.mimas" + + mod = importlib.import_module(f"{implementors}.implementors.tasks") + mod = importlib.reload(mod) + + classes = [] + for importer, modname, ispkg in pkgutil.walk_packages( + path=mod.__path__, prefix=mod.__name__ + ".", onerror=lambda x: None + ): + try: + class_name = modname.split(".")[-1] + classes.append(load(modname, class_name)) + except AttributeError: + logger.error( + f"Implementor for '{modname}' module does not contain class '{class_name}'" + ) + except Exception: + logger.exception(f"Could not load mimas task class {mod}") + + if not classes: + logger.warning("No processing classes found") + + for cls in classes: + try: + instance = cls() + tasks.extend(instance.do_run(scenario)) + except Exception: + logger.exception(f"Could not run mimas task class {cls.__name__}") + + return tasks diff --git a/src/zocalo/mimas/implementors/__init__.py b/src/zocalo/mimas/implementors/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/zocalo/mimas/implementors/tasks/__init__.py b/src/zocalo/mimas/implementors/tasks/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/zocalo/mimas/implementors/tasks/bl.py b/src/zocalo/mimas/implementors/tasks/bl.py new file mode 100644 index 0000000..115ccc8 --- /dev/null +++ b/src/zocalo/mimas/implementors/tasks/bl.py @@ -0,0 +1,34 @@ +from zocalo.mimas.tasks import Tasks +from zocalo.mimas.classes import ( + MimasExperimentType, + MimasRecipeInvocation, + MimasISPyBJobInvocation, + MimasEvent, +) + + +class bl(Tasks): + beamline = "bl" + + def run(self, scenario): + tasks = [] + + if scenario.event == MimasEvent.START: + if scenario.experimenttype == MimasExperimentType.ENERGY_SCAN: + tasks.append( + MimasRecipeInvocation(dcid=scenario.dcid, recipe="exafs-qa") + ) + + if scenario.event == MimasEvent.END: + if scenario.experimenttype == MimasExperimentType.XRF_MAP: + tasks.append( + MimasISPyBJobInvocation( + dcid=scenario.dcid, + autostart=True, + displayname="PyMCA Fitter", + recipe="pymca", + source="automatic", + ) + ) + + return tasks diff --git a/src/zocalo/mimas/implementors/triggers/__init__.py b/src/zocalo/mimas/implementors/triggers/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/zocalo/mimas/implementors/triggers/test.py b/src/zocalo/mimas/implementors/triggers/test.py new file mode 100644 index 0000000..159e3f0 --- /dev/null +++ b/src/zocalo/mimas/implementors/triggers/test.py @@ -0,0 +1,17 @@ +from zocalo.trigger import Trigger, TriggerResponse + + +class test(Trigger): + name = "TestTrigger" + + def run(self): + self._jobid = 12 + # self._add_job("test processing", "myrecipe") + # params = { + # "param1": 42 + # } + # self._add_parameters(params) + + self._trigger_job({"testid": 42}) + + return TriggerResponse(success=True, return_value=self._jobid) diff --git a/src/zocalo/mimas/recipes/mimas.json b/src/zocalo/mimas/recipes/mimas.json new file mode 100644 index 0000000..fdb5d11 --- /dev/null +++ b/src/zocalo/mimas/recipes/mimas.json @@ -0,0 +1,51 @@ +{ + "1": [ + "Recipe for the business logic handling of a data collection start and end event" + ], + "1": { + "service": "Mimas Business Logic", + "queue": "mimas", + "parameters": { + "dcid": "{ispyb_dcid}", + "event": "{event}", + "beamline": "{ispyb_beamline}", + "detectorclass": "{ispyb_detectorclass}", + "experimenttype": "{ispyb_experimenttype}", + "preferred_processing": "{$REPLACE:ispyb_preferred_processing}", + "run_status": "{ispyb_dc_info[runStatus]}", + "space_group": "{$REPLACE:ispyb_space_group}", + "sweep_list": "{$REPLACE:ispyb_related_sweeps}", + "unit_cell": "{$REPLACE:ispyb_unit_cell}" + }, + "passthrough": { + "param": 1 + }, + "output": { + "dispatcher": 2, + "ispyb": 3 + } + }, + "2": { + "service": "Immediate recipe invocation", + "queue": "processing_recipe" + }, + "3": { + "service": "db connector", + "queue": "ispyb", + "parameters": { "ispyb_command": "create_ispyb_job" }, + "output": { + "trigger": 4, + "held": 5 + } + }, + "4": { + "service": "Trigger for created ispyb processing jobs", + "queue": "processing_recipe" + }, + "5": { + "service": "Collector for held ispyb processing jobs", + "queue": "mimas.held", + "output": 4 + }, + "start": [[1, []]] +} diff --git a/src/zocalo/mimas/recipes/trigger.json b/src/zocalo/mimas/recipes/trigger.json new file mode 100644 index 0000000..fdd1909 --- /dev/null +++ b/src/zocalo/mimas/recipes/trigger.json @@ -0,0 +1,14 @@ +{ + "1": ["Recipe to test Trigger service"], + "1": { + "service": "Trigger", + "queue": "trigger", + "parameters": { + "target": "test", + "dcid": "{ispyb_dcid}", + "comment": "Test process trigger", + "automatic": true + } + }, + "start": [[1, []]] +} diff --git a/src/zocalo/mimas/tasks.py b/src/zocalo/mimas/tasks.py new file mode 100644 index 0000000..8a62d4c --- /dev/null +++ b/src/zocalo/mimas/tasks.py @@ -0,0 +1,43 @@ +import logging +from typing import List, Union +from abc import abstractmethod, ABC + +from zocalo.mimas.classes import ( + MimasScenario, + MimasRecipeInvocation, + MimasISPyBJobInvocation, +) + +logger = logging.getLogger(__name__) + + +class Tasks(ABC): + beamline = None + event = None + enabled = True + + def do_run( + self, scenario: MimasScenario, + ) -> List[Union[MimasRecipeInvocation, MimasISPyBJobInvocation]]: + + if not self.enabled: + logger.info("Class is currently disabled") + return + + if self.beamline is not None: + if scenario.beamline != self.beamline: + logger.info("No beamline match, skipping tasks") + return [] + + if self.event is not None: + if scenario.event != self.event: + logger.info("No event match, skipping tasks") + return [] + + return self.run(scenario) + + @abstractmethod + def run( + self, scenario: MimasScenario, + ) -> List[Union[MimasRecipeInvocation, MimasISPyBJobInvocation]]: + pass diff --git a/src/zocalo/service/mimas.py b/src/zocalo/service/mimas.py new file mode 100644 index 0000000..23bbf4a --- /dev/null +++ b/src/zocalo/service/mimas.py @@ -0,0 +1,215 @@ +import workflows.recipe +from workflows.services.common_service import CommonService + +from zocalo.mimas.core import run +from zocalo.mimas.classes import ( + MimasEvent, + MimasExperimentType, + MimasRunStatus, + MimasDCClass, + MimasISPyBSweep, + MimasISPyBUnitCell, + MimasISPyBSpaceGroup, + MimasDetectorClass, + MimasScenario, + MimasRecipeInvocation, + MimasISPyBJobInvocation, + validate, + zocalo_message, +) + + +class Mimas(CommonService): + """ + Business logic component. Given a data collection id and some description + of event circumstances (beamline, experiment description, start or end of + scan) this service decides what recipes should be run with what settings. + """ + + # Human readable service name + _service_name = "Mimas" + + # Logger name + _logger_name = "zocalo.services.mimas" + + def initializing(self): + """Subscribe to the mimas queue. Received messages must be acknowledged.""" + self.log.info("Mimas starting") + + workflows.recipe.wrap_subscribe( + self._transport, + "mimas", + self.process, + acknowledgement=True, + log_extender=self.extend_log, + ) + + def _extract_scenario(self, step): + dcid = step.get("dcid") + if not dcid or not dcid.isnumeric(): + return f"Invalid Mimas request rejected (dcid = {dcid!r})" + + event = step.get("event") + if not isinstance(event, str): + event = repr(event) + try: + event = MimasEvent[event.upper()] + except KeyError: + return f"Invalid Mimas request rejected (Event = {event!r})" + + experimenttype = step.get("experimenttype") + if not experimenttype or not isinstance(experimenttype, str): + return ( + f"Invalid Mimas request rejected (experimenttype = {experimenttype!r})" + ) + + try: + experimenttype_safe = experimenttype.replace(" ", "_") + experimenttype_mimas = MimasExperimentType[experimenttype_safe.upper()] + except KeyError: + self.log.warning( + f"Invalid Mimas request (Experiment type = {experimenttype!r})" + ) + experimenttype_mimas = MimasExperimentType.UNDEFINED + + dc_class = step.get("dc_class") + if isinstance(dc_class, dict): + # legacy format + if dc_class["grid"]: + dc_class_mimas = MimasDCClass.GRIDSCAN + elif dc_class["screen"]: + dc_class_mimas = MimasDCClass.SCREENING + elif dc_class["rotation"]: + dc_class_mimas = MimasDCClass.ROTATION + else: + dc_class_mimas = MimasDCClass.UNDEFINED + else: + try: + dc_class_mimas = MimasDCClass[dc_class.upper()] + except (KeyError, AttributeError): + self.log.warning( + f"Invalid Mimas request (Data collection class = {dc_class!r})" + ) + dc_class_mimas = MimasDCClass.UNDEFINED + + run_status = step.get("run_status").lower() + if "success" in run_status: + run_status_mimas = MimasRunStatus.SUCCESS + elif "fail" in run_status: + run_status_mimas = MimasRunStatus.FAILURE + else: + run_status_mimas = MimasRunStatus.UNKNOWN + + sweep_list = tuple( + MimasISPyBSweep(*info) for info in (step.get("sweep_list") or []) + ) + + cell = step.get("unit_cell") + if cell: + cell = MimasISPyBUnitCell(*cell) + else: + cell = None + + spacegroup = step.get("space_group") + if spacegroup: + spacegroup = MimasISPyBSpaceGroup(spacegroup) + self.log.info(spacegroup) + try: + validate(spacegroup) + except ValueError: + self.log.warning( + f"Invalid spacegroup for dcid {dcid}: {spacegroup}", exc_info=True + ) + spacegroup = None + else: + spacegroup = None + + detectorclass = { + "eiger": MimasDetectorClass.EIGER, + "pilatus": MimasDetectorClass.PILATUS, + }.get(step.get("detectorclass", "").lower()) + + return MimasScenario( + dcid=int(dcid), + experimenttype=experimenttype_mimas, + dcclass=dc_class_mimas, + event=event, + beamline=step.get("beamline"), + proposalcode=step.get("proposalcode"), + runstatus=run_status_mimas, + spacegroup=spacegroup, + unitcell=cell, + getsweepslistfromsamedcg=sweep_list, + preferred_processing=step.get("preferred_processing"), + detectorclass=detectorclass, + ) + + def process(self, rw, header, message): + """Process an incoming event.""" + + # Pass incoming event information into Mimas scenario object + scenario = self._extract_scenario(rw.recipe_step["parameters"]) + if isinstance(scenario, str): + self.log.error(scenario) + rw.transport.nack(header) + return + + # Validate scenario + try: + validate(scenario, expectedtype=MimasScenario) + except ValueError: + self.log.error("Invalid Mimas request rejected", exc_info=True) + rw.transport.nack(header) + return + + txn = rw.transport.transaction_begin() + rw.set_default_channel("dispatcher") + + self.log.debug("Evaluating %r", scenario) + things_to_do = run(scenario, self.config.mimas.get("implementors")) + + passthrough = rw.recipe_step.get("passthrough", {}) + for ttd in things_to_do: + try: + validate( + ttd, expectedtype=(MimasRecipeInvocation, MimasISPyBJobInvocation), + ) + except ValueError: + self.log.error("Invalid Mimas response detected", exc_info=True) + rw.transport.nack(header) + rw.transport.transaction_abort(txn) + return + + self.log.info("Running: %r", ttd) + try: + ttd_zocalo = zocalo_message(ttd) + except ValueError: + self.log.error(f"Error zocalizing Mimas object {ttd!r}", exc_info=True) + rw.transport.nack(header) + rw.transport.transaction_abort(txn) + return + + if passthrough: + if isinstance(ttd, MimasRecipeInvocation): + for param, value in passthrough.items(): + ttd_zocalo["parameters"][param] = value + else: + # TODO: Must be a better way to deal with this + ttd_zocalo["triggervariables"] = list( + ttd_zocalo["triggervariables"] + ) + for param, value in passthrough.items(): + ttd_zocalo["triggervariables"].append( + {"key": param, "value": value} + ) + ttd_zocalo["triggervariables"] = tuple( + ttd_zocalo["triggervariables"] + ) + + if isinstance(ttd, MimasRecipeInvocation): + rw.send(ttd_zocalo, transaction=txn) + else: + rw.send_to("ispyb", ttd_zocalo, transaction=txn) + + rw.transport.ack(header, transaction=txn) + rw.transport.transaction_commit(txn) diff --git a/src/zocalo/service/trigger.py b/src/zocalo/service/trigger.py new file mode 100644 index 0000000..16f48bd --- /dev/null +++ b/src/zocalo/service/trigger.py @@ -0,0 +1,108 @@ +import importlib +import os +import pkgutil + +try: + import ispyb +except ImportError: + print("Error: Trigger service requires ispyb to be installed") + ispyb = None + + +import workflows.recipe +from workflows.services.common_service import CommonService + + +class Trigger(CommonService): + """A service that creates and runs downstream processing jobs.""" + + # Human readable service name + _service_name = "Trigger" + + # Logger name + _logger_name = "zocalo.services.trigger" + + def initializing(self): + """Subscribe to the trigger queue. Received messages must be acknowledged.""" + workflows.recipe.wrap_subscribe( + self._transport, + "trigger", + self.trigger, + acknowledgement=True, + log_extender=self.extend_log, + ) + + if ispyb: + self.ispyb = ispyb.open(os.environ["ISPYB_CREDENTIALS"]) + else: + self.ispyb = None + + def trigger(self, rw, header, message): + """Forward the trigger message to a specific trigger function.""" + # Extract trigger target from the recipe + params = rw.recipe_step.get("parameters", {}) + target = params.get("target") + if not target: + self.log.error("No trigger target defined in recipe") + rw.transport.nack(header) + return + + implementors = self.config.mimas.get("implementors", "zocalo.mimas") + + mod = importlib.import_module(f"{implementors}.implementors.triggers") + mod = importlib.reload(mod) + + modules = {} + for importer, modname, ispkg in pkgutil.walk_packages( + path=mod.__path__, prefix=mod.__name__ + ".", onerror=lambda x: None + ): + class_name = modname.split(".")[-1] + modules[class_name] = modname + + if target not in modules: + self.log.error("Unknown target %s defined in recipe", target) + rw.transport.nack(header) + return + + txn = rw.transport.transaction_begin() + rw.set_default_channel("output") + + def parameters(parameter, replace_variables=True): + if isinstance(message, dict): + base_value = message.get(parameter, params.get(parameter)) + else: + base_value = params.get(parameter) + if ( + not replace_variables + or not base_value + or not isinstance(base_value, str) + or "$" not in base_value + ): + return base_value + for key in rw.environment: + if "$" + key in base_value: + base_value = base_value.replace("$" + key, str(rw.environment[key])) + return base_value + + mod = importlib.import_module(modules[target]) + mod = importlib.reload(mod) + cls = getattr(mod, target) + instance = cls( + self.ispyb, + rw=rw, + header=header, + message=message, + parameters=parameters, + transaction=txn, + ) + + result = instance.run() + if result.success: + rw.send({"result": result.return_value}, transaction=txn) + rw.transport.ack(header, transaction=txn) + else: + rw.transport.transaction_abort(txn) + rw.transport.nack(header) + return + + rw.transport.transaction_commit(txn) diff --git a/src/zocalo/system_test/tests/mimas.py b/src/zocalo/system_test/tests/mimas.py new file mode 100644 index 0000000..df5f76c --- /dev/null +++ b/src/zocalo/system_test/tests/mimas.py @@ -0,0 +1,59 @@ +from workflows.recipe import Recipe + +from zocalo.system_test.common import CommonSystemTest + + +class MimasService(CommonSystemTest): + """Tests for the Mimas service.""" + + def test_mimas(self): + """Run a Mimas scenario""" + + params = self.config["mimas"] + recipe = { + 1: { + "service": "Mimas", + "queue": "mimas", + "parameters": { + "dcid": "1", + "event": params["event"], + "beamline": params["beamline"], + "experimenttype": params["experimenttype"], + "proposalcode": params["proposalcode"], + "dc_class": params["dc_class"], + "run_status": params["run_status"], + }, + "output": {"dispatcher": 2, "ispyb": 2}, + }, + 2: {"service": "System Test", "queue": "transient.system_test"}, + "start": [(1, [])], + } + recipe = Recipe(recipe) + recipe.validate() + + self.send_message( + queue=recipe[1]["queue"], + message={ + "payload": recipe["start"][0][1], + "recipe": recipe.recipe, + "recipe-pointer": "1", + "environment": {"ID": self.uuid}, + }, + headers={"workflows-recipe": True}, + ) + + self.expect_recipe_message( + environment={"ID": self.uuid}, + recipe=recipe, + recipe_path=[1], + recipe_pointer=2, + payload={ + "recipes": [params["expected_recipe"]], + "parameters": {"ispyb_dcid": 1}, + }, + timeout=30, + ) + + +if __name__ == "__main__": + MimasService().validate() diff --git a/src/zocalo/trigger.py b/src/zocalo/trigger.py new file mode 100644 index 0000000..c5f921f --- /dev/null +++ b/src/zocalo/trigger.py @@ -0,0 +1,79 @@ +import logging +from typing import NamedTuple, Any, List +from abc import abstractmethod, ABC + + +logger = logging.getLogger(__name__) + + +class TriggerResponse(NamedTuple): + success: bool + return_value: Any + + +class Trigger(ABC): + name = None + + _jobid = None + _dcid = None + + def __init__(self, ispyb, rw, header, message, parameters, transaction): + self._ispyb = ispyb + self._rw = rw + self._parameters = parameters + self._dcid = parameters("dcid") + + @property + def parameters(self): + return self._parameters + + @abstractmethod + def run(self) -> TriggerResponse: + pass + + def _add_job(self, display_name: str, recipe: str) -> int: + """Create an ISPyB ProcessingJob""" + jp = self._ispyb.mx_processing.get_job_params() + jp["automatic"] = bool(self.parameters("automatic")) + jp["comments"] = self.parameters("comment") + jp["datacollectionid"] = self._dcid + jp["display_name"] = display_name + jp["recipe"] = recipe + return self._ispyb.mx_processing.upsert_job(list(jp.values())) + + def _add_image_sweep(self, dcid: int, start_image: int, end_image: int) -> int: + """Create an ISPyB ProcessingJobImageSweep""" + if not self._jobid: + logger.error("No jobid defined") + return + + jisp = self._ispyb.mx_processing.get_job_image_sweep_params() + jisp["datacollectionid"] = dcid + jisp["start_image"] = start_image + jisp["end_image"] = end_image + jisp["job_id"] = self._jobid + return self._ispyb.mx_processing.upsert_job_image_sweep(list(jisp.values())) + + def _add_parameters(self, params: dict) -> List[int]: + """Add ISPyB ProcessingJobParameters""" + if not self._jobid: + logger.error("No jobid defined") + return + + jppids = [] + for key, value in params.items(): + jpp = self._ispyb.mx_processing.get_job_parameter_params() + jpp["job_id"] = self._jobid + jpp["parameter_key"] = key + jpp["parameter_value"] = value + jppids.append( + self._ispyb.mx_processing.upsert_job_parameter(list(jpp.values())) + ) + + return jppids + + def _trigger_job(self, parameters={}): + """Trigger a recipe with an ISPyB ProcessingJob id""" + message = {"recipes": [], "parameters": {"ispyb_process": self._jobid}} + message["parameters"].update(parameters) + self._rw.transport.send("processing_recipe", message) diff --git a/tests/test_mimas_scenario.py b/tests/test_mimas_scenario.py new file mode 100644 index 0000000..631d8ba --- /dev/null +++ b/tests/test_mimas_scenario.py @@ -0,0 +1,49 @@ +import functools + +from zocalo.mimas.core import run +from zocalo.mimas.classes import ( + MimasScenario, + MimasExperimentType, + MimasRunStatus, + MimasEvent, + zocalo_command_line, + validate, +) + + +def get_zocalo_commands(scenario): + commands = set() + actions = run(scenario) + for a in actions: + validate(a) + commands.add(zocalo_command_line(a).strip()) + return commands + + +def test_bl_start(): + dcid = 5918093 + scenario = functools.partial( + MimasScenario, + dcid=dcid, + experimenttype=MimasExperimentType.ENERGY_SCAN, + beamline="bl", + proposalcode="ev", + runstatus=MimasRunStatus.SUCCESS, + ) + assert get_zocalo_commands(scenario(event=MimasEvent.START)) == { + f"zocalo.go -r exafs-qa {dcid}", + } + + +def test_bl_end(): + dcid = 5918093 + scenario = functools.partial( + MimasScenario, + dcid=dcid, + experimenttype=MimasExperimentType.XRF_MAP, + beamline="bl", + runstatus=MimasRunStatus.SUCCESS, + ) + assert get_zocalo_commands(scenario(event=MimasEvent.END)) == { + f"ispyb.job --new --dcid={dcid} --source=automatic --recipe=pymca --display='PyMCA Fitter' --trigger", + } diff --git a/tests/test_mimas_service.py b/tests/test_mimas_service.py new file mode 100644 index 0000000..0c0b2a2 --- /dev/null +++ b/tests/test_mimas_service.py @@ -0,0 +1,83 @@ +from unittest import mock + +import workflows.transport.common_transport +from workflows.recipe.wrapper import RecipeWrapper + +import zocalo.configuration +from zocalo.service.mimas import Mimas + +sample_configuration = """ +version: 1 +mimas: + plugin: mimas +environments: + default: live + live: + mimas: mimas +""" + + +def generate_recipe_message(parameters, extra_params={}): + """Helper function for tests.""" + return { + "recipe": { + 1: { + "service": "mimas business logic", + "queue": "mimas", + "parameters": parameters, + "output": {"dispatcher": 2, "ispyb": 3}, + **extra_params, + }, + 2: {"service": "dispatcher", "queue": "transient.output"}, + 3: {"service": "ispyb", "queue": "transient.output"}, + "start": [(1, [])], + }, + "recipe-pointer": 1, + "recipe-path": [], + "environment": { + "ID": mock.sentinel.GUID, + "source": mock.sentinel.source, + "timestamp": mock.sentinel.timestamp, + }, + "payload": mock.sentinel.payload, + } + + +def test_mimas_service(mocker, tmp_path): + mock_transport = mock.Mock() + zc = zocalo.configuration.from_string(sample_configuration) + zc.activate() + mimas = Mimas(environment={"config": zc}) + setattr(mimas, "_transport", mock_transport) + mimas.initializing() + + dcid = 12345 + testid = 12 + + t = mock.create_autospec(workflows.transport.common_transport.CommonTransport) + m = generate_recipe_message( + parameters={ + "dcid": f"{dcid}", + "event": "start", + "beamline": "bl", + "experimenttype": "Energy scan", + "run_status": "Datacollection Successful", + }, + extra_params={"passthrough": {"testid": testid}}, + ) + rw = RecipeWrapper(message=m, transport=t) + send_to = mocker.spy(rw, "send_to") + + mimas.process(rw, {"some": "header"}, mock.sentinel.message) + send_to.assert_has_calls( + [ + mock.call( + "dispatcher", + { + "recipes": ["exafs-qa"], + "parameters": {"ispyb_dcid": dcid, "testid": testid}, + }, + transaction=mock.ANY, + ), + ] + ) diff --git a/tests/test_mimas_validation.py b/tests/test_mimas_validation.py new file mode 100644 index 0000000..be0f371 --- /dev/null +++ b/tests/test_mimas_validation.py @@ -0,0 +1,203 @@ +import dataclasses +import itertools + +import pytest + +from zocalo.mimas.classes import ( + MimasScenario, + MimasEvent, + MimasExperimentType, + MimasRunStatus, + MimasDetectorClass, + MimasISPyBUnitCell, + MimasISPyBSpaceGroup, + MimasISPyBSweep, + MimasISPyBParameter, + MimasRecipeInvocation, + MimasISPyBJobInvocation, + validate, +) + + +def test_validation_of_unknown_objects(): + for failing_object in ( + 5, + "string", + b"bytestring", + None, + True, + False, + [], + {}, + dict(), + ): + with pytest.raises(ValueError): + validate(failing_object) + + +def test_validation_of_scenario(): + valid_scenario = MimasScenario( + dcid=1, + experimenttype=MimasExperimentType.OSC, + event=MimasEvent.START, + beamline="i03", + unitcell=MimasISPyBUnitCell(a=10, b=10.0, c=10, alpha=90.0, beta=90, gamma=90), + spacegroup=MimasISPyBSpaceGroup("P41212"), + preferred_processing=None, + runstatus=MimasRunStatus.SUCCESS, + detectorclass=MimasDetectorClass.PILATUS, + ) + validate(valid_scenario) + + # replacing individual values should fail validation + for key, value in [ + ("dcid", "banana"), + ("experimenttype", None), + ("experimenttype", 1), + ("event", MimasRecipeInvocation(dcid=1, recipe="invalid")), + ("getsweepslistfromsamedcg", MimasRecipeInvocation(dcid=1, recipe="invalid"),), + ( + "getsweepslistfromsamedcg", + (MimasRecipeInvocation(dcid=1, recipe="invalid"),), + ), + ("getsweepslistfromsamedcg", MimasISPyBSweep(dcid=1, start=1, end=100),), + ("getsweepslistfromsamedcg", ""), + ("getsweepslistfromsamedcg", None), + ("unitcell", False), + ("unitcell", (10, 10, 10, 90, 90, 90)), + ("unitcell", MimasRecipeInvocation(dcid=1, recipe="invalid"),), + ("detectorclass", "ADSC"), + ]: + print(f"testing {key}: {value}") + invalid_scenario = dataclasses.replace(valid_scenario, **{key: value}) + with pytest.raises(ValueError): + validate(invalid_scenario) + + +def test_validation_of_recipe_invocation(): + valid_invocation = MimasRecipeInvocation(dcid=1, recipe="string") + validate(valid_invocation) + + # replacing individual values should fail validation + for key, value in [ + ("dcid", "banana"), + ("recipe", MimasRecipeInvocation(dcid=1, recipe="invalid")), + ("recipe", ""), + ("recipe", None), + ]: + print(f"testing {key}: {value}") + invalid_invocation = dataclasses.replace(valid_invocation, **{key: value}) + with pytest.raises(ValueError): + validate(invalid_invocation) + + +def test_validation_of_ispyb_invocation(): + valid_invocation = MimasISPyBJobInvocation( + dcid=1, + autostart=True, + comment="", + displayname="", + parameters=(MimasISPyBParameter(key="test", value="valid"),), + recipe="string", + source="automatic", + sweeps=(MimasISPyBSweep(dcid=1, start=1, end=100),), + triggervariables=(), + ) + validate(valid_invocation) + + # replacing individual values should fail validation + for key, value in [ + ("dcid", "banana"), + ("autostart", "banana"), + ("parameters", MimasRecipeInvocation(dcid=1, recipe="invalid")), + ("parameters", (MimasRecipeInvocation(dcid=1, recipe="invalid"),)), + ("parameters", MimasISPyBParameter(key="test", value="invalid")), + ("parameters", ""), + ("parameters", None), + ("recipe", MimasRecipeInvocation(dcid=1, recipe="invalid")), + ("recipe", ""), + ("recipe", None), + ("sweeps", MimasRecipeInvocation(dcid=1, recipe="invalid")), + ("sweeps", (MimasRecipeInvocation(dcid=1, recipe="invalid"),)), + ("sweeps", MimasISPyBSweep(dcid=1, start=1, end=100)), + ("sweeps", ""), + ("sweeps", None), + ]: + print(f"testing {key}: {value}") + invalid_invocation = dataclasses.replace(valid_invocation, **{key: value}) + with pytest.raises(ValueError): + validate(invalid_invocation) + + +def test_validation_of_ispyb_parameters(): + valid = MimasISPyBParameter(key="key", value="value") + validate(valid) + + # replacing individual values should fail validation + for key, value in [ + ("key", ""), + ("key", 5), + ("key", None), + ("key", False), + ("value", 5), + ("value", None), + ("value", False), + ]: + print(f"testing {key}: {value}") + invalid = dataclasses.replace(valid, **{key: value}) + with pytest.raises(ValueError): + validate(invalid) + + +def test_validation_of_ispyb_sweeps(): + valid = MimasISPyBSweep(dcid=1, start=10, end=100) + validate(valid) + + # replacing individual values should fail validation + for key, value in [ + ("dcid", ""), + ("dcid", "1"), + ("dcid", None), + ("dcid", 0), + ("start", ""), + ("start", "5"), + ("start", False), + ("start", -3), + ("end", ""), + ("end", "5"), + ("end", False), + ("end", -3), + ("end", 5), + ]: + print(f"testing {key}: {value}") + invalid = dataclasses.replace(valid, **{key: value}) + with pytest.raises(ValueError): + validate(invalid) + + +def test_validation_of_ispyb_unit_cells(): + valid = MimasISPyBUnitCell(a=10, b=11, c=12, alpha=90, beta=91.0, gamma=92) + validate(valid) + assert valid.string == "10,11,12,90,91.0,92" + + # replacing individual values should fail validation + for key, value in itertools.chain( + itertools.product( + ("a", "b", "c", "alpha", "beta", "gamma"), (-10, 0, "", False) + ), + [("alpha", 180), ("beta", 180), ("gamma", 180)], + ): + print(f"testing {key}: {value}") + invalid = dataclasses.replace(valid, **{key: value}) + with pytest.raises(ValueError): + validate(invalid) + + +def test_validataion_of_ispyb_space_groups(): + valid = MimasISPyBSpaceGroup(symbol="P 41 21 2") + validate(valid) + assert valid.string == "P41212" + + invalid = MimasISPyBSpaceGroup(symbol="P 5") + with pytest.raises(ValueError): + validate(invalid) diff --git a/tests/test_mimas_zocalo_messages.py b/tests/test_mimas_zocalo_messages.py new file mode 100644 index 0000000..bd0b7c4 --- /dev/null +++ b/tests/test_mimas_zocalo_messages.py @@ -0,0 +1,67 @@ +from zocalo.mimas.classes import ( + MimasRecipeInvocation, + MimasISPyBJobInvocation, + MimasISPyBParameter, + MimasISPyBSweep, + MimasISPyBUnitCell, + MimasISPyBSpaceGroup, + zocalo_message, +) + + +def test_transformation_of_recipe_invocation(): + valid_invocation = MimasRecipeInvocation(dcid=1, recipe="string") + zocdata = zocalo_message(valid_invocation) + assert isinstance(zocdata, dict) + + +def test_validation_of_ispyb_invocation(): + valid_invocation = MimasISPyBJobInvocation( + dcid=1, + autostart=True, + comment="", + displayname="", + parameters=(MimasISPyBParameter(key="test", value="valid"),), + recipe="string", + source="automatic", + sweeps=(MimasISPyBSweep(dcid=1, start=1, end=100),), + triggervariables=(), + ) + zocdata = zocalo_message(valid_invocation) + assert isinstance(zocdata, dict) + + +def test_validation_of_ispyb_parameters(): + valid = MimasISPyBParameter(key="key", value="value") + zocdata = zocalo_message(valid) + assert isinstance(zocdata, dict) + + zoclist = zocalo_message([valid, valid]) + assert isinstance(zoclist, list) + assert len(zoclist) == 2 + assert zoclist[0] == zocdata + assert zoclist[1] == zocdata + + +def test_validation_of_ispyb_sweeps(): + valid = MimasISPyBSweep(dcid=1, start=10, end=100) + zocdata = zocalo_message(valid) + assert isinstance(zocdata, dict) + + zoclist = zocalo_message((valid, valid)) + assert isinstance(zoclist, tuple) + assert len(zoclist) == 2 + assert zoclist[0] == zocdata + assert zoclist[1] == zocdata + + +def test_validation_of_ispyb_unit_cells(): + valid = MimasISPyBUnitCell(a=10, b=11, c=12, alpha=90, beta=91.0, gamma=92) + zocdata = zocalo_message(valid) + assert zocdata == (10, 11, 12, 90, 91.0, 92) + + +def test_validation_of_ispyb_space_groups(): + valid = MimasISPyBSpaceGroup(symbol="P 41 21 2") + zocdata = zocalo_message(valid) + assert zocdata == "P41212" diff --git a/tests/test_trigger.py b/tests/test_trigger.py new file mode 100644 index 0000000..3f1092f --- /dev/null +++ b/tests/test_trigger.py @@ -0,0 +1,84 @@ +from unittest import mock + +import workflows.transport.common_transport +from workflows.recipe.wrapper import RecipeWrapper + +import zocalo.configuration +from zocalo.service.trigger import Trigger + +sample_configuration = """ +version: 1 +mimas: + plugin: mimas +environments: + default: live + live: + mimas: mimas +""" + +def generate_recipe_message(parameters, extra_params={}): + """Helper function for tests.""" + return { + "recipe": { + 1: { + "service": "trigger", + "queue": "trigger", + "parameters": parameters, + **extra_params, + }, + "start": [(1, [])], + }, + "recipe-pointer": 1, + "recipe-path": [], + "environment": { + "ID": mock.sentinel.GUID, + "source": mock.sentinel.source, + "timestamp": mock.sentinel.timestamp, + }, + "payload": mock.sentinel.payload, + } + + +def test_trigger_service(mocker): + mock_transport = mock.Mock() + zc = zocalo.configuration.from_string(sample_configuration) + zc.activate() + trigger = Trigger(environment={"config": zc}) + setattr(trigger, "_transport", mock_transport) + trigger.initializing() + + t = mock.create_autospec(workflows.transport.common_transport.CommonTransport) + m = generate_recipe_message( + parameters={ + "target": "test", + "dcid": "12345", + "comment": "Test process", + "automatic": True, + }, + extra_params={"testid": 42}, + ) + rw = RecipeWrapper(message=m, transport=t) + trigger.trigger(rw, {"some": "header"}, mock.sentinel.message) + t.send.assert_has_calls( + [ + mock.call( + "processing_recipe", + {"recipes": [], "parameters": {"ispyb_process": 12, "testid": 42}}, + ), + ] + ) + + +def test_trigger_service_invalid_target(mocker): + mock_transport = mock.Mock() + zc = zocalo.configuration.from_string(sample_configuration) + zc.activate() + trigger = Trigger(environment={"config": zc}) + setattr(trigger, "_transport", mock_transport) + trigger.initializing() + + t = mock.create_autospec(workflows.transport.common_transport.CommonTransport) + m = generate_recipe_message(parameters={"target": "invalid", "dcid": "12345"}) + rw = RecipeWrapper(message=m, transport=t) + trigger.trigger(rw, {"some": "header"}, mock.sentinel.message) + t.send.assert_has_calls([])