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

experimaestro/core/types.py

@@ -1,7 +1,18 @@
 from abc import ABC, abstractmethod
+from dataclasses import dataclass
 import inspect
 import sys
-from typing import Set, TypeVar, Union, Dict, Iterator, List, get_args, get_origin
+from typing import (
+    Set,
+    TypeVar,
+    Union,
+    Dict,
+    Iterator,
+    List,
+    Optional,
+    get_args,
+    get_origin,
+)
 from collections import ChainMap
 from pathlib import Path
 import typing
@@ -13,15 +24,27 @@ from enum import Enum
 import ast
 import textwrap
 
-if sys.version_info.major == 3 and sys.version_info.minor < 9:
-    from typing_extensions import _AnnotatedAlias, get_type_hints
-else:
-    from typing import _AnnotatedAlias, get_type_hints
+from typing import _AnnotatedAlias, get_type_hints
 
 if typing.TYPE_CHECKING:
     from experimaestro.scheduler.base import Job
     from experimaestro.launchers import Launcher
     from experimaestro.core.objects import Config
+    from experimaestro.core.subparameters import Subparameters
+
+
+@dataclass
+class DeprecationInfo:
+    """Information about a deprecated configuration type."""
+
+    #: The original identifier before deprecation
+    original_identifier: "Identifier"
+
+    #: The target configuration class to convert to
+    target: type
+
+    #: If True, creating an instance immediately converts to the target type
+    replace: bool = False
 
 
 class Identifier:
@@ -221,6 +244,9 @@ class ObjectType(Type):
         self._title = None
         self.submit_hooks = set()
 
+        # Warning flag for non-resumable task directory cleanup
+        self.warned_clean_not_resumable = False
+
         # --- Get the identifier
         if identifier is None and hasattr(tp, "__xpmid__"):
             __xpmid__ = getattr(tp, "__xpmid__")
@@ -279,7 +305,20 @@ class ObjectType(Type):
         self.__initialized__ = False
         self._runtype = None
         self.annotations = []
-        self._deprecated = False
+        self._deprecation: Optional[DeprecationInfo] = None
+
+        # --- Value class (for external value types, e.g., nn.Module subclasses)
+        self._original_type: type = tp  # Keep reference to original config class
+
+        # --- Subparameters for partial identifier computation
+        self._subparameters: Dict[str, "Subparameters"] = {}
+
+    def set_value_type(self, value_class: type) -> None:
+        """Register an explicit value class for this configuration.
+
+        The value class will be used when creating instances via .instance().
+        """
+        self.value_type = value_class
 
     def addAnnotation(self, annotation):
         assert not self.__initialized__
@@ -335,15 +374,18 @@ class ObjectType(Type):
         # Add task
         if self.taskcommandfactory is not None:
             self.task = self.taskcommandfactory(self)
-        elif issubclass(self.value_type, Task):
+        elif issubclass(self._original_type, Task):
             self.task = self.getpythontaskcommand()
 
         # Add arguments from type hints
+        # Use _original_type since value_type may have been overridden by set_value_type
         from .arguments import TypeAnnotation
 
-        if hasattr(self.value_type, "__annotations__"):
-            typekeys = set(self.value_type.__dict__.get("__annotations__", {}).keys())
-            hints = get_type_hints(self.value_type, include_extras=True)
+        if hasattr(self._original_type, "__annotations__"):
+            typekeys = set(
+                self._original_type.__dict__.get("__annotations__", {}).keys()
+            )
+            hints = get_type_hints(self._original_type, include_extras=True)
             for key, typehint in hints.items():
                 # Filter out hints from parent classes
                 if key in typekeys:
@@ -356,17 +398,27 @@ class ObjectType(Type):
                             try:
                                 self.addArgument(
                                     options.create(
-                                        key, self.value_type, typehint.__args__[0]
+                                        key, self._original_type, typehint.__args__[0]
                                     )
                                 )
                             except Exception:
                                 logger.error(
                                     "while adding argument %s of %s",
                                     key,
-                                    self.value_type,
+                                    self._original_type,
                                 )
                                 raise
 
+        # Collect subparameters from class attributes
+        from .subparameters import Subparameters as SubparametersClass
+
+        for name, value in self._original_type.__dict__.items():
+            if isinstance(value, SubparametersClass):
+                # Auto-set name from attribute name if not already set
+                if value.name is None:
+                    value.name = name
+                self._subparameters[name] = value
+
     def name(self):
         return f"{self.value_type.__module__}.{self.value_type.__qualname__}"
 
