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

experimaestro/scheduler/jobs.py

@@ -1,7 +1,6 @@
 import asyncio
-import time
 from collections import ChainMap
-import enum
+from datetime import datetime
 from functools import cached_property
 import itertools
 from pathlib import Path
@@ -13,60 +12,42 @@ from experimaestro.core.objects import Config, ConfigWalkContext, WatchedOutput
 from experimaestro.notifications import LevelInformation, Reporter
 
 # from experimaestro.scheduler.base import Scheduler
-from experimaestro.scheduler.dependencies import Dependency, DependencyStatus, Resource
+from experimaestro.scheduler.dependencies import Dependency, Resource
 from experimaestro.scheduler.workspace import RunMode, Workspace
-from experimaestro.locking import Lock, LockError, Locks
+from experimaestro.scheduler.interfaces import (
+    BaseJob,
+    JobState,
+    JobStateUnscheduled,
+    JobStateWaiting,
+    JobStateReady,
+    JobStateScheduled,
+    JobStateRunning,
+    JobStateDone,
+    JobStateError,
+    JobFailureStatus,
+)
+from experimaestro.locking import Lock
 from experimaestro.utils import logger
 
 if TYPE_CHECKING:
     from experimaestro.connectors import Process
     from experimaestro.launchers import Launcher
+    from experimaestro.scheduler.experiment import experiment
 
 
-class JobState(enum.Enum):
-    # Job is not yet scheduled
-    UNSCHEDULED = 0
-
-    # Job is waiting for dependencies to be done
-    WAITING = 1
-
-    # Job is ready to run
-    READY = 2
-
-    # Job is scheduled (e.g. slurm)
-    SCHEDULED = 3
-
-    # Job is running
-    RUNNING = 4
-
-    # Job is done (finished)
-    DONE = 5
-
-    # Job failed (finished)
-    ERROR = 6
-
-    def notstarted(self):
-        return self.value <= JobState.READY.value
-
-    def running(self):
-        return (
-            self.value == JobState.RUNNING.value
-            or self.value == JobState.SCHEDULED.value
-        )
-
-    def finished(self):
-        return self.value >= JobState.DONE.value
-
-
-class JobFailureStatus(enum.Enum):
-    #: Job failed
-    DEPENDENCY = 0
-
-    #: Job dependency failed
-    FAILED = 1
-
-    #: Memory
-    MEMORY = 2
+# Re-export JobState for backward compatibility
+__all__ = [
+    "JobState",
+    "JobStateUnscheduled",
+    "JobStateWaiting",
+    "JobStateReady",
+    "JobStateScheduled",
+    "JobStateRunning",
+    "JobStateDone",
+    "JobStateError",
+    "JobFailureStatus",
+    "Job",
+]
 
 
 class JobLock(Lock):
@@ -85,22 +66,43 @@ class JobDependency(Dependency):
     def __init__(self, job):
         super().__init__(job)
 
-    def status(self) -> DependencyStatus:
-        if self.origin.state == JobState.DONE:
-            return DependencyStatus.OK
-        elif self.origin.state == JobState.ERROR:
-            return DependencyStatus.FAIL
-        return DependencyStatus.WAIT
+    async def aio_lock(self, timeout: float = 0):
+        """Acquire lock on job dependency by waiting for job to complete
 
-    def lock(self):
-        return JobLock(self.origin)
+        Args:
+            timeout: Must be 0 (wait indefinitely) for job dependencies
 
+        Raises:
+            ValueError: If timeout is not 0
+            RuntimeError: If the job has not been submitted or if it failed
+        """
+        if timeout != 0:
+            raise ValueError(
+                "Job dependencies only support timeout=0 (wait indefinitely)"
+            )
 
-class Job(Resource):
+        # Wait for the job to finish
+        if self.origin._future is None:
+            raise RuntimeError(f"Job {self.origin} has no future - not submitted")
+        await asyncio.wrap_future(self.origin._future)
+
+        # Check if the job succeeded
+        if self.origin.state != JobState.DONE:
+            raise RuntimeError(
+                f"Dependency job {self.origin.identifier} failed with state "
+                f"{self.origin.state} for {self.target.identifier}"
+            )
+
+        # Job succeeded, acquire and return the lock
+        lock = JobLock(self.origin)
+        lock.acquire()
+        return lock
+
+
+class Job(BaseJob, Resource):
     """A job is a resource that is produced by the execution of some code"""
 
     # Set by the scheduler
-    _readyEvent: Optional[asyncio.Event]
     _future: Optional["concurrent.futures.Future"]
 
     def __init__(
@@ -110,6 +112,7 @@ class Job(Resource):
         workspace: Workspace = None,
         launcher: "Launcher" = None,
         run_mode: RunMode = RunMode.NORMAL,
+        max_retries: Optional[int] = None,
     ):
         from experimaestro.scheduler.base import Scheduler
 
