experimaestro 0.23.0__zip → 0.24.0__zip

{experimaestro-0.23.0 → experimaestro-0.24.0}/src/experimaestro/core/objects.py

@@ -2,6 +2,12 @@
 
 from functools import cached_property
 import json
+
+try:
+    from types import NoneType
+except Exception:
+    # compatibility: python-3.8
+    NoneType = type(None)
 from termcolor import cprint
 import os
 from pathlib import Path
@@ -14,7 +20,6 @@ import inspect
 import importlib
 from typing import (
     Any,
-    Callable,
     ClassVar,
     Dict,
     List,
@@ -169,7 +174,7 @@ class HashComputer:
             for key, value in items:
                 self.update(key, subparam=subparam)
                 self.update(value, subparam=subparam)
-        elif isinstance(value, TaskOutput):
+        elif isinstance(value, ConfigWrapper):
             # Add the task ID...
             self.update(value.__xpm__.task, subparam=subparam)
 
@@ -278,8 +283,14 @@ def updatedependencies(
     elif isinstance(value, (list, set)):
         for el in value:
             updatedependencies(dependencies, el, path, taskids)
-    elif isinstance(value, TaskOutput):
-        dependencies.add(value.__xpm__.task.__xpm__.dependency())
+    elif isinstance(value, ConfigWrapper):
+        # Add the base value (if any)
+        if value.__xpm__.base is not None:
+            value.__xpm__.base.__xpm__.updatedependencies(dependencies, path, taskids)
+
+        # Add the task (if any)
+        if value.__xpm__.task is not None:
+            dependencies.add(value.__xpm__.task.__xpm__.dependency())
     elif isinstance(value, (dict,)):
         for key, val in value.items():
             updatedependencies(dependencies, key, path, taskids)
@@ -419,9 +430,13 @@ class ConfigWalk:
                     result[key] = self(value)
             return result
 
-        if isinstance(x, TaskOutput):
+        if isinstance(x, ConfigWrapper):
             # Process task if different
-            if self.recurse_task and x.__xpm__.task is not x.__unwrap__():
+            if (
+                x.__xpm__.task is not None
+                and self.recurse_task
+                and x.__xpm__.task is not x.__unwrap__()
+            ):
                 self(x.__xpm__.task)
 
             # Processed the wrapped config
@@ -430,6 +445,9 @@ class ConfigWalk:
         if isinstance(x, (float, int, str, Path, Enum)):
             return x
 
+        if isinstance(x, Proxy):
+            return self(x.__unwrap__())
+
         raise NotImplementedError(f"Cannot handle a value of type {type(x)}")
 
 
@@ -448,7 +466,7 @@ class ConfigInformation:
     """Forces this configuration to be a meta-parameter"""
 
     # Set to true when loading from JSON
-    LOADING = False
+    LOADING: ClassVar[bool] = False
 
     def __init__(self, pyobject: "TypeConfig"):
         # The underlying pyobject and XPM type
@@ -477,6 +495,7 @@ class ConfigInformation:
         self._meta = None
 
     def set_meta(self, value: Optional[bool]):
+        """Sets the meta flag"""
         assert not self._sealed, "Configuration is sealed"
         self._meta = value
 
@@ -672,12 +691,17 @@ class ConfigInformation:
             hook(job)
 
     def submit(
-        self, workspace: "Workspace", launcher: "Launcher", run_mode=None
-    ) -> "TaskOutput":
+        self, workspace: "Workspace", launcher: "Launcher", *, run_mode=None, pre=None
+    ) -> "ConfigWrapper":
         from experimaestro.scheduler import experiment, JobContext
         from experimaestro.scheduler.workspace import RunMode
 
+        # --- Handle default values
+
+        pre = pre or []
+
         # --- Prepare the object
+
         if self.job:
             raise Exception("task %s was already submitted" % self)
         if not self.xpmtype.task:
@@ -685,6 +709,7 @@ class ConfigInformation:
 
         # --- Submit the job
 
+        # Creates a new job
         self.job = self.xpmtype.task(
             self.pyobject, launcher=launcher, workspace=workspace, run_mode=run_mode
         )
@@ -759,16 +784,20 @@ class ConfigInformation:
                 hints = get_type_hints(self.pyobject.config)
                 config = hints["return"](**config)
             config.__xpm__.validate()
-            self._taskoutput = TaskOutput(config, self.pyobject)
+            self._taskoutput = ConfigWrapper.__create_taskoutput__(
+                config, self.pyobject
+            )
 
         # New way to handle outputs
         elif hasattr(self.pyobject, "taskoutputs"):
             value = self.pyobject.taskoutputs()
-            self._taskoutput = TaskOutput(value, self.pyobject)
+            self._taskoutput = ConfigWrapper.__create_taskoutput__(value, self.pyobject)
 
         # Otherwise, the output is just the config
         else:
-            self._taskoutput = TaskOutput(self.pyobject, self.pyobject)
+            self._taskoutput = ConfigWrapper.__create_taskoutput__(
+                self.pyobject, self.pyobject
+            )
 
         return self._taskoutput
 
@@ -810,20 +839,13 @@ class ConfigInformation:
                 "value": value.name,
             }
 
