experimaestro 1.5.6__py3-none-any.whl → 1.5.8__py3-none-any.whl

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 experimaestro import LauncherRegistry, RunMode, experiment
-from experimaestro.experiments.configuration import ConfigurationBase
-from experimaestro.exceptions import HandledException
-from experimaestro.settings import 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"
@@ -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)
@@ -231,27 +228,8 @@ def experiments_cli(  # noqa: C901
     )
 
     # Define the workspace
-    workdir = Path(workdir) if workdir else None
-
-    if workspace:
-        ws_env = get_workspace(workspace)
-        if ws_env is None:
-            raise RuntimeError("No workspace named %s", workspace)
-
-        logging.info("Using workspace %s", ws_env.id)
-        if workdir:
-            # Overrides working directory
-            logging.info(" override working directory: %s", workdir)
-            ws_env.path = workdir
-        else:
-            workdir = ws_env.path
-    elif workdir:
-        logging.info("Using workdir %s", workdir)
-        ws_env = workdir
-    else:
-        ws_env = get_workspace()
-        assert ws_env is not None, "No workdir or workspace defined, and no default"
-        logging.info("Using default workspace %s", ws_env.id)
+    ws_env = find_workspace(workdir=workdir, workspace=workspace)
+    workdir = ws_env.path
 
     logging.info("Using working directory %s", str(workdir.resolve()))
 

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/launchers/slurm/base.py

@@ -262,7 +262,9 @@ class SlurmProcessBuilder(ProcessBuilder):
         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()

experimaestro/notifications.py

@@ -44,6 +44,12 @@ class LevelInformation:
         return f"[{self.level}] {self.desc} {int(self.progress*1000)/10}%"
 
 
+class ListenerInformation:
+    def __init__(self, url: str):
+        self.url = url
+        self.error_count = 0
+
+
 class Reporter(threading.Thread):
     NOTIFICATION_FOLDER = ".notifications"
 
@@ -59,7 +65,7 @@ class Reporter(threading.Thread):
         super().__init__(daemon=True)
         self.path = path / Reporter.NOTIFICATION_FOLDER
         self.path.mkdir(exist_ok=True)
-        self.urls: Dict[str, str] = {}
+        self.urls: Dict[str, ListenerInformation] = {}
 
         # Last check of notification URLs
         self.lastcheck = 0
@@ -80,7 +86,7 @@ class Reporter(threading.Thread):
             self.cv.notifyAll()
 
     @staticmethod
-    def isfatal_httperror(e: Exception) -> bool:
+    def isfatal_httperror(e: Exception, info: ListenerInformation) -> bool:
         """Returns True if this HTTP error indicates that the server won't recover"""
         if isinstance(e, HTTPError):
             if e.code >= 400 and e.code < 500:
@@ -90,6 +96,13 @@ class Reporter(threading.Thread):
                 return True
             if isinstance(e.reason, socket.gaierror) and e.reason.errno == -2:
                 return True
+            if isinstance(e.reason, TimeoutError):
+                info.error_count += 1
+
+            # Too many errors
+            if info.error_count > 3:
+                logger.info("Too many errors with %s", info.error_count)
+                return True
 
         return False
 
@@ -100,8 +113,8 @@ class Reporter(threading.Thread):
         mtime = os.path.getmtime(self.path)
         if mtime > self.lastcheck:
             for f in self.path.iterdir():
-                self.urls[f.name] = f.read_text().strip()
-                logger.info("Added new notification URL: %s", self.urls[f.name])
+                self.urls[f.name] = ListenerInformation(f.read_text().strip())
+                logger.info("Added new notification URL: %s", self.urls[f.name].url)
                 f.unlink()
 
             self.lastcheck = os.path.getmtime(self.path)
@@ -128,7 +141,9 @@ class Reporter(threading.Thread):
                         params = level.report()
 
                         # Go over all URLs
-                        for key, baseurl in self.urls.items():
+                        for key, info in self.urls.items():
+                            baseurl = info.url
+
                             url = "{}/progress?{}".format(
                                 baseurl, urllib.parse.urlencode(params)
                             )
@@ -147,7 +162,7 @@ class Reporter(threading.Thread):
                                     url,
                                     e,
                                 )
-                                if Reporter.isfatal_httperror(e):
+                                if Reporter.isfatal_httperror(e, info):
                                     toremove.append(key)
 
                 # Removes unvalid URLs
@@ -165,7 +180,8 @@ class Reporter(threading.Thread):
             self.check_urls()
             if self.urls:
                 # Go over all URLs
