Add starting point for SmartEM integration

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,13 @@
 from sqlmodel import select
 from werkzeug.utils import secure_filename
 
+try:
+    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 +82,7 @@
                 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 +155,15 @@
         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:
+                acquisition_data = AcquisitionData(name=visit)
+                async with clientsession.post(
+                    f"{machine_config.smartem_api_url}/acquisitions",
+                    AcquisitionData.model_json_schema(),
+                ) as response:
+                    acquisition_data = await response.json()
+                acquisition_uuid = acquisition_data.uuid
             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 +178,8 @@
                     "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,47 @@ 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)
+
+
 @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 = ""

src/murfey/util/instrument_models.py

@@ -13,3 +13,5 @@ class MultigridWatcherSpec(BaseModel):
     destination_overrides: Dict[Path, str] = {}
     rsync_restarts: List[str] = []
     visit_end_time: Optional[datetime] = None
+    acquisition_uuid: Optional[str] = None
+    serialem: bool = False

src/murfey/util/models.py

@@ -208,6 +208,7 @@ class MultigridWatcherSetup(BaseModel):
     source: Path
     destination_overrides: Dict[Path, str] = {}
     rsync_restarts: List[str] = []
+    serialem: bool = False
 
 
 class Token(BaseModel):

src/murfey/util/route_manifest.yaml

@@ -975,6 +975,13 @@ murfey.server.api.session_control.spa_router:
         type: int
     methods:
       - GET
+  - path: /session_control/spa/sessions/{session_id}/register_atlas
+    function: register_atlas
+    path_params:
+      - name: session_id
+        type: int
+    methods:
+      - POST
   - path: /session_control/spa/sessions/{session_id}/make_atlas_jpg
     function: make_atlas_jpg
     path_params:
@@ -1085,6 +1092,13 @@ murfey.server.api.session_info.router:
         type: str
     methods:
       - GET
+  - path: /session_info/instruments/{instrument_name}/smartem
+    function: check_smartem_availability
+    path_params:
+      - name: instrument_name
+        type: str
+    methods:
+      - GET
   - path: /session_info/instruments/{instrument_name}/visits_raw
     function: get_current_visits
     path_params: