experimaestro 0.22.0__zip → 0.24.0__zip

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

@@ -1,8 +1,10 @@
+from abc import ABC
 import typing_extensions
 
 from experimaestro.core.types import ObjectType
 import experimaestro
 import io
+from experimaestro.launchers import Launcher
 from experimaestro.scheduler.base import Job
 
 from experimaestro.scheduler.workspace import RunMode
@@ -20,7 +22,18 @@ from experimaestro.core.types import (
 from experimaestro.utils import logger as logger
 from functools import cached_property as cached_property
 from pathlib import Path
-from typing import Any, ClassVar, Dict, List, Optional, Set, TypeVar, Union, overload
+from typing import (
+    Any,
+    Callable,
+    ClassVar,
+    Dict,
+    List,
+    Optional,
+    Set,
+    TypeVar,
+    Union,
+    overload,
+)
 
 T = TypeVar("T", bound="Config")
 
@@ -62,7 +75,7 @@ class TaggedValue:
 
 def add_to_path(p) -> Generator[None, None, None]: ...
 
-class GenerationContext:
+class ConfigWalkContext:
     @property
     def path(self) -> None: ...
     def __init__(self) -> None: ...
@@ -83,9 +96,9 @@ class ConfigProcessing:
     def map(self, k: str): ...
     def __call__(self, x): ...
 
-class GenerationConfigProcessing(ConfigProcessing):
+class ConfigWalk(ConfigProcessing):
     context: Incomplete
-    def __init__(self, context: GenerationContext) -> None: ...
+    def __init__(self, context: ConfigWalkContext) -> None: ...
     def list(self, i: int): ...
     def map(self, k: str): ...
 
@@ -108,7 +121,7 @@ class ConfigInformation:
     def xpmvalues(self, generated: bool = ...) -> Generator[Incomplete, None, None]: ...
     def tags(self): ...
     def validate(self) -> None: ...
-    def seal(self, context: GenerationContext): ...
+    def seal(self, context: ConfigWalkContext): ...
     @property
     def identifier(self) -> Identifier: ...
     def dependency(self): ...
@@ -141,13 +154,13 @@ class ConfigInformation:
         save_directory: Optional[Path] = ...,
     ) -> Config: ...
 
-    class FromPython(GenerationConfigProcessing):
+    class FromPython(ConfigWalk):
         objects: Incomplete
-        def __init__(self, context: GenerationContext) -> None: ...
+        def __init__(self, context: ConfigWalkContext) -> None: ...
         def preprocess(self, config: Config): ...
         def postprocess(self, config: Config, values: Dict[str, Any]): ...
 
-    def fromConfig(self, context: GenerationContext): ...
+    def fromConfig(self, context: ConfigWalkContext): ...
     def add_dependencies(self, *dependencies) -> None: ...
 
 def clone(v): ...
@@ -162,13 +175,14 @@ class TypeConfig:
     def __arguments__(self): ...
     def tags(self): ...
     def add_dependencies(self, *dependencies): ...
-    def instance(self, context: GenerationContext = ...) -> T: ...
+    def instance(self, context: ConfigWalkContext = ...) -> T: ...
     def submit(
         self,
         *,
         workspace: Incomplete | None = ...,
         launcher: Incomplete | None = ...,
-        run_mode: RunMode = ...
+        run_mode: RunMode = ...,
+        pre: List[Task] = []
     ): ...
     def stdout(self): ...
     def stderr(self): ...
@@ -183,7 +197,7 @@ class Config:
     __xpmtype__: ClassVar[ObjectType]
     __xpm__: ConfigInformation
     @classmethod
-    def __getxpmtype__(cls): ...
+    def __getxpmtype__(cls) -> ObjectType: ...
     def __getnewargs_ex__(self): ...
     @classmethod
     def c(cls, **kwargs) -> T: ...

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.22.0 → experimaestro-0.24.0}/src/experimaestro/core/types.py

@@ -1,6 +1,7 @@
+from abc import ABC, abstractmethod
 import inspect
 import sys
-from typing import Union, Dict, Iterator, List
+from typing import Set, Union, Dict, Iterator, List
 from collections import ChainMap
 from pathlib import Path
 import typing
@@ -17,6 +18,11 @@ if sys.version_info.major == 3 and sys.version_info.minor < 9:
 else:
     from typing import _AnnotatedAlias, get_type_hints
 
+if typing.TYPE_CHECKING:
+    from experimaestro.scheduler.base import Job
+    from experimaestro.launchers import Launcher
+    from experimaestro.core.objects import Config
+
 
 class Identifier:
     def __init__(self, name: str):
