Fix for multiple smartem grid registration

src/murfey/client/contexts/atlas.py

@@ -78,27 +78,6 @@ def post_transfer_serialem(
                 logger.info(
                     f"Registered data collection group for atlas {str(transferred_file.stem)!r}"
                 )
-                capture_post(
-                    base_url=str(environment.url.geturl()),
-                    router_name="session_control.spa_router",
-                    function_name="register_atlas",
-                    token=self._token,
-                    instrument_name=environment.instrument_name,
-                    session_id=environment.murfey_session,
-                    data={
-                        "tag": str(transferred_file.parent / transferred_file.stem),
-                        "name": transferred_file.stem,
-                        "acquisition_uuid": environment.acquisition_uuid,
-                        "storage_folder": str(
-                            _atlas_destination(
-                                environment,
-                                source,
-                                Path(self._machine_config.get("rsync_basepath", "")),
-                            )
-                            / "atlas"
-                        ),
-                    },
-                )
 
     def post_transfer_epu(
         self,
@@ -171,8 +150,7 @@ def post_transfer_epu(
                     "atlas": str(transferred_atlas_jpg).replace("//", "/"),
                     "sample": sample,
                     "atlas_pixel_size": atlas_pixel_size,
-                    "create_smartem_grid": bool(environment.acquisition_uuid),
-                    "acquisition_uuid": environment.acquisition_uuid,
+                    "create_smartem_grid": False,
                 }
                 capture_post(
                     base_url=str(environment.url.geturl()),
@@ -204,20 +182,36 @@ def post_transfer_epu(
                 return
 
             # Make sure a dcg is requested before doing grid squares
-            capture_post(
-                base_url=str(environment.url.geturl()),
-                router_name="workflow.router",
-                function_name="register_dc_group",
-                token=self._token,
-                instrument_name=environment.instrument_name,
-                visit_name=environment.visit,
-                session_id=environment.murfey_session,
-                data={
-                    "experiment_type_id": 44,  # Atlas
-                    "tag": str(transferred_file.parent),
-                    "sample": sample,
-                },
-            )
+            source = _get_source(transferred_file, environment)
+            atlas_mrc_glob = list(transferred_file.parent.glob("Atlas_*.mrc"))
+
+            if source:
+                if atlas_mrc_glob:
+                    atlas_mrc = atlas_mrc_glob[0]
+                    transferred_atlas: str | Path = _atlas_destination(
+                        environment,
+                        source,
+                        Path(self._machine_config.get("rsync_basepath", "")),
+                    ) / atlas_mrc.relative_to(source.parent)
+                else:
+                    transferred_atlas = ""
+                capture_post(
+                    base_url=str(environment.url.geturl()),
+                    router_name="workflow.router",
+                    function_name="register_dc_group",
+                    token=self._token,
+                    instrument_name=environment.instrument_name,
+                    visit_name=environment.visit,
+                    session_id=environment.murfey_session,
+                    data={
+                        "experiment_type_id": 44,  # Atlas
+                        "tag": str(transferred_file.parent),
+                        "atlas": str(transferred_atlas).replace("//", "/"),
+                        "sample": sample,
+                        "create_smartem_grid": bool(environment.acquisition_uuid),
+                        "acquisition_uuid": environment.acquisition_uuid,
+                    },
+                )
             # Register all grid squares on this atlas
             for gs, pos_data in gs_pix_positions.items():
                 if pos_data:
@@ -255,5 +249,15 @@ def post_transfer_epu(
                         "acquisition_uuid": environment.acquisition_uuid,
                         "register_grid": True,
                         "tag": str(transferred_file.parent),
+                        "storage_folder": str(
+                            _atlas_destination(
+                                environment,
+                                source,
+                                Path(self._machine_config.get("rsync_basepath", "")),
+                            )
+                            / "atlas"
+                            if source
+                            else ""
+                        ),
                     },
                 )

src/murfey/client/contexts/spa_metadata.py

@@ -228,21 +228,6 @@ def post_transfer_epu(
                                     "angle": pos_data[6],
                                 },
                             )
-                    if gs_pix_positions:
-                        capture_post(
-                            base_url=str(environment.url.geturl()),
-                            router_name="session_control.spa_router",
-                            function_name="register_atlas",
-                            token=self._token,
-                            instrument_name=environment.instrument_name,
-                            session_id=environment.murfey_session,
-                            data={
-                                "name": f"{environment.visit}-slot-{environment.samples[images_disc].sample}",
-                                "acquisition_uuid": environment.acquisition_uuid,
-                                "register_grid": True,
-                                "tag": dcg_tag,
-                            },
-                        )
 
         elif (
             transferred_file.suffix == ".dm"

src/murfey/server/api/workflow.py

@@ -56,7 +56,7 @@
 from murfey.server.ispyb import DB as ispyb_db, get_proposal_id
 from murfey.server.murfey_db import murfey_db
 from murfey.util import sanitise
-from murfey.util.config import get_machine_config
+from murfey.util.config import MachineConfig, get_machine_config
 from murfey.util.db import (
     AutoProcProgram,
     DataCollection,
@@ -112,28 +112,10 @@ class DCGroupParameters(BaseModel):
     acquisition_uuid: Optional[str] = None
 
 
-@router.post(
-    "/visits/{visit_name}/sessions/{session_id}/register_data_collection_group"
-)
-def register_dc_group(
-    visit_name: str,
-    session_id: MurfeySessionID,
-    dcg_params: DCGroupParameters,
-    db: SQLModelSession = murfey_db,
-):
-    ispyb_proposal_code = visit_name[:2]
-    ispyb_proposal_number = visit_name.split("-")[0][2:]
-    ispyb_visit_number = visit_name.split("-")[-1]
-    instrument_name = (
-        db.exec(select(MurfeySession).where(MurfeySession.id == session_id))
-        .one()
-        .instrument_name
-    )
-    logger.info(f"Registering data collection group on microscope {instrument_name}")
-    machine_config = get_machine_config(instrument_name=instrument_name)[
-        instrument_name
-    ]
-    smartem_grid_uuid = None
+def _register_smartem_grid(
+    machine_config: MachineConfig, dcg_params: DCGroupParameters, visit_name: str
+) -> str:
+    smartem_grid_uuid = ""
     if SMARTEM_ACTIVE and dcg_params.acquisition_uuid:
         if machine_config.smartem_api_url:
             try:
@@ -167,6 +149,31 @@ def register_dc_group(
 
             except Exception:
                 logger.warning("Failed to register SmartEM grid", exc_info=True)
+    return smartem_grid_uuid
+
+
+@router.post(
+    "/visits/{visit_name}/sessions/{session_id}/register_data_collection_group"
+)
+def register_dc_group(
+    visit_name: str,
+    session_id: MurfeySessionID,
+    dcg_params: DCGroupParameters,
+    db: SQLModelSession = murfey_db,
+):
+    ispyb_proposal_code = visit_name[:2]
+    ispyb_proposal_number = visit_name.split("-")[0][2:]
+    ispyb_visit_number = visit_name.split("-")[-1]
+    instrument_name = (
+        db.exec(select(MurfeySession).where(MurfeySession.id == session_id))
+        .one()
+        .instrument_name
+    )
+    logger.info(f"Registering data collection group on microscope {instrument_name}")
+    machine_config = get_machine_config(instrument_name=instrument_name)[
+        instrument_name
+    ]
+
     if (
         dcg_murfey := db.exec(
             select(DataCollectionGroup)
@@ -185,6 +192,15 @@ def register_dc_group(
     ):
         # Either switching atlas for a common (atlas or processing) tag
         # Or registering a new atlas-type dcg for a sample that is already present
+        smartem_grid_uuid = ""
+        for dcg in dcg_murfey:
+            if dcg.smartem_grid_uuid:
+                smartem_grid_uuid = dcg.smartem_grid_uuid
+                break
+        else:
+            smartem_grid_uuid = _register_smartem_grid(
+                machine_config, dcg_params, visit_name
+            )
         for dcg_instance in dcg_murfey:
             # Update all instances in case there are multiple processing runs
             # Skip Scaup registration if sample is unchanged
@@ -300,6 +316,10 @@ def register_dc_group(
         )
     ).all():
         # Case where we switch from atlas to processing
+        if not dcg_murfey[0].smartem_grid_uuid:
+            dcg_murfey[0].smartem_grid_uuid = _register_smartem_grid(
+                machine_config, dcg_params, visit_name
+            )
         original_tag = dcg_murfey[0].tag
         dcg_murfey[0].tag = dcg_params.tag or dcg_murfey[0].tag
         if murfey.server._transport_object:
@@ -321,6 +341,9 @@ def register_dc_group(
             db.add(grid_square)
         db.commit()
     else:
+        smartem_grid_uuid = _register_smartem_grid(
+            machine_config, dcg_params, visit_name
+        )
         dcg_parameters = {
             "start_time": str(datetime.now()),
             "experiment_type_id": dcg_params.experiment_type_id,

tests/client/contexts/test_atlas.py

@@ -20,14 +20,12 @@ def test_atlas_context_mrc(mock_capture_post, tmp_path):
         url=urlparse("http://localhost:8000"),
         client_id=0,
         sources=[tmp_path / "cm12345-6"],
-        default_destinations={
-            tmp_path / "cm12345-6": f"{tmp_path}/destination/cm12345-6"
-        },
+        default_destinations={tmp_path / "cm12345-6": "destination/cm12345-6"},
         instrument_name="m01",
         visit="cm12345-6",
         murfey_session=1,
     )
-    context = AtlasContext("tomo", tmp_path, {}, "token")
+    context = AtlasContext("tomo", tmp_path, {"rsync_basepath": "/base"}, "token")
 
     atlas_mrc = tmp_path / "cm12345-6/Supervisor_atlas/Sample2/Atlas/Atlas_1.mrc"
     atlas_mrc.parent.mkdir(parents=True)
@@ -41,7 +39,7 @@ def test_atlas_context_mrc(mock_capture_post, tmp_path):
         token="token",
         instrument_name="m01",
         session_id=1,
-        data={"path": f"{tmp_path}/destination/{atlas_mrc.relative_to(tmp_path)}"},
+        data={"path": f"/base/destination/{atlas_mrc.relative_to(tmp_path)}"},
     )
 
 
@@ -51,14 +49,12 @@ def test_atlas_context_xml(mock_capture_post, tmp_path):
         url=urlparse("http://localhost:8000"),
         client_id=0,
         sources=[tmp_path / "cm12345-6"],
-        default_destinations={
-            tmp_path / "cm12345-6": f"{tmp_path}/destination/cm12345-6"
-        },
+        default_destinations={tmp_path / "cm12345-6": "destination/cm12345-6"},
         instrument_name="m01",
         visit="cm12345-6",
         murfey_session=1,
     )
-    context = AtlasContext("tomo", tmp_path, {}, "token")
+    context = AtlasContext("tomo", tmp_path, {"rsync_basepath": "/base"}, "token")
 
     atlas_pixel_size = 4.6
     atlas_xml = tmp_path / "cm12345-6/Supervisor_atlas/Sample2/Atlas/Atlas_1.xml"
@@ -73,11 +69,10 @@ def test_atlas_context_xml(mock_capture_post, tmp_path):
     dcg_data = {
         "experiment_type_id": 44,  # Atlas
         "tag": str(atlas_xml.parent),
-        "atlas": f"{tmp_path}/destination/{atlas_xml.relative_to(tmp_path).with_suffix('.jpg')}",
+        "atlas": f"/base/destination/{atlas_xml.relative_to(tmp_path).with_suffix('.jpg')}",
         "sample": 2,
         "atlas_pixel_size": atlas_pixel_size * 7.8,
         "create_smartem_grid": False,
-        "acquisition_uuid": None,
     }
     mock_capture_post.assert_called_once_with(
         base_url="http://localhost:8000",
@@ -97,9 +92,7 @@ def test_atlas_context_dm(mock_capture_post, tmp_path):
         url=urlparse("http://localhost:8000"),
         client_id=0,
         sources=[tmp_path / "cm12345-6"],
-        default_destinations={
-            tmp_path / "cm12345-6": f"{tmp_path}/destination/cm12345-6"
-        },
+        default_destinations={tmp_path / "cm12345-6": "destination/cm12345-6"},
         instrument_name="m01",
         visit="cm12345-6",
         murfey_session=1,
@@ -109,6 +102,7 @@ def test_atlas_context_dm(mock_capture_post, tmp_path):
     # Write sample dm file
     atlas_dm = tmp_path / "cm12345-6/Supervisor_atlas/Sample2/Atlas/Atlas.dm"
     atlas_dm.parent.mkdir(parents=True)
+    (tmp_path / "cm12345-6/Supervisor_atlas/Sample2/Atlas/Atlas_01.mrc").touch()
     grid_square_values = (
         "<value><b:PositionOnTheAtlas>"
         "<c:Center><d:x>1200</d:x><d:y>1500</d:y></c:Center>"
@@ -136,7 +130,8 @@ def test_atlas_context_dm(mock_capture_post, tmp_path):
             "</_items></TilesEfficient></Atlas></AtlasSessionXml>"
         )
 
-    context = AtlasContext("tomo", tmp_path, {}, "token")
+    context = AtlasContext("tomo", tmp_path, {"rsync_basepath": "/base"}, "token")
+    assert context._machine_config.get("rsync_basepath") == "/base"
     context.post_transfer(atlas_dm, environment=env)
 
     assert mock_capture_post.call_count == 6
@@ -152,6 +147,9 @@ def test_atlas_context_dm(mock_capture_post, tmp_path):
             "experiment_type_id": 44,  # Atlas
             "tag": str(atlas_dm.parent),
             "sample": 2,
+            "atlas": "/base/destination/cm12345-6/Supervisor_atlas/Sample2/Atlas/Atlas_01.mrc",
+            "create_smartem_grid": True,
+            "acquisition_uuid": "uuid1",
         },
     )
     mock_capture_post.assert_any_call(
@@ -186,5 +184,6 @@ def test_atlas_context_dm(mock_capture_post, tmp_path):
             "acquisition_uuid": "uuid1",
             "register_grid": True,
             "tag": str(atlas_dm.parent),
+            "storage_folder": "/base/destination/atlas",
         },
     )