experimaestro 2.0.0b4__py3-none-any.whl → 2.0.0b8__py3-none-any.whl

experimaestro/scheduler/services.py

@@ -1,9 +1,14 @@
 import abc
 from enum import Enum
-import functools
 import logging
 import threading
-from typing import Set
+from pathlib import Path
+from typing import Callable, Optional, Set, TYPE_CHECKING
+
+from experimaestro.scheduler.interfaces import BaseService
+
+if TYPE_CHECKING:
+    from experimaestro.scheduler.experiment import Experiment
 
 logger = logging.getLogger(__name__)
 
@@ -28,7 +33,7 @@ class ServiceState(Enum):
     STOPPING = 3
 
 
-class Service:
+class Service(BaseService):
     """An experiment service
 
     Services can be associated with an experiment. They send
@@ -46,32 +51,89 @@ class Service:
         self._listeners: Set[ServiceListener] = set()
         self._listeners_lock = threading.Lock()
 
+    def set_experiment(self, xp: "Experiment") -> None:
+        """Called when the service is added to an experiment.
+
+        Override this method to access the experiment context (e.g., workdir).
+        The default implementation does nothing.
+
+        Args:
+            xp: The experiment this service is being added to.
+        """
+        pass
+
     def state_dict(self) -> dict:
-        """Return a dictionary representation for serialization.
+        """Return parameters needed to recreate this service.
+
+        Subclasses should override this to return constructor arguments.
+        Path values are automatically serialized and restored (with
+        translation for remote monitoring).
+
+        Example::
+
+            def state_dict(self):
+                return {
+                    "log_dir": self.log_dir,  # Path is auto-handled
+                    "name": self.name,
+                }
+
+        Returns:
+            Dict with constructor kwargs (no need to include __class__).
+        """
+        return {}
 
-        Subclasses should override this to include any parameters needed
-        to recreate the service. The base implementation returns the
-        class module and name.
+    def _full_state_dict(self) -> dict:
+        """Get complete state_dict including __class__ for serialization."""
+        d = self.state_dict()
+        d["__class__"] = f"{self.__class__.__module__}.{self.__class__.__name__}"
+        return d
+
+    @staticmethod
+    def serialize_state_dict(data: dict) -> dict:
+        """Serialize a state_dict, converting Path objects to serializable format.
+
+        This is called automatically when storing services. Path values are
+        converted to {"__path__": "/path/string"} format.
+
+        Args:
+            data: Raw state_dict from service (should include __class__)
 
         Returns:
-            Dict with '__class__' key and any additional kwargs.
+            Serializable dictionary with paths converted
         """
-        return {
-            "__class__": f"{self.__class__.__module__}.{self.__class__.__name__}",
-        }
+        result = {}
+        for k, v in data.items():
+            if isinstance(v, Path):
+                result[k] = {"__path__": str(v)}
+            else:
+                result[k] = v
+        return result
 
     @staticmethod
-    def from_state_dict(data: dict) -> "Service":
+    def from_state_dict(
+        data: dict, path_translator: Optional[Callable[[str], Path]] = None
+    ) -> "Service":
         """Recreate a service from a state dictionary.
 
         Args:
-            data: Dictionary from :meth:`state_dict`
+            data: Dictionary from :meth:`state_dict` (may be serialized)
+            path_translator: Optional function to translate remote paths to local.
+                Used by remote clients to map paths to local cache.
 
         Returns:
             A new Service instance, or raises if the class cannot be loaded.
+
+        Raises:
+            ValueError: If __unserializable__ is True or __class__ is missing
         """
         import importlib
 
+        # Check if service is marked as unserializable
+        if data.get("__unserializable__"):
+            raise ValueError(
+                f"Service cannot be recreated: {data.get('__reason__', 'unknown reason')}"
+            )
+
         class_path = data.get("__class__")
         if not class_path:
             raise ValueError("Missing '__class__' in service state_dict")
@@ -80,8 +142,22 @@ class Service:
         module = importlib.import_module(module_name)
         cls = getattr(module, class_name)
 
-        # Remove __class__ and pass remaining as kwargs
-        kwargs = {k: v for k, v in data.items() if k != "__class__"}
+        # Build kwargs, detecting and translating paths automatically
+        kwargs = {}
+        for k, v in data.items():
+            if k.startswith("__"):
+                continue  # Skip special keys
+            if isinstance(v, dict) and "__path__" in v:
+                # Serialized path - deserialize with optional translation
+                path_str = v["__path__"]
+                if path_translator:
+                    kwargs[k] = path_translator(path_str)
+                else:
+                    kwargs[k] = Path(path_str)
+            else:
+                kwargs[k] = v
+
+        logger.debug("Creating %s with kwargs: %s", cls.__name__, kwargs)
         return cls(**kwargs)
 
     def add_listener(self, listener: ServiceListener):
@@ -158,6 +234,8 @@ class WebService(Service):
         self.url = None
         self.thread = None
         self._stop_event = threading.Event()
