experimaestro 2.0.0a8__py3-none-any.whl → 2.0.0b4__py3-none-any.whl

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/metaloader.py

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

experimaestro/mypy.py

@@ -1,15 +1,460 @@
-from mypy.plugin import Plugin
+"""Mypy plugin for experimaestro.
 
+This plugin provides type hints support for experimaestro's Config system,
+particularly for the Config.C pattern and proper parameter type inference.
 
-class ExperimaestroPlugin(Plugin):
-    """Just do nothing for now"""
+The plugin handles:
+- Config.C, Config.XPMConfig, Config.XPMValue class properties
+- Adding __init__ with proper Param field signatures
+- Adding ConfigMixin to the class hierarchy for method access
+- Handling task_outputs return type for submit()
 
-    def get_class_decorator_hook(self, tada):
-        pass
+Usage in mypy.ini or pyproject.toml:
+    [mypy]
+    plugins = experimaestro.mypy
+
+Or in pyproject.toml:
+    [tool.mypy]
+    plugins = ["experimaestro.mypy"]
+"""
+
+from __future__ import annotations
+
+from typing import Callable, List, Optional
+
+from mypy.nodes import (
+    TypeInfo,
+    Var,
+    Argument,
+    ARG_NAMED_OPT,
+    ARG_NAMED,
+)
+from mypy.plugin import Plugin, ClassDefContext
+from mypy.plugins.common import add_attribute_to_class, add_method_to_class
+from mypy.types import (
+    Instance,
+    TypeType,
+    NoneType,
+)
+from mypy.mro import calculate_mro, MroError
+
+# Full names of Config and its subclasses that need C/XPMConfig attributes
+CONFIG_FULLNAMES = {
+    "experimaestro.core.objects.config.Config",
+    "experimaestro.core.objects.config.LightweightTask",
+    "experimaestro.core.objects.config.Task",
+    "experimaestro.core.objects.config.ResumableTask",
+    "experimaestro.Config",
+    "experimaestro.Task",
+    "experimaestro.LightweightTask",
+    "experimaestro.ResumableTask",
+    "experimaestro.core.objects.Config",
+    "experimaestro.core.objects.Task",
+    "experimaestro.core.objects.LightweightTask",
+    "experimaestro.core.objects.ResumableTask",
+}
+
+# ConfigMixin full name for method inheritance
+CONFIGMIXIN_FULLNAME = "experimaestro.core.objects.config.ConfigMixin"
+
+# Full names for Param annotations (required by default)
+PARAM_FULLNAMES = {
+    "experimaestro.core.arguments.Param",
+    "experimaestro.Param",
+}
+
+# Full names for Meta/Option annotations (always optional, ignored in identifier)
+META_FULLNAMES = {
+    "experimaestro.core.arguments.Meta",
+    "experimaestro.Meta",
+    "experimaestro.core.arguments.Option",
+    "experimaestro.Option",
+}
+
+# Full names for Constant annotations (excluded from __init__)
+CONSTANT_FULLNAMES = {
+    "experimaestro.core.arguments.Constant",
+    "experimaestro.Constant",
+}
+
+
+def is_config_subclass(info: TypeInfo) -> bool:
+    """Check if a TypeInfo represents a Config subclass.
+
+    Args:
+        info: The TypeInfo to check
+
+    Returns:
+        True if the type is Config or a subclass of Config
+    """
+    if info.fullname in CONFIG_FULLNAMES:
+        return True
+    for base in info.mro:
+        if base.fullname in CONFIG_FULLNAMES:
+            return True
+    return False
+
+
+# Fields to skip when building __init__ signature
+SKIP_FIELDS = {
+    "C",
+    "XPMConfig",
+    "XPMValue",
+    "__xpm__",
+    "__xpmtype__",
+    "__xpmid__",
+    "_deprecated_from",
+}
+
+
+def _is_config_class(base: TypeInfo) -> bool:
+    """Check if a TypeInfo is a Config subclass.
+
+    Returns True for user-defined Config subclasses.
+    """
+    for mro_base in base.mro:
+        if mro_base.fullname in CONFIG_FULLNAMES:
+            return True
+    return False
+
+
+def _get_annotation_type_str(name: str, base: TypeInfo) -> Optional[str]:
+    """Get the type annotation string for a field.
+
+    Tries multiple sources to find the original annotation:
+    1. The AST unanalyzed_type (preserves original)
+    2. The variable's type string
+    """
+    # Check the AST first to get unanalyzed types
+    if base.defn is not None:
+        for stmt in base.defn.defs.body:
+            from mypy.nodes import AssignmentStmt
+
+            if isinstance(stmt, AssignmentStmt):
+                for lvalue in stmt.lvalues:
+                    from mypy.nodes import NameExpr
+
+                    if isinstance(lvalue, NameExpr) and lvalue.name == name:
+                        # Try unanalyzed_type first (preserves the original annotation)
+                        if stmt.unanalyzed_type is not None:
+                            return str(stmt.unanalyzed_type)
+                        # Fall back to analyzed type
+                        if stmt.type is not None:
+                            return str(stmt.type)
+
+    # Fall back to checking the symbol's type
+    if name in base.names:
+        sym = base.names[name]
+        if sym.node is not None and isinstance(sym.node, Var):
+            var = sym.node
+            if var.type is not None:
+                return str(var.type)
+
+    return None
+
+
+def _is_constant_field(name: str, base: TypeInfo) -> bool:
+    """Check if a field is declared as Constant[T].
+
+    Constant fields should be excluded from __init__.
+    """
+    type_str = _get_annotation_type_str(name, base)
+    if type_str is None:
+        return False
+
+    # Normalize type string - remove optional markers (?)
+    # mypy represents types like "Constant?[str?]"
+    type_lower = type_str.lower().replace("?", "")
+
+    # Check for Constant annotation in the type string
+    if "constant[" in type_lower:
+        return True
+    for fullname in CONSTANT_FULLNAMES:
+        if fullname.lower() in type_lower:
+            return True
+    return False
+
+
+def _is_meta_field(name: str, base: TypeInfo) -> bool:
+    """Check if a field is declared as Meta[T] or Option[T].
+
+    Meta fields should always be optional in __init__.
+    """
+    type_str = _get_annotation_type_str(name, base)
+    if type_str is None:
+        return False
+
+    # Normalize type string - remove optional markers (?)
+    # mypy represents types like "Meta?[Path?]"
+    type_lower = type_str.lower().replace("?", "")
+
+    # Check for Meta/Option annotation in the type string
+    if "meta[" in type_lower or "option[" in type_lower:
+        return True
+    for fullname in META_FULLNAMES:
+        if fullname.lower() in type_lower:
+            return True
+    return False
+
+
+def _get_param_fields(info: TypeInfo) -> List[tuple]:
+    """Extract Param and Meta fields from a class and its bases.
+
+    Returns list of (name, type, has_default) tuples.
+
+    Only includes fields from Config subclasses to avoid picking up
+    attributes from other base classes like nn.Module.
+    Excludes Constant fields which should not be in __init__.
+    """
+    fields = []
+    seen = set()
+
+    # Walk MRO to get inherited fields (in reverse to get proper order)
+    for base in reversed(info.mro):
+        if base.fullname == "builtins.object":
+            continue
+        if base.fullname in CONFIG_FULLNAMES:
+            # Skip Config/Task base classes - we only want user-defined fields
+            continue
+        if base.fullname == CONFIGMIXIN_FULLNAME:
+            # Skip ConfigMixin - it has methods, not params
+            continue
+
+        # Only include fields from Config subclasses
+        # This skips bases like nn.Module that don't inherit from Config
+        if not _is_config_class(base):
+            continue
+
+        for name, sym in base.names.items():
+            if name in seen or name in SKIP_FIELDS:
+                continue
+            if sym.node is None or not isinstance(sym.node, Var):
+                continue
+
+            var = sym.node
+            if var.type is None:
+                continue
+
+            # Skip private/dunder fields
+            if name.startswith("_"):
+                continue
 
