experimaestro 1.5.3__py3-none-any.whl → 1.6.0__py3-none-any.whl

experimaestro/core/objects.py

@@ -583,6 +583,9 @@ class ConfigInformation:
         # State information
         self.job = None
 
+        #: True when this configuration was loaded from disk
+        self.loaded = False
+
         # Explicitely added dependencies
         self.dependencies = []
 
@@ -846,8 +849,8 @@ class ConfigInformation:
                 dependencies, path + ["__init_tasks__"], taskids
             )
 
-        # Check for an associated task
-        if self.task:
+        # Check for an associated task (and not loaded)
+        if self.task and not self.loaded:
             if id(self.task) not in taskids:
                 taskids.add(id(self.task))
                 dependencies.add(self.task.__xpm__.dependency())
@@ -916,14 +919,14 @@ class ConfigInformation:
 
         # --- Submit the job
 
+        # Sets the init tasks
+        self.init_tasks = init_tasks
+
         # Creates a new job
         self.job = self.xpmtype.task(
             self.pyobject, launcher=launcher, workspace=workspace, run_mode=run_mode
         )
 
-        # Sets the init tasks
-        self.init_tasks = init_tasks
-
         # Validate the object
         job_context = JobContext(self.job)
         self.validate_and_seal(job_context)
@@ -979,6 +982,9 @@ class ConfigInformation:
                 elif self.job.failedpath.is_file():
                     color = "light_red"
                     cprint(f"[failed] {s}", color, file=sys.stderr)
+                elif self.job.pidpath.is_file():
+                    color = "blue"
+                    cprint(f"[running] {s}", color, file=sys.stderr)
                 else:
                     color = "light_blue"
                     cprint(f"[not run] {s}", color, file=sys.stderr)
@@ -1382,12 +1388,13 @@ class ConfigInformation:
             else:
                 o.__init__()
                 xpminfo = o.__xpm__  # type: ConfigInformation
+                xpminfo.loaded = True
 
                 meta = definition.get("meta", None)
                 if meta:
                     xpminfo._meta = meta
                 if xpminfo.xpmtype.task is not None:
-                    o.__xpm__.job = object()
+                    xpminfo.job = object()
 
             # Set the fields
             for name, value in definition["fields"].items():

experimaestro/core/types.py

@@ -325,12 +325,17 @@ class ObjectType(Type):
 
         # Get the module
         module = inspect.getmodule(self.originaltype)
-        if getattr(module, "__file__", None) is None:
+        self._module = module.__name__
+        self._package = module.__package__
+
+        if self._module and self._package:
             self._file = None
         else:
             self._file = Path(inspect.getfile(self.originaltype)).absolute()
-        self._module = module.__name__
-        self._package = module.__package__
+
+        assert (
+            self._module and self._package
+        ) or self._file, f"Could not detect module/file for {self.originaltype}"
 
         # The class of the object
 

experimaestro/experiments/cli.py

@@ -1,27 +1,29 @@
+import imp
+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,29 @@ 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.pythonpath = set()
 
-    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))
+    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")
 
-    return data[::-1]
+        with yaml_file.open("rt") as fp:
+            _data = yaml.full_load(fp)
+        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.pythonpath.add(path.resolve())
+            else:
+                self.pythonpath.add((yaml_file.parent / path).resolve())
 
 
 @click.option("--debug", is_flag=True, help="Print debug information")
@@ -105,11 +118,20 @@ def load(yaml_file: Path):
 @click.option(
     "--file", "xp_file", 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 +150,55 @@ 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
+    pythonpath = list(conf_loader.pythonpath)
+    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 pythonpath:
+                pythonpath.append(xp_file.parent)
+            logging.info("Using python path: %s", ", ".join(str(s) for s in pythonpath))
+
+    assert (
+        module_name or xp_file
+    ), "Either the module name or experiment file should be given"
 
     # --- Set some options
 
@@ -165,32 +206,35 @@ 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
-
-    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())
+    # --- Finds the "run" function
 
-    code = compile(source, filename=the__file__, mode="exec")
-    _locals: Dict[str, Any] = {}
-
-    sys.path.append(str(xp_file.parent.absolute()))
     try:
-        exec(code, _locals, _locals)
+        for path in pythonpath:
+            sys.path.append(str(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),
+                )
+        else:
+            # Module
+            mod = importlib.import_module(module_name)
+
+        helper = getattr(mod, "run", None)
     finally:
-        sys.path.pop()
+        pass
 
     # --- ... and runs it
-    helper = _locals.get("run", None)
     if helper is None:
-        raise ValueError(f"Could not find run function in {the__file__}")
+        raise ValueError(f"Could not find run function in {xp_file}")
 
     if not isinstance(helper, ExperimentHelper):
         helper = ExperimentHelper(helper)
@@ -205,9 +249,6 @@ def experiments_cli(  # noqa: C901
     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)
