experimaestro 2.0.0b8__py3-none-any.whl → 2.0.0b17__py3-none-any.whl

experimaestro/experiments/cli.py

@@ -5,7 +5,7 @@ import json
 import logging
 import sys
 from pathlib import Path
-from typing import Any, List, Optional, Protocol, Tuple
+from typing import TYPE_CHECKING, Any, List, Optional, Protocol, Tuple
 
 import click
 import omegaconf
@@ -13,16 +13,21 @@ import yaml
 from omegaconf import OmegaConf, SCMode
 from termcolor import cprint
 
-from experimaestro import LauncherRegistry, RunMode, experiment
 from experimaestro.exceptions import HandledException
 from experimaestro.experiments.configuration import ConfigurationBase
+from experimaestro.launcherfinder.registry import LauncherRegistry
+from experimaestro.scheduler.workspace import RunMode
 from experimaestro.settings import find_workspace
+from experimaestro.utils.logging import setup_logging
+
+if TYPE_CHECKING:
+    from experimaestro.scheduler.experiment import experiment
 
 
 class ExperimentHelper:
     """Helper for experiments"""
 
-    xp: experiment
+    xp: "experiment"
     """The experiment object"""
 
     #: Run function
@@ -52,7 +57,7 @@ class ExperimentHelper:
 class ExperimentCallable(Protocol):
     """Protocol for the run function"""
 
-    def __call__(self, helper: ExperimentHelper, configuration: Any): ...  # noqa: E704
+    def __call__(self, helper: ExperimentHelper, configuration: Any): ...
 
 
 class ConfigurationLoader:
@@ -74,6 +79,11 @@ class ConfigurationLoader:
             if not path.is_absolute():
                 _data["file"] = str((yaml_file.parent / path).resolve())
 
+        if "pre_experiment" in _data:
+            path = Path(_data["pre_experiment"])
+            if not path.is_absolute():
+                _data["pre_experiment"] = str((yaml_file.parent / path).resolve())
+
         if "module" in _data:
             # Keeps track of the YAML file where the module was defined
             self.yaml_module_file = yaml_file
@@ -93,6 +103,18 @@ class ConfigurationLoader:
 
 @click.option("--debug", is_flag=True, help="Print debug information")
 @click.option("--show", is_flag=True, help="Print configuration and exits")
+@click.option(
+    "--watcher",
+    type=click.Choice(["auto", "polling", "inotify", "fsevents", "kqueue", "windows"]),
+    default="auto",
+    help="Filesystem watcher type (auto=platform default, polling=network mounts)",
+)
+@click.option(
+    "--polling-interval",
+    type=float,
+    default=1.0,
+    help="Polling interval in seconds (only for --watcher=polling)",
+)
 @click.option(
     "--env",
     help="Define one environment variable",
@@ -103,8 +125,8 @@ class ConfigurationLoader:
     "--host",
     type=str,
     default=None,
-    help="Server hostname (default to localhost,"
-    " not suitable if your jobs are remote)",
+    help="[DEPRECATED] Server hostname (use --web instead)",
+    hidden=True,
 )
 @click.option(
     "--run-mode",
@@ -116,20 +138,31 @@ class ConfigurationLoader:
     "--xpm-config-dir",
     type=Path,
     default=None,
-    help="Path for the experimaestro config directory "
-    "(if not specified, use $HOME/.config/experimaestro)",
+    help="Path for the experimaestro config directory (if not specified, use $HOME/.config/experimaestro)",
 )
 @click.option(
     "--port",
     type=int,
     default=None,
-    help="Port for monitoring (can be defined in the settings.yaml file)",
+    help="[DEPRECATED] Port for monitoring (use --web instead)",
+    hidden=True,
+)
+@click.option(
+    "--web",
+    is_flag=True,
+    help="Start web server for monitoring (use settings.yaml for host/port config)",
 )
 @click.option(
     "--console",
     is_flag=True,
     help="Launch Textual console UI for monitoring with logs",
 )