-        elif isinstance(value, SerializedTaskOutput):
-            # Reference to a serialized object
-            return {
-                "type": "serialized",
-                "value": id(value.__xpm__.serialized.loader),
-                "path": [c.toJSON() for c in value.__xpm__.path],
-            }
-
-        elif isinstance(value, TaskOutput):
+        elif isinstance(value, ConfigWrapper):
             return {
                 "type": "python",
                 "value": id(value.__unwrap__()),
                 # We add the task for identifier computation
                 "task": id(value.__xpm__.task),
+                "path": [c.toJSON() for c in value.__xpm__.path or []],
             }
 
         elif isinstance(value, Config):
@@ -892,25 +914,18 @@ class ConfigInformation:
     def __collect_objects__(value, objects: List[Dict], context: SerializationContext):
         """Serialize all needed configuration objects, looking at sub
         configurations if necessary"""
-        # objects
-        if isinstance(value, SerializedTaskOutput):
-            loader = value.__xpm__.serialized.loader
-            if id(loader) not in context.serialized:
-                objects.append(
-                    {
-                        "id": id(loader),
-                        "serialized": True,
-                        "module": loader.__class__.__module__,
-                        "type": loader.__class__.__qualname__,
-                        "value": loader.toJSON(),
-                    }
-                )
-            return
-
         # Unwrap if needed
-        if isinstance(value, TaskOutput):
+        if isinstance(value, ConfigWrapper):
             # We will need to output the task configuration objects
-            ConfigInformation.__collect_objects__(value.__xpm__.task, objects, context)
+            if value.__xpm__.task is not None:
+                ConfigInformation.__collect_objects__(
+                    value.__xpm__.task, objects, context
+                )
+
+            if value.__xpm__.base is not None:
+                ConfigInformation.__collect_objects__(
+                    value.__xpm__.base, objects, context
+                )
 
             # Unwrap the value to output it
             value = value.__unwrap__()
@@ -1040,7 +1055,7 @@ class ConfigInformation:
                 if not as_instance:
                     if task_id := value.get("task", None):
                         task = objects[task_id]
-                        return TaskOutput(obj, task)
+                        return ConfigWrapper.__create_taskoutput__(obj, task)
                 return obj
 
             if value["type"] == "serialized":
@@ -1102,6 +1117,7 @@ class ConfigInformation:
         o = None
         objects = {}
         import experimaestro.taskglobals as taskglobals
+        from .serializers import SerializedConfig
 
         for definition in definitions:
             module_name = definition["module"]
