API setup to pass visit end times to multigrid controller for use with rsyncers

Author: d-j-hatton

src/murfey/client/multigrid_control.py

@@ -48,6 +48,8 @@ class MultigridController:
     data_collection_parameters: dict = field(default_factory=lambda: {})
     token: str = ""
     _machine_config: dict = field(default_factory=lambda: {})
+    visit_end_time: Optional[datetime] = None
+    end_time_grace_period: int = 0
 
     def __post_init__(self):
         if self.token:
@@ -277,6 +279,8 @@ def _start_rsyncer(
                 stop_callback=self._rsyncer_stopped,
                 do_transfer=self.do_transfer,
                 remove_files=remove_files,
+                end_time=self.visit_end_time,
+                grace_period=self.end_time_grace_period,
             )
 
             def rsync_result(update: RSyncerUpdate):

src/murfey/instrument_server/api.py

@@ -160,6 +160,8 @@ def setup_multigrid_watcher(
         token=tokens.get(session_id, "token"),
         data_collection_parameters=data_collection_parameters.get(label, {}),
         rsync_restarts=watcher_spec.rsync_restarts,
+        visit_end_time=watcher_spec.visit_end_time,
+        end_time_grace_period=watcher_spec.grace_period,
     )
     watcher_spec.source.mkdir(exist_ok=True)
     machine_config = requests.get(

src/murfey/server/api/instrument.py

@@ -98,9 +98,8 @@ async def setup_multigrid_watcher(
     session_id: MurfeySessionID, watcher_spec: MultigridWatcherSetup, db=murfey_db
 ):
     data = {}
-    instrument_name = (
-        db.exec(select(Session).where(Session.id == session_id)).one().instrument_name
-    )
+    session = db.exec(select(Session).where(Session.id == session_id)).one()
+    instrument_name = session.instrument_name
     machine_config = get_machine_config(instrument_name=instrument_name)[
         instrument_name
     ]
@@ -130,6 +129,8 @@ async def setup_multigrid_watcher(
                         str(k): v for k, v in watcher_spec.destination_overrides.items()
                     },
                     "rsync_restarts": watcher_spec.rsync_restarts,
+                    "visit_end_time": session.visit_end_time,
+                    "grace_period": machine_config.grace_period,
                 },
                 headers={
                     "Authorization": f"Bearer {instrument_server_tokens[session_id]['access_token']}"

src/murfey/util/config.py

@@ -71,6 +71,8 @@ class MachineConfig(BaseModel, extra=Extra.allow):  # type: ignore
 
     notifications_queue: str = "pato_notification"
 
+    grace_period: int = 0
+
 
 def from_file(config_file_path: Path, instrument: str = "") -> Dict[str, MachineConfig]:
     with open(config_file_path, "r") as config_stream:

src/murfey/util/db.py

@@ -3,6 +3,7 @@
 of the sessions that Murfey is overseeing, along with the relationships between them.
 """
 
+from datetime import datetime
 from typing import List, Optional
 
 import sqlalchemy
@@ -48,6 +49,7 @@ class Session(SQLModel, table=True):  # type: ignore
     started: bool = Field(default=False)
     current_gain_ref: str = Field(default="")
     instrument_name: str = Field(default="")
+    visit_end_time: Optional[datetime] = Field(default=None)
 
     # CLEM Workflow
 

src/murfey/util/instrument_models.py

@@ -1,5 +1,6 @@
+from datetime import datetime
 from pathlib import Path
-from typing import Dict, List
+from typing import Dict, List, Optional
 
 from pydantic import BaseModel
 
@@ -15,3 +16,5 @@ class MultigridWatcherSpec(BaseModel):
     skip_existing_processing: bool = False
     destination_overrides: Dict[Path, str] = {}
     rsync_restarts: List[str] = []
+    visit_end_time: Optional[datetime] = None
+    grace_period: int = 0