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

experimaestro/core/objects/config.py

@@ -30,13 +30,14 @@ from typing import (
 )
 import sys
 import experimaestro
-from experimaestro.utils import logger
+from experimaestro.utils import logger, get_caller_location
 from experimaestro.core.types import DeprecatedAttribute, ObjectType, TypeVarType
 from ..context import SerializationContext, SerializedPath, SerializedPathLoader
 
 if TYPE_CHECKING:
     from ..callbacks import TaskEventListener
     from ..identifier import Identifier
+    from ..subparameters import Subparameters
     from experimaestro.scheduler.base import Job
     from experimaestro.scheduler.workspace import RunMode
     from experimaestro.launchers import Launcher
@@ -55,6 +56,26 @@ T = TypeVar("T", bound="Config")
 
 
 DependentMarker = Callable[["Config"], None]
+"""Type alias for dependency marker functions.
+
+A DependentMarker is a callable that marks a configuration as a dependency
+of another configuration. Used in ``task_outputs()`` and dynamic output methods
+to establish task dependencies.
+
+Example::
+
+    class Learn(Task):
+        model: Param[Model]
+
+        def task_outputs(self, dep: DependentMarker):
+            return dep(Checkpoint.C(model=self.model, path=self.checkpoint_path))
+
+    class Validation(Config):
+        model: Param[Model]
+
+        def checkpoint(self, dep: DependentMarker, *, step: int) -> Checkpoint:
+            return dep(Checkpoint.C(model=self.model, step=step))
+"""
 
 
 def updatedependencies(
@@ -90,9 +111,6 @@ NOT_SET = object()
 
 @define()
 class WatchedOutput:
-    #: The enclosing job
-    job: "Job"
-
     #: The configuration containing the watched output
     config: "ConfigInformation"
 
@@ -105,6 +123,9 @@ class WatchedOutput:
     #: The callback to call (with the output of the previous method)
     callback: Callable
 
+    #: The enclosing job (set when registered with scheduler)
+    job: Optional["Job"] = None
+
 
 def get_generated_paths(
     v: Union["ConfigMixin", list, dict],
@@ -142,6 +163,22 @@ def get_generated_paths(
     return paths
 
 
+@define
+class TaskStub:
+    """Stub for a task that was not loaded during partial loading.
+
+    This is used when loading configurations from disk (e.g., HuggingFace)
+    where the task code may have changed or is not available. The stub stores
+    the identifier and typename so the information is preserved.
+    """
+
+    identifier: "Identifier"
+    """The experimaestro identifier of the task"""
+
+    typename: str
+    """The type name of the task (e.g., 'mymodule.MyTask')"""
+
+
 class ConfigInformation:
     """Holds experimaestro information for a config (or task) instance"""
 
@@ -158,7 +195,9 @@ class ConfigInformation:
         self.values = {}
 
         # Meta-informations
-        self._tags: dict[str, Any] = {}
+        # Tags are stored as {name: (value, source_location)}
+        # where source_location is "file:line" string for error reporting
+        self._tags: dict[str, tuple[Any, str]] = {}
         self._initinfo = ""
 
         self._taskoutput = None
@@ -192,6 +231,9 @@ class ConfigInformation:
         self._identifier = None
         """The configuration identifier (cached when sealed)"""
 
+        self._partial_identifiers: Dict[str, "Identifier"] = {}
+        """Cached partial identifiers (keyed by subparameters name)"""
+
         self._validated = False
         self._sealed = False
         self._meta = None
@@ -331,8 +373,17 @@ class ConfigInformation:
             f" (current typevars bindings: {self.concrete_typevars})"
         )
 
-    def addtag(self, name, value):
-        self._tags[name] = value
+    def addtag(self, name, value, source: str = None):
+        """Add a tag with optional source location for error reporting
+
+        Args:
+            name: Tag name
+            value: Tag value
+            source: Source location string (file:line). If None, captured from caller.
+        """
+        if source is None:
+            source = get_caller_location(skip_frames=1)
+        self._tags[name] = (value, source)
 
     def xpmvalues(self, generated=False):
         """Returns an iterarator over arguments and associated values"""
@@ -344,11 +395,29 @@ class ConfigInformation:
         class TagFinder(ConfigWalk):
             def __init__(self):
                 super().__init__(recurse_task=True)
-                self.tags = {}
+                # Store {name: (value, source)} for conflict detection
+                self.tags_with_source: dict[str, tuple[Any, str]] = {}
 
             def postprocess(self, stub, config: Config, values):
-                self.tags.update(config.__xpm__._tags)
-                return self.tags
+                for name, (value, source) in config.__xpm__._tags.items():
+                    if name in self.tags_with_source:
+                        existing_value, existing_source = self.tags_with_source[name]
+                        if existing_value != value:
+                            logger.warning(
+                                "Tag '%s' has conflicting values: "
+                                "'%s' (set at %s) vs '%s' (set at %s). "
+                                "Using the latter value.",
+                                name,
+                                existing_value,
+                                existing_source,
+                                value,
+                                source,
+                            )
+                    self.tags_with_source[name] = (value, source)
+                # Return just the values (without source info)
+                return {
+                    name: value for name, (value, _) in self.tags_with_source.items()
+                }
 
         return TagFinder()(self.pyobject)
 
@@ -489,6 +558,33 @@ class ConfigInformation:
             self._identifier = identifier
         return identifier
 
+    def get_partial_identifier(self, subparameters: "Subparameters") -> "Identifier":
+        """Get the partial identifier for a given subparameters instance.
+
+        Partial identifiers exclude certain parameter groups, allowing
+        configurations that differ only in those groups to share the same
+        partial identifier (and thus the same partial directory).
+
+        Args:
+            subparameters: The Subparameters instance defining which groups
+                to include/exclude.
+
+        Returns:
+            The partial identifier for this configuration.
+        """
+        from ..identifier import IdentifierComputer
+
+        name = subparameters.name
+        if name in self._partial_identifiers:
+            return self._partial_identifiers[name]
+
+        identifier = IdentifierComputer.compute_partial(self.pyobject, subparameters)
+
+        if self._sealed:
+            self._partial_identifiers[name] = identifier
+
+        return identifier
+
     def dependency(self):
         """Returns a dependency"""
         from experimaestro.scheduler import JobDependency
@@ -563,10 +659,20 @@ class ConfigInformation:
 
         :param method: The method to watch
         :param callback: The callback
+
+        :raises TypeError: If the task is not a ResumableTask
         """
-        watched = WatchedOutput(
-            self, method.__self__, method.__name__, method, callback
-        )
+        # Only ResumableTask can have dynamic outputs - regular tasks
+        # have their directories cleaned up, losing the output file
+        if not isinstance(self.pyobject, ResumableTask):
+            raise TypeError(
+                f"Only ResumableTask can use watch_output. "
+                f"{self.xpmtype} is not a ResumableTask. "
+                "Dynamic outputs require the task directory to be preserved "
+                "across restarts, which only ResumableTask provides."
+            )
+
+        watched = WatchedOutput(method.__self__, method.__name__, method, callback)
         self.watched_outputs.append(watched)
         if self.job:
             self.job.watch_output(watched)
@@ -587,6 +693,7 @@ class ConfigInformation:
         *,
         run_mode=None,
         init_tasks: List["LightweightTask"] = [],
+        max_retries: Optional[int] = None,
     ):
         from experimaestro.scheduler import experiment, JobContext
         from experimaestro.scheduler.workspace import RunMode
@@ -606,7 +713,11 @@ class ConfigInformation:
 
         # Creates a new job
         self.job = self.xpmtype.task(
-            self.pyobject, launcher=launcher, workspace=workspace, run_mode=run_mode
+            self.pyobject,
+            launcher=launcher,
+            workspace=workspace,
+            run_mode=run_mode,
+            max_retries=max_retries,
         )
 
         # Validate the object
@@ -644,7 +755,6 @@ class ConfigInformation:
         ) or RunMode.NORMAL
         if run_mode == RunMode.NORMAL:
             TaskEventListener.connect(experiment.CURRENT)
-            experiment.CURRENT.submit(self.job)
             other = experiment.CURRENT.submit(self.job)
             if other:
                 # Our job = previously submitted job
@@ -844,7 +954,6 @@ class ConfigInformation:
 
         The format is an array of objects
         {
-            "tags: [ LIST_OF_TAGS ],
             "workspace": FOLDERPATH,
             "version": 2,
             "objects": [
@@ -864,6 +973,10 @@ class ConfigInformation:
 
         The last object is the one that is serialized
 
+        Note: Tags are no longer stored in params.json. They are managed by the
+        experiment state provider (scoped to job_id, experiment_id, run_id) and
+        also stored in experiment state.json for full experiment details.
+
         Arguments:
             out {io.TextIOBase} -- The output stream
             context {[type]} -- the command context
@@ -871,7 +984,6 @@ class ConfigInformation:
         json.dump(
             {
                 "workspace": str(context.workspace.path.absolute()),
-                "tags": {key: value for key, value in self.tags().items()},
                 "version": 2,
                 "experimaestro": experimaestro.__version__,
                 "objects": self.__get_objects__([], context),
@@ -896,12 +1008,18 @@ class ConfigInformation:
         path: Union[str, Path, SerializedPathLoader],
         as_instance: bool = False,
         return_tasks: bool = False,
+        partial_loading: Optional[bool] = None,
     ) -> "Config":
         """Deserialize a configuration
 
         :param path: The filesystem Path to use, or a way to download the
             information through a function taking two arguments
         :param as_instance: Return an instance
+        :param return_tasks: Return init tasks instead of executing them
+        :param partial_loading: If True, skip loading task references. If None
+            (default), partial_loading is enabled when as_instance is True.
+            This is useful when loading configurations from disk (e.g.,
+            HuggingFace) where the task code may have changed.
         :return: a Config object, its instance or a tuple (instance, init_tasks) is return_tasks is True
         """
         # Load
@@ -926,6 +1044,7 @@ class ConfigInformation:
             as_instance=as_instance,
             data_loader=data_loader,
             return_tasks=return_tasks,
+            partial_loading=partial_loading,
         )
 
     @staticmethod
@@ -973,6 +1092,7 @@ class ConfigInformation:
         as_instance=True,
         save_directory: Optional[Path] = None,
         discard_id: bool = False,
+        partial_loading: Optional[bool] = None,
     ) -> "ConfigMixin": ...
 
     @overload
@@ -983,6 +1103,7 @@ class ConfigInformation:
         return_tasks=True,
         save_directory: Optional[Path] = None,
         discard_id: bool = False,
+        partial_loading: Optional[bool] = None,
     ) -> Tuple["Config", List["LightweightTask"]]: ...
 
     @overload
@@ -992,23 +1113,115 @@ class ConfigInformation:
         as_instance=False,
         save_directory: Optional[Path] = None,
         discard_id: bool = False,
+        partial_loading: Optional[bool] = None,
     ) -> "Config": ...
 
+    @staticmethod
+    def _get_field_refs(value: Any) -> Set[int]:
+        """Recursively extract object references from a serialized field value"""
+        refs: Set[int] = set()
+        if isinstance(value, dict):
+            if value.get("type") == "python":
+                refs.add(value["value"])
+            else:
+                for v in value.values():
+                    refs.update(ConfigInformation._get_field_refs(v))
+        elif isinstance(value, list):
+            for v in value:
+                refs.update(ConfigInformation._get_field_refs(v))
+        return refs
+
+    @staticmethod
+    def _compute_skipped_ids(definitions: List[Dict]) -> Set[int]:
+        """Compute IDs of objects only reachable through task references.
+
+        When partial_loading is enabled, we skip loading task references and
+        any objects that are only used by those tasks. This method computes
+        which object IDs should be skipped by finding objects reachable from
+        the main object (last definition) without following task references.
+
+        :param definitions: List of object definitions from JSON
+        :return: Set of object IDs to skip loading
+        """
+        # Build index of definitions by ID
+        def_by_id = {d["id"]: d for d in definitions}
+
+        # Compute reachable objects from main object (last definition)
+        # without going through task references
+        main_defn = definitions[-1]
+        main_id = main_defn["id"]
+        reachable: Set[int] = set()
+        to_visit = [main_id]
+
+        # Also include init-tasks as reachable (needed for as_instance/return_tasks)
+        for init_task_id in main_defn.get("init-tasks", []):
+            to_visit.append(init_task_id)
+
+        while to_visit:
+            obj_id = to_visit.pop()
+            if obj_id in reachable:
+                continue
+            reachable.add(obj_id)
+
+            defn = def_by_id.get(obj_id)
+            if defn is None:
+                continue
+
+            # Add field references (not task reference)
+            for field_value in defn.get("fields", {}).values():
+                for ref_id in ConfigInformation._get_field_refs(field_value):
+                    if ref_id not in reachable:
+                        to_visit.append(ref_id)
+
+            # Note: we intentionally skip defn["task"] to avoid loading tasks
+
+        # All objects not reachable should be skipped
+        all_ids = {d["id"] for d in definitions}
+        return all_ids - reachable
+
     @staticmethod
     def load_objects(  # noqa: C901
         definitions: List[Dict],
         as_instance=True,
         data_loader: Optional[SerializedPathLoader] = None,
         discard_id: bool = False,
+        partial_loading: bool = False,
     ):
-        """Load the objects"""
+        """Load the objects
+
+        :param definitions: List of object definitions from JSON
+        :param as_instance: Return instances instead of configs
+        :param data_loader: Function to load data files
+        :param discard_id: If True, don't use the stored identifier
+        :param partial_loading: If True, skip loading task references. This is
+            useful when loading configurations from disk (e.g., HuggingFace)
+            where the task code may have changed.
+        """
         o = None
         objects = {}
         import experimaestro.taskglobals as taskglobals
         from ..identifier import Identifier
 
+        # Compute which objects to skip when partial_loading
+        skipped_ids = (
+            ConfigInformation._compute_skipped_ids(definitions)
+            if partial_loading
+            else set()
+        )
+
         # Loop over all the definitions and create objects
         for definition in definitions:
+            obj_id = definition["id"]
+
+            # Skip objects that are only reachable through task references
+            if obj_id in skipped_ids:
+                # Create a TaskStub for skipped task objects
+                objects[obj_id] = TaskStub(
+                    identifier=Identifier.from_state_dict(definition["identifier"]),
+                    typename=definition.get("typename", definition["type"]),
+                )
+                continue
+
             module_name = definition["module"]
 
             # Avoids problem when runing module
@@ -1041,12 +1254,18 @@ class ConfigInformation:
                 o = cls.__new__(cls)
             else:
                 o = cls.XPMConfig.__new__(cls.XPMConfig)
-            assert definition["id"] not in objects, "Duplicate id %s" % definition["id"]
-            objects[definition["id"]] = o
+            assert obj_id not in objects, "Duplicate id %s" % obj_id
+            objects[obj_id] = o
 
         # Now that objects have been created, fill in the fields
         for definition in definitions:
-            o = objects[definition["id"]]
+            obj_id = definition["id"]
+
+            # Skip processing skipped objects (they are TaskStubs)
+            if obj_id in skipped_ids:
+                continue
+
+            o = objects[obj_id]
             xpmtype = o.__getxpmtype__()  # type: ObjectType
 
             # If instance...
@@ -1136,13 +1355,20 @@ class ConfigInformation:
         data_loader: Optional[SerializedPathLoader] = None,
         discard_id: bool = False,
         return_tasks: bool = False,
+        partial_loading: Optional[bool] = None,
     ):
+        # Determine effective partial_loading: as_instance implies partial_loading
+        effective_partial_loading = (
+            partial_loading if partial_loading is not None else as_instance
+        )
+
         # Get the objects
         objects = ConfigInformation.load_objects(
             definitions,
             as_instance=as_instance,
             data_loader=data_loader,
             discard_id=discard_id,
+            partial_loading=effective_partial_loading,
         )
 
         # Get the last one
@@ -1258,9 +1484,98 @@ class ConfigMixin:
     """The __xpm__ object contains all instance specific information about a
     configuration/task"""
 
+    # Set when this instance was created via a deprecated config with replace=True
+    _deprecated_from: "ConfigMixin | None" = None
+
+    def __new__(cls, **kwargs):
+        """Create a new config instance, handling deprecated replacements."""
+        xpmtype = cls.__xpmtype__
+
+        # Check if this is a deprecated type with replace=True
+        if xpmtype._deprecation is not None and xpmtype._deprecation.replace:
+            # Create the deprecated instance normally
+            instance = object.__new__(cls)
+            # Initialize it
+            ConfigMixin.__init__(instance, **kwargs)
+            # Convert to the new type
+            converted = instance.__convert__()
+            # Mark that this came from a deprecated config
+            converted._deprecated_from = instance
+            return converted
+
+        # Normal creation
+        return object.__new__(cls)
+
+    def __getattribute__(self, name: str):
+        """Get an attribute, handling XPM arguments specially.
+
+        We use __getattribute__ instead of __getattr__ because default values
+        like `b: Param[X] = None` create class attributes that would prevent
+        __getattr__ from being called.
+        """
+        # Get __xpm__ without recursion
+        try:
+            xpm = object.__getattribute__(self, "__xpm__")
+        except AttributeError:
+            # During early init, __xpm__ may not exist yet
+            return object.__getattribute__(self, name)
+
+        # Check if this is an XPM argument - parameters take precedence
+        xpmtype = object.__getattribute__(self, "__xpmtype__")
+        if name in xpmtype.arguments:
+            return xpm.get(name)
+
+        # Fall back to normal lookup (methods, etc.)
+        return object.__getattribute__(self, name)
+
+    def __setattr__(self, name: str, value):
+        """Set an attribute, handling XPM arguments specially."""
+        # Allow setting internal attributes directly
+        if name in ("__xpm__", "_deprecated_from"):
+            object.__setattr__(self, name, value)
+            return
+
+        # Check if we have __xpm__ yet (might not during early init)
+        xpm = self.__dict__.get("__xpm__")
+        if xpm is None:
+            object.__setattr__(self, name, value)
+            return
+
+        # Check if this is an XPM argument
+        xpmtype = self.__xpmtype__
+        if name in xpmtype.arguments:
+            # Handle TaggedValue: extract value and add tag
+            if isinstance(value, TaggedValue):
+                actual_value = value.value
+                source = get_caller_location(skip_frames=1)
+                xpm.addtag(name, actual_value, source=source)
+                xpm.set(name, actual_value)
+            else:
+                xpm.set(name, value)
+            return
+
+        # Check for deprecated replacement warning
+        deprecated_from = self.__dict__.get("_deprecated_from")
+        if deprecated_from is not None:
+            deprecated_xpmtype = deprecated_from.__xpmtype__
+            if name in deprecated_xpmtype.arguments:
+                logger.warning(
+                    f"Attribute '{name}' was in deprecated config "
+                    f"{deprecated_xpmtype.identifier} but is not in "
+                    f"{xpmtype.identifier}. The value is being discarded."
+                )
+                return  # Don't set the attribute
+
+        # Normal attribute setting
+        object.__setattr__(self, name, value)
+
     def __init__(self, **kwargs):
         """Initialize the configuration with the given parameters"""
 
+        # Skip if already initialized (can happen with deprecated replace=True)
+        if hasattr(self, "__xpm__"):
+            return
+
         # Add configuration
         xpmtype = self.__xpmtype__
 
@@ -1294,7 +1609,8 @@ class ConfigMixin:
             # Special case of a tagged value
             if isinstance(value, TaggedValue):
                 value = value.value
-                self.__xpm__._tags[name] = value
+                # Use _initinfo as source since tag is set at config creation
+                self.__xpm__.addtag(name, value, source=xpm._initinfo)
 
             # Really set the value
             xpm.set(name, value)
@@ -1312,7 +1628,9 @@ class ConfigMixin:
         )
 
     def tag(self, name, value):
-        self.__xpm__.addtag(name, value)
+        # Capture caller's location and pass to addtag
+        source = get_caller_location(skip_frames=1)
+        self.__xpm__.addtag(name, value, source=source)
         return self
 
     def __eq__(self, other):
@@ -1372,16 +1690,22 @@ class ConfigMixin:
         launcher=None,
         run_mode: "RunMode" = None,
         init_tasks: List["LightweightTask"] = [],
+        max_retries: Optional[int] = None,
     ):
         """Submit this task
 
         :param workspace: the workspace, defaults to None
         :param launcher: The launcher, defaults to None
         :param run_mode: Run mode (if None, uses the workspace default)
+        :param max_retries: Maximum number of retries for resumable tasks that timeout (default: from workspace settings or 3)
         :return: an object object
         """
         return self.__xpm__.submit(
-            workspace, launcher, run_mode=run_mode, init_tasks=init_tasks
+            workspace,
+            launcher,
+            run_mode=run_mode,
+            init_tasks=init_tasks,
+            max_retries=max_retries,
         )
 
     def stdout(self):
@@ -1419,6 +1743,59 @@ class ConfigMixin:
         # Add other dependencies
         self.__xpm__.add_dependencies(*other.__xpm__.dependencies)
 
+    def __rmatmul__(self, other: "ConfigMixin") -> "ConfigMixin":
+        """Right-associative composition operator: B() @ A(x=1) is equivalent to B(a=A(x=1))
+
+        For expression `other @ self`, finds the unique parameter in `other` that
+        accepts `self`'s type and sets it. Returns `self` to enable right-associative
+        chaining: `Outer() @ Middle() @ Inner()` builds Outer(middle=Middle(inner=Inner())).
+
+        The chain is evaluated left-to-right by Python, but returns the inner config
+        so each step adds the current result as a parameter to the next outer config.
+
+        :param other: The outer configuration that will receive self as a parameter
+        :return: self (the inner configuration) to continue the chain
+        """
+        if not isinstance(other, ConfigMixin):
+            return NotImplemented
+
+        # Find parameters in 'other' that can accept self's type
+        self_type = self.__xpmtype__.value_type
+        matching_params = []
+
+        for name, argument in other.__xpmtype__.arguments.items():
+            # Get the expected type for this argument
+            arg_type = argument.type
+            if hasattr(arg_type, "value_type"):
+                # It's an ObjectType wrapper
+                expected_type = arg_type.value_type
+            elif hasattr(arg_type, "__origin__"):
+                # Generic type like Optional[X] or List[X]
+                continue  # Skip complex types for now
+            elif isinstance(arg_type, type):
+                expected_type = arg_type
+            else:
+                continue
+
+            # Check if self's type is compatible
+            if isinstance(expected_type, type) and issubclass(self_type, expected_type):
+                matching_params.append(name)
+
+        if len(matching_params) == 0:
+            raise ValueError(
+                f"No parameter in {other.__xpmtype__} accepts type {self_type.__name__}"
+            )
+        if len(matching_params) > 1:
+            raise ValueError(
+                f"Ambiguous composition: parameters {matching_params} in "
+                f"{other.__xpmtype__} all accept type {self_type.__name__}"
+            )
+
+        # Set the parameter on 'other'
+        param_name = matching_params[0]
+        other.__xpm__.set(param_name, self)
+        return other
+
 
 class Config:
     """Base type for all objects in python interface"""
@@ -1442,6 +1819,68 @@ class Config:
         """Alias for XPMConfig"""
         return cls.XPMConfig
 
+    @classproperty
+    def XPMValue(cls):
+        """Get the value class for this configuration.
+
+        Returns the explicitly registered value class, or the base config class
+        if no value class was registered.
+        """
+        return cls.__getxpmtype__().value_type
+
+    @classmethod
+    def value_class(cls):
+        """Decorator to register an external value class for this configuration.
+
+        This allows declaring a separate class that will be used when creating
+        instances, which is useful to avoid initializing resources (e.g., PyTorch)
+        when only configuring.
+
+        .. code-block:: python
+
+            class Model(Config):
+                hidden_size: Param[int]
+
+            @Model.value_class()
+            class TorchModel(Model, nn.Module):
+                def __init__(self):
+                    super().__init__()
+                    self.layer = nn.Linear(self.hidden_size, self.hidden_size)
+
+        The value class must be a subclass of the configuration class
+        and a subclass of parent configuration value classes (if any).
+        """
+
+        def decorator(value_class: type) -> type:
+            xpmtype = cls.__getxpmtype__()
+
+            # Check that value class is a subclass of the config class
+            if not issubclass(value_class, cls):
+                raise TypeError(
+                    f"Value class {value_class.__name__} must be a subclass of "
+                    f"{cls.__name__}"
+                )
+
+            # Check that value class inherits from parent value classes
+            for base in cls.__bases__:
+                if base is Config or not issubclass(base, Config):
+                    continue
+                parent_xpmtype = base.__getxpmtype__()
+                # Check if parent has an explicit value type (different from original)
+                if parent_xpmtype.value_type is not parent_xpmtype._original_type:
+                    parent_value = parent_xpmtype.value_type
+                    if not issubclass(value_class, parent_value):
+                        raise TypeError(
+                            f"Value class {value_class.__name__} must be a subclass of "
+                            f"parent value class {parent_value.__name__}"
+                        )
+
+            # Register the value class
+            xpmtype.set_value_type(value_class)
+            return value_class
+
+        return decorator
+
     @classmethod
     def __getxpmtype__(cls) -> "ObjectType":
         """Get (and create if necessary) the Object type associated
@@ -1497,6 +1936,40 @@ class Config:
             fp.flush()
 
 
+class InstanceConfig(Config):
+    """Base class for configurations where instance identity matters.
+
+    When a Config class derives from InstanceConfig instead of Config,
+    instances are distinguished based on their object identity when used
+    in containers. This enables distinguishing between shared and separate
+    instances even when all parameters are identical.
+
+    Example:
+        >>> class SubModel(InstanceConfig):
+        ...     value: Param[int] = 100
+        >>> class MainModel(Config):
+        ...     m1: Param[SubModel]
+        ...     m2: Param[SubModel]
+        >>>
+        >>> sm1 = SubModel.C()
+        >>> sm2 = SubModel.C()  # Same params, different instance
+        >>>
+        >>> # Shared instance (same object used twice)
+        >>> shared = MainModel.C(m1=sm1, m2=sm1)
+        >>>
+        >>> # Separate instances (different objects)
+        >>> separate = MainModel.C(m1=sm1, m2=sm2)
+        >>>
+        >>> # Different identifiers: shared vs separate
+        >>> shared.__identifier__() != separate.__identifier__()
+
+    The instance order is determined by the traversal order during
+    identifier computation, ensuring reproducibility.
+    """
+
+    pass
+
+
 class LightweightTask(Config):
     """A task that can be run before or after a real task to modify its behaviour"""
 
@@ -1559,11 +2032,55 @@ def copyconfig(config: Config, **kwargs):
 
 
 def setmeta(config: Config, flag: bool):
-    """Flags the configuration as a meta-parameter"""
+    """Force a configuration to be treated as a meta-parameter.
+
+    When a configuration is marked as meta, it is excluded from the
+    identifier computation of its parent configuration.
+
+    Example::
+
+        class Ensemble(Config):
+            model1: Param[Model]
+            model2: Param[Model]
+
+        # Mark model2 as meta - it won't affect the ensemble's identifier
+        model2 = setmeta(Model.C(...), True)
+        ensemble = Ensemble.C(model1=model1, model2=model2)
+
+    :param config: The configuration to mark
+    :param flag: True to mark as meta, False to include in identifier
+    :return: The same configuration (for chaining)
+    """
     config.__xpm__.set_meta(flag)
     return config
 
 
+class ResumableTask(Task):
+    """Base class for resumable/checkpointable tasks
+
+    Resumable tasks can be restarted if they are stopped by a time limit
+    (e.g., SLURM job timeout). The task directory and dynamic outputs are
+    preserved across restarts to allow checkpoint recovery.
+    """
+
+    def remaining_time(self) -> Optional[float]:
+        """Returns the remaining time in seconds before the job times out.
+
+        This is useful for checkpointing before hitting a time limit
+        (e.g., SLURM walltime).
+
+        Returns:
+            The remaining time in seconds, or None if:
+            - There is no time limit
+            - The launcher doesn't support querying remaining time
+            - The task is not running
+        """
+        launcher_info = taskglobals.Env.instance().launcher_info
+        if launcher_info is None:
+            return None
+        return launcher_info.remaining_time()
+
+
 def cache(fn, name: str):
     def __call__(config, *args, **kwargs):
         import experimaestro.taskglobals as taskglobals