@@ -1184,7 +1200,10 @@ class ConfigInformation:
                             assert isinstance(v, Path), "Excepted Path, got {type(v)}"
 
                     if as_instance:
+                        # Unwrap the value if needed
+                        v = unwrap(v)
                         setattr(o, name, v)
+
                         assert (
                             getattr(o, name) is v
                         ), f"Problem with deserialization {name} of {o.__class__}"
@@ -1194,6 +1213,8 @@ class ConfigInformation:
                 if as_instance:
                     # Calls post-init
                     o.__post_init__()
+                    if isinstance(o, SerializedConfig):
+                        o.initialize()
                 else:
                     # Seal and set the identifier
                     if not discard_id:
@@ -1247,6 +1268,13 @@ class ConfigInformation:
             # Call __post_init__
             o.__post_init__()
 
+            # Process a serialized configuration
+            from .serializers import SerializedConfig
+
+            if isinstance(o, SerializedConfig):
+                o.initialize()
+                o = o.__unwrap__()
+
             return o
 
     def fromConfig(self, context: ConfigWalkContext):
@@ -1406,15 +1434,18 @@ class TypeConfig:
             ), f"{context.__class__} is not an instance of ConfigWalkContext"
         return self.__xpm__.fromConfig(context)  # type: ignore
 
-    def submit(self, *, workspace=None, launcher=None, run_mode: "RunMode" = None):
+    def submit(
+        self, *, workspace=None, launcher=None, run_mode: "RunMode" = None, pre=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)
-        :return: a :py:class:TaskOutput object
+        :param pre: Pre-tasks to execute before
+        :return: a :py:class:ConfigWrapper object
         """
-        return self.__xpm__.submit(workspace, launcher, run_mode=run_mode)
+        return self.__xpm__.submit(workspace, launcher, run_mode=run_mode, pre=pre)
 
     def stdout(self):
         return self.__xpm__.job.stdout
@@ -1515,9 +1546,6 @@ class Task(Config):
         raise NotImplementedError()
 
 
-# --- Output proxy
-
-
 class Proxy:
     """A proxy for a value"""
 
@@ -1525,7 +1553,29 @@ class Proxy:
         raise NotImplementedError()
 
 
+def unwrap(v: Any):
+    """Unwrap all proxies"""
+    while isinstance(v, Proxy):
+        v = v.__unwrap__()
+    return v
+
+
+class AttrAccessor:
+    """Access an attribute"""
+
+    def __init__(self, key: str):
+        self.key = key
+
+    def get(self, value):
+        return getattr(value, self.key)
+
+    def toJSON(self):
+        return {"type": "attr", "name": self.key}
+
+
 class ItemAccessor:
+    """Access an array item"""
+
     def __init__(self, key: Any):
         self.key = key
 
@@ -1536,49 +1586,65 @@ class ItemAccessor:
         return value.__getitem__(self.key)
 
 
-class AttrAccessor:
-    def __init__(self, key: Any, default: Any):
-        self.key = key
-        self.default = default
-
-    def get(self, value):
-        return getattr(value, self.key, self.default)
-
-    def toJSON(self):
-        return {"type": "attr", "name": self.key}
+class ConfigWrapperInfo:
+    """Global information about the Configuration wrapper"""
 
+    def __init__(
+        self,
+        value: Any,
+        *,
+        task: Optional[Task] = None,
+        parent: Optional["ConfigWrapperInfo"] = None,
+        base: Optional[Config] = None,
+        path: Optional[List[Path]] = None,
+    ):
+        # Current value
+        self.value = value
+        self.parent = parent
 
-class Serialized:
-    """Simple serialization object"""
+        # The task
+        if isinstance(value, Task):
+            self.task = value
+        else:
+            self.task = task
 
-    def __init__(self, value):
-        self.value = value
+        # Holds serialized config information
+        from .serializers import SerializedConfig
 
-    def toJSON(self):
-        return self.value
+        if isinstance(value, SerializedConfig):
+            self.base = value
+            self.value = value.config
+            self.path = [AttrAccessor("config")]
+        else:
+            self.base = base
+            self.path = path
 
+    def __state_dict__(self):
+        return {"task": self.task, "base": self.base, "path": self.path}
 
-class SerializedConfig:
-    """A serializable configuration
+    def __getitem__(self, key: Any):
+        kwargs = self.__state_dict__()
+        value = self.value.__getitem__(key)
 
