experimaestro 1.11.1__py3-none-any.whl → 2.0.0b4__py3-none-any.whl

experimaestro/generators.py

@@ -1,10 +1,13 @@
 import inspect
 from pathlib import Path
 from abc import ABC, abstractmethod
-from typing import Callable, Union
+from typing import Callable, Union, TYPE_CHECKING
 from experimaestro.core.arguments import ArgumentOptions, TypeAnnotation
 from experimaestro.core.objects import ConfigWalkContext, Config
 
+if TYPE_CHECKING:
+    from experimaestro.core.subparameters import Subparameters
+
 
 class Generator(ABC):
     """Base class for all generators"""
@@ -15,25 +18,63 @@ class Generator(ABC):
         return False
 
     @abstractmethod
-    def __call__(self, context: ConfigWalkContext, config: Config):
-        ...
+    def __call__(self, context: ConfigWalkContext, config: Config): ...
 
 
 class PathGenerator(Generator):
-    """Generates a path"""
+    """Generate paths within the task directory.
+
+    Use ``PathGenerator`` with ``field(default_factory=...)`` to create
+    paths relative to the task's working directory.
+
+    Example::
+
+        class MyTask(Task):
+            output: Meta[Path] = field(default_factory=PathGenerator("results.json"))
+            model: Meta[Path] = field(default_factory=PathGenerator("model.pt"))
+
+    For shared directories across related tasks, use with subparameters::
+
+        training_group = param_group("training")
+
+        class Train(Task):
+            epochs: Param[int] = field(groups=[training_group])
+            checkpoint: Meta[Path] = field(
+                default_factory=PathGenerator(
+                    "model.pt",
+                    subparameters=subparameters(exclude=[training_group])
+                )
+            )
+
+    :param path: Relative path within the task directory. Can be a string,
+        Path, or callable that takes (context, config) and returns a Path.
+    :param subparameters: Optional subparameters for partial directory sharing.
+        When provided, the path is generated in a shared partial directory.
+    """
 
     def __init__(
-        self, path: Union[str, Path, Callable[[ConfigWalkContext, Config], Path]]
+        self,
+        path: Union[str, Path, Callable[[ConfigWalkContext, Config], Path]] = "",
+        *,
+        partial: "Subparameters" = None,
     ):
         self.path = path
+        self.partial = partial
 
     def __call__(self, context: ConfigWalkContext, config: Config):
-        if inspect.isfunction(self.path):
-            path = context.currentpath() / self.path(context, config)  # type: Path
+        # Determine base path: partial directory or job directory
+        if self.partial is not None:
+            base_path = context.partial_path(self.partial, config)
         else:
-            path = context.currentpath() / Path(self.path)
+            base_path = context.currentpath()
 
-        return path
+        # Generate the final path
+        if inspect.isfunction(self.path):
+            return base_path / self.path(context, config)
+        elif self.path:
+            return base_path / Path(self.path)
+        else:
+            return base_path
 
     def isoutput(self):
         return True

experimaestro/huggingface.py

@@ -82,4 +82,6 @@ class ExperimaestroHFHub(ModelHubMixin):
                 )
                 return hf_path
 
-        return ConfigInformation.deserialize(data_loader, as_instance=as_instance)
+        return ConfigInformation.deserialize(
+            data_loader, as_instance=as_instance, partial_loading=True
+        )

experimaestro/launcherfinder/parser.py

@@ -105,6 +105,35 @@ class Visitor(PTNodeVisitor):
 
 
 def parse(expr: str):
+    """Parse a requirement specification string into a HostRequirement object.
+
+    The specification string describes hardware requirements for running a task.
+    Multiple alternatives can be specified using ``|`` (OR), and requirements
+    within an alternative are combined using ``&`` (AND).
+
+    **Syntax elements:**
+
+    - ``duration=<N><unit>``: Job duration (units: h/hours, d/days, m/mins)
+    - ``cpu(mem=<size>, cores=<N>)``: CPU requirements
+    - ``cuda(mem=<size>) * <N>``: GPU requirements (memory and count)
+    - Memory sizes: ``<N>G``, ``<N>GiB``, ``<N>M``, ``<N>MiB``
+
+    :param expr: The requirement specification string
+    :return: A :class:`~experimaestro.launcherfinder.specs.HostRequirement` object
+
+    **Example:**
+
+    .. code-block:: python
+
+        from experimaestro.launcherfinder.parser import parse
+
+        # Request 2 GPUs with 32GB each, 700GB RAM, for 40 hours
+        # OR 4 GPUs with 32GB each for 50 hours
+        req = parse(
+            "duration=40h & cpu(mem=700GiB) & cuda(mem=32GiB) * 2"
+            " | duration=50h & cpu(mem=700GiB) & cuda(mem=32GiB) * 4"
+        )
+    """
     parser = ParserPython(grammar, syntax_classes={"StrMatch": SuppressStrMatch})
     parse_tree = parser.parse(expr)
     return visit_parse_tree(parse_tree, Visitor(debug=False))

