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

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

@@ -92,7 +92,7 @@ class TaskRunner:
             logger.info("Finished cleanup")
 
     def handle_error(self, code, frame_type):
-        logger.info("Finished with code %d", code)
+        logger.info("Error handler: finished with code %d", code)
         self.failedpath.write_text(str(code))
         self.cleanup()
         logger.info("Exiting")
@@ -100,8 +100,18 @@ class TaskRunner:
 
     def run(self):
         atexit.register(self.cleanup)
-        signal.signal(signal.SIGTERM, self.handle_error)
-        signal.signal(signal.SIGINT, self.handle_error)
+        sigterm_handler = signal.signal(signal.SIGTERM, self.handle_error)
+        sigint_handler = signal.signal(signal.SIGINT, self.handle_error)
+
+        def remove_signal_handlers(remove_cleanup=True):
+            """Removes cleanup in forked processes"""
+            signal.signal(signal.SIGTERM, sigterm_handler)
+            signal.signal(signal.SIGINT, sigint_handler)
+            atexit.unregister(self.cleanup)
+
+        if sys.platform != "win32":
+            os.register_at_fork(after_in_child=remove_signal_handlers)
+
         try:
             workdir = self.scriptpath.parent
             os.chdir(workdir)
@@ -129,6 +139,10 @@ class TaskRunner:
                 self.started = True
                 run(workdir / "params.json")
 
+                # ... remove the handlers
+                logger.info("Task ended successfully")
+                remove_signal_handlers(remove_cleanup=False)
+
                 # Everything went OK
                 sys.exit(0)
         except Exception:
@@ -137,6 +151,10 @@ class TaskRunner:
 
         except SystemExit as e:
             if e.code == 0:
+                # Normal exit, just create the ".done" file
                 self.donepath.touch()
+
+                # ... and finish the exit process
+                raise
             else:
                 self.handle_error(e.code, None)

experimaestro/scheduler/base.py

@@ -1,5 +1,6 @@
 from collections import ChainMap
 from functools import cached_property
+import logging
 import os
 from pathlib import Path
 from shutil import rmtree
@@ -13,13 +14,12 @@ from experimaestro.exceptions import HandledException
 from experimaestro.notifications import LevelInformation, Reporter
 from typing import Dict
 from experimaestro.scheduler.services import Service
-from experimaestro.settings import get_settings
+from experimaestro.settings import WorkspaceSettings, get_settings
 
 
 from experimaestro.core.objects import Config, ConfigWalkContext
 from experimaestro.utils import logger
 from experimaestro.locking import Locks, LockError, Lock
-from .environment import Environment
 from .workspace import RunMode, Workspace
 from .dependencies import Dependency, DependencyStatus, Resource
 import concurrent.futures
@@ -180,7 +180,7 @@ class Job(Resource):
         return ChainMap(
             {},
             self.launcher.environ if self.launcher else {},
-            self.workspace.environment.environ if self.workspace else {},
+            self.workspace.env if self.workspace else {},
         )
 
     @property
@@ -508,6 +508,12 @@ class Scheduler:
         job.scheduler = self
         self.waitingjobs.add(job)
 
+        # Check that we don't have a completed job in
+        # alternate directories
+        for jobspath in experiment.current().alt_jobspaths:
+            # FIXME: check if done
+            pass
+
         # Creates a link into the experiment folder
         path = experiment.current().jobspath / job.relpath
         path.parent.mkdir(parents=True, exist_ok=True)