-    This can be used to define a loading mechanism when instanciating the
-    configuration
-    """
+        if self.path:
+            kwargs["path"].append(ItemAccessor(key))
 
-    pyobject: Config
-    """The configuration that will be serialized"""
+        return ConfigWrapperInfo(value, **kwargs)
 
-    def __init__(self, pyobject: Config, loader: Callable[[Path], Config]):
-        self.pyobject = pyobject
-        self.loader = loader
+    def __getattr__(self, key: str, *default) -> Any:
+        kwargs = self.__state_dict__()
+        value = getattr(self.value, key, *default)
+        if self.path:
+            kwargs["path"].append(AttrAccessor(key, *default))
+        return ConfigWrapperInfo(value, parent=self, **kwargs)
 
+    def wrap(self):
+        """Wrap a value if needed"""
+        if isinstance(self.value, (str, int, float, Path, bool, NoneType)):
+            return self.value
 
-class TaskOutputInfo:
-    def __init__(self, task: Task):
-        self.task = task
-        self.value = None
-        self.path = None
-        self.serialized = None
+        # Returns a new config wrapper
+        return ConfigWrapper(self)
 
     @property
     def identifier(self):
@@ -1589,92 +1655,64 @@ class TaskOutputInfo:
         return self.task.__xpm__.job
 
     def tags(self):
+        """Returns the tags of the task"""
         tags = self.task.__xpm__.tags()
         return tags
 
     def stdout(self):
+        """Returns the standard output of the associated task"""
         return self.task.__xpm__.job.stdout
 
     def stderr(self):
+        """Returns the standard error of the associated task"""
         return self.task.__xpm__.job.stderr
 
     def wait(self):
+        """Wait for the task to end
+
+        :return: True if the task completed without error
+        """
         from experimaestro.scheduler import JobState
 
         return self.task.__xpm__.job.wait() == JobState.DONE
 
 
-class TaskOutput(Proxy):
+# Cleanup into just one
+class ConfigWrapper(Proxy):
     """Task proxy
 
     This is used when accessing properties *after* having submitted a task,