+        self._start_lock = threading.Lock()
+        self._running_event: Optional[threading.Event] = None
 
     def should_stop(self) -> bool:
         """Check if the service should stop.
@@ -173,21 +251,46 @@ class WebService(Service):
         """Get the URL of this web service, starting it if needed.
 
         If the service is not running, this method will start it and
-        block until the URL is available.
+        block until the URL is available. If the service is already
+        starting or running, returns the existing URL.
 
         :return: The URL where this service can be accessed
+        :raises RuntimeError: If called while service is stopping
         """
-        if self.state == ServiceState.STOPPED:
-            self._stop_event.clear()
-            self.state = ServiceState.STARTING
-            self.running = threading.Event()
-            self.serve()
-
-        # Wait until the server is ready
-        self.running.wait()
-        self.state = ServiceState.RUNNING
+        with self._start_lock:
+            if self.state == ServiceState.STOPPING:
+                raise RuntimeError("Cannot start service while it is stopping")
+
+            if self.state == ServiceState.RUNNING:
+                logger.debug("Service already running, returning existing URL")
+                return self.url
+
+            if self.state == ServiceState.STOPPED:
+                logger.info(
+                    "Starting service %s (id=%s)", self.__class__.__name__, id(self)
+                )
+                self._stop_event.clear()
+                self.state = ServiceState.STARTING
+                self._running_event = threading.Event()
+                self.serve()
+            else:
+                logger.info(
+                    "Service %s (id=%s) already starting, waiting for it",
+                    self.__class__.__name__,
+                    id(self),
+                )
+
+            # State is STARTING - wait for it to be ready
+            running_event = self._running_event
+
+        # Wait outside the lock to avoid blocking other callers
+        if running_event:
+            running_event.wait()
+            # Set state to RUNNING (this will notify listeners)
+            with self._start_lock:
+                if self.state == ServiceState.STARTING:
+                    self.state = ServiceState.RUNNING
 
-        # Returns the URL
         return self.url
 
     def stop(self, timeout: float = 2.0):
@@ -199,10 +302,21 @@ class WebService(Service):
 
         :param timeout: Seconds to wait for graceful shutdown before forcing
         """
-        if self.state == ServiceState.STOPPED:
-            return
+        with self._start_lock:
+            if self.state == ServiceState.STOPPED:
+                return
+
+            if self.state == ServiceState.STARTING:
+                # Wait for service to finish starting before stopping
+                running_event = self._running_event
+            else:
+                running_event = None
 
-        self.state = ServiceState.STOPPING
+            self.state = ServiceState.STOPPING
+
+        # Wait for starting to complete if needed (outside lock to avoid deadlock)
+        if running_event is not None:
+            running_event.wait()
 
         # Signal the service to stop
         self._stop_event.set()
@@ -215,8 +329,10 @@ class WebService(Service):
             if self.thread.is_alive():
                 self._force_stop_thread()
 
-        self.url = None
-        self.state = ServiceState.STOPPED
+        with self._start_lock:
+            self.url = None
+            self._running_event = None
+            self.state = ServiceState.STOPPED
 
     def _force_stop_thread(self):
         """Attempt to forcefully stop the service thread.
@@ -254,12 +370,22 @@ class WebService(Service):
         This method creates a daemon thread that calls :meth:`_serve`.
         """
         self.thread = threading.Thread(
-            target=functools.partial(self._serve, self.running),
+            target=self._serve_wrapper,
             name=f"service[{self.id}]",
         )
         self.thread.daemon = True
         self.thread.start()
 
+    def _serve_wrapper(self):
+        """Wrapper for _serve that handles state transitions."""
+        running_event = self._running_event
+        try:
+            self._serve(running_event)
+        finally:
+            # Ensure the event is set even if _serve fails
+            if running_event and not running_event.is_set():
+                running_event.set()
+
     @abc.abstractmethod
     def _serve(self, running: threading.Event):
         """Start the web server (implement in subclasses).

experimaestro/scheduler/state_db.py

@@ -14,7 +14,9 @@ Key design:
 - Database instance is passed explicitly to avoid global state
 """
 
+import logging
 from pathlib import Path