-    def get_customize_class_mro_hook(self, tada):
+            # Skip Constant fields - they should not be in __init__
+            if _is_constant_field(name, base):
+                continue
+
+            # Meta fields are always optional
+            # Param fields are optional only if they have a default
+            is_meta = _is_meta_field(name, base)
+            has_default = var.has_explicit_value or is_meta
+
+            seen.add(name)
+            fields.append((name, var.type, has_default))
+
+    return fields
+
+
+def _add_init_method(ctx: ClassDefContext) -> None:
+    """Add an __init__ method with proper Param field signatures."""
+    info = ctx.cls.info
+
+    # Get all Param fields from this class and bases
+    fields = _get_param_fields(info)
+
+    # Build __init__ arguments
+    args = []
+    for name, field_type, has_default in fields:
+        # All experimaestro params are keyword-only
+        # Fields with defaults are optional
+        kind = ARG_NAMED_OPT if has_default else ARG_NAMED
+
+        # Create argument
+        arg = Argument(
+            variable=Var(name, field_type),
+            type_annotation=field_type,
+            initializer=None,
+            kind=kind,
+        )
+        args.append(arg)
+
+    # Add the __init__ method if we have any args
+    if args:
+        add_method_to_class(
+            ctx.api,
+            ctx.cls,
+            "__init__",
+            args,
+            NoneType(),
+        )
+
+
+def _get_task_outputs_return_type(info: TypeInfo) -> Optional[Instance]:
+    """Check if the class has a task_outputs method and return its return type.
+
+    If the class defines task_outputs, submit() should return that type instead
+    of Self.
+    """
+    # Look for task_outputs method in the class
+    if "task_outputs" in info.names:
+        sym = info.names["task_outputs"]
+        if sym.node is not None:
+            # Try to get the return type from the method signature
+            from mypy.nodes import FuncDef
+
+            if isinstance(sym.node, FuncDef):
+                ret_type = sym.node.type
+                if ret_type is not None:
+                    from mypy.types import CallableType
+
+                    if isinstance(ret_type, CallableType):
+                        return ret_type.ret_type
+    return None
+
+
+def _add_configmixin_to_bases(ctx: ClassDefContext) -> None:
+    """Add ConfigMixin to the class bases if not already present.
+
+    This allows mypy to see all ConfigMixin methods on Config subclasses.
+    """
+    info = ctx.cls.info
+
+    # Check if ConfigMixin is already in the MRO
+    for base in info.mro:
+        if base.fullname == CONFIGMIXIN_FULLNAME:
+            return  # Already has ConfigMixin
+
+    # Try to look up ConfigMixin
+    try:
+        configmixin_sym = ctx.api.lookup_fully_qualified_or_none(CONFIGMIXIN_FULLNAME)
+        if configmixin_sym is None or not isinstance(configmixin_sym.node, TypeInfo):
+            return
+
+        configmixin_info = configmixin_sym.node
+        configmixin_instance = Instance(configmixin_info, [])
+
+        # Add ConfigMixin to bases if not already present
+        configmixin_in_bases = any(
+            isinstance(b, Instance) and b.type.fullname == CONFIGMIXIN_FULLNAME
+            for b in info.bases
+        )
+        if not configmixin_in_bases:
+            info.bases.append(configmixin_instance)
+
+            # Recalculate MRO
+            try:
+                calculate_mro(info)
+            except MroError:
+                # If MRO calculation fails, remove the base we added
+                info.bases.pop()
+    except Exception:
+        # If lookup fails, continue without adding ConfigMixin
         pass
 
 
