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

experimaestro/scheduler/experiment.py

@@ -1,7 +1,9 @@
 import asyncio
+import json
 import logging
 import os
 from pathlib import Path
+import time
 from shutil import rmtree
 from typing import Any, Dict, Optional, TypeVar, Union
 
@@ -9,7 +11,7 @@ from experimaestro.core.objects import WatchedOutput
 from experimaestro.exceptions import HandledException
 
 from experimaestro.scheduler.signal_handler import SIGNAL_HANDLER
-from experimaestro.scheduler.jobs import Job, JobFailureStatus
+from experimaestro.scheduler.jobs import Job
 from experimaestro.scheduler.services import Service
 from experimaestro.scheduler.workspace import RunMode, Workspace
 from experimaestro.settings import WorkspaceSettings, get_settings
@@ -24,15 +26,54 @@ class FailedExperiment(HandledException):
     pass
 
 
+class DatabaseListener:
+    """Listener that updates job state in the database"""
+
+    def __init__(self, state_provider, experiment_id: str, run_id: str):
+        self.state_provider = state_provider
+        self.experiment_id = experiment_id
+        self.run_id = run_id
+
+    def job_submitted(self, job):
+        # Already handled in experiment.add_job()
+        pass
+
+    def job_state(self, job):
+        """Update job state in database"""
+        self.state_provider.update_job_state(job, self.experiment_id, self.run_id)
+
+    def service_add(self, service):
+        """Register service in database"""
+        from experimaestro.scheduler.services import Service
+
+        state_dict = Service.serialize_state_dict(service._full_state_dict())
+        self.state_provider.register_service(
+            service.id,
+            self.experiment_id,
+            self.run_id,
+            service.description(),
+            state_dict=json.dumps(state_dict),
+        )
+
+    def service_state_changed(self, service):
+        """Called when service state changes (runtime only, not persisted)"""
+        # Service state is managed at runtime, not persisted to DB
+        pass
+
+
 class experiment:
-    """Main experiment object
+    """Context manager for running experiments.
+
+    Creates a workspace, manages task submission, and optionally starts
+    a web server for monitoring.
 
-    It is a context object, i.e. an experiment is run with
+    Example::
 
-    ```py
-        with experiment(...) as xp:
-            ...
-    ```
+        from experimaestro import experiment
+
+        with experiment("./workdir", "my-experiment", port=12345) as xp:
+            task = MyTask.C(param=42).submit()
+            result = task.wait()
     """
 
     #: Current experiment
@@ -57,6 +98,7 @@ class experiment:
         token: Optional[str] = None,
         run_mode: Optional[RunMode] = None,
         launcher=None,
+        register_signals: bool = True,
     ):
         """
         :param env: an environment -- or a working directory for a local
@@ -73,9 +115,11 @@ class experiment:
 
         :param run_mode: The run mode for the experiment (normal, generate run
             files, dry run)
