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

experimaestro/experiments/cli.py

@@ -1,4 +1,4 @@
-import imp
+import datetime
 import importlib
 import inspect
 import json
@@ -59,7 +59,8 @@ class ExperimentCallable(Protocol):
 class ConfigurationLoader:
     def __init__(self):
         self.yamls = []
-        self.pythonpath = set()
+        self.python_path = set()
+        self.yaml_module_file: None | Path = None
 
     def load(self, yaml_file: Path):
         """Loads a YAML file, and parents one if they exist"""
@@ -68,6 +69,16 @@ class ConfigurationLoader:
 
         with yaml_file.open("rt") as fp:
             _data = yaml.full_load(fp)
+
+        if "file" in _data:
+            path = Path(_data["file"])
+            if not path.is_absolute():
+                _data["file"] = str((yaml_file.parent / path).resolve())
+
+        if "module" in _data:
+            # Keeps track of the YAML file where the module was defined
+            self.yaml_module_file = yaml_file
+
         if parent := _data.get("parent", None):
             self.load(yaml_file.parent / parent)
 
@@ -76,9 +87,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")
@@ -116,7 +127,10 @@ class ConfigurationLoader:
     help="Port for monitoring (can be defined in the settings.yaml file)",
 )
 @click.option(
-    "--file", "xp_file", help="The file containing the main experimental code"
+    "--file",
+    "xp_file",
+    type=Path,
+    help="The file containing the main experimental code",
 )
 @click.option(
     "--module-name", "module_name", help="Module containing the experimental code"
@@ -181,7 +195,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 +206,11 @@ 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.absolute())
+            logging.info(
+                "Using python path: %s", ", ".join(str(s) for s in python_path)
+            )
 
     assert (
         module_name or xp_file
@@ -209,40 +225,59 @@ 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))
 
+    # --- Adds automatically the experiment module if not found
+    if module_name and conf_loader.yaml_module_file:
+        try:
+            importlib.import_module(module_name)
+        except ModuleNotFoundError:
+            # Try to setup a path
+            path = conf_loader.yaml_module_file.resolve()
+            for _ in range(len(module_name.split("."))):
+                path = path.parent
+
+            logging.info("Appending %s to python path", path)
+            sys.path.append(str(path))
+            python_path.append(path)
+
     if xp_file:
         if not xp_file.exists() and xp_file.suffix != ".py":
             xp_file = xp_file.with_suffix(".py")
         xp_file: Path = Path(yaml_file).parent / xp_file
-        with open(xp_file) as src:
-            module_name = xp_file.with_suffix("").name
-            mod = imp.load_module(
-                module_name,
-                src,
-                str(xp_file.absolute()),
-                (".py", "r", imp.PY_SOURCE),
-            )
+        module_name = xp_file.with_suffix("").name
+        spec = importlib.util.spec_from_file_location(
+            module_name, str(xp_file.absolute())
+        )
+        mod = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(mod)
     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)
 
     # --- ... and runs it
     if helper is None:
-        raise ValueError(f"Could not find run function in {xp_file}")
+        raise click.ClickException(
+            f"Could not find run function in {xp_file if xp_file else module_name}"
+        )
 
     if not isinstance(helper, ExperimentHelper):
         helper = ExperimentHelper(helper)
 
     parameters = inspect.signature(helper.callable).parameters
     list_parameters = list(parameters.values())
-    assert len(list_parameters) == 2, (
-        "Callable function should only "
-        f"have two arguments (got {len(list_parameters)})"
-    )
+    if len(list_parameters) != 2:
+        raise click.ClickException(
+            f"run in {xp_file if xp_file else module_name} function should only "
+            f"have two arguments (got {len(list_parameters)}), "
+        )
 
     schema = list_parameters[1].annotation
     omegaconf_schema = OmegaConf.structured(schema())
