experimaestro 1.11.1__py3-none-any.whl → 2.0.0b4__py3-none-any.whl

experimaestro/scheduler/services.py

@@ -1,9 +1,12 @@
 import abc
 from enum import Enum
 import functools
+import logging
 import threading
 from typing import Set
 
+logger = logging.getLogger(__name__)
+
 
 class ServiceListener:
     """A service listener"""
@@ -13,6 +16,12 @@ class ServiceListener:
 
 
 class ServiceState(Enum):
+    """State of a service lifecycle.
+
+    Services transition through these states:
+    STOPPED -> STARTING -> RUNNING -> STOPPING -> STOPPED
+    """
+
     STOPPED = 0
     STARTING = 1
     RUNNING = 2
@@ -24,27 +33,72 @@ class Service:
 
     Services can be associated with an experiment. They send
     notifications to service listeners.
+
+    To support restarting services from monitor mode, subclasses should
+    override :meth:`state_dict` to return the data needed to recreate
+    the service, and implement :meth:`from_state_dict` to recreate it.
     """
 
     id: str
     _state: ServiceState = ServiceState.STOPPED
 
     def __init__(self):
-        self.listeners: Set[ServiceListener] = set()
+        self._listeners: Set[ServiceListener] = set()
+        self._listeners_lock = threading.Lock()
+
+    def state_dict(self) -> dict:
+        """Return a dictionary representation for serialization.
+
+        Subclasses should override this to include any parameters needed
+        to recreate the service. The base implementation returns the
+        class module and name.
+
+        Returns:
+            Dict with '__class__' key and any additional kwargs.
+        """
+        return {
+            "__class__": f"{self.__class__.__module__}.{self.__class__.__name__}",
+        }
+
+    @staticmethod
+    def from_state_dict(data: dict) -> "Service":
+        """Recreate a service from a state dictionary.
+
+        Args:
+            data: Dictionary from :meth:`state_dict`
+
+        Returns:
+            A new Service instance, or raises if the class cannot be loaded.
+        """
+        import importlib
+
+        class_path = data.get("__class__")
+        if not class_path:
+            raise ValueError("Missing '__class__' in service state_dict")
+
+        module_name, class_name = class_path.rsplit(".", 1)
+        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__"}
+        return cls(**kwargs)
 
     def add_listener(self, listener: ServiceListener):
         """Adds a listener
 
         :param listener: The listener to add
         """
-        self.listeners.add(listener)
+        with self._listeners_lock:
+            self._listeners.add(listener)
 
     def remove_listener(self, listener: ServiceListener):
         """Removes a listener
 
         :param listener: The listener to remove
         """
-        self.listeners.remove(listener)
+        with self._listeners_lock:
+            self._listeners.discard(listener)
 
     def description(self):
         return ""
@@ -58,35 +112,147 @@ class Service:
         # Set the state
         self._state = state
 
-        for listener in self.listeners:
-            listener.service_state_changed(self)
+        # Notify listeners with thread-safe snapshot
+        with self._listeners_lock:
+            listeners_snapshot = list(self._listeners)
+
+        for listener in listeners_snapshot:
+            try:
+                listener.service_state_changed(self)
+            except Exception:
+                logger.exception("Error notifying listener %s", listener)
 
 
 class WebService(Service):
-    """Web service"""
+    """Base class for web-based experiment services.
+
+    Web services provide HTTP endpoints that can be accessed through the
+    experimaestro web interface. When an experiment is running with a port
+    configured, web services are automatically proxied through the main
+    experimaestro server.
+
+    To implement a web service:
+
+    1. Subclass ``WebService``
+    2. Set a unique ``id`` class attribute
+    3. Implement the :meth:`_serve` method to start your web server
+    4. Set ``self.url`` and call ``running.set()`` when ready
+    5. Optionally check ``self.should_stop()`` to handle graceful shutdown
+
+    Example::
+
+        class MyWebService(WebService):
+            id = "myservice"
+
+            def _serve(self, running: threading.Event):
+                # Start your web server
+                self.url = "http://localhost:8080"
+                running.set()
+                # Keep serving, checking for stop signal
+                while not self.should_stop():
+                    time.sleep(1)
+    """
 
     def __init__(self):
         super().__init__()
         self.url = None
+        self.thread = None
+        self._stop_event = threading.Event()
+
+    def should_stop(self) -> bool:
+        """Check if the service should stop.
+
+        Subclasses can call this in their _serve loop to check for
+        graceful shutdown requests.
+
+        :return: True if stop() has been called
+        """
+        return self._stop_event.is_set()
 
     def get_url(self):
+        """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.
+
+        :return: The URL where this service can be accessed
+        """
         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
 
         # Returns the URL
         return self.url
 
-    def stop(self):
-        ...
+    def stop(self, timeout: float = 2.0):
+        """Stop the web service.
+
+        This method signals the service to stop and waits for the thread
+        to terminate. If the thread doesn't stop gracefully within the
+        timeout, it attempts to forcefully terminate it.
+
+        :param timeout: Seconds to wait for graceful shutdown before forcing
+        """
+        if self.state == ServiceState.STOPPED:
+            return
+
+        self.state = ServiceState.STOPPING
+
+        # Signal the service to stop
+        self._stop_event.set()
+
+        # Wait for the thread to finish
+        if self.thread is not None and self.thread.is_alive():
+            self.thread.join(timeout=timeout)
+
+            # If thread is still alive, try to terminate it forcefully
+            if self.thread.is_alive():
+                self._force_stop_thread()
+
+        self.url = None
+        self.state = ServiceState.STOPPED
+
+    def _force_stop_thread(self):
+        """Attempt to forcefully stop the service thread.
+
+        This uses ctypes to raise an exception in the thread. It's not
+        guaranteed to work (e.g., if the thread is blocked in C code),
+        but it's the best we can do in Python.
+        """
+        import ctypes
+
+        if self.thread is None or not self.thread.is_alive():
+            return
+
+        thread_id = self.thread.ident
+        if thread_id is None:
+            return
+
+        # Raise SystemExit in the target thread
+        res = ctypes.pythonapi.PyThreadState_SetAsyncExc(
+            ctypes.c_ulong(thread_id), ctypes.py_object(SystemExit)
+        )
+
+        if res == 0:
+            # Thread ID was invalid
+            pass
+        elif res > 1:
+            # Multiple threads affected - reset
+            ctypes.pythonapi.PyThreadState_SetAsyncExc(
+                ctypes.c_ulong(thread_id), ctypes.c_long(0)
+            )
 
     def serve(self):
-        import threading
+        """Start the web service in a background thread.
 
