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

experimaestro/launcherfinder/registry.py

@@ -1,27 +1,15 @@
-# Launcher registry
+# Configuration registers
+
+from contextlib import contextmanager
+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
-import pkg_resources
-import humanfriendly
-import yaml
-from yaml import Loader, Dumper
+from omegaconf import DictConfig, OmegaConf, SCMode
+from importlib.metadata import entry_points
 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, RequirementUnion
 
 if typing.TYPE_CHECKING:
     from experimaestro.launchers import Launcher
@@ -32,80 +20,30 @@ 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()
+        cfg = OmegaConf.load(fp)
+        return OmegaConf.to_container(
+            OmegaConf.merge(cfg, schema), structured_config_mode=SCMode.INSTANTIATE
+        )
 
 
-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}",
-    )
+@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:
@@ -132,72 +70,47 @@ 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"):
+        for entry_point in entry_points(group="experimaestro.connectors"):
             entry_point.load().init_registry(self)
 
-        for entry_point in pkg_resources.iter_entry_points("experimaestro.tokens"):
+        for entry_point in entry_points(group="experimaestro.tokens"):
             entry_point.load().init_registry(self)
 
         # 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
-        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 {}
+        self.connectors = load_yaml(
+            self.connectors_schema, basepath / "connectors.yaml"
         )
+        self.tokens = load_yaml(self.tokens_schema, basepath / "tokens.yaml")
 
-    def register_launcher(self, identifier: str, cls: Type[YAMLDataClass]):
-        add_path_resolvers(
-            self.LauncherLoader, [identifier, None], cls, dumper=self.Dumper
-        )
-
-    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 +140,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
@@ -237,25 +150,24 @@ 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:
-                if launcher := self.find_launcher_fn(*specs, tags):
+            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
 
-        # 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/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/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/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,14 +256,22 @@ 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("slurm sbatch command: %s", builder.command)
+        logger.info(
+            "slurm sbatch command: %s", " ".join(f'"{s}"' for s in builder.command)
+        )
         handler = OutputCaptureHandler()
         builder.stdout = Redirect.pipe(handler)
         builder.stderr = Redirect.inherit()
@@ -425,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/mkdocs/base.py

@@ -4,12 +4,11 @@ See https://www.mkdocs.org/user-guide/plugins/ for plugin API documentation
 """
 
 from collections import defaultdict
-import functools
 import re
 from experimaestro.mkdocs.annotations import shoulddocument
 import requests
 from urllib.parse import urljoin
-from experimaestro.core.types import ObjectType, Type
+from experimaestro.core.types import ObjectType
 import mkdocs
 from pathlib import Path
 from typing import Dict, Iterator, List, Optional, Set, Tuple, Type as TypingType
@@ -76,7 +75,7 @@ class ObjectLatticeNode:
         return f"node({self.objecttype.identifier})"
 
     def isAncestor(self, other):
-        return issubclass(self.objecttype.configtype, other.objecttype.configtype)
+        return issubclass(self.objecttype.config_type, other.objecttype.config_type)
 
     def _addChild(self, child: "ObjectLatticeNode"):
         child.parents.add(self)
@@ -321,7 +320,7 @@ class Documentation(mkdocs.plugins.BasePlugin):
 
             for node in self.lattice.iter_all():
                 if node.objecttype is not None:
-                    member = node.objecttype.basetype
+                    member = node.objecttype.value_type
                     qname = f"{member.__module__}.{member.__qualname__}"
                     path = self.type2path[qname]
 
@@ -354,7 +353,7 @@ class Documentation(mkdocs.plugins.BasePlugin):
         # Now, sort according to descendant/ascendant relationship or name
         nodes = set()
         for _node in cfgs:
-            if issubclass(_node.objecttype.configtype, xpmtype.configtype):
+            if issubclass(_node.objecttype.config_type, xpmtype.config_type):
                 nodes.add(_node)
 
         # Removes so they are not generated twice
@@ -443,11 +442,10 @@ class Documentation(mkdocs.plugins.BasePlugin):
                 lines.append("\n\n")
 
             for name, argument in xpminfo.arguments.items():
-
                 if isinstance(argument.type, ObjectType):
-                    basetype = argument.type.basetype
+                    value_type = argument.type.value_type
                     typestr = self.getlink(
-                        page.url, f"{basetype.__module__}.{basetype.__qualname__}"
+                        page.url, f"{value_type.__module__}.{value_type.__qualname__}"
                     )
                 else:
                     typestr = argument.type.name()