-def plugin(version: str):
+def _add_submit_method(ctx: ClassDefContext) -> None:
+    """Add submit() method that returns Self (or task_outputs return type).
+
+    The actual submit() signature from ConfigMixin:
+        def submit(self, *, workspace=None, launcher=None, run_mode=None,
+                   init_tasks=[], max_retries=None)
+    """
+    info = ctx.cls.info
+
+    # Check if the class has task_outputs
+    task_outputs_type = _get_task_outputs_return_type(info)
+
+    # submit() returns task_outputs return type if defined, otherwise Self
+    if task_outputs_type is not None:
+        return_type = task_outputs_type
+    else:
+        return_type = Instance(info, [])
+
+    # Build submit() arguments - all optional kwargs
+    from mypy.types import AnyType, TypeOfAny
+
+    any_type = AnyType(TypeOfAny.explicit)
+    submit_args = []
+    for arg_name in ("workspace", "launcher", "run_mode", "init_tasks", "max_retries"):
+        arg = Argument(
+            variable=Var(arg_name, any_type),
+            type_annotation=any_type,
+            initializer=None,
+            kind=ARG_NAMED_OPT,
+        )
+        submit_args.append(arg)
+
+    add_method_to_class(
+        ctx.api,
+        ctx.cls,
+        "submit",
+        submit_args,
+        return_type,
+    )
+
+
+def _process_config_class(ctx: ClassDefContext) -> None:
+    """Process a Config subclass to add type hints.
+
+    This adds:
+    - ConfigMixin to the class hierarchy for method access
+    - C, XPMConfig, XPMValue as class attributes returning Type[Self]
+    - An __init__ method with proper Param field signatures
+    - A submit() method that returns Self (or task_outputs return type)
+    """
+    info = ctx.cls.info
+
+    # Add ConfigMixin to bases for method access (tag, instance, etc.)
+    _add_configmixin_to_bases(ctx)
+
+    # Create Type[Self] for this class
+    class_type = Instance(info, [])
+    type_type = TypeType(class_type)
+
+    # Add C, XPMConfig, XPMValue as class attributes returning the class type
+    for attr_name in ("C", "XPMConfig", "XPMValue"):
+        if attr_name not in info.names:
+            add_attribute_to_class(
+                ctx.api,
+                ctx.cls,
+                attr_name,
+                type_type,
+            )
+
+    # Add __init__ with proper field signatures
+    _add_init_method(ctx)
+
+    # Add submit() method that returns Self (or task_outputs type)
+    _add_submit_method(ctx)
+
+
+class ExperimaestroPlugin(Plugin):
+    """Mypy plugin for experimaestro type hints.
+
+    This plugin handles:
+    - Converting @classproperty decorated methods to proper class attributes
+    - Type inference for Config.C and Config.XPMConfig patterns
+    - Adding __init__ methods with proper Param field signatures
+    """
+
+    def get_base_class_hook(
+        self, fullname: str
+    ) -> Callable[[ClassDefContext], None] | None:
+        """Hook called when a class inherits from Config.
+
+        This allows us to process classproperty attributes and add __init__.
+        """
+        if fullname in CONFIG_FULLNAMES:
+            return _process_config_class
+        return None
+
+
+def plugin(_version: str):
+    """Entry point for mypy plugin.
+
+    Args:
+        _version: The mypy version string (unused but required by mypy API)
+
+    Returns:
+        The ExperimaestroPlugin class
+    """
     return ExperimaestroPlugin