# Copyright Advanced Micro Devices, Inc., or its affiliates. # SPDX-License-Identifier: MIT """Tests for optimize.py CLI interface. Usage: python3 -m pytest tests/test_optimize.py \\ --hipblaslt-path /path/to/hipblaslt \\ --workload /path/to/workload.yaml \\ [--hw gfx942] ``--hw`` sets ``configure.py --architecture`` (default: gfx950). Hint: add -rs to see skip reasons, or -v for verbose output. ``TestPipelineFast`` exercises ``geko.pipeline.run_optimize`` with Tensile/GPU work mocked; it runs in CI without hipBLASLt. ``TestIntegration`` runs configure then real optimize (long-running; needs a built hipBLASLt and GPU). """ import json import subprocess import sys import warnings from pathlib import Path from unittest.mock import MagicMock, patch import pytest import yaml from geko import pipeline as pipeline_mod from geko.pipeline import run_optimize ROOT = Path(__file__).resolve().parents[1] # Configure step matches test_configure.py; optimize can take much longer (Tensile / GA). _CONFIGURE_TIMEOUT_S = 600 _OPTIMIZE_TIMEOUT_S = 7200 def _patch_optimize_config(path, n_gen_target=10, pop_size_target=128): """Rewrite one generated optimization config with faster GA settings. Config format is stable: Backend -> Config -> {n_gen, pop_size}. """ with path.open("r") as f: data = yaml.safe_load(f) if data is None: return backend_cfg = data.setdefault("Backend", {}).setdefault("Config", {}) backend_cfg["n_gen"] = n_gen_target backend_cfg["pop_size"] = pop_size_target with path.open("w") as f: yaml.safe_dump(data, f, sort_keys=False) # --------------------------------------------------------------------------- # Fast pipeline test (mocked optim / merge / analyze) # --------------------------------------------------------------------------- class TestPipelineFast: """Call ``run_optimize`` with heavy steps stubbed; no real tuning or benchmarks.""" def test_main_completes_with_mocks(self, tmp_path): """Full control-flow: validation, state flags, mocked run → merge → analyze.""" hip = tmp_path / "hipblaslt_stub" hip.mkdir() workdir = tmp_path / "workdir" (workdir / "optimizations").mkdir(parents=True) from geko.schemas import RunState dummy_input = tmp_path / "workload_stub.yaml" dummy_input.write_text("# stub\n") state = RunState.create(dummy_input) state.configured = True state.dump(workdir / "run_state.json") merged = MagicMock() def _dump_libs(target): p = Path(target) p.mkdir(parents=True, exist_ok=True) (p / "merged_stub.yaml").write_text( "library_name: stub\nlibrary_type: tensile\nsolutions: []\nsizes: []\n" ) merged.dump.side_effect = _dump_libs mock_run = MagicMock() mock_analyze = MagicMock(return_value=None) mock_merge = MagicMock(return_value=merged) with patch.object(pipeline_mod.optim, "run", mock_run), patch.object( pipeline_mod.optim.utils, "check_progress", side_effect=[(1, 1, 0), (1, 1, 0)], ), patch.object(pipeline_mod.library, "merge_solutions", mock_merge), patch.object( pipeline_mod.optim, "analyze", mock_analyze, ): run_optimize(str(hip), workdir=str(workdir), devices=[0], verbose=0) mock_run.assert_called_once() ca = mock_run.call_args assert ca.args[0] == Path(hip) assert ca.args[1] == workdir / "optimizations" assert ca.kwargs["devices"] == [0] mock_merge.assert_called_once() mock_analyze.assert_called_once() with (workdir / "run_state.json").open() as f: final_state = json.load(f) assert final_state["configured"] is True assert final_state["optimized"] is True # --------------------------------------------------------------------------- # End-to-end integration test # --------------------------------------------------------------------------- @pytest.mark.slow class TestIntegration: """Run configure.py then optimize.py against a real workload and hipBLASLt repo. Usage: python3 -m pytest tests/test_optimize.py \\ --hipblaslt-path ~/rocm-libraries/projects/hipblaslt \\ --workload tests/test_data/workload.yaml \\ [--hw gfx942] """ def test_full_cli_run(self, hipblaslt_path, workload_path, hw_arch, tmp_path): """Configure then optimize subprocesses complete (uses conftest CLI options).""" if hipblaslt_path is None or workload_path is None: msg = "Skipped: requires --hipblaslt-path and --workload CLI options" warnings.warn(msg, stacklevel=1) pytest.skip(msg) hip = Path(hipblaslt_path) workload = Path(workload_path) if not hip.is_dir(): msg = f"Skipped: hipBLASLt path not found: {hip}" warnings.warn(msg, stacklevel=1) pytest.skip(msg) if not workload.is_file(): msg = f"Skipped: workload file not found: {workload}" warnings.warn(msg, stacklevel=1) pytest.skip(msg) workdir = tmp_path / "workdir" client_build = (tmp_path / "build_tmp").resolve() run_cwd = str(tmp_path.resolve()) # Absolute paths so subprocess cwd=tmp_path does not break relative CLI args. configure_py = str((ROOT / "scripts" / "configure.py").resolve()) optimize_py = str((ROOT / "scripts" / "optimize.py").resolve()) hip_abs = str(hip.resolve()) workload_abs = str(workload.resolve()) workdir_abs = str(workdir.resolve()) cfg = subprocess.run( [ sys.executable, configure_py, workload_abs, "--hipblaslt", hip_abs, "--workdir", workdir_abs, "--architecture", hw_arch, ], cwd=run_cwd, capture_output=True, text=True, timeout=_CONFIGURE_TIMEOUT_S, ) assert cfg.returncode == 0, ( f"configure.py exited with code {cfg.returncode}\n" f"--- stdout ---\n{cfg.stdout[-2000:]}\n" f"--- stderr ---\n{cfg.stderr[-2000:]}" ) assert workdir.is_dir() assert (workdir / "optimizations").is_dir() config_files = pipeline_mod.optim.utils.list_optimization_configs(workdir / "optimizations") assert len(config_files) > 0, "Expected config files but workdir is empty" for f in config_files: # Patch config so it runs faster, independent of original defaults. _patch_optimize_config(Path(f), n_gen_target=10, pop_size_target=128) opt = subprocess.run( [ sys.executable, optimize_py, "--hipblaslt", hip_abs, "--workdir", workdir_abs, "--client_build_dir", str(client_build), "--devices", "0", "--n_slots", "1", ], cwd=run_cwd, capture_output=True, text=True, timeout=_OPTIMIZE_TIMEOUT_S, ) assert opt.returncode == 0, ( f"optimize.py exited with code {opt.returncode}\n" f"--- stdout ---\n{opt.stdout[-2000:]}\n" f"--- stderr ---\n{opt.stderr[-2000:]}" ) state_path = workdir / "run_state.json" assert state_path.is_file() with state_path.open() as f: state = json.load(f) assert state.get("configured") is True assert state.get("optimized") is True generated = list(workdir.rglob("*")) assert len(generated) > 0, "Expected output files but workdir is empty"