@@ -378,7 +430,8 @@ class ObjectType(Type):
         self.__initialize__()
 
         # Get description from documentation
-        __doc__ = self.value_type.__dict__.get("__doc__", None)
+        # Use _original_type since value_type may have been overridden
+        __doc__ = self._original_type.__dict__.get("__doc__", None)
         if __doc__:
             parseddoc = parse(__doc__)
             self._title = parseddoc.short_description
@@ -407,24 +460,56 @@ class ObjectType(Type):
 
                 argname = None
 
-    def deprecate(self):
-        if len(self.value_type.__bases__) != 1:
-            raise RuntimeError(
-                "Deprecated configurations must have "
-                "only one parent (the new configuration)"
-            )
-        assert not self._deprecated, "Already deprecated"
-
-        # Uses the parent identifier (and saves the deprecated one for path updates)
-        self._deprecated_identifier = self.identifier
-        parent = self.value_type.__bases__[0].__getxpmtype__()
-        self.identifier = parent.identifier
-        self._deprecated = True
+    def deprecate(self, target=None, replace: bool = False):
+        """Mark this configuration type as deprecated.
+
+        Args:
+            target: Optional target configuration class. If provided, uses
+                    target's identifier. If None, uses parent class's identifier
+                    (legacy behavior requiring single inheritance).
+            replace: If True, creating an instance of this class immediately
+                    returns a converted instance of the target class.
+
+        When a target is specified, the deprecated class should define a
+        __convert__ method that returns an equivalent target configuration.
+        The identifier is computed from the converted configuration.
+        """
+        assert self._deprecation is None, "Already deprecated"
+
+        # Save the deprecated identifier for migration tools (fix_deprecated)
+        original_identifier = self.identifier
+
+        if target is not None:
+            # New mechanism: explicit target class
+            target_xpmtype = target.__getxpmtype__()
+            self.identifier = target_xpmtype.identifier
+            deprecation_target = target
+        else:
+            # Legacy mechanism: parent class is the target
+            if len(self.value_type.__bases__) != 1:
+                raise RuntimeError(
+                    "Deprecated configurations must have "
+                    "only one parent (the new configuration)"
+                )
+            parent = self.value_type.__bases__[0].__getxpmtype__()
+            self.identifier = parent.identifier
+            deprecation_target = self.value_type.__bases__[0]
+
+        self._deprecation = DeprecationInfo(
+            original_identifier=original_identifier,
+            target=deprecation_target,
+            replace=replace,
+        )
 
     @property
     def deprecated(self) -> bool:
         """Returns true if this type is deprecated"""
-        return self._deprecated
+        return self._deprecation is not None
+
+    @property
+    def _deprecated_identifier(self) -> Optional["Identifier"]:
+        """Returns the original identifier before deprecation (for backwards compatibility)"""
+        return self._deprecation.original_identifier if self._deprecation else None
 
     @property
     def description(self) -> str:
@@ -442,23 +527,72 @@ class ObjectType(Type):
         return self._arguments
 
     def addArgument(self, argument: Argument):
+        # Check if this argument overrides a parent argument
+        # _arguments is a ChainMap where maps[0] is current class, maps[1:] are parents
+        parent_argument = None
+        for parent_map in self._arguments.maps[1:]:
+            if argument.name in parent_map:
+                parent_argument = parent_map[argument.name]
+                break
+
+        if parent_argument is not None:
+            # Check type compatibility (child type should be subtype of parent type)
+            self._check_override_type_compatibility(argument, parent_argument)
+
+            # Warn if overrides flag is not set
+            if not argument.overrides:
+                logger.warning(
+                    "Parameter '%s' in %s overrides parent parameter from %s. "
+                    "Use field(overrides=True) to suppress this warning.",
+                    argument.name,
+                    self._original_type.__qualname__,
+                    (
+                        parent_argument.objecttype._original_type.__qualname__
+                        if parent_argument.objecttype
+                        else "unknown"
+                    ),
+                )
+
         self._arguments[argument.name] = argument
         argument.objecttype = self
 
