From 85fefb40f4308a9e39d1e4b7b78af5261ee6576a Mon Sep 17 00:00:00 2001 From: Zeng Date: Sat, 10 Oct 2026 18:03:18 +0800 Subject: [PATCH 1/3] fix(profile): keep the inherited PYTHONPATH on profile servers Magpie launches the server with os.environ overlaid by the materialized config's benchmark.envs, so a PYTHONPATH in the config replaces the inherited one. Two profile-flow sites prepended to PYTHONPATH reading only the config: the TraceLens kernel_shape_tool step in materialization and the start-up shim in ProfileExecutor._inject_host_probe. When the config carried no PYTHONPATH they wrote one holding only their own entry, and the profile server lost the operator's private framework trees, the image's PYTHONPATH and the container wrapper paths. In 14 days of sessions, 89 of 93 roofline configs carried such a PYTHONPATH (71 shim-only, 16 shape tool plus shim on SGLang), while baseline/explore/integrate configs never set one, so only the profile server ran on a different import path. Both sites now go through prepend_launch_pythonpath, which builds on the config value when the key is present (the config stays authoritative) and on the inherited value otherwise, puts the new entry first, dedups and drops empty parts. The shim is therefore also moved to the front when it is already present further down, since Python imports only the first sitecustomize on the path. Co-Authored-By: Claude --- .../tests/test_framework_rewrite_evidence.py | 45 +++++++++++++++++ .../tests/test_workload_envs_branches_unit.py | 48 +++++++++++++++++++ .../actions/executors/_workload_envs.py | 16 ++++++- .../orchestrator/actions/executors/profile.py | 8 ++-- 4 files changed, 111 insertions(+), 6 deletions(-) diff --git a/src/hyperloom/inference_optimizer/tests/test_framework_rewrite_evidence.py b/src/hyperloom/inference_optimizer/tests/test_framework_rewrite_evidence.py index 0fea308cd8..6408155f88 100644 --- a/src/hyperloom/inference_optimizer/tests/test_framework_rewrite_evidence.py +++ b/src/hyperloom/inference_optimizer/tests/test_framework_rewrite_evidence.py @@ -1386,6 +1386,50 @@ def test_probe_injection_is_idempotent(tmp_path, monkeypatch): assert entries.count(str(evidence.probe_asset_dir())) == 1 +def _injected_pythonpath(tmp_path: Path, envs: dict[str, Any]) -> str: + """Arm the start-up shim on a fresh config carrying ``envs`` and return the resulting PYTHONPATH.""" + import yaml + + from hyperloom.orchestrator.actions.executors.profile import ProfileExecutor + + config = tmp_path / "profile.yaml" + _write_profile_config(config, envs) + ProfileExecutor()._inject_host_probe(config, tmp_path / "ws") + return yaml.safe_load(config.read_text(encoding="utf-8"))["benchmark"]["envs"]["PYTHONPATH"] + + +def test_probe_injection_keeps_the_inherited_pythonpath(tmp_path, monkeypatch): + """Magpie replaces the launch PYTHONPATH with the config's, so a config without one carries the inherited one.""" + monkeypatch.setenv("PYTHONPATH", os.pathsep.join(["/a", "/b"])) + shim = str(evidence.probe_asset_dir()) + assert _injected_pythonpath(tmp_path, {"TP": 8}) == os.pathsep.join([shim, "/a", "/b"]) + + +def test_probe_injection_does_not_append_the_inherited_pythonpath_to_the_config_one(tmp_path, monkeypatch): + """A config that carries PYTHONPATH is authoritative over the inherited value.""" + monkeypatch.setenv("PYTHONPATH", os.pathsep.join(["/a", "/b"])) + shim = str(evidence.probe_asset_dir()) + assert _injected_pythonpath(tmp_path, {"PYTHONPATH": "/c"}) == os.pathsep.join([shim, "/c"]) + + +def test_probe_injection_moves_a_present_shim_to_the_front(tmp_path, monkeypatch): + """Python imports only the first sitecustomize, so a shim already on the path but not first is moved up.""" + import yaml + + from hyperloom.orchestrator.actions.executors.profile import ProfileExecutor + + monkeypatch.delenv("PYTHONPATH", raising=False) + shim = str(evidence.probe_asset_dir()) + config = tmp_path / "profile.yaml" + _write_profile_config(config, {"PYTHONPATH": os.pathsep.join(["/c", shim, "", "/d"])}) + executor = ProfileExecutor() + executor._inject_host_probe(config, tmp_path / "ws") + first = yaml.safe_load(config.read_text(encoding="utf-8"))["benchmark"]["envs"]["PYTHONPATH"] + assert first == os.pathsep.join([shim, "/c", "/d"]) + executor._inject_host_probe(config, tmp_path / "ws") + assert yaml.safe_load(config.read_text(encoding="utf-8"))["benchmark"]["envs"]["PYTHONPATH"] == first + + def test_probe_injection_respects_the_off_switch(tmp_path, monkeypatch): """With the probe switched off only the start-up shim goes in, without the probe's environment.""" import yaml @@ -1393,6 +1437,7 @@ def test_probe_injection_respects_the_off_switch(tmp_path, monkeypatch): from hyperloom.orchestrator.actions.executors.profile import ProfileExecutor monkeypatch.setenv(evidence.ENABLE_ENV, "0") + monkeypatch.delenv("PYTHONPATH", raising=False) config = tmp_path / "profile.yaml" _write_profile_config(config, {"TP": 8}) assert ProfileExecutor()._inject_host_probe(config, tmp_path / "ws") == "" diff --git a/src/hyperloom/inference_optimizer/tests/test_workload_envs_branches_unit.py b/src/hyperloom/inference_optimizer/tests/test_workload_envs_branches_unit.py index ff6dee9fd7..3534b826ba 100644 --- a/src/hyperloom/inference_optimizer/tests/test_workload_envs_branches_unit.py +++ b/src/hyperloom/inference_optimizer/tests/test_workload_envs_branches_unit.py @@ -623,6 +623,54 @@ def test_profile_manual_max_iters_above_cap_warns(monkeypatch, tmp_path, caplog) assert any("exceeds the serialization-safe cap" in r.message for r in caplog.records) +def _arm_sglang_shape_tool(monkeypatch, tmp_path) -> Path: + """Route the SGLang profile through the kernel_shape_tool path and return the tool dir.""" + from hyperloom.orchestrator.actions.executors import _server_patcher + + tool_dir = tmp_path / "kernel_shape_tool" + tool_dir.mkdir() + _clear_env(monkeypatch) + monkeypatch.setenv("INFERENCE_OPTIMIZER_DISABLE_TP_CLAMP", "1") + monkeypatch.setenv("HYPERLOOM_ENABLE_PATCH", "0") + monkeypatch.delenv("HYPERLOOM_PROFILE_SHAPE_DISCOVERY", raising=False) + monkeypatch.setattr(_server_patcher, "resolve_sglang_shape_mode", lambda: "sitecustomize") + monkeypatch.setattr(_server_patcher, "kernel_shape_tool_dir", lambda: tool_dir) + return tool_dir + + +def test_profile_shape_tool_keeps_the_inherited_pythonpath(monkeypatch, tmp_path): + # Magpie replaces the launch env's PYTHONPATH with the config's, so a config without one must carry the + # inherited entries, or the profile server loses the operator's and the image's import paths. + import os + + from hyperloom.orchestrator.actions.executors import _framework_rewrite_evidence as evidence + from hyperloom.orchestrator.actions.executors.profile import ProfileExecutor + + tool_dir = _arm_sglang_shape_tool(monkeypatch, tmp_path) + monkeypatch.setenv("PYTHONPATH", os.pathsep.join(["/a", "/b"])) + src = _write(tmp_path / "cfg.yaml", envs={"PROFILE": "1"}) + out = we.materialize_config_with_envs(src, tmp_path / "out") + envs = yaml.safe_load(out.read_text())["benchmark"]["envs"] + assert envs["TRACELENS_SHAPE_DISCOVERY"] == "1" + assert envs["PYTHONPATH"] == os.pathsep.join([str(tool_dir), "/a", "/b"]) + + # The profile shim lands in front of the tool, keeping the inherited tail. + ProfileExecutor()._inject_host_probe(out, tmp_path / "ws") + envs = yaml.safe_load(out.read_text())["benchmark"]["envs"] + assert envs["PYTHONPATH"] == os.pathsep.join([str(evidence.probe_asset_dir()), str(tool_dir), "/a", "/b"]) + + +def test_profile_shape_tool_prepends_to_the_config_pythonpath(monkeypatch, tmp_path): + # A config that carries PYTHONPATH is authoritative: the inherited value is not appended to it. + import os + + tool_dir = _arm_sglang_shape_tool(monkeypatch, tmp_path) + monkeypatch.setenv("PYTHONPATH", os.pathsep.join(["/a", "/b"])) + src = _write(tmp_path / "cfg.yaml", envs={"PROFILE": "1", "PYTHONPATH": "/c"}) + bench = _materialize(src, tmp_path / "out") + assert bench["envs"]["PYTHONPATH"] == os.pathsep.join([str(tool_dir), "/c"]) + + def test_quality_ref_variant_compares(monkeypatch, tmp_path): # A non-baseline scriptable variant must COMPARE against the operator reference and must NOT write. _clear_env(monkeypatch) diff --git a/src/hyperloom/orchestrator/actions/executors/_workload_envs.py b/src/hyperloom/orchestrator/actions/executors/_workload_envs.py index 5de77f5a20..3b18cab74d 100644 --- a/src/hyperloom/orchestrator/actions/executors/_workload_envs.py +++ b/src/hyperloom/orchestrator/actions/executors/_workload_envs.py @@ -756,6 +756,19 @@ def _resolve_framework_repo_path( return "" +def prepend_launch_pythonpath(envs: dict[str, Any], *entries: str) -> None: + """Put ``entries`` first on the ``PYTHONPATH`` the materialized config hands the server launch. + + Magpie launches with ``os.environ`` overlaid by ``benchmark.envs``, so a config ``PYTHONPATH`` replaces the + inherited one. A config without the key therefore builds on the inherited value, or prepending would drop the + operator's and the image's import paths; a config with the key stays authoritative. An entry already on the + path moves to the front, and empty parts are dropped. + """ + base = envs["PYTHONPATH"] if "PYTHONPATH" in envs else os.environ.get("PYTHONPATH") + parts = [*entries, *str(base or "").strip().split(os.pathsep)] + envs["PYTHONPATH"] = os.pathsep.join(dict.fromkeys(part for part in parts if part)) + + #: Source-root envs Hyperloom writes into a vLLM launch; vLLM reports them as unknown. VLLM_SOURCE_ROOT_ENVS: tuple[str, ...] = ("FRAMEWORK_REPO_PATH", "VLLM_REPO_PATH", "VLLM_DIR") @@ -1792,8 +1805,7 @@ def materialize_config_with_envs( extra_body.setdefault("detailed_annotations", True) _tool_dir = kernel_shape_tool_dir() if _shape_disc and _tool_dir is not None: - _existing_pp = str(envs.get("PYTHONPATH", "")).strip() - envs["PYTHONPATH"] = f"{_tool_dir}{os.pathsep}{_existing_pp}" if _existing_pp else str(_tool_dir) + prepend_launch_pythonpath(envs, str(_tool_dir)) envs["TRACELENS_SHAPE_DISCOVERY"] = "1" else: envs["TRACELENS_SHAPE_DISCOVERY"] = "0" diff --git a/src/hyperloom/orchestrator/actions/executors/profile.py b/src/hyperloom/orchestrator/actions/executors/profile.py index c2c10ab100..f3c3834bb8 100644 --- a/src/hyperloom/orchestrator/actions/executors/profile.py +++ b/src/hyperloom/orchestrator/actions/executors/profile.py @@ -848,6 +848,7 @@ def _inject_host_probe(self, config_path: Path, output_dir: Path) -> str: or ``""`` when the probe is not armed. """ from . import _framework_rewrite_evidence as _evidence + from ._workload_envs import prepend_launch_pythonpath asset_dir = _evidence.probe_asset_dir() if not asset_dir.is_dir(): @@ -868,10 +869,9 @@ def _inject_host_probe(self, config_path: Path, output_dir: Path) -> str: envs = bench.setdefault("envs", {}) if not isinstance(envs, dict): return "" - current = str(envs.get("PYTHONPATH", "") or "").strip() - entry = str(asset_dir) - if entry not in current.split(os.pathsep): - envs["PYTHONPATH"] = f"{entry}{os.pathsep}{current}" if current else entry + # First on the path even when already present: Python imports only the first sitecustomize, and the + # shim's chains to the next one. + prepend_launch_pythonpath(envs, str(asset_dir)) for key, value in probe_env.items(): envs[key] = value try: From 1b85d7451933b27ddda52f84080e3f08a96dc242 Mon Sep 17 00:00:00 2001 From: Zeng Date: Sun, 11 Oct 2026 13:57:16 +0800 Subject: [PATCH 2/3] fix(executors): keep the inherited PYTHONPATH under the reference overlay The reference recipe's overlay_pythonpath was prepended to the config's PYTHONPATH only. Magpie and the bypass launcher overlay benchmark.envs on os.environ, so a config without PYTHONPATH launched with the overlay as its whole import path and lost the operator's and the image's entries. The join also hardcoded ":" instead of os.pathsep. The PR's prepend_launch_pythonpath already owns this rule for the profile sites (config value when the key is present, inherited value otherwise, new entries first, deduplicated, empty parts dropped). It becomes prepend_launch_path(envs, var, *entries) so the same rule serves PATH and LD_LIBRARY_PATH too; the shape-tool and shim sites move to it and the reference overlay goes through it. Overlay validation is unchanged. Co-Authored-By: Claude --- .../tests/test_workload_envs_branches_unit.py | 21 ++++++++++++++++ .../actions/executors/_workload_envs.py | 24 +++++++++---------- .../orchestrator/actions/executors/profile.py | 4 ++-- 3 files changed, 35 insertions(+), 14 deletions(-) diff --git a/src/hyperloom/inference_optimizer/tests/test_workload_envs_branches_unit.py b/src/hyperloom/inference_optimizer/tests/test_workload_envs_branches_unit.py index 3534b826ba..b1a55da9f7 100644 --- a/src/hyperloom/inference_optimizer/tests/test_workload_envs_branches_unit.py +++ b/src/hyperloom/inference_optimizer/tests/test_workload_envs_branches_unit.py @@ -671,6 +671,27 @@ def test_profile_shape_tool_prepends_to_the_config_pythonpath(monkeypatch, tmp_p assert bench["envs"]["PYTHONPATH"] == os.pathsep.join([str(tool_dir), "/c"]) +@pytest.mark.parametrize( + ("config_envs", "tail"), + [({}, ["/a", "/b"]), ({"PYTHONPATH": "/c"}, ["/c"])], + ids=["inherited", "config-authoritative"], +) +def test_reference_overlay_builds_on_the_launch_pythonpath(monkeypatch, tmp_path, config_envs, tail): + # The reference overlay is prepended like every launch path entry: on the inherited value when the config has + # no PYTHONPATH, on the config's otherwise. The overlay is relative so it carries no drive colon on any host. + import os + + _clear_env(monkeypatch) + (tmp_path / "overlay").mkdir() + (tmp_path / "overlay" / "sitecustomize.py").write_text("", encoding="utf-8") + monkeypatch.chdir(tmp_path) + monkeypatch.setattr(we, "resolve_reference_launch", lambda: ("", {}, {"overlay_pythonpath": "overlay"})) + monkeypatch.setenv("PYTHONPATH", os.pathsep.join(["/a", "/b"])) + src = _write(tmp_path / "cfg.yaml", envs=config_envs) + bench = _materialize(src, tmp_path / "out") + assert bench["envs"]["PYTHONPATH"] == os.pathsep.join(["overlay", *tail]) + + def test_quality_ref_variant_compares(monkeypatch, tmp_path): # A non-baseline scriptable variant must COMPARE against the operator reference and must NOT write. _clear_env(monkeypatch) diff --git a/src/hyperloom/orchestrator/actions/executors/_workload_envs.py b/src/hyperloom/orchestrator/actions/executors/_workload_envs.py index 3b18cab74d..3f6f41f2fe 100644 --- a/src/hyperloom/orchestrator/actions/executors/_workload_envs.py +++ b/src/hyperloom/orchestrator/actions/executors/_workload_envs.py @@ -756,17 +756,18 @@ def _resolve_framework_repo_path( return "" -def prepend_launch_pythonpath(envs: dict[str, Any], *entries: str) -> None: - """Put ``entries`` first on the ``PYTHONPATH`` the materialized config hands the server launch. - - Magpie launches with ``os.environ`` overlaid by ``benchmark.envs``, so a config ``PYTHONPATH`` replaces the - inherited one. A config without the key therefore builds on the inherited value, or prepending would drop the - operator's and the image's import paths; a config with the key stays authoritative. An entry already on the - path moves to the front, and empty parts are dropped. +def prepend_launch_path(envs: dict[str, Any], var: str, *entries: str) -> None: + """Put ``entries`` first on the path-list env ``var`` the materialized config hands the server launch. + + Magpie and the bypass launcher start the server with ``os.environ`` overlaid by ``benchmark.envs``, so a config + value replaces the inherited one. A config without the key therefore builds on the inherited value, or + prepending would drop the operator's and the image's entries (import paths, ``PATH``, ``/opt/rocm/lib``); a + config with the key stays authoritative. An entry already on the path moves to the front, and empty parts are + dropped. """ - base = envs["PYTHONPATH"] if "PYTHONPATH" in envs else os.environ.get("PYTHONPATH") + base = envs[var] if var in envs else os.environ.get(var) parts = [*entries, *str(base or "").strip().split(os.pathsep)] - envs["PYTHONPATH"] = os.pathsep.join(dict.fromkeys(part for part in parts if part)) + envs[var] = os.pathsep.join(dict.fromkeys(part for part in parts if part)) #: Source-root envs Hyperloom writes into a vLLM launch; vLLM reports them as unknown. @@ -1805,7 +1806,7 @@ def materialize_config_with_envs( extra_body.setdefault("detailed_annotations", True) _tool_dir = kernel_shape_tool_dir() if _shape_disc and _tool_dir is not None: - prepend_launch_pythonpath(envs, str(_tool_dir)) + prepend_launch_path(envs, "PYTHONPATH", str(_tool_dir)) envs["TRACELENS_SHAPE_DISCOVERY"] = "1" else: envs["TRACELENS_SHAPE_DISCOVERY"] = "0" @@ -1886,8 +1887,7 @@ def materialize_config_with_envs( if reference_controls.get("overlay_pythonpath"): from hyperloom.common.overlay import validate_overlay_pythonpath - overlay = validate_overlay_pythonpath(reference_controls["overlay_pythonpath"]) - envs["PYTHONPATH"] = overlay + (f":{envs['PYTHONPATH']}" if envs.get("PYTHONPATH") else "") + prepend_launch_path(envs, "PYTHONPATH", validate_overlay_pythonpath(reference_controls["overlay_pythonpath"])) recipe_overwritten = launcher_overwritten_envs(bench) if server_args: # Merge into (not overwrite) the framework env so the profile path's diff --git a/src/hyperloom/orchestrator/actions/executors/profile.py b/src/hyperloom/orchestrator/actions/executors/profile.py index f3c3834bb8..4f91f19ec7 100644 --- a/src/hyperloom/orchestrator/actions/executors/profile.py +++ b/src/hyperloom/orchestrator/actions/executors/profile.py @@ -848,7 +848,7 @@ def _inject_host_probe(self, config_path: Path, output_dir: Path) -> str: or ``""`` when the probe is not armed. """ from . import _framework_rewrite_evidence as _evidence - from ._workload_envs import prepend_launch_pythonpath + from ._workload_envs import prepend_launch_path asset_dir = _evidence.probe_asset_dir() if not asset_dir.is_dir(): @@ -871,7 +871,7 @@ def _inject_host_probe(self, config_path: Path, output_dir: Path) -> str: return "" # First on the path even when already present: Python imports only the first sitecustomize, and the # shim's chains to the next one. - prepend_launch_pythonpath(envs, str(asset_dir)) + prepend_launch_path(envs, "PYTHONPATH", str(asset_dir)) for key, value in probe_env.items(): envs[key] = value try: From af9d9099477043a0a99872f68e601ccbbb2244cd Mon Sep 17 00:00:00 2001 From: Zeng Date: Sun, 11 Oct 2026 13:57:17 +0800 Subject: [PATCH 3/3] fix(executors): keep inherited launch paths at every grid prepend site apply_runtime_override (path_prefix, entrypoint_bin_dir, pythonpath_prefix, pythonpath_prefixes, ld_library_path_prefix) and the variant overlay_pythonpath in _build_variant_yaml prepended onto the config value only. When the config carried no such key, the launch got a PATH holding a single directory, an LD_LIBRARY_PATH without /opt/rocm/lib, or a PYTHONPATH without the operator's and the image's import paths, because the launcher lets a config key replace the inherited value. Three of the sites also joined with ":" rather than os.pathsep. Every site now goes through prepend_launch_path, which builds on the inherited value when the config lacks the key and keeps the config authoritative when it has it; _prepend_path_entry is deleted. The per-site validation and warn-and-drop behaviour is unchanged. Behaviour change: an entry already on the path is moved to the front instead of being left where it was, and PYTHONPATH prefixes are now deduplicated. Attempt PATH, LD_LIBRARY_PATH and PYTHONPATH entries are meant to win, so moving them up is the intended order. Co-Authored-By: Claude --- .../test_grid_runner_helpers_coverage_unit.py | 22 ++++++ .../test_grid_runner_runtime_override.py | 75 ++++++++++++++++--- .../actions/executors/_grid_runner.py | 30 +++----- 3 files changed, 97 insertions(+), 30 deletions(-) diff --git a/src/hyperloom/inference_optimizer/tests/test_grid_runner_helpers_coverage_unit.py b/src/hyperloom/inference_optimizer/tests/test_grid_runner_helpers_coverage_unit.py index 25a662ce9b..66a15fdff2 100644 --- a/src/hyperloom/inference_optimizer/tests/test_grid_runner_helpers_coverage_unit.py +++ b/src/hyperloom/inference_optimizer/tests/test_grid_runner_helpers_coverage_unit.py @@ -337,6 +337,28 @@ def test_build_variant_yaml_prepends_legit_overlay(tmp_path: Path) -> None: assert envs.get("PYTHONPATH", "").startswith(str(overlay)) +@pytest.mark.parametrize( + ("config_envs", "tail"), + [({}, ["/inh/a", "/inh/b"]), ({"PYTHONPATH": "/cfg"}, ["/cfg"])], + ids=["inherited", "config-authoritative"], +) +def test_build_variant_yaml_overlay_builds_on_the_launch_pythonpath(tmp_path, monkeypatch, config_envs, tail) -> None: + # The launch overlays benchmark.envs on os.environ, so a config without PYTHONPATH carries the inherited one. + # The overlay is relative so the path carries no drive colon on any host. + import os + + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("PYTHONPATH", os.pathsep.join(["/inh/a", "/inh/b"])) + (tmp_path / "overlay").mkdir() + base = tmp_path / "base.yaml" + base.write_text(yaml.safe_dump({"benchmark": {"framework": "sglang", "envs": config_envs}}), encoding="utf-8") + variant = _variant("v1") + variant.overlay_pythonpath = "overlay" # type: ignore[attr-defined] + out = gr._build_variant_yaml(base, "", variant, output_subdir=tmp_path / "v1") + envs = yaml.safe_load(out.read_text(encoding="utf-8"))["benchmark"]["envs"] + assert envs["PYTHONPATH"] == os.pathsep.join(["overlay", *tail]) + + def test_build_variant_yaml_drops_unsafe_overlay(tmp_path: Path) -> None: # SWSPLAT-42358: a ``:``-joined / traversal / non-existent overlay is dropped (would otherwise smuggle extra # PYTHONPATH entries or escape). diff --git a/src/hyperloom/inference_optimizer/tests/test_grid_runner_runtime_override.py b/src/hyperloom/inference_optimizer/tests/test_grid_runner_runtime_override.py index f4528aeaae..9e2b565538 100644 --- a/src/hyperloom/inference_optimizer/tests/test_grid_runner_runtime_override.py +++ b/src/hyperloom/inference_optimizer/tests/test_grid_runner_runtime_override.py @@ -34,17 +34,46 @@ def test_path_prefix_prepended(): def test_path_prefix_not_duplicated(): - envs = {"PATH": "/attempt/bin:/usr/bin"} + envs = {"PATH": os.pathsep.join(["/attempt/bin", "/usr/bin"])} apply_runtime_override(envs, {"path_prefix": "/attempt/bin"}) assert envs["PATH"].count("/attempt/bin") == 1 +def test_path_prefix_moves_a_present_entry_to_the_front(): + envs = {"PATH": os.pathsep.join(["/usr/bin", "/attempt/bin"])} + apply_runtime_override(envs, {"path_prefix": "/attempt/bin"}) + assert envs["PATH"] == os.pathsep.join(["/attempt/bin", "/usr/bin"]) + + +def test_path_prefix_keeps_the_inherited_path(monkeypatch): + """The launch overlays ``benchmark.envs`` on ``os.environ``, so a config without PATH must carry the inherited one.""" + monkeypatch.setenv("PATH", os.pathsep.join(["/usr/local/bin", "/usr/bin"])) + envs: dict[str, str] = {} + apply_runtime_override(envs, {"path_prefix": "/attempt/bin"}) + assert envs["PATH"] == os.pathsep.join(["/attempt/bin", "/usr/local/bin", "/usr/bin"]) + + def test_pythonpath_prefix_prepended(tmp_path): envs: dict[str, str] = {} apply_runtime_override(envs, {"pythonpath_prefix": str(tmp_path)}) assert str(tmp_path) in envs.get("PYTHONPATH", "") +def test_pythonpath_prefix_keeps_the_inherited_pythonpath(monkeypatch): + monkeypatch.setenv("PYTHONPATH", os.pathsep.join(["/inh/a", "/inh/b"])) + envs: dict[str, str] = {} + apply_runtime_override(envs, {"pythonpath_prefix": "/attempt/pkg"}) + assert envs["PYTHONPATH"] == os.pathsep.join(["/attempt/pkg", "/inh/a", "/inh/b"]) + + +def test_pythonpath_prefix_prepends_to_the_config_pythonpath(monkeypatch): + """A config that carries PYTHONPATH is authoritative over the inherited value.""" + monkeypatch.setenv("PYTHONPATH", os.pathsep.join(["/inh/a", "/inh/b"])) + envs = {"PYTHONPATH": "/cfg"} + apply_runtime_override(envs, {"pythonpath_prefix": "/attempt/pkg"}) + assert envs["PYTHONPATH"] == os.pathsep.join(["/attempt/pkg", "/cfg"]) + + def test_pythonpath_unsafe_colon_dropped(): envs: dict[str, str] = {} apply_runtime_override(envs, {"pythonpath_prefix": "/a:/b"}) @@ -125,28 +154,46 @@ def test_no_runtime_override_field_is_noop(tmp_path): # compiled-artifact runtime prefixes -def test_pythonpath_prefixes_multi_entry_ordered_prepended(): +def test_pythonpath_prefixes_multi_entry_ordered_prepended(monkeypatch): + monkeypatch.setenv("PYTHONPATH", "/inh/a") envs = {"PYTHONPATH": "/opt/venv"} apply_runtime_override(envs, {"pythonpath_prefixes": ["/a/pkg", "/b/pkg"]}) - assert envs["PYTHONPATH"] == "/a/pkg:/b/pkg:/opt/venv" + assert envs["PYTHONPATH"] == os.pathsep.join(["/a/pkg", "/b/pkg", "/opt/venv"]) -def test_pythonpath_prefixes_drops_unsafe_entry(): +def test_pythonpath_prefixes_keep_the_inherited_pythonpath(monkeypatch): + monkeypatch.setenv("PYTHONPATH", os.pathsep.join(["/inh/a", "/inh/b"])) + envs: dict[str, str] = {} + apply_runtime_override(envs, {"pythonpath_prefixes": ["/a/pkg", "/b/pkg"]}) + assert envs["PYTHONPATH"] == os.pathsep.join(["/a/pkg", "/b/pkg", "/inh/a", "/inh/b"]) + + +def test_pythonpath_prefixes_drops_unsafe_entry(monkeypatch): + monkeypatch.delenv("PYTHONPATH", raising=False) envs = {} apply_runtime_override(envs, {"pythonpath_prefixes": ["/good", "../evil"]}) assert envs.get("PYTHONPATH") == "/good" -def test_ld_library_path_prefix_prepends_and_preserves_rocm(): +def test_ld_library_path_prefix_prepends_and_preserves_rocm(monkeypatch): + monkeypatch.setenv("LD_LIBRARY_PATH", "/inh/lib") envs = {"LD_LIBRARY_PATH": "/opt/rocm/lib"} apply_runtime_override(envs, {"ld_library_path_prefix": ["/attempt/lib"]}) - assert envs["LD_LIBRARY_PATH"] == "/attempt/lib:/opt/rocm/lib" + assert envs["LD_LIBRARY_PATH"] == os.pathsep.join(["/attempt/lib", "/opt/rocm/lib"]) + + +def test_ld_library_path_prefix_keeps_the_inherited_rocm_lib(monkeypatch): + monkeypatch.setenv("LD_LIBRARY_PATH", "/opt/rocm/lib") + envs: dict[str, str] = {} + apply_runtime_override(envs, {"ld_library_path_prefix": ["/attempt/lib"]}) + assert envs["LD_LIBRARY_PATH"] == os.pathsep.join(["/attempt/lib", "/opt/rocm/lib"]) -def test_ld_library_path_prefix_multi_entry_order(): +def test_ld_library_path_prefix_multi_entry_order(monkeypatch): + monkeypatch.delenv("LD_LIBRARY_PATH", raising=False) envs = {} apply_runtime_override(envs, {"ld_library_path_prefix": ["/a", "/b"]}) - assert envs["LD_LIBRARY_PATH"] == "/a:/b" + assert envs["LD_LIBRARY_PATH"] == os.pathsep.join(["/a", "/b"]) def test_runtime_env_merged(): @@ -162,10 +209,18 @@ def test_runtime_env_rejects_reserved_and_blocked_keys(): assert envs == {"OK_KEY": "v"} -def test_entrypoint_bin_dir_prepended_to_path(): +def test_entrypoint_bin_dir_prepended_to_path(monkeypatch): + monkeypatch.setenv("PATH", "/inh/bin") envs = {"PATH": "/usr/bin"} apply_runtime_override(envs, {"entrypoint_bin_dir": "/attempt/bin"}) - assert envs["PATH"] == "/attempt/bin:/usr/bin" + assert envs["PATH"] == os.pathsep.join(["/attempt/bin", "/usr/bin"]) + + +def test_entrypoint_bin_dir_keeps_the_inherited_path(monkeypatch): + monkeypatch.setenv("PATH", os.pathsep.join(["/usr/local/bin", "/usr/bin"])) + envs: dict[str, str] = {} + apply_runtime_override(envs, {"entrypoint_bin_dir": "/attempt/bin"}) + assert envs["PATH"] == os.pathsep.join(["/attempt/bin", "/usr/local/bin", "/usr/bin"]) def test_entrypoint_bin_dir_unsafe_dropped(): diff --git a/src/hyperloom/orchestrator/actions/executors/_grid_runner.py b/src/hyperloom/orchestrator/actions/executors/_grid_runner.py index ef72bcd368..8ef7cc98dc 100644 --- a/src/hyperloom/orchestrator/actions/executors/_grid_runner.py +++ b/src/hyperloom/orchestrator/actions/executors/_grid_runner.py @@ -323,14 +323,6 @@ def _is_safe_path_entry(entry: str) -> bool: ) -def _prepend_path_entry(envs: dict[str, str], var: str, entry: str) -> None: - """Prepend a single already-validated entry onto a ``:``-joined env var.""" - parts = [p for p in str(envs.get(var, "") or "").split(":") if p] - if entry not in parts: - parts.insert(0, entry) - envs[var] = ":".join(parts) - - # Reserved keys carried by dedicated override fields; never accepted via runtime_env. _RUNTIME_ENV_RESERVED: frozenset[str] = frozenset({"PATH", "PYTHONPATH", "LD_LIBRARY_PATH"}) @@ -339,35 +331,34 @@ def apply_runtime_override(envs: dict[str, str], override: dict[str, Any]) -> No """Inject an attempt runtime override into the materialized YAML envs dict.""" if not override: return + from ._workload_envs import prepend_launch_path + path_prefix = str(override.get("path_prefix") or "").strip() if path_prefix: - _prepend_path_entry(envs, "PATH", path_prefix) + prepend_launch_path(envs, "PATH", path_prefix) entrypoint_bin = str(override.get("entrypoint_bin_dir") or "").strip() if entrypoint_bin: if _is_safe_path_entry(entrypoint_bin): - _prepend_path_entry(envs, "PATH", entrypoint_bin) + prepend_launch_path(envs, "PATH", entrypoint_bin) else: log.warning("apply_runtime_override: dropping unsafe entrypoint_bin_dir %r", entrypoint_bin) pp_prefix = str(override.get("pythonpath_prefix") or "").strip() if pp_prefix: if _is_safe_path_entry(pp_prefix): - _cur_pp = str(envs.get("PYTHONPATH", "") or "") - envs["PYTHONPATH"] = f"{pp_prefix}:{_cur_pp}" if _cur_pp else pp_prefix + prepend_launch_path(envs, "PYTHONPATH", pp_prefix) else: log.warning("apply_runtime_override: dropping unsafe pythonpath_prefix %r", pp_prefix) # Multi-entry prefixes: prepend in order so the first entry wins, ahead of any single-dir pythonpath_prefix and # the inherited value. for entry in reversed(_coerce_prefix_list(override.get("pythonpath_prefixes"))): if _is_safe_path_entry(entry): - _cur_pp = str(envs.get("PYTHONPATH", "") or "") - envs["PYTHONPATH"] = f"{entry}:{_cur_pp}" if _cur_pp else entry + prepend_launch_path(envs, "PYTHONPATH", entry) else: log.warning("apply_runtime_override: dropping unsafe pythonpath entry %r", entry) - # Native loader path: prepend attempt entries while preserving any inherited entries (e.g. /opt/rocm/lib), which - # _prepend_path_entry never drops. + # Native loader path: prepend attempt entries while preserving any inherited entries (e.g. /opt/rocm/lib). for entry in reversed(_coerce_prefix_list(override.get("ld_library_path_prefix"))): if _is_safe_path_entry(entry): - _prepend_path_entry(envs, "LD_LIBRARY_PATH", entry) + prepend_launch_path(envs, "LD_LIBRARY_PATH", entry) else: log.warning("apply_runtime_override: dropping unsafe ld_library_path entry %r", entry) _runtime_env = override.get("runtime_env") @@ -474,7 +465,7 @@ def _build_variant_yaml( envs.pop(str(k), None) for k, v in variant.extra_envs.items(): envs[str(k)] = str(v) - from ._workload_envs import pin_mlperf_round_concurrency + from ._workload_envs import pin_mlperf_round_concurrency, prepend_launch_path pin_mlperf_round_concurrency(envs) # The launcher re-exports these unconditionally, so a value carried here is @@ -496,8 +487,7 @@ def _build_variant_yaml( and Path(_overlay).is_dir() ) if _overlay_ok: - _cur_pp = str(envs.get("PYTHONPATH", "") or "") - envs["PYTHONPATH"] = f"{_overlay}:{_cur_pp}" if _cur_pp else _overlay + prepend_launch_path(envs, "PYTHONPATH", _overlay) else: log.warning( "grid: dropping unsafe overlay_pythonpath %r (not a single "