experimaestro/launcherfinder/registry.py

@@ -6,7 +6,7 @@ from typing import ClassVar, Dict, Optional, Set, Type, Union
 from pathlib import Path
 import typing
 from omegaconf import DictConfig, OmegaConf, SCMode
-import pkg_resources
+from importlib.metadata import entry_points
 from experimaestro.utils import logger
 from .base import ConnectorConfiguration, TokenConfiguration
 from .specs import HostRequirement, RequirementUnion
@@ -75,10 +75,10 @@ class LauncherRegistry:
         self.find_launcher_fn = None
 
         # Use entry points for connectors and launchers
-        for entry_point in pkg_resources.iter_entry_points("experimaestro.connectors"):
+        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

experimaestro/launchers/__init__.py

@@ -31,7 +31,19 @@ SubmitListener = Callable[[Job], None]
 
 
 class Launcher(ABC):
-    """A launcher"""
+    """Base class for task launchers.
+
+    Launchers are responsible for executing tasks on a compute resource.
+    They work with a :class:`~experimaestro.connectors.Connector` to
+    access the target system and manage process execution.
+
+    Subclasses include:
+
+    - :class:`~experimaestro.launchers.direct.DirectLauncher`: Local execution
+    - :class:`~experimaestro.launchers.slurm.SlurmLauncher`: SLURM cluster
+
+    :param connector: The connector to use for accessing the compute resource
+    """
 
     submit_listeners: List[SubmitListener]
 