@@ -252,36 +287,51 @@ def experiments_cli(  # noqa: C901
             configuration = OmegaConf.merge(omegaconf_schema, configuration)
         except omegaconf.errors.ConfigKeyError as e:
             cprint(f"Error in configuration:\n\n{e}", "red", file=sys.stderr)
-            sys.exit(1)
+            raise click.ClickException("Error in configuration")
 
     if show:
         print(json.dumps(OmegaConf.to_container(configuration)))  # noqa: T201
         sys.exit(0)
 
     # Move to an object container
-    configuration = OmegaConf.to_container(
+    xp_configuration: ConfigurationBase = OmegaConf.to_container(
         configuration, structured_config_mode=SCMode.INSTANTIATE
     )
 
     # Define the workspace
     ws_env = find_workspace(workdir=workdir, workspace=workspace)
+
     workdir = ws_env.path
 
-    logging.info("Using working directory %s", str(workdir.resolve()))
+    # --- Sets up the experiment ID
 
     # --- Runs the experiment
+    if xp_configuration.add_timestamp:
+        timestamp = datetime.datetime.now().strftime("%Y%m%d-%H%M")
+        experiment_id = f"""{xp_configuration.id}-{timestamp}"""
+    else:
+        experiment_id = xp_configuration.id
+
+    logging.info(
+        "Running experiment %s working directory %s",
+        experiment_id,
+        str(workdir.resolve()),
+    )
     with experiment(
-        ws_env, configuration.id, host=host, port=port, run_mode=run_mode
+        ws_env, experiment_id, host=host, port=port, run_mode=run_mode
     ) as xp:
         # Set up the environment
         # (1) global settings (2) workspace settings and (3) command line settings
         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
-            helper.run(list(args), configuration)
+            helper.run(list(args), xp_configuration)
 
             # ... and wait
             xp.wait()

experimaestro/experiments/configuration.py

@@ -51,3 +51,6 @@ class ConfigurationBase:
 
     description: str = ""
     """Description of the experiment"""
+
+    add_timestamp: bool = False
+    """Adds a timestamp YYYY_MM_DD-HH_MM to the experiment ID"""

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/parser.py

@@ -23,7 +23,7 @@ class SuppressStrMatch(StrMatch):
 
 
 def mem_spec():
-    return "mem", "=", RegExMatch(r"\d+(G|M)?")
+    return "mem", "=", RegExMatch(r"\d+(GiB|MiB|G|M)?")
 
 
 def cores_spec():
@@ -51,7 +51,12 @@ def cpu():
 
 
 def duration():
-    return "duration", "=", RegExMatch(r"\d+"), RegExMatch(r"h(ours)?|d(ays)?")
+    return (
+        "duration",
+        "=",
+        RegExMatch(r"\d+"),
+        RegExMatch(r"h(ours?)?|d(ays?)?|m(ins?)?"),
+    )
 
 
 def one_spec():
@@ -67,7 +72,7 @@ def grammar():
 
 class Visitor(PTNodeVisitor):
     def visit_grammar(self, node, children):
-        return [child for child in children]
+        return specs.RequirementUnion(*[child for child in children])
 
     def visit_one_spec(self, node, children):
         return reduce(lambda x, el: x & el, children)

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,9 +8,8 @@ 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
+from .specs import HostRequirement, RequirementUnion
 
 if typing.TYPE_CHECKING:
     from experimaestro.launchers import Launcher
@@ -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
@@ -73,18 +83,22 @@ class LauncherRegistry:
 
         # Register the find launcher function if it exists
         launchers_py = basepath / "launchers.py"
+        print(f"basepath {launchers_py}")
         if launchers_py.is_file():
             logger.info("Loading %s", launchers_py)
 
             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(
@@ -136,17 +150,22 @@ class LauncherRegistry:
         # Parse specs
         from .parser import parse
 
-        specs = []
+        specs = RequirementUnion()
         for spec in input_specs:
             if isinstance(spec, str):
-                specs.extend(parse(spec))
+                specs.add(parse(spec))
             else:
-                specs.append(spec)
+                specs.add(spec)
 
         # Use launcher function
+        from experimaestro.launchers import Launcher
+
         if self.find_launcher_fn is not None:
-            for spec in specs:
+            for spec in specs.requirements:
                 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/launcherfinder/specs.py

@@ -1,4 +1,6 @@
+from abc import ABC, abstractmethod
 import logging
+import math
 from attr import Factory
 from attrs import define
 from copy import copy, deepcopy
@@ -20,9 +22,6 @@ class CudaSpecification:
     min_memory: int = 0
     """Minimum request memory (in bytes)"""
 
-    def __lt__(self, other: "CudaSpecification"):
-        return self.memory < other.memory
-
     def match(self, spec: "CudaSpecification"):
         """Returns True if the specification matches this host"""
         return (self.memory >= spec.memory) and (self.min_memory <= spec.memory)
@@ -30,7 +29,7 @@ class CudaSpecification:
     def __repr__(self):
         return (
             f"CUDA({self.model} "
-            f"{format_size(self.memory)}/{format_size(self.min_memory)})"
+            f"max={format_size(self.memory, binary=True)}/min={format_size(self.min_memory, binary=True)})"
         )
 
 
@@ -48,8 +47,15 @@ class CPUSpecification:
     cpu_per_gpu: int = 0
     """Number of CPU per GPU (0 if not defined)"""
 
-    def __lt__(self, other: "CPUSpecification"):
-        return self.memory < other.memory and self.cores < other.cores
+    def __repr__(self):
+        return (
+            f"CPU("
+            f"mem={format_size(self.memory, binary=True)}, cores={self.cores}"
+            ")"
+        )
+
+    def match(self, other: "CPUSpecification"):
+        return (self.memory >= other.memory) and (self.cores >= other.cores)
 
     def total_memory(self, gpus: int = 0):
         return max(
@@ -91,10 +97,11 @@ class MatchRequirement:
     requirement: "HostSimpleRequirement"
 
 
-class HostRequirement:
+class HostRequirement(ABC):
     """A requirement must be a disjunction of host requirements"""
 
     requirements: List["HostSimpleRequirement"]
+    """List of requirements (by order of priority)"""
 
     def __init__(self) -> None:
         self.requirements = []
@@ -105,6 +112,12 @@ class HostRequirement:
     def match(self, host: HostSpecification) -> Optional[MatchRequirement]:
         raise NotImplementedError()
 
+    @abstractmethod
+    def multiply_duration(self, coefficient: float) -> "HostRequirement":
+        """Returns a new HostRequirement with a duration multiplied by the
+        provided coefficient"""
+        ...
+
 
 class RequirementUnion(HostRequirement):
     """Ordered list of simple host requirements -- the first one is the priority"""
@@ -114,6 +127,16 @@ class RequirementUnion(HostRequirement):
     def __init__(self, *requirements: "HostSimpleRequirement"):
         self.requirements = list(requirements)
 
+    def add(self, requirement: "HostRequirement"):
+        match requirement:
+            case HostSimpleRequirement():
+                self.requirements.extend(*requirement.requirements)
+            case RequirementUnion():
+                self.requirements.append(requirement)
+            case _:
+                raise RuntimeError("Cannot handle type %s", type(requirement))
+        return self
+
     def match(self, host: HostSpecification) -> Optional[MatchRequirement]:
         """Returns the matched requirement (if any)"""
 
@@ -128,12 +151,23 @@ class RequirementUnion(HostRequirement):
 
         return argmax
 
+    def multiply_duration(self, coefficient: float) -> "RequirementUnion":
+        return RequirementUnion(
+            *[r.multiply_duration(coefficient) for r in self.requirements]
+        )
+
+    def __repr__(self):
+        return " | ".join(repr(r) for r in self.requirements)
+
 
 class HostSimpleRequirement(HostRequirement):
     """Simple host requirement"""
 
     cuda_gpus: List["CudaSpecification"]
+    """Specification for CUDA gpus"""
+
     cpu: "CPUSpecification"
+    """Specification for CPU"""
 
     duration: int
     """Requested duration (in seconds)"""
@@ -141,6 +175,11 @@ class HostSimpleRequirement(HostRequirement):
     def __repr__(self):
         return f"Req(cpu={self.cpu}, cuda={self.cuda_gpus}, duration={self.duration})"
 
+    def multiply_duration(self, coefficient: float) -> "HostSimpleRequirement":
+        r = HostSimpleRequirement(self)
+        r.duration = math.ceil(self.duration * coefficient)
+        return r
+
     def __init__(self, *reqs: "HostSimpleRequirement"):
         self.cuda_gpus = []
         self.cpu = CPUSpecification(0, 0)
@@ -158,7 +197,7 @@ class HostSimpleRequirement(HostRequirement):
         self.cpu.cores = max(req.cpu.cores, self.cpu.cores)
         self.duration = max(req.duration, self.duration)
         self.cuda_gpus.extend(req.cuda_gpus)
-        self.cuda_gpus.sort()
+        self.cuda_gpus.sort(key=lambda cuda: -cuda.memory)
 
     def match(self, host: HostSpecification) -> Optional[MatchRequirement]:
         if self.cuda_gpus:
@@ -182,7 +221,7 @@ class HostSimpleRequirement(HostRequirement):
             )
             return None
 
-        if host.cpu < self.cpu:
+        if not host.cpu.match(self.cpu):
             return None
 
         if host.max_duration > 0 and self.duration > host.max_duration:
@@ -197,7 +236,7 @@ class HostSimpleRequirement(HostRequirement):
         _self = deepcopy(self)
         for _ in range(count - 1):
             _self.cuda_gpus.extend(self.cuda_gpus)
-        _self.cuda_gpus.sort()
+        _self.cuda_gpus.sort(key=lambda cuda: -cuda.memory)
 
         return _self
 

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 (
@@ -86,7 +87,7 @@ class SlurmProcessWatcher(threading.Thread):
     WATCHERS: Dict[Tuple[Tuple[str, Any]], "SlurmProcessWatcher"] = {}
 
     def __init__(self, launcher: "SlurmLauncher"):
-        super().__init__()
+        super().__init__(daemon=True)
         self.launcher = launcher
         self.count = 1
         self.jobs: Dict[str, SlurmJobState] = {}
@@ -183,10 +184,10 @@ class BatchSlurmProcess(Process):
                 if state and state.finished():
                     return 0 if state.slurm_state == "COMPLETED" else 1
 
-    async def aio_state(self):
+    async def aio_state(self, timeout: float | None = None) -> ProcessState:
         def check():
             with SlurmProcessWatcher.get(self.launcher) as watcher:
-                jobinfo = watcher.getjob(self.jobid)
+                jobinfo = watcher.getjob(self.jobid, timeout=timeout)
                 return jobinfo.state if jobinfo else ProcessState.SCHEDULED
 
         return await asyncThreadcheck("slurm.aio_isrunning", check)
@@ -211,7 +212,7 @@ class BatchSlurmProcess(Process):
 
         # Checks that the process is running
         with SlurmProcessWatcher.get(launcher) as watcher:
-            logger.info("Checking SLURM job %s", process.jobid)
+            logger.debug("Checking SLURM job %s", process.jobid)
             jobinfo = watcher.getjob(process.jobid, timeout=0.1)
             if jobinfo and jobinfo.state.running:
                 logger.debug(
@@ -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/mkdocs/__init__.py

@@ -1 +1 @@
-# from .base import Documentation
+from .base import Documentation  # noqa: F401

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
@@ -110,6 +109,7 @@ class Reporter(threading.Thread):
         return any(level.modified(self) for level in self.levels)
 
     def check_urls(self):
+        """Check whether we have new schedulers to notify"""
         mtime = os.path.getmtime(self.path)
         if mtime > self.lastcheck:
             for f in self.path.iterdir():
@@ -222,6 +222,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