diff --git a/src/murfey/client/contexts/atlas.py b/src/murfey/client/contexts/atlas.py index 4905efd82..9ccfda43d 100644 --- a/src/murfey/client/contexts/atlas.py +++ b/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 "" + ), }, ) diff --git a/src/murfey/client/contexts/spa_metadata.py b/src/murfey/client/contexts/spa_metadata.py index b48129646..027701cab 100644 --- a/src/murfey/client/contexts/spa_metadata.py +++ b/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" diff --git a/src/murfey/server/api/workflow.py b/src/murfey/server/api/workflow.py index 3f3a72358..7f5c9a996 100644 --- a/src/murfey/server/api/workflow.py +++ b/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, diff --git a/tests/client/contexts/test_atlas.py b/tests/client/contexts/test_atlas.py index d34b2ef3c..a3ca5122f 100644 --- a/tests/client/contexts/test_atlas.py +++ b/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 = ( "" "12001500" @@ -136,7 +130,8 @@ def test_atlas_context_dm(mock_capture_post, tmp_path): "" ) - 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", }, )