experimaestro 0.22.0__zip → 0.24.0__zip

{experimaestro-0.22.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,12 +20,12 @@ import inspect
 import importlib
 from typing import (
     Any,
-    Callable,
     ClassVar,
     Dict,
     List,
     Optional,
     Set,
+    Tuple,
     Type,
     TypeVar,
     Union,
@@ -35,6 +41,7 @@ from experimaestro.core.types import DeprecatedAttribute, ObjectType
 from .context import SerializationContext, SerializedPath, SerializedPathLoader
 
 if TYPE_CHECKING:
+    from experimaestro.scheduler.base import Job
     from experimaestro.scheduler.workspace import RunMode
     from experimaestro.launchers import Launcher
     from experimaestro.scheduler import Workspace
@@ -167,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)
 
@@ -276,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)
@@ -307,7 +320,7 @@ def add_to_path(p):
         sys.path = old_path
 
 
-class GenerationContext:
+class ConfigWalkContext:
     """Context when generating values in configurations"""
 
     @property
@@ -342,33 +355,38 @@ class GenerationContext:
 NOT_SET = object()
 
 
-class ConfigProcessing:
+class ConfigWalk:
     """Allows to perform an operation on all nested configurations"""
 
-    def __init__(self, recurse_task=False):
+    def __init__(self, context: ConfigWalkContext = None, recurse_task=False):
         """
 
-        Parameters:
-            recurse_task: Recurse into linked tasks
+        :param recurse_task: Recurse into linked tasks
+        :param context: The context, by default only tracks the position in the
+            config tree
         """
         self.recurse_task = recurse_task
+        self.context = ConfigWalkContext() if context is None else context
 
         # Stores already visited nodes
         self.visited = {}
 
-    def preprocess(self, config: "Config"):
+    def preprocess(self, config: "Config") -> Tuple[bool, Any]:
+        """Returns a tuple boolean/value
+
+        The boolean value is used to stop the processing if False.
+        The value is returned
+        """
         return True, None
 
     def postprocess(self, config: "Config", values: Dict[str, Any]):
         return config
 
-    @contextmanager
     def list(self, i: int):
-        yield i
+        return self.context.push(str(i))
 
-    @contextmanager
     def map(self, k: str):
-        yield k
+        return self.context.push(k)
 
     def __call__(self, x):
         if isinstance(x, Config):
@@ -412,9 +430,13 @@ class ConfigProcessing:
                     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
@@ -423,19 +445,10 @@ class ConfigProcessing:
         if isinstance(x, (float, int, str, Path, Enum)):
             return x
 
-        raise NotImplementedError(f"Cannot handle a value of type {type(x)}")
-
-
-class GenerationConfigProcessing(ConfigProcessing):
-    def __init__(self, context: GenerationContext, recurse_task=False):
-        super().__init__(recurse_task=recurse_task)
-        self.context = context
-
-    def list(self, i: int):
-        return self.context.push(str(i))
+        if isinstance(x, Proxy):
+            return self(x.__unwrap__())
 
-    def map(self, k: str):
-        return self.context.push(k)
+        raise NotImplementedError(f"Cannot handle a value of type {type(x)}")
 
 
 def getqualattr(module, qualname):
@@ -453,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
@@ -464,6 +477,7 @@ class ConfigInformation:
         # Meta-informations
         self._tags = {}
         self._initinfo = ""
+        self.submit_hooks = set()
 
         # Generated task
         self._taskoutput = None
@@ -481,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
 
@@ -534,7 +549,7 @@ class ConfigInformation:
                 yield argument, self.values[argument.name]
 
     def tags(self):
-        class TagFinder(ConfigProcessing):
+        class TagFinder(ConfigWalk):
             def __init__(self):
                 super().__init__(recurse_task=True)
                 self.tags = {}
@@ -578,7 +593,7 @@ class ConfigInformation:
                     )
                     raise
 
-    def seal(self, context: GenerationContext):
+    def seal(self, context: ConfigWalkContext):
         """Seal the object, generating values when needed,
         before scheduling the associated job(s)
 
@@ -586,7 +601,7 @@ class ConfigInformation:
             - context: the generation context
         """
 