-    to keep track of the dependencies
+    to keep track of the dependencies, and/or as an accessor when dealing with
+    a serialized config
     """
 
-    def __init__(self, value: Any, task: Union[Task, TaskOutputInfo]):
-        self.__xpm__ = (
-            task if isinstance(task, TaskOutputInfo) else TaskOutputInfo(task)
-        )
-        self.__xpm__.value = value
-
-    def _wrap(self, value):
-        if isinstance(value, SerializedConfig):
-            return SerializedTaskOutput(value.pyobject, value, self.__xpm__.task, [])
-
-        if isinstance(value, (str, int, float, Path, bool)):
-            # No need to wrap if direct
-            return value
+    def __init__(self, info: ConfigWrapperInfo):
+        self.__xpm__ = info
 
-        return TaskOutput(value, self.__xpm__.task)
+    @staticmethod
+    def __create_taskoutput__(value: Any, task: Task):
+        return ConfigWrapperInfo(value, task=task).wrap()
 
     def __getitem__(self, key: Any):
-        return self._wrap(self.__xpm__.value.__getitem__(key))
+        return self.__xpm__[key].wrap()
 
-    def __getattr__(self, key: str, default=None) -> Any:
-        return self._wrap(getattr(self.__xpm__.value, key, default))
+    def __getattr__(self, key: str, *default) -> Any:
+        return self.__xpm__.__getattr__(key, *default).wrap()
 
     def __unwrap__(self):
         return self.__xpm__.value
 
     def __call__(self, *args, **kwargs):
         assert callable(self.__xpm__.value), "Attribute is not a function"
-        __self__ = TaskOutput(self.__xpm__.value.__self__, self.__xpm__.task)
-        return self.__xpm__.value.__func__(__self__, *args, **kwargs)
-
 
-class SerializedTaskOutput(TaskOutput):
-    """Used when serializing a configuration
-
-    Here, we need to keep track of the path to the value we need
-    """
-
-    def __init__(
-        self, value, serialized: SerializedConfig, task: Task, path: List[Any]
-    ):
-        super().__init__(value, task)
-        self.__xpm__.serialized = serialized
-        self.__xpm__.path = path
-
-    def __getitem__(self, key: Any):
-        value = self.__xpm__.value.__getitem__(key)
-        return SerializedTaskOutput(
-            value, self.serialized, self.__xpm__.task, self.path + [ItemAccessor(key)]
-        )
-
-    def __getattr__(self, key: str, default=None) -> Any:
-        value = getattr(self.__xpm__.value, key, default)
-        return SerializedTaskOutput(
-            value,
-            self.__xpm__.serialized,
-            self.__xpm__.task,
-            self.__xpm__.path + [AttrAccessor(key, default)],
-        )
+        __self__ = self.__xpm__.parent.wrap()
+        return self.__xpm__.value.__func__(__self__, *args, **kwargs)
 
 
 # --- Utility functions
 
 
-def copyconfig(config_or_output: Union[Config, TaskOutput], **kwargs):
+def copyconfig(config_or_output: Union[Config, ConfigWrapper], **kwargs):
     """Copy a configuration or task output
 
     Useful to modify a configuration that can be potentially
