experimaestro 1.5.0__tar.gz → 1.5.1__tar.gz

{experimaestro-1.5.0 → experimaestro-1.5.1}/PKG-INFO

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

{experimaestro-1.5.0 → experimaestro-1.5.1}/pyproject.toml

@@ -19,7 +19,7 @@ include = [
     "src/experimaestro/sphinx/static/experimaestro.css",
     "src/experimaestro/mkdocs/style.css"
 ]
-version = "1.5.0"
+version = "1.5.1"
 repository = "https://github.com/experimaestro/experimaestro-python"
 documentation = "https://experimaestro-python.readthedocs.io/"
 
@@ -123,7 +123,7 @@ warn_unused_ignores = true
 
 [tool.commitizen]
 name = "cz_conventional_commits"
-version = "1.5.0"
+version = "1.5.1"
 changelog_start_rev = "0.15.0"
 tag_format = "v$version"
 update_changelog_on_bump = true

{experimaestro-1.5.0 → experimaestro-1.5.1}/src/experimaestro/experiments/cli.py

@@ -1,4 +1,5 @@
 import inspect
+import itertools
 import json
 import logging
 import sys
@@ -11,7 +12,7 @@ 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 experimaestro.settings import get_settings, get_workspace
 from omegaconf import OmegaConf, SCMode
 from termcolor import cprint
 
