diff --git a/docs/source/experiments.md b/docs/source/experiments.md index 816d8660..622eaf67 100644 --- a/docs/source/experiments.md +++ b/docs/source/experiments.md @@ -22,6 +22,17 @@ with experiment("/path/to/workspace", "my-experiment", port=12345) as xp: model = Learn.C(epochs=100).submit() ``` +Task workers load the task module from its file and do not inherit the +`sys.path` of the launching script. If your tasks import modules *by name* at +run time (plugin registries, factories, sibling packages), give the roots of +these packages with `python_path` — the counterpart of the `pythonpath` key of +the [YAML configuration](experiments/cli-runner.md): + +```python +with experiment("/path/to/workspace", "my-experiment", python_path=["src"]) as xp: + ... +``` + or through the command line, where the experiment code is a `run(helper, cfg)` function and its parameters come from a configuration file: diff --git a/src/experimaestro/scheduler/experiment.py b/src/experimaestro/scheduler/experiment.py index 38ce0c13..a6659bf8 100644 --- a/src/experimaestro/scheduler/experiment.py +++ b/src/experimaestro/scheduler/experiment.py @@ -160,6 +160,7 @@ def __init__( launcher=None, register_signals: bool = True, project_paths: Optional[list[Path]] = None, + python_path: list[Path | str] | None = None, wait_for_quit: bool = False, dirty_git: DirtyGitAction = DirtyGitAction.WARN, no_db: bool = False, @@ -187,6 +188,15 @@ def __init__( :param project_paths: Paths to the project files (for git info). If not provided, will be inferred from the caller's location. + :param python_path: Extra import paths prepended to ``sys.path`` of task + workers, searched in the order given. Task workers reconstruct the + task module by file path and do NOT inherit the launching + interpreter's ``sys.path``; pass the roots of any packages your + tasks import *by name* at run time (plugins, factories, sibling + packages) here. Equivalent to exporting ``PYTHONPATH`` before + submitting, and the counterpart of the ``pythonpath`` YAML key of + ``run-experiment``. + :param wait_for_quit: Deprecated, no longer used. Web server is no longer started automatically. @@ -248,6 +258,14 @@ def __init__( run_mode = run_mode or RunMode.NORMAL self.workspace = Workspace(settings, env, launcher=launcher, run_mode=run_mode) + # Import paths made available to task workers. The script builder + # inserts each one at the front of their sys.path, so we add them in + # reverse order to keep the given order as the search order. + if python_path: + self.workspace.python_path.extend( + str(Path(p).resolve()) for p in reversed(python_path) + ) + # Store experiment name for ID references self.name = name diff --git a/src/experimaestro/tests/test_python_path.py b/src/experimaestro/tests/test_python_path.py new file mode 100644 index 00000000..c8e8acfd --- /dev/null +++ b/src/experimaestro/tests/test_python_path.py @@ -0,0 +1,67 @@ +"""Regression test for ``experiment(python_path=...)``. + +Task workers reconstruct the task module by file path and do not inherit the +launching interpreter's ``sys.path``. A task that imports a module *by name* at +run time therefore fails unless that module's root is made available to the +worker. ``experiment(python_path=[...])`` is the programmatic way to do so. +""" + +import importlib +import sys + +import pytest + +from experimaestro import Param, Task, experiment +from experimaestro.scheduler import JobState + +_PROBE_PKG = "xpm_python_path_probe" + + +class DynamicImportProbe(Task): + """Imports ``module`` by name inside the worker (needs worker sys.path).""" + + module: Param[str] + + def execute(self): + from pathlib import Path + + mod = importlib.import_module(self.module) + Path("probe.txt").write_text(mod.VALUE) + + +def _write_probe_package(root): + pkg = root / _PROBE_PKG + pkg.mkdir(parents=True) + (pkg / "__init__.py").write_text("VALUE = 'imported-via-python-path'\n") + return root + + +def test_python_path_makes_module_importable_in_worker(tmp_path): + pkgroot = _write_probe_package(tmp_path / "pkgs") + # The probe package must not be importable in the parent either. + assert _PROBE_PKG not in sys.modules + + with experiment( + tmp_path / "xp", + "python_path_test", + no_environmental_impact=True, + python_path=[pkgroot], + ): + job = DynamicImportProbe.C(module=_PROBE_PKG).submit() + assert job.__xpm__.job.wait() == JobState.DONE + + +@pytest.mark.xfail( + reason="without python_path the worker cannot import the module", strict=True +) +def test_without_python_path_the_worker_cannot_import(tmp_path, monkeypatch): + # Make sure the worker cannot find the package through the environment + monkeypatch.delenv("PYTHONPATH", raising=False) + _write_probe_package(tmp_path / "pkgs") + sys.path_importer_cache.clear() + + with experiment( + tmp_path / "xp", "python_path_control", no_environmental_impact=True + ): + job = DynamicImportProbe.C(module=_PROBE_PKG).submit() + assert job.__xpm__.job.wait() == JobState.DONE