+from typing import Tuple
 from peewee import (
     Model,
     SqliteDatabase,
@@ -30,6 +32,11 @@ from peewee import (
 from datetime import datetime
 import fasteners
 
+logger = logging.getLogger("xpm.state_db")
+
+# Database schema version - increment when schema changes require resync
+CURRENT_DB_VERSION = 3
+
 
 class BaseModel(Model):
     """Base model for workspace database tables
@@ -84,6 +91,7 @@ class ExperimentRunModel(BaseModel):
         started_at: When this run started
         ended_at: When this run completed (null if still active)
         status: Run status (active, completed, failed, abandoned)
+        hostname: Host where the experiment was launched (null for old runs)
     """
 
     experiment_id = CharField(index=True)
@@ -91,6 +99,7 @@ class ExperimentRunModel(BaseModel):
     started_at = DateTimeField(default=datetime.now)
     ended_at = DateTimeField(null=True)
     status = CharField(default="active", index=True)
+    hostname = CharField(null=True)
 
     class Meta:
         table_name = "experiment_runs"
@@ -108,11 +117,13 @@ class WorkspaceSyncMetadata(BaseModel):
         id: Always "workspace" (single row table)
         last_sync_time: When last sync completed
         sync_interval_minutes: Minimum interval between syncs
+        db_version: Schema version for migration detection
     """
 
     id = CharField(primary_key=True, default="workspace")
     last_sync_time = DateTimeField(null=True)
     sync_interval_minutes = IntegerField(default=5)
+    db_version = IntegerField(default=1)
 
     class Meta:
         table_name = "workspace_sync_metadata"
@@ -219,26 +230,23 @@ class ServiceModel(BaseModel):
     """Service information linked to specific experiment run
 
     Services are tied to a specific run of an experiment via (experiment_id, run_id).
+    Services are only added or removed, not updated - state is managed at runtime.
 
     Fields:
         service_id: Unique identifier for the service
         experiment_id: ID of the experiment this service belongs to
         run_id: ID of the run this service belongs to
         description: Human-readable description
-        state: Service state (e.g., "running", "stopped")
         state_dict: JSON serialized state_dict for service recreation
-        created_at: When service was created
-        updated_at: Timestamp of last update
+        created_at: When service was registered
     """
 
     service_id = CharField()
     experiment_id = CharField(index=True)
     run_id = CharField(index=True)
     description = TextField(default="")
-    state = CharField()
     state_dict = TextField(default="{}")  # JSON for service recreation
     created_at = DateTimeField(default=datetime.now)
-    updated_at = DateTimeField(default=datetime.now)
 
     class Meta:
         table_name = "services"
@@ -311,7 +319,7 @@ ALL_MODELS = [
 
 def initialize_workspace_database(
     db_path: Path, read_only: bool = False
-) -> SqliteDatabase:
+) -> Tuple[SqliteDatabase, bool]:
     """Initialize a workspace database connection with proper configuration
 
     Creates and configures a SQLite database connection for the workspace.
@@ -325,7 +333,9 @@ def initialize_workspace_database(
         read_only: If True, open database in read-only mode
 
     Returns:
-        Configured SqliteDatabase instance
+        Tuple of (SqliteDatabase instance, needs_resync flag)
+        The needs_resync flag is True when the database schema version is outdated
+        and a full resync from disk is required.
     """
     # Ensure parent directory exists (unless read-only)
     if not read_only:
@@ -336,6 +346,8 @@ def initialize_workspace_database(
     lock_path = db_path.parent / f".{db_path.name}.init.lock"
     lock = fasteners.InterProcessLock(str(lock_path))
 
+    needs_resync = False
+
     # Acquire lock (blocking) - only one process can initialize at a time
     with lock:
         # Create database connection
@@ -364,18 +376,55 @@ def initialize_workspace_database(
         if not read_only:
             db.create_tables(ALL_MODELS, safe=True)
 
+            # Check database version for migration - use raw SQL since column may not exist
+            current_version = 0
+            try:
+                cursor = db.execute_sql(
+                    "SELECT db_version FROM workspace_sync_metadata WHERE id='workspace'"
+                )
+                row = cursor.fetchone()
+                if row is not None:
+                    current_version = row[0]
+                if current_version < CURRENT_DB_VERSION:
+                    needs_resync = True
+            except OperationalError:
+                # Column doesn't exist - add it and trigger resync
+                needs_resync = True
+                try:
+                    db.execute_sql(
+                        "ALTER TABLE workspace_sync_metadata "
+                        "ADD COLUMN db_version INTEGER DEFAULT 1"
+                    )
+                except OperationalError:
+                    pass  # Column may already exist
+
+            # Run schema migrations for older databases
+            if current_version < 2:
+                # Migration v1 -> v2: Add hostname column to experiment_runs table
+                try:
+                    db.execute_sql(
+                        "ALTER TABLE experiment_runs ADD COLUMN hostname VARCHAR(255) NULL"
+                    )
+                    logger.info("Added hostname column to experiment_runs table")
+                except OperationalError:
+                    pass  # Column already exists
+
             # Initialize WorkspaceSyncMetadata with default row if not exists
             # Use try/except to handle race condition (shouldn't happen with lock, but be safe)
             try:
                 WorkspaceSyncMetadata.get_or_create(
                     id="workspace",
-                    defaults={"last_sync_time": None, "sync_interval_minutes": 5},
+                    defaults={
+                        "last_sync_time": None,
+                        "sync_interval_minutes": 5,
+                        "db_version": 1,
+                    },
                 )
             except (IntegrityError, OperationalError):
                 # If get_or_create fails, the row likely already exists
                 pass
 
-    return db
+    return db, needs_resync
 
 
 def close_workspace_database(db: SqliteDatabase):