@@ -224,26 +265,19 @@ def experiments_cli(  # noqa: C901
         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()))
 
     # --- Runs the experiment
     with experiment(
-        workdir, configuration.id, host=host, port=port, run_mode=run_mode
+        ws_env, configuration.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)
 
         try:

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"""
 

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

@@ -1,27 +1,15 @@
-# Launcher registry
+# Configuration registers
+
+from typing import ClassVar, Dict, Optional, Set, Type, Union
 
-from dataclasses import dataclass
-import itertools
-from types import new_class
-from typing import ClassVar, Dict, List, Optional, Set, Type, Union
-from experimaestro import Annotated
 from pathlib import Path
 import typing
+from omegaconf import DictConfig, OmegaConf, SCMode
 import pkg_resources
-import humanfriendly
-import yaml
-from yaml import Loader, Dumper
 from experimaestro.utils import logger
-from experimaestro.utils.yaml import (
-    Initialize,
-    YAMLDataClass,
-    YAMLException,
-    YAMLList,
-    add_path_resolvers,
-)
 
-from .base import LauncherConfiguration, ConnectorConfiguration, TokenConfiguration
-from .specs import CPUSpecification, CudaSpecification, HostRequirement
+from .base import ConnectorConfiguration, TokenConfiguration
+from .specs import HostRequirement
 
 if typing.TYPE_CHECKING:
     from experimaestro.launchers import Launcher
@@ -32,80 +20,20 @@ class LauncherNotFoundError(Exception):
     pass
 
 
-@dataclass
-class GPU(YAMLDataClass):
-    """Represents a GPU"""
-
-    model: str
-    count: int
-    memory: Annotated[int, Initialize(humanfriendly.parse_size)]
-
-    def to_spec(self):
-        return [CudaSpecification(self.memory, self.model) for _ in range(self.count)]
-
-
-class GPUList(YAMLList[GPU]):
-    """Represents a list of GPUs"""
-
-    def __repr__(self):
-        return f"GPUs({super().__repr__()})"
-
-    def to_spec(self) -> List[CudaSpecification]:
-        return list(itertools.chain(*[gpu.to_spec() for gpu in self]))
-
-
-@dataclass
-class CPU(YAMLDataClass):
-    """Represents a CPU"""
-
-    memory: Annotated[int, Initialize(humanfriendly.parse_size)] = 0
-    cores: int = 1
-
-    def to_spec(self):
-        return CPUSpecification(self.memory, self.cores)
-
-
-@dataclass
-class Host(YAMLDataClass):
-    name: str
-    gpus: List[GPU]
-    launchers: List[str]
-
-
-Launchers = Dict[str, List[LauncherConfiguration]]
 Connectors = Dict[str, Dict[str, ConnectorConfiguration]]
 Tokens = Dict[str, Dict[str, TokenConfiguration]]
 
 
-def new_loader(name: str) -> Type[Loader]:
-    return new_class("LauncherLoader", (yaml.FullLoader,))  # type: ignore
-
-
-def load_yaml(loader_cls: Type[Loader], path: Path):
+def load_yaml(schema, path: Path):
     if not path.is_file():
-        return None
+        return {}
 
-    logger.warning(
-        "Using YAML file to configure launchers is deprecated. Please remove %s using launchers.py",
-        path,
-    )
     logger.debug("Loading %s", path)
     with path.open("rt") as fp:
-        loader = loader_cls(fp)
-        try:
-            return loader.get_single_data()
-        finally:
-            loader.dispose()
-
-
-def unknown_error(loader: Loader, node):
-    raise YAMLException(
-        "",
-        node.start_mark.name,
-        node.start_mark.line,
-        node.start_mark.column,
-        f"No handler defined for key {node}",
-    )
+        cfg = OmegaConf.load(fp)
+        return OmegaConf.to_container(
+            OmegaConf.merge(cfg, schema), structured_config_mode=SCMode.INSTANTIATE
+        )
 
 
 class LauncherRegistry:
@@ -132,27 +60,14 @@ class LauncherRegistry:
         LauncherRegistry.CURRENT_CONFIG_DIR = config_dir
 
     def __init__(self, basepath: Path):
-        self.LauncherLoader: Type[Loader] = new_loader("LauncherLoader")
-        self.ConnectorLoader: Type[Loader] = new_loader("ConnectorLoader")
-        self.TokenLoader: Type[Loader] = new_loader("TokenLoader")
-        self.Dumper: Type[Dumper] = new_class("CustomDumper", (Dumper,), {})
+        self.connectors_schema = DictConfig({})
+        self.tokens_schema = DictConfig({})
         self.find_launcher_fn = None
 
-        # Add safeguards
-        add_path_resolvers(
-            self.LauncherLoader,
-            [],
-            Dict[str, LauncherConfiguration],
-            dumper=self.Dumper,
-        )
-
         # Use entry points for connectors and launchers
         for entry_point in pkg_resources.iter_entry_points("experimaestro.connectors"):
             entry_point.load().init_registry(self)
 
-        for entry_point in pkg_resources.iter_entry_points("experimaestro.launchers"):
-            entry_point.load().init_registry(self)
-
         for entry_point in pkg_resources.iter_entry_points("experimaestro.tokens"):
             entry_point.load().init_registry(self)
 
@@ -172,32 +87,16 @@ class LauncherRegistry:
                 logger.warn("No find_launcher() function was found in %s", launchers_py)
 
         # Read the configuration file
-        launchers: Launchers = (
-            load_yaml(self.LauncherLoader, basepath / "launchers.yaml") or {}
-        )
-        self.launchers = sorted(
-            itertools.chain(*launchers.values()), key=lambda launcher: -launcher.weight
-        )
-
-        self.connectors: Connectors = (
-            load_yaml(self.ConnectorLoader, basepath / "connectors.yaml") or {}
-        )
-        self.tokens: Tokens = (
-            load_yaml(self.TokenLoader, basepath / "tokens.yaml") or {}
-        )
-
-    def register_launcher(self, identifier: str, cls: Type[YAMLDataClass]):
-        add_path_resolvers(
-            self.LauncherLoader, [identifier, None], cls, dumper=self.Dumper
+        self.connectors = load_yaml(
+            self.connectors_schema, basepath / "connectors.yaml"
         )
+        self.tokens = load_yaml(self.tokens_schema, basepath / "tokens.yaml")
 
-    def register_connector(self, identifier: str, cls: Type[YAMLDataClass]):
-        add_path_resolvers(
-            self.ConnectorLoader, [identifier, None], cls, dumper=self.Dumper
-        )
+    def register_connector(self, identifier: str, cls: Type):
+        self.connectors_schema.merge_with({identifier: cls})
 
-    def register_token(self, identifier: str, cls: Type[YAMLDataClass]):
-        add_path_resolvers(self.TokenLoader, [identifier], cls, dumper=self.Dumper)
+    def register_token(self, identifier: str, cls: Type):
+        self.tokens_schema.merge_with({identifier: cls})
 
     def getToken(self, identifier: str) -> "Token":
         for tokens in self.tokens.values():
@@ -227,7 +126,7 @@ class LauncherRegistry:
             tags: Restrict the launchers to those containing one of the specified tags
         """
 
