Add check-in for experiment clients

src/wei/experiment_client.py

@@ -2,6 +2,7 @@
 
 import json
 import time
+import traceback
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Union
 
@@ -24,6 +25,7 @@
 )
 from wei.types.exceptions import WorkflowFailedException
 from wei.types.experiment_types import Experiment, ExperimentDesign
+from wei.utils import threaded_daemon
 
 
 class ExperimentClient:
@@ -33,7 +35,7 @@ def __init__(
         self,
         server_host: str,
         server_port: str,
-        experiment_name: str,
+        experiment_name: Optional[str] = None,
         experiment_id: Optional[str] = None,
         campaign_id: Optional[str] = None,
         description: Optional[str] = None,
@@ -64,7 +66,7 @@ def __init__(
             List of email addresses to send notifications at the end of the experiment
 
         log_experiment_end_on_exit: bool
-            Whether to log the end of the experiment when exiting the experiment application context (only relevant if using the Experiment as a context manager)
+            Whether to log the end of the experiment when cleaning up the experiment
         """
 
         self.server_host = server_host
@@ -117,15 +119,35 @@ def __init__(
             raise Exception(
                 "Timed out while attempting to connect with WEI server. Check that your server is running and the server_host and server_port are correct."
             )
+        self.check_in()
+
+    @threaded_daemon
+    def check_in(self):
+        """Checks in with the server to let it know the experiment is still running"""
+        while True:
+            time.sleep(10)
+            try:
+                response = requests.post(
+                    f"http://{self.server_host}:{self.server_port}/experiments/{self.experiment_id}/check_in"
+                )
+                if not response.ok:
+                    response.raise_for_status()
+            except Exception:
+                traceback.print_exc()
+                pass
+
+    def __del__(self):
+        """Logs the end of the experiment when cleaning up the experiment, if log_experiment_end_on_exit is True"""
+        if self.log_experiment_end_on_exit:
+            self.log_experiment_end()
 
     def __enter__(self):
         """Creates the experiment application context"""
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         """Logs the end of the experiment when exiting the experiment application context, if log_experiment_end_on_exit is True"""
-        if self.log_experiment_end_on_exit:
-            self.log_experiment_end()
+        pass
 
     def _register_experiment(self) -> None:
         """Registers a new experiment with the server

src/wei/routers/experiment_routes.py

@@ -3,6 +3,7 @@
 """
 
 import json
+from datetime import datetime
 from typing import Dict
 
 from fastapi import APIRouter
@@ -27,6 +28,17 @@ async def event_return(experiment_id: str) -> Dict[str, Event]:
     return events
 
 
+@router.post("/{experiment_id}/check_in")
+async def check_in(experiment_id: str) -> None:
+    """Returns a simple pong response"""
+    experiment = state_manager.get_experiment(experiment_id)
+    experiment.check_in_timestamp = datetime.now()
+    with state_manager.wc_state_lock():
+        state_manager.set_experiment(experiment)
+
+    return
+
+
 @router.get("/{experiment_id}/log")
 async def log_return(experiment_id: str) -> str:
     """Returns the log for a given experiment"""

src/wei/types/experiment_types.py

@@ -1,5 +1,6 @@
 """Types related to experiments"""
 
+from datetime import datetime
 from typing import List, Optional
 
 from pydantic import AliasChoices, Field
@@ -37,3 +38,6 @@ class Experiment(ExperimentDesign):
     experiment_id: str = Field(default_factory=ulid_factory)
     """ID of the experiment"""
     experiment_directory: Optional[PathLike] = None
+    """The directory where the experiment is stored on disk"""
+    check_in_timestamp: Optional[datetime] = None
+    """The last time the experiment client checked in"""