Add starting point for SmartEM integration (#748)

If installed and configured then starting a multigrid controller also adds an Acquisition in SmartEM and triggers the acquisition creation event.

The atlas context is expanded to behave differently under different acquisition software. In the case that SerialEM is being used it will "register" atlases that it sees, meaning they will be registered with SmartEM if installed and configured. The SerialEM context is not accessed by discovery at the moment, as the others are, but is explicitly set in the API calls that trigger the multigrid controller setup. Note this means that this context will be used for all atlas analysers spawned by that controller.

Author: d-j-hatton

pyproject.toml

@@ -68,6 +68,9 @@ server = [
     "stomp-py>8.1.1", # 8.1.1 (released 2024-04-06) doesn't work with our project
     "zocalo>=1",
 ]
+smartem = [
+    "smartem-decisions[backend]",
+]
 [project.urls]
 Bug-Tracker = "https://github.com/DiamondLightSource/python-murfey/issues"
 Documentation = "https://github.com/DiamondLightSource/python-murfey"

src/murfey/client/analyser.py

@@ -39,6 +39,7 @@ def __init__(
         environment: MurfeyInstanceEnvironment | None = None,
         force_mdoc_metadata: bool = False,
         limited: bool = False,
+        serialem: bool = False,
     ):
         super().__init__()
         self._basepath = basepath_local.absolute()
@@ -52,6 +53,7 @@ def __init__(
         self._environment = environment
         self._force_mdoc_metadata = force_mdoc_metadata
         self._token = token
+        self._serialem = serialem
         self.parameters_model: (
             Type[ProcessingParametersSPA] | Type[ProcessingParametersTomo] | None
         ) = None
@@ -138,7 +140,9 @@ def _find_context(self, file_path: Path) -> bool:
 
         # Tomography and SPA workflow checks
         if "atlas" in file_path.parts:
-            self._context = AtlasContext("epu", self._basepath, self._token)
+            self._context = AtlasContext(
+                "serialem" if self._serialem else "epu", self._basepath, self._token
+            )
             return True
 
         if "Metadata" in file_path.parts or file_path.name == "EpuSession.dm":

src/murfey/client/contexts/atlas.py

@@ -28,7 +28,40 @@ def post_transfer(
             environment=environment,
             **kwargs,
         )
+        if self._acquisition_software == "serialem":
+            self.post_transfer_serialem(
+                transferred_file, environment=environment, **kwargs
+            )
+        else:
+            self.post_transfer_epu(transferred_file, environment=environment, **kwargs)
 
+    def post_transfer_serialem(
+        self,
+        transferred_file: Path,
+        environment: Optional[MurfeyInstanceEnvironment] = None,
+        **kwargs,
+    ):
+        if environment and transferred_file.suffix == ".mrc":
+            source = _get_source(transferred_file, environment)
+            if source:
+                capture_post(
+                    base_url=str(environment.url.geturl()),
+                    router_name="session_control.spa_router",
+                    function_name="register_atlas",
+                    token=self._token,
+                    session_id=environment.murfey_session,
+                    data={
+                        "name": transferred_file.stem,
+                        "acquisition_uuid": environment.acquisition_uuid,
+                    },
+                )
+
+    def post_transfer_epu(
+        self,
+        transferred_file: Path,
+        environment: Optional[MurfeyInstanceEnvironment] = None,
+        **kwargs,
+    ):
         if (
             environment
             and "Atlas_" in transferred_file.stem

src/murfey/client/instance_environment.py

@@ -53,6 +53,7 @@ class MurfeyInstanceEnvironment(BaseModel):
     murfey_session: Optional[int] = None
     samples: Dict[Path, SampleInfo] = {}
     rsync_url: str = ""
+    acquisition_uuid: Optional[str] = None
 
     model_config = ConfigDict(arbitrary_types_allowed=True)
 

src/murfey/client/multigrid_control.py

@@ -48,6 +48,8 @@ class MultigridController:
     analysers: Dict[Path, Analyser] = field(default_factory=lambda: {})
     data_collection_parameters: dict = field(default_factory=lambda: {})
     token: str = ""
+    serialem: bool = False
+    acquisition_uuid: Optional[str] = None
     _machine_config: dict = field(default_factory=lambda: {})
     visit_end_time: Optional[datetime] = None
 
@@ -72,6 +74,7 @@ def __post_init__(self):
             symmetry=self.data_collection_parameters.get("symmetry"),
             eer_fractionation=self.data_collection_parameters.get("eer_fractionation"),
             instrument_name=self.instrument_name,
+            acquisition_uuid=self.acquisition_uuid,
         )
         self._machine_config = get_machine_config_client(
             str(self._environment.url.geturl()),
@@ -449,6 +452,7 @@ def rsync_result(update: RSyncerUpdate):
                 environment=self._environment if not self.dummy_dc else None,
                 force_mdoc_metadata=self.force_mdoc_metadata,
                 limited=limited,
+                serialem=self.serialem,
             )
             self.analysers[source].subscribe(self._start_dc)
             self.analysers[source].start()

src/murfey/instrument_server/api.py

@@ -179,6 +179,8 @@ def setup_multigrid_watcher(
         data_collection_parameters=data_collection_parameters.get(label, {}),
         rsync_restarts=watcher_spec.rsync_restarts,
         visit_end_time=watcher_spec.visit_end_time,
+        acquisition_uuid=watcher_spec.acquisition_uuid,
+        serialem=watcher_spec.serialem,
     )
     # Make child directories, if specified
     watcher_spec.source.mkdir(exist_ok=True)

src/murfey/server/api/instrument.py

@@ -13,6 +13,14 @@
 from sqlmodel import select
 from werkzeug.utils import secure_filename
 
+try:
+    from smartem_backend.api_client import EntityConverter
+    from smartem_common.schemas import AcquisitionData
+
+    SMARTEM_ACTIVE = True
+except ImportError:
+    SMARTEM_ACTIVE = False
+
 import murfey.server.prometheus as prom
 from murfey.server.api.auth import (
     MurfeyInstrumentNameFrontend as MurfeyInstrumentName,
@@ -75,6 +83,7 @@ async def activate_instrument_server_for_session(
                 success = response.status == 200
                 instrument_server_token = await response.json()
                 instrument_server_tokens[session_id] = instrument_server_token
+
     if success:
         log.info("Handshake successful")
     else:
@@ -147,6 +156,25 @@ async def setup_multigrid_watcher(
         session = db.exec(select(Session).where(Session.id == session_id)).one()
         visit = session.visit
         async with aiohttp.ClientSession() as clientsession:
+            acquisition_uuid = None
+            if SMARTEM_ACTIVE and machine_config.smartem_api_url:
+                log.info("registering an acquisition with smartem")
+                try:
+                    acquisition_data = EntityConverter.acquisition_to_request(
+                        AcquisitionData(name=visit)
+                    )
+                    async with clientsession.post(
+                        f"{machine_config.smartem_api_url}/acquisitions",
+                        json=acquisition_data.model_dump(),
+                    ) as response:
+                        acquisition_response_data = await response.json()
+                    acquisition_uuid = acquisition_response_data["uuid"]
+                except Exception:
+                    log.warning(
+                        "failed to register acquisition with smartem", exc_info=True
+                    )
+            else:
+                log.info("smartem not configured")
             async with clientsession.post(
                 f"{machine_config.instrument_server_url}{url_path_for('api.router', 'setup_multigrid_watcher', session_id=session_id)}",
                 json={
@@ -161,6 +189,8 @@ async def setup_multigrid_watcher(
                     "visit_end_time": (
                         str(session.visit_end_time) if session.visit_end_time else None
                     ),
+                    "acquisition_uuid": acquisition_uuid,
+                    "serialem": watcher_spec.serialem,
                 },
                 headers={
                     "Authorization": f"Bearer {instrument_server_tokens[session_id]['access_token']}"

src/murfey/server/api/session_control.py

@@ -10,6 +10,14 @@
 from sqlalchemy import func
 from sqlmodel import select
 
+try:
+    from smartem_backend.api_client import SmartEMAPIClient
+    from smartem_common.schemas import AtlasData
+
+    SMARTEM_ACTIVE = True
+except ImportError:
+    SMARTEM_ACTIVE = False
+
 import murfey.server.prometheus as prom
 from murfey.server import _transport_object
 from murfey.server.api.auth import (
@@ -349,6 +357,49 @@ def get_foil_hole(
     return _get_foil_hole(session_id, fh_name, db)
 
 
+class AtlasRegistration(BaseModel):
+    name: str
+    acqusition_uuid: str
+
+
+@spa_router.post("/sessions/{session_id}/register_atlas")
+def register_atlas(
+    session_id: MurfeySessionID,
+    atlas_registration_data: AtlasRegistration,
+    db=murfey_db,
+):
+    if SMARTEM_ACTIVE:
+        session = db.exec(select(Session).where(Session.id == session_id)).one()
+        machine_config = get_machine_config(session.instrument_name)[
+            session.instrument_name
+        ]
+        if machine_config.smartem_api_url:
+            smartem_client = SmartEMAPIClient(
+                base_url=machine_config.smartem_api_url, logger=logger
+            )
+            possible_grids = smartem_client.get_acquisition_grids(
+                atlas_registration_data.acqusition_uuid
+            )
+            grid_uuid = None
+            for grid in possible_grids:
+                if grid.name == atlas_registration_data.name.replace("_atlas", ""):
+                    grid_uuid = grid.uuid
+                    break
+            if grid_uuid is not None:
+                atlas_data = AtlasData(
+                    id=atlas_registration_data.name,
+                    acquisition_data=datetime.now(),
+                    storage_folder="",
+                    name=atlas_registration_data.name,
+                    tiles=[],
+                    gridsquare_positions=None,
+                    grid_uuid=grid_uuid,
+                )
+                smartem_client.create_grid_atlas(atlas_data)
+    else:
+        logger.info("smartem deactivated so did not register atlas")
+
+
 @spa_router.post("/sessions/{session_id}/make_atlas_jpg")
 def make_atlas_jpg(
     session_id: MurfeySessionID, atlas_mrc: StringOfPathModel, db=murfey_db

src/murfey/server/api/session_info.py

@@ -81,6 +81,12 @@ def machine_info_by_instrument(
     return get_machine_config(instrument_name)[instrument_name]
 
 
+@router.get("/instruments/{instrument_name}/smartem")
+def check_smartem_availability(instrument_name: str):
+    machine_config = get_machine_config(instrument_name)[instrument_name]
+    return {"available": bool(machine_config.smartem_api_url)}
+
+
 @router.get("/instruments/{instrument_name}/visits_raw", response_model=List[Visit])
 def get_current_visits(instrument_name: MurfeyInstrumentName, db=ispyb_db):
     logger.debug(

src/murfey/util/config.py

@@ -106,6 +106,7 @@ class MachineConfig(BaseModel):  # type: ignore
     murfey_url: str = "http://localhost:8000"
     frontend_url: str = "http://localhost:3000"
     instrument_server_url: str = "http://localhost:8001"
+    smartem_api_url: str = ""
 
     # Messaging queues
     failure_queue: str = ""