@@ -128,47 +131,74 @@ class Job(Resource):
         self.name = str(self.type.identifier).rsplit(".", 1)[-1]
 
         self.scheduler: Optional["Scheduler"] = None
+        self.experiments: List["experiment"] = []  # Experiments this job belongs to
         self.config = config
         self.state: JobState = JobState.UNSCHEDULED
 
-        #: If a job has failed, indicates the failure status
-        self.failure_status: JobFailureStatus = None
-
         # Dependencies
         self.dependencies: Set[Dependency] = set()  # as target
 
-        # Watched outputs
-        self.watched_outputs = {}
-        for watched in config.__xpm__.watched_outputs:
-            self.watch_output(watched)
+        # Check if this is a resumable task
+        from experimaestro.core.objects import ResumableTask
+
+        self.resumable = isinstance(config, ResumableTask)
+
+        # Retry configuration for resumable tasks
+        # Use workspace setting if max_retries is not specified
+        if max_retries is None and self.workspace:
+            max_retries = self.workspace.workspace_settings.max_retries
+        self.max_retries = max_retries if max_retries is not None else 3
+        self.retry_count = 0
+
+        # Watched outputs (stored for deferred registration with scheduler)
+        self.watched_outputs: List["WatchedOutput"] = list(
+            config.__xpm__.watched_outputs
+        )
 
         # Process
         self._process = None
-        self.unsatisfied = 0
 
         # Meta-information
         self.starttime: Optional[float] = None
         self.submittime: Optional[float] = None
         self.endtime: Optional[float] = None
+        self.exit_code: Optional[int] = None
         self._progress: List[LevelInformation] = []
         self.tags = config.tags()
 
     def watch_output(self, watched: "WatchedOutput"):
-        """Monitor task outputs
+        """Add a watched output to this job.
 
         :param watched: A description of the watched output
         """
-        self.scheduler.xp.watch_output(watched)
+        self.watched_outputs.append(watched)
+
+    def register_watched_outputs(self):
+        """Register all watched outputs with the scheduler.
+
+        This should be called after the job is submitted and has a scheduler.
+        """
+        from experimaestro.scheduler.experiment import experiment
 
-    def task_output_update(self, subpath: Path):
-        """Notification of an updated task output"""
-        if watcher := self.watched_outputs.get(subpath, None):
-            watcher.update()
+        xp = experiment.current()
+        for watched in self.watched_outputs:
+            # Set the job reference so the watcher knows where to look
+            watched.job = self
+            xp.watch_output(watched)
 
     def done_handler(self):
-        """The task has been completed"""
-        for watcher in self.watched_outputs.values():
-            watcher.update()
+        """The task has been completed.
+
+        Ensures all remaining task output events are processed by explicitly
+        reading the task outputs file. This is necessary because file system
+        watchers may have latency, and we need to process all outputs before
+        the experiment can exit.
+        """
+        if not self.watched_outputs:
+            return
+
+        for xp in self.experiments:
+            xp.taskOutputsWorker.process_job_outputs(self)
 
     def __str__(self):
         return "Job[{}]".format(self.identifier)
@@ -177,6 +207,57 @@ class Job(Resource):
         assert self._future, "Cannot wait a not submitted job"
         return self._future.result()
 
+    def set_state(self, new_state: JobState):
+        """Set the job state and update experiment statistics
+
+        This method should be called instead of direct state assignment
+        to ensure experiment statistics (unfinishedJobs, failedJobs) are
+        properly updated.
+
+        :param new_state: The new job state
+        """
+        old_state = self.state
+        self.state = new_state
+
+        # Helper to determine if a state should be "counted" in unfinishedJobs
+        # A job is counted when it's been submitted and hasn't finished yet
+        def is_counted(state):
+            return state != JobState.UNSCHEDULED and not state.finished()
+
+        # Update experiment statistics based on state transition
+        for xp in self.experiments:
+            # Handle transitions in/out of "counted" state
+            if is_counted(new_state) and not is_counted(old_state):
+                # Job is now being tracked (new submission or resubmit)
+                xp.unfinishedJobs += 1
+                logger.debug(
+                    "Job %s submitted, unfinished jobs for %s: %d",
+                    self.identifier[:8],
+                    xp.workdir.name,
+                    xp.unfinishedJobs,
+                )
+            elif not is_counted(new_state) and is_counted(old_state):
+                # Job is no longer being tracked (finished)
+                xp.unfinishedJobs -= 1
+                logger.debug(
+                    "Job %s finished, unfinished jobs for %s: %d",
+                    self.identifier[:8],
+                    xp.workdir.name,
+                    xp.unfinishedJobs,
+                )
+
+            # Handle error state
+            if new_state.is_error() and not old_state.is_error():
+                xp.failedJobs[self.identifier] = self
+
+            # Handle recovery from error (e.g., resubmit)
+            if old_state.is_error() and not new_state.is_error():
+                xp.failedJobs.pop(self.identifier, None)
+
+        # Notify listeners via scheduler's thread-safe mechanism
+        if self.scheduler:
+            self.scheduler.notify_job_state(self)
+
     @cached_property
     def python_path(self) -> Iterator[str]:
         """Returns an iterator over python path"""
@@ -220,8 +301,8 @@ class Job(Resource):
             self._progress[-1].desc = desc
         self._progress[-1].progress = value
 
-        for listener in self.scheduler.listeners:
-            listener.job_state(self)
+        # Notify listeners via scheduler's thread-safe mechanism
+        self.scheduler.notify_job_state(self)
 
     def add_notification_server(self, server):
         """Adds a notification server"""
@@ -271,6 +352,16 @@ class Job(Resource):
     def identifier(self):
         return self.config.__xpm__.identifier.all.hex()
 
+    @property
+    def task_id(self) -> str:
+        """Task class identifier (for BaseJob interface)"""
+        return str(self.type.identifier)
+
+    @property
+    def locator(self) -> str:
+        """Full task locator (for BaseJob interface)"""
+        return self.identifier
+
     def prepare(self, overwrite=False):
         """Prepare all files before starting a task
 
