experimaestro 1.6.1__py3-none-any.whl → 1.15.2__py3-none-any.whl

experimaestro/run.py

@@ -8,6 +8,7 @@ import json
 from typing import List
 import fasteners
 from experimaestro.notifications import progress, report_eoj
+from experimaestro.utils.multiprocessing import delayed_shutdown
 from .core.types import ObjectType
 from experimaestro.utils import logger
 from experimaestro.core.objects import ConfigInformation
@@ -96,6 +97,7 @@ class TaskRunner:
         self.failedpath.write_text(str(code))
         self.cleanup()
         logger.info("Exiting")
+        delayed_shutdown(60, exit_code=code)
         sys.exit(1)
 
     def run(self):
@@ -140,10 +142,10 @@ class TaskRunner:
                 run(workdir / "params.json")
 
                 # ... remove the handlers
-                logger.info("Task ended successfully")
                 remove_signal_handlers(remove_cleanup=False)
 
                 # Everything went OK
+                logger.info("Task ended successfully")
                 sys.exit(0)
         except Exception:
             logger.exception("Got exception while running")

experimaestro/scheduler/base.py

@@ -1,12 +1,22 @@
 from collections import ChainMap
 from functools import cached_property
+import itertools
 import logging
 import os
 from pathlib import Path
 from shutil import rmtree
 import threading
 import time
-from typing import Any, List, Optional, Set, TypeVar, Union, TYPE_CHECKING
+from typing import (
+    Any,
+    Iterator,
+    List,
+    Optional,
+    Set,
+    TypeVar,
+    Union,
+    TYPE_CHECKING,
+)
 import enum
 import signal
 import asyncio
@@ -17,9 +27,10 @@ from experimaestro.scheduler.services import Service
 from experimaestro.settings import WorkspaceSettings, get_settings
 
 
-from experimaestro.core.objects import Config, ConfigWalkContext
+from experimaestro.core.objects import Config, ConfigWalkContext, WatchedOutput
 from experimaestro.utils import logger
 from experimaestro.locking import Locks, LockError, Lock
+from experimaestro.utils.asyncio import asyncThreadcheck
 from .workspace import RunMode, Workspace
 from .dependencies import Dependency, DependencyStatus, Resource
 import concurrent.futures
@@ -110,7 +121,7 @@ class JobDependency(Dependency):
 
 
 class Job(Resource):
-    """A job is a resouce that is produced by the execution of some code"""
+    """A job is a resource that is produced by the execution of some code"""
 
     # Set by the scheduler
     _readyEvent: Optional[asyncio.Event]
@@ -148,6 +159,11 @@ class Job(Resource):
         # Dependencies
         self.dependencies: Set[Dependency] = set()  # as target
 
+        # Watched outputs
+        self.watched_outputs = {}
+        for watched in config.__xpm__.watched_outputs:
+            self.watch_output(watched)
+
         # Process
         self._process = None
         self.unsatisfied = 0
@@ -159,6 +175,23 @@ class Job(Resource):
         self._progress: List[LevelInformation] = []
         self.tags = config.tags()
 
+    def watch_output(self, watched: "WatchedOutput"):
+        """Monitor task outputs
+
+        :param watched: A description of the watched output
+        """
+        self.scheduler.xp.watch_output(watched)
+
+    def task_output_update(self, subpath: Path):
+        """Notification of an updated task output"""
+        if watcher := self.watched_outputs.get(subpath, None):
+            watcher.update()
+
+    def done_handler(self):
+        """The task has been completed"""
+        for watcher in self.watched_outputs.values():
+            watcher.update()
+
     def __str__(self):
         return "Job[{}]".format(self.identifier)
 
@@ -166,6 +199,11 @@ class Job(Resource):
         assert self._future, "Cannot wait a not submitted job"
         return self._future.result()
 
