experimaestro 1.5.7__py3-none-any.whl → 1.5.9__py3-none-any.whl

experimaestro/cli/jobs.py

@@ -90,8 +90,14 @@ def process(
                     if kill:
                         if perform:
                             process = info.getprocess()
-                            print("KILLING", process)
-                            process.kill()
+                            if process is None:
+                                cprint(
+                                    "internal error – no process could be retrieved",
+                                    "red",
+                                )
+                            else:
+                                cprint(f"KILLING {process}", "light_red")
+                                process.kill()
                         else:
                             print("KILLING (not performing)", process)
                     print(

experimaestro/connectors/ssh.py

@@ -3,7 +3,7 @@ from pathlib import Path, _posix_flavour
 import io
 import os
 import re
-from experimaestro.launcherfinder import LauncherRegistry, YAMLDataClass
+from experimaestro.launcherfinder import LauncherRegistry
 from fabric import Connection
 from invoke import Promise
 import invoke.exceptions
@@ -132,7 +132,7 @@ class SshPath(Path):
 
 
 @dataclass
-class SshConfiguration(YAMLDataClass):
+class SshConfiguration:
     hostname: str
 
     def create(self, registry: LauncherRegistry):

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,26 +1,28 @@
+import imp
 import inspect
 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 omegaconf import OmegaConf, SCMode
+from termcolor import cprint
+
 from experimaestro import LauncherRegistry, RunMode, experiment
-from experimaestro.experiments.configuration import ConfigurationBase
 from experimaestro.exceptions import HandledException
+from experimaestro.experiments.configuration import ConfigurationBase
 from experimaestro.settings import find_workspace
-from omegaconf import OmegaConf, SCMode
-from termcolor import cprint
 
 
 class ExperimentHelper:
     """Helper for experiments"""
 
-    # The experiment
     xp: experiment
+    """The experiment object"""
 
     #: Run function
     callable: "ExperimentCallable"
@@ -175,28 +177,23 @@ def experiments_cli(  # noqa: C901
     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())
-
-    code = compile(source, filename=the__file__, mode="exec")
-    _locals: Dict[str, Any] = {}
+    xp_file: Path = Path(yaml_file).parent / xp_file
 
-    sys.path.append(str(xp_file.parent.absolute()))
+    # --- Finds the "run" function
     try:
-        exec(code, _locals, _locals)
+        sys.path.append(str(xp_file.parent.absolute()))
+        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)
+            )
+            helper = getattr(mod, "run", None)
     finally:
         sys.path.pop()
 
     # --- ... 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)

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

experimaestro/scheduler/base.py

@@ -889,7 +889,7 @@ class experiment:
         """Shortcut to set the environment value"""
         if override or name not in self.workspace.env:
             logging.info("Setting environment: %s=%s", name, value)
-            self.worskpace.env[name] = value
+            self.workspace.env[name] = value
 
     def token(self, name: str, count: int):
         """Returns a token for this experiment

experimaestro/tests/launchers/config_slurm/launchers.py

@@ -0,0 +1,25 @@
+from experimaestro.launcherfinder.specs import (
+    CPUSpecification,
+    CudaSpecification,
+    HostRequirement,
+    HostSpecification,
+)
+from experimaestro.launchers.slurm.base import SlurmLauncher, SlurmOptions
+
+GIGA = 1024**3
+
+
+def find_launcher(requirements: HostRequirement, tags: set[str] = set()):
+    host = HostSpecification(
+        cpu=CPUSpecification(cores=16, memory=32 * GIGA),
+        max_duration=3600 * 24 * 10,
+        cuda=[CudaSpecification(memory=32 * GIGA) for _ in range(4)],
+    )
+    if match := requirements.match(host):
+        return SlurmLauncher(
+            options=SlurmOptions(
+                gpus_per_node=len(match.requirement.cuda_gpus),
+                partition="hard,electronic",
+                constraint="(A6000&GPU2&GPUM48G)|(A6000&GPU3&GPUM48G)|(RTX&GPU4&GPUM48G)",
+            )
+        )

experimaestro/tests/test_findlauncher.py

@@ -81,7 +81,7 @@ def slurm_constraint_split(constraint: str):
 def test_findlauncher_slurm():
     path = ResourcePathWrapper.create(f"{__package__ }.launchers", "config_slurm")
 
-    assert (path / "launchers.yaml").is_file()
+    assert (path / "launchers.py").is_file()
 
     registry = LauncherRegistry(path)
     launcher = registry.find("""duration=4 days & cuda(mem=24G) * 2""")

experimaestro/tokens.py

@@ -1,10 +1,12 @@
-"""Tokens are special types of dependency controlling the access to 
+"""Tokens are special types of dependency controlling the access to
 a computational resource (e.g. number of launched jobs, etc.)
 """
 
 from dataclasses import dataclass
 import sys
 from pathlib import Path
+
+from omegaconf import DictConfig
 from experimaestro.core.objects import Config
 import fasteners
 import threading
@@ -14,7 +16,6 @@ from typing import Dict
 from experimaestro.launcherfinder.base import TokenConfiguration
 
 from experimaestro.launcherfinder.registry import LauncherRegistry
-from experimaestro.utils.yaml import YAMLDict
 
 from .ipc import ipcom
 from .locking import Lock, LockError
@@ -87,7 +88,7 @@ class TokenFile:
         try:
             self.path = path
             with path.open("rt") as fp:
-                count, self.uri = [l.strip() for l in fp.readlines()]
+                count, self.uri = [line.strip() for line in fp.readlines()]
                 self.count = int(count)
         except Exception:
             logging.exception("Error while reading %s", self.path)
@@ -183,7 +184,10 @@ class CounterToken(Token, FileSystemEventHandler):
 
     @staticmethod
     def init_registry(registry: LauncherRegistry):
-        registry.register_token("countertoken", CounterTokenConfiguration)
+        registry.register_token(
+            "countertoken",
+            DictConfig({}, key_type=str, element_type=CounterConfiguration),
+        )
 
     def __init__(self, name: str, path: Path, count: int, force=True):
         """[summary]
@@ -456,7 +460,3 @@ class CounterConfiguration(TokenConfiguration):
         from experimaestro.connectors.local import LocalConnector
 
         return LocalConnector.instance().createtoken(identifier, self.tokens)
-
-
-class CounterTokenConfiguration(YAMLDict[CounterConfiguration]):
-    pass

experimaestro/utils/resources.py

@@ -1,11 +1,12 @@
 from contextlib import contextmanager
+from os import PathLike
 from pathlib import Path
 from typing import Union
 from importlib import resources
 from experimaestro.compat import cached_property
 
 
-class ResourcePathWrapper:
+class ResourcePathWrapper(PathLike):
     """Simple wrapper for resource path"""
 
     def __init__(self, path: Path):
@@ -30,6 +31,9 @@ class ResourcePathWrapper:
     def is_file(self):
         return resources.is_resource(self.package, self.name)
 
+    def __fspath__(self):
+        return resources.path(self.package, self.name).__fspath__()
+
     @contextmanager
     def open(self, *args, **kwargs):
         with resources.path(self.package, self.name) as path:

{experimaestro-1.5.7.dist-info → experimaestro-1.5.9.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: experimaestro
-Version: 1.5.7
+Version: 1.5.9
 Summary: "Experimaestro is a computer science experiment manager"
 Home-page: https://github.com/experimaestro/experimaestro-python
 License: GPL-3