@@ -278,113 +369,12 @@ class Job(Resource):
         """
         pass
 
-    async def aio_start(self, sched_dependency_lock, notification_server=None):
-        """Start the job with core job starting logic
-
-        This method contains the core logic for starting a job that was previously
-        located in Scheduler.aio_start(). It handles job locking, dependency
-        acquisition, directory setup, and job execution while using the scheduler's
-        coordination lock to prevent race conditions between multiple jobs.
-
-        :param sched_dependency_lock: The scheduler's dependency lock for coordination
-            between jobs to prevent race conditions during dependency acquisition
-        :param notification_server: Optional notification server from the experiment
-            for job progress reporting
-        :return: JobState.DONE if job completed successfully, JobState.ERROR if job
-            failed during execution, or None if dependencies couldn't be locked
-            (signals WAITING state to scheduler)
-        :raises Exception: Various exceptions during job execution, dependency locking,
-            or process creation
-        """
-        # We first lock the job before proceeding
-        assert self.launcher is not None
-
-        with Locks() as locks:
-            logger.debug("[starting] Locking job %s", self)
-            async with self.launcher.connector.lock(self.lockpath):
-                logger.debug("[starting] Locked job %s", self)
-
-                state = None
-                try:
-                    logger.debug(
-                        "Starting job %s with %d dependencies",
-                        self,
-                        len(self.dependencies),
-                    )
-
-                    # Individual dependency lock acquisition
-                    # We use the scheduler-wide lock to avoid cross-jobs race conditions
-                    async with sched_dependency_lock:
-                        for dependency in self.dependencies:
-                            try:
-                                locks.append(dependency.lock().acquire())
-                            except LockError:
-                                logger.warning(
-                                    "Could not lock %s, aborting start for job %s",
-                                    dependency,
-                                    self,
-                                )
-                                dependency.check()
-                                return None  # Signal to scheduler that dependencies couldn't be locked
-
-                    # Dependencies have been locked, we can start the job
-                    self.starttime = time.time()
-
-                    # Creates the main directory
-                    directory = self.path
-                    logger.debug("Making directories job %s...", directory)
-                    if not directory.is_dir():
-                        directory.mkdir(parents=True, exist_ok=True)
-
-                    # Sets up the notification URL
-                    if notification_server is not None:
-                        self.add_notification_server(notification_server)
-
-                except Exception:
-                    logger.warning("Error while locking job", exc_info=True)
-                    return None  # Signal waiting state to scheduler
-
-                try:
-                    # Runs the job
-                    process = await self.aio_run()
-                except Exception:
-                    logger.warning("Error while starting job", exc_info=True)
-                    return JobState.ERROR
-
-            try:
-                if isinstance(process, JobState):
-                    state = process
-                    logger.debug("Job %s ended (state %s)", self, state)
-                else:
-                    logger.debug("Waiting for job %s process to end", self)
-
-                    code = await process.aio_code()
-                    logger.debug("Got return code %s for %s", code, self)
-
-                    # Check the file if there is no return code
-                    if code is None:
-                        # Case where we cannot retrieve the code right away
-                        if self.donepath.is_file():
-                            code = 0
-                        else:
-                            code = int(self.failedpath.read_text())
-
-                    logger.debug("Job %s ended with code %s", self, code)
-                    state = JobState.DONE if code == 0 else JobState.ERROR
-
-            except JobError:
-                logger.warning("Error while running job")
-                state = JobState.ERROR
-
-            except Exception:
-                logger.warning(
-                    "Error while running job (in experimaestro)", exc_info=True
-                )
-                state = JobState.ERROR
-        return state
+    async def aio_run(self) -> "Process":
+        """Actually run the code
 