+    @cached_property
+    def python_path(self) -> Iterator[str]:
+        """Returns an iterator over python path"""
+        return itertools.chain(self.workspace.python_path)
+
     @cached_property
     def environ(self):
         """Returns the job environment
@@ -219,7 +257,7 @@ class Job(Resource):
         return self.state == JobState.READY
 
     @property
-    def jobpath(self):
+    def jobpath(self) -> Path:
         """Deprecated, use `path`"""
         return self.workspace.jobspath / self.relpath
 
@@ -227,6 +265,14 @@ class Job(Resource):
     def path(self) -> Path:
         return self.workspace.jobspath / self.relpath
 
+    @property
+    def experimaestro_path(self) -> Path:
+        return (self.path / ".experimaestro").resolve()
+
+    @cached_property
+    def task_outputs_path(self) -> Path:
+        return self.experimaestro_path / "task-outputs.jsonl"
+
     @property
     def relpath(self):
         identifier = self.config.__xpm__.identifier
@@ -436,7 +482,7 @@ class Scheduler:
         self.jobs: Dict[str, "Job"] = {}
 
         # List of jobs
-        self.waitingjobs = set()
+        self.waitingjobs: Set[Job] = set()
 
         # Listeners
         self.listeners: Set[Listener] = set()
@@ -459,14 +505,17 @@ class Scheduler:
 
     def submit(self, job: Job) -> Optional[Job]:
         # Wait for the future containing the submitted job
+        logger.debug("Registering the job %s within the scheduler", job)
         otherFuture = asyncio.run_coroutine_threadsafe(
             self.aio_registerJob(job), self.loop
         )
         other = otherFuture.result()
+        logger.debug("Job already submitted" if other else "First submission")
         if other:
             return other
 
         job._future = asyncio.run_coroutine_threadsafe(self.aio_submit(job), self.loop)
+        return None
 
     def prepare(self, job: Job):
         """Prepares the job for running"""
@@ -598,9 +647,13 @@ class Scheduler:
         if job.state != JobState.DONE:
             self.xp.failedJobs[job.identifier] = job
 
+        # Process all remaining tasks outputs
+        await asyncThreadcheck("End of job processing", job.done_handler)
+
         # Decrement the number of unfinished jobs and notify
         self.xp.unfinishedJobs -= 1
         async with self.xp.central.exitCondition:
+            logging.debug("Updated number of unfinished jobs")
             self.xp.central.exitCondition.notify_all()
 
         job.endtime = time.time()
@@ -688,6 +741,7 @@ class Scheduler:
                     code = await process.aio_code()
                     logger.debug("Got return code %s for %s", code, job)
 
+                    # Check the file if there is no return code
                     if code is None:
                         # Case where we cannot retrieve the code right away
                         if job.donepath.is_file():
@@ -766,6 +820,7 @@ class experiment:
         """
 
         from experimaestro.server import Server
+        from experimaestro.scheduler import Listener
 
         settings = get_settings()
         if not isinstance(env, WorkspaceSettings):
@@ -782,6 +837,7 @@ class experiment:
         self.xplock = None
         self.old_experiment = None
         self.services: Dict[str, Service] = {}
+        self._job_listener: Optional[Listener] = None
 
         # Get configuration settings
 
@@ -853,6 +909,7 @@ class experiment:
             assert self.central is not None
             async with self.central.exitCondition:
                 self.exitMode = True
+                logging.debug("Setting exit mode to true")
                 self.central.exitCondition.notify_all()
 
         assert self.central is not None and self.central.loop is not None
@@ -863,10 +920,22 @@ class experiment:
 
         async def awaitcompletion():
             assert self.central is not None
+            logger.debug("Waiting to exit scheduler...")
             async with self.central.exitCondition:
                 while True:
-                    if self.unfinishedJobs == 0 or self.exitMode:
+                    if self.exitMode:
+                        break
+
+                    # If we have still unfinished jobs or possible new tasks, wait
+                    logger.debug(
+                        "Checking exit condition: unfinished jobs=%d, task output queue size=%d",
+                        self.unfinishedJobs,
+                        self.taskOutputQueueSize,
+                    )
+                    if self.unfinishedJobs == 0 and self.taskOutputQueueSize == 0:
                         break
+
+                    # Wait for more news...
                     await self.central.exitCondition.wait()
 
                 if self.failedJobs:
@@ -898,6 +967,8 @@ class experiment:
         return self.workspace.connector.createtoken(name, count)
 
     def __enter__(self):
+        from .dynamic_outputs import TaskOutputsWorker
+
         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__()
