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

experimaestro/core/objects/config_walk.py

@@ -29,6 +29,26 @@ class ConfigWalkContext:
             return self.path / self._configpath
         return self.path
 
+    def partial_path(self, subparameters, config) -> Path:
+        """Returns the partial directory path for a given subparameters instance.
+
+        This method should be overridden in subclasses that have access to
+        workspace information (like JobContext).
+
+        Args:
+            subparameters: The Subparameters instance defining which groups to exclude
+            config: The configuration to compute the partial identifier for
+
+        Returns:
+            The partial directory path.
+
+        Raises:
+            NotImplementedError: If the context doesn't support partial paths.
+        """
+        raise NotImplementedError(
+            "Partial paths require a context with workspace information (like JobContext)"
+        )
+
     @contextmanager
     def push(self, key: str):
         """Push a new key to contextualize paths"""

experimaestro/core/serialization.py

@@ -31,10 +31,14 @@ def json_object(context: SerializationContext, value: Any, objects=[]):
 
 
 def state_dict(context: SerializationContext, obj: Any):
-    """Returns a state dictionary of the object
+    """Convert an object to a state dictionary for serialization.
+
+    Returns a dictionary representation that can be serialized to JSON
+    and later restored with :func:`from_state_dict`.
 
     :param context: The serialization context