@@ -219,21 +220,30 @@ def experiments_cli(  # noqa: C901
         sys.exit(0)
 
     # Move to an object container
-    configuration: schema = OmegaConf.to_container(
+    configuration = OmegaConf.to_container(
         configuration, structured_config_mode=SCMode.INSTANTIATE
     )
 
     # Get the working directory
-    if workdir is None or not Path(workdir).is_dir():
-        workdir = get_workspace(workdir).path.expanduser().resolve()
-        logging.info("Using working directory %s", workdir)
+    settings = get_settings()
+    ws_env = {}
+    workdir = Path(workdir) if workdir else None
+    if (workdir is None) or (not workdir.is_dir()):
+        logging.info("Searching for workspace %s", workdir)
+        ws_settings = get_workspace(str(workdir))
+        workdir = ws_settings.path.expanduser()
+        ws_env = ws_settings.env
+
+    logging.info("Using working directory %s", str(workdir.resolve()))
 
     # --- Runs the experiment
     with experiment(
         workdir, configuration.id, host=host, port=port, run_mode=run_mode
     ) as xp:
         # Set up the environment
-        for key, value in env:
+        # (1) global settings (2) workspace settings and (3) command line settings
+        for key, value in itertools.chain(settings.env.items(), ws_env.items(), env):
+            logging.info("Setting environment: %s=%s", key, value)
             xp.setenv(key, value)
 
         try:

{experimaestro-1.5.0 → experimaestro-1.5.1}/src/experimaestro/launcherfinder/parser.py

@@ -54,8 +54,12 @@ def duration():
     return "duration", "=", RegExMatch(r"\d+"), RegExMatch(r"h(ours)?|d(ays)?")
 
 
+def one_spec():
+    return OneOrMore(OrderedChoice([duration, cuda, cpu]), sep="&")
+
+
 def grammar():
-    return OneOrMore(OrderedChoice([duration, cuda, cpu]), sep="&"), EndOfFile()
+    return OneOrMore(one_spec, sep="|"), EndOfFile()
 
 
 # ---- Visitor
@@ -63,6 +67,9 @@ def grammar():
 
 class Visitor(PTNodeVisitor):
     def visit_grammar(self, node, children):
+        return [child for child in children]
+
+    def visit_one_spec(self, node, children):
         return reduce(lambda x, el: x & el, children)
 
     def visit_duration(self, node, children):

{experimaestro-1.5.0 → experimaestro-1.5.1}/src/experimaestro/launcherfinder/registry.py

@@ -219,7 +219,7 @@ class LauncherRegistry:
         raise AssertionError(f"No connector with identifier {identifier}")
 
     def find(
-        self, *specs: Union[HostRequirement, str], tags: Set[str] = set()
+        self, *input_specs: Union[HostRequirement, str], tags: Set[str] = set()
     ) -> Optional["Launcher"]:
         """ "
         Arguments:
@@ -237,7 +237,12 @@ class LauncherRegistry:
         # Parse specs
         from .parser import parse
 
-        specs = [parse(spec) if isinstance(spec, str) else spec for spec in specs]
+        specs = []
+        for spec in input_specs:
+            if isinstance(spec, str):
+                specs.extend(parse(spec))
+            else:
+                specs.append(spec)
 
         # Use launcher function
         if self.find_launcher_fn is not None:
@@ -247,8 +252,6 @@ class LauncherRegistry:
 
         # We have registered launchers
         for spec in specs:
-            if isinstance(spec, str):
-                spec = parse(spec)
             for handler in self.launchers:
                 if (not tags) or any((tag in tags) for tag in handler.tags):
                     if launcher := handler.get(self, spec):

{experimaestro-1.5.0 → experimaestro-1.5.1}/src/experimaestro/launchers/slurm/base.py

@@ -92,6 +92,8 @@ class SlurmProcessWatcher(threading.Thread):
         self.jobs: Dict[str, SlurmJobState] = {}
 
         self.cv = ThreadingCondition()
+        self.fetched_event = threading.Event()
+        self.updating_jobs = threading.Lock()
         self.start()
 
     @staticmethod
@@ -109,10 +111,18 @@ class SlurmProcessWatcher(threading.Thread):
         with watcher.cv:
             watcher.cv.notify()
 
-    def getjob(self, jobid):
+    def getjob(self, jobid, timeout=None):
         """Allows to share the calls to sacct"""
+
+        # Ensures that we have fetched at least once
+        self.fetched_event.wait()
+
+        # Waits that jobs are refreshed (with a timeout)
         with self.cv:
-            self.cv.wait()
+            self.cv.wait(timeout=timeout)
+
+        # Ensures jobs are not updated right now
+        with self.updating_jobs:
             return self.jobs.get(jobid)
 
     def run(self):
@@ -129,9 +139,9 @@ class SlurmProcessWatcher(threading.Thread):
             builder.stdout = Redirect.pipe(handler)
             builder.environ = self.launcher.launcherenv
             logger.debug("Checking SLURM state with sacct")
-            builder.start()
+            process = builder.start()
 
-            with self.cv:
+            with self.updating_jobs:
                 self.jobs = {}
                 output = handler.output.decode("utf-8")
                 for line in output.split("\n"):
@@ -143,7 +153,11 @@ class SlurmProcessWatcher(threading.Thread):
                             logger.debug("Parsed line: %s", line)
                         except ValueError:
                             logger.error("Could not parse line %s", line)
+            process.kill()
+
+            with self.cv:
                 logger.debug("Jobs %s", self.jobs)
+                self.fetched_event.set()
                 self.cv.notify_all()
 
                 self.cv.wait_for(
@@ -193,7 +207,18 @@ class BatchSlurmProcess(Process):
     def fromspec(cls, connector: Connector, spec: Dict[str, Any]):
         options = {k: v for k, v in spec.get("options", ())}
         launcher = SlurmLauncher(connector=connector, **options)
-        return BatchSlurmProcess(launcher, spec["pid"])
+        process = BatchSlurmProcess(launcher, spec["pid"])
+
+        # Checks that the process is running
+        with SlurmProcessWatcher.get(launcher) as watcher:
+            logger.info("Checking SLURM job %s", process.jobid)
+            jobinfo = watcher.getjob(process.jobid, timeout=0.1)
+            if jobinfo and jobinfo.state.running:
+                logger.debug(
+                    "SLURM job is running (%s), returning process", process.jobid
+                )
+                return process
+        return None
 
 
 def addstream(command: List[str], option: str, redirect: Redirect):

{experimaestro-1.5.0 → experimaestro-1.5.1}/src/experimaestro/scriptbuilder.py

@@ -91,24 +91,28 @@ class PythonScriptBuilder:
         logger.debug("Writing script %s", scriptpath)
         with scriptpath.open("wt") as out:
             out.write("#!{}\n".format(self.pythonpath))
-            out.write("# Experimaestro generated task\n")
+            out.write("# Experimaestro generated task\n\n")
+            out.write("""import logging\nlogging.basicConfig(level=logging.INFO)\n\n""")
+
+            out.write("\nif __name__ == '__main__':\n\n" "")
 
             # --- Checks locks right away
 
-            out.write("""import logging\nlogging.basicConfig(level=logging.INFO)\n\n""")
-            out.write("""from experimaestro.run import TaskRunner\nimport os\n\n""")
+            out.write(
+                """    from experimaestro.run import TaskRunner\n    import os\n\n"""
+            )
 
-            out.write("lockfiles = [\n")
+            out.write("    lockfiles = [\n")
             for path in self.lockfiles:
-                out.write(f"   '''{relpath(path)}''',\n")
-            out.write("]\n")
+                out.write(f"       '''{relpath(path)}''',\n")
+            out.write("    ]\n")
 
             for name, value in job.environ.items():
-                out.write(f"""os.environ["{name}"] = "{shquote(value)}"\n""")
+                out.write(f"""    os.environ["{name}"] = "{shquote(value)}"\n""")
             out.write("\n")
 
             out.write(
-                f"""TaskRunner("{shquote(connector.resolve(scriptpath))}","""
+                f"""    TaskRunner("{shquote(connector.resolve(scriptpath))}","""
                 """ lockfiles).run()\n"""
             )
 

{experimaestro-1.5.0 → experimaestro-1.5.1}/src/experimaestro/settings.py

@@ -1,9 +1,9 @@
 import os
 from omegaconf import OmegaConf
-from dataclasses import dataclass, field
+from dataclasses import field, dataclass
 from functools import lru_cache
 from pathlib import Path
-from typing import Optional, List
+from typing import Dict, Optional, List
 
 
 @dataclass
@@ -29,12 +29,18 @@ class WorkspaceSettings:
     path: Path
     """The workspace path"""
 
+    env: Dict[str, str] = field(default_factory=dict)
+    """Workspace specific environment variables"""
+
 
 @dataclass
 class Settings:
     server: ServerSettings = field(default_factory=ServerSettings)
     workspaces: List[WorkspaceSettings] = field(default_factory=list)
 
+    env: Dict[str, str] = field(default_factory=dict)
+    """Default environment variables"""
+
 
 @lru_cache()
 def get_settings(path: Optional[Path] = None) -> Settings:

{experimaestro-1.5.0 → experimaestro-1.5.1}/src/experimaestro/tests/test_findlauncher.py

@@ -60,7 +60,7 @@ def test_findlauncher_specs_gpu_mem():
 
 
 def test_findlauncher_parse():
-    r = parse("""duration=4 d & cuda(mem=4G) * 2 & cpu(mem=400M, cores=4)""")
+    (r,) = parse("""duration=4 d & cuda(mem=4G) * 2 & cpu(mem=400M, cores=4)""")
     assert isinstance(r, HostSimpleRequirement)
 
     assert len(r.cuda_gpus) == 2