@@ -926,6 +997,7 @@ class experiment:
         global SIGNAL_HANDLER
         # Number of unfinished jobs
         self.unfinishedJobs = 0
+        self.taskOutputQueueSize = 0
 
         # List of failed jobs
         self.failedJobs: Dict[str, Job] = {}
@@ -934,6 +1006,8 @@ class experiment:
         self.exitMode = False
 
         self.central = SchedulerCentral.create(self.scheduler.name)
+        self.taskOutputsWorker = TaskOutputsWorker(self)
+        self.taskOutputsWorker.start()
 
         SIGNAL_HANDLER.add(self)
 
@@ -942,6 +1016,7 @@ class experiment:
         return self
 
     def __exit__(self, exc_type, exc_value, traceback):
+        logger.debug("Exiting scheduler context")
         # If no exception and normal run mode, remove old "jobs"
         if self.workspace.run_mode == RunMode.NORMAL:
             if exc_type is None and self.jobsbakpath.is_dir():
@@ -967,8 +1042,13 @@ class experiment:
                 service.stop()
 
             if self.central is not None:
+                logger.info("Stopping scheduler event loop")
                 self.central.loop.stop()
 
+            if self.taskOutputsWorker is not None:
+                logger.info("Stopping tasks outputs worker")
+                self.taskOutputsWorker.queue.put(None)
+
             self.central = None
             self.workspace.__exit__(exc_type, exc_value, traceback)
             if self.xplock:
@@ -977,8 +1057,36 @@ 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
+            logging.info("Saving the experiment state")
+            from experimaestro.scheduler.state import ExperimentState
+
+            ExperimentState.save(
+                self.workdir / "state.json", self.scheduler.jobs.values()
+            )
+
+    async def update_task_output_count(self, delta: int):
+        """Change in the number of task outputs to process"""
+        async with self.central.exitCondition:
+            self.taskOutputQueueSize += delta
+            logging.debug(
+                "Updating queue size with %d => %d", delta, self.taskOutputQueueSize
+            )
+            if self.taskOutputQueueSize == 0:
+                self.central.exitCondition.notify_all()
+
+    def watch_output(self, watched: "WatchedOutput"):
+        """Watch an output
+
+        :param watched: The watched output specification
+        """
+
+        self.taskOutputsWorker.watch_output(watched)
+
     def add_service(self, service: ServiceClass) -> ServiceClass:
         """Adds a service (e.g. tensorboard viewer) to the experiment
 

experimaestro/scheduler/dynamic_outputs.py

@@ -0,0 +1,184 @@
+"""Handles dynamic task outputs"""
+
+import asyncio
+import json
+import logging
+import queue
+import threading
+from collections import defaultdict
+from functools import cached_property
+from pathlib import Path
+from typing import Callable, TYPE_CHECKING
+
+from watchdog.events import FileSystemEventHandler
+
+from experimaestro.ipc import ipcom
+from experimaestro.utils import logger
+
+from .base import Job, experiment
+
+if TYPE_CHECKING:
+    from experimaestro.core.objects import WatchedOutput
+
+
+class TaskOutputCallbackHandler:
+    def __init__(self, converter: Callable):
+        pass
+
+
+class TaskOutputs(FileSystemEventHandler):
+    """Represent and monitors dynamic outputs generated by one task"""
+
+    #: Global dictionary for handles
+    HANDLERS: dict[Path, "TaskOutputs"] = {}
+
+    #: Global lock to access current HANDLERS
+    LOCK = threading.Lock()
+
+    def create(job: Job):
+        with TaskOutputs.LOCK:
+            if instance := TaskOutputs.get(job.task_outputs_path, None):
+                return instance
+
+            instance = TaskOutputs(job.task_outputs_path)
+            TaskOutputs[job.task_outputs_path] = instance
+            return instance
+
+    def __init__(self, path: Path):
+        """Monitors an event path"""
+        logger.debug("Watching dynamic task outputs in %s", path)
+        self.path = path
+        self.handle = None
+        self.count = 0
+        self.lock = threading.Lock()
+        self.listeners: dict[str, dict[Callable, set[Callable]]] = defaultdict(
+            lambda: defaultdict(set)
+        )
+
+        #: The events registered so far
+        self.events = []
+
+    def __enter__(self):
+        """Starts monitoring task outputs"""
+        self.job.task_outputs_path.parent.mkdir(parents=True, exist_ok=True)
+        with self.lock:
+            if self.handle is None:
+                assert self.count == 0
+                self.handle = ipcom().fswatch(self, self.path.parent, False)
+            self.count += 1
+        return self
+
+    def __exit__(self, *args):
+        """Stops monitoring task outputs"""
+        with self.lock:
+            self.count -= 1
+            if self.count == 0:
+                ipcom().fsunwatch(self.handle)
+                self.fh.close()
+
+                self.handle = None
+                self._fh = None
+
+    def watch_output(self, watched: "WatchedOutput"):
+        """Add a new listener"""
+        key = f"{watched.config.__identifier__}/{watched.method_name}"
+        with self.lock:
+            # Process events so far
+            listener = self.listeners[key].get(watched.method, None)
+            if listener is None:
+                listener = TaskOutputCallbackHandler(watched.method)
+
+            # Register
+            self.listeners[key][watched.method].add(watched.callback)
+
+    #
+    # --- Events
+    #
+
+    @cached_property
+    def fh(self):
+        if self._fh is None:
+            self._fh = self.path.open("rt")
+        return self._fh
+
+    def on_modified(self, event):
+        self.handle(Path(event.src_path))
+
+    def on_created(self, event):
+        self.handle(Path(event.src_path))
+
+    def handle(self, path: Path):
+        if path != self.path:
+            return
+
+        with self.lock:
+            logger.debug("[TASK OUTPUT] Handling task output for %s", self.path)
+
+            while json_line := self.fh.readline():
+                # Read the event
+                event = json.loads(json_line)
+                logger.debug("Event: %s", event)
+
+                # FIXME: move elsewhere
+                # # Process the event
+                # event = self.config_method(
+                #     self.job.config.__xpm__.mark_output,
+                #     *event["args"],
+                #     **event["kwargs"],
+                # )
+
+                self.events.append(event)
+                # self.job.scheduler.xp.taskOutputsWorker.add(self, event)
+
+
+class TaskOutputsWorker(threading.Thread):
+    """This worker process dynamic output queue for one experiment"""
+
+    def __init__(self, xp: experiment):
+        super().__init__(name="task outputs worker", daemon=True)
+        self.queue = queue.Queue()
+        self.xp = xp
+
+    def watch_output(self, watched: "WatchedOutput"):
+        """Watch an output
+
+        :param watched: The watched output specification
+        """
+        logger.debug("Registering task output listener %s", watched)
+
+        # path = watched.job.tasks_output_path
+        TaskOutputs.create(watched.job).watch_output(watched)
+
+    def add(self, watcher, event):
+        asyncio.run_coroutine_threadsafe(
+            self.xp.update_task_output_count(1),
+            self.xp.scheduler.loop,
+        ).result()
+        self.queue.put((watcher, event))
+
+    def run(self):
+        logging.debug("Starting output listener queue")
+        while True:
+            # Get the next element in the queue
+            element = self.queue.get()
+            if element is None:
+                # end of processing
+                break
+
+            # Call all the listeners
+            logging.debug("Got one event: %s", element)
+            watcher, event = element
+            for listener in watcher.listeners:
+                try:
+                    logger.debug("Calling listener [%s] with %s", listener, event)
+                    listener(event)
+                    logger.debug(
+                        "[done] Calling listener [%s] with %s", listener, event
+                    )
+                except Exception:
+                    logging.exception("Exception while calling the listener")
+                self.queue.task_done()
+
+                asyncio.run_coroutine_threadsafe(
+                    self.xp.update_task_output_count(-1), self.xp.scheduler.loop
+                ).result()

experimaestro/scheduler/state.py

@@ -0,0 +1,75 @@
+from dataclasses import dataclass, field
+import json
+from pathlib import Path
+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.settings import find_workspace
+from experimaestro.core.serialization import from_state_dict, save_definition
+
+
+@dataclass
+class JobInformation:
+    id: str
+    path: Path
+    task: Task
+    tags: dict[str, str]
+    depends_on: list["JobInformation"] = field(default_factory=list)
+
+    def __post_init__(self):
+        self.path = Path(self.path)
+
+
+class ExperimentState:
+    def __init__(self, workdir: Path, name: str):
+        path = workdir / "xp" / name / "state.json"
+        with path.open("rt") as fh:
+            content = json.load(fh)
+
+        self.states: dict[str, JobInformation] = {
+            state_dict["id"]: JobInformation(**state_dict)
+            for state_dict in from_state_dict(content, as_instance=False)
+        }
+
+        for state in self.states.values():
+            state.depends_on = [self.states[key] for key in state.depends_on]
+
+    def get_jobs(self, task_class: Type[Task]) -> list[JobInformation]:
+        if task_class is None:
+            return list(self.data.values())
+
+        tasks = []
+        for job_state in self.states.values():
+            if isinstance(job_state.task, task_class):
+                tasks.append(job_state)
+        return tasks
+
+    @staticmethod
+    def save(path: Path, jobs: Iterable[Job]):
+        save_definition(
+            [
+                {
+                    "id": job.identifier,
+                    "path": str(job.path),
+                    "task": job.config,
+                    "tags": job.config.__xpm__.tags(),
+                    "depends_on": list(
+                        dep.origin.identifier
+                        for dep in job.dependencies
+                        if isinstance(dep, JobDependency)
+                    ),
+                }
+                for job in jobs
+            ],
+            SerializationContext(),
+            path,
+        )
+
+
+def get_experiment(
+    name: str, *, workspace: Optional[str] = None, workdir: Optional[Path] = None
+) -> ExperimentState:
+    ws = find_workspace(workspace=workspace, workdir=workdir)
+    return ExperimentState(ws.path, name)

experimaestro/scheduler/workspace.py

@@ -2,7 +2,7 @@ from collections import ChainMap
 from enum import Enum
 from functools import cached_property
 from pathlib import Path
-from typing import Optional
+from typing import Iterator, Optional
 from experimaestro.settings import WorkspaceSettings, Settings
 
 
@@ -46,6 +46,7 @@ class Workspace:
             path = path.absolute()
         self.path = path
         self.run_mode = run_mode
+        self.python_path = []
         from ..launchers import Launcher
 
         self.launcher = launcher or Launcher.get(path)

experimaestro/scriptbuilder.py

@@ -51,6 +51,8 @@ class PythonScriptBuilder:
         self.lockfiles: List[Path] = []
         self.notificationURL: Optional[str] = None
         self.command: Optional[AbstractCommand] = None
+
+        # This is used to serialize the full process identifier on disk
         self.processtype = "local"
 
     def write(self, job: CommandLineJob):
@@ -63,7 +65,7 @@ class PythonScriptBuilder:
             job {CommandLineJob} -- [description]
 
         Returns:
-            [type] -- [description]
+            str -- The script path on disk
         """
         assert isinstance(
             job, CommandLineJob
@@ -94,6 +96,7 @@ class PythonScriptBuilder:
             out.write("# Experimaestro generated task\n\n")
             out.write(
                 """import logging\n"""
+                """import sys\n"""
                 """logging.basicConfig(level=logging.INFO, """
                 """format='%(levelname)s:%(process)d:%(asctime)s [%(name)s] %(message)s', datefmt='%y-%m-%d %H:%M:%S')\n\n"""
             )
@@ -112,9 +115,17 @@ class PythonScriptBuilder:
             out.write("    ]\n")
 
             for name, value in job.environ.items():
-                out.write(f"""    os.environ["{name}"] = "{shquote(value)}"\n""")
+                if name == "PYTHONPATH":
+                    # Handles properly python path
+                    for path in value.split(":"):
+                        out.write(f"""    sys.path.insert(0, "{shquote(path)}")\n""")
+                else:
+                    out.write(f"""    os.environ["{name}"] = "{shquote(value)}"\n""")
             out.write("\n")
 
+            for path in job.python_path:
+                out.write(f"""    sys.path.insert(0, "{shquote(str(path))}")\n""")
+
             out.write(
                 f"""    TaskRunner("{shquote(connector.resolve(scriptpath))}","""
                 """ lockfiles).run()\n"""