feat(experiment): add python_path to make modules importable in task workers

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:
 

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
 

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