@@ -719,7 +725,7 @@ class experiment:
     ```
     """
 
-    # Current experiment
+    #: Current experiment
     CURRENT: Optional["experiment"] = None
 
     @staticmethod
@@ -733,7 +739,7 @@ class experiment:
 
     def __init__(
         self,
-        env: Union[Path, str, Environment],
+        env: Union[Path, str, WorkspaceSettings],
         name: str,
         *,
         host: Optional[str] = None,
@@ -761,16 +767,13 @@ class experiment:
 
         from experimaestro.server import Server
 
-        if isinstance(env, Environment):
-            self.environment = env
-        else:
-            self.environment = Environment(workdir=env)
+        settings = get_settings()
+        if not isinstance(env, WorkspaceSettings):
+            env = WorkspaceSettings(id=None, path=Path(env))
 
         # Creates the workspace
         run_mode = run_mode or RunMode.NORMAL
-        self.workspace = Workspace(
-            self.environment, launcher=launcher, run_mode=run_mode
-        )
+        self.workspace = Workspace(settings, env, launcher=launcher, run_mode=run_mode)
 
         # Mark the directory has an experimaestro folder
         self.workdir = self.workspace.experimentspath / name
@@ -780,7 +783,7 @@ class experiment:
         self.old_experiment = None
         self.services: Dict[str, Service] = {}
 
-        settings = get_settings()
+        # Get configuration settings
 
         if host is not None:
             settings.server.host = host
@@ -832,6 +835,12 @@ class experiment:
         """Return the directory in which results can be stored for this experiment"""
         return self.workdir / "jobs"
 
+    @property
+    def alt_jobspaths(self):
+        """Return potential other directories"""
+        for alt_workdir in self.workspace.alt_workdirs:
+            yield alt_workdir / "jobs"
+
     @property
     def jobsbakpath(self):
         """Return the directory in which results can be stored for this experiment"""
@@ -876,9 +885,11 @@ class experiment:
         future = asyncio.run_coroutine_threadsafe(awaitcompletion(), self.loop)
         return future.result()
 
-    def setenv(self, name, value):
+    def setenv(self, name, value, override=True):
         """Shortcut to set the environment value"""
-        self.environment.setenv(name, value)
+        if override or name not in self.workspace.env:
+            logging.info("Setting environment: %s=%s", name, value)
+            self.workspace.env[name] = value
 
     def token(self, name: str, count: int):
         """Returns a token for this experiment

experimaestro/scheduler/workspace.py

@@ -1,9 +1,9 @@
+from collections import ChainMap
 from enum import Enum
+from functools import cached_property
 from pathlib import Path
-from typing import Optional, TYPE_CHECKING
-
-if TYPE_CHECKING:
-    from experimaestro.scheduler.environment import Environment
+from typing import Optional
+from experimaestro.settings import WorkspaceSettings, Settings
 
 
 class RunMode(str, Enum):
@@ -25,15 +25,22 @@ class Workspace:
     """
 
     CURRENT = None
-    environment: "Environment"
+    settings: "Settings"
+    worspace: "WorkspaceSettings"
 
     """Creates a workspace for experiments"""
 
     def __init__(
-        self, environment: "Environment", launcher=None, run_mode: RunMode = None
+        self,
+        settings: "Settings",
+        workspace_settings: "WorkspaceSettings",
+        launcher=None,
+        run_mode: RunMode = None,
     ):
-        self.environment = environment
-        path = environment.workdir
+        self.settings = settings
+        self.workspace_settings = workspace_settings
+
+        path = self.workspace_settings.path
         self.notificationURL: Optional[str] = None
         if isinstance(path, Path):
             path = path.absolute()
@@ -43,6 +50,8 @@ class Workspace:
 
         self.launcher = launcher or Launcher.get(path)
 
+        self.env = ChainMap({}, workspace_settings.env, settings.env)
+
     def __enter__(self):
         self.old_workspace = Workspace.CURRENT
         Workspace.CURRENT = self
@@ -50,6 +59,15 @@ class Workspace:
     def __exit__(self, *args):
         Workspace.CURRENT = self.old_workspace
 
+    @cached_property
+    def alt_workspaces(self):
+        for ws_id in self.workspace_settings.alt_workspaces:
+            yield self.settings.workspaces[ws_id]
+
+    @property
+    def alt_workdirs(self):
+        yield from map(lambda ws: ws.path, self.workspace_settings.alt_workspaces)
+
     @property
     def connector(self):
         """Returns the default connector"""

experimaestro/scriptbuilder.py

@@ -92,7 +92,11 @@ class PythonScriptBuilder:
         with scriptpath.open("wt") as out:
             out.write("#!{}\n".format(self.pythonpath))
             out.write("# Experimaestro generated task\n\n")
-            out.write("""import logging\nlogging.basicConfig(level=logging.INFO)\n\n""")
+            out.write(
+                """import logging\n"""
+                """logging.basicConfig(level=logging.INFO, """
+                """format='%(levelname)s:%(process)d:%(asctime)s [%(name)s] %(message)s', datefmt='%y-%m-%d %H:%M:%S')\n\n"""
+            )
 
             out.write("\nif __name__ == '__main__':\n\n" "")
 

experimaestro/settings.py

@@ -1,9 +1,10 @@
 import os
-from omegaconf import OmegaConf
+from omegaconf import OmegaConf, SCMode
 from dataclasses import field, dataclass
 from functools import lru_cache
 from pathlib import Path
 from typing import Dict, Optional, List
+import logging
 
 
 @dataclass
@@ -23,15 +24,23 @@ class ServerSettings:
 
 @dataclass
 class WorkspaceSettings:
+    """Defines the workspace"""
+
     id: str
     """The workspace identifier"""
 
-    path: Path
+    path: Path = field()
     """The workspace path"""
 
     env: Dict[str, str] = field(default_factory=dict)
     """Workspace specific environment variables"""
 
+    alt_workspaces: List[str] = field(default_factory=list)
+    """Alternative workspaces to find jobs or experiments"""
+
+    def __post_init__(self):
+        self.path = self.path.expanduser().resolve()
+
 
 @dataclass
 class Settings:
@@ -51,13 +60,17 @@ def get_settings(path: Optional[Path] = None) -> Settings:
 
         path = path or Path("~/.config/experimaestro/settings.yaml").expanduser()
         if not path.is_file():
-            return schema
+            return OmegaConf.to_container(
+                schema, structured_config_mode=SCMode.INSTANTIATE
+            )
 
         conf = OmegaConf.load(path)
-        return OmegaConf.merge(schema, conf)
+        return OmegaConf.to_container(
+            OmegaConf.merge(schema, conf), structured_config_mode=SCMode.INSTANTIATE
+        )
 
 
-def get_workspace(id: Optional[str]) -> WorkspaceSettings:
+def get_workspace(id: Optional[str] = None) -> Optional[WorkspaceSettings]:
     """Return the workspace settings given an id (or None for the default one)"""
     workspaces = get_settings().workspaces
     if workspaces:
@@ -68,3 +81,28 @@ def get_workspace(id: Optional[str]) -> 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_ssh.py

@@ -1,14 +1,18 @@
+import sys
+import pytest
 from experimaestro.connectors.ssh import SshPath
 
 # --- Test SSH path and SSH path manipulation
 
 
+@pytest.mark.skipif(sys.version_info >= (3, 12), reason="requires python3.10 or higher")
 def test_absolute():
     path = SshPath("ssh://host//a/path")
     assert path.host == "host"
     assert path.is_absolute()
 
 
+@pytest.mark.skipif(sys.version_info >= (3, 12), reason="requires python3.10 or higher")
 def test_relative():
     path = SshPath("ssh://host")
     assert path.host == "host"
@@ -17,6 +21,7 @@ def test_relative():
     assert not path.is_absolute()
 
 
+@pytest.mark.skipif(sys.version_info >= (3, 12), reason="requires python3.10 or higher")
 def test_relative_withpath():
     path = SshPath("ssh://host/relative/path")
     assert path.host == "host"
@@ -24,6 +29,7 @@ def test_relative_withpath():
     assert not path.is_absolute()
 
 
+@pytest.mark.skipif(sys.version_info >= (3, 12), reason="requires python3.10 or higher")
 def test_relative_absolute():
     path = SshPath("ssh://host") / "/absolute/path"
     assert path.host == "host"
@@ -31,6 +37,7 @@ def test_relative_absolute():
     assert path.is_absolute()
 
 
+@pytest.mark.skipif(sys.version_info >= (3, 12), reason="requires python3.10 or higher")
 def test_relative_compose():
     path = SshPath("ssh://host/abc") / "relative/path"
     assert path.host == "host"

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)

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.4.dist-info → experimaestro-1.6.0.dist-info}/METADATA

@@ -1,8 +1,7 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.3
 Name: experimaestro
-Version: 1.5.4
+Version: 1.6.0
 Summary: "Experimaestro is a computer science experiment manager"
-Home-page: https://github.com/experimaestro/experimaestro-python
 License: GPL-3
 Keywords: experiment manager
 Author: Benjamin Piwowarski
@@ -20,13 +19,13 @@ Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Requires-Dist: arpeggio (>=2,<3)
 Requires-Dist: attrs (>=23.1.0,<24.0.0)
 Requires-Dist: click (>=8)
 Requires-Dist: decorator (>=5,<6)
 Requires-Dist: docstring-parser (>=0.15,<0.16)
-Requires-Dist: fabric (>=3,<4)
 Requires-Dist: fasteners (>=0.19,<0.20)
 Requires-Dist: flask (>=2.3,<3.0)
 Requires-Dist: flask-socketio (>=5.3,<6.0)
@@ -41,7 +40,7 @@ Requires-Dist: pyparsing (>=3.1,<4.0)
 Requires-Dist: pytools (>=2023.1.1,<2024.0.0)
 Requires-Dist: pyyaml (>=6.0.1,<7.0.0)
 Requires-Dist: requests (>=2.31,<3.0)
-Requires-Dist: rpyc (>=5,<6)
+Requires-Dist: rpyc (>=5,<7)
 Requires-Dist: sortedcontainers (>=2.4,<3.0)
 Requires-Dist: termcolor (>=2.3)
 Requires-Dist: tqdm (>=4.66.1,<5.0.0)