-        # The the attribute for the config type
-        setattr(
-            self.config_type,
-            argument.name,
-            property(
-                lambda _self: _self.__xpm__.get(argument.name),
-                lambda _self, value: _self.__xpm__.set(argument.name, value),
-            ),
-        )
-
         # Check default value
         if argument.default is not None:
             argument.type.validate(argument.default)
 
+    def _check_override_type_compatibility(
+        self, child_arg: Argument, parent_arg: Argument
+    ):
+        """Check that the child argument type is compatible with the parent type.
+
+        For Config types, the child type should be a subtype of the parent type
+        (covariant). For other types, we check for exact match.
+        """
+        child_type = child_arg.type
+        parent_type = parent_arg.type
+
+        # Check if both are ObjectType (Config types)
+        if isinstance(child_type, ObjectType) and isinstance(parent_type, ObjectType):
+            child_pytype = child_type.value_type
+            parent_pytype = parent_type.value_type
+
+            # Check if child is a subtype of parent
+            if not issubclass(child_pytype, parent_pytype):
+                raise TypeError(
+                    f"Parameter '{child_arg.name}' type {child_pytype.__qualname__} "
+                    f"is not a subtype of parent type {parent_pytype.__qualname__}. "
+                    f"Override types must be subtypes of the parent type."
+                )
+        elif type(child_type) is not type(parent_type):
+            # For non-Config types, check for exact type match
+            # Different type classes (e.g., IntType vs StrType) are incompatible
+            raise TypeError(
+                f"Parameter '{child_arg.name}' type {type(child_type).__name__} "
+                f"is not compatible with parent type {type(parent_type).__name__}. "
+                f"Override types must be the same type or a subtype."
+            )
+        # Same type class is allowed (e.g., both are IntType)
+
     def getArgument(self, key: str) -> Argument:
         self.__initialize__()
         return self._arguments[key]
@@ -466,7 +600,10 @@ class ObjectType(Type):
     def parents(self) -> Iterator["ObjectType"]:
         from .objects import Config, Task
 
-        for tp in self.value_type.__bases__:
+        # Use _original_type to avoid issues when value_type has been
+        # overridden by set_value_type (the value class would create
+        # circular references since it inherits from the config class)
+        for tp in self._original_type.__bases__:
             if issubclass(tp, Config) and tp not in [Config, Task]:
                 yield tp.__xpmtype__
 

experimaestro/exceptions.py

@@ -1,2 +1,28 @@
 class HandledException(Exception):
     pass
+
+
+class GracefulTimeout(Exception):
+    """Exception raised to signal a graceful timeout in resumable tasks.
+
+    Raise this exception when a task needs to checkpoint and exit before
+    a time limit (e.g., SLURM walltime). The task will be marked for retry
+    rather than as failed.
+
+    Example::
+
+        ```python
+            class LongTraining(ResumableTask):
+                def execute(self):
+                    for epoch in range(self.epochs):
+                        remaining = self.remaining_time()
+                        if remaining is not None and remaining < 300:
+                            save_checkpoint(self.checkpoint, epoch)
+                            raise GracefulTimeout("Not enough time for another epoch")
+                        train_one_epoch()
+        ```
+    """
+
+    def __init__(self, message: str = "Task stopped gracefully before timeout"):
+        self.message = message
+        super().__init__(message)

experimaestro/experiments/cli.py

@@ -52,8 +52,7 @@ class ExperimentHelper:
 class ExperimentCallable(Protocol):
     """Protocol for the run function"""
 
-    def __call__(self, helper: ExperimentHelper, configuration: Any):
-        ...
+    def __call__(self, helper: ExperimentHelper, configuration: Any): ...  # noqa: E704
 
 
 class ConfigurationLoader:
@@ -126,6 +125,11 @@ class ConfigurationLoader:
     default=None,
     help="Port for monitoring (can be defined in the settings.yaml file)",
 )
