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

experimaestro/scheduler/interfaces.py

@@ -16,13 +16,91 @@ to enable unified access in the TUI and other monitoring tools.
 import enum
 import json
 import logging
+from dataclasses import dataclass, field
 from datetime import datetime
 from pathlib import Path
-from typing import Dict, List, Optional
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+
+if TYPE_CHECKING:
+    from experimaestro.scheduler.transient import TransientMode
 
 logger = logging.getLogger("xpm.interfaces")
 
 
+@dataclass
+class ExperimentJobInformation:
+    """Lightweight job information for experiment state serialization
+
+    This class contains the minimal job metadata stored in status.json and jobs.jsonl.
+    Full job state (progress, state changes, etc.) comes from events.jsonl replay
+    or from the state provider.
+    """
+
+    job_id: str
+    task_id: str
+    tags: Dict[str, str] = field(default_factory=dict)
+    timestamp: Optional[float] = None
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Serialize to dictionary for JSON"""
+        return {
+            "job_id": self.job_id,
+            "task_id": self.task_id,
+            "tags": self.tags,
+            "timestamp": self.timestamp,
+        }
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, Any]) -> "ExperimentJobInformation":
+        """Create from dictionary"""
+        return cls(
+            job_id=d["job_id"],
+            task_id=d["task_id"],
+            tags=d.get("tags", {}),
+            timestamp=d.get("timestamp"),
+        )
+
+
+def serialize_timestamp(ts: Optional[Union[float, datetime, str]]) -> Optional[str]:
+    """Serialize timestamp to ISO format string for DB/network storage
+
+    Handles:
+    - None: returns None
+    - float/int: Unix timestamp, converts to ISO format
+    - datetime: converts to ISO format
+    - str: returns as-is (already serialized)
+    """
+    if ts is None:
+        return None
+    if isinstance(ts, str):
+        return ts  # Already serialized
+    if isinstance(ts, (int, float)):
+        return datetime.fromtimestamp(ts).isoformat()
+    if isinstance(ts, datetime):
+        return ts.isoformat()
+    return str(ts)
+
+
+def deserialize_timestamp(ts: Optional[Union[float, str]]) -> Optional[float]:
+    """Deserialize timestamp from ISO format string to Unix timestamp
+
+    Handles:
+    - None: returns None
+    - float/int: returns as-is (already a Unix timestamp)
+    - str: parses ISO format and converts to Unix timestamp
+    """
+    if ts is None:
+        return None
+    if isinstance(ts, (int, float)):
+        return float(ts)
+    if isinstance(ts, str):
+        try:
+            return datetime.fromisoformat(ts).timestamp()
+        except ValueError:
+            return None
+    return None
+
+
 # =============================================================================
 # Job State Classes
 # =============================================================================
@@ -188,6 +266,19 @@ class JobFailureStatus(enum.Enum):
     TIMEOUT = 3
 
 
+class ExperimentStatus(enum.Enum):
+    """Status of an experiment run"""
+
+    #: Experiment is currently running
+    RUNNING = "running"
+
+    #: Experiment completed successfully
+    DONE = "done"
+
+    #: Experiment failed
+    FAILED = "failed"
+
+
 class JobStateError(JobState):
     """Job has failed
 
@@ -225,7 +316,7 @@ class JobStateError(JobState):
         return True
 
 
-# FIXME: Get rid of those
+# NOTE: Consider removing these singleton instances in a future refactor
 # Create singleton instances for backward compatibility
 # These can be used in comparisons: if state == JobState.DONE: ...
 JobState.UNSCHEDULED = JobStateUnscheduled()
@@ -264,30 +355,33 @@ class BaseJob:
     Attributes:
         identifier: Unique identifier for the job (hash)
         task_id: Task class identifier (string)
-        locator: Full task locator (identifier)
         path: Path to job directory
         state: Current job state (JobState object or compatible)
         submittime: When job was submitted (Unix timestamp or None)
         starttime: When job started running (Unix timestamp or None)
         endtime: When job finished (Unix timestamp or None)
         progress: List of progress updates
-        tags: Dictionary of tag key-value pairs
         exit_code: Process exit code (optional)
         retry_count: Number of retries