-    async def aio_run(self):
-        """Actually run the code"""
+        Returns:
+            A Process instance representing the running job
+        """
         raise NotImplementedError(f"Method aio_run not implemented in {self.__class__}")
 
     async def aio_process(self) -> Optional["Process"]:
@@ -420,33 +410,28 @@ class Job(Resource):
     def stderr(self) -> Path:
         return self.jobpath / ("%s.err" % self.name)
 
+    def rotate_logs(self) -> None:
+        """Rotate log files before restarting a task.
+
+        Renames non-empty stdout and stderr files with a timestamp suffix
+        (e.g., job.20231215143022.out) to preserve logs from previous runs.
+        """
+        timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
+
+        for log_path in [self.stdout, self.stderr]:
+            if log_path.exists() and log_path.stat().st_size > 0:
+                # Extract extension (.out or .err)
+                ext = log_path.suffix
+                # Create new name with timestamp before extension
+                new_name = f"{log_path.stem}.{timestamp}{ext}"
+                new_path = log_path.parent / new_name
+                logger.info("Rotating log file %s -> %s", log_path.name, new_name)
+                log_path.rename(new_path)
+
     @property
     def basepath(self) -> Path:
         return self.jobpath / self.name
 
-    def dependencychanged(self, dependency, oldstatus, status):
-        """Called when a dependency has changed"""
-
-        def value(s):
-            return 1 if s == DependencyStatus.OK else 0
-
-        self.unsatisfied -= value(status) - value(oldstatus)
-
-        logger.debug("Job %s: unsatisfied %d", self, self.unsatisfied)
-
-        if status == DependencyStatus.FAIL:
-            # Job completed
-            if not self.state.finished():
-                self.state = JobState.ERROR
-                self.failure_status = JobFailureStatus.DEPENDENCY
-                self._readyEvent.set()
-
-        if self.unsatisfied == 0:
-            logger.info("Job %s is ready to run", self)
-            # We are ready
-            self.state = JobState.READY
-            self._readyEvent.set()
-
     def finalState(self) -> "concurrent.futures.Future[JobState]":
         assert self._future is not None
         return self._future
@@ -469,6 +454,31 @@ class JobContext(ConfigWalkContext):
     def task(self):
         return self.job.config
 
+    def partial_path(self, subparameters, config) -> Path:
+        """Returns the partial directory path for a given subparameters instance.
+
+        The partial path structure is:
+        WORKSPACE/partials/TASK_ID/SUBPARAM_NAME/PARTIAL_ID/
+
+        Args:
+            subparameters: The Subparameters instance defining which groups to exclude
+            config: The configuration to compute the partial identifier for
+
+        Returns:
+            The partial directory path.
+        """
+        # Compute partial identifier
+        partial_id = config.__xpm__.get_partial_identifier(subparameters)
+
+        # Build partial directory path
+        task_id = str(config.__xpmtype__.identifier)
+        return (
+            self.job.workspace.partialspath
+            / task_id
+            / subparameters.name
+            / partial_id.all.hex()
+        )
+
 
 class JobError(Exception):
     def __init__(self, code):

experimaestro/scheduler/remote/__init__.py

@@ -0,0 +1,31 @@
+"""Remote monitoring support for experimaestro
+
+This package provides SSH-based remote monitoring capabilities for experiments.
+
+Main components:
+- SSHStateProviderServer: JSON-RPC server that wraps WorkspaceStateProvider
+- SSHStateProviderClient: Client that connects via SSH and implements StateProvider interface
+- RemoteFileSynchronizer: Rsync-based file synchronization
+
+Usage:
+    # On remote host (run via SSH):
+    from experimaestro.scheduler.remote.server import SSHStateProviderServer
+    server = SSHStateProviderServer(workspace_path)
+    server.start()
+
+    # On local host:
+    from experimaestro.scheduler.remote.client import SSHStateProviderClient
+    client = SSHStateProviderClient(host="server", remote_workspace="/path")
+    client.connect()
+    experiments = client.get_experiments()
+"""
+
+from experimaestro.scheduler.remote.server import SSHStateProviderServer
+from experimaestro.scheduler.remote.client import SSHStateProviderClient
+from experimaestro.scheduler.remote.sync import RemoteFileSynchronizer
+
+__all__ = [
+    "SSHStateProviderServer",
+    "SSHStateProviderClient",
+    "RemoteFileSynchronizer",
+]