Mimas: more flexible/exensible handling of scenario (#131)

Use an entry points mechanism to allow enable a more flexible and extensible
approach to handling scenarios in mimas.

The simplest mechanism is to register one (or more) functions that take a
scenario and return a list of MimasRecipeInvocations or
MimasISPyBJobInvocations:

def handle_scenarios(scenario: MimasScenario) -> List[Invocation]:
    return [
        MimasRecipeInvocation(...),
        MimasISPyBJobInvocation(...),
        ...
    ]

This function will handle every scenario passed to the mimas service.
Alternatively, a more modular approach can be used, using the
@match_specification decorator to filter scenarios that match a given
specification before passing them to the function:

from dlstbx.mimas.specification import BeamlineSpecification, DCClassSpecification

is_i99 = BeamlineSpecification("i99")
is_rotation = DCClassSpecification(MimasDCClass.ROTATION)

@match_specification(is_i99 & is_rotation)
def handle_i99_rotation(scenario: MimasScenario) -> List[Invocation]:
    return [
        MimasRecipeInvocation(...),
        MimasISPyBJobInvocation(...),
        ...
    ]

Various commonly used specifications are provided, such as
BeamlineSpecification, DCClassSpecification, DetectorClassSpecification,
EventSpecification andVisitSpecification. More complex custom specifications
may be built by inheriting from ScenarioSpecification. Specifications can be
combined using bitwise operators (&, |, ~).

Author: rjgildea

setup.py

@@ -160,6 +160,27 @@
     "zocalo.dlq.rabbitmq = dlstbx.health_checks.rabbitmq:check_rabbitmq_dlq",
 ]
 
+mimas_scenario_handlers = [
+    "cloud = dlstbx.mimas.cloud:handle_cloud",
+    "eiger_screening = dlstbx.mimas.core:handle_eiger_screening",
+    "eiger_start = dlstbx.mimas.core:handle_eiger_start",
+    "eiger_end = dlstbx.mimas.core:handle_eiger_end",
+    "i19_pilatus_start = dlstbx.mimas.i19:handle_i19_start_pilatus",
+    "i19_eiger_start = dlstbx.mimas.i19:handle_i19_start_eiger",
+    "i19_pilatus_end = dlstbx.mimas.i19:handle_i19_end_pilatus",
+    "i19_eiger_end = dlstbx.mimas.i19:handle_i19_end_eiger",
+    "i19_end = dlstbx.mimas.i19:handle_i19_end",
+    "pilatus_end = dlstbx.mimas.core:handle_pilatus_end",
+    "pilatus_gridscan_start = dlstbx.mimas.core:handle_pilatus_gridscan_start",
+    "pilatus_not_gridscan_start = dlstbx.mimas.core:handle_pilatus_not_gridscan_start",
+    "pilatus_screening = dlstbx.mimas.core:handle_pilatus_screening",
+    "rotation_end = dlstbx.mimas.core:handle_rotation_end",
+    "vmxi_end = dlstbx.mimas.vmxi:handle_vmxi_end",
+    "vmxi_gridscan = dlstbx.mimas.vmxi:handle_vmxi_gridscan",
+    "vmxi_rotation = dlstbx.mimas.vmxi:handle_vmxi_rotation_scan",
+    "vmxi_start = dlstbx.mimas.vmxi:handle_vmxi_start",
+]
+
 
 def get_git_revision():
     """Try to obtain the current git revision number"""
@@ -219,6 +240,7 @@ def get_git_revision():
             "thumbnail = dlstbx.services.images:thumbnail",
         ],
         "zocalo.wrappers": sorted(known_wrappers),
+        "zocalo.mimas.handlers": sorted(mimas_scenario_handlers),
     },
     packages=find_packages("src"),
     package_dir={"": "src"},

src/dlstbx/cli/mimas.py

@@ -12,7 +12,6 @@
 import argparse
 
 import dlstbx.ispybtbx