-        if len(self.launchers) == 0 and self.find_launcher_fn is None:
+        if self.find_launcher_fn is None:
             logger.info("No launchers.yaml file: using local host ")
             from experimaestro.launchers.direct import DirectLauncher
             from experimaestro.connectors.local import LocalConnector
@@ -250,12 +149,6 @@ class LauncherRegistry:
                 if launcher := self.find_launcher_fn(spec, tags):
                     return launcher
 
-        # We have registered launchers
-        for spec in specs:
-            for handler in self.launchers:
-                if (not tags) or any((tag in tags) for tag in handler.tags):
-                    if launcher := handler.get(self, spec):
-                        return launcher
         return None
 
 

experimaestro/launchers/direct.py

@@ -1,15 +1,3 @@
-from dataclasses import dataclass, field
-from functools import cached_property
-from typing import Dict, List, Optional
-from experimaestro.launcherfinder import (
-    LauncherConfiguration,
-    LauncherRegistry,
-    HostRequirement,
-)
-from experimaestro.launcherfinder.registry import CPU, GPUList, YAMLDataClass
-from experimaestro.launcherfinder.specs import (
-    HostSpecification,
-)
 from experimaestro.scriptbuilder import PythonScriptBuilder
 from . import Launcher
 
@@ -18,40 +6,5 @@ class DirectLauncher(Launcher):
     def scriptbuilder(self):
         return PythonScriptBuilder()
 
-    @staticmethod
-    def init_registry(registry: LauncherRegistry):
-        registry.register_launcher("local", DirectLauncherConfiguration)
-
     def __str__(self):
         return f"DirectLauncher({self.connector})"
-
-
-@dataclass
-class DirectLauncherConfiguration(YAMLDataClass, LauncherConfiguration):
-    connector: str = "connector"
-    cpu: CPU = field(default_factory=CPU)
-    gpus: GPUList = field(default_factory=GPUList)
-    tokens: Optional[Dict[str, int]] = None
-    tags: List[str] = field(default_factory=lambda: [])
-    weight: int = 0
-    disable: bool = False
-
-    @cached_property
-    def spec(self) -> HostSpecification:
-        return HostSpecification(cpu=self.cpu.to_spec(), cuda=self.gpus.to_spec())
-
-    def get(
-        self, registry: LauncherRegistry, requirement: "HostRequirement"
-    ) -> Optional[Launcher]:
-        if requirement.match(self.spec):
-            launcher = DirectLauncher(connector=registry.getConnector(self.connector))
-            if self.tokens:
-                for token_identifier, count in self.tokens.items():
-                    token = registry.getToken(token_identifier)
-                    # TODO: handle the case where this is not a CounterToken
-                    launcher.addListener(
-                        lambda job: job.dependencies.add(token.dependency(count))
-                    )
-            return launcher
-
-        return None