experimaestro 1.11.1__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:
@@ -203,18 +226,14 @@ class ObjectType(Type):
     """ObjectType contains class-level information about
     experimaestro configurations and tasks
 
-    :param objecttype: The python Type of the associated object
-    :param configtype: The python Type of the configuration object that uses
-        property for arguments
+    :param value_type: The Python type of the associated object
+    :param config_type: The Python type of the configuration object
     """
 
-    # Those entries should not be copied in the __dict__
-    FORBIDDEN_KEYS = set(("__dict__", "__weakref__"))
-
     def __init__(
         self,
         tp: type,
-        identifier: Union[str, Identifier] = None,
+        identifier: Union[str, Identifier, None] = None,
     ):
         """Creates a type"""
         from .objects import Config, ConfigMixin
@@ -225,7 +244,10 @@ class ObjectType(Type):
         self._title = None
         self.submit_hooks = set()
 
-        # Get the identifier
+        # 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__")
             if isinstance(__xpmid__, Identifier):
@@ -250,58 +272,53 @@ class ObjectType(Type):
         # --- Creates the config type and not config type
 
         self.originaltype = tp
-        if not issubclass(tp, Config):
-            # Adds Config as a base class if not present
-            __bases__ = () if tp.__bases__ == (object,) else tp.__bases__
-            __dict__ = dict(tp.__dict__)
-
-            __dict__ = {
-                key: value
-                for key, value in tp.__dict__.items()
-                if key not in ObjectType.FORBIDDEN_KEYS
-            }
-            self.basetype = type(tp.__name__, (Config,) + __bases__, __dict__)
-            self.basetype.__module__ = tp.__module__
-            self.basetype.__qualname__ = tp.__qualname__
-        else:
-            self.basetype = tp
+        assert issubclass(tp, Config)
+        self.value_type = tp
 
         # --- Create the type-specific configuration class (XPMConfig)
         __configbases__ = tuple(
-            s.__getxpmtype__().configtype
+            s.__getxpmtype__().config_type
             for s in tp.__bases__
             if issubclass(s, Config) and (s is not Config)
         ) or (ConfigMixin,)
 
-        *tp_qual, tp_name = self.basetype.__qualname__.split(".")
-        self.configtype = type(
-            f"{tp_name}.XPMConfig", __configbases__ + (self.basetype,), {}
+        *tp_qual, tp_name = self.value_type.__qualname__.split(".")
+        self.config_type = type(
+            f"{tp_name}.XPMConfig", __configbases__ + (self.value_type,), {}
         )
-        self.configtype.__qualname__ = ".".join(tp_qual + [self.configtype.__name__])
-        self.configtype.__module__ = tp.__module__
+        self.config_type.__qualname__ = ".".join(tp_qual + [self.config_type.__name__])
+        self.config_type.__module__ = tp.__module__
 
-        # Return type is used by tasks to change the output
-        if hasattr(self.basetype, "task_outputs") or False:
+        # --- Get the return type
+        if hasattr(self.value_type, "task_outputs") or False:
             self.returntype = get_type_hints(
-                getattr(self.basetype, "task_outputs")
+                getattr(self.value_type, "task_outputs")
             ).get("return", typing.Any)
         else:
-            self.returntype = self.basetype
+            self.returntype = self.value_type
 
-        # Registers ourselves
-        self.basetype.__xpmtype__ = self
-        self.configtype.__xpmtype__ = self
+        # --- Registers ourselves
+        self.value_type.__xpmtype__ = self
+        self.config_type.__xpmtype__ = self
 
-        # Other initializations
+        # --- Other initializations
         self.__initialized__ = False
         self._runtype = None
         self.annotations = []
-        self._deprecated = False
+        self._deprecation: Optional[DeprecationInfo] = None
 
-    @property
-    def objecttype(self):
-        """Returns the object type"""
-        return self.basetype.XPMValue
+        # --- 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__
@@ -357,15 +374,18 @@ class ObjectType(Type):
         # Add task
         if self.taskcommandfactory is not None:
             self.task = self.taskcommandfactory(self)
-        elif issubclass(self.basetype, 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.basetype, "__annotations__"):
-            typekeys = set(self.basetype.__dict__.get("__annotations__", {}).keys())
-            hints = get_type_hints(self.basetype, 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:
@@ -378,19 +398,29 @@ class ObjectType(Type):
                             try:
                                 self.addArgument(
                                     options.create(
-                                        key, self.objecttype, typehint.__args__[0]
+                                        key, self._original_type, typehint.__args__[0]
                                     )
                                 )
                             except Exception:
                                 logger.error(
                                     "while adding argument %s of %s",
                                     key,
-                                    self.objecttype,
+                                    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.basetype.__module__}.{self.basetype.__qualname__}"
+        return f"{self.value_type.__module__}.{self.value_type.__qualname__}"
 
     def __parsedoc__(self):
         """Parse the documentation"""
@@ -400,7 +430,8 @@ class ObjectType(Type):
         self.__initialize__()
 
         # Get description from documentation
-        __doc__ = self.basetype.__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
@@ -429,24 +460,56 @@ class ObjectType(Type):
 
                 argname = None
 
-    def deprecate(self):
-        if len(self.basetype.__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.basetype.__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:
@@ -454,7 +517,7 @@ class ObjectType(Type):
         return self._description
 
     @property
-    def title(self) -> Dict[str, Argument]:
+    def title(self) -> str:
         self.__parsedoc__()
         return self._title or str(self.identifier)
 
@@ -464,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.configtype,
-            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]
@@ -488,7 +600,10 @@ class ObjectType(Type):
     def parents(self) -> Iterator["ObjectType"]:
         from .objects import Config, Task
 
-        for tp in self.basetype.__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__
 
@@ -504,7 +619,7 @@ class ObjectType(Type):
         if not isinstance(value, Config):
             raise ValueError(f"{value} is not an experimaestro type or task")
 
-        types = self.basetype
+        types = self.value_type
 
         if not isinstance(value, types):
             raise ValueError(
@@ -519,7 +634,7 @@ class ObjectType(Type):
 
     def fullyqualifiedname(self) -> str:
         """Returns the fully qualified (Python) name"""
-        return f"{self.basetype.__module__}.{self.basetype.__qualname__}"
+        return f"{self.value_type.__module__}.{self.value_type.__qualname__}"
 
 
 class TypeProxy:

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