-                for key, baseurl in self.urls.items():
+                for key, info in self.urls.items():
+                    baseurl = info.url
                     url = "{}?status=eoj".format(baseurl)
                     try:
                         with urlopen(url) as _:
@@ -243,7 +259,7 @@ class xpm_tqdm(std_tqdm):
 
     def update(self, n=1):
         result = super().update(n)
-        if self.total is not None:
+        if self.total is not None and self.total > 0:
             progress(self.n / self.total, level=self.pos, console=False)
         return result
 

experimaestro/run.py

@@ -155,7 +155,6 @@ class TaskRunner:
                 self.donepath.touch()
 
                 # ... and finish the exit process
-                logger.info("This is the end (TODO: remove this line)")
                 raise
             else:
                 self.handle_error(e.code, None)

experimaestro/scheduler/base.py

@@ -803,11 +803,6 @@ class experiment:
             else None
         )
 
-        # Copy environment variable from main (but do not
-        # override)
-        for key, value in settings.env.items():
-            self.setenv(key, value, override=False)
-
         if os.environ.get("XPM_ENABLEFAULTHANDLER", "0") == "1":
             import faulthandler
 

experimaestro/settings.py

@@ -4,6 +4,7 @@ from dataclasses import field, dataclass
 from functools import lru_cache
 from pathlib import Path
 from typing import Dict, Optional, List
+import logging
 
 
 @dataclass
@@ -59,7 +60,9 @@ def get_settings(path: Optional[Path] = None) -> Settings:
 
         path = path or Path("~/.config/experimaestro/settings.yaml").expanduser()
         if not path.is_file():
-            return schema.to_object()
+            return OmegaConf.to_container(
+                schema, structured_config_mode=SCMode.INSTANTIATE
+            )
 
         conf = OmegaConf.load(path)
         return OmegaConf.to_container(
@@ -78,3 +81,28 @@ def get_workspace(id: Optional[str] = None) -> Optional[WorkspaceSettings]:
                 return workspace
 
     return None
+
+
+def find_workspace(*, workspace: Optional[str] = None, workdir: Optional[Path] = None):
+    """Find workspace"""
+    workdir = Path(workdir) if workdir else None
+
+    if workspace:
+        ws_env = get_workspace(workspace)
+        if ws_env is None:
+            raise RuntimeError("No workspace named %s", workspace)
+
+        logging.info("Using workspace %s", ws_env.id)
+        if workdir:
+            # Overrides working directory
+            logging.info(" override working directory: %s", workdir)
+            ws_env.path = workdir
+    elif workdir:
+        logging.info("Using workdir %s", workdir)
+        ws_env = WorkspaceSettings("", workdir)
+    else:
+        ws_env = get_workspace()
+        assert ws_env is not None, "No workdir or workspace defined, and no default"
+        logging.info("Using default workspace %s", ws_env.id)
+
+    return ws_env

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/tests/test_tags.py

@@ -2,6 +2,7 @@ from typing import Dict
 from pathlib import Path
 from experimaestro import (
     tag,
+    LightweightTask,
     config,
     argument,
     Config,
@@ -69,6 +70,40 @@ def test_inneroutput():
     assert evaluate.__xpm__.tags() == {"hello": "world"}
 
 
+def test_tags_init_tasks():
+    """Test tags within init tasks"""
+
+    class MyTask(Task):
+        pass
+
+    class InitTask(LightweightTask):
+        pass
+
+    class MyConfig(Config):
+        pass
+
+    class TaskWithOutput(Task):
+        x: Param[MyConfig]
+
+        def task_outputs(self, dep) -> MyConfig:
+            return dep(MyConfig())
+
+    init_task = InitTask().tag("hello", "world")
+    task = MyTask()
+    result = task.submit(run_mode=RunMode.DRY_RUN, init_tasks=[init_task])
+    assert result.tags() == {"hello": "world"}
+
+    other_task = TaskWithOutput(x=MyConfig().tag("hello", "world"))
+    assert other_task.tags() == {"hello": "world"}
+
+    result = other_task.submit(run_mode=RunMode.DRY_RUN)
+    assert isinstance(result, MyConfig)
+    assert result.tags() == {"hello": "world"}
+
+    result = MyTask().submit(run_mode=RunMode.DRY_RUN, init_tasks=[result])
+    assert result.tags() == {"hello": "world"}
+
+
 class TaskDirectoryContext(DirectoryContext):
     def __init__(self, task, path):
         super().__init__(path)