+        transient: Transient mode (NONE, TRANSIENT, or REMOVE)
     """
 
     identifier: str
     task_id: str
-    locator: str
     path: Path
     state: JobState
     submittime: Optional[float]
     starttime: Optional[float]
     endtime: Optional[float]
     progress: List[Dict]
-    tags: Dict[str, str]
     exit_code: Optional[int]
     retry_count: int
+    transient: "TransientMode"
+
+    @property
+    def locator(self) -> str:
+        """Full task locator (identifier): {task_id}/{identifier}"""
+        return f"{self.task_id}/{self.identifier}"
 
     # -------------------------------------------------------------------------
     # Static path computation (for use without a job instance)
@@ -304,9 +398,9 @@ class BaseJob:
         return job_path / ".experimaestro"
 
     @staticmethod
-    def get_metadata_path(job_path: Path) -> Path:
-        """Get metadata file path for a job path"""
-        return job_path / ".experimaestro" / "information.json"
+    def get_status_path(job_path: Path) -> Path:
+        """Get status file path for a job path"""
+        return job_path / ".experimaestro" / "status.json"
 
     @staticmethod
     def get_pidfile(job_path: Path, scriptname: str) -> Path:
@@ -338,9 +432,9 @@ class BaseJob:
         return BaseJob.get_xpm_dir(self.path)
 
     @property
-    def metadata_path(self) -> Path:
-        """Path to the job metadata file"""
-        return BaseJob.get_metadata_path(self.path)
+    def status_path(self) -> Path:
+        """Path to the job status file"""
+        return BaseJob.get_status_path(self.path)
 
     @property
     def pidfile(self) -> Path:
@@ -358,117 +452,265 @@ class BaseJob:
         return BaseJob.get_failedfile(self.path, self.scriptname)
 
     # -------------------------------------------------------------------------
-    # Metadata I/O
+    # State I/O (unified state_dict pattern)
     # -------------------------------------------------------------------------
 
-    def write_metadata(self, **extra_fields) -> None:
-        """Write or update job metadata in .experimaestro/information.json file
-
-        Automatically extracts metadata from job attributes (identifier, state,
-        submittime, starttime, endtime, retry_count) and writes to the metadata file.
+    def state_dict(self) -> Dict[str, Any]:
+        """Get job state as dictionary (single source of truth)
 
-        Performs atomic write using temp file + rename. If metadata exists,
-        new fields are merged with existing ones. Updates last_updated timestamp.
+        This is the canonical representation of job state used for both
+        serialization to status files and network communication.
 
-        Args:
-            **extra_fields: Optional extra fields (e.g., launcher, launcher_job_id, exit_code)
+        Returns:
+            Dictionary with all job state fields
         """