-        class Sealer(GenerationConfigProcessing):
+        class Sealer(ConfigWalk):
             def preprocess(self, config: Config):
                 return not config.__xpm__._sealed, config
 
@@ -607,9 +622,9 @@ class ConfigInformation:
 
         Internal API - do not use
         """
-        context = GenerationContext()
+        context = ConfigWalkContext()
 
-        class Unsealer(GenerationConfigProcessing):
+        class Unsealer(ConfigWalk):
             def preprocess(self, config: Config):
                 return config.__xpm__._sealed, config
 
@@ -658,13 +673,35 @@ class ConfigInformation:
                 logger.error("While setting %s", path + [argument.name])
                 raise
 
+    def apply_submit_hooks(self, job: "Job"):
+        """Apply configuration hooks"""
+        context = ConfigWalkContext()
+
+        class HookGatherer(ConfigWalk):
+            def __init__(self, *args, **kwargs):
+                super().__init__(*args, **kwargs)
+                self.hooks = set()
+
+            def postprocess(self, config: "Config", values: Dict[str, Any]):
+                self.hooks.update(config.__xpm__.submit_hooks)
+
+        gatherer = HookGatherer(context, recurse_task=False)
+        gatherer(self.pyobject)
+        for hook in gatherer.hooks:
+            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:
@@ -672,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
         )
@@ -696,7 +734,7 @@ class ConfigInformation:
             experiment.CURRENT.workspace if experiment.CURRENT else None
         )
 
-        # Call onSubmit
+        # Call onSubmit hooks
         launcher = (
             launcher
             or (workspace and workspace.launcher)
@@ -705,6 +743,9 @@ class ConfigInformation:
         if launcher:
             launcher.onSubmit(self.job)
 
+        # Apply submit hooks
+        self.apply_submit_hooks(self.job)
+
         # Add job dependencies
         self.updatedependencies(self.job.dependencies, [], set([id(self.pyobject)]))
 
@@ -743,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
 
@@ -794,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):
@@ -876,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__()
@@ -1024,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":
@@ -1086,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"]
@@ -1168,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__}"
@@ -1178,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:
@@ -1194,8 +1231,8 @@ class ConfigInformation:
 
         return o
 
-    class FromPython(GenerationConfigProcessing):
-        def __init__(self, context: GenerationContext):
+    class FromPython(ConfigWalk):
+        def __init__(self, context: ConfigWalkContext):
             super().__init__(context)
             self.objects = {}
 
@@ -1231,9 +1268,16 @@ 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: GenerationContext):
+    def fromConfig(self, context: ConfigWalkContext):
         """Generate an instance given the current configuration"""
         self.validate()
         processor = ConfigInformation.FromPython(context)
@@ -1378,7 +1422,7 @@ class TypeConfig:
         self.__xpm__.add_dependencies(*dependencies)
         return self
 
-    def instance(self, context: GenerationContext = None) -> T:
+    def instance(self, context: ConfigWalkContext = None) -> T:
         """Return an instance with the current values"""
         if context is None:
             from experimaestro.xpmutils import EmptyContext
@@ -1386,19 +1430,22 @@ class TypeConfig:
             context = EmptyContext()
         else:
             assert isinstance(
-                context, GenerationContext
-            ), f"{context.__class__} is not an instance of GenerationContext"
+                context, ConfigWalkContext
+            ), 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
@@ -1434,7 +1481,7 @@ class Config:
     configuration/task"""
 
     @classmethod
-    def __getxpmtype__(cls):
+    def __getxpmtype__(cls) -> "ObjectType":
         """Get (and create if necessary) the Object type of this"""
         xpmtype = cls.__dict__.get("__xpmtype__", None)
         if xpmtype is None:
@@ -1499,9 +1546,6 @@ class Task(Config):
         raise NotImplementedError()
 
 
-# --- Output proxy
-
-
 class Proxy:
     """A proxy for a value"""
 
@@ -1509,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
 
@@ -1520,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):
@@ -1573,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
@@ -1666,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)
@@ -1686,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):