-import dlstbx.mimas.core
 
 _readable = {
     dlstbx.mimas.MimasEvent.START: "start of data collection",
@@ -104,7 +103,7 @@ def run(args=None):
 
     for dcid in args.dcids:
         for scenario in get_scenarios(dcid):
-            actions = dlstbx.mimas.core.run(scenario)
+            actions = dlstbx.mimas.handle_scenario(scenario)
             print(f"At the {_readable.get(scenario.event)} {dcid}:")
             for a in sorted(actions, key=lambda a: str(type(a)) + " " + a.recipe):
                 try:

src/dlstbx/mimas/__init__.py

@@ -1,10 +1,15 @@
+from __future__ import annotations
+
 import dataclasses
 import enum
 import functools
 import numbers
-from typing import Optional, Tuple
+from typing import Callable, List, Optional, Tuple, Union
 
 import gemmi
+import pkg_resources
+
+from dlstbx.mimas.specification import BaseSpecification
 
 MimasDCClass = enum.Enum("MimasDCClass", "GRIDSCAN ROTATION SCREENING UNDEFINED")
 
@@ -373,3 +378,34 @@ def _(mimasobject: MimasISPyBJobInvocation):
             *triggervars,
         )
     )
+
+
+Invocation = Union[MimasISPyBJobInvocation, MimasRecipeInvocation]
+
+
+def match_specification(specification: BaseSpecification):
+    def outer_wrapper(handler: Callable):
+        @functools.wraps(handler)
+        def inner_wrapper(scenario: MimasScenario) -> List[Invocation]:
+            if specification.is_satisfied_by(scenario):
+                return handler(scenario)
+            return []
+
+        return inner_wrapper
+
+    return outer_wrapper
+
+
+@functools.lru_cache
+def _get_handlers() -> dict[str, Callable]:
+    return {
+        e.name: e.load()
+        for e in pkg_resources.iter_entry_points("zocalo.mimas.handlers")
+    }
+
+
+def handle_scenario(scenario: MimasScenario) -> List[Invocation]:
+    tasks: List[Invocation] = []
+    for handler in _get_handlers().values():
+        tasks.extend(handler(scenario))
+    return tasks

src/dlstbx/mimas/cloud.py

@@ -0,0 +1,61 @@
+from __future__ import annotations
+
+from typing import List
+
+from dlstbx import mimas
+from dlstbx.mimas.core import is_end, is_rotation, xia2_dials_absorption_params
+from dlstbx.mimas.specification import BeamlineSpecification, VisitSpecification
+
+CLOUD_VISITS = {
+    "cm",
+    "nt28218",
+    "mx",
+}
+
+
+is_cloud = (
+    VisitSpecification(CLOUD_VISITS)
+    & is_end
+    & is_rotation
+    & BeamlineSpecification("i03")
+)
+
+
+@mimas.match_specification(is_cloud)
+def handle_cloud(
+    scenario: mimas.MimasScenario,
+) -> List[mimas.Invocation]:
+    return [
+        # xia2-dials
+        mimas.MimasISPyBJobInvocation(
+            DCID=scenario.DCID,
+            autostart=True,
+            recipe="autoprocessing-xia2-dials-eiger-cloud",
+            source="automatic",
+            parameters=(
+                mimas.MimasISPyBParameter(
+                    key="resolution.cc_half_significance_level", value="0.1"
+                ),
+                *xia2_dials_absorption_params(scenario),
+            ),
+        ),
+        # xia2-3dii
+        mimas.MimasISPyBJobInvocation(
+            DCID=scenario.DCID,
+            autostart=True,
+            recipe="autoprocessing-xia2-3dii-eiger-cluster",
+            source="automatic",
+            parameters=(
+                mimas.MimasISPyBParameter(
+                    key="resolution.cc_half_significance_level", value="0.1"
+                ),
+            ),
+        ),
+        # autoPROC
+        mimas.MimasISPyBJobInvocation(
+            DCID=scenario.DCID,
+            autostart=True,
+            recipe="autoprocessing-autoPROC-eiger-cluster",
+            source="automatic",
+        ),
+    ]