+@click.option(
+    "--no-db",
+    "no_db",
+    is_flag=True,
+    help="Disable database state tracking for this experiment",
+)
 @click.option(
     "--file",
     "xp_file",
@@ -166,7 +199,9 @@ def experiments_cli(  # noqa: C901
     xp_file: str,
     host: str,
     port: int,
+    web: bool,
     console: bool,
+    no_db: bool,
     xpm_config_dir: Path,
     workdir: Optional[Path],
     workspace: Optional[str],
@@ -178,13 +213,40 @@ def experiments_cli(  # noqa: C901
     module_name: Optional[str],
     args: List[str],
     show: bool,
+    watcher: str,
+    polling_interval: float,
     debug: bool,
 ):
     """Run an experiment"""
-
-    # --- Set the logger
-    logging.getLogger().setLevel(logging.DEBUG if debug else logging.INFO)
-    logging.getLogger("xpm.hash").setLevel(logging.INFO)
+    import warnings
+
+    # --- Set the logger with colors if outputting to terminal
+    setup_logging(debug=debug)
+
+    # --- Configure filesystem watcher type
+    from experimaestro.ipc import IPCom, WatcherType
+
+    if watcher != "auto":
+        IPCom.set_watcher_type(WatcherType(watcher), polling_interval)
+    elif polling_interval != 1.0:
+        # If polling interval is specified but watcher is auto, use polling
+        IPCom.set_watcher_type(WatcherType.POLLING, polling_interval)
+
+    # --- Warn about deprecated options
+    if host is not None:
+        warnings.warn(
+            "The '--host' option is deprecated. Use '--web' flag instead. "
+            "Configure host in settings.yaml.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+    if port is not None:
+        warnings.warn(
+            "The '--port' option is deprecated. Use '--web' flag instead. "
+            "Configure port in settings.yaml.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
 
     # --- Loads the YAML
     conf_loader = ConfigurationLoader()
@@ -207,9 +269,9 @@ def experiments_cli(  # noqa: C901
     if xp_file is None:
         xp_file = configuration.get("file", None)
         if xp_file:
-            assert (
-                not module_name
-            ), "Module name and experiment file are mutually exclusive options"
+            assert not module_name, (
+                "Module name and experiment file are mutually exclusive options"
+            )
             xp_file = Path(xp_file)
             if not python_path:
                 python_path.append(xp_file.parent.absolute())
@@ -217,9 +279,9 @@ def experiments_cli(  # noqa: C901
                 "Using python path: %s", ", ".join(str(s) for s in python_path)
             )
 
-    assert (
-        module_name or xp_file
-    ), "Either the module name or experiment file should be given"
+    assert module_name or xp_file, (
+        "Either the module name or experiment file should be given"
+    )
 
     # --- Set some options
 
@@ -233,6 +295,27 @@ def experiments_cli(  # noqa: C901
     for path in python_path:
         sys.path.append(str(path))
 
+    # --- Execute pre-experiment script if specified
+    pre_experiment = configuration.get("pre_experiment", None)
+    if pre_experiment:
+        pre_exp_path = Path(pre_experiment)
+        if pre_exp_path.exists():
+            logging.info("Executing pre-experiment script: %s", pre_exp_path)
+            try:
+                spec = importlib.util.spec_from_file_location(
+                    "pre_experiment", str(pre_exp_path.absolute())
+                )
+                pre_mod = importlib.util.module_from_spec(spec)
+                spec.loader.exec_module(pre_mod)
+            except Exception as e:
+                raise click.ClickException(
+                    f"Failed to execute pre-experiment script '{pre_exp_path}': {e}"
+                )
+        else:
+            raise click.ClickException(
+                f"Pre-experiment script not found: {pre_exp_path}"
+            )
+
     # --- Adds automatically the experiment module if not found
     if module_name and conf_loader.yaml_module_file:
         try:
@@ -285,6 +368,14 @@ def experiments_cli(  # noqa: C901
         )
 
     schema = list_parameters[1].annotation
+    if isinstance(schema, str):
+        # Get the schema from the module
+        schema = getattr(mod.__dict__, schema, None)
+        assert schema is not None, (
+            f"Could not find schema {list_parameters[1].annotation} "
+            f"in module {module_name}"
+        )
+
     omegaconf_schema = OmegaConf.structured(schema())
 
     if omegaconf_schema is not None:
@@ -323,20 +414,48 @@ def experiments_cli(  # noqa: C901
         str(workdir.resolve()),
     )
 
+    # Determine project path for git info
+    project_paths = []
+    if xp_file:
+        project_paths.append(xp_file.resolve().parent)
+    elif hasattr(mod, "__file__") and mod.__file__:
+        project_paths.append(Path(mod.__file__).resolve().parent)
+
     # Define the experiment execution function
-    def run_experiment_code(xp_holder=None, xp_ready_event=None, register_signals=True):
+    def run_experiment_code(
+        xp_holder=None, xp_ready_event=None, register_signals=True, in_thread=False
+    ):
         """Run the experiment code - optionally storing xp in xp_holder"""
         try:
+            from experimaestro.scheduler.experiment import experiment
+            from experimaestro.settings import get_settings
+
             with experiment(
                 ws_env,
                 experiment_id,
-                host=host,
-                port=port,
                 run_mode=run_mode,
                 register_signals=register_signals,
+                project_paths=project_paths,
+                dirty_git=xp_configuration.dirty_git,
+                no_db=no_db,
             ) as xp:
                 if xp_holder is not None:
                     xp_holder["xp"] = xp
+
+                # Start web server if requested
+                if web and run_mode == RunMode.NORMAL:
+                    settings = get_settings()
+                    xp.scheduler.start_server(
+                        settings.server,
+                        workspace=xp.workspace,
+                        wait_for_quit=False,
+                    )
+                    logging.info(
+                        "Web server started at http://%s:%d",
+                        settings.server.host or "localhost",
+                        settings.server.port or 12345,
+                    )
+
                 if xp_ready_event is not None:
                     xp_ready_event.set()  # Signal that xp is ready
 
@@ -357,7 +476,11 @@ def experiments_cli(  # noqa: C901
                 # ... and wait
                 xp.wait()
 
-        except HandledException:
+        except HandledException as e:
+            if in_thread:
+                # Re-raise to preserve exception info for the main thread
+                raise
+            cprint(f"Experiment failed: {e}", "red", file=sys.stderr)
             sys.exit(1)
 
     # Console mode is only available in NORMAL run mode
@@ -366,47 +489,91 @@ def experiments_cli(  # noqa: C901
         logging.warning("--console is ignored when run_mode is not NORMAL")
 
     if use_console:
-        # Run experiment in background thread, console UI in main thread
+        # Start TUI first, then run experiment in background thread
+        # This ensures all logs (including startup) are captured
         import threading
         from experimaestro.tui import ExperimentTUI
 
+        # Initialize multiprocessing resource tracker before Textual takes over
+        # terminal file descriptors. This prevents "bad value(s) in fds_to_keep"
+        # errors when code in the background thread (e.g., tqdm in torchvision)
+        # tries to use multiprocessing.
+        try:
+            from multiprocessing import resource_tracker
+
+            resource_tracker.ensure_running()
+        except Exception:
+            pass  # Best effort - may not be needed on all systems
+
         xp_holder = {"xp": None}
         exception_holder = {"exception": None}
-        xp_ready = threading.Event()
 
-        def run_in_thread():
+        # Create TUI first in deferred mode (no state_provider yet)
+        # This allows capturing all logs from experiment startup
+        tui_app = ExperimentTUI(show_logs=True)
+
+        def run_experiment_in_thread():
+            """Run experiment and connect state provider to TUI when ready"""
             try:
-                # Don't register signals in background thread
-                run_experiment_code(xp_holder, xp_ready, register_signals=False)
-                # Add a test message after experiment completes
+                from experimaestro.scheduler.experiment import experiment as exp_context
+                from experimaestro.settings import get_settings
+
+                with exp_context(
+                    ws_env,
+                    experiment_id,
+                    run_mode=run_mode,
+                    register_signals=False,  # TUI handles signals
+                    project_paths=project_paths,
+                    dirty_git=xp_configuration.dirty_git,
+                    no_db=no_db,
+                ) as xp:
+                    xp_holder["xp"] = xp
+
+                    # Connect TUI to the experiment's scheduler
+                    tui_app.call_from_thread(tui_app.set_state_provider, xp.scheduler)
+
+                    # Start web server if requested
+                    if web:
+                        settings = get_settings()
+                        xp.scheduler.start_server(
+                            settings.server,
+                            workspace=xp.workspace,
+                            wait_for_quit=False,
+                        )
+                        logging.info(
+                            "Web server started at http://%s:%d",
+                            settings.server.host or "localhost",
+                            settings.server.port or 12345,
+                        )
+
+                    logging.info("Experiment started")
+
+                    # Set up the environment
+                    for key, value in env:
+                        xp.setenv(key, value)
+
+                    # Sets the python path
+                    xp.workspace.python_path.extend(python_path)
+
+                    # Run the experiment
+                    helper.xp = xp
+                    helper.run(list(args), xp_configuration)
+
+                    # ... and wait
+                    xp.wait()
+
                 logging.info("Experiment thread completed")
-            except Exception as e:
+
+            except BaseException as e:
+                # Use BaseException to also catch SystemExit from sys.exit()
                 exception_holder["exception"] = e
-                xp_ready.set()  # Signal even on error
 
         # Start experiment in background thread
-        exp_thread = threading.Thread(target=run_in_thread, daemon=True)
+        exp_thread = threading.Thread(target=run_experiment_in_thread, daemon=True)
         exp_thread.start()
 
-        # Wait for experiment to start (up to 30 seconds)
-        if not xp_ready.wait(timeout=30.0):
-            cprint("Timeout waiting for experiment to start", "red", file=sys.stderr)
-            sys.exit(1)
-
-        if xp_holder["xp"] is None:
-            cprint("Failed to start experiment", "red", file=sys.stderr)
-            if exception_holder["exception"]:
-                raise exception_holder["exception"]
-            sys.exit(1)
-
-        # Run TUI in main thread (handles signals via Textual)
-        tui_app = ExperimentTUI(
-            workdir=workdir,
-            state_provider=xp_holder["xp"].state_provider,
-            show_logs=True,
-        )
-
         try:
+            # Run TUI in main thread (handles signals via Textual)
             # Textual automatically captures stdout/stderr via Print events
             tui_app.run()
         finally:

experimaestro/experiments/configuration.py

@@ -1,3 +1,4 @@
+from enum import Enum
 from omegaconf import MISSING
 from typing import Optional, List
 import attr
@@ -8,6 +9,19 @@ except ImportError:
     from typing_extensions import dataclass_transform
 
 
+class DirtyGitAction(str, Enum):
+    """Action to take when the git repository has uncommitted changes"""
+
+    IGNORE = "ignore"
+    """Don't check or warn about dirty git state"""
+
+    WARN = "warn"
+    """Warn about dirty git state (default)"""
+
+    ERROR = "error"
+    """Raise an error if git is dirty"""
+
+
 @dataclass_transform(kw_only_default=True)
 def configuration(*args, **kwargs):
     """Method to define keyword only dataclasses
@@ -40,6 +54,13 @@ class ConfigurationBase:
     parent: Optional[str] = None
     """Relative path of a YAML file that should be merged"""
 
+    pre_experiment: Optional[str] = None
+    """Relative path to a Python file to execute before importing the experiment.
+
+    This is useful for setting environment variables or mocking modules to speed up
+    the experiment setup phase (e.g., mocking torch.compile or torch.nn).
+    The actual job execution will use real modules."""
+
     title: str = ""
     """Short description of the experiment"""
 
@@ -54,3 +75,6 @@ class ConfigurationBase:
 
     add_timestamp: bool = False
     """Adds a timestamp YYYY_MM_DD-HH_MM to the experiment ID"""
+
+    dirty_git: DirtyGitAction = DirtyGitAction.WARN
+    """Action when git repository has uncommitted changes: ignore, warn (default), error"""

experimaestro/generators.py

@@ -6,7 +6,7 @@ from experimaestro.core.arguments import ArgumentOptions, TypeAnnotation
 from experimaestro.core.objects import ConfigWalkContext, Config
 
 if TYPE_CHECKING:
-    from experimaestro.core.subparameters import Subparameters
+    from experimaestro.core.partial import Partial
 
 
 class Generator(ABC):
@@ -33,7 +33,7 @@ class PathGenerator(Generator):
             output: Meta[Path] = field(default_factory=PathGenerator("results.json"))
             model: Meta[Path] = field(default_factory=PathGenerator("model.pt"))
 
-    For shared directories across related tasks, use with subparameters::
+    For shared directories across related tasks, use with partial::
 
         training_group = param_group("training")
 
@@ -42,13 +42,13 @@ class PathGenerator(Generator):
             checkpoint: Meta[Path] = field(
                 default_factory=PathGenerator(
                     "model.pt",
-                    subparameters=subparameters(exclude=[training_group])
+                    partial=partial(exclude=[training_group])
                 )
             )
 
     :param path: Relative path within the task directory. Can be a string,
         Path, or callable that takes (context, config) and returns a Path.
-    :param subparameters: Optional subparameters for partial directory sharing.
+    :param partial: Optional partial for partial directory sharing.
         When provided, the path is generated in a shared partial directory.
     """
 
@@ -56,7 +56,7 @@ class PathGenerator(Generator):
         self,
         path: Union[str, Path, Callable[[ConfigWalkContext, Config], Path]] = "",
         *,
-        partial: "Subparameters" = None,
+        partial: "Partial" = None,
     ):
         self.path = path
         self.partial = partial

experimaestro/ipc.py

@@ -1,5 +1,6 @@
 """IPC utilities"""
 
+from enum import Enum
 from typing import Optional
 from pathlib import Path
 import os
@@ -11,16 +12,132 @@ from watchdog.observers.api import ObservedWatch
 from watchdog.events import FileSystemEventHandler
 
 
+class WatcherType(str, Enum):
+    """Available filesystem watcher types"""
+
+    AUTO = "auto"
+    """Use the best available watcher for the platform (default)"""
+
+    POLLING = "polling"
+    """Platform-independent polling (works on network mounts)"""
+
+    INOTIFY = "inotify"
+    """Linux inotify (Linux 2.6.13+ only)"""
+
+    FSEVENTS = "fsevents"
+    """macOS FSEvents (macOS only)"""
+
+    KQUEUE = "kqueue"
+    """BSD/macOS kqueue (less scalable for deep directories)"""
+
+    WINDOWS = "windows"
+    """Windows API (Windows only)"""
+
+
+def _create_observer(watcher_type: WatcherType, polling_interval: float = 1.0):
+    """Create an observer of the specified type
+
+    Args:
+        watcher_type: The type of filesystem watcher to use
+        polling_interval: Polling interval in seconds (for polling watcher)
+
+    Returns:
+        An observer instance
+
+    Raises:
+        ImportError: If the requested watcher type is not available on this platform
+        ValueError: If watcher_type is invalid
+    """
+    match watcher_type:
+        case WatcherType.AUTO:
+            return Observer()
+
+        case WatcherType.POLLING:
+            from watchdog.observers.polling import PollingObserver
+
+            return PollingObserver(timeout=polling_interval)
+
+        case WatcherType.INOTIFY:
+            from watchdog.observers.inotify import InotifyObserver
+
+            return InotifyObserver()
+
+        case WatcherType.FSEVENTS:
+            from watchdog.observers.fsevents import FSEventsObserver
+
+            return FSEventsObserver()
+
+        case WatcherType.KQUEUE:
+            from watchdog.observers.kqueue import KqueueObserver
+
+            return KqueueObserver()
+
+        case WatcherType.WINDOWS:
+            from watchdog.observers.read_directory_changes import WindowsApiObserver
+
+            return WindowsApiObserver()
+
+        case _:
+            raise ValueError(f"Unknown watcher type: {watcher_type}")
+
+
 class IPCom:
     """IPC async thread"""
 
     INSTANCE: Optional["IPCom"] = None
+    # Testing mode: use polling observer with small interval
+    TESTING_MODE: bool = False
+    POLLING_INTERVAL: float = 0.01
+    # Watcher type configuration
+    WATCHER_TYPE: WatcherType = WatcherType.AUTO
 
     def __init__(self):
-        self.observer = Observer()
+        if IPCom.TESTING_MODE:
+            from watchdog.observers.polling import PollingObserver
+
+            self.observer = PollingObserver(timeout=IPCom.POLLING_INTERVAL)
+        else:
+            self.observer = _create_observer(IPCom.WATCHER_TYPE, IPCom.POLLING_INTERVAL)
         self.observer.start()
         self.pid = os.getpid()
 
+    @classmethod
+    def set_watcher_type(cls, watcher_type: WatcherType, polling_interval: float = 1.0):
+        """Set the filesystem watcher type
+
+        Args:
+            watcher_type: The type of watcher to use
+            polling_interval: Polling interval in seconds (for polling watcher)
+
+        Note:
+            This must be called before the first IPCom instance is created.
+            If an instance already exists, it will be reset.
+        """
+        cls.WATCHER_TYPE = watcher_type
+        cls.POLLING_INTERVAL = polling_interval
+        # Reset instance to apply new settings
+        if cls.INSTANCE is not None:
+            cls.INSTANCE.observer.stop()
+            cls.INSTANCE.observer.join(timeout=5)
+            cls.INSTANCE = None
+        logger.info("Set watcher type to %s", watcher_type.value)
+
+    @classmethod
+    def set_testing_mode(cls, enabled: bool = True, polling_interval: float = 0.01):
+        """Enable testing mode with polling observer
+
+        Args:
+            enabled: Whether to enable testing mode
+            polling_interval: Polling interval in seconds (default 0.01)
+        """
+        cls.TESTING_MODE = enabled
+        cls.POLLING_INTERVAL = polling_interval
+        # Reset instance to apply new settings
+        if cls.INSTANCE is not None:
+            cls.INSTANCE.observer.stop()
+            cls.INSTANCE.observer.join(timeout=5)
+            cls.INSTANCE = None
+
     def fswatch(
         self, watcher: FileSystemEventHandler, path: Path, recursive=False
     ) -> ObservedWatch:

experimaestro/launcherfinder/__init__.py

@@ -1,6 +1,6 @@
-# flake8: noqa: F401
+# ruff: noqa: F401
 
-from .base import *
+from .base import ConnectorConfiguration, TokenConfiguration
 from .specs import (
     cpu,
     cuda_gpu,

experimaestro/launcherfinder/registry.py

@@ -59,9 +59,9 @@ class LauncherRegistry:
             ).expanduser()
 
         if LauncherRegistry.CURRENT_CONFIG_DIR not in LauncherRegistry.INSTANCES:
-            LauncherRegistry.INSTANCES[
-                LauncherRegistry.CURRENT_CONFIG_DIR
-            ] = LauncherRegistry(LauncherRegistry.CURRENT_CONFIG_DIR)
+            LauncherRegistry.INSTANCES[LauncherRegistry.CURRENT_CONFIG_DIR] = (
+                LauncherRegistry(LauncherRegistry.CURRENT_CONFIG_DIR)
+            )
 
         return LauncherRegistry.INSTANCES[LauncherRegistry.CURRENT_CONFIG_DIR]
 
@@ -83,7 +83,6 @@ class LauncherRegistry:
 
         # Register the find launcher function if it exists
         launchers_py = basepath / "launchers.py"
-        print(f"basepath {launchers_py}")
         if launchers_py.is_file():
             logger.info("Loading %s", launchers_py)
 
@@ -163,9 +162,9 @@ class LauncherRegistry:
         if self.find_launcher_fn is not None:
             for spec in specs.requirements:
                 if launcher := self.find_launcher_fn(spec, tags):
-                    assert isinstance(
-                        launcher, Launcher
-                    ), "f{self.find_launcher_fn} did not return a Launcher but {type(launcher)}"
+                    assert isinstance(launcher, Launcher), (
+                        "f{self.find_launcher_fn} did not return a Launcher but {type(launcher)}"
+                    )
                     return launcher
 
         return None