@@ -143,7 +149,42 @@ class DeprecatedAttribute:
         self.fn(instance, value)
 
 
+class SubmitHook(ABC):
+    """Hook called before the job is submitted to the scheduler
+
+    This allows modifying e.g. the run environnement
+    """
+
+    @abstractmethod
+    def __call__(self, job: "Job", launcher: "Launcher"):
+        ...
+
+    @abstractmethod
+    def __spec__(self):
+        """Returns an identifier tuple for hashing/equality"""
+        ...
+
+    def __eq__(self, other):
+        if other.__class__ is not self.__class__:
+            return False
+        return self.__spec__ == other.__spec__
+
+    def __hash__(self):
+        return hash((self.__class__, self.__spec__))
+
+
+def submit_hook_decorator(hook: SubmitHook):
+    def decorator(cls: typing.Type["Config"]):
+        cls.__getxpmtype__().submit_hooks.add(hook)
+        return cls
+
+    return decorator
+
+
 class ObjectType(Type):
+    submit_hooks: Set[SubmitHook]
+    """Hooks associated with this configuration"""
+
     """ObjectType contains class-level information about
     experimaestro configurations and tasks
 
@@ -167,6 +208,7 @@ class ObjectType(Type):
         self.taskcommandfactory = None
         self.task = None
         self._title = None
+        self.submit_hooks = set()
 
         # Get the identifier
         if identifier is None and hasattr(tp, "__xpmid__"):
@@ -219,7 +261,7 @@ class ObjectType(Type):
         self.configtype.__module__ = tp.__module__
 
         # Create the type-specific object class
-        # (now, the same as basetype - TODO: remove references)
+        # (now, the same as basetype - but in the future, remove references)
         self.objecttype = self.basetype  # type: type
         self.basetype._ = self.configtype
 

{experimaestro-0.22.0 → experimaestro-0.24.0}/src/experimaestro/generators.py

@@ -1,15 +1,16 @@
 import inspect
 from pathlib import Path
-from typing import Callable, List, Tuple, Union
+from typing import Callable, Union
 from experimaestro.core.arguments import ArgumentOptions, TypeAnnotation
-from experimaestro.core.objects import GenerationContext, Config
+from experimaestro.core.objects import ConfigWalkContext, Config
 
 
 class Generator:
     """Base class for all generators"""
 
     def isoutput(self):
-        """Returns True if this generator is a task output (e.g. generates a path within the job folder)"""
+        """Returns True if this generator is a task output (e.g. generates a
+        path within the job folder)"""
         return False
 
 
@@ -17,11 +18,11 @@ class PathGenerator(Generator):
     """Generates a path"""
 
     def __init__(
-        self, path: Union[str, Path, Callable[[GenerationContext, Config], Path]]
+        self, path: Union[str, Path, Callable[[ConfigWalkContext, Config], Path]]
     ):
         self.path = path
 
-    def __call__(self, context: GenerationContext, config: Config):
+    def __call__(self, context: ConfigWalkContext, config: Config):
         if inspect.isfunction(self.path):
             path = context.currentpath() / self.path(context, config)  # type: Path
         else:
@@ -34,7 +35,7 @@ class PathGenerator(Generator):
 
 
 class pathgenerator(TypeAnnotation):
-    def __init__(self, value: Union[str, Callable[[GenerationContext, Config], str]]):
+    def __init__(self, value: Union[str, Callable[[ConfigWalkContext, Config], str]]):
         self.value = value
 
     def annotate(self, options: ArgumentOptions):

{experimaestro-0.22.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.22.0 → experimaestro-0.24.0}/src/experimaestro/launchers/__init__.py

@@ -1,25 +1,36 @@
 from pathlib import Path, PosixPath
-from typing import Callable, Dict, List, Optional, Protocol
+from typing import Callable, Dict, List, Optional
 from experimaestro.commandline import AbstractCommand, Job, CommandLineJob
 from experimaestro.connectors import Connector
 from experimaestro.connectors.local import ProcessBuilder, LocalConnector
 from experimaestro.connectors.ssh import SshPath, SshConnector
+from abc import ABC, abstractmethod
 
 
-class ScriptBuilder:
-    lockfiles: List[Path]
-    command: "AbstractCommand"
+class ScriptBuilder(ABC):
     """A script builder is responsible for generating the script
     used to launch a command line job"""
 
+    lockfiles: List[Path]
+    """The files that must be locked before starting the job"""
+
+    command: "AbstractCommand"
+    """Command to be run"""
+
+    @abstractmethod
     def write(self, job: CommandLineJob) -> Path:
-        raise NotImplementedError()
+        """Write the commmand line job
+
+        :params job: The job to be written
+        """
+        ...
 
 
 SubmitListener = Callable[[Job], None]
+"""Listen to job submissions"""
 
 
-class Launcher:
+class Launcher(ABC):
     """A launcher"""
 
     submit_listeners: List[SubmitListener]
@@ -36,9 +47,10 @@ class Launcher:
     def setNotificationURL(self, url: Optional[str]):
         self.notificationURL = url
 
+    @abstractmethod
     def scriptbuilder(self) -> ScriptBuilder:
         """Returns a script builder"""
-        raise NotImplementedError(f"For class {self.__class__}")
+        ...
 
     def addListener(self, listener: SubmitListener):
         self.submit_listeners.append(listener)

{experimaestro-0.22.0 → experimaestro-0.24.0}/src/experimaestro/scheduler/base.py

@@ -1,4 +1,4 @@
-from collections import defaultdict
+from collections import ChainMap, defaultdict
 import os
 from pathlib import Path
 from shutil import rmtree
@@ -14,7 +14,7 @@ from experimaestro.scheduler.services import Service
 from experimaestro.settings import get_settings
 
 
-from experimaestro.core.objects import Config, GenerationContext
+from experimaestro.core.objects import Config, ConfigWalkContext
 from experimaestro.utils import logger
 from experimaestro.locking import Locks, LockError, Lock
 from experimaestro.tokens import ProcessCounterToken
@@ -151,6 +151,23 @@ class Job(Resource):
         assert self._future, "Cannot wait a not submitted job"
         return self._future.result()
 
+    @property
+    def environ(self):
+        """Returns the job environment
+
+        It is made of (by order of priority):
+
+        1. The job environment
+        1. The launcher environment
+        1. The workspace environment
+
+        """
+        return ChainMap(
+            {},
+            self.launcher.environ if self.launcher else {},
+            self.workspace.environment.environ,
+        )
+
     @property
     def progress(self):
         return self._progress
@@ -292,7 +309,7 @@ class Job(Resource):
         return self._future
 
 
-class JobContext(GenerationContext):
+class JobContext(ConfigWalkContext):
     def __init__(self, job: Job):
         super().__init__()
         self.job = job
@@ -839,6 +856,7 @@ class experiment:
     def __enter__(self):
         logger.info("Locking experiment %s", self.xplockpath)
         self.xplock = self.workspace.connector.lock(self.xplockpath, 0).__enter__()
+        logger.info("Experiment locked")
 
         # Move old jobs into "jobs.bak"
         if self.workspace.run_mode == RunMode.NORMAL:

{experimaestro-0.22.0 → experimaestro-0.24.0}/src/experimaestro/server/__init__.py

@@ -143,10 +143,15 @@ def proxy_response(base_url: str, request: Request, path: str):
 
 
 def start_app(server: "Server"):
+    logging.debug("Starting Flask server...")
     app = Flask("experimaestro")
-    socketio = SocketIO(app, path="/api", async_mode="eventlet")
+
+    logging.debug("Starting Flask server (SocketIO)...")
+    socketio = SocketIO(app, path="/api", async_mode="gevent")
     listener = Listener(server.scheduler, socketio)
 
+    logging.debug("Starting Flask server (setting up socketio)...")
+
     @socketio.on("connect")
     def handle_connect():
         if server.token != request.cookies.get("experimaestro_token", None):
@@ -183,6 +188,8 @@ def start_app(server: "Server"):
         if process is not None:
             process.kill()
 
+    logging.debug("Starting Flask server (setting up routes)...")
+
     @app.route("/services/<path:path>", methods=["GET", "POST"])
     def route_service(path):
         service, *path = path.split("/", 1)
@@ -256,8 +263,9 @@ def start_app(server: "Server"):
 
     # Start the app
     if server.port is None or server.port == 0:
+        logging.info("Searching for an available port")
         sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-        sock.bind(("localhost", 0))
+        sock.bind(("", 0))
         server.port = sock.getsockname()[1]
         sock.close()
 

{experimaestro-0.22.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)
 
 

{experimaestro-0.22.0 → experimaestro-0.24.0}/src/experimaestro/tests/test_instance.py

@@ -1,6 +1,7 @@
 from typing import Optional
 from experimaestro import config, Param, Config
 from experimaestro.core.objects import TypeConfig
+from experimaestro.core.serializers import SerializedConfig
 
 
 @config()
@@ -31,27 +32,29 @@ def test_simple_instance():
     assert isinstance(b.a, A.__xpmtype__.basetype)
 
 
-@config()
-class SerializedConfig:
-    x: Param[int] = 1
-    pass
+class Model(Config):
+    def __post_init__(self):
+        self.initialized = False
+
 
+class Trainer(Config):
+    model: Param[Model]
 
-class TestSerialization:
-    """Test that a config can be serialized during execution"""
 
-    def test_instance(self):
-        import pickle
+class SerializedModel(SerializedConfig):
+    def initialize(self):
+        self.config.initialized = True
 
-        a = SerializedConfig(x=2).instance()
-        assert not isinstance(a, TypeConfig)
-        assert isinstance(a, SerializedConfig)
 
-        s_a = pickle.dumps(a)
+def test_instance_serialized():
+    model = SerializedModel(config=Model())
+    trainer = Trainer(model=model)
+    instance = trainer.instance()
 
-        deserialized = pickle.loads(s_a)
-        assert not isinstance(deserialized, TypeConfig)
-        assert deserialized.x == 2
+    assert isinstance(
+        instance.model, Model
+    ), f"The model is not a Model but a {type(instance.model).__qualname__}"
+    assert instance.model.initialized, "The model was not initialized"
 
 
 class ConfigWithOptional(Config):

experimaestro-0.24.0/src/experimaestro/tests/test_outputs.py

@@ -0,0 +1,50 @@
+"""Test for task outputs"""
+
+from experimaestro import Config, Task, Param, ConfigWrapper
+from experimaestro.scheduler.workspace import RunMode
+
+
+class B(Config):
+    x: Param[int] = 1
+
+
+class A(Config):
+    b: Param[B]
+
+
+class Main(Task):
+    a: Param[A]
+
+    def taskoutputs(self):
+        return self.a, {
+            "a": self.a,
+        }
+
+    def execute(self):
+        print(self.a.b.x)  # noqa: T201
+
+
+class MainB(Task):
+    b: Param[B]
+
+    def execute(self):
+        pass
+
+
+def test_output_taskoutput():
+    a = A(b=B())
+    output, ioutput = Main(a=a).submit(run_mode=RunMode.DRY_RUN)
+
+    assert isinstance(output, ConfigWrapper), "outputs should be task proxies"
+
+    # Direct
+    Main(a=output)
+
+    # Via getattr
+    Main(a=A(b=output.b))
+
+    # Via getitem
+    Main(a=ioutput["a"])
+
+    # Now, submits
+    Main(a=output).submit(run_mode=RunMode.DRY_RUN)

{experimaestro-0.22.0 → experimaestro-0.24.0}/src/experimaestro/tests/test_progress.py

@@ -1,13 +1,11 @@
 # Test that progress notification work
 from copy import copy
-import logging
 from pathlib import Path
 import time
 import fasteners
 from typing import List, Tuple, Union
-from experimaestro.commandline import CommandLineJob
 from experimaestro import Task, Annotated, pathgenerator, progress, tqdm
-from experimaestro.core.objects import TaskOutput, logger
+from experimaestro.core.objects import ConfigWrapper, logger
 from experimaestro.notifications import LevelInformation
 from experimaestro.scheduler import Job, Listener
 from queue import Queue
@@ -74,7 +72,7 @@ def test_progress_basic():
         listener = ProgressListener()
         xp.scheduler.addlistener(listener)
 
-        out = ProgressingTask().submit()  # type: TaskOutput
+        out: ConfigWrapper = ProgressingTask().submit()
         path = out.path  # type: Path
         job = out.__xpm__.job
 
@@ -87,9 +85,9 @@ def test_progress_basic():
         for v in progresses:
             writeprogress(path, v)
             if v < 1:
-                l = listener.progresses.get()[0]
-                logger.info("Got %s", l)
-                assert l.progress == v
+                info = listener.progresses.get()[0]
+                logger.info("Got %s", info)
+                assert info.progress == v
 
 
 def test_progress_multiple():
@@ -105,7 +103,7 @@ def test_progress_multiple():
         listener1 = ProgressListener()
         xp1.scheduler.addlistener(listener1)
 
-        out = ProgressingTask().submit()  # type: TaskOutput
+        out = ProgressingTask().submit()  # type: ConfigWrapper
         path = out.path  # type: Path
         job = out.__xpm__.job
 
@@ -219,7 +217,7 @@ def test_progress_nested():
         listener = ProgressListener()
         xp.scheduler.addlistener(listener)
 
-        out = NestedProgressingTask().submit()  # type: TaskOutput
+        out = NestedProgressingTask().submit()  # type: ConfigWrapper
         job = out.__xpm__.job
         path = out.path  # type: Path