-    :param obj: the object to serialize
+    :param obj: The object to serialize
+    :return: A dictionary with 'objects' and 'data' keys
     """
     objects: list[Any] = []
     data = json_object(context, obj, objects)
@@ -48,14 +52,19 @@ def save_definition(obj: Any, context: SerializationContext, path: Path):
 
 
 def save(obj: Any, save_directory: Optional[Path]):
-    """Saves an object into a disk file
+    """Save a configuration to a directory.
+
+    The serialization process stores the configuration in "definition.json"
+    and copies any files or folders registered as DataPath parameters.
+
+    Example::
 
-    The serialization process also stores in the given folder the different
-    files or folders that are registered as Path parameters (or
-    meta-parameters).
+        config = MyConfig.C(data_path=Path("/data/file.txt"))
+        save(config, Path("/output/saved_config"))
 
+    :param obj: The configuration to save
     :param save_directory: The directory in which the object and its data will
-        be saved (by default, the object is saved in "definition.json")
+        be saved (object is saved in "definition.json")
     """
     context = SerializationContext(save_directory=save_directory)
     save_definition(
@@ -89,19 +98,32 @@ def from_state_dict(
     state: Dict[str, Any],
     path: Union[None, str, Path, SerializedPathLoader] = None,
     *,
-    as_instance: bool = False
+    as_instance: bool = False,
+    partial_loading: Optional[bool] = None,
 ):
-    """Load an object from a state dictionary
-
-    :param state: The state
-    :param path: A directory or a function that transforms relative file path
-        into absolute ones
-    :param as_instance: returns instances instead of configuration objects
+    """Load an object from a state dictionary.
+
+    Restores a configuration from a dictionary previously created by
+    :func:`state_dict`.
+
+    :param state: The state dictionary to load from
+    :param path: Directory containing data files, or a function that resolves
+        relative paths to absolute ones
+    :param as_instance: If True, return an instance instead of a config
+    :param partial_loading: If True, skip loading task references. If None
+        (default), partial_loading is enabled when as_instance is True.
+    :return: The loaded configuration or instance
     """
+    # Determine effective partial_loading: as_instance implies partial_loading
+    effective_partial_loading = (
+        partial_loading if partial_loading is not None else as_instance
+    )
+
     objects = ConfigInformation.load_objects(
         state["objects"],
         as_instance=as_instance,
         data_loader=get_data_loader(path),
+        partial_loading=effective_partial_loading,
     )
 
     return ConfigInformation._objectFromParameters(state["data"], objects)
@@ -110,46 +132,72 @@ def from_state_dict(
 def load(
     path: Union[str, Path, SerializedPathLoader],
     as_instance: bool = False,
+    partial_loading: Optional[bool] = None,
 ) -> Tuple[Any, List["LightweightTask"]]:
-    """Load data from disk
+    """Load a configuration from a directory.
+
+    Restores a configuration previously saved with :func:`save`.
+
+    Example::
 
-    :param path: A directory or a function that transforms relative file path
-        into absolute ones
-    :param as_instance: returns instances instead of configuration objects
+        config = load(Path("/output/saved_config"))
+
+    :param path: Directory containing the saved configuration, or a function
+        that resolves relative paths to absolute ones
+    :param as_instance: If True, return an instance instead of a config
+    :param partial_loading: If True, skip loading task references. If None
+        (default), partial_loading is enabled when as_instance is True.
+    :return: The loaded configuration or instance
     """
     data_loader = get_data_loader(path)
 
     with data_loader("definition.json").open("rt") as fh:
         content = json.load(fh)
-    return from_state_dict(content, as_instance=as_instance)
+    return from_state_dict(
+        content, as_instance=as_instance, partial_loading=partial_loading
+    )
 
 
 def from_task_dir(
     path: Union[str, Path, SerializedPathLoader],
     as_instance: bool = False,
+    partial_loading: Optional[bool] = None,
 ):
-    """Loads a task object"""
+    """Load a task configuration from a task directory.
+
+    Loads the task parameters from a job directory (containing params.json).
+    This is useful for reloading task configurations after execution.
+
+    :param path: Task directory containing params.json, or a function that
+        resolves relative paths to absolute ones
+    :param as_instance: If True, return an instance instead of a config
+    :param partial_loading: If True, skip loading task references. If None
+        (default), partial_loading is enabled when as_instance is True.
+    :return: The loaded task configuration or instance
+    """
     data_loader = get_data_loader(path)
     with data_loader("params.json").open("rt") as fh:
         content = json.load(fh)
 
     content["data"] = {"type": "python", "value": content["objects"][-1]["id"]}
 
-    return from_state_dict(content, as_instance=as_instance)
+    return from_state_dict(
+        content, as_instance=as_instance, partial_loading=partial_loading
+    )
 
 
 def serialize(
     obj: Any, save_directory: Path, *, init_tasks: list["LightweightTask"] = []
 ):
-    """Saves an object into a disk file, including initialization tasks
+    """Serialize a configuration to a directory with initialization tasks.
 
-    The serialization process also stores in the given folder the different
-    files or folders that are registered as Path parameters (or
-    meta-parameters).
+    Similar to :func:`save`, but also stores lightweight initialization tasks
+    that should be run when the configuration is deserialized.
 
+    :param obj: The configuration to serialize
     :param save_directory: The directory in which the object and its data will
-        be saved (by default, the object is saved in "definition.json")
-    :param init_tasks: The optional
+        be saved (object is saved in "definition.json")
+    :param init_tasks: List of lightweight tasks to run on deserialization
     """
     context = SerializationContext(save_directory=save_directory)
     save_definition((obj, init_tasks), context, save_directory / "definition.json")
@@ -158,20 +206,29 @@ def serialize(
 def deserialize(
     path: Union[str, Path, SerializedPathLoader],
     as_instance: bool = False,
+    partial_loading: Optional[bool] = None,
 ) -> tuple[Any, List["LightweightTask"]] | Any:
-    """Load data from disk, and initialize the object
-
-    :param path: A directory or a function that transforms relative file path
-        into absolute ones
-    :param as_instance: returns instances instead of configuration objects
-    :returns: either the object (as_instance is true), or a tuple
+    """Deserialize a configuration from a directory.
+
+    Restores a configuration previously saved with :func:`serialize`.
+    When ``as_instance=True``, runs any stored initialization tasks.
+
+    :param path: Directory containing the serialized configuration, or a function
+        that resolves relative paths to absolute ones
+    :param as_instance: If True, return an instance and run init tasks
+    :param partial_loading: If True, skip loading task references. If None
+        (default), partial_loading is enabled when as_instance is True.
+    :return: The configuration/instance (if as_instance), or tuple of
+        (configuration, init_tasks)
     """
     data_loader = get_data_loader(path)
 
     with data_loader("definition.json").open("rt") as fh:
         content = json.load(fh)
 
-    object, init_tasks = from_state_dict(content, data_loader, as_instance=as_instance)
+    object, init_tasks = from_state_dict(
+        content, data_loader, as_instance=as_instance, partial_loading=partial_loading
+    )
 
     if as_instance:
         for init_task in init_tasks:

experimaestro/core/subparameters.py

@@ -0,0 +1,164 @@
+"""Subparameters for partial identifier computation.
+
+This module provides the `subparameters` function and `Subparameters` class
+for defining parameter subsets that compute partial identifiers. This enables
+sharing directories (like checkpoints) across tasks that differ only in
+excluded parameter groups.
+
+Example:
+    iter_group = param_group("iter")
+
+    class Learn(Task):
+        checkpoints = subparameters(exclude_groups=[iter_group])
+
+        max_iter: Param[int] = field(groups=[iter_group])
+        learning_rate: Param[float]
+
+        # Path will be in WORKSPACE/partials/TASK_ID/checkpoints/PARTIAL_ID/
+        checkpoints_path: Param[Path] = field(
+            default_factory=PathGenerator(partial=checkpoints)
+        )
+"""
+
+from dataclasses import dataclass, field as dataclass_field
+from typing import Set, Optional
+
+
+@dataclass(frozen=True)
+class ParameterGroup:
+    """A parameter group with a name
+
+    The name is just for reference, what is important is the identity of the
+    object. The dataclass is frozen to make it hashable.
+    """
+
+    name: str
+
+
+def param_group(name: str) -> ParameterGroup:
+    """Create a parameter group for use with subparameters.
+
+    Parameter groups allow computing partial identifiers that exclude
+    certain parameters, enabling shared directories across related tasks.
+
+    Example::
+
+        training_group = param_group("training")
+
+        class MyTask(Task):
+            model_size: Param[int]
+            learning_rate: Param[float] = field(groups=[training_group])
+
+    :param name: Unique name for this parameter group
+    :return: A ParameterGroup object
+    """
+    return ParameterGroup(name)
+
+
+@dataclass
+class Subparameters:
+    """Defines a subset of parameters for partial identifier computation.
+
+    A Subparameters instance defines which parameter groups to include or exclude
+    when computing a partial identifier. This enables sharing directories
+    (like checkpoints) across experiments that only differ in excluded groups.
+
+    The inclusion/exclusion logic follows these rules:
+    1. If `exclude_all` is True, all parameters are excluded by default
+    2. Parameters in `exclude_groups` are excluded
+    3. Parameters with no group are excluded if `exclude_no_group` is True
+    4. Parameters in `include_groups` are always included (overrides exclusion)
+
+    Attributes:
+        exclude_groups: Set of group names to exclude from identifier computation
+        include_groups: Set of group names to always include (overrides exclusion)
+        exclude_no_group: If True, exclude parameters with no group assigned
+        exclude_all: If True, exclude all parameters by default
+        name: The name of this parameter set (auto-set from class attribute name)
+    """
+
+    #: Set of group names to exclude from identifier computation
+    exclude_groups: Set[ParameterGroup] = dataclass_field(default_factory=set)
+
+    #: Set of group names to always include (overrides exclusion)
+    include_groups: Set[ParameterGroup] = dataclass_field(default_factory=set)
+
+    #: If True, exclude parameters with no group assigned
+    exclude_no_group: bool = False
+
+    #: If True, exclude all parameters by default (use include_groups to select)
+    exclude_all: bool = False
+
+    #: Name of this parameter set (auto-set from class attribute)
+    name: Optional[ParameterGroup] = None
+
+    def __post_init__(self):
+        # Ensure groups are sets
+        if not isinstance(self.exclude_groups, set):
+            self.exclude_groups = set(self.exclude_groups)
+        if not isinstance(self.include_groups, set):
+            self.include_groups = set(self.include_groups)
+
+    def is_excluded(self, groups: Set[ParameterGroup]) -> bool:
+        """Check if a parameter with the given groups should be excluded.
+
+        Args:
+            groups: The set of groups the parameter belongs to (empty if no groups).
+
+        Returns:
+            True if the parameter should be excluded from partial identifier.
+        """
+        # Include always overrides exclude - if any group is in include_groups
+        if groups and (groups & self.include_groups):
+            return False
+
+        # Check exclusion rules
+        if self.exclude_all:
+            return True
+        if not groups and self.exclude_no_group:
+            return True
+        if groups and (groups & self.exclude_groups):
+            return True
+
+        return False
+
+
+def subparameters(
+    *,
+    exclude_groups: list[ParameterGroup] | None = None,
+    include_groups: list[ParameterGroup] | None = None,
+    exclude_no_group: bool = False,
+    exclude_all: bool = False,
+) -> Subparameters:
+    """Create a subparameters specification for partial identifier computation.
+
+    Subparameters allow tasks to share directories when they differ only
+    in certain parameter groups (e.g., training hyperparameters).
+
+    Example::
+
+        training_group = param_group("training")
+
+        class Train(Task):
+            model: Param[Model]
+            epochs: Param[int] = field(groups=[training_group])
+
+            checkpoint: Meta[Path] = field(
+                default_factory=PathGenerator(
+                    "model.pt",
+                    subparameters=subparameters(exclude=[training_group])
+                )
+            )
+
+    :param exclude_groups: Parameter groups to exclude from identifier
+    :param include_groups: Parameter groups to always include (overrides exclusion)
+    :param exclude_no_group: If True, exclude parameters with no group assigned
+    :param exclude_all: If True, exclude all parameters by default
+    :return: A Subparameters object
+    """
+    return Subparameters(
+        exclude_groups=set(exclude_groups or []),
+        include_groups=set(include_groups or []),
+        exclude_no_group=exclude_no_group,
+        exclude_all=exclude_all,
+    )