-        # Ensure .experimaestro directory exists
-        self.xpm_dir.mkdir(parents=True, exist_ok=True)
-        metadata_path = self.metadata_path
-
-        # Read existing metadata
-        existing = {}
-        if metadata_path.exists():
-            try:
-                with metadata_path.open("r") as f:
-                    existing = json.load(f)
-            except Exception as e:
-                logger.warning(
-                    "Failed to read existing metadata from %s: %s", metadata_path, e
-                )
-
-        # Build metadata from job attributes
-        fields = {
+        failure_reason = None
+        if (
+            self.state
+            and self.state.is_error()
+            and hasattr(self.state, "failure_reason")
+        ):
+            fr = self.state.failure_reason
+            if fr is not None:
+                failure_reason = fr.name
+
+        return {
             "job_id": self.identifier,
             "task_id": self.task_id,
+            "path": str(self.path) if self.path else None,
             "state": self.state.name if self.state else None,
+            "failure_reason": failure_reason,
+            "submitted_time": self.submittime,
+            "started_time": self.starttime,
+            "ended_time": self.endtime,
+            "exit_code": self.exit_code,
+            "retry_count": self.retry_count,
+            "progress": [
+                p.to_dict() if hasattr(p, "to_dict") else p
+                for p in (self.progress or [])
+            ],
+            "process": self.process_state_dict(),
         }
 
-        # Add timing information if available
-        if self.submittime is not None:
-            fields["submitted_time"] = self.submittime
-        if self.starttime is not None:
-            fields["started_time"] = self.starttime
-        if self.endtime is not None:
-            fields["ended_time"] = self.endtime
+    def process_state_dict(self) -> dict | None:
+        """Get process state as dictionary. Override in subclasses."""
+        return None
+
+
+# =============================================================================
+# Base Experiment Interface
+# =============================================================================
 
-        # Add exit code if available
-        if self.exit_code is not None:
-            fields["exit_code"] = self.exit_code
 
-        # Add retry count
-        if hasattr(self, "retry_count"):
-            fields["retry_count"] = self.retry_count
+class BaseExperiment:
+    """Base interface for experiment information
 
-        # Merge with extra fields (for launcher info, exit_code, etc.)
-        fields.update(extra_fields)
+    This class defines the interface for experiment data. Both live experiment
+    instances and MockExperiment instances should provide these attributes.
+
+    Core attributes:
+        workdir: Path to run directory (experiments/{exp-id}/{run-id}/)
+        run_id: Run identifier
+
+    State tracking (replaces StatusData):
+        jobs: Dict mapping job_id to BaseJob
+        services: Dict mapping service_id to BaseService
+        tags: Dict mapping job_id to tag dict
+        dependencies: Dict mapping job_id to list of dependency job_ids
+        events_count: Number of events processed
+        hostname: Hostname where experiment runs
+        started_at: Start timestamp
+        ended_at: End timestamp (None if running)
+    """
 
-        # Merge with existing and update timestamp
-        existing.update(fields)
-        existing["last_updated"] = datetime.now().timestamp()
+    # Status file version
+    STATUS_VERSION = 1
 
-        # Atomic write
-        temp_path = metadata_path.with_suffix(".json.tmp")
+    workdir: Path
+    run_id: str
+
+    @property
+    def experiment_id(self) -> str:
+        """Experiment identifier derived from workdir structure"""
+        # workdir is experiments/{exp-id}/{run-id}, so parent.name is exp-id
+        return self.workdir.parent.name
+
+    @property
+    def run_dir(self) -> Path:
+        """Path to run directory (same as workdir)"""
+        return self.workdir
+
+    @property
+    def status(self) -> "ExperimentStatus":
+        """Experiment status - override in subclasses"""
+        raise NotImplementedError
+
+    # State tracking properties (abstract - must be implemented by subclasses)
+
+    @property
+    def jobs(self) -> Dict[str, "BaseJob"]:
+        """Jobs in this experiment"""
+        raise NotImplementedError
+
+    @property
+    def services(self) -> Dict[str, "BaseService"]:
+        """Services in this experiment"""
+        raise NotImplementedError
+
+    @property
+    def tags(self) -> Dict[str, Dict[str, str]]:
+        """Tags for jobs"""
+        raise NotImplementedError
+
+    @property
+    def dependencies(self) -> Dict[str, List[str]]:
+        """Job dependencies"""
+        raise NotImplementedError
+
+    @property
+    def events_count(self) -> int:
+        """Number of events processed"""
+        raise NotImplementedError
+
+    @property
+    def hostname(self) -> Optional[str]:
+        """Hostname where experiment runs"""
+        raise NotImplementedError
+
+    @property
+    def started_at(self) -> Optional[float]:
+        """Start timestamp"""
+        raise NotImplementedError
+
+    @property
+    def ended_at(self) -> Optional[float]:
+        """End timestamp (None if running)"""
+        raise NotImplementedError
+
+    # Computed properties
+
+    @property
+    def total_jobs(self) -> int:
+        """Total number of jobs"""
+        return len(self.jobs)
+
+    @property
+    def finished_jobs(self) -> int:
+        """Number of finished jobs"""
+        return sum(1 for j in self.jobs.values() if j.state == JobState.DONE)
+
+    @property
+    def failed_jobs(self) -> int:
+        """Number of failed jobs"""
+        return sum(1 for j in self.jobs.values() if j.state.is_error())
+
+    def get_services(self) -> List["BaseService"]:
+        """Get services for this experiment as a list"""
+        return list(self.services.values())
+
+    @staticmethod
+    def get_status_path(run_dir: Path) -> Path:
+        """Get status file path for a run directory"""
+        return run_dir / "status.json"
+
+    def state_dict(self) -> Dict[str, Any]:
+        """Get experiment state as dictionary (single source of truth)
+
+        This is the canonical representation of experiment state used for both
+        serialization to status files and network communication.
+
+        Note: Jobs are not included here - they are stored in jobs.jsonl.
+        """
         try:
-            with temp_path.open("w") as f:
-                json.dump(existing, f, indent=2)
-            temp_path.replace(metadata_path)
-            logger.debug("Wrote metadata to %s: %s", metadata_path, list(fields.keys()))
-        except Exception as e:
-            logger.error("Failed to write metadata to %s: %s", metadata_path, e)
-            if temp_path.exists():
-                temp_path.unlink()
-            raise
-
-    def read_metadata(self) -> Optional[dict]:
-        """Read job metadata from .experimaestro/information.json file
+            status_value = self.status.value
+        except NotImplementedError:
+            status_value = None
+
+        return {
+            "version": self.STATUS_VERSION,
+            "experiment_id": self.experiment_id,
+            "run_id": self.run_id,
+            "status": status_value,
+            "events_count": self.events_count,
+            "hostname": self.hostname,
+            "started_at": self.started_at,
+            "ended_at": self.ended_at,
+            "finished_jobs": self.finished_jobs,
+            "failed_jobs": self.failed_jobs,
+            "services": {k: v.full_state_dict() for k, v in self.services.items()},
+        }
 
-        Returns:
-            Dictionary of metadata fields, or None if file doesn't exist
+    def write_status(self) -> None:
+        """Write status.json to disk (calls state_dict internally)
+
+        Uses file locking to ensure atomic writes across processes.
         """
-        metadata_path = self.metadata_path
-        if not metadata_path.exists():
-            return None
+        import fasteners
 
-        try:
-            with metadata_path.open("r") as f:
-                return json.load(f)
-        except Exception as e:
-            logger.warning("Failed to read metadata from %s: %s", metadata_path, e)
-            return None
+        run_dir = self.run_dir
+        if run_dir is None:
+            return
 
+        status_path = run_dir / "status.json"
+        status_path.parent.mkdir(parents=True, exist_ok=True)
+        lock_path = status_path.parent / f".{status_path.name}.lock"
+        lock = fasteners.InterProcessLock(str(lock_path))
 
-# =============================================================================
-# Base Experiment Interface
-# =============================================================================
+        data = self.state_dict()
+        data["last_updated"] = datetime.now().isoformat()
 
+        with lock:
+            temp_path = status_path.with_suffix(".json.tmp")
+            with temp_path.open("w") as f:
+                json.dump(data, f, indent=2)
+            temp_path.replace(status_path)
 
-class BaseExperiment:
-    """Base interface for experiment information
 
-    This class defines the interface for experiment data. Both live experiment
-    instances and database-loaded MockExperiment instances should provide these attributes.
+class BaseService:
+    """Base interface for service information
+
+    This class defines the interface for service data. Both live Service instances
+    and MockService instances should provide these attributes and methods.
 
     Attributes:
-        workdir: Path to experiment directory
-        current_run_id: Current/latest run ID (or None)
+        id: Unique identifier for the service
+        state: Current service state (ServiceState enum or compatible)
     """
 
-    workdir: Path
-    current_run_id: Optional[str]
+    id: str
 
     @property
-    def experiment_id(self) -> str:
-        """Experiment identifier derived from workdir name"""
-        return self.workdir.name
+    def state(self):
+        """Current service state"""
+        raise NotImplementedError
+
+    def description(self) -> str:
+        """Human-readable description of the service"""
+        raise NotImplementedError
+
+    def state_dict(self) -> dict:
+        """Return service state for serialization/recreation"""
+        return {}
+
+    def full_state_dict(self) -> Dict[str, Any]:
+        """Get service state as dictionary for JSON serialization.
+
+        This method properly serializes Path objects and other non-JSON types.
+        """
+        return {
+            "service_id": self.id,
+            "description": self.description(),
+            "class": f"{self.__class__.__module__}.{self.__class__.__name__}",
+            "state_dict": self.state_dict(),
+        }
+
+    def to_service(self) -> "BaseService":
+        """Convert to a live Service instance.
+
+        For live Service instances, returns self.
+        For MockService instances, tries to recreate the service from config.
+
+        Returns:
+            A live Service instance, or self if conversion is not possible
+        """
+        return self  # Default: return self (for live services)