@@ -1682,7 +1720,7 @@ def copyconfig(config_or_output: Union[Config, TaskOutput], **kwargs):
     a task output).
     """
 
-    if isinstance(config_or_output, TaskOutput):
+    if isinstance(config_or_output, ConfigWrapper):
         output = config_or_output
         config = config_or_output.__unwrap__()
         assert isinstance(config, Config)
@@ -1702,7 +1740,7 @@ def copyconfig(config_or_output: Union[Config, TaskOutput], **kwargs):
         return copy
 
     # wrap in Task output
-    return TaskOutput(copy, output.__xpm__)
+    return ConfigWrapper(copy, output.__xpm__)
 
 
 def setmeta(config: Config, flag: bool):

{experimaestro-0.23.0 → experimaestro-0.24.0}/src/experimaestro/core/objects.pyi

@@ -181,7 +181,8 @@ class TypeConfig:
         *,
         workspace: Incomplete | None = ...,
         launcher: Incomplete | None = ...,
-        run_mode: RunMode = ...
+        run_mode: RunMode = ...,
+        pre: List[Task] = []
     ): ...
     def stdout(self): ...
     def stderr(self): ...

experimaestro-0.24.0/src/experimaestro/core/serializers.py

@@ -0,0 +1,52 @@
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import List, TypeVar
+
+from experimaestro import Param
+
+from .objects import Config, Proxy
+from .arguments import DataPath
+
+T = TypeVar("T")
+
+
+class SerializedConfig(Config, Proxy, ABC):
+    """A serializable configuration
+
+    This can be used to define a loading mechanism when instanciating the
+    configuration
+    """
+
+    config: Param[Config]
+    """The configuration that will be serialized"""
+
+    registered: List[Config]
+    """(execution only) List of configurations that use this serialized config"""
+
+    def __post_init__(self):
+        super().__post_init__()
+        self.registered = []
+
+    def register(self, config: Config):
+        self.registered.append(config)
+
+    def __unwrap__(self):
+        return self.config
+
+    @abstractmethod
+    def initialize(self):
+        """Initialize the object
+
+        This might imply loading saved data (e.g. learned models)
+        """
+        ...
+
+
+class PathBasedSerializedConfig(SerializedConfig):
+    """A path based serialized configuration
+
+    The most common case it to have
+    """
+
+    path: DataPath[Path]
+    """Path containing the data"""

{experimaestro-0.23.0 → experimaestro-0.24.0}/src/experimaestro/huggingface.py

@@ -1,6 +1,6 @@
 from pathlib import Path
 from typing import Optional, Union
-from experimaestro import Config, TaskOutput
+from experimaestro import Config, ConfigWrapper
 from experimaestro.core.context import SerializedPath
 from experimaestro.core.objects import ConfigInformation
 from huggingface_hub import ModelHubMixin, hf_hub_download, snapshot_download
@@ -11,7 +11,7 @@ class ExperimaestroHFHub(ModelHubMixin):
     """Defines models that can be uploaded/downloaded from the Hub"""
 
     def __init__(
-        self, config: Union[Config, TaskOutput], variant: Optional[str] = None
+        self, config: Union[Config, ConfigWrapper], variant: Optional[str] = None
     ):
         self.config = config if isinstance(config, Config) else config.__unwrap__()
         self.variant = variant

{experimaestro-0.23.0 → experimaestro-0.24.0}/src/experimaestro/tests/test_identifier.py

@@ -17,7 +17,8 @@ from experimaestro import (
     Annotated,
     Task,
 )
-from experimaestro.core.objects import ConfigInformation, TaskOutput, setmeta
+from experimaestro.core.objects import ConfigInformation, ConfigWrapper, setmeta
+from experimaestro.core.serializers import SerializedConfig
 from experimaestro.scheduler.workspace import RunMode
 
 
@@ -56,7 +57,7 @@ class Values:
 
 
 def getidentifier(x):
-    if isinstance(x, TaskOutput):
+    if isinstance(x, ConfigWrapper):
         return x.__xpm__.identifier.all
     return x.__xpm__.identifier.all
 
@@ -387,6 +388,32 @@ def test_identifier_meta_default_array():
     )
 
 
+# --- Check ConfigWrapper
+
+
+class Model(Config):
+    def __post_init__(self):
+        self.initialized = False
+
+
+class Trainer(Config):
+    model: Param[Model]
+
+
+class SerializedModel(SerializedConfig):
+    def initialize(self):
+        self.config.initialized = True
+
+
+def test_identifier_serialized_config():
+    trainer1 = Trainer(model=Model())
+    trainer2 = Trainer(model=SerializedModel(config=Model()))
+    assert_notequal(trainer1, trainer2)
+
+
+# --- Check configuration reloads
+
+
 def check_reload(config):
     old_identifier = config.__xpm__.identifier.all
 
@@ -412,14 +439,14 @@ def test_identifier_reload_config():
     check_reload(IdentifierReloadConfig(id="123"))
 
 
-class IdentifierReloadTaskOutput(Task):
+class IdentifierReloadConfigWrapper(Task):
     id: Param[str]
 
     def taskoutputs(self):
         return IdentifierReloadConfig(id=self.id)
 
 
-class IdentifierReloadTaskOutputDerived(Config):
+class IdentifierReloadConfigWrapperDerived(Config):
     task: Param[IdentifierReloadConfig]
 
 
@@ -427,8 +454,8 @@ def test_identifier_reload_taskoutput():
     """When using a task output, the identifier should not be different"""
 
     # Creates the configuration
-    task = IdentifierReloadTaskOutput(id="123").submit(run_mode=RunMode.DRY_RUN)
-    config = IdentifierReloadTaskOutputDerived(task=task)
+    task = IdentifierReloadConfigWrapper(id="123").submit(run_mode=RunMode.DRY_RUN)
+    config = IdentifierReloadConfigWrapperDerived(task=task)
     check_reload(config)