+        This method creates a daemon thread that calls :meth:`_serve`.
+        """
         self.thread = threading.Thread(
             target=functools.partial(self._serve, self.running),
             name=f"service[{self.id}]",
@@ -95,9 +261,17 @@ class WebService(Service):
         self.thread.start()
 
     @abc.abstractmethod
-    def _server(self, running: threading.Event):
-        """Starts the web service
+    def _serve(self, running: threading.Event):
+        """Start the web server (implement in subclasses).
+
+        This method should:
+
+        1. Start your web server
+        2. Set ``self.url`` to the service URL
+        3. Call ``running.set()`` to signal readiness
+        4. Keep the server running (this runs in a background thread)
+        5. Optionally check ``self.should_stop()`` for graceful shutdown
 
-        :param running: signals that `self.url` is set
+        :param running: Event to signal when ``self.url`` is set
         """
         ...

experimaestro/scheduler/signal_handler.py

@@ -0,0 +1,32 @@
+import signal
+from typing import Set
+from experimaestro.scheduler import experiment
+from experimaestro.utils import logger
+
+
+class SignalHandler:
+    def __init__(self):
+        self.experiments: Set["experiment"] = set()
+        self.original_sigint_handler = None
+
+    def add(self, xp: "experiment"):
+        if not self.experiments:
+            self.original_sigint_handler = signal.getsignal(signal.SIGINT)
+
+            signal.signal(signal.SIGINT, self)
+
+        self.experiments.add(xp)
+
+    def remove(self, xp):
+        self.experiments.remove(xp)
+        if not self.experiments:
+            signal.signal(signal.SIGINT, self.original_sigint_handler)
+
+    def __call__(self, signum, frame):
+        """SIGINT signal handler"""
+        logger.warning("Signal received")
+        for xp in self.experiments:
+            xp.stop()
+
+
+SIGNAL_HANDLER = SignalHandler()