@@ -69,6 +81,19 @@ class Launcher(ABC):
         By default, returns the associated connector builder"""
         return self.connector.processbuilder()
 
+    @abstractmethod
+    def launcher_info_code(self) -> str:
+        """Returns Python code to set up launcher info during task execution.
+
+        This code is inserted into the generated task script to set up
+        launcher-specific information (like LauncherInformation for
+        querying remaining time).
+
+        Returns:
+            Python code as a string, or empty string if no setup needed.
+        """
+        ...
+
     @staticmethod
     def get(path: Path):
         """Get a default launcher for a given path"""

experimaestro/launchers/direct.py

@@ -3,8 +3,20 @@ from . import Launcher
 
 
 class DirectLauncher(Launcher):
+    """Launcher that runs tasks directly as local processes.
+
+    This is the default launcher that executes tasks on the local machine
+    without any job scheduler. Tasks are run as Python subprocesses.
+
+    :param connector: The connector to use (defaults to LocalConnector)
+    """
+
     def scriptbuilder(self):
         return PythonScriptBuilder()
 
+    def launcher_info_code(self) -> str:
+        """Returns empty string as local launcher has no time limits."""
+        return ""
+
     def __str__(self):
         return f"DirectLauncher({self.connector})"

experimaestro/launchers/slurm/base.py

@@ -7,6 +7,7 @@ from typing import (
     List,
     Optional,
     Tuple,
+    TYPE_CHECKING,
     get_type_hints,
 )
 from experimaestro.connectors.local import LocalConnector
@@ -20,7 +21,7 @@ from experimaestro.launcherfinder.registry import (
 from experimaestro.utils import ThreadingCondition
 from experimaestro.tests.connectors.utils import OutputCaptureHandler
 from experimaestro.utils.asyncio import asyncThreadcheck
-from experimaestro.compat import cached_property
+from functools import cached_property
 from experimaestro.launchers import Launcher
 from experimaestro.scriptbuilder import PythonScriptBuilder
 from experimaestro.connectors import (
@@ -32,8 +33,131 @@ from experimaestro.connectors import (
     RedirectType,
 )
 
+if TYPE_CHECKING:
+    from experimaestro.scheduler.jobs import JobState
+
 logger = logging.getLogger("xpm.slurm")
 
+# Cached job end time (absolute timestamp).
+# Only used when a task is running within a SLURM job.
+_slurm_job_end_time: Optional[float] = None
+
+
+class SlurmLauncherInformation:
+    """Launcher information for SLURM jobs, used during task execution."""
+
+    def __init__(self, binpath: str = "/usr/bin"):
+        self.binpath = Path(binpath)
+
+    def remaining_time(self) -> Optional[float]:
+        """Returns the remaining time in seconds before the SLURM job times out.
+
+        Uses the SLURM_JOB_ID environment variable to query squeue for the
+        remaining time. The job end time is cached on first call.
+
+        Returns:
+            The remaining time in seconds, or None if no time limit.
+        """
+        import os
+        import time
+
+        global _slurm_job_end_time
+
+        # Use cached end time if available
+        if _slurm_job_end_time is not None:
+            remaining = _slurm_job_end_time - time.time()
+            return max(0.0, remaining)
+
+        # Query SLURM for remaining time and compute end time
+        job_id = os.environ.get("SLURM_JOB_ID")
+        if not job_id:
+            logger.debug("No SLURM_JOB_ID in environment, cannot get remaining time")
+            return None
+
+        remaining_seconds = self._query_remaining_time(job_id)
+        if remaining_seconds is None:
+            return None
+
+        # Cache the absolute end time
+        _slurm_job_end_time = time.time() + remaining_seconds
+        return remaining_seconds
+
+    def _query_remaining_time(self, job_id: str) -> Optional[float]:
+        """Query SLURM for remaining time of a job."""
+        import subprocess
+
+        try:
+            result = subprocess.run(
+                [
+                    f"{self.binpath}/squeue",
+                    "--job",
+                    job_id,
+                    "--format=%L",
+                    "--noheader",
+                ],
+                capture_output=True,
+                text=True,
+                timeout=30,
+            )
+
+            if result.returncode != 0:
+                logger.warning(
+                    "squeue returned error code %d: %s",
+                    result.returncode,
+                    result.stderr,
+                )
+                return None
+
+            time_str = result.stdout.strip()
+            if not time_str or time_str == "UNLIMITED":
+                return None
+
+            return self._parse_slurm_time(time_str)
+        except subprocess.TimeoutExpired:
+            logger.warning("Timeout querying squeue for remaining time")
+            return None
+        except Exception as e:
+            logger.warning("Error querying SLURM remaining time: %s", e)
+            return None
+
+    @staticmethod
+    def _parse_slurm_time(time_str: str) -> Optional[float]:
+        """Parse SLURM time format to seconds.
+
+        SLURM time format can be:
+        - D-HH:MM:SS (days-hours:minutes:seconds)
+        - HH:MM:SS (hours:minutes:seconds)
+        - MM:SS (minutes:seconds)
+        - SS (seconds)
+
+        Returns:
+            Time in seconds, or None if parsing fails
+        """
+        try:
+            days = 0
+            if "-" in time_str:
+                days_str, time_str = time_str.split("-", 1)
+                days = int(days_str)
+
+            parts = time_str.split(":")
+            if len(parts) == 3:
+                hours, minutes, seconds = int(parts[0]), int(parts[1]), int(parts[2])
+            elif len(parts) == 2:
+                hours = 0
+                minutes, seconds = int(parts[0]), int(parts[1])
+            elif len(parts) == 1:
+                hours = 0
+                minutes = 0
+                seconds = int(parts[0])
+            else:
+                logger.warning("Could not parse SLURM time: %s", time_str)
+                return None
+
+            return float(days * 86400 + hours * 3600 + minutes * 60 + seconds)
+        except (ValueError, IndexError) as e:
+            logger.warning("Could not parse SLURM time '%s': %s", time_str, e)
+            return None
+
 
 class SlurmJobState:
     start: str
@@ -176,14 +300,34 @@ class BatchSlurmProcess(Process):
     def __init__(self, launcher: "SlurmLauncher", jobid: str):
         self.launcher = launcher
         self.jobid = jobid
+        self._last_state: Optional[SlurmJobState] = None
 
     def wait(self):
         with SlurmProcessWatcher.get(self.launcher) as watcher:
             while True:
                 state = watcher.getjob(self.jobid)
                 if state and state.finished():
+                    self._last_state = state
                     return 0 if state.slurm_state == "COMPLETED" else 1
 
+    def get_job_state(self, code: int) -> "JobState":
+        """Convert SLURM exit code to JobState, detecting timeouts"""
+        from experimaestro.scheduler.jobs import (
+            JobState,
+            JobStateError,
+            JobFailureStatus,
+        )
+
+        if code == 0:
+            return JobState.DONE
+
+        # Check if this was a SLURM timeout
+        if self._last_state and self._last_state.slurm_state == "TIMEOUT":
+            logger.info("SLURM job %s timed out", self.jobid)
+            return JobStateError(JobFailureStatus.TIMEOUT)
+
+        return JobState.ERROR
+
     async def aio_state(self, timeout: float | None = None) -> ProcessState:
         def check():
             with SlurmProcessWatcher.get(self.launcher) as watcher:
@@ -432,7 +576,7 @@ class SlurmLauncher(Launcher):
     def scriptbuilder(self):
         """Returns the script builder
 
-        We assume *nix, but should be changed to PythonScriptBuilder when working
+        We assume Unix, but should be changed to PythonScriptBuilder when working
         """
         return SlurmScriptBuilder(self)
 
@@ -442,6 +586,14 @@ class SlurmLauncher(Launcher):
         By default, returns the associated connector builder"""
         return SlurmProcessBuilder(self)
 
+    def launcher_info_code(self) -> str:
+        """Returns Python code to set up launcher info during task execution."""
+        return (
+            "    from experimaestro.launchers.slurm import SlurmLauncherInformation\n"
+            "    from experimaestro import taskglobals\n"
+            f'    taskglobals.Env.instance().launcher_info = SlurmLauncherInformation(binpath="{self.binpath}")\n'
+        )
+
 
 class SlurmScriptBuilder(PythonScriptBuilder):
     def __init__(self, launcher: SlurmLauncher, pythonpath=None):

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()

experimaestro/mkdocs/metaloader.py

@@ -2,7 +2,6 @@
 # when building documentation
 
 import sys
-import re
 import importlib.abc
 import importlib.machinery