diff --git a/src/perf_helper_materialization.py b/src/perf_helper_materialization.py index 1ce7d04d..c3cef62b 100644 --- a/src/perf_helper_materialization.py +++ b/src/perf_helper_materialization.py @@ -44,7 +44,7 @@ ) MARK_STARTS = (MARK_START, OLD_MARK_START, LEGACY_MARK_START) MARK_END = "# <<< AKA-GENERATED <<<" -FUNC_ANCHOR = "def _measure_cuda_event_fallback(" +VLLM_BLOCK_ANCHOR = "import sys as _aka_sys" VLLM_HELPER_SYMBOLS = ("_measure_cuda_event_fallback", "_benchmark_cuda_graph_or_events") ROCMBENCH_HELPER_STUB = '''"""Generated at workspace setup from src/tools/perf/performance_utils_pytest.py. @@ -129,7 +129,7 @@ def canonical_aka_helper(root: Path = ROOT) -> str: def canonical_vllm_block(root: Path = ROOT) -> str: text = (root / "src/tools/perf/vllm_cuda_graph_block.py").read_text() - return text[text.index(FUNC_ANCHOR) :] + return text[text.index(VLLM_BLOCK_ANCHOR) :] def replace_marked_region(current: str, block: str) -> str | None: diff --git a/src/tools/perf/vllm_cuda_graph_block.py b/src/tools/perf/vllm_cuda_graph_block.py index 316cfe0d..8b068c8a 100644 --- a/src/tools/perf/vllm_cuda_graph_block.py +++ b/src/tools/perf/vllm_cuda_graph_block.py @@ -4,6 +4,17 @@ in the sibling materialized ``_aka_benchmark.py`` module. """ +import sys as _aka_sys +from pathlib import Path as _AkaPath + + +# Task-provided forge drivers are copied to the workspace root and load this +# runner by file path. Python does not add a file-loaded module's directory to +# sys.path, so make the sibling materialized helper discoverable explicitly. +_AKA_HELPER_DIR = str(_AkaPath(__file__).resolve().parent) +if _AKA_HELPER_DIR not in _aka_sys.path: + _aka_sys.path.insert(0, _AKA_HELPER_DIR) + def _measure_cuda_event_fallback(fn, repetition, prepare_fn=None): try: diff --git a/tests/test_perf_helper_materialization.py b/tests/test_perf_helper_materialization.py index 910ef247..9dc52278 100644 --- a/tests/test_perf_helper_materialization.py +++ b/tests/test_perf_helper_materialization.py @@ -57,6 +57,45 @@ def test_materializes_vllm_adapter_and_sibling_helper(tmp_path): assert (scripts / AKA_HELPER_FILE_NAME).read_text() == canonical_aka_helper(ROOT) +def test_file_loaded_vllm_runner_finds_sibling_helper_from_workspace_root(tmp_path): + scripts = tmp_path / "scripts" + scripts.mkdir() + runner = scripts / "task_runner.py" + runner.write_text(f"{MARK_START}\n{VLLM_HELPER_STUB_BLOCK}{MARK_END}\n") + (tmp_path / "config.yaml").write_text( + "performance_command:\n - python3 scripts/task_runner.py performance\n" + ) + materialize_perf_helpers_in_workspace(tmp_path) + + # Match forge's invocation shape: forge_driver.py runs at the workspace + # root and imports scripts/task_runner.py with spec_from_file_location(). + probe = """ +import importlib.util +from pathlib import Path + +runner = Path("scripts/task_runner.py").resolve() +scripts = str(runner.parent) +assert scripts not in __import__("sys").path +spec = importlib.util.spec_from_file_location("_forge_task_runner", runner) +module = importlib.util.module_from_spec(spec) +spec.loader.exec_module(module) +helper = importlib.util.find_spec("_aka_benchmark") +assert helper is not None +print(Path(helper.origin).resolve()) +""" + result = subprocess.run( + [sys.executable, "-S", "-c", probe], + cwd=tmp_path, + capture_output=True, + text=True, + check=False, + env={"PYTHONPATH": ""}, + ) + + assert result.returncode == 0, result.stderr or result.stdout + assert Path(result.stdout.strip()) == (scripts / AKA_HELPER_FILE_NAME).resolve() + + def test_materializes_helper_beside_eval_tools_entrypoint(tmp_path): eval_tools = tmp_path / "eval_tools" eval_tools.mkdir()