+
+        :param register_signals: Whether to register signal handlers (default: True).
+            Set to False when running in a background thread.
         """
 
-        from experimaestro.server import Server
         from experimaestro.scheduler import Listener, Scheduler
 
         settings = get_settings()
@@ -94,6 +138,7 @@ class experiment:
         self.old_experiment = None
         self.services: Dict[str, Service] = {}
         self._job_listener: Optional[Listener] = None
+        self._register_signals = register_signals
 
         # Get configuration settings
 
@@ -106,14 +151,14 @@ class experiment:
         if token is not None:
             settings.server.token = token
 
-        # Create the scheduler
-        self.scheduler = Scheduler.create(self, name)
-        self.server = (
-            Server(self.scheduler, settings.server)
-            if (settings.server.port is not None and settings.server.port >= 0)
-            and self.workspace.run_mode == RunMode.NORMAL
-            else None
-        )
+        # Use singleton scheduler
+        self.scheduler = Scheduler.instance()
+
+        # Determine if we need a server
+        self._needs_server = (
+            settings.server.port is not None and settings.server.port >= 0
+        ) and self.workspace.run_mode == RunMode.NORMAL
+        self._server_settings = settings.server if self._needs_server else None
 
         if os.environ.get("XPM_ENABLEFAULTHANDLER", "0") == "1":
             import faulthandler
@@ -137,6 +182,11 @@ class experiment:
         assert self.scheduler is not None, "No scheduler defined"
         return self.scheduler.loop
 
+    @property
+    def server(self):
+        """Access the server via the scheduler"""
+        return self.scheduler.server if self.scheduler else None
+
     @property
     def resultspath(self):
         """Return the directory in which results can be stored for this experiment"""
@@ -158,6 +208,83 @@ class experiment:
         """Return the directory in which results can be stored for this experiment"""
         return self.workdir / "jobs.bak"
 
+    @property
+    def jobs_jsonl_path(self):
+        """Return the path to the jobs.jsonl file for this experiment"""
+        return self.workdir / "jobs.jsonl"
+
+    @property
+    def services_json_path(self):
+        """Return the path to the services.json file for this experiment"""
+        return self.workdir / "services.json"
+
+    def _write_services_json(self):
+        """Write all services to services.json file"""
+        from experimaestro.scheduler.services import Service
+
+        services_data = {}
+        for service_id, service in self.services.items():
+            # Get state_dict from service (includes __class__ for recreation)
+            # and serialize paths to JSON-compatible format
+            service_state = Service.serialize_state_dict(service._full_state_dict())
+            # Add runtime state info
+            service_state.update(
+                {
+                    "service_id": service_id,
+                    "description": service.description(),
+                    "state": service.state.name,
+                    "url": getattr(service, "url", None),
+                    "timestamp": time.time(),
+                }
+            )
+            services_data[service_id] = service_state
+
+        with self.services_json_path.open("w") as f:
+            json.dump(services_data, f, indent=2)
+
+    def add_job(self, job: "Job"):
+        """Register a job and its tags to jobs.jsonl file and database
+
+        Note: For NEW jobs, the unfinishedJobs counter is updated by
+        job.set_state() when the state transitions from UNSCHEDULED.
+        For jobs already running, we increment here since no state
+        transition will occur.
+        """
+        from experimaestro.scheduler.interfaces import JobState
+
+        if self in job.experiments:
+            # Do not double register
+            return
+
+        # Track which experiments this job belongs to
+        job.experiments.append(self)
+
+        # If job is already being tracked (not UNSCHEDULED and not finished),
+        # increment unfinishedJobs since no state transition will trigger it
+        if job.state != JobState.UNSCHEDULED and not job.state.finished():
+            self.unfinishedJobs += 1
+            logging.debug(
+                "Job %s already running, unfinished jobs for %s: %d",
+                job.identifier[:8],
+                self.workdir.name,
+                self.unfinishedJobs,
+            )
+
+        record = {
+            "job_id": job.identifier,
+            "task_id": str(job.type.identifier),
+            "tags": dict(job.tags.items()) if job.tags else {},
+            "timestamp": time.time(),
+        }
+
+        with self.jobs_jsonl_path.open("a") as f:
+            f.write(json.dumps(record) + "\n")
+
+        # Also register in database for TUI/monitoring (only in NORMAL mode)
+        if self._db_listener is not None:
+            experiment_id = self.workdir.name
+            self.state_provider.update_job_submitted(job, experiment_id, self.run_id)
+
     def stop(self):
         """Stop the experiment as soon as possible"""
 
@@ -196,9 +323,24 @@ class experiment:
 
                 if self.failedJobs:
                     # Show some more information
+                    from experimaestro.scheduler.jobs import (
+                        JobStateError,
+                        JobFailureStatus,
+                    )
+
                     count = 0
                     for job in self.failedJobs.values():
-                        if job.failure_status != JobFailureStatus.DEPENDENCY:
+                        # Skip dependency failures - only log direct failures
+                        if isinstance(job.state, JobStateError):
+                            if job.state.failure_reason != JobFailureStatus.DEPENDENCY:
+                                count += 1
+                                logger.error(
+                                    "Job %s failed, check the log file %s",
+                                    job.relpath,
+                                    job.stderr,
+                                )
+                        else:
+                            # Should not happen, but count it anyway
                             count += 1
                             logger.error(
                                 "Job %s failed, check the log file %s",
@@ -224,12 +366,18 @@ class experiment:
 
     def __enter__(self):
         from .dynamic_outputs import TaskOutputsWorker
+        from experimaestro.utils.environment import save_environment_info
 
         if self.workspace.run_mode != RunMode.DRY_RUN:
             logger.info("Locking experiment %s", self.xplockpath)
             self.xplock = self.workspace.connector.lock(self.xplockpath, 0).__enter__()
             logger.info("Experiment locked")
 
+        # Capture and save environment info (git info for editable packages + all package versions)
+        if self.workspace.run_mode == RunMode.NORMAL:
+            env_info_path = self.workdir / "environment.json"
+            save_environment_info(env_info_path)
+
         # Move old jobs into "jobs.bak"
         if self.workspace.run_mode == RunMode.NORMAL:
             self.jobsbakpath.mkdir(exist_ok=True)
@@ -244,12 +392,43 @@ class experiment:
                         target.parent.mkdir(parents=True, exist_ok=True)
                         p.rename(target)
 
-        if self.server:
-            self.server.start()
+        # Register experiment with scheduler
+        self.scheduler.register_experiment(self)
+
+        # Start server via scheduler if needed
+        if self._needs_server:
+            self.scheduler.start_server(self._server_settings, workspace=self.workspace)
 
         self.workspace.__enter__()
         (self.workspace.path / ".__experimaestro__").touch()
 
+        # Initialize workspace state provider (singleton per workspace path)
+        # Use read_only mode when not in NORMAL run mode to prevent DB changes
+        from .state_provider import WorkspaceStateProvider
+
+        is_normal_mode = self.workspace.run_mode == RunMode.NORMAL
+        self.state_provider = WorkspaceStateProvider.get_instance(
+            self.workspace.path,
+            read_only=not is_normal_mode,
+            sync_on_start=False,  # Experiments don't sync on start
+        )
+
+        # Register experiment in database and create a run (only in NORMAL mode)
+        experiment_id = self.workdir.name
+        self._db_listener = None
+        if is_normal_mode:
+            self.state_provider.ensure_experiment(experiment_id)
+            self.run_id = self.state_provider.create_run(experiment_id)
+
+            # Add database listener to update job state in database
+            self._db_listener = DatabaseListener(
+                self.state_provider, experiment_id, self.run_id
+            )
+            self.scheduler.addlistener(self._db_listener)
+        else:
+            # In non-NORMAL modes, use a placeholder run_id
+            self.run_id = None
+
         # Number of unfinished jobs
         self.unfinishedJobs = 0
         self.taskOutputQueueSize = 0
@@ -260,11 +439,12 @@ class experiment:
         # Exit mode when catching signals
         self.exitMode = False
 
-        self.scheduler.start_scheduler()
+        # Note: scheduler is already running as singleton
         self.taskOutputsWorker = TaskOutputsWorker(self)
         self.taskOutputsWorker.start()
 
-        SIGNAL_HANDLER.add(self)
+        if self._register_signals:
+            SIGNAL_HANDLER.add(self)
 
         self.old_experiment = experiment.CURRENT
         experiment.CURRENT = self
@@ -288,17 +468,33 @@ class experiment:
                 )
             else:
                 self.wait()
+
+            # Wait for all pending notifications to be processed
+            # before removing listeners
+            self.scheduler.wait_for_notifications()
         finally:
-            SIGNAL_HANDLER.remove(self)
+            if self._register_signals:
+                SIGNAL_HANDLER.remove(self)
 
             # Stop services
             for service in self.services.values():
                 logger.info("Closing service %s", service.description())
                 service.stop()
 
-            if self.scheduler is not None:
-                logger.info("Stopping scheduler event loop")
-                self.scheduler.loop.stop()
+            # Unregister experiment from scheduler
+            self.scheduler.unregister_experiment(self)
+
+            # Remove database listener and mark run as completed (only in NORMAL mode)
+            if self._db_listener is not None:
+                self.scheduler.removelistener(self._db_listener)
+
+                # Mark run as completed in database
+                experiment_id = self.workdir.name
+                status = "failed" if exc_type else "completed"
+                self.state_provider.complete_run(experiment_id, self.run_id, status)
+
+            # Note: Don't stop scheduler - it's shared!
+            # Note: Don't stop server - it runs in daemon mode until program exit
 
             if self.taskOutputsWorker is not None:
                 logger.info("Stopping tasks outputs worker")
@@ -310,9 +506,6 @@ class experiment:
 
             # Put back old experiment as current one
             experiment.CURRENT = self.old_experiment
-            if self.server:
-                logger.info("Stopping web server")
-                self.server.stop()
 
         if self.workspace.run_mode == RunMode.NORMAL:
             # Write the state
@@ -345,13 +538,45 @@ class experiment:
         """Adds a service (e.g. tensorboard viewer) to the experiment
 
         :param service: A service instance
-        :return: The same service instance
+        :return: The same service instance (or existing service if already added)
         """
+        existing = self.services.get(service.id)
+        if existing is not None:
+            if existing is service:
+                # Same service instance added twice - just return it
+                logger.debug("Service %s already added, ignoring duplicate", service.id)
+                return service
+            else:
+                # Different service with same id - warn and replace
+                logger.warning(
+                    "Replacing service %s (old id=%s, new id=%s)",
+                    service.id,
+                    id(existing),
+                    id(service),
+                )
+
         self.services[service.id] = service
-        for listener in self.scheduler.listeners:
-            listener.service_add(service)
+
+        # Allow service to access experiment context
+        service.set_experiment(self)
+
+        # Register database listener for state changes
+        service.add_listener(self._db_listener)
+
+        # Register file listener for state changes (writes to services.json)
+        service.add_listener(self)
+
+        self.scheduler.notify_service_add(service)
+
+        # Write services.json file
+        self._write_services_json()
+
         return service
 
+    def service_state_changed(self, service):
+        """Called when a service state changes - update services.json"""
+        self._write_services_json()
+
     def save(self, obj: Any, name: str = "default"):
         """Serializes configurations.