experimaestro 1.5.1__py3-none-any.whl → 2.0.0a8__py3-none-any.whl

experimaestro/experiments/cli.py

@@ -1,27 +1,29 @@
+import datetime
+import importlib
 import inspect
-import itertools
 import json
 import logging
 import sys
 from pathlib import Path
-from typing import Any, List, Optional, Protocol, Tuple, Dict
+from typing import Any, List, Optional, Protocol, Tuple
 
 import click
 import omegaconf
 import yaml
-from experimaestro import LauncherRegistry, RunMode, experiment
-from experimaestro.experiments.configuration import ConfigurationBase
-from experimaestro.exceptions import HandledException
-from experimaestro.settings import get_settings, get_workspace
 from omegaconf import OmegaConf, SCMode
 from termcolor import cprint
 
+from experimaestro import LauncherRegistry, RunMode, experiment
+from experimaestro.exceptions import HandledException
+from experimaestro.experiments.configuration import ConfigurationBase
+from experimaestro.settings import find_workspace
+
 
 class ExperimentHelper:
     """Helper for experiments"""
 
-    # The experiment
     xp: experiment
+    """The experiment object"""
 
     #: Run function
     callable: "ExperimentCallable"
@@ -54,18 +56,40 @@ class ExperimentCallable(Protocol):
         ...
 
 