+@click.option(
+    "--console",
+    is_flag=True,
+    help="Launch Textual console UI for monitoring with logs",
+)
 @click.option(
     "--file",
     "xp_file",
@@ -162,6 +166,7 @@ def experiments_cli(  # noqa: C901
     xp_file: str,
     host: str,
     port: int,
+    console: bool,
     xpm_config_dir: Path,
     workdir: Optional[Path],
     workspace: Optional[str],
@@ -298,43 +303,120 @@ def experiments_cli(  # noqa: C901
         configuration, structured_config_mode=SCMode.INSTANTIATE
     )
 
-    # Define the workspace
-    ws_env = find_workspace(workdir=workdir, workspace=workspace)
-
-    workdir = ws_env.path
-
     # --- Sets up the experiment ID
-
-    # --- Runs the experiment
     if xp_configuration.add_timestamp:
         timestamp = datetime.datetime.now().strftime("%Y%m%d-%H%M")
         experiment_id = f"""{xp_configuration.id}-{timestamp}"""
     else:
         experiment_id = xp_configuration.id
 
+    # Define the workspace (may auto-select based on experiment_id triggers)
+    ws_env = find_workspace(
+        workdir=workdir, workspace=workspace, experiment_id=experiment_id
+    )
+
+    workdir = ws_env.path
+
     logging.info(
         "Running experiment %s working directory %s",
         experiment_id,
         str(workdir.resolve()),
     )
-    with experiment(
-        ws_env, experiment_id, host=host, port=port, run_mode=run_mode
-    ) as xp:
-        # Set up the environment
-        # (1) global settings (2) workspace settings and (3) command line settings
-        for key, value in env:
-            xp.setenv(key, value)
-
-        # Sets the python path
-        xp.workspace.python_path.extend(python_path)
 
+    # Define the experiment execution function
+    def run_experiment_code(xp_holder=None, xp_ready_event=None, register_signals=True):
+        """Run the experiment code - optionally storing xp in xp_holder"""
         try:
-            # Run the experiment
-            helper.xp = xp
-            helper.run(list(args), xp_configuration)
-
-            # ... and wait
-            xp.wait()
+            with experiment(
+                ws_env,
+                experiment_id,
+                host=host,
+                port=port,
+                run_mode=run_mode,
+                register_signals=register_signals,
+            ) as xp:
+                if xp_holder is not None:
+                    xp_holder["xp"] = xp
+                if xp_ready_event is not None:
+                    xp_ready_event.set()  # Signal that xp is ready
+
+                # Test logging from experiment thread
+                logging.info("Experiment started in background thread")
+
+                # Set up the environment
+                for key, value in env:
+                    xp.setenv(key, value)
+
+                # Sets the python path
+                xp.workspace.python_path.extend(python_path)
+
+                # Run the experiment
+                helper.xp = xp
+                helper.run(list(args), xp_configuration)
+
+                # ... and wait
+                xp.wait()
 
         except HandledException:
             sys.exit(1)
+
+    if console:
+        # Run experiment in background thread, console UI in main thread
+        import threading
+        from experimaestro.tui import ExperimentTUI
+
+        xp_holder = {"xp": None}
+        exception_holder = {"exception": None}
+        xp_ready = threading.Event()
+
+        def run_in_thread():
+            try:
+                # Don't register signals in background thread
+                run_experiment_code(xp_holder, xp_ready, register_signals=False)
+                # Add a test message after experiment completes
+                logging.info("Experiment thread completed")
+                print("Experiment thread print test")
+            except Exception as e:
+                exception_holder["exception"] = e
+                xp_ready.set()  # Signal even on error
+
+        # Start experiment in background thread
+        exp_thread = threading.Thread(target=run_in_thread, daemon=True)
+        exp_thread.start()
+
+        # Wait for experiment to start (up to 30 seconds)
+        if not xp_ready.wait(timeout=30.0):
+            cprint("Timeout waiting for experiment to start", "red", file=sys.stderr)
+            sys.exit(1)
+
+        if xp_holder["xp"] is None:
+            cprint("Failed to start experiment", "red", file=sys.stderr)
+            if exception_holder["exception"]:
+                raise exception_holder["exception"]
+            sys.exit(1)
+
+        # Run TUI in main thread (handles signals via Textual)
+        tui_app = ExperimentTUI(
+            workdir=workdir,
+            state_provider=xp_holder["xp"].state_provider,
+            show_logs=True,
+        )
+
+        try:
+            # Textual automatically captures stdout/stderr via Print events
+            tui_app.run()
+        finally:
+            # TUI exited (user pressed q or Ctrl+C) - stop the experiment
+            if xp_holder["xp"]:
+                xp_holder["xp"].stop()
+
+            # Wait for experiment thread to finish
+            exp_thread.join(timeout=5.0)
+
+        # Handle exceptions
+        if exception_holder["exception"]:
+            raise exception_holder["exception"]
+
+    else:
+        # Normal mode without TUI - run directly
+        run_experiment_code()

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/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"""