experimaestro/scheduler/state.py

@@ -5,7 +5,7 @@ from typing import Iterable, Optional, Type
 from experimaestro import Task
 
 from experimaestro.core.context import SerializationContext
-from experimaestro.scheduler.base import Job, JobDependency
+from experimaestro.scheduler.jobs import Job, JobDependency
 from experimaestro.settings import find_workspace
 from experimaestro.core.serialization import from_state_dict, save_definition
 

experimaestro/scheduler/state_db.py

@@ -0,0 +1,388 @@
+"""Database models for experiment state persistence
+
+This module provides peewee ORM models for storing job and service state
+in a workspace-level SQLite database. The workspace has a single database
+file (.experimaestro/workspace.db) with WAL mode enabled for concurrent
+read/write access.
+
+Key design:
+- One database per workspace at: workdir/.experimaestro/workspace.db
+- Experiments can be run multiple times, each run tracked separately
+- Jobs and services are scoped to (experiment_id, run_id)
+- Tags are scoped to (job_id, experiment_id, run_id) - fixes GH #128
+- Current state and progress stored in JobModel - no history tracking
+- Database instance is passed explicitly to avoid global state
+"""
+
+from pathlib import Path
+from peewee import (
+    Model,
+    SqliteDatabase,
+    CharField,
+    FloatField,
+    IntegerField,
+    TextField,
+    DateTimeField,
+    CompositeKey,
+    IntegrityError,
+    OperationalError,
+)
+from datetime import datetime
+import fasteners
+
+
+class BaseModel(Model):
+    """Base model for workspace database tables
+
+    Models are unbound by default. Use database.bind_ctx() when querying:
+
+        with workspace.workspace_db.bind_ctx([ExperimentModel, JobModel, ...]):
+            experiments = ExperimentModel.select()
+
+    Or use the convenience method bind_models() defined below.
+    """
+
+    class Meta:
+        database = None  # Unbound - will be bound when used
+
+
+class ExperimentModel(BaseModel):
+    """Experiment metadata - tracks experiment definitions
+
+    An experiment can be run multiple times. This table tracks the experiment
+    itself and points to the current/latest run.
+
+    Fields:
+        experiment_id: Unique identifier for the experiment
+        current_run_id: Points to the current/latest run (null if no runs yet)
+        created_at: When experiment was first created
+        updated_at: When experiment was last modified (for incremental queries)
+
+    Note: Experiment path is derivable: {workspace}/xp/{experiment_id}
+    """
+
+    experiment_id = CharField(primary_key=True)
+    current_run_id = CharField(null=True)
+    created_at = DateTimeField(default=datetime.now)
+    updated_at = DateTimeField(default=datetime.now, index=True)
+
+    class Meta:
+        table_name = "experiments"
+
+
+class ExperimentRunModel(BaseModel):
+    """Individual experiment runs
+
+    Each time an experiment is executed, a new run is created.
+    Runs are identified by (experiment_id, run_id) composite key.
+
+    run_id format: timestamp-based like "20250120_143022" or sequential counter
+
+    Fields:
+        experiment_id: ID of the experiment this run belongs to
+        run_id: Unique ID for this run (timestamp or sequential)
+        started_at: When this run started
+        ended_at: When this run completed (null if still active)
+        status: Run status (active, completed, failed, abandoned)
+    """
+
+    experiment_id = CharField(index=True)
+    run_id = CharField(index=True)
+    started_at = DateTimeField(default=datetime.now)
+    ended_at = DateTimeField(null=True)
+    status = CharField(default="active", index=True)
+
+    class Meta:
+        table_name = "experiment_runs"
+        primary_key = CompositeKey("experiment_id", "run_id")
+        indexes = ((("experiment_id", "started_at"), False),)  # For finding latest run
+
+
+class WorkspaceSyncMetadata(BaseModel):
+    """Workspace-level metadata for disk sync tracking
+
+    Single-row table to track when the last disk sync occurred.
+    Used to throttle sync operations and prevent excessive disk scanning.
+
+    Fields:
+        id: Always "workspace" (single row table)
+        last_sync_time: When last sync completed
+        sync_interval_minutes: Minimum interval between syncs
+    """
+
+    id = CharField(primary_key=True, default="workspace")
+    last_sync_time = DateTimeField(null=True)
+    sync_interval_minutes = IntegerField(default=5)
+
+    class Meta:
+        table_name = "workspace_sync_metadata"
+
+
+class JobModel(BaseModel):
+    """Job information linked to specific experiment run
+
+    Jobs are tied to a specific run of an experiment via (experiment_id, run_id).
+    The same job can appear in multiple runs with different states/tags.
+
+    Fields:
+        job_id: Unique identifier for the job (from task identifier)
+        experiment_id: ID of the experiment this job belongs to
+        run_id: ID of the run this job belongs to
+        task_id: Task class identifier
+        locator: Full task locator (identifier)
+        state: Current job state (e.g., "unscheduled", "waiting", "running", "done", "error")
+        failure_reason: Optional failure reason for error states (e.g., "TIMEOUT", "DEPENDENCY")
+        submitted_time: When job was submitted (Unix timestamp)
+        started_time: When job started running (Unix timestamp)
+        ended_time: When job finished (Unix timestamp)
+        progress: JSON-encoded list of progress updates
+        updated_at: When job was last modified (for incremental queries)
+
+    Note: Job path is derivable: {workspace}/jobs/{task_id}/{job_id}
+    Note: Tags are stored in separate JobTagModel table (run-scoped)
+    Note: Dependencies are NOT stored in DB (available in state.json only)
+    """
+
+    job_id = CharField(index=True)
+    experiment_id = CharField(index=True)
+    run_id = CharField(index=True)
+    task_id = CharField(index=True)
+    locator = CharField()
+    state = CharField(default="unscheduled", index=True)
+    failure_reason = CharField(null=True)
+    submitted_time = FloatField(null=True)
+    started_time = FloatField(null=True)
+    ended_time = FloatField(null=True)
+    progress = TextField(default="[]")
+    updated_at = DateTimeField(default=datetime.now, index=True)
+
+    class Meta:
+        table_name = "jobs"
+        primary_key = CompositeKey("job_id", "experiment_id", "run_id")
+        indexes = (
+            (
+                ("experiment_id", "run_id", "state"),
+                False,
+            ),  # Query jobs by run and state
+            (
+                ("experiment_id", "run_id", "task_id"),
+                False,
+            ),  # Query jobs by run and task
+            (
+                ("experiment_id", "run_id", "updated_at"),
+                False,
+            ),  # Query jobs by run and update time
+        )
+
+
+class JobTagModel(BaseModel):
+    """Job tags for efficient searching (fixes GH #128)
+
+    **FIX FOR GH ISSUE #128**: Tags are now experiment-run-dependent, not job-dependent.
+    The same job in different experiment runs can have different tags, because tags
+    are scoped to the (job_id, experiment_id, run_id) combination.
+
+    Tags are stored as key-value pairs in a separate table for efficient indexing.
+    Each job can have multiple tags within an experiment run context.
+
+    Key change from old behavior:
+    - OLD: Tags were global per job_id (broken - same job in different experiments/runs shared tags)
+    - NEW: Tags are scoped per (job_id, experiment_id, run_id) - same job can have different tags in different runs
+
+    Fields:
+        job_id: ID of the job
+        experiment_id: ID of the experiment
+        run_id: ID of the run
+        tag_key: Tag name
+        tag_value: Tag value
+    """
+
+    job_id = CharField(index=True)
+    experiment_id = CharField(index=True)
+    run_id = CharField(index=True)
+    tag_key = CharField(index=True)
+    tag_value = CharField(index=True)
+
+    class Meta:
+        table_name = "job_tags"
+        primary_key = CompositeKey("job_id", "experiment_id", "run_id", "tag_key")
+        indexes = (
+            (("tag_key", "tag_value"), False),  # For tag-based queries
+            (
+                ("experiment_id", "run_id", "tag_key"),
+                False,
+            ),  # For experiment run tag queries
+        )
+
+
+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).
+
+    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
+    """
+
+    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"
+        primary_key = CompositeKey("service_id", "experiment_id", "run_id")
+
+
+class PartialModel(BaseModel):
+    """Partial directory tracking for subparameters
+
+    Tracks partial directories that are shared across jobs with different
+    parameter values (but same partial identifier). These directories are
+    at WORKSPACE/partials/TASK_ID/SUBPARAM_NAME/PARTIAL_ID/ (reconstructible).
+
+    Fields:
+        partial_id: Hex hash of the partial identifier
+        task_id: Task class identifier
+        subparameters_name: Name of the subparameters definition
+        created_at: When this partial directory was first created
+    """
+
+    partial_id = CharField(primary_key=True)
+    task_id = CharField(index=True)
+    subparameters_name = CharField(index=True)
+    created_at = DateTimeField(default=datetime.now)
+
+    class Meta:
+        table_name = "partials"
+        indexes = ((("task_id", "subparameters_name"), False),)
+
+
+class JobPartialModel(BaseModel):
+    """Links jobs to partial directories they use
+
+    Tracks which jobs reference which partial directories. This enables
+    cleanup of orphan partials when all referencing jobs are deleted.
+
+    A job can use multiple partials (different subparameters definitions),
+    and a partial can be used by multiple jobs.
+
+    Fields:
+        job_id: ID of the job using this partial
+        experiment_id: ID of the experiment
+        run_id: ID of the run
+        partial_id: ID of the partial directory being used
+    """
+
+    job_id = CharField(index=True)
+    experiment_id = CharField(index=True)
+    run_id = CharField(index=True)
+    partial_id = CharField(index=True)
+
+    class Meta:
+        table_name = "job_partials"
+        primary_key = CompositeKey("job_id", "experiment_id", "run_id", "partial_id")
+        indexes = ((("partial_id",), False),)  # For finding jobs using a partial
+
+
+# List of all models for binding
+ALL_MODELS = [
+    ExperimentModel,
+    ExperimentRunModel,
+    WorkspaceSyncMetadata,
+    JobModel,
+    JobTagModel,
+    ServiceModel,
+    PartialModel,
+    JobPartialModel,
+]
+
+
+def initialize_workspace_database(
+    db_path: Path, read_only: bool = False
+) -> SqliteDatabase:
+    """Initialize a workspace database connection with proper configuration
+
+    Creates and configures a SQLite database connection for the workspace.
+    Models must be bound to this database before querying.
+
+    Uses file-based locking to prevent multiple processes from initializing
+    the database simultaneously, which could cause SQLite locking issues.
+
+    Args:
+        db_path: Path to the workspace SQLite database file
+        read_only: If True, open database in read-only mode
+
+    Returns:
+        Configured SqliteDatabase instance
+    """
+    # Ensure parent directory exists (unless read-only)
+    if not read_only:
+        db_path.parent.mkdir(parents=True, exist_ok=True)
+
+    # Use file-based lock to prevent concurrent initialization from multiple processes
+    # This prevents SQLite locking issues during table creation
+    lock_path = db_path.parent / f".{db_path.name}.init.lock"
+    lock = fasteners.InterProcessLock(str(lock_path))
+
+    # Acquire lock (blocking) - only one process can initialize at a time
+    with lock:
+        # Create database connection
+        # check_same_thread=False allows the connection to be used from multiple threads
+        # This is safe with WAL mode and proper locking
+        db = SqliteDatabase(
+            str(db_path),
+            pragmas={
+                "journal_mode": "wal",  # Write-Ahead Logging for concurrent reads
+                "foreign_keys": 1,  # Enable foreign key constraints
+                "ignore_check_constraints": 0,
+                "synchronous": 1,  # NORMAL mode (balance safety/speed)
+                "busy_timeout": 5000,  # Wait up to 5 seconds for locks
+            },
+            check_same_thread=False,
+        )
+
+        if read_only:
+            # Set query-only mode for read-only access
+            db.execute_sql("PRAGMA query_only = ON")
+
+        # Bind all models to this database
+        db.bind(ALL_MODELS)
+
+        # Create tables if they don't exist (only in write mode)
+        if not read_only:
+            db.create_tables(ALL_MODELS, safe=True)
+
+            # 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},
+                )
+            except (IntegrityError, OperationalError):
+                # If get_or_create fails, the row likely already exists
+                pass
+
+    return db
+
+
+def close_workspace_database(db: SqliteDatabase):
+    """Close a workspace database connection
+
+    Args:
+        db: The database connection to close
+    """
+    if db and not db.is_closed():
+        db.close()