-def load(yaml_file: Path):
-    """Loads a YAML file, and parents one if they exist"""
-    if not yaml_file.exists() and yaml_file.suffix != ".yaml":
-        yaml_file = yaml_file.with_suffix(".yaml")
+class ConfigurationLoader:
+    def __init__(self):
+        self.yamls = []
+        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"""
+        if not yaml_file.exists() and yaml_file.suffix != ".yaml":
+            yaml_file = yaml_file.with_suffix(".yaml")
+
+        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())
 
-    with yaml_file.open("rt") as fp:
-        _data = yaml.full_load(fp)
-    data = [_data]
-    if parent := _data.get("parent", None):
-        data.extend(load(yaml_file.parent / parent))
+        if "module" in _data:
+            # Keeps track of the YAML file where the module was defined
+            self.yaml_module_file = yaml_file
 
-    return data[::-1]
+        if parent := _data.get("parent", None):
+            self.load(yaml_file.parent / parent)
+
+        self.yamls.append(_data)
+
+        for path in _data.get("pythonpath", []):
+            path = Path(path)
+            if path.is_absolute():
+                self.python_path.add(path.resolve())
+            else:
+                self.python_path.add((yaml_file.parent / path).resolve())
 
 
 @click.option("--debug", is_flag=True, help="Print debug information")
@@ -103,13 +127,25 @@ def load(yaml_file: Path):
     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"
+)
+@click.option(
+    "--workspace",
+    type=str,
+    default=None,
+    help="Workspace ID (reads from settings.yaml in experimaestro config)",
 )
 @click.option(
     "--workdir",
     type=str,
     default=None,
-    help="Working directory - if None, uses the default XPM " "working directory",
+    help="Working environment",
 )
 @click.option("--conf", "-c", "extra_conf", type=str, multiple=True)
 @click.option(
@@ -128,36 +164,57 @@ def experiments_cli(  # noqa: C901
     port: int,
     xpm_config_dir: Path,
     workdir: Optional[Path],
+    workspace: Optional[str],
     env: List[Tuple[str, str]],
     run_mode: RunMode,
     extra_conf: List[str],
     pre_yaml: List[str],
     post_yaml: List[str],
+    module_name: Optional[str],
     args: List[str],
     show: bool,
     debug: bool,
 ):
     """Run an experiment"""
+
     # --- Set the logger
     logging.getLogger().setLevel(logging.DEBUG if debug else logging.INFO)
     logging.getLogger("xpm.hash").setLevel(logging.INFO)
 
     # --- Loads the YAML
-    yamls = []
+    conf_loader = ConfigurationLoader()
     for y in pre_yaml:
-        yamls.extend(load(Path(y)))
-    yamls.extend(load(Path(yaml_file)))
+        conf_loader.load(Path(y))
+    conf_loader.load(Path(yaml_file))
     for y in post_yaml:
-        yamls.extend(load(Path(y)))
+        conf_loader.load(Path(y))
+
+    # --- Merge the YAMLs
+    configuration = OmegaConf.merge(*conf_loader.yamls)
+    if extra_conf:
+        configuration.merge_with(OmegaConf.from_dotlist(extra_conf))
 
     # --- Get the XP file
-    if xp_file is None:
-        for data in yamls[::-1]:
-            if xp_file := data.get("file", None):
-                break
+    python_path = list(conf_loader.python_path)
+    if module_name is None:
+        module_name = configuration.get("module", None)
 
-        if xp_file is None:
-            raise ValueError("No experiment file given")
+    if xp_file is None:
+        xp_file = configuration.get("file", None)
+        if xp_file:
+            assert (
+                not module_name
+            ), "Module name and experiment file are mutually exclusive options"
+            xp_file = Path(xp_file)
+            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
+    ), "Either the module name or experiment file should be given"
 
     # --- Set some options
 
@@ -165,91 +222,116 @@ def experiments_cli(  # noqa: C901
         assert xpm_config_dir.is_dir()
         LauncherRegistry.set_config_dir(xpm_config_dir)
 
-    # --- Loads the XP file
-    xp_file = Path(xp_file)
-    if not xp_file.exists() and xp_file.suffix != ".py":
-        xp_file = xp_file.with_suffix(".py")
-    xp_file = Path(yaml_file).parent / xp_file
+    # --- Finds the "run" function
 
-    with open(xp_file, "r") as f:
-        source = f.read()
-    if sys.version_info < (3, 9):
-        the__file__ = str(xp_file)
-    else:
-        the__file__ = str(xp_file.absolute())
+    # Modifies the Python path
+    for path in python_path:
+        sys.path.append(str(path))
 
-    code = compile(source, filename=the__file__, mode="exec")
-    _locals: Dict[str, Any] = {}
+    # --- 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
+        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
+        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
 
-    sys.path.append(str(xp_file.parent.absolute()))
-    try:
-        exec(code, _locals, _locals)
-    finally:
-        sys.path.pop()
+    helper = getattr(mod, "run", None)
 
     # --- ... and runs it
-    helper = _locals.get("run", None)
     if helper is None:
-        raise ValueError(f"Could not find run function in {the__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())
 
-    configuration = OmegaConf.merge(*yamls)
-    if extra_conf:
-        configuration.merge_with(OmegaConf.from_dotlist(extra_conf))
     if omegaconf_schema is not None:
         try:
             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
     )
 
-    # Get the working directory
-    settings = get_settings()
-    ws_env = {}
-    workdir = Path(workdir) if workdir else None
-    if (workdir is None) or (not workdir.is_dir()):
-        logging.info("Searching for workspace %s", workdir)
-        ws_settings = get_workspace(str(workdir))
-        workdir = ws_settings.path.expanduser()
-        ws_env = ws_settings.env
+    # 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(
-        workdir, 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 itertools.chain(settings.env.items(), ws_env.items(), env):
-            logging.info("Setting environment: %s=%s", key, value)
+        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

@@ -1,5 +1,5 @@
 from omegaconf import MISSING
-from typing import Optional
+from typing import Optional, List
 import attr
 
 try:
@@ -31,6 +31,12 @@ class ConfigurationBase:
     file: str = "experiment"
     """Relative path of the file containing a run function"""
 
+    module: Optional[str] = None
+    """Relative path of the file containing a run function"""
+
+    pythonpath: Optional[List[str]] = None
+    """Python path relative to the parent directory of the YAML file"""
+
     parent: Optional[str] = None
     """Relative path of a YAML file that should be merged"""
 
@@ -45,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/__init__.py

@@ -10,5 +10,5 @@ from .specs import (
     HostRequirement,
     MatchRequirement,
 )
-from .registry import find_launcher, LauncherRegistry, YAMLDataClass
+from .registry import find_launcher, LauncherRegistry
 from .parser import parse

experimaestro/launcherfinder/base.py

@@ -1,33 +1,17 @@
-from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, List, Optional
+from typing import TYPE_CHECKING
 
-from experimaestro.utils.yaml import YAMLDataClass
-from .specs import HostRequirement
 
 if TYPE_CHECKING:
-    from experimaestro.launchers import Launcher
     from experimaestro.connectors import Connector
     from experimaestro.tokens import Token
     from .registry import LauncherRegistry
 
 
-class LauncherConfiguration:
-    tags: List[str]
-    weight: int
-
-    """Generic class for a launcher configuration"""
-
-    def get(
-        self, registry: "LauncherRegistry", requirement: HostRequirement
-    ) -> Optional["Launcher"]:
-        raise NotImplementedError(f"For {self.__class__}")
-
-
 class ConnectorConfiguration:
     def create(self, registry: "LauncherRegistry") -> "Connector":
         raise NotImplementedError(f"For {self.__class__}")
 
 
-class TokenConfiguration(YAMLDataClass):
+class TokenConfiguration:
     def create(self, registry: "LauncherRegistry", identifier: str) -> "Token":
         raise NotImplementedError(f"For {self.__class__}")

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)