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

experimaestro/experiments/cli.py

@@ -59,7 +59,7 @@ class ExperimentCallable(Protocol):
 class ConfigurationLoader:
     def __init__(self):
         self.yamls = []
-        self.pythonpath = set()
+        self.python_path = set()
 
     def load(self, yaml_file: Path):
         """Loads a YAML file, and parents one if they exist"""
@@ -76,9 +76,9 @@ class ConfigurationLoader:
         for path in _data.get("pythonpath", []):
             path = Path(path)
             if path.is_absolute():
-                self.pythonpath.add(path.resolve())
+                self.python_path.add(path.resolve())
             else:
-                self.pythonpath.add((yaml_file.parent / path).resolve())
+                self.python_path.add((yaml_file.parent / path).resolve())
 
 
 @click.option("--debug", is_flag=True, help="Print debug information")
@@ -181,7 +181,7 @@ def experiments_cli(  # noqa: C901
         configuration.merge_with(OmegaConf.from_dotlist(extra_conf))
 
     # --- Get the XP file
-    pythonpath = list(conf_loader.pythonpath)
+    python_path = list(conf_loader.python_path)
     if module_name is None:
         module_name = configuration.get("module", None)
 
@@ -192,9 +192,13 @@ def experiments_cli(  # noqa: C901
                 not module_name
             ), "Module name and experiment file are mutually exclusive options"
             xp_file = Path(xp_file)
-            if not pythonpath:
-                pythonpath.append(xp_file.parent)
-            logging.info("Using python path: %s", ", ".join(str(s) for s in pythonpath))
+            if not python_path:
+                python_path.append(xp_file.parent)
+            logging.info(
+                "Using python path: %s", ", ".join(str(s) for s in python_path)
+            )
+    else:
+        xp_file = Path(xp_file)
 
     assert (
         module_name or xp_file
@@ -209,7 +213,7 @@ def experiments_cli(  # noqa: C901
     # --- Finds the "run" function
 
     # Modifies the Python path
-    for path in pythonpath:
+    for path in python_path:
         sys.path.append(str(path))
 
     if xp_file:
@@ -226,7 +230,11 @@ def experiments_cli(  # noqa: C901
             )
     else:
         # Module
-        mod = importlib.import_module(module_name)
+        try:
+            mod = importlib.import_module(module_name)
+        except ModuleNotFoundError as e:
+            logging.error("Module not found: %s with python path %s", e, sys.path)
+            raise
 
     helper = getattr(mod, "run", None)
 
@@ -265,6 +273,7 @@ def experiments_cli(  # noqa: C901
 
     # Define the workspace
     ws_env = find_workspace(workdir=workdir, workspace=workspace)
+
     workdir = ws_env.path
 
     logging.info("Using working directory %s", str(workdir.resolve()))
@@ -278,6 +287,9 @@ def experiments_cli(  # noqa: C901
         for key, value in env:
             xp.setenv(key, value)
 
+        # Sets the python path
+        xp.workspace.python_path.extend(python_path)
+
         try:
             # Run the experiment
             helper.xp = xp

experimaestro/generators.py

@@ -1,11 +1,12 @@
 import inspect
 from pathlib import Path
+from abc import ABC, abstractmethod
 from typing import Callable, Union
 from experimaestro.core.arguments import ArgumentOptions, TypeAnnotation
 from experimaestro.core.objects import ConfigWalkContext, Config
 
 
-class Generator:
+class Generator(ABC):
     """Base class for all generators"""
 
     def isoutput(self):
@@ -13,6 +14,10 @@ class Generator:
         path within the job folder)"""
         return False
 
+    @abstractmethod
+    def __call__(self, context: ConfigWalkContext, config: Config):
+        ...
+
 
 class PathGenerator(Generator):
     """Generates a path"""

experimaestro/ipc.py

@@ -7,6 +7,7 @@ import sys
 import logging
 from .utils import logger
 from watchdog.observers import Observer
+from watchdog.observers.api import ObservedWatch
 from watchdog.events import FileSystemEventHandler
 
 
@@ -20,7 +21,9 @@ class IPCom:
         self.observer.start()
         self.pid = os.getpid()
 
-    def fswatch(self, watcher: FileSystemEventHandler, path: Path, recursive=False):
+    def fswatch(
+        self, watcher: FileSystemEventHandler, path: Path, recursive=False
+    ) -> ObservedWatch:
         if not self.observer.is_alive():
             logging.error("Observer is not alive")
 

experimaestro/launcherfinder/registry.py

@@ -1,5 +1,6 @@
 # Configuration registers
 
+from contextlib import contextmanager
 from typing import ClassVar, Dict, Optional, Set, Type, Union
 
 from pathlib import Path
@@ -7,7 +8,6 @@ import typing
 from omegaconf import DictConfig, OmegaConf, SCMode
 import pkg_resources
 from experimaestro.utils import logger
-
 from .base import ConnectorConfiguration, TokenConfiguration
 from .specs import HostRequirement
 
@@ -36,6 +36,16 @@ def load_yaml(schema, path: Path):
         )
 
 
+@contextmanager
+def ensure_enter(fp):
+    """Behaves as a resource, whether it is one or not"""
+    if hasattr(fp, "__enter__"):
+        with fp as _fp:
+            yield _fp
+    else:
+        yield fp
+
+
 class LauncherRegistry:
     INSTANCES: ClassVar[Dict[Path, "LauncherRegistry"]] = {}
     CURRENT_CONFIG_DIR: ClassVar[Optional[Path]] = None
@@ -78,13 +88,16 @@ class LauncherRegistry:
 
             from importlib import util
 
-            spec = util.spec_from_file_location("xpm_launchers_conf", launchers_py)
-            module = util.module_from_spec(spec)
-            spec.loader.exec_module(module)
+            with ensure_enter(launchers_py.__fspath__()) as fp:
+                spec = util.spec_from_file_location("xpm_launchers_conf", fp)
+                module = util.module_from_spec(spec)
+                spec.loader.exec_module(module)
 
             self.find_launcher_fn = getattr(module, "find_launcher", None)
             if self.find_launcher_fn is None:
-                logger.warn("No find_launcher() function was found in %s", launchers_py)
+                logger.warning(
+                    "No find_launcher() function was found in %s", launchers_py
+                )
 
         # Read the configuration file
         self.connectors = load_yaml(
@@ -144,9 +157,14 @@ class LauncherRegistry:
                 specs.append(spec)
 
         # Use launcher function
+        from experimaestro.launchers import Launcher
+
         if self.find_launcher_fn is not None:
             for spec in specs:
                 if launcher := self.find_launcher_fn(spec, tags):
+                    assert isinstance(
+                        launcher, Launcher
+                    ), "f{self.find_launcher_fn} did not return a Launcher but {type(launcher)}"
                     return launcher
 
         return None

experimaestro/launchers/slurm/base.py

@@ -11,6 +11,7 @@ from typing import (
 )
 from experimaestro.connectors.local import LocalConnector
 import re
+from shlex import quote as shquote
 from contextlib import contextmanager
 from dataclasses import dataclass
 from experimaestro.launcherfinder.registry import (
@@ -235,15 +236,15 @@ class SlurmProcessBuilder(ProcessBuilder):
         super().__init__()
         self.launcher = launcher
 
-    def start(self) -> BatchSlurmProcess:
+    def start(self, task_mode: bool = False) -> BatchSlurmProcess:
         """Start the process"""
         builder = self.launcher.connector.processbuilder()
-        builder.workingDirectory = self.workingDirectory
         builder.environ = self.launcher.launcherenv
         builder.detach = False
 
         if not self.detach:
             # Simplest case: we wait for the output
+            builder.workingDirectory = self.workingDirectory
             builder.command = [f"{self.launcher.binpath}/srun"]
             builder.command.extend(self.launcher.options.args())
             builder.command.extend(self.command)
@@ -255,11 +256,17 @@ class SlurmProcessBuilder(ProcessBuilder):
             return builder.start()
 
         builder.command = [f"{self.launcher.binpath}/sbatch", "--parsable"]
-        builder.command.extend(self.launcher.options.args())
 
-        addstream(builder.command, "-e", self.stderr)
-        addstream(builder.command, "-o", self.stdout)
-        addstream(builder.command, "-i", self.stdin)
+        if not task_mode:
+            # Use command line parameters when not running a task
+            builder.command.extend(self.launcher.options.args())
+
+            if self.workingDirectory:
+                workdir = self.launcher.connector.resolve(self.workingDirectory)
+                builder.command.append(f"--chdir={workdir}")
+            addstream(builder.command, "-e", self.stderr)
+            addstream(builder.command, "-o", self.stdout)
+            addstream(builder.command, "-i", self.stdin)
 
         builder.command.extend(self.command)
         logger.info(
@@ -427,12 +434,43 @@ class SlurmLauncher(Launcher):
 
         We assume *nix, but should be changed to PythonScriptBuilder when working
         """
-        builder = PythonScriptBuilder()
-        builder.processtype = "slurm"
-        return builder
+        return SlurmScriptBuilder(self)
 
     def processbuilder(self) -> SlurmProcessBuilder:
         """Returns the process builder for this launcher
 
         By default, returns the associated connector builder"""
         return SlurmProcessBuilder(self)
+
+
+class SlurmScriptBuilder(PythonScriptBuilder):
+    def __init__(self, launcher: SlurmLauncher, pythonpath=None):
+        super().__init__(pythonpath)
+        self.launcher = launcher
+        self.processtype = "slurm"
+
+    def write(self, job):
+        py_path = super().write(job)
+        main_path = py_path.parent
+
+        def relpath(path: Path):
+            return shquote(self.launcher.connector.resolve(path, main_path))
+
+        # Writes the sbatch shell script containing all the options
+        sh_path = job.jobpath / ("%s.sh" % job.name)
+        with sh_path.open("wt") as out:
+            out.write("""#!/bin/sh\n\n""")
+
+            workdir = self.launcher.connector.resolve(main_path)
+            out.write(f"#SBATCH --chdir={shquote(workdir)}\n")
+            out.write(f"""#SBATCH --error={relpath(job.stderr)}\n""")
+            out.write(f"""#SBATCH --output={relpath(job.stdout)}\n""")
+
+            for arg in self.launcher.options.args():
+                out.write(f"""#SBATCH {arg}\n""")
+
+            # We finish by the call to srun
+            out.write(f"""\nsrun ./{relpath(py_path)}\n\n""")
+
+        self.launcher.connector.setExecutable(sh_path, True)
+        return sh_path

experimaestro/notifications.py

@@ -78,7 +78,6 @@ class Reporter(threading.Thread):
 
         self.progress_threshold = 0.01
         self.cv = threading.Condition()
-        self.start()
 
     def stop(self):
         self.stopping = True
@@ -222,6 +221,7 @@ class Reporter(threading.Thread):
             taskpath = TaskEnv.instance().taskpath
             assert taskpath is not None, "Task path is not defined"
             Reporter.INSTANCE = Reporter(taskpath)
+            Reporter.INSTANCE.start()
         return Reporter.INSTANCE
 
 

experimaestro/run.py

@@ -140,10 +140,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,10 +505,12 @@ 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
 
@@ -598,9 +646,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 +740,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():
@@ -853,6 +906,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 +917,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 +964,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 +994,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 +1003,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 +1013,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 +1039,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 +1054,27 @@ 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()
 
+    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