Send cryolo model when flushing spa (#598)

Also removes cryolo model api endpoint

Author: stephen-riggs

src/murfey/server/api/spa.py

@@ -27,7 +27,6 @@
 import murfey.server.prometheus as prom
 from murfey.server import _transport_object
 from murfey.server.api.auth import MurfeySessionID, validate_token
-from murfey.server.api.spa import _cryolo_model_path
 from murfey.server.feedback import (
     _murfey_id,
     check_tilt_series_mc,
@@ -58,7 +57,11 @@
     TiltSeries,
 )
 from murfey.util.models import ProcessingParametersSPA, ProcessingParametersTomo
-from murfey.util.processing_params import default_spa_parameters, motion_corrected_mrc
+from murfey.util.processing_params import (
+    cryolo_model_path,
+    default_spa_parameters,
+    motion_corrected_mrc,
+)
 from murfey.util.tomo import midpoint
 
 logger = getLogger("murfey.server.api.workflow")
@@ -457,7 +460,7 @@ async def request_spa_preprocessing(
                 ),
                 "do_icebreaker_jobs": default_spa_parameters.do_icebreaker_jobs,
                 "cryolo_model_weights": str(
-                    _cryolo_model_path(visit_name, instrument_name)
+                    cryolo_model_path(visit_name, instrument_name)
                 ),
             },
         }

src/murfey/server/api/workflow.py

@@ -22,7 +22,6 @@
 import murfey.server.api.prometheus
 import murfey.server.api.session_control
 import murfey.server.api.session_info
-import murfey.server.api.spa
 import murfey.server.api.websocket
 import murfey.server.api.workflow
 from murfey.server import template_files
@@ -90,7 +89,6 @@ class Settings(BaseSettings):
 app.include_router(murfey.server.api.workflow.correlative_router)
 app.include_router(murfey.server.api.workflow.spa_router)
 app.include_router(murfey.server.api.workflow.tomo_router)
-app.include_router(murfey.server.api.spa.router)
 app.include_router(murfey.server.api.clem.router)
 
 app.include_router(murfey.server.api.prometheus.router)

src/murfey/server/main.py

@@ -1,10 +1,12 @@
+from datetime import datetime
+from functools import lru_cache
 from pathlib import Path
 from typing import Literal, Optional
 
 from pydantic import BaseModel
 from werkzeug.utils import secure_filename
 
-from murfey.util.config import MachineConfig
+from murfey.util.config import MachineConfig, get_machine_config
 
 
 def motion_corrected_mrc(
@@ -33,6 +35,23 @@ def motion_corrected_mrc(
     return Path("/".join(secure_filename(p) for p in mrc_out.parts))
 
 
+@lru_cache(maxsize=5)
+def cryolo_model_path(visit: str, instrument_name: str) -> Path:
+    machine_config = get_machine_config(instrument_name=instrument_name)[
+        instrument_name
+    ]
+    if machine_config.model_search_directory:
+        visit_directory = (
+            machine_config.rsync_basepath / str(datetime.now().year) / visit
+        )
+        possible_models = list(
+            (visit_directory / machine_config.model_search_directory).glob("*.h5")
+        )
+        if possible_models:
+            return sorted(possible_models, key=lambda x: x.stat().st_ctime)[-1]
+    return machine_config.default_model
+
+
 class CLEMAlignAndMergeParameters(BaseModel):
     crop_to_n_frames: Optional[int] = 50
     align_self: Literal["enabled", ""] = "enabled"

src/murfey/util/processing_params.py

@@ -980,14 +980,6 @@ murfey.server.api.session_info.tomo_router:
         type: str
     methods:
       - GET
-murfey.server.api.spa.router:
-  - path: /sessions/{session_id}/cryolo_model
-    function: get_cryolo_model_path
-    path_params:
-      - name: session_id
-        type: int
-    methods:
-      - GET
 murfey.server.api.websocket.ws:
   - path: /ws/test/{client_id}
     function: websocket_endpoint

src/murfey/util/route_manifest.yaml

@@ -24,7 +24,7 @@
 from murfey.util.db import Session as MurfeySession
 from murfey.util.db import SPAFeedbackParameters, SPARelionParameters
 from murfey.util.models import FoilHoleParameters, GridSquareParameters
-from murfey.util.processing_params import default_spa_parameters
+from murfey.util.processing_params import cryolo_model_path, default_spa_parameters
 from murfey.util.spa_metadata import (
     GridSquareInfo,
     foil_hole_data,
@@ -314,15 +314,12 @@ def flush_spa_preprocess(message: dict, murfey_db: Session, demo: bool = False)
     ).all()
     if not stashed_files:
         return True
-    instrument_name = (
-        murfey_db.exec(
-            select(MurfeySession).where(MurfeySession.id == message["session_id"])
-        )
-        .one()
-        .instrument_name
-    )
-    machine_config = get_machine_config(instrument_name=instrument_name)[
-        instrument_name
+
+    murfey_session = murfey_db.exec(
+        select(MurfeySession).where(MurfeySession.id == message["session_id"])
+    ).one()
+    machine_config = get_machine_config(instrument_name=murfey_session.instrument_name)[
+        murfey_session.instrument_name
     ]
     recipe_name = machine_config.recipes.get("em-spa-preprocess", "em-spa-preprocess")
     collected_ids = murfey_db.exec(
@@ -424,6 +421,11 @@ def flush_spa_preprocess(message: dict, murfey_db: Session, demo: bool = False)
                     else f.eer_fractionation_file
                 ),
                 "do_icebreaker_jobs": default_spa_parameters.do_icebreaker_jobs,
+                "cryolo_model_weights": str(
+                    cryolo_model_path(
+                        murfey_session.visit, murfey_session.instrument_name
+                    )
+                ),
                 "foil_hole_id": foil_hole_id,
             },
         }