From 8b3b96e3dca9f4c076ec8a29ee44bc2f3670295f Mon Sep 17 00:00:00 2001 From: Vincent Ouyang Date: Fri, 7 Aug 2026 08:10:16 +0000 Subject: [PATCH 01/11] Add targeted kernel trace evidence and quality receipts --- Magpie/benchmark_config.yaml.example | 13 + Magpie/main.py | 19 + Magpie/modes/benchmark/__init__.py | 3 +- Magpie/modes/benchmark/benchmarker.py | 82 ++- Magpie/modes/benchmark/config.py | 51 ++ Magpie/modes/benchmark/quality.py | 197 +++++ Magpie/modes/benchmark/result.py | 38 + Magpie/modes/benchmark/targeted_trace.py | 54 ++ Magpie/modes/benchmark/workspace.py | 27 +- Magpie/scripts/benchmark/atom_mi300x.sh | 5 +- Magpie/scripts/benchmark/atom_mi355x.sh | 5 +- .../benchmark/magpie_bench_remote_compat.sh | 56 ++ Magpie/scripts/benchmark/sglang_mi300x.sh | 6 +- Magpie/scripts/benchmark/sglang_mi355x.sh | 5 +- Magpie/scripts/benchmark/vllm_mi300x.sh | 6 +- Magpie/scripts/benchmark/vllm_mi355x.sh | 5 +- Magpie/targeted_trace/README.md | 125 ++++ Magpie/targeted_trace/__init__.py | 44 ++ Magpie/targeted_trace/capture.py | 311 ++++++++ Magpie/targeted_trace/cli.py | 106 +++ Magpie/targeted_trace/config.py | 116 +++ Magpie/targeted_trace/postprocess.py | 327 +++++++++ Magpie/targeted_trace/sampling.py | 43 ++ Magpie/targeted_trace/schema.py | 690 ++++++++++++++++++ Magpie/targeted_trace/serialization.py | 215 ++++++ Magpie/targeted_trace/torch_profiler.py | 398 ++++++++++ Magpie/targeted_trace/writer.py | 331 +++++++++ README.md | 4 + docs/how-to/benchmarking/benchmark.md | 22 +- .../benchmark_vllm_qwen3_next_80b_fp8.yaml | 7 + pyproject.toml | 8 +- .../fixtures/targeted_trace/torch_trace.json | 47 ++ tests/test_benchmark_support.py | 157 ++++ tests/test_targeted_trace.py | 588 +++++++++++++++ 34 files changed, 4096 insertions(+), 15 deletions(-) create mode 100644 Magpie/modes/benchmark/quality.py create mode 100644 Magpie/modes/benchmark/targeted_trace.py create mode 100644 Magpie/targeted_trace/README.md create mode 100644 Magpie/targeted_trace/__init__.py create mode 100644 Magpie/targeted_trace/capture.py create mode 100644 Magpie/targeted_trace/cli.py create mode 100644 Magpie/targeted_trace/config.py create mode 100644 Magpie/targeted_trace/postprocess.py create mode 100644 Magpie/targeted_trace/sampling.py create mode 100644 Magpie/targeted_trace/schema.py create mode 100644 Magpie/targeted_trace/serialization.py create mode 100644 Magpie/targeted_trace/torch_profiler.py create mode 100644 Magpie/targeted_trace/writer.py create mode 100644 tests/fixtures/targeted_trace/torch_trace.json create mode 100644 tests/test_targeted_trace.py diff --git a/Magpie/benchmark_config.yaml.example b/Magpie/benchmark_config.yaml.example index 439f783..abd6b2a 100644 --- a/Magpie/benchmark_config.yaml.example +++ b/Magpie/benchmark_config.yaml.example @@ -20,6 +20,10 @@ benchmark: # Run mode: docker (default), local, or ray run_mode: docker + + # Evidence lane: profiler-free measurement or non-reward diagnostic. + # "auto" infers diagnostic when any heavyweight profiler is enabled. + run_kind: measurement # Environment variables (passed to container/process) envs: @@ -36,6 +40,15 @@ benchmark: enabled: false system_profiler: enabled: false + + # Diagnostic-only selected-kernel evidence. Requires torch_profiler=true. + targeted_trace: + enabled: false + backend: torch_profiler + run_seed: magpie-targeted-trace + sample_rate: 1.0 + max_records_per_shard: 100000 + targets: [] # Timeout (seconds) timeout_seconds: 3600 diff --git a/Magpie/main.py b/Magpie/main.py index ca71690..88d5b09 100644 --- a/Magpie/main.py +++ b/Magpie/main.py @@ -1045,6 +1045,7 @@ def run_benchmark(args, config: Dict[str, Any]) -> int: "framework": args.framework, "model": args.model, "precision": args.precision, + "run_kind": args.run_kind or "auto", "envs": { "TP": args.tp, "CONC": args.concurrency, @@ -1077,6 +1078,9 @@ def run_benchmark(args, config: Dict[str, Any]) -> int: run_mode = getattr(args, "run_mode", None) if run_mode: benchmark_cfg["run_mode"] = run_mode + run_kind = getattr(args, "run_kind", None) + if run_kind: + benchmark_cfg["run_kind"] = run_kind # Get benchmark settings from framework config bench_settings = config.get("benchmark", {}) @@ -1156,6 +1160,10 @@ def create_parser() -> argparse.ArgumentParser: subparsers = parser.add_subparsers(dest="mode", help="Evaluation mode") + from .targeted_trace.cli import add_targeted_trace_parser + + add_targeted_trace_parser(subparsers) + # Analyze subcommand analyze_parser = subparsers.add_parser( "analyze", help="Analyze kernel(s) - requires testcase" @@ -1268,6 +1276,13 @@ def create_parser() -> argparse.ArgumentParser: help="Execution mode: 'docker' (default) runs inside a container; " "'local' runs directly on the host (useful inside pods/containers)" ) + benchmark_parser.add_argument( + "--run-kind", + choices=["measurement", "diagnostic"], + default=None, + help="Evidence lane: measurement rejects profilers; diagnostic artifacts " + "are never reward eligible", + ) benchmark_parser.add_argument( "--docker-image", type=str, help="Override Docker image" ) @@ -1371,6 +1386,10 @@ def main() -> int: return run_compare(args, config) elif args.mode == "benchmark": return run_benchmark(args, config) + elif args.mode == "targeted-trace": + from .targeted_trace.cli import run_targeted_trace + + return run_targeted_trace(args) else: logger.error(f"Unknown mode: {args.mode}") return 1 diff --git a/Magpie/modes/benchmark/__init__.py b/Magpie/modes/benchmark/__init__.py index d647476..28ca44a 100644 --- a/Magpie/modes/benchmark/__init__.py +++ b/Magpie/modes/benchmark/__init__.py @@ -26,6 +26,7 @@ from .image_selector import ImageSelector from .inferencex import InferenceXManager, ensure_inferencex_available from .gap_analysis import GapAnalyzer, GapAnalysisResult +from ...targeted_trace.config import TargetedTraceConfig __all__ = [ "BenchmarkMode", @@ -36,6 +37,7 @@ "TorchProfilerConfig", "SystemProfilerConfig", "GapAnalysisConfig", + "TargetedTraceConfig", "GapAnalyzer", "GapAnalysisResult", "WorkspaceManager", @@ -44,4 +46,3 @@ "ensure_inferencex_available", ] - diff --git a/Magpie/modes/benchmark/benchmarker.py b/Magpie/modes/benchmark/benchmarker.py index ad937b8..15e903e 100644 --- a/Magpie/modes/benchmark/benchmarker.py +++ b/Magpie/modes/benchmark/benchmarker.py @@ -32,7 +32,9 @@ from .config import BenchmarkConfig from .image_selector import ImageSelector from .inferencex import ensure_inferencex_available +from .quality import parse_lm_eval_quality from .result import BenchmarkResult, LatencyMetrics, ResultParser, ThroughputMetrics +from .targeted_trace import run_targeted_trace_analysis from .tracelens import TraceLensAnalyzer from .tracelens_inference import ( TraceLensInferencePipeline, @@ -43,6 +45,10 @@ logger = logging.getLogger(__name__) + +def _env_truthy(value: Any) -> bool: + return str(value).strip().lower() in {"1", "true", "yes", "on"} + # Scripts shipping with Magpie that honor MAGPIE_RUN_PHASE (server/client split). MAGPIE_BUILTIN_SCRIPTS = frozenset( { @@ -88,6 +94,7 @@ def __init__( container_writable=config.run_mode == "docker", ) self._task_id: Optional[str] = None + self._resolved_docker_image: Optional[str] = None def run(self, task_id: Optional[str] = None) -> BenchmarkResult: """ @@ -200,6 +207,8 @@ def run(self, task_id: Optional[str] = None) -> BenchmarkResult: workspace_dir=str(workspace), execution_time=time.time() - start_time, profiling_enabled=self.config.profiler.torch_profiler.enabled, + run_kind=self.config.run_kind, + reward_eligible=self.config.reward_eligible, ) result.errors.append(f"TraceLens runtime image setup failed: {e}") result.tracelens_analysis = { @@ -240,6 +249,8 @@ def run(self, task_id: Optional[str] = None) -> BenchmarkResult: workspace_dir=str(workspace), execution_time=time.time() - start_time, profiling_enabled=self.config.profiler.torch_profiler.enabled, + run_kind=self.config.run_kind, + reward_eligible=self.config.reward_eligible, ) result.errors.append(f"TraceLens preprocess failed: {e}") result.tracelens_analysis = { @@ -304,6 +315,7 @@ def run(self, task_id: Optional[str] = None) -> BenchmarkResult: self._cleanup_server_processes(self.config.framework) else: docker_image = self._select_image() + self._resolved_docker_image = docker_image docker_cmd = self._build_docker_command( docker_image=docker_image, workspace=workspace, @@ -324,6 +336,8 @@ def run(self, task_id: Optional[str] = None) -> BenchmarkResult: result.framework = self.config.framework result.model = self.config.model result.profiling_enabled = self.config.profiler.torch_profiler.enabled + result.run_kind = self.config.run_kind + result.reward_eligible = self.config.reward_eligible # Add GPU monitor stats if gpu_monitor_stats is not None: @@ -355,6 +369,21 @@ def run(self, task_id: Optional[str] = None) -> BenchmarkResult: # Without this the gate enforcement would be silently dropped. if not parsed.success: result.success = False + + if not self.config.is_scriptable: + quality_requested = _env_truthy( + self.config.envs.get("RUN_EVAL", False) + ) + result.quality_gate = parse_lm_eval_quality( + workspace, + requested=quality_requested, + ) + if quality_requested and result.quality_gate.get("passed") is not True: + result.success = False + result.errors.append( + "Serving quality evidence gate failed: " + f"{result.quality_gate.get('errors', [])}" + ) else: result.success = False mode_label = "locally" if self.config.is_local else "inside container" @@ -398,12 +427,32 @@ def run(self, task_id: Optional[str] = None) -> BenchmarkResult: if self.config.profiler.torch_profiler.enabled: torch_trace_dir = workspace / "torch_trace" # Recursive: atom writes per-rank traces under rank_/ subdirs - trace_files = list(torch_trace_dir.rglob("*.json.gz")) if torch_trace_dir.is_dir() else [] + trace_files = ( + sorted(torch_trace_dir.rglob("*.json.gz")) + + sorted(torch_trace_dir.rglob("*.json")) + if torch_trace_dir.is_dir() + else [] + ) has_traces = len(trace_files) > 0 if not result.success or not has_traces: if not has_traces: - logger.warning("No torch trace files found, skipping trace analysis / gap analysis") + logger.warning( + "No torch trace files found, skipping trace analysis / " + "gap analysis" + ) + if self.config.profiler.targeted_trace.enabled: + message = ( + "TargetedKernelTrace requested but no Torch profiler " + "trace files were produced" + ) + result.targeted_trace = { + "valid": False, + "reward_eligible": False, + "issues": [message], + } + result.errors.append(message) + result.success = False else: logger.warning("Benchmark failed, skipping trace analysis / gap analysis") else: @@ -411,6 +460,35 @@ def run(self, task_id: Optional[str] = None) -> BenchmarkResult: result.kernel_summary = kernels result.top_bottlenecks = [k.name for k in kernels[:10]] + if self.config.profiler.targeted_trace.enabled: + try: + result.targeted_trace = run_targeted_trace_analysis( + config=self.config, + trace_files=trace_files, + workspace=workspace, + run_id=self._task_id or workspace.name, + resolved_image=( + self._resolved_docker_image + or self.config.docker_image + ), + ) + if not result.targeted_trace["valid"]: + message = ( + "TargetedKernelTrace did not produce valid target " + "evidence" + ) + result.errors.append(message) + result.success = False + except Exception as exc: + message = f"TargetedKernelTrace adaptation failed: {exc}" + result.targeted_trace = { + "valid": False, + "reward_eligible": False, + "issues": [str(exc)], + } + result.errors.append(message) + result.success = False + if self.config.profiler.tracelens.enabled: if is_tracelens_inference_enabled(self.config): tracelens_result = self._run_tracelens_inference_analysis( diff --git a/Magpie/modes/benchmark/config.py b/Magpie/modes/benchmark/config.py index d62c740..e990b40 100644 --- a/Magpie/modes/benchmark/config.py +++ b/Magpie/modes/benchmark/config.py @@ -11,6 +11,8 @@ from enum import Enum from typing import Any, Dict, List, Optional +from ...targeted_trace.config import TargetedTraceConfig + class BenchmarkFramework(Enum): """Supported benchmark frameworks.""" @@ -353,12 +355,14 @@ class ProfilerConfig: system_profiler: System profiler settings (default disabled) tracelens: TraceLens trace analysis settings (default disabled) gpu_monitor: GPU hardware monitoring settings (default enabled) + targeted_trace: Diagnostic selected-kernel evidence settings """ torch_profiler: TorchProfilerConfig = field(default_factory=TorchProfilerConfig) system_profiler: SystemProfilerConfig = field(default_factory=SystemProfilerConfig) tracelens: TraceLensConfig = field(default_factory=TraceLensConfig) gpu_monitor: GPUMonitorConfig = field(default_factory=GPUMonitorConfig) + targeted_trace: TargetedTraceConfig = field(default_factory=TargetedTraceConfig) def to_dict(self) -> Dict[str, Any]: """Convert to dictionary.""" @@ -367,6 +371,7 @@ def to_dict(self) -> Dict[str, Any]: "system_profiler": self.system_profiler.to_dict(), "tracelens": self.tracelens.to_dict(), "gpu_monitor": self.gpu_monitor.to_dict(), + "targeted_trace": self.targeted_trace.to_dict(), } @classmethod @@ -376,6 +381,7 @@ def from_dict(cls, data: Dict[str, Any]) -> "ProfilerConfig": sys_cfg = data.get("system_profiler", {}) tracelens_cfg = data.get("tracelens", {}) gpu_monitor_cfg = data.get("gpu_monitor", {}) + targeted_trace_cfg = data.get("targeted_trace", {}) return cls( torch_profiler=TorchProfilerConfig.from_dict(torch_cfg) if torch_cfg @@ -389,6 +395,9 @@ def from_dict(cls, data: Dict[str, Any]) -> "ProfilerConfig": gpu_monitor=GPUMonitorConfig.from_dict(gpu_monitor_cfg) if gpu_monitor_cfg else GPUMonitorConfig(), + targeted_trace=TargetedTraceConfig.from_dict(targeted_trace_cfg) + if targeted_trace_cfg + else TargetedTraceConfig(), ) @@ -669,6 +678,7 @@ class BenchmarkConfig: model: Model name or path (e.g., "meta-llama/Llama-2-7b-hf") precision: Model precision ("fp8", "fp16", "bf16", "fp4") run_mode: Execution mode - "docker" (default), "local", or "ray" + run_kind: "measurement", "diagnostic", or inferred "auto" envs: Environment variables for benchmark (TP, CONC, ISL, OSL, etc.) profiler: Profiler configuration docker_image: Override automatic image selection @@ -687,6 +697,9 @@ class BenchmarkConfig: # Execution mode: "docker", "local", or "ray" run_mode: str = "docker" + # Evidence lane. ``auto`` becomes diagnostic when a heavy profiler is on. + run_kind: str = "auto" + # Environment variables for benchmark envs: Dict[str, Any] = field(default_factory=dict) @@ -740,6 +753,13 @@ def __post_init__(self): f"Unsupported run_mode: {self.run_mode}. Use 'docker', 'local', or 'ray'." ) + self.run_kind = self.run_kind.lower().strip() + if self.run_kind not in ("auto", "measurement", "diagnostic"): + raise ValueError( + "Unsupported run_kind: " + f"{self.run_kind}. Use 'auto', 'measurement', or 'diagnostic'." + ) + # ``xdit`` is server-less (scriptable) with no Docker image, so it must # run locally. Reject docker/ray here so benchmark_images.yaml stays a # pure Docker-image mapping and the scriptable contract is explicit. @@ -766,6 +786,29 @@ def __post_init__(self): # Convert gap_analysis dict to GapAnalysisConfig if needed if isinstance(self.gap_analysis, dict): self.gap_analysis = GapAnalysisConfig.from_dict(self.gap_analysis) + + heavy_diagnostics = bool( + self.profiler.torch_profiler.enabled + or self.profiler.system_profiler.enabled + or self.profiler.tracelens.enabled + or self.profiler.targeted_trace.enabled + or self.gap_analysis.enabled + ) + if self.run_kind == "auto": + self.run_kind = "diagnostic" if heavy_diagnostics else "measurement" + if self.run_kind == "measurement" and heavy_diagnostics: + raise ValueError( + "run_kind='measurement' requires torch_profiler, system_profiler, " + "TraceLens, gap analysis, and targeted_trace to be disabled" + ) + if ( + self.profiler.targeted_trace.enabled + and not self.profiler.torch_profiler.enabled + ): + raise ValueError( + "profiler.targeted_trace backend 'torch_profiler' requires " + "profiler.torch_profiler.enabled=true" + ) # Convert gpu_selection dict to GpuSelectionConfig if needed if isinstance(self.gpu_selection, dict): @@ -863,6 +906,12 @@ def is_server_lifecycle(self) -> bool: """Reuse a shared inference server across local benchmark tasks.""" return self.server_lifecycle is not None and bool(self.server_lifecycle.enabled) + @property + def reward_eligible(self) -> bool: + """Only profiler-free measurement runs can contribute reward.""" + + return self.run_kind == "measurement" + def to_dict(self) -> Dict[str, Any]: """Convert to dictionary.""" d: Dict[str, Any] = { @@ -870,6 +919,7 @@ def to_dict(self) -> Dict[str, Any]: "model": self.model, "precision": self.precision, "run_mode": self.run_mode, + "run_kind": self.run_kind, "envs": self.envs, "profiler": self.profiler.to_dict(), "gap_analysis": self.gap_analysis.to_dict(), @@ -927,6 +977,7 @@ def from_dict(cls, data: Dict[str, Any]) -> "BenchmarkConfig": model=data.get("model", ""), precision=data.get("precision", "fp8"), run_mode=data.get("run_mode", "docker"), + run_kind=data.get("run_kind", "auto"), envs=data.get("envs", {}), profiler=profiler, gap_analysis=gap_analysis, diff --git a/Magpie/modes/benchmark/quality.py b/Magpie/modes/benchmark/quality.py new file mode 100644 index 0000000..e75cb3c --- /dev/null +++ b/Magpie/modes/benchmark/quality.py @@ -0,0 +1,197 @@ +"""Parse persistent lm-eval artifacts into a bounded serving quality receipt.""" + +from __future__ import annotations + +import hashlib +import json +import math +from pathlib import Path +from typing import Any, Dict, List, Mapping, Optional, Tuple + + +PRIMARY_METRICS = ( + "exact_match,strict-match", + "exact_match,flexible-extract", + "exact_match,none", + "exact_match", + "acc_norm,none", + "acc,none", + "acc_norm", + "acc", + "pass@1,none", + "pass@1", +) +MAX_REPORTED_ARTIFACTS = 256 +MAX_REPORTED_TASKS = 256 +MAX_METRICS_PER_TASK = 64 + + +def _file_sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _numeric_metrics(data: Mapping[str, Any]) -> Dict[str, float]: + metrics: Dict[str, float] = {} + for name, value in data.items(): + lowered = str(name).lower() + if "stderr" in lowered or isinstance(value, bool): + continue + if isinstance(value, (int, float)): + number = float(value) + if math.isfinite(number): + metrics[str(name)] = number + return metrics + + +def _primary(metrics: Mapping[str, float]) -> Tuple[Optional[str], Optional[float]]: + for name in PRIMARY_METRICS: + if name in metrics: + return name, metrics[name] + if metrics: + name = sorted(metrics)[0] + return name, metrics[name] + return None, None + + +def _reject_json_constant(value: str) -> None: + raise ValueError(f"non-finite JSON number is forbidden: {value}") + + +def _reject_duplicate_keys(pairs: List[Tuple[str, Any]]) -> Dict[str, Any]: + result: Dict[str, Any] = {} + for key, value in pairs: + if key in result: + raise ValueError(f"duplicate JSON object key: {key}") + result[key] = value + return result + + +def parse_lm_eval_quality(workspace: Path, *, requested: bool) -> Dict[str, Any]: + """Return an auditable quality gate from ``workspace/lm_eval``. + + ``passed`` means the requested evaluation produced parseable task metrics; it + is an evidence-completeness gate, not an absolute accuracy threshold. Apex can + compare the exposed primary metrics between baseline and candidate runs. + """ + + workspace = Path(workspace) + eval_dir = workspace / "lm_eval" + result_files = ( + sorted(eval_dir.rglob("results*.json")) if eval_dir.is_dir() else [] + ) + artifact_files = ( + sorted(path for path in eval_dir.rglob("*") if path.is_file()) + if eval_dir.is_dir() + else [] + ) + relative_artifacts = [ + str(path.relative_to(workspace)) + for path in artifact_files[:MAX_REPORTED_ARTIFACTS] + ] + result_receipts: List[Dict[str, Any]] = [] + for path in result_files[:MAX_REPORTED_ARTIFACTS]: + try: + result_receipts.append( + { + "path": str(path.relative_to(workspace)), + "size_bytes": path.stat().st_size, + "sha256": _file_sha256(path), + } + ) + except OSError: + pass + + if not result_files: + status = "missing" if requested else "not_requested" + missing_errors = ( + ["RUN_EVAL was requested but no lm_eval/results*.json artifact exists"] + if requested + else [] + ) + return { + "kind": "lm_eval", + "requested": requested, + "status": status, + "passed": False if requested else None, + "evidence_present": False, + "tasks": {}, + "artifacts": relative_artifacts, + "artifact_count": len(artifact_files), + "artifacts_truncated": len(artifact_files) > MAX_REPORTED_ARTIFACTS, + "result_artifact_receipts": result_receipts, + "result_artifact_count": len(result_files), + "result_artifacts_truncated": ( + len(result_files) > MAX_REPORTED_ARTIFACTS + ), + "errors": missing_errors, + "error_count": len(missing_errors), + "errors_truncated": False, + } + + tasks: Dict[str, Dict[str, Any]] = {} + errors: List[str] = [] + for path in result_files: + try: + data = json.loads( + path.read_text(encoding="utf-8"), + parse_constant=_reject_json_constant, + object_pairs_hook=_reject_duplicate_keys, + ) + except (OSError, json.JSONDecodeError, ValueError) as exc: + errors.append(f"{path.name}: {exc}") + continue + raw_results = data.get("results") if isinstance(data, Mapping) else None + if not isinstance(raw_results, Mapping): + errors.append(f"{path.name}: missing results object") + continue + for task_name, task_result in raw_results.items(): + if not isinstance(task_result, Mapping): + errors.append(f"{path.name}:{task_name}: result is not an object") + continue + metrics = _numeric_metrics(task_result) + primary_name, primary_value = _primary(metrics) + if primary_name is None: + errors.append(f"{path.name}:{task_name}: no numeric quality metric") + continue + reported_metrics = dict( + sorted(metrics.items())[:MAX_METRICS_PER_TASK] + ) + if primary_name not in reported_metrics: + reported_metrics[primary_name] = primary_value + tasks[str(task_name)] = { + "primary_metric": primary_name, + "value": primary_value, + "metrics": reported_metrics, + "metric_count": len(metrics), + "metrics_truncated": len(metrics) > MAX_METRICS_PER_TASK, + "source": str(path.relative_to(workspace)), + } + + passed = bool(tasks) and not errors + task_items = sorted(tasks.items()) + reported_tasks = dict(task_items[:MAX_REPORTED_TASKS]) + return { + "kind": "lm_eval", + "requested": requested, + "status": "passed" if passed else "invalid", + "passed": passed, + "evidence_present": bool(tasks), + "tasks": reported_tasks, + "task_count": len(tasks), + "tasks_truncated": len(tasks) > MAX_REPORTED_TASKS, + "artifacts": relative_artifacts, + "artifact_count": len(artifact_files), + "artifacts_truncated": len(artifact_files) > MAX_REPORTED_ARTIFACTS, + "result_artifact_receipts": result_receipts, + "result_artifact_count": len(result_files), + "result_artifacts_truncated": ( + len(result_files) > MAX_REPORTED_ARTIFACTS + ), + "errors": errors[:MAX_REPORTED_ARTIFACTS], + "error_count": len(errors), + "errors_truncated": len(errors) > MAX_REPORTED_ARTIFACTS, + } diff --git a/Magpie/modes/benchmark/result.py b/Magpie/modes/benchmark/result.py index d2a42b3..973da8a 100644 --- a/Magpie/modes/benchmark/result.py +++ b/Magpie/modes/benchmark/result.py @@ -143,6 +143,9 @@ class BenchmarkResult: # Gap analysis results gap_analysis: Optional[Dict[str, Any]] = None + + # TargetedKernelTrace manifest and bounded streaming validation summary + targeted_trace: Optional[Dict[str, Any]] = None # GPU hardware monitoring (temperature, frequency, power) gpu_monitor: Optional[Dict[str, Any]] = None @@ -151,6 +154,8 @@ class BenchmarkResult: workspace_dir: str = "" execution_time: float = 0.0 profiling_enabled: bool = False + run_kind: str = "" + reward_eligible: bool = False # Errors errors: List[str] = field(default_factory=list) @@ -180,10 +185,13 @@ def to_dict(self) -> Dict[str, Any]: "top_bottlenecks": self.top_bottlenecks, "tracelens_analysis": self.tracelens_analysis, "gap_analysis": self.gap_analysis, + "targeted_trace": self.targeted_trace, "gpu_monitor": self.gpu_monitor, "workspace_dir": self.workspace_dir, "execution_time": self.execution_time, "profiling_enabled": self.profiling_enabled, + "run_kind": self.run_kind, + "reward_eligible": self.reward_eligible, "errors": self.errors, } # Scriptable (server-less) extras — e.g. xDiT diffusion. Only emit when @@ -206,6 +214,8 @@ def get_summary(self) -> str: f"{'=' * 60}", f"Model: {self.model}", f"Status: {'SUCCESS' if self.success else 'FAILED'}", + f"Run kind: {self.run_kind or 'unspecified'}", + f"Reward eligible: {'yes' if self.reward_eligible else 'no'}", ] if self.throughput: @@ -275,6 +285,34 @@ def get_summary(self) -> str: if len(top_kernels) > 5: lines.append(f" ... and {len(top_kernels) - 5} more") + if self.targeted_trace: + trace = self.targeted_trace + coverage = trace.get("coverage", {}) + lines.extend( + [ + "", + "TargetedKernelTrace:", + f" Valid: {trace.get('valid', False)}", + f" Seen/written/dropped: {coverage.get('seen', 0)}/" + f"{coverage.get('written', 0)}/{coverage.get('dropped', 0)}", + ] + ) + + if self.quality_gate: + gate = self.quality_gate + lines.extend( + [ + "", + "Quality Evidence:", + f" Status: {gate.get('status', 'unknown')}", + ] + ) + for task, task_result in list(gate.get("tasks", {}).items())[:5]: + lines.append( + f" {task}: {task_result.get('primary_metric')}=" + f"{task_result.get('value')}" + ) + if self.gpu_monitor: lines.extend(["", "GPU Hardware Monitoring:"]) gm = self.gpu_monitor diff --git a/Magpie/modes/benchmark/targeted_trace.py b/Magpie/modes/benchmark/targeted_trace.py new file mode 100644 index 0000000..b393757 --- /dev/null +++ b/Magpie/modes/benchmark/targeted_trace.py @@ -0,0 +1,54 @@ +"""Benchmark-workspace adapter for the generic TargetedKernelTrace contract.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any, Dict, Iterable, Optional + +from ...targeted_trace.postprocess import postprocess_trace_dir +from ...targeted_trace.torch_profiler import adapt_torch_profiler_traces +from .config import BenchmarkConfig + + +def run_targeted_trace_analysis( + *, + config: BenchmarkConfig, + trace_files: Iterable[Path], + workspace: Path, + run_id: str, + resolved_image: Optional[str] = None, +) -> Dict[str, Any]: + """Materialize and validate selected target evidence for one diagnostic run.""" + + targeted_dir = Path(workspace) / "targeted_trace" + manifest = adapt_torch_profiler_traces( + sorted(Path(path) for path in trace_files), + targeted_dir, + config=config.profiler.targeted_trace, + run_id=run_id, + framework=config.framework, + framework_version=str(config.envs.get("FRAMEWORK_VERSION", "")) or None, + image=resolved_image or config.docker_image, + provenance={ + "model": config.model, + "precision": config.precision, + "gpu_arch": config.gpu_arch, + "benchmark_script": config.benchmark_script, + "run_kind": config.run_kind, + }, + ) + summary_path = targeted_dir / "summary.json" + summary = postprocess_trace_dir(targeted_dir, output_path=summary_path) + coverage = manifest.coverage.to_dict() + return { + "valid": bool(summary["valid"] and coverage["written"] > 0), + "reward_eligible": False, + "manifest_path": str(targeted_dir / "manifest.json"), + "summary_path": str(summary_path), + "coverage": coverage, + "events": summary["events"], + "integrity_failures_by_reason": summary[ + "integrity_failures_by_reason" + ], + "issues": summary["issues"], + } diff --git a/Magpie/modes/benchmark/workspace.py b/Magpie/modes/benchmark/workspace.py index ae45f60..c8e6f4b 100644 --- a/Magpie/modes/benchmark/workspace.py +++ b/Magpie/modes/benchmark/workspace.py @@ -29,6 +29,7 @@ class WorkspaceManager: - config.yaml: Configuration snapshot - torch_trace/: PyTorch profiler output - system_profile/: System profiler output (rocprof/ncu) + - targeted_trace/: TargetedKernelTrace manifest, shards, and summary - inferencex_result.json: InferenceX raw output - server.log: Server logs - benchmark_report.json: Magpie summary report @@ -75,15 +76,22 @@ def create(self, config: Optional[Dict[str, Any]] = None) -> Path: # Create subdirectories torch_trace_path = workspace_path / "torch_trace" system_profile_path = workspace_path / "system_profile" + targeted_trace_path = workspace_path / "targeted_trace" torch_trace_path.mkdir(exist_ok=True) system_profile_path.mkdir(exist_ok=True) + targeted_trace_path.mkdir(exist_ok=True) if self.container_writable: # Docker may run with user-namespace remapping or an NFS # root-squash policy. In either case container root is not the # workspace owner on the host, so ordinary 0755/0775 directories # reject server logs and profiler traces. - for path in (workspace_path, torch_trace_path, system_profile_path): + for path in ( + workspace_path, + torch_trace_path, + system_profile_path, + targeted_trace_path, + ): path.chmod(0o777) # Save configuration snapshot @@ -123,6 +131,14 @@ def system_profile_dir(self) -> Optional[Path]: if self._workspace_path: return self._workspace_path / "system_profile" return None + + @property + def targeted_trace_dir(self) -> Optional[Path]: + """Get the TargetedKernelTrace artifact directory.""" + + if self._workspace_path: + return self._workspace_path / "targeted_trace" + return None def get_result_file_path(self, filename: str = "inferencex_result.json") -> Optional[Path]: """Get path for a result file in workspace.""" @@ -185,6 +201,7 @@ def collect_results(self) -> Dict[str, Any]: "inferencex_result": None, "torch_trace_files": [], "system_profile_files": [], + "targeted_trace_files": [], "server_log": None, } @@ -210,6 +227,14 @@ def collect_results(self) -> Dict[str, Any]: results["system_profile_files"] = [ str(f) for f in system_profile_dir.iterdir() if f.is_file() ] + + targeted_trace_dir = self._workspace_path / "targeted_trace" + if targeted_trace_dir.exists(): + results["targeted_trace_files"] = [ + str(path) + for path in targeted_trace_dir.rglob("*") + if path.is_file() + ] # Read server log server_log = self._workspace_path / "server.log" diff --git a/Magpie/scripts/benchmark/atom_mi300x.sh b/Magpie/scripts/benchmark/atom_mi300x.sh index 693ee8a..954080d 100644 --- a/Magpie/scripts/benchmark/atom_mi300x.sh +++ b/Magpie/scripts/benchmark/atom_mi300x.sh @@ -136,7 +136,7 @@ if [[ "$PHASE" == "client" || "$PHASE" == "all" ]]; then fi # After throughput, run evaluation only if RUN_EVAL is true -if [[ "$PHASE" != "server" && "${RUN_EVAL}" = "true" ]]; then +if [[ "$PHASE" != "server" && "${RUN_EVAL,,}" = "true" ]]; then if [[ -n "${BENCHMARK_BASE_URL:-}" ]]; then if declare -F magpie_run_eval_remote_direct &>/dev/null; then magpie_run_eval_remote_direct || exit $? @@ -144,8 +144,11 @@ if [[ "$PHASE" != "server" && "${RUN_EVAL}" = "true" ]]; then echo "[atom_mi300x] RUN_EVAL=true with BENCHMARK_BASE_URL but magpie_run_eval_remote_direct shim not available; skipping eval (results gate will see accuracy=None)." fi else + magpie_mark_lm_eval_start || exit $? run_eval --framework lm-eval --port "$PORT" --concurrent-requests $CONC || exit $? + magpie_preserve_lm_eval_artifacts || exit $? append_lm_eval_summary + magpie_preserve_lm_eval_artifacts || exit $? fi fi set +x diff --git a/Magpie/scripts/benchmark/atom_mi355x.sh b/Magpie/scripts/benchmark/atom_mi355x.sh index 12ca251..51cfc80 100644 --- a/Magpie/scripts/benchmark/atom_mi355x.sh +++ b/Magpie/scripts/benchmark/atom_mi355x.sh @@ -131,7 +131,7 @@ if [[ "$PHASE" == "client" || "$PHASE" == "all" ]]; then fi fi -if [[ "$PHASE" != "server" && "${RUN_EVAL}" = "true" ]]; then +if [[ "$PHASE" != "server" && "${RUN_EVAL,,}" = "true" ]]; then if [[ -n "${BENCHMARK_BASE_URL:-}" ]]; then if declare -F magpie_run_eval_remote_direct &>/dev/null; then magpie_run_eval_remote_direct || exit $? @@ -139,8 +139,11 @@ if [[ "$PHASE" != "server" && "${RUN_EVAL}" = "true" ]]; then echo "[atom_mi355x] RUN_EVAL=true with BENCHMARK_BASE_URL but magpie_run_eval_remote_direct shim not available; skipping eval (results gate will see accuracy=None)." fi else + magpie_mark_lm_eval_start || exit $? run_eval --framework lm-eval --port "$PORT" --concurrent-requests $CONC || exit $? + magpie_preserve_lm_eval_artifacts || exit $? append_lm_eval_summary + magpie_preserve_lm_eval_artifacts || exit $? fi fi set +x diff --git a/Magpie/scripts/benchmark/magpie_bench_remote_compat.sh b/Magpie/scripts/benchmark/magpie_bench_remote_compat.sh index 29cb305..2b4b619 100644 --- a/Magpie/scripts/benchmark/magpie_bench_remote_compat.sh +++ b/Magpie/scripts/benchmark/magpie_bench_remote_compat.sh @@ -167,3 +167,59 @@ magpie_run_eval_remote_direct() { fi return $rc } + +############################################################################### +# magpie_preserve_lm_eval_artifacts +# +# InferenceX's append_lm_eval_summary may move files out of EVAL_RESULT_DIR into +# its repository working directory. Preserve the raw evaluation evidence under +# the Magpie workspace both before and after that helper runs. The fallback only +# copies known lm-eval artifact names from the working-directory root. +############################################################################### +magpie_mark_lm_eval_start() { + local result_dir="${RESULT_DIR:-${WORKSPACE_DIR:-/workspace}}" + local inferx_root="${MAGPIE_INFERENCEX_ROOT:-$(pwd)}" + local snapshot="${result_dir%/}/.lm_eval_preexisting.tsv" + mkdir -p "$result_dir" || return 1 + # Snapshot the exact pre-run artifact identities. Timestamp-only filtering is + # racy on filesystems whose timestamp resolution cannot distinguish the + # marker from an immediately-created lm-eval result. + find "$inferx_root" -maxdepth 1 -type f \ + \( -name 'results*.json' -o -name 'samples*.jsonl' \ + -o -name '*lm_eval*summary*.json' \) \ + -printf '%f\t%s\t%T@\n' | LC_ALL=C sort > "$snapshot" || return 1 + touch "${result_dir%/}/.lm_eval_started" || return 1 +} + +magpie_preserve_lm_eval_artifacts() { + local result_dir="${RESULT_DIR:-${WORKSPACE_DIR:-/workspace}}" + local out_dir="${result_dir%/}/lm_eval" + local marker="${result_dir%/}/.lm_eval_started" + local snapshot="${result_dir%/}/.lm_eval_preexisting.tsv" + mkdir -p "$out_dir" || return 1 + + if [[ -n "${EVAL_RESULT_DIR:-}" && -d "${EVAL_RESULT_DIR}" ]]; then + local source_real out_real + source_real=$(readlink -f "${EVAL_RESULT_DIR}" 2>/dev/null || true) + out_real=$(readlink -f "$out_dir" 2>/dev/null || true) + if [[ -n "$source_real" && "$source_real" != "$out_real" ]]; then + cp -a "${EVAL_RESULT_DIR}"/. "$out_dir"/ || return 1 + fi + fi + + local artifact signature + while IFS= read -r -d '' artifact; do + if [[ -e "$marker" && -f "$snapshot" ]]; then + signature=$(find "$artifact" -maxdepth 0 -printf '%f\t%s\t%T@\n') + if grep -Fqx -- "$signature" "$snapshot"; then + continue + fi + fi + cp -a "$artifact" "$out_dir"/ || return 1 + done < <( + find "${MAGPIE_INFERENCEX_ROOT:-$(pwd)}" -maxdepth 1 -type f \ + \( -name 'results*.json' -o -name 'samples*.jsonl' \ + -o -name '*lm_eval*summary*.json' \) \ + -print0 + ) +} diff --git a/Magpie/scripts/benchmark/sglang_mi300x.sh b/Magpie/scripts/benchmark/sglang_mi300x.sh index 1776cc9..5c70105 100644 --- a/Magpie/scripts/benchmark/sglang_mi300x.sh +++ b/Magpie/scripts/benchmark/sglang_mi300x.sh @@ -135,7 +135,7 @@ if [[ "$PHASE" == "client" || "$PHASE" == "all" ]]; then fi fi -if [[ "$PHASE" != "server" && "${RUN_EVAL}" = "true" ]]; then +if [[ "$PHASE" != "server" && "${RUN_EVAL,,}" = "true" ]]; then if [[ -n "${BENCHMARK_BASE_URL:-}" ]]; then if declare -F magpie_run_eval_remote_direct &>/dev/null; then magpie_run_eval_remote_direct || exit $? @@ -143,9 +143,11 @@ if [[ "$PHASE" != "server" && "${RUN_EVAL}" = "true" ]]; then echo "[sglang_mi300x] RUN_EVAL=true with BENCHMARK_BASE_URL but magpie_run_eval_remote_direct shim not available; skipping eval (results gate will see accuracy=None)." fi else + magpie_mark_lm_eval_start || exit $? run_eval --framework lm-eval --port "$PORT" --concurrent-requests $CONC || exit $? + magpie_preserve_lm_eval_artifacts || exit $? append_lm_eval_summary + magpie_preserve_lm_eval_artifacts || exit $? fi fi set +x - diff --git a/Magpie/scripts/benchmark/sglang_mi355x.sh b/Magpie/scripts/benchmark/sglang_mi355x.sh index 69ff396..2e64d69 100644 --- a/Magpie/scripts/benchmark/sglang_mi355x.sh +++ b/Magpie/scripts/benchmark/sglang_mi355x.sh @@ -130,7 +130,7 @@ if [[ "$PHASE" == "client" || "$PHASE" == "all" ]]; then fi fi -if [[ "$PHASE" != "server" && "${RUN_EVAL}" = "true" ]]; then +if [[ "$PHASE" != "server" && "${RUN_EVAL,,}" = "true" ]]; then if [[ -n "${BENCHMARK_BASE_URL:-}" ]]; then if declare -F magpie_run_eval_remote_direct &>/dev/null; then magpie_run_eval_remote_direct || exit $? @@ -138,8 +138,11 @@ if [[ "$PHASE" != "server" && "${RUN_EVAL}" = "true" ]]; then echo "[sglang_mi355x] RUN_EVAL=true with BENCHMARK_BASE_URL but magpie_run_eval_remote_direct shim not available; skipping eval (results gate will see accuracy=None)." fi else + magpie_mark_lm_eval_start || exit $? run_eval --framework lm-eval --port "$PORT" --concurrent-requests $CONC || exit $? + magpie_preserve_lm_eval_artifacts || exit $? append_lm_eval_summary + magpie_preserve_lm_eval_artifacts || exit $? fi fi set +x diff --git a/Magpie/scripts/benchmark/vllm_mi300x.sh b/Magpie/scripts/benchmark/vllm_mi300x.sh index f68460b..1bf78b1 100644 --- a/Magpie/scripts/benchmark/vllm_mi300x.sh +++ b/Magpie/scripts/benchmark/vllm_mi300x.sh @@ -143,7 +143,7 @@ if [[ "$PHASE" == "client" || "$PHASE" == "all" ]]; then fi # After throughput, run evaluation only if RUN_EVAL is true -if [[ "$PHASE" != "server" && "${RUN_EVAL}" = "true" ]]; then +if [[ "$PHASE" != "server" && "${RUN_EVAL,,}" = "true" ]]; then if [[ -n "${BENCHMARK_BASE_URL:-}" ]]; then if declare -F magpie_run_eval_remote_direct &>/dev/null; then magpie_run_eval_remote_direct || exit $? @@ -151,9 +151,11 @@ if [[ "$PHASE" != "server" && "${RUN_EVAL}" = "true" ]]; then echo "[vllm_mi300x] RUN_EVAL=true with BENCHMARK_BASE_URL but magpie_run_eval_remote_direct shim not available; skipping eval (results gate will see accuracy=None)." fi else + magpie_mark_lm_eval_start || exit $? run_eval --framework lm-eval --port "$PORT" --concurrent-requests $CONC || exit $? + magpie_preserve_lm_eval_artifacts || exit $? append_lm_eval_summary + magpie_preserve_lm_eval_artifacts || exit $? fi fi set +x - diff --git a/Magpie/scripts/benchmark/vllm_mi355x.sh b/Magpie/scripts/benchmark/vllm_mi355x.sh index 1df07f2..c4e6f41 100644 --- a/Magpie/scripts/benchmark/vllm_mi355x.sh +++ b/Magpie/scripts/benchmark/vllm_mi355x.sh @@ -139,7 +139,7 @@ if [[ "$PHASE" == "client" || "$PHASE" == "all" ]]; then fi fi -if [[ "$PHASE" != "server" && "${RUN_EVAL}" = "true" ]]; then +if [[ "$PHASE" != "server" && "${RUN_EVAL,,}" = "true" ]]; then if [[ -n "${BENCHMARK_BASE_URL:-}" ]]; then if declare -F magpie_run_eval_remote_direct &>/dev/null; then magpie_run_eval_remote_direct || exit $? @@ -147,8 +147,11 @@ if [[ "$PHASE" != "server" && "${RUN_EVAL}" = "true" ]]; then echo "[vllm_mi355x] RUN_EVAL=true with BENCHMARK_BASE_URL but magpie_run_eval_remote_direct shim not available; skipping eval (results gate will see accuracy=None)." fi else + magpie_mark_lm_eval_start || exit $? run_eval --framework lm-eval --port "$PORT" --concurrent-requests $CONC || exit $? + magpie_preserve_lm_eval_artifacts || exit $? append_lm_eval_summary + magpie_preserve_lm_eval_artifacts || exit $? fi fi set +x diff --git a/Magpie/targeted_trace/README.md b/Magpie/targeted_trace/README.md new file mode 100644 index 0000000..6278ad8 --- /dev/null +++ b/Magpie/targeted_trace/README.md @@ -0,0 +1,125 @@ +# TargetedKernelTrace + +`Magpie.targeted_trace` is Magpie's versioned acquisition contract for kernel +evidence needed by optimization agents. It is deliberately independent of Apex, +TraceLens internals, container tags, and image-specific source registries. + +## Ownership and run separation + +- Magpie owns acquisition, integrity receipts, deterministic sampling, and loss + accounting. +- TraceLens remains the authoritative system/GPU analysis layer. +- Consumers such as Apex read `manifest.json` plus checksummed shard artifacts. +- Targeted traces are diagnostic artifacts and always carry + `reward_eligible: false`. A benchmark explicitly declared as `measurement` + rejects Torch profiler, TraceLens, system profiler, gap analysis, and targeted + tracing. + +## Artifact contract + +Each `shards/trace_pid_rank.jsonl` contains: + +1. A header with run/rank/PID, sampling seed/rate, and finite record budget. +2. Typed event envelopes with a monotonic sequence and chained SHA-256 checksum. +3. An end sentinel with `seen`, `sampled`, `written`, `dropped`, and + `dropped_by_reason` counters. + +`manifest.json` records schema/version, targets, provenance, aggregate coverage, +and per-shard file/checksum receipts. Unsupported schema versions fail fast. +Postprocessing reads one JSONL line at a time and reports corrupt or missing tails; +it never silently skips them. + +## Evidence fidelity + +The explicit `TargetedTraceRecorder` API captures Python-visible Triton and HIP +wrapper calls: launch source/hash, Python grid, tensor shape/dtype/stride, named +scalars, constexpr values, and meta parameters. It does not read tensor contents or +store raw data pointers. + +The Torch profiler adapter contributes runtime symbol/grid/block/stream/duration, +rank/stage/graph context, correlation IDs, and any tensor metadata present in the +trace. Missing fields remain null/empty with warnings; symbol/count/order is not +treated as a globally stable CPU-to-GPU join. + +Sampling uses only `{run_seed, stable_event_key}` through SHA-256. Python's +process-randomized `hash()` and mutable PRNG state are not used. + +## Benchmark configuration + +Target selection uses portable glob patterns rather than a hardcoded image map: + +```yaml +benchmark: + run_kind: diagnostic + profiler: + torch_profiler: + enabled: true + targeted_trace: + enabled: true + backend: torch_profiler + run_seed: qwen-diagnostic-1 + sample_rate: 0.1 + max_records_per_shard: 10000 + targets: + - target_id: aiter.fused_moe + name_patterns: + - "*fused_moe*" + package: aiter +``` + +The benchmark workspace receives `targeted_trace/manifest.json`, shards, and +`summary.json`; `benchmark_report.json` contains the bounded summary and manifest +path. + +## Standalone CLI + +Convert existing Torch profiler traces: + +```bash +magpie targeted-trace adapt-torch \ + --trace-dir ./torch_trace \ + --target-config targets.yaml \ + --output-dir ./targeted_trace \ + --run-id diagnostic-001 \ + --framework vllm +``` + +Validate and aggregate artifacts: + +```bash +magpie targeted-trace postprocess \ + --trace-dir ./targeted_trace \ + --output ./targeted_trace/summary.json \ + --strict +``` + +## Runtime probes + +Framework integration code can use the generic API at a known launch boundary: + +```python +from Magpie.targeted_trace import TargetedTraceRecorder + +with TargetedTraceRecorder( + "/workspace/targeted_trace", + run_id="diagnostic-001", + run_seed="qwen-diagnostic-1", + framework="vllm", + rank=0, +) as trace: + trace.record_triton_launch( + target_id="aiter.fused_moe", + kernel_name="fused_moe_kernel", + args=(x, weights), + positional_names=("x", "weights"), + kwargs={"BLOCK_SIZE": 256}, + constexpr_names=("BLOCK_SIZE",), + grid=(128, 1, 1), + source_path=__file__, + source_line=42, + ) +``` + +This module intentionally does not ship an image registry or mutate framework +packages. Source discovery and temporary instrumentation are separate adapters; +the durable evidence format remains generic. diff --git a/Magpie/targeted_trace/__init__.py b/Magpie/targeted_trace/__init__.py new file mode 100644 index 0000000..34ba2d6 --- /dev/null +++ b/Magpie/targeted_trace/__init__.py @@ -0,0 +1,44 @@ +"""TargetedKernelTrace acquisition contract and adapters.""" + +from .capture import TargetedTraceRecorder +from .config import TargetSpec, TargetedTraceConfig +from .postprocess import postprocess_trace_dir, validate_shard +from .schema import ( + SCHEMA_NAME, + SCHEMA_VERSION, + LaunchSemantics, + RuntimeEvidence, + ShardCounters, + ShardReceipt, + SourceEvidence, + TargetedTraceManifest, + TargetedTraceRecord, + TensorEvidence, + TraceContext, + TraceIdentity, + TraceValidationError, +) +from .torch_profiler import adapt_torch_profiler_traces, iter_trace_events + +__all__ = [ + "SCHEMA_NAME", + "SCHEMA_VERSION", + "LaunchSemantics", + "RuntimeEvidence", + "ShardCounters", + "ShardReceipt", + "SourceEvidence", + "TargetSpec", + "TargetedTraceConfig", + "TargetedTraceManifest", + "TargetedTraceRecord", + "TargetedTraceRecorder", + "TensorEvidence", + "TraceContext", + "TraceIdentity", + "TraceValidationError", + "adapt_torch_profiler_traces", + "iter_trace_events", + "postprocess_trace_dir", + "validate_shard", +] diff --git a/Magpie/targeted_trace/capture.py b/Magpie/targeted_trace/capture.py new file mode 100644 index 0000000..edca93c --- /dev/null +++ b/Magpie/targeted_trace/capture.py @@ -0,0 +1,311 @@ +"""Generic Python runtime capture API for Triton and Python-visible HIP launches.""" + +from __future__ import annotations + +import os +import time +from collections import defaultdict +from pathlib import Path +from typing import Any, Iterable, Mapping, Optional, Sequence + +from .sampling import stable_key +from .schema import ( + RuntimeEvidence, + TargetedTraceRecord, + TraceContext, + TraceIdentity, +) +from .serialization import invocation_semantics, source_evidence +from .writer import TraceShardWriter, default_shard_path, merge_runtime_manifest + + +class TargetedTraceRecorder: + """Capture semantic launch evidence through an explicit, framework-neutral API. + + The recorder does not patch packages or dispatchers. Callers inject these + methods at a known Python-visible launch/wrapper boundary, keeping framework + discovery separate from the durable artifact contract. + """ + + def __init__( + self, + output_dir: Path, + *, + run_id: str, + run_seed: str, + framework: str, + rank: int = 0, + pid: Optional[int] = None, + world_size: Optional[int] = None, + framework_version: Optional[str] = None, + execution_mode: str = "unknown", + stage: str = "unknown", + sample_rate: float = 1.0, + max_records: int = 100_000, + ) -> None: + self.output_dir = Path(output_dir) + self.run_id = run_id + self.framework = framework + self.rank = rank + self.pid = os.getpid() if pid is None else pid + self.world_size = world_size + self.framework_version = framework_version + self.execution_mode = execution_mode + self.stage = stage + self._occurrences: defaultdict[str, int] = defaultdict(int) + self._targets: dict[tuple[str, str], dict[str, Any]] = {} + self._manifest_written = False + self.writer = TraceShardWriter( + default_shard_path(self.output_dir, rank=rank, pid=self.pid), + run_id=run_id, + rank=rank, + pid=self.pid, + run_seed=run_seed, + sample_rate=sample_rate, + max_records=max_records, + header_metadata={ + "framework": framework, + "framework_version": framework_version, + "world_size": world_size, + "capture_backend": "python_runtime", + }, + ) + + def _record( + self, + *, + kind: str, + target_id: str, + kernel_name: str, + args: Sequence[Any], + kwargs: Mapping[str, Any], + grid: Any, + source_path: Optional[str], + source_line: Optional[int], + source_function: Optional[str], + positional_names: Optional[Sequence[str]], + meta_names: Iterable[str], + constexpr_names: Iterable[str], + variant_id: str, + package: Optional[str], + image: Optional[str], + source_hashes: Optional[Mapping[str, str]], + provenance_hashes: Optional[Mapping[str, str]], + runtime: Optional[RuntimeEvidence], + stage: Optional[str], + graph_id: Optional[str], + ) -> bool: + target_key = (target_id, variant_id) + if target_key not in self._targets: + self._targets[target_key] = { + "target_id": target_id, + "variant_id": variant_id, + "name_patterns": [], + "package": package, + "image": image, + "source": ( + { + "path": source_path, + "line": source_line, + "function": source_function, + } + if source_path + else None + ), + "source_hashes": dict(source_hashes or {}), + "provenance_hashes": dict(provenance_hashes or {}), + } + patterns = self._targets[target_key]["name_patterns"] + if kernel_name not in patterns: + patterns.append(kernel_name) + try: + semantics, warnings = invocation_semantics( + args=args, + kwargs=kwargs, + positional_names=positional_names, + meta_names=meta_names, + constexpr_names=constexpr_names, + python_grid=grid, + source=source_evidence( + source_path, + line=source_line, + function=source_function, + ), + ) + base_parts = { + "kind": kind, + "target_id": target_id, + "kernel_name": kernel_name, + "source": semantics.source.to_dict() if semantics.source else None, + "tensors": [ + { + "name": tensor.name, + "shape": list(tensor.shape), + "dtype": tensor.dtype, + "stride": list(tensor.stride) if tensor.stride else None, + } + for tensor in semantics.tensors + ], + "scalars": dict(semantics.named_scalars), + "constexpr": dict(semantics.constexpr), + "meta": dict(semantics.meta), + "grid": semantics.python_grid, + } + base_token = stable_key(base_parts) + occurrence = self._occurrences[base_token] + self._occurrences[base_token] += 1 + event_key = stable_key(base_parts, occurrence=occurrence) + runtime_evidence = runtime or RuntimeEvidence(gpu_symbol=kernel_name) + record = TargetedTraceRecord( + kind=kind, + stable_event_key=event_key, + identity=TraceIdentity( + run_id=self.run_id, + target_id=target_id, + variant_id=variant_id, + package=package, + image=image, + source_hashes=dict(source_hashes or {}), + provenance_hashes=dict(provenance_hashes or {}), + ), + context=TraceContext( + framework=self.framework, + framework_version=self.framework_version, + rank=self.rank, + pid=self.pid, + world_size=self.world_size, + stage=stage or self.stage, + execution_mode=self.execution_mode, + graph_id=graph_id, + ), + semantics=semantics, + runtime=runtime_evidence, + timestamp_ns=time.time_ns(), + warnings=warnings, + ) + return self.writer.submit(record) + except Exception: + self.writer.note_failed_observation("serialization_error") + return False + + def record_triton_launch( + self, + *, + target_id: str, + kernel_name: str, + args: Sequence[Any] = (), + kwargs: Optional[Mapping[str, Any]] = None, + grid: Any = None, + source_path: Optional[str] = None, + source_line: Optional[int] = None, + source_function: Optional[str] = None, + positional_names: Optional[Sequence[str]] = None, + meta_names: Iterable[str] = (), + constexpr_names: Iterable[str] = (), + variant_id: str = "baseline", + package: Optional[str] = None, + image: Optional[str] = None, + source_hashes: Optional[Mapping[str, str]] = None, + provenance_hashes: Optional[Mapping[str, str]] = None, + runtime: Optional[RuntimeEvidence] = None, + stage: Optional[str] = None, + graph_id: Optional[str] = None, + ) -> bool: + """Record one Python-visible Triton launch.""" + + return self._record( + kind="triton_launch", + target_id=target_id, + kernel_name=kernel_name, + args=args, + kwargs=kwargs or {}, + grid=grid, + source_path=source_path, + source_line=source_line, + source_function=source_function, + positional_names=positional_names, + meta_names=meta_names, + constexpr_names=constexpr_names, + variant_id=variant_id, + package=package, + image=image, + source_hashes=source_hashes, + provenance_hashes=provenance_hashes, + runtime=runtime, + stage=stage, + graph_id=graph_id, + ) + + def record_python_hip_launch( + self, + *, + target_id: str, + kernel_name: str, + args: Sequence[Any] = (), + kwargs: Optional[Mapping[str, Any]] = None, + grid: Any = None, + source_path: Optional[str] = None, + source_line: Optional[int] = None, + source_function: Optional[str] = None, + positional_names: Optional[Sequence[str]] = None, + meta_names: Iterable[str] = (), + constexpr_names: Iterable[str] = (), + variant_id: str = "baseline", + package: Optional[str] = None, + image: Optional[str] = None, + source_hashes: Optional[Mapping[str, str]] = None, + provenance_hashes: Optional[Mapping[str, str]] = None, + runtime: Optional[RuntimeEvidence] = None, + stage: Optional[str] = None, + graph_id: Optional[str] = None, + ) -> bool: + """Record one Python wrapper call that launches a HIP/custom op.""" + + return self._record( + kind="python_hip_launch", + target_id=target_id, + kernel_name=kernel_name, + args=args, + kwargs=kwargs or {}, + grid=grid, + source_path=source_path, + source_line=source_line, + source_function=source_function, + positional_names=positional_names, + meta_names=meta_names, + constexpr_names=constexpr_names, + variant_id=variant_id, + package=package, + image=image, + source_hashes=source_hashes, + provenance_hashes=provenance_hashes, + runtime=runtime, + stage=stage, + graph_id=graph_id, + ) + + def close(self): + """Close the underlying shard and return its receipt.""" + + receipt = self.writer.close() + if not self._manifest_written: + merge_runtime_manifest( + self.output_dir / "manifest.json", + run_id=self.run_id, + receipt=receipt, + targets=list(self._targets.values()), + provenance={ + "framework": self.framework, + "framework_version": self.framework_version, + "world_size": self.world_size, + "capture_backend": "python_runtime", + }, + ) + self._manifest_written = True + return receipt + + def __enter__(self) -> "TargetedTraceRecorder": + return self + + def __exit__(self, exc_type: object, exc: object, tb: object) -> None: + self.close() diff --git a/Magpie/targeted_trace/cli.py b/Magpie/targeted_trace/cli.py new file mode 100644 index 0000000..abb2b6f --- /dev/null +++ b/Magpie/targeted_trace/cli.py @@ -0,0 +1,106 @@ +"""CLI adapter for offline Torch-profiler conversion and trace validation.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any, Dict, Iterable, List + +import yaml # type: ignore[import-untyped] + +from .config import TargetedTraceConfig +from .postprocess import postprocess_trace_dir +from .torch_profiler import adapt_torch_profiler_traces + + +def add_targeted_trace_parser(subparsers: Any) -> None: + """Register the ``targeted-trace`` CLI without coupling it to benchmark code.""" + + parser = subparsers.add_parser( + "targeted-trace", + help="Adapt or validate diagnostic TargetedKernelTrace artifacts", + ) + commands = parser.add_subparsers(dest="targeted_trace_command", required=True) + + adapt = commands.add_parser( + "adapt-torch", help="Stream selected events from Torch profiler traces" + ) + adapt.add_argument("--trace", type=Path, action="append", default=[]) + adapt.add_argument( + "--trace-dir", + type=Path, + help="Recursively discover .json/.json.gz Torch profiler traces", + ) + adapt.add_argument("--target-config", type=Path, required=True) + adapt.add_argument("--output-dir", "-o", type=Path, required=True) + adapt.add_argument("--run-id", required=True) + adapt.add_argument("--framework", required=True) + adapt.add_argument("--framework-version") + adapt.add_argument("--image") + + postprocess = commands.add_parser( + "postprocess", help="Stream-validate shards and write a bounded summary" + ) + postprocess.add_argument("--trace-dir", type=Path, required=True) + postprocess.add_argument("--output", "-o", type=Path) + postprocess.add_argument("--strict", action="store_true") + + +def _load_target_config(path: Path) -> TargetedTraceConfig: + data = yaml.safe_load(path.read_text(encoding="utf-8")) or {} + if not isinstance(data, dict): + raise ValueError("target config root must be an object") + raw = data.get("targeted_trace", data) + if not isinstance(raw, dict): + raise ValueError("targeted_trace config must be an object") + raw = dict(raw) + raw["enabled"] = True + return TargetedTraceConfig.from_dict(raw) + + +def _discover_traces(explicit: Iterable[Path], trace_dir: Path | None) -> List[Path]: + paths = [Path(path) for path in explicit] + if trace_dir is not None: + paths.extend(trace_dir.rglob("*.json")) + paths.extend(trace_dir.rglob("*.json.gz")) + unique = sorted({path.resolve() for path in paths if path.is_file()}) + if not unique: + raise ValueError("no Torch profiler trace files were found") + return unique + + +def run_targeted_trace(args: argparse.Namespace) -> int: + """Execute a parsed targeted-trace command.""" + + if args.targeted_trace_command == "adapt-torch": + try: + config = _load_target_config(args.target_config) + traces = _discover_traces(args.trace, args.trace_dir) + manifest = adapt_torch_profiler_traces( + traces, + args.output_dir, + config=config, + run_id=args.run_id, + framework=args.framework, + framework_version=args.framework_version, + image=args.image, + ) + except (OSError, ValueError) as exc: + print(f"Error: {exc}") + return 1 + print(json.dumps(manifest.to_dict(), indent=2, sort_keys=True)) + return 0 + if args.targeted_trace_command == "postprocess": + try: + summary = postprocess_trace_dir( + args.trace_dir, + output_path=args.output, + strict=args.strict, + ) + except (OSError, ValueError) as exc: + print(f"Error: {exc}") + return 1 + print(json.dumps(summary, indent=2, sort_keys=True)) + return 0 if summary["valid"] else 1 + raise ValueError(f"unknown targeted trace command: {args.targeted_trace_command}") diff --git a/Magpie/targeted_trace/config.py b/Magpie/targeted_trace/config.py new file mode 100644 index 0000000..6ba921d --- /dev/null +++ b/Magpie/targeted_trace/config.py @@ -0,0 +1,116 @@ +"""User-facing configuration for generic target selection and trace budgets.""" + +from __future__ import annotations + +import fnmatch +from dataclasses import dataclass, field +from typing import Any, Dict, List, Mapping, Optional + + +@dataclass(frozen=True) +class TargetSpec: + """A target selected by portable symbol patterns, not image registries.""" + + target_id: str + name_patterns: tuple[str, ...] + variant_id: str = "baseline" + package: Optional[str] = None + source: Optional[Mapping[str, Any]] = None + source_hashes: Mapping[str, str] = field(default_factory=dict) + provenance_hashes: Mapping[str, str] = field(default_factory=dict) + + def __post_init__(self) -> None: + if not self.target_id.strip(): + raise ValueError("target_id must be non-empty") + if not self.name_patterns or any(not pattern for pattern in self.name_patterns): + raise ValueError(f"target {self.target_id!r} requires name_patterns") + + def matches(self, symbol: str) -> bool: + """Return whether a profiler/runtime symbol belongs to this target.""" + + return any(fnmatch.fnmatchcase(symbol, pattern) for pattern in self.name_patterns) + + def to_dict(self) -> Dict[str, Any]: + return { + "target_id": self.target_id, + "name_patterns": list(self.name_patterns), + "variant_id": self.variant_id, + "package": self.package, + "source": dict(self.source) if self.source else None, + "source_hashes": dict(self.source_hashes), + "provenance_hashes": dict(self.provenance_hashes), + } + + @classmethod + def from_dict(cls, data: Mapping[str, Any]) -> "TargetSpec": + patterns = data.get("name_patterns", data.get("names", [])) + if isinstance(patterns, str): + patterns = [patterns] + return cls( + target_id=str(data.get("target_id", data.get("id", ""))), + name_patterns=tuple(str(item) for item in patterns), + variant_id=str(data.get("variant_id", "baseline")), + package=str(data["package"]) if data.get("package") is not None else None, + source=dict(data["source"]) if isinstance(data.get("source"), Mapping) else None, + source_hashes=dict(data.get("source_hashes", {})), + provenance_hashes=dict(data.get("provenance_hashes", {})), + ) + + +@dataclass +class TargetedTraceConfig: + """Diagnostic-only targeted trace settings.""" + + enabled: bool = False + backend: str = "torch_profiler" + run_seed: str = "magpie-targeted-trace" + sample_rate: float = 1.0 + max_records_per_shard: int = 100_000 + targets: List[TargetSpec] = field(default_factory=list) + + def __post_init__(self) -> None: + self.backend = self.backend.lower().strip() + if self.backend not in {"torch_profiler"}: + raise ValueError( + f"unsupported targeted trace backend {self.backend!r}; " + "use 'torch_profiler'" + ) + if not 0.0 <= self.sample_rate <= 1.0: + raise ValueError("targeted_trace.sample_rate must be between 0 and 1") + if self.max_records_per_shard < 0: + raise ValueError( + "targeted_trace.max_records_per_shard must be non-negative" + ) + converted: List[TargetSpec] = [] + for target in self.targets: + converted.append( + TargetSpec.from_dict(target) if isinstance(target, Mapping) else target + ) + self.targets = converted + if self.enabled and not self.targets: + raise ValueError("enabled targeted_trace requires at least one target") + + def to_dict(self) -> Dict[str, Any]: + return { + "enabled": self.enabled, + "backend": self.backend, + "run_seed": self.run_seed, + "sample_rate": self.sample_rate, + "max_records_per_shard": self.max_records_per_shard, + "targets": [target.to_dict() for target in self.targets], + } + + @classmethod + def from_dict(cls, data: Mapping[str, Any]) -> "TargetedTraceConfig": + return cls( + enabled=bool(data.get("enabled", False)), + backend=str(data.get("backend", "torch_profiler")), + run_seed=str(data.get("run_seed", "magpie-targeted-trace")), + sample_rate=float(data.get("sample_rate", 1.0)), + max_records_per_shard=int(data.get("max_records_per_shard", 100_000)), + targets=[ + TargetSpec.from_dict(item) + for item in data.get("targets", []) + if isinstance(item, Mapping) + ], + ) diff --git a/Magpie/targeted_trace/postprocess.py b/Magpie/targeted_trace/postprocess.py new file mode 100644 index 0000000..6764cfb --- /dev/null +++ b/Magpie/targeted_trace/postprocess.py @@ -0,0 +1,327 @@ +"""Streaming integrity validation and bounded aggregation for trace shards.""" + +from __future__ import annotations + +import hashlib +import json +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Callable, Dict, List, Mapping, Optional + +from .schema import ( + ZERO_CHECKSUM, + ShardCounters, + ShardReceipt, + TargetedTraceManifest, + TargetedTraceRecord, + TraceValidationError, + canonical_json, + validate_envelope, +) + + +MAX_JSONL_LINE_BYTES = 16 * 1024 * 1024 + + +def _reject_json_constant(value: str) -> None: + raise ValueError(f"non-finite JSON number is forbidden: {value}") + + +@dataclass +class ShardValidationResult: + """Validation outcome for one streamed shard.""" + + path: str + valid: bool = False + complete: bool = False + event_count: int = 0 + byte_count: int = 0 + file_sha256: str = "" + sequence_end: int = -1 + chain_checksum: str = ZERO_CHECKSUM + counters: Optional[ShardCounters] = None + rank: Optional[int] = None + pid: Optional[int] = None + issues: List[str] = field(default_factory=list) + + def to_dict(self) -> Dict[str, Any]: + return { + "path": self.path, + "valid": self.valid, + "complete": self.complete, + "event_count": self.event_count, + "byte_count": self.byte_count, + "file_sha256": self.file_sha256, + "sequence_end": self.sequence_end, + "chain_checksum": self.chain_checksum, + "counters": self.counters.to_dict() if self.counters else None, + "rank": self.rank, + "pid": self.pid, + "issues": list(self.issues), + } + + +def _receipt_mismatches( + result: ShardValidationResult, receipt: ShardReceipt +) -> List[str]: + mismatches: List[str] = [] + checks = { + "rank": (result.rank, receipt.rank), + "pid": (result.pid, receipt.pid), + "sequence_end": (result.sequence_end, receipt.sequence_end), + "chain_checksum": (result.chain_checksum, receipt.chain_checksum), + "file_sha256": (result.file_sha256, receipt.file_sha256), + "byte_count": (result.byte_count, receipt.byte_count), + "complete": (result.complete, receipt.complete), + } + for name, (actual, expected) in checks.items(): + if actual != expected: + mismatches.append( + f"receipt {name} mismatch: expected {expected!r}, got {actual!r}" + ) + if result.counters and result.counters.to_dict() != receipt.counters.to_dict(): + mismatches.append("receipt counters mismatch") + return mismatches + + +def validate_shard( + path: Path, + *, + expected_receipt: Optional[ShardReceipt] = None, + on_event: Optional[Callable[[TargetedTraceRecord], None]] = None, +) -> ShardValidationResult: + """Validate *path* one line at a time, including sequence/checksum/sentinel.""" + + path = Path(path) + result = ShardValidationResult(path=str(path)) + expected_sequence = 0 + previous_checksum = ZERO_CHECKSUM + file_hash = hashlib.sha256() + saw_header = False + saw_end = False + + try: + stream = path.open("rb") + except OSError as exc: + result.issues.append(f"open failed: {exc}") + return result + + with stream: + for line_number, raw_line in enumerate(stream, 1): + result.byte_count += len(raw_line) + file_hash.update(raw_line) + if len(raw_line) > MAX_JSONL_LINE_BYTES: + result.issues.append( + f"line {line_number}: exceeds {MAX_JSONL_LINE_BYTES} byte limit" + ) + break + if saw_end: + result.issues.append(f"line {line_number}: data after end sentinel") + break + try: + raw = json.loads(raw_line, parse_constant=_reject_json_constant) + except (UnicodeDecodeError, json.JSONDecodeError, ValueError) as exc: + result.issues.append(f"line {line_number}: corrupt JSON tail: {exc}") + break + if not isinstance(raw, Mapping): + result.issues.append(f"line {line_number}: envelope is not an object") + break + try: + envelope = validate_envelope( + raw, + expected_sequence=expected_sequence, + previous_checksum=previous_checksum, + ) + except (TraceValidationError, TypeError, ValueError, OverflowError) as exc: + result.issues.append(f"line {line_number}: {exc}") + break + + record_type = envelope["record_type"] + payload = envelope["payload"] + if expected_sequence == 0 and record_type != "header": + result.issues.append("line 1: first envelope is not a header") + break + if record_type == "header": + if saw_header or expected_sequence != 0: + result.issues.append(f"line {line_number}: duplicate header") + break + saw_header = True + try: + result.rank = int(payload["rank"]) + result.pid = int(payload["pid"]) + except (KeyError, TypeError, ValueError): + result.issues.append(f"line {line_number}: invalid header rank/pid") + break + elif record_type == "event": + if not saw_header: + result.issues.append(f"line {line_number}: event before header") + break + try: + record = TargetedTraceRecord.from_dict(payload) + except (TraceValidationError, TypeError, ValueError, OverflowError) as exc: + result.issues.append(f"line {line_number}: invalid event: {exc}") + break + if record.context.rank != result.rank or record.context.pid != result.pid: + result.issues.append( + f"line {line_number}: event rank/pid differs from header" + ) + break + result.event_count += 1 + if on_event: + on_event(record) + elif record_type == "end": + if not saw_header: + result.issues.append(f"line {line_number}: end before header") + break + try: + result.counters = ShardCounters.from_dict(payload["counters"]) + except (KeyError, TypeError, TraceValidationError) as exc: + result.issues.append(f"line {line_number}: invalid end counters: {exc}") + break + if result.counters.written != result.event_count: + result.issues.append( + "end counters written does not match event envelope count" + ) + break + saw_end = True + + result.sequence_end = expected_sequence + result.chain_checksum = str(envelope["checksum"]) + previous_checksum = result.chain_checksum + expected_sequence += 1 + + result.file_sha256 = file_hash.hexdigest() + result.complete = saw_end + if not saw_header: + result.issues.append("missing header") + if not saw_end: + result.issues.append("missing end sentinel") + if expected_receipt is not None: + result.issues.extend(_receipt_mismatches(result, expected_receipt)) + result.valid = not result.issues and result.complete + return result + + +def postprocess_trace_dir( + trace_dir: Path, + *, + output_path: Optional[Path] = None, + strict: bool = False, +) -> Dict[str, Any]: + """Validate/aggregate a targeted trace directory without loading its events.""" + + trace_dir = Path(trace_dir) + manifest_path = trace_dir / "manifest.json" + manifest: Optional[TargetedTraceManifest] = None + manifest_error: Optional[str] = None + if manifest_path.is_file(): + try: + raw_manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + if not isinstance(raw_manifest, Mapping): + raise TraceValidationError("manifest root must be an object") + manifest = TargetedTraceManifest.from_dict(raw_manifest) + except (OSError, json.JSONDecodeError, TraceValidationError) as exc: + manifest_error = str(exc) + else: + manifest_error = "manifest.json is missing" + + receipts: Dict[str, ShardReceipt] = {} + if manifest is not None: + for receipt in manifest.shards: + receipt_path = Path(receipt.path) + receipts[str(receipt_path.resolve())] = receipt + receipts[receipt_path.name] = receipt + + aggregates: Dict[str, Dict[str, int]] = { + "by_target": {}, + "by_kind": {}, + "by_rank": {}, + } + + def observe(record: TargetedTraceRecord) -> None: + target_id = record.identity.target_id + aggregates["by_target"][target_id] = ( + aggregates["by_target"].get(target_id, 0) + 1 + ) + aggregates["by_kind"][record.kind] = ( + aggregates["by_kind"].get(record.kind, 0) + 1 + ) + rank = str(record.context.rank) + aggregates["by_rank"][rank] = aggregates["by_rank"].get(rank, 0) + 1 + + shard_paths = sorted((trace_dir / "shards").glob("*.jsonl")) + if not shard_paths: + shard_paths = sorted(trace_dir.glob("*.jsonl")) + validations: List[ShardValidationResult] = [] + for path in shard_paths: + receipt = receipts.get(str(path.resolve())) or receipts.get(path.name) + validations.append( + validate_shard(path, expected_receipt=receipt, on_event=observe) + ) + + counter_items = [item.counters for item in validations if item.counters] + coverage = ( + ShardCounters.aggregate(counter_items).to_dict() + if counter_items + else ShardCounters().to_dict() + ) + issues: List[str] = [] + if manifest_error: + issues.append(f"manifest: {manifest_error}") + if manifest is not None: + adapter_warnings = manifest.provenance.get("adapter_warnings", []) + if isinstance(adapter_warnings, list): + issues.extend(f"acquisition: {warning}" for warning in adapter_warnings) + if manifest is not None: + observed_names = {Path(item.path).name for item in validations} + declared_names = {Path(item.path).name for item in manifest.shards} + for receipt in manifest.shards: + if Path(receipt.path).name not in observed_names: + issues.append(f"manifest shard missing: {receipt.path}") + for undeclared in sorted(observed_names - declared_names): + issues.append(f"undeclared trace shard: {undeclared}") + for item in validations: + issues.extend(f"{item.path}: {issue}" for issue in item.issues) + if not shard_paths: + issues.append(f"no trace shards found under {trace_dir}") + + integrity_failures: Dict[str, int] = {} + for issue in issues: + if "corrupt JSON tail" in issue: + reason = "corrupt_tail" + elif "missing end sentinel" in issue: + reason = "missing_end_sentinel" + elif "checksum" in issue: + reason = "checksum" + elif "receipt" in issue: + reason = "receipt" + elif issue.startswith("acquisition:"): + reason = "acquisition" + elif issue.startswith("manifest:"): + reason = "manifest" + else: + reason = "other" + integrity_failures[reason] = integrity_failures.get(reason, 0) + 1 + + summary: Dict[str, Any] = { + "schema_name": manifest.schema_name if manifest else None, + "schema_version": manifest.schema_version if manifest else None, + "run_id": manifest.run_id if manifest else None, + "valid": not issues and all(item.valid for item in validations), + "streaming": True, + "coverage": coverage, + "events": aggregates, + "integrity_failures_by_reason": dict(sorted(integrity_failures.items())), + "shards": [item.to_dict() for item in validations], + "issues": issues, + } + if output_path is not None: + output_path = Path(output_path) + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text( + json.dumps(summary, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + if strict and issues: + raise TraceValidationError("; ".join(issues)) + return summary diff --git a/Magpie/targeted_trace/sampling.py b/Magpie/targeted_trace/sampling.py new file mode 100644 index 0000000..3944530 --- /dev/null +++ b/Magpie/targeted_trace/sampling.py @@ -0,0 +1,43 @@ +"""Deterministic sampling keyed only by run seed and stable event key.""" + +from __future__ import annotations + +import hashlib +from typing import Any, Mapping + +from .schema import canonical_json + + +def sampling_value(run_seed: str, stable_event_key: str) -> float: + """Map ``{run_seed, stable_event_key}`` to a reproducible value in [0, 1). + + Python's process-randomized ``hash()`` and mutable PRNG state are deliberately + excluded so the same semantic event makes the same decision across processes + and reruns. + """ + + digest = hashlib.sha256( + f"{run_seed}\x00{stable_event_key}".encode("utf-8") + ).digest() + return int.from_bytes(digest[:8], "big") / float(1 << 64) + + +def should_sample(run_seed: str, stable_event_key: str, sample_rate: float) -> bool: + """Return the deterministic sampling decision for one event.""" + + if not 0.0 <= sample_rate <= 1.0: + raise ValueError(f"sample_rate must be between 0 and 1, got {sample_rate}") + if sample_rate == 0.0: + return False + if sample_rate == 1.0: + return True + return sampling_value(run_seed, stable_event_key) < sample_rate + + +def stable_key(parts: Mapping[str, Any], *, occurrence: int = 0) -> str: + """Build a stable event key from canonical semantic parts and occurrence.""" + + if occurrence < 0: + raise ValueError("occurrence must be non-negative") + payload = {"parts": dict(parts), "occurrence": occurrence} + return hashlib.sha256(canonical_json(payload).encode("utf-8")).hexdigest() diff --git a/Magpie/targeted_trace/schema.py b/Magpie/targeted_trace/schema.py new file mode 100644 index 0000000..5628347 --- /dev/null +++ b/Magpie/targeted_trace/schema.py @@ -0,0 +1,690 @@ +"""Versioned data contract for targeted kernel trace artifacts. + +The schema intentionally describes evidence, not a particular framework patching +mechanism. Runtime probes, Torch profiler adapters, and future TraceLens adapters +all emit the same records and shard receipts. +""" + +from __future__ import annotations + +import hashlib +import json +import re +from dataclasses import asdict, dataclass, field +from datetime import datetime, timezone +from typing import Any, Dict, Iterable, List, Mapping, Optional, Sequence, Tuple + + +SCHEMA_NAME = "magpie.targeted-kernel-trace" +SCHEMA_VERSION = "1.0.0" +ZERO_CHECKSUM = "0" * 64 +ENVELOPE_TYPES = frozenset({"header", "event", "end"}) +_SHA256 = re.compile(r"^[0-9a-f]{64}$") + + +class TraceValidationError(ValueError): + """Raised when a trace artifact violates the semantic contract.""" + + +def utc_now() -> str: + """Return an RFC3339 timestamp in UTC.""" + + return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z") + + +def canonical_json(value: Any) -> str: + """Serialize *value* deterministically for hashing and golden fixtures.""" + + return json.dumps( + value, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=False, + allow_nan=False, + ) + + +def sha256_json(value: Any) -> str: + """Return the SHA-256 digest of canonical JSON.""" + + return hashlib.sha256(canonical_json(value).encode("utf-8")).hexdigest() + + +def _require_nonempty(value: str, field_name: str) -> None: + if not isinstance(value, str) or not value.strip(): + raise TraceValidationError(f"{field_name} must be a non-empty string") + + +def _require_sha256(value: str, field_name: str) -> None: + if not isinstance(value, str) or not _SHA256.fullmatch(value): + raise TraceValidationError( + f"{field_name} must be a lowercase SHA-256 hex string" + ) + + +def _json_copy(value: Any) -> Any: + """Validate JSON compatibility and return a detached canonical copy.""" + + try: + return json.loads(canonical_json(value)) + except (TypeError, ValueError) as exc: + raise TraceValidationError(f"value is not canonical JSON: {exc}") from exc + + +@dataclass(frozen=True) +class SourceEvidence: + """Python-visible launch or wrapper source evidence.""" + + path: str + line: Optional[int] = None + function: Optional[str] = None + sha256: Optional[str] = None + + def __post_init__(self) -> None: + _require_nonempty(self.path, "source.path") + if self.line is not None and self.line <= 0: + raise TraceValidationError("source.line must be positive") + if self.sha256 is not None: + _require_sha256(self.sha256, "source.sha256") + + def to_dict(self) -> Dict[str, Any]: + return asdict(self) + + @classmethod + def from_dict(cls, data: Mapping[str, Any]) -> "SourceEvidence": + return cls( + path=str(data.get("path", "")), + line=int(data["line"]) if data.get("line") is not None else None, + function=( + str(data["function"]) if data.get("function") is not None else None + ), + sha256=str(data["sha256"]) if data.get("sha256") is not None else None, + ) + + +@dataclass(frozen=True) +class TensorEvidence: + """Host-visible tensor metadata; tensor contents and raw pointers are excluded.""" + + name: str + shape: Tuple[Any, ...] + dtype: str + stride: Optional[Tuple[Any, ...]] = None + device: Optional[str] = None + layout: Optional[str] = None + requires_grad: Optional[bool] = None + + def __post_init__(self) -> None: + _require_nonempty(self.name, "tensor.name") + _require_nonempty(self.dtype, "tensor.dtype") + _json_copy(list(self.shape)) + if self.stride is not None: + _json_copy(list(self.stride)) + + def to_dict(self) -> Dict[str, Any]: + data = asdict(self) + data["shape"] = list(self.shape) + if self.stride is not None: + data["stride"] = list(self.stride) + return data + + @classmethod + def from_dict(cls, data: Mapping[str, Any]) -> "TensorEvidence": + stride = data.get("stride") + return cls( + name=str(data.get("name", "")), + shape=tuple(data.get("shape", [])), + dtype=str(data.get("dtype", "unknown")), + stride=tuple(stride) if isinstance(stride, (list, tuple)) else None, + device=str(data["device"]) if data.get("device") is not None else None, + layout=str(data["layout"]) if data.get("layout") is not None else None, + requires_grad=( + bool(data["requires_grad"]) + if data.get("requires_grad") is not None + else None + ), + ) + + +@dataclass(frozen=True) +class TraceIdentity: + """Stable identity and provenance for one target observation.""" + + run_id: str + target_id: str + variant_id: str = "baseline" + package: Optional[str] = None + image: Optional[str] = None + source_hashes: Mapping[str, str] = field(default_factory=dict) + provenance_hashes: Mapping[str, str] = field(default_factory=dict) + + def __post_init__(self) -> None: + _require_nonempty(self.run_id, "identity.run_id") + _require_nonempty(self.target_id, "identity.target_id") + _require_nonempty(self.variant_id, "identity.variant_id") + _json_copy(dict(self.source_hashes)) + _json_copy(dict(self.provenance_hashes)) + + def to_dict(self) -> Dict[str, Any]: + data = asdict(self) + data["source_hashes"] = dict(self.source_hashes) + data["provenance_hashes"] = dict(self.provenance_hashes) + return data + + @classmethod + def from_dict(cls, data: Mapping[str, Any]) -> "TraceIdentity": + return cls( + run_id=str(data.get("run_id", "")), + target_id=str(data.get("target_id", "")), + variant_id=str(data.get("variant_id", "baseline")), + package=str(data["package"]) if data.get("package") is not None else None, + image=str(data["image"]) if data.get("image") is not None else None, + source_hashes=dict(data.get("source_hashes", {})), + provenance_hashes=dict(data.get("provenance_hashes", {})), + ) + + +@dataclass(frozen=True) +class TraceContext: + """Execution context used for rank/stage/graph attribution.""" + + framework: str + rank: int + pid: int + framework_version: Optional[str] = None + stage: str = "unknown" + execution_mode: str = "unknown" + graph_id: Optional[str] = None + world_size: Optional[int] = None + + def __post_init__(self) -> None: + _require_nonempty(self.framework, "context.framework") + if self.rank < 0: + raise TraceValidationError("context.rank must be non-negative") + if self.pid < 0: + raise TraceValidationError("context.pid must be non-negative") + if self.world_size is not None and self.world_size <= 0: + raise TraceValidationError("context.world_size must be positive") + if self.world_size is not None and self.rank >= self.world_size: + raise TraceValidationError("context.rank must be less than world_size") + + def to_dict(self) -> Dict[str, Any]: + return asdict(self) + + @classmethod + def from_dict(cls, data: Mapping[str, Any]) -> "TraceContext": + return cls( + framework=str(data.get("framework", "unknown")), + rank=int(data.get("rank", 0)), + pid=int(data.get("pid", 0)), + framework_version=( + str(data["framework_version"]) + if data.get("framework_version") is not None + else None + ), + stage=str(data.get("stage", "unknown")), + execution_mode=str(data.get("execution_mode", "unknown")), + graph_id=( + str(data["graph_id"]) if data.get("graph_id") is not None else None + ), + world_size=( + int(data["world_size"]) + if data.get("world_size") is not None + else None + ), + ) + + +@dataclass(frozen=True) +class LaunchSemantics: + """Python-visible invocation semantics needed to reproduce a launch.""" + + source: Optional[SourceEvidence] = None + tensors: Tuple[TensorEvidence, ...] = () + named_scalars: Mapping[str, Any] = field(default_factory=dict) + constexpr: Mapping[str, Any] = field(default_factory=dict) + meta: Mapping[str, Any] = field(default_factory=dict) + python_grid: Any = None + + def __post_init__(self) -> None: + _json_copy(dict(self.named_scalars)) + _json_copy(dict(self.constexpr)) + _json_copy(dict(self.meta)) + _json_copy(self.python_grid) + + def to_dict(self) -> Dict[str, Any]: + return { + "source": self.source.to_dict() if self.source else None, + "tensors": [tensor.to_dict() for tensor in self.tensors], + "named_scalars": _json_copy(dict(self.named_scalars)), + "constexpr": _json_copy(dict(self.constexpr)), + "meta": _json_copy(dict(self.meta)), + "python_grid": _json_copy(self.python_grid), + } + + @classmethod + def from_dict(cls, data: Mapping[str, Any]) -> "LaunchSemantics": + source = data.get("source") + return cls( + source=( + SourceEvidence.from_dict(source) + if isinstance(source, Mapping) + else None + ), + tensors=tuple( + TensorEvidence.from_dict(item) + for item in data.get("tensors", []) + if isinstance(item, Mapping) + ), + named_scalars=dict(data.get("named_scalars", {})), + constexpr=dict(data.get("constexpr", {})), + meta=dict(data.get("meta", {})), + python_grid=data.get("python_grid"), + ) + + +@dataclass(frozen=True) +class RuntimeEvidence: + """Profiler/runtime evidence. Unknown fields remain ``None``, never zero.""" + + cpu_uid: Optional[str] = None + correlation_id: Optional[str] = None + gpu_uid: Optional[str] = None + gpu_symbol: Optional[str] = None + grid: Optional[Tuple[int, ...]] = None + block: Optional[Tuple[int, ...]] = None + stream: Optional[str] = None + duration_us: Optional[float] = None + timestamp_us: Optional[float] = None + + def __post_init__(self) -> None: + if self.duration_us is not None and self.duration_us < 0: + raise TraceValidationError("runtime.duration_us must be non-negative") + + def to_dict(self) -> Dict[str, Any]: + data = asdict(self) + if self.grid is not None: + data["grid"] = list(self.grid) + if self.block is not None: + data["block"] = list(self.block) + return data + + @classmethod + def from_dict(cls, data: Mapping[str, Any]) -> "RuntimeEvidence": + grid = data.get("grid") + block = data.get("block") + return cls( + cpu_uid=str(data["cpu_uid"]) if data.get("cpu_uid") is not None else None, + correlation_id=( + str(data["correlation_id"]) + if data.get("correlation_id") is not None + else None + ), + gpu_uid=str(data["gpu_uid"]) if data.get("gpu_uid") is not None else None, + gpu_symbol=( + str(data["gpu_symbol"]) + if data.get("gpu_symbol") is not None + else None + ), + grid=( + tuple(int(value) for value in grid) + if isinstance(grid, Sequence) + and not isinstance(grid, (str, bytes, bytearray)) + else None + ), + block=( + tuple(int(value) for value in block) + if isinstance(block, Sequence) + and not isinstance(block, (str, bytes, bytearray)) + else None + ), + stream=str(data["stream"]) if data.get("stream") is not None else None, + duration_us=( + float(data["duration_us"]) + if data.get("duration_us") is not None + else None + ), + timestamp_us=( + float(data["timestamp_us"]) + if data.get("timestamp_us") is not None + else None + ), + ) + + +@dataclass(frozen=True) +class TargetedTraceRecord: + """One semantic and/or runtime observation for a selected target.""" + + kind: str + stable_event_key: str + identity: TraceIdentity + context: TraceContext + semantics: LaunchSemantics = field(default_factory=LaunchSemantics) + runtime: RuntimeEvidence = field(default_factory=RuntimeEvidence) + timestamp_ns: Optional[int] = None + warnings: Tuple[str, ...] = () + + def __post_init__(self) -> None: + _require_nonempty(self.kind, "record.kind") + _require_nonempty(self.stable_event_key, "record.stable_event_key") + if self.timestamp_ns is not None and self.timestamp_ns < 0: + raise TraceValidationError("record.timestamp_ns must be non-negative") + + def to_dict(self) -> Dict[str, Any]: + return { + "kind": self.kind, + "stable_event_key": self.stable_event_key, + "identity": self.identity.to_dict(), + "context": self.context.to_dict(), + "semantics": self.semantics.to_dict(), + "runtime": self.runtime.to_dict(), + "timestamp_ns": self.timestamp_ns, + "warnings": list(self.warnings), + } + + @classmethod + def from_dict(cls, data: Mapping[str, Any]) -> "TargetedTraceRecord": + try: + identity = data["identity"] + context = data["context"] + except KeyError as exc: + raise TraceValidationError(f"record missing {exc.args[0]}") from exc + if not isinstance(identity, Mapping) or not isinstance(context, Mapping): + raise TraceValidationError("record identity/context must be objects") + semantics = data.get("semantics", {}) + runtime = data.get("runtime", {}) + return cls( + kind=str(data.get("kind", "")), + stable_event_key=str(data.get("stable_event_key", "")), + identity=TraceIdentity.from_dict(identity), + context=TraceContext.from_dict(context), + semantics=LaunchSemantics.from_dict( + semantics if isinstance(semantics, Mapping) else {} + ), + runtime=RuntimeEvidence.from_dict( + runtime if isinstance(runtime, Mapping) else {} + ), + timestamp_ns=( + int(data["timestamp_ns"]) + if data.get("timestamp_ns") is not None + else None + ), + warnings=tuple(str(item) for item in data.get("warnings", [])), + ) + + +@dataclass +class ShardCounters: + """Loss-accounting counters for one PID/rank shard.""" + + seen: int = 0 + sampled: int = 0 + written: int = 0 + dropped: int = 0 + dropped_by_reason: Dict[str, int] = field(default_factory=dict) + + def note_drop(self, reason: str) -> None: + _require_nonempty(reason, "drop reason") + self.dropped += 1 + self.dropped_by_reason[reason] = self.dropped_by_reason.get(reason, 0) + 1 + + def validate(self) -> None: + values = (self.seen, self.sampled, self.written, self.dropped) + if any(not isinstance(value, int) or value < 0 for value in values): + raise TraceValidationError("coverage counters must be non-negative integers") + if sum(self.dropped_by_reason.values()) != self.dropped: + raise TraceValidationError("dropped_by_reason does not sum to dropped") + if self.seen != self.written + self.dropped: + raise TraceValidationError("seen must equal written + dropped") + unsampled = self.dropped_by_reason.get("sampling", 0) + if self.sampled != self.written + self.dropped - unsampled: + raise TraceValidationError( + "sampled must equal written + non-sampling drops" + ) + + def to_dict(self) -> Dict[str, Any]: + self.validate() + return { + "seen": self.seen, + "sampled": self.sampled, + "written": self.written, + "dropped": self.dropped, + "dropped_by_reason": dict(sorted(self.dropped_by_reason.items())), + } + + @classmethod + def from_dict(cls, data: Mapping[str, Any]) -> "ShardCounters": + counters = cls( + seen=int(data.get("seen", 0)), + sampled=int(data.get("sampled", 0)), + written=int(data.get("written", 0)), + dropped=int(data.get("dropped", 0)), + dropped_by_reason={ + str(key): int(value) + for key, value in dict(data.get("dropped_by_reason", {})).items() + }, + ) + counters.validate() + return counters + + @classmethod + def aggregate(cls, counters: Iterable["ShardCounters"]) -> "ShardCounters": + result = cls() + for item in counters: + result.seen += item.seen + result.sampled += item.sampled + result.written += item.written + result.dropped += item.dropped + for reason, count in item.dropped_by_reason.items(): + result.dropped_by_reason[reason] = ( + result.dropped_by_reason.get(reason, 0) + count + ) + result.validate() + return result + + +@dataclass(frozen=True) +class ShardReceipt: + """Integrity and coverage receipt for a completed shard.""" + + path: str + rank: int + pid: int + sequence_end: int + chain_checksum: str + file_sha256: str + byte_count: int + counters: ShardCounters + complete: bool = True + + def __post_init__(self) -> None: + _require_nonempty(self.path, "receipt.path") + if self.rank < 0 or self.pid < 0 or self.sequence_end < 0: + raise TraceValidationError("receipt rank/pid/sequence are invalid") + _require_sha256(self.chain_checksum, "receipt.chain_checksum") + _require_sha256(self.file_sha256, "receipt.file_sha256") + if self.byte_count < 0: + raise TraceValidationError("receipt.byte_count must be non-negative") + self.counters.validate() + + def to_dict(self) -> Dict[str, Any]: + return { + "path": self.path, + "rank": self.rank, + "pid": self.pid, + "sequence_end": self.sequence_end, + "chain_checksum": self.chain_checksum, + "file_sha256": self.file_sha256, + "byte_count": self.byte_count, + "counters": self.counters.to_dict(), + "complete": self.complete, + } + + @classmethod + def from_dict(cls, data: Mapping[str, Any]) -> "ShardReceipt": + counters = data.get("counters", {}) + return cls( + path=str(data.get("path", "")), + rank=int(data.get("rank", 0)), + pid=int(data.get("pid", 0)), + sequence_end=int(data.get("sequence_end", 0)), + chain_checksum=str(data.get("chain_checksum", "")), + file_sha256=str(data.get("file_sha256", "")), + byte_count=int(data.get("byte_count", 0)), + counters=ShardCounters.from_dict( + counters if isinstance(counters, Mapping) else {} + ), + complete=bool(data.get("complete", False)), + ) + + +@dataclass(frozen=True) +class TargetedTraceManifest: + """Run-level manifest consumed by Apex and other evidence clients.""" + + run_id: str + acquisition_backend: str + targets: Tuple[Mapping[str, Any], ...] + shards: Tuple[ShardReceipt, ...] + provenance: Mapping[str, Any] = field(default_factory=dict) + created_at: str = field(default_factory=utc_now) + pass_kind: str = "diagnostic" + reward_eligible: bool = False + schema_name: str = SCHEMA_NAME + schema_version: str = SCHEMA_VERSION + + def __post_init__(self) -> None: + _require_nonempty(self.run_id, "manifest.run_id") + _require_nonempty(self.acquisition_backend, "manifest.acquisition_backend") + if self.schema_name != SCHEMA_NAME or self.schema_version != SCHEMA_VERSION: + raise TraceValidationError( + f"unsupported schema {self.schema_name}@{self.schema_version}" + ) + if self.pass_kind != "diagnostic": + raise TraceValidationError("targeted trace artifacts must be diagnostic") + if self.reward_eligible: + raise TraceValidationError("targeted trace artifacts cannot be reward eligible") + _json_copy([dict(item) for item in self.targets]) + _json_copy(dict(self.provenance)) + receipt_keys = [ + (receipt.rank, receipt.pid, receipt.path) for receipt in self.shards + ] + if len(receipt_keys) != len(set(receipt_keys)): + raise TraceValidationError("manifest contains duplicate shard receipts") + + @property + def coverage(self) -> ShardCounters: + return ShardCounters.aggregate(receipt.counters for receipt in self.shards) + + def to_dict(self) -> Dict[str, Any]: + return { + "schema_name": self.schema_name, + "schema_version": self.schema_version, + "run_id": self.run_id, + "created_at": self.created_at, + "pass_kind": self.pass_kind, + "reward_eligible": self.reward_eligible, + "acquisition_backend": self.acquisition_backend, + "targets": [_json_copy(dict(item)) for item in self.targets], + "provenance": _json_copy(dict(self.provenance)), + "coverage": self.coverage.to_dict(), + "shards": [receipt.to_dict() for receipt in self.shards], + } + + @classmethod + def from_dict(cls, data: Mapping[str, Any]) -> "TargetedTraceManifest": + schema_name = str(data.get("schema_name", "")) + schema_version = str(data.get("schema_version", "")) + if schema_name != SCHEMA_NAME or schema_version != SCHEMA_VERSION: + raise TraceValidationError( + f"unsupported schema {schema_name or ''}@" + f"{schema_version or ''}; expected " + f"{SCHEMA_NAME}@{SCHEMA_VERSION}" + ) + raw_targets = data.get("targets") + raw_shards = data.get("shards") + if not isinstance(raw_targets, list) or any( + not isinstance(item, Mapping) for item in raw_targets + ): + raise TraceValidationError("manifest targets must be a list of objects") + if not isinstance(raw_shards, list) or any( + not isinstance(item, Mapping) for item in raw_shards + ): + raise TraceValidationError("manifest shards must be a list of objects") + manifest = cls( + run_id=str(data.get("run_id", "")), + acquisition_backend=str(data.get("acquisition_backend", "")), + targets=tuple(dict(item) for item in raw_targets), + shards=tuple(ShardReceipt.from_dict(item) for item in raw_shards), + provenance=dict(data.get("provenance", {})), + created_at=str(data.get("created_at", utc_now())), + pass_kind=str(data.get("pass_kind", "")), + reward_eligible=bool(data.get("reward_eligible", True)), + schema_name=schema_name, + schema_version=schema_version, + ) + expected_coverage = data.get("coverage") + if not isinstance(expected_coverage, Mapping): + raise TraceValidationError("manifest coverage must be an object") + if ShardCounters.from_dict(expected_coverage).to_dict() != manifest.coverage.to_dict(): + raise TraceValidationError("manifest coverage does not match shard receipts") + return manifest + + +def build_envelope( + *, record_type: str, sequence: int, previous_checksum: str, payload: Mapping[str, Any] +) -> Dict[str, Any]: + """Build one checksummed JSONL envelope.""" + + if record_type not in ENVELOPE_TYPES: + raise TraceValidationError(f"unknown envelope record_type: {record_type}") + if sequence < 0: + raise TraceValidationError("envelope.sequence must be non-negative") + _require_sha256(previous_checksum, "previous_checksum") + body = { + "schema_name": SCHEMA_NAME, + "schema_version": SCHEMA_VERSION, + "record_type": record_type, + "sequence": sequence, + "previous_checksum": previous_checksum, + "payload": _json_copy(dict(payload)), + } + body["checksum"] = sha256_json(body) + return body + + +def validate_envelope( + data: Mapping[str, Any], *, expected_sequence: int, previous_checksum: str +) -> Dict[str, Any]: + """Validate one shard envelope and return a detached dictionary.""" + + schema_name = data.get("schema_name") + schema_version = data.get("schema_version") + if schema_name != SCHEMA_NAME or schema_version != SCHEMA_VERSION: + raise TraceValidationError( + f"unsupported envelope schema {schema_name}@{schema_version}" + ) + record_type = data.get("record_type") + if record_type not in ENVELOPE_TYPES: + raise TraceValidationError(f"invalid record_type {record_type!r}") + if data.get("sequence") != expected_sequence: + raise TraceValidationError( + f"sequence mismatch: expected {expected_sequence}, got {data.get('sequence')}" + ) + if data.get("previous_checksum") != previous_checksum: + raise TraceValidationError("previous checksum mismatch") + checksum = data.get("checksum") + if not isinstance(checksum, str): + raise TraceValidationError("missing envelope checksum") + body = dict(data) + body.pop("checksum", None) + try: + observed_checksum = sha256_json(body) + except (TypeError, ValueError, OverflowError) as exc: + raise TraceValidationError("envelope is not canonical JSON") from exc + if observed_checksum != checksum: + raise TraceValidationError("envelope checksum mismatch") + if not isinstance(data.get("payload"), Mapping): + raise TraceValidationError("envelope payload must be an object") + return _json_copy(dict(data)) diff --git a/Magpie/targeted_trace/serialization.py b/Magpie/targeted_trace/serialization.py new file mode 100644 index 0000000..bc29177 --- /dev/null +++ b/Magpie/targeted_trace/serialization.py @@ -0,0 +1,215 @@ +"""Side-effect-free serialization of Python-visible launch metadata.""" + +from __future__ import annotations + +import hashlib +from functools import lru_cache +from collections.abc import Mapping, Sequence +from pathlib import Path +from typing import Any, Dict, Iterable, List, Optional, Tuple + +from .schema import LaunchSemantics, SourceEvidence, TensorEvidence + + +MAX_DEPTH = 3 +MAX_ITEMS = 64 + + +def is_trace_unsafe_proxy(value: Any) -> bool: + """Return whether inspecting *value* can perturb Torch tracing/compilation.""" + + typ = type(value) + module = getattr(typ, "__module__", "") + name = getattr(typ, "__name__", "") + markers = ("torch.fx", "proxy_tensor", "fake_tensor", "torch._subclasses") + return any(marker in module for marker in markers) or "Proxy" in name or "FakeTensor" in name + + +def is_tensor_like(value: Any) -> bool: + """Recognize tensors without importing Torch or reading device memory.""" + + return ( + not is_trace_unsafe_proxy(value) + and hasattr(value, "shape") + and hasattr(value, "dtype") + and callable(getattr(value, "stride", None)) + ) + + +def _dimension(value: Any) -> Any: + try: + return int(value) + except (TypeError, ValueError): + return str(value) + + +def tensor_evidence(name: str, value: Any) -> TensorEvidence: + """Capture shape/dtype/stride while never touching tensor contents.""" + + shape = tuple(_dimension(item) for item in list(getattr(value, "shape", ()))) + stride: Optional[Tuple[Any, ...]] + try: + stride = tuple(_dimension(item) for item in list(value.stride())) + except Exception: + stride = None + requires_grad: Optional[bool] + try: + requires_grad = bool(getattr(value, "requires_grad")) + except Exception: + requires_grad = None + return TensorEvidence( + name=name, + shape=shape, + dtype=str(getattr(value, "dtype", "unknown")), + stride=stride, + device=( + str(getattr(value, "device")) + if getattr(value, "device", None) is not None + else None + ), + layout=( + str(getattr(value, "layout")) + if getattr(value, "layout", None) is not None + else None + ), + requires_grad=requires_grad, + ) + + +def json_safe(value: Any, *, depth: int = 0) -> Any: + """Convert host metadata to bounded deterministic JSON without raw pointers.""" + + if depth > MAX_DEPTH: + return {"type": type(value).__name__, "truncated": True} + if value is None or isinstance(value, (bool, int, str)): + return value + if isinstance(value, float): + if value != value or value in (float("inf"), float("-inf")): + return str(value) + return value + if is_trace_unsafe_proxy(value): + typ = type(value) + return { + "type": getattr(typ, "__name__", "proxy"), + "module": getattr(typ, "__module__", ""), + "unavailable": "torch_tracing_proxy", + } + if is_tensor_like(value): + return tensor_evidence("value", value).to_dict() + if isinstance(value, Mapping): + items = list(value.items()) + result = { + str(key): json_safe(item, depth=depth + 1) + for key, item in items[:MAX_ITEMS] + } + if len(items) > MAX_ITEMS: + result["_truncated"] = len(items) - MAX_ITEMS + return result + if isinstance(value, tuple): + values = list(value) + return { + "type": "tuple", + "items": [json_safe(item, depth=depth + 1) for item in values[:MAX_ITEMS]], + **({"truncated": len(values) - MAX_ITEMS} if len(values) > MAX_ITEMS else {}), + } + if isinstance(value, Sequence) and not isinstance(value, (bytes, bytearray)): + values = list(value) + return [json_safe(item, depth=depth + 1) for item in values[:MAX_ITEMS]] + if callable(value): + return { + "type": "callable", + "module": getattr(value, "__module__", ""), + "qualified_name": getattr(value, "__qualname__", getattr(value, "__name__", "")), + } + return { + "type": type(value).__name__, + "module": type(value).__module__, + } + + +def source_evidence( + path: Optional[str], *, line: Optional[int] = None, function: Optional[str] = None +) -> Optional[SourceEvidence]: + """Build source evidence and hash a readable file without requiring it.""" + + if not path: + return None + digest = _cached_source_sha256(str(path)) + return SourceEvidence(path=str(path), line=line, function=function, sha256=digest) + + +@lru_cache(maxsize=256) +def _cached_source_sha256(path: str) -> Optional[str]: + """Hash each immutable launch source at most once per tracing process.""" + + digest = hashlib.sha256() + try: + with Path(path).open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + except OSError: + return None + return digest.hexdigest() + + +def invocation_semantics( + *, + args: Sequence[Any], + kwargs: Mapping[str, Any], + positional_names: Optional[Sequence[str]] = None, + meta_names: Iterable[str] = (), + constexpr_names: Iterable[str] = (), + python_grid: Any = None, + source: Optional[SourceEvidence] = None, +) -> Tuple[LaunchSemantics, Tuple[str, ...]]: + """Split invocation values into tensors, scalars, constexpr, and meta evidence.""" + + names = list(positional_names or ()) + tensors: List[TensorEvidence] = [] + named_scalars: Dict[str, Any] = {} + meta: Dict[str, Any] = {} + constexpr: Dict[str, Any] = {} + warnings: List[str] = [] + meta_set = set(meta_names) + constexpr_set = set(constexpr_names) + + for index, value in enumerate(args): + name = names[index] if index < len(names) else f"arg{index}" + if is_trace_unsafe_proxy(value): + warnings.append(f"{name}:torch_tracing_proxy") + if is_tensor_like(value): + tensors.append(tensor_evidence(name, value)) + continue + serialized = json_safe(value) + if name in constexpr_set: + constexpr[name] = serialized + elif name in meta_set: + meta[name] = serialized + else: + named_scalars[name] = serialized + + for name, value in kwargs.items(): + if is_trace_unsafe_proxy(value): + warnings.append(f"{name}:torch_tracing_proxy") + if is_tensor_like(value): + tensors.append(tensor_evidence(str(name), value)) + continue + serialized = json_safe(value) + if name in constexpr_set: + constexpr[str(name)] = serialized + elif name in meta_set: + meta[str(name)] = serialized + else: + named_scalars[str(name)] = serialized + + return ( + LaunchSemantics( + source=source, + tensors=tuple(tensors), + named_scalars=named_scalars, + constexpr=constexpr, + meta=meta, + python_grid=json_safe(python_grid), + ), + tuple(warnings), + ) diff --git a/Magpie/targeted_trace/torch_profiler.py b/Magpie/targeted_trace/torch_profiler.py new file mode 100644 index 0000000..baffb45 --- /dev/null +++ b/Magpie/targeted_trace/torch_profiler.py @@ -0,0 +1,398 @@ +"""Streaming adapter from PyTorch/Chrome profiler traces to the trace contract.""" + +from __future__ import annotations + +import gzip +import json +import re +from collections import defaultdict +from pathlib import Path +from typing import Any, Dict, Iterable, Iterator, List, Mapping, Optional, Sequence, TextIO, Tuple + +from .config import TargetSpec, TargetedTraceConfig +from .sampling import stable_key +from .schema import ( + LaunchSemantics, + RuntimeEvidence, + SourceEvidence, + TargetedTraceManifest, + TargetedTraceRecord, + TensorEvidence, + TraceContext, + TraceIdentity, +) +from .writer import TraceShardWriter, default_shard_path, write_manifest + + +TRACE_EVENTS_KEY = '"traceEvents"' +RANK_RE = re.compile(r"(?:^|[/_.-])rank[_-]?(\d+)(?:[/_.-]|$)", re.IGNORECASE) + + +def _open_trace(path: Path) -> TextIO: + if path.suffix == ".gz": + return gzip.open(path, "rt", encoding="utf-8") + return path.open("rt", encoding="utf-8") + + +def iter_trace_events(path: Path, *, chunk_size: int = 64 * 1024) -> Iterator[Mapping[str, Any]]: + """Yield ``traceEvents`` entries without materializing the full trace object.""" + + decoder = json.JSONDecoder() + with _open_trace(Path(path)) as stream: + buffer = "" + array_started = False + eof = False + while not array_started: + chunk = stream.read(chunk_size) + if not chunk: + raise ValueError(f"traceEvents array not found in {path}") + buffer += chunk + key_index = buffer.find(TRACE_EVENTS_KEY) + if key_index < 0: + buffer = buffer[-len(TRACE_EVENTS_KEY) :] + continue + array_index = buffer.find("[", key_index + len(TRACE_EVENTS_KEY)) + while array_index < 0: + chunk = stream.read(chunk_size) + if not chunk: + raise ValueError(f"traceEvents array is truncated in {path}") + buffer += chunk + array_index = buffer.find("[", key_index + len(TRACE_EVENTS_KEY)) + buffer = buffer[array_index + 1 :] + array_started = True + + while True: + buffer = buffer.lstrip() + if buffer.startswith(","): + buffer = buffer[1:].lstrip() + if buffer.startswith("]"): + return + if not buffer and eof: + raise ValueError(f"traceEvents array is truncated in {path}") + try: + value, end = decoder.raw_decode(buffer) + except json.JSONDecodeError: + chunk = stream.read(chunk_size) + if not chunk: + if eof: + raise ValueError(f"invalid traceEvents JSON in {path}") + eof = True + else: + buffer += chunk + continue + buffer = buffer[end:] + if isinstance(value, Mapping): + yield value + + +def _first(args: Mapping[str, Any], names: Sequence[str]) -> Any: + lowered = {str(key).lower(): value for key, value in args.items()} + for name in names: + if name.lower() in lowered: + return lowered[name.lower()] + return None + + +def _vector(args: Mapping[str, Any], prefix: str) -> Optional[Tuple[int, ...]]: + direct = _first(args, [prefix, f"{prefix} size", f"{prefix}_size"]) + if isinstance(direct, (list, tuple)): + try: + return tuple(int(item) for item in direct) + except (TypeError, ValueError): + return None + values: List[int] = [] + for axis in ("x", "y", "z"): + value = _first(args, [f"{prefix}_{axis}", f"{prefix} {axis}"]) + if value is None: + if values: + values.append(1) + continue + try: + values.append(int(value)) + except (TypeError, ValueError): + return None + return tuple(values) if values else None + + +def _rank(event: Mapping[str, Any], path: Path) -> int: + args = event.get("args", {}) + if isinstance(args, Mapping): + value = _first(args, ["rank", "global rank", "distributed rank"]) + if value is not None: + try: + return max(0, int(value)) + except (TypeError, ValueError): + pass + match = RANK_RE.search(path.as_posix()) + return int(match.group(1)) if match else 0 + + +def _stage(event: Mapping[str, Any], path: Path) -> str: + args = event.get("args", {}) + if isinstance(args, Mapping): + value = _first(args, ["stage", "inference stage", "phase"]) + if value is not None: + return str(value).lower() + lowered = path.as_posix().lower() + for stage in ("prefilldecode", "decode", "prefill", "mixed"): + if stage in lowered: + return stage + return "unknown" + + +def _tensor_semantics( + args: Mapping[str, Any], +) -> Tuple[Tuple[TensorEvidence, ...], Tuple[str, ...]]: + dims = _first(args, ["Input Dims", "Input Shapes", "shapes"]) + dtypes = _first(args, ["Input type", "Input Types", "dtypes"]) + strides = _first(args, ["Input Strides", "strides"]) + if not isinstance(dims, (list, tuple)): + return (), ("torch_profiler_missing_tensor_shapes",) + if dims and not isinstance(dims[0], (list, tuple)): + dims = [dims] + dtype_items = list(dtypes) if isinstance(dtypes, (list, tuple)) else [] + stride_items = list(strides) if isinstance(strides, (list, tuple)) else [] + tensors: List[TensorEvidence] = [] + warnings: List[str] = [] + for index, shape in enumerate(dims): + if not isinstance(shape, (list, tuple)): + continue + dtype = str(dtype_items[index]) if index < len(dtype_items) else "unknown" + stride_value = stride_items[index] if index < len(stride_items) else None + stride = tuple(stride_value) if isinstance(stride_value, (list, tuple)) else None + if stride is None: + warnings.append(f"arg{index}:torch_profiler_missing_stride") + tensors.append( + TensorEvidence( + name=f"arg{index}", + shape=tuple(shape), + dtype=dtype, + stride=stride, + ) + ) + return tuple(tensors), tuple(warnings) + + +def _source(spec: TargetSpec) -> Optional[SourceEvidence]: + if not spec.source: + return None + return SourceEvidence.from_dict(spec.source) + + +def _is_kernel_event(event: Mapping[str, Any]) -> bool: + if str(event.get("ph", "X")) not in {"X", ""}: + return False + category = str(event.get("cat", "")).lower() + args = event.get("args", {}) + return ( + any(token in category for token in ("kernel", "gpu", "cuda", "hip")) + or isinstance(args, Mapping) + and _first(args, ["stream", "grid", "grid_x", "device"]) is not None + ) + + +def _matching_targets(symbol: str, targets: Sequence[TargetSpec]) -> Iterator[TargetSpec]: + for target in targets: + if target.matches(symbol): + yield target + + +def adapt_torch_profiler_traces( + trace_paths: Iterable[Path], + output_dir: Path, + *, + config: TargetedTraceConfig, + run_id: str, + framework: str, + framework_version: Optional[str] = None, + image: Optional[str] = None, + provenance: Optional[Mapping[str, Any]] = None, +) -> TargetedTraceManifest: + """Stream selected Torch profiler events into checksummed targeted shards.""" + + if not config.enabled: + raise ValueError("targeted trace adapter requires config.enabled=true") + output_dir = Path(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + writers: Dict[Tuple[int, int], TraceShardWriter] = {} + occurrences: defaultdict[str, int] = defaultdict(int) + input_paths = [Path(path) for path in trace_paths] + non_capture_paths = [ + path + for path in input_paths + if "capture_traces" not in {part.lower() for part in path.parts} + ] + selected_paths = sorted(non_capture_paths or input_paths) + adapter_warnings: List[str] = [] + + def get_writer(rank: int, pid: int) -> TraceShardWriter: + key = (rank, pid) + if key not in writers: + writers[key] = TraceShardWriter( + default_shard_path(output_dir, rank=rank, pid=pid), + run_id=run_id, + rank=rank, + pid=pid, + run_seed=config.run_seed, + sample_rate=config.sample_rate, + max_records=config.max_records_per_shard, + header_metadata={ + "framework": framework, + "framework_version": framework_version, + "capture_backend": "torch_profiler", + }, + ) + return writers[key] + + for trace_path in selected_paths: + try: + events = iter_trace_events(trace_path) + for event in events: + if not _is_kernel_event(event): + continue + symbol = str(event.get("name", "")) + matches = list(_matching_targets(symbol, config.targets)) + if not matches: + continue + rank = _rank(event, trace_path) + try: + pid = max(0, int(event.get("pid", 0))) + except (TypeError, ValueError): + pid = 0 + writer = get_writer(rank, pid) + args = event.get("args", {}) + args = args if isinstance(args, Mapping) else {} + tensors, tensor_warnings = _tensor_semantics(args) + runtime = RuntimeEvidence( + cpu_uid=( + str(_first(args, ["cpu uid", "external id"])) + if _first(args, ["cpu uid", "external id"]) is not None + else None + ), + correlation_id=( + str(_first(args, ["correlation", "correlation id"])) + if _first(args, ["correlation", "correlation id"]) + is not None + else None + ), + gpu_uid=stable_key( + { + "name": symbol, + "rank": rank, + "tid": event.get("tid"), + "ts": event.get("ts"), + } + ), + gpu_symbol=symbol, + grid=_vector(args, "grid"), + block=_vector(args, "block"), + stream=( + str(_first(args, ["stream", "stream id"])) + if _first(args, ["stream", "stream id"]) is not None + else None + ), + duration_us=( + float(event["dur"]) + if event.get("dur") is not None + else None + ), + timestamp_us=( + float(event["ts"]) + if event.get("ts") is not None + else None + ), + ) + graph_id = _first(args, ["graph id", "graph_id", "cuda graph id"]) + execution_mode = "graph" if graph_id is not None else "unknown" + for target in matches: + base_parts = { + "target_id": target.target_id, + "symbol": symbol, + "rank": rank, + "stage": _stage(event, trace_path), + "grid": list(runtime.grid) if runtime.grid else None, + "block": list(runtime.block) if runtime.block else None, + "tensors": [ + { + "shape": list(tensor.shape), + "dtype": tensor.dtype, + "stride": ( + list(tensor.stride) if tensor.stride else None + ), + } + for tensor in tensors + ], + } + token = stable_key(base_parts) + occurrence = occurrences[token] + occurrences[token] += 1 + warnings = list(tensor_warnings) + if target.source is None: + warnings.append("torch_profiler_missing_launch_source") + if runtime.correlation_id is None: + warnings.append("torch_profiler_missing_runtime_correlation") + try: + record = TargetedTraceRecord( + kind="torch_profiler_kernel", + stable_event_key=stable_key( + base_parts, occurrence=occurrence + ), + identity=TraceIdentity( + run_id=run_id, + target_id=target.target_id, + variant_id=target.variant_id, + package=target.package, + image=image, + source_hashes=target.source_hashes, + provenance_hashes=target.provenance_hashes, + ), + context=TraceContext( + framework=framework, + framework_version=framework_version, + rank=rank, + pid=pid, + stage=_stage(event, trace_path), + execution_mode=execution_mode, + graph_id=( + str(graph_id) if graph_id is not None else None + ), + ), + semantics=LaunchSemantics( + source=_source(target), + tensors=tensors, + meta={ + "profiler_category": str(event.get("cat", "")), + }, + ), + runtime=runtime, + timestamp_ns=( + int(float(event["ts"]) * 1000) + if event.get("ts") is not None + else None + ), + warnings=tuple(warnings), + ) + writer.submit(record) + except Exception: + writer.note_failed_observation("serialization_error") + except (OSError, ValueError, json.JSONDecodeError) as exc: + adapter_warnings.append(f"{trace_path}: {exc}") + continue + + receipts = tuple(writer.close() for _, writer in sorted(writers.items())) + manifest = TargetedTraceManifest( + run_id=run_id, + acquisition_backend="torch_profiler", + targets=tuple(target.to_dict() for target in config.targets), + shards=receipts, + provenance={ + **dict(provenance or {}), + "framework": framework, + "framework_version": framework_version, + "image": image, + "input_traces": [str(path) for path in selected_paths], + "adapter_warnings": adapter_warnings, + }, + ) + write_manifest(output_dir / "manifest.json", manifest) + return manifest diff --git a/Magpie/targeted_trace/writer.py b/Magpie/targeted_trace/writer.py new file mode 100644 index 0000000..19c2a62 --- /dev/null +++ b/Magpie/targeted_trace/writer.py @@ -0,0 +1,331 @@ +"""Loss-accounted per-rank shard writer for targeted trace records.""" + +from __future__ import annotations + +import hashlib +import json +import os +import threading +from pathlib import Path +from typing import Any, Mapping, Optional + +from .sampling import should_sample +from .schema import ( + ZERO_CHECKSUM, + ShardCounters, + ShardReceipt, + TargetedTraceManifest, + TargetedTraceRecord, + TraceValidationError, + build_envelope, + canonical_json, + utc_now, +) + + +class TraceShardWriter: + """Write one checksummed JSONL shard. + + Header and end-sentinel envelopes are outside the event budget. Every + observed target event becomes either one written event or one named drop, + making sampling/cap/serialization loss visible to downstream consumers. + """ + + def __init__( + self, + path: Path, + *, + run_id: str, + rank: int, + pid: int, + run_seed: str, + sample_rate: float = 1.0, + max_records: int = 100_000, + header_metadata: Optional[Mapping[str, Any]] = None, + ) -> None: + if rank < 0 or pid < 0: + raise ValueError("rank and pid must be non-negative") + if max_records < 0: + raise ValueError("max_records must be non-negative") + if not 0.0 <= sample_rate <= 1.0: + raise ValueError("sample_rate must be between 0 and 1") + self.path = Path(path) + self.run_id = run_id + self.rank = rank + self.pid = pid + self.run_seed = str(run_seed) + self.sample_rate = float(sample_rate) + self.max_records = int(max_records) + self.counters = ShardCounters() + self._sequence = 0 + self._previous_checksum = ZERO_CHECKSUM + self._file_hash = hashlib.sha256() + self._byte_count = 0 + self._closed = False + self._io_failed = False + self._lock = threading.Lock() + self.path.parent.mkdir(parents=True, exist_ok=True) + # A shard is an append-once artifact. Reusing a PID/rank path would + # silently mix or erase evidence, so require a fresh output directory. + self._file = self.path.open("xb") + try: + self._write_envelope( + "header", + { + "run_id": run_id, + "rank": rank, + "pid": pid, + "run_seed": self.run_seed, + "sample_rate": self.sample_rate, + "max_records": self.max_records, + "metadata": dict(header_metadata or {}), + }, + ) + except Exception: + self._file.close() + raise + + def _write_envelope(self, record_type: str, payload: Mapping[str, Any]) -> None: + envelope = build_envelope( + record_type=record_type, + sequence=self._sequence, + previous_checksum=self._previous_checksum, + payload=payload, + ) + raw = (canonical_json(envelope) + "\n").encode("utf-8") + written = self._file.write(raw) + if written is None: + written = 0 + self._file_hash.update(raw[:written]) + self._byte_count += written + if written != len(raw): + raise OSError(f"short trace shard write: {written}/{len(raw)} bytes") + self._previous_checksum = str(envelope["checksum"]) + self._sequence += 1 + + def submit(self, record: TargetedTraceRecord) -> bool: + """Observe and conditionally write *record*; return whether it was written.""" + + with self._lock: + if self._closed: + raise RuntimeError("cannot submit to a closed trace shard") + self.counters.seen += 1 + + if not should_sample( + self.run_seed, record.stable_event_key, self.sample_rate + ): + self.counters.note_drop("sampling") + return False + + self.counters.sampled += 1 + if self._io_failed: + self.counters.note_drop("io_error") + return False + if self.counters.written >= self.max_records: + self.counters.note_drop("cap") + return False + if record.identity.run_id != self.run_id: + self.counters.note_drop("invalid_record") + return False + if record.context.rank != self.rank or record.context.pid != self.pid: + self.counters.note_drop("invalid_record") + return False + + try: + payload = record.to_dict() + except Exception: + self.counters.note_drop("serialization_error") + return False + + try: + self._write_envelope("event", payload) + except (OSError, ValueError, TypeError): + self._io_failed = True + self.counters.note_drop("io_error") + return False + self.counters.written += 1 + return True + + def note_failed_observation(self, reason: str, *, sampled: bool = True) -> None: + """Account for an event that failed before a typed record could be built.""" + + with self._lock: + if self._closed: + raise RuntimeError("cannot update a closed trace shard") + self.counters.seen += 1 + if sampled: + self.counters.sampled += 1 + self.counters.note_drop(reason) + else: + self.counters.note_drop("sampling") + + def close(self) -> ShardReceipt: + """Write the end sentinel, fsync, and return an integrity receipt.""" + + with self._lock: + if self._closed: + if not hasattr(self, "_receipt"): + raise RuntimeError("closed trace shard has no receipt") + return self._receipt + + self.counters.validate() + complete = not self._io_failed + if complete: + try: + self._write_envelope( + "end", + { + "run_id": self.run_id, + "rank": self.rank, + "pid": self.pid, + "counters": self.counters.to_dict(), + "end_reason": "complete", + }, + ) + self._file.flush() + os.fsync(self._file.fileno()) + except OSError: + complete = False + self._io_failed = True + self._file.close() + self._closed = True + self._receipt = ShardReceipt( + path=str(self.path), + rank=self.rank, + pid=self.pid, + sequence_end=max(0, self._sequence - 1), + chain_checksum=self._previous_checksum, + file_sha256=self._file_hash.hexdigest(), + byte_count=self._byte_count, + counters=self.counters, + complete=complete, + ) + return self._receipt + + def __enter__(self) -> "TraceShardWriter": + return self + + def __exit__(self, exc_type: object, exc: object, tb: object) -> None: + self.close() + + +def write_manifest(path: Path, manifest: TargetedTraceManifest) -> None: + """Atomically write a run manifest.""" + + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + tmp_path = path.with_name(f".{path.name}.{os.getpid()}.tmp") + tmp_path.write_text( + canonical_json(manifest.to_dict()) + "\n", + encoding="utf-8", + ) + os.replace(tmp_path, path) + + +def merge_runtime_manifest( + path: Path, + *, + run_id: str, + receipt: ShardReceipt, + targets: list[Mapping[str, Any]], + provenance: Mapping[str, Any], +) -> TargetedTraceManifest: + """Lock/merge one runtime shard into a multi-process run manifest.""" + + import fcntl + + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + lock_path = path.with_suffix(path.suffix + ".lock") + with lock_path.open("a+b") as lock_file: + fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX) + existing: Optional[TargetedTraceManifest] = None + if path.is_file(): + raw = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(raw, Mapping): + raise TraceValidationError("runtime manifest root must be an object") + existing = TargetedTraceManifest.from_dict(raw) + if existing.run_id != run_id: + raise TraceValidationError( + f"runtime manifest run mismatch: {existing.run_id} != {run_id}" + ) + + merged_provenance = dict(existing.provenance) if existing else {} + for key, value in provenance.items(): + if ( + key in merged_provenance + and merged_provenance[key] is not None + and value is not None + and merged_provenance[key] != value + ): + raise TraceValidationError( + f"runtime manifest provenance conflict for {key!r}" + ) + if value is not None: + merged_provenance[key] = value + + target_map: dict[tuple[str, str], dict[str, Any]] = {} + for raw_target in [*(existing.targets if existing else ()), *targets]: + target = dict(raw_target) + key = ( + str(target.get("target_id", "")), + str(target.get("variant_id", "baseline")), + ) + if key not in target_map: + target_map[key] = target + continue + current = target_map[key] + for field_name in ("package", "image", "source"): + old_value = current.get(field_name) + new_value = target.get(field_name) + if ( + old_value is not None + and new_value is not None + and old_value != new_value + ): + raise TraceValidationError( + f"runtime target {key!r} conflicts on {field_name}" + ) + if old_value is None and new_value is not None: + current[field_name] = new_value + current["name_patterns"] = sorted( + { + *current.get("name_patterns", []), + *target.get("name_patterns", []), + } + ) + for field_name in ("source_hashes", "provenance_hashes"): + merged_hashes = dict(current.get(field_name, {})) + for hash_name, digest in dict(target.get(field_name, {})).items(): + if hash_name in merged_hashes and merged_hashes[hash_name] != digest: + raise TraceValidationError( + f"runtime target {key!r} conflicts on {field_name}." + f"{hash_name}" + ) + merged_hashes[hash_name] = digest + current[field_name] = merged_hashes + + receipt_map = { + (item.rank, item.pid, Path(item.path).name): item + for item in (existing.shards if existing else ()) + } + receipt_map[(receipt.rank, receipt.pid, Path(receipt.path).name)] = receipt + manifest = TargetedTraceManifest( + run_id=run_id, + acquisition_backend="python_runtime", + targets=tuple(target_map[key] for key in sorted(target_map)), + shards=tuple(receipt_map[key] for key in sorted(receipt_map)), + provenance=merged_provenance, + created_at=existing.created_at if existing else utc_now(), + ) + write_manifest(path, manifest) + fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN) + return manifest + + +def default_shard_path(output_dir: Path, *, rank: int, pid: int) -> Path: + """Return the conventional unique shard path.""" + + if rank < 0 or pid < 0: + raise TraceValidationError("rank and pid must be non-negative") + return Path(output_dir) / "shards" / f"trace_pid{pid}_rank{rank}.jsonl" diff --git a/README.md b/README.md index b987275..d7445ed 100644 --- a/README.md +++ b/README.md @@ -10,6 +10,7 @@ A lightweight, general-purpose framework for evaluating GPU kernel correctness a - **Hardware Control**: Hardware-aware kernel evaluation under controlled execution settings - **Auto GPU Selection**: Benchmark mode picks idle GPU(s) before launching (AMD + NVIDIA) - **Trace Analysis**: TraceLens integration with per-stage roofline summaries for benchmark traces +- **Targeted Kernel Evidence**: Deterministic sampling, checksummed per-rank shards, and streaming validation - **MCP Server**: Model Context Protocol integration for AI agents - **Structured Reports**: JSON output for pipeline integration @@ -60,6 +61,9 @@ magpie compare --kernel-config examples/ck_grouped_gemm_compare.yaml # Benchmark vLLM (see examples/benchmarks/*.yaml) magpie benchmark --benchmark-config examples/benchmarks/benchmark_vllm_dsr1.yaml +# Validate an existing TargetedKernelTrace artifact +magpie targeted-trace postprocess --trace-dir results/.../targeted_trace --strict + # GPU / toolchain summary magpie --gpu-info diff --git a/docs/how-to/benchmarking/benchmark.md b/docs/how-to/benchmarking/benchmark.md index 4664491..f058012 100644 --- a/docs/how-to/benchmarking/benchmark.md +++ b/docs/how-to/benchmarking/benchmark.md @@ -73,11 +73,16 @@ results/benchmark_vllm_/ ├── config.yaml # Snapshot of benchmark configuration ├── container_stdout.log # Container stdout ├── container_stderr.log # Container stderr -├── inferencex_result.json # Raw InferenceX output +├── inferencex_result.json # Raw InferenceX output +├── lm_eval/ # Preserved serving accuracy artifacts (RUN_EVAL=true) ├── torch_trace/ # Raw torch profiler traces │ ├── *-rank-0.*.pt.trace.json.gz │ ├── *-rank-1.*.pt.trace.json.gz │ └── ... +├── targeted_trace/ # Diagnostic selected-kernel evidence (if enabled) +│ ├── manifest.json # Schema/provenance/coverage and shard receipts +│ ├── summary.json # Streaming integrity/coverage summary +│ └── shards/ # Checksummed PID/rank JSONL shards ├── gap_analysis/ # Gap analysis output (if enabled) │ ├── gap_analysis.csv # Merged kernel stats across all ranks │ ├── gap_analysis_rank0.csv # Per-rank kernel stats @@ -103,7 +108,20 @@ machine-readable `params_json` for matched TraceLens `param:*` metadata. ## Benchmark report -The primary summary file is **`benchmark_report.json`**, written to the run workspace directory. It aggregates throughput, latency, and optional `gap_analysis` and `tracelens_analysis` sections. A typical shape (abbreviated, with `...` marking elided values): +The primary summary file is **`benchmark_report.json`**, written to the run workspace directory. It aggregates throughput, latency, and optional `gap_analysis` and `tracelens_analysis` sections. + +Every report declares `run_kind` and `reward_eligible`. A +`run_kind: measurement` run rejects heavy profilers; diagnostic runs and all +TargetedKernelTrace artifacts have `reward_eligible: false`. When `RUN_EVAL=true`, +raw lm-eval files remain under `lm_eval/` and `quality_gate` exposes each task's +primary metric. Missing or invalid requested accuracy evidence fails the benchmark. + +Targeted trace selection uses portable symbol glob patterns under +`profiler.targeted_trace.targets`; it does not depend on a fixed container-image +registry. See `Magpie/targeted_trace/README.md` for its artifact contract and +standalone conversion commands. + +A typical report shape (abbreviated, with `...` marking elided values): ```text { diff --git a/examples/benchmarks/benchmark_vllm_qwen3_next_80b_fp8.yaml b/examples/benchmarks/benchmark_vllm_qwen3_next_80b_fp8.yaml index 3bcdb57..82f0fb5 100644 --- a/examples/benchmarks/benchmark_vllm_qwen3_next_80b_fp8.yaml +++ b/examples/benchmarks/benchmark_vllm_qwen3_next_80b_fp8.yaml @@ -14,6 +14,7 @@ benchmark: model: Qwen/Qwen3-Next-80B-A3B-Instruct-FP8 precision: fp8 run_mode: docker + run_kind: diagnostic # Environment variables envs: @@ -27,6 +28,12 @@ benchmark: # AMD performance flags VLLM_ROCM_USE_AITER: 1 + + # Preserve a real correctness signal alongside throughput. Magpie stores + # raw lm-eval artifacts under the benchmark workspace and exposes a bounded + # task/metric summary in benchmark_report.json. + RUN_EVAL: "true" + MAGPIE_EVAL_TASKS: gsm8k profiler: torch_profiler: diff --git a/pyproject.toml b/pyproject.toml index 83aa7f7..15a100e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -73,7 +73,13 @@ where = ["."] include = ["Magpie*"] [tool.setuptools.package-data] -Magpie = ["*.yaml", "*.yaml.example", "*.json", "scripts/benchmark/*.sh"] +Magpie = [ + "*.yaml", + "*.yaml.example", + "*.json", + "scripts/benchmark/*.sh", + "targeted_trace/README.md", +] [tool.black] line-length = 88 diff --git a/tests/fixtures/targeted_trace/torch_trace.json b/tests/fixtures/targeted_trace/torch_trace.json new file mode 100644 index 0000000..1a1cb56 --- /dev/null +++ b/tests/fixtures/targeted_trace/torch_trace.json @@ -0,0 +1,47 @@ +{ + "schemaVersion": 1, + "traceEvents": [ + { + "name": "cpu_only_event", + "cat": "cpu_op", + "ph": "X", + "pid": 17, + "tid": 4, + "ts": 90.0, + "dur": 2.0, + "args": {} + }, + { + "name": "triton_fused_moe_kernel", + "cat": "kernel", + "ph": "X", + "pid": 17, + "tid": 8, + "ts": 100.0, + "dur": 12.5, + "args": { + "rank": 1, + "stage": "decode", + "stream": 7, + "correlation id": 99, + "grid": [32, 2, 1], + "block": [256, 1, 1], + "graph id": "graph-4", + "Input Dims": [[4, 128], [128, 256]], + "Input type": ["bf16", "fp8"], + "Input Strides": [[128, 1], [256, 1]] + } + }, + { + "name": "unselected_kernel", + "cat": "gpu_kernel", + "ph": "X", + "pid": 17, + "tid": 8, + "ts": 120.0, + "dur": 1.0, + "args": {"rank": 1, "stream": 7} + } + ], + "displayTimeUnit": "ns" +} diff --git a/tests/test_benchmark_support.py b/tests/test_benchmark_support.py index 8c88b97..053d82d 100644 --- a/tests/test_benchmark_support.py +++ b/tests/test_benchmark_support.py @@ -20,6 +20,7 @@ ) from Magpie.modes.benchmark.image_selector import ImageSelector from Magpie.modes.benchmark.result import BenchmarkResult, ResultParser +from Magpie.modes.benchmark.quality import parse_lm_eval_quality from Magpie.modes.benchmark.tracelens_inference import ( SGLANG_SHAPE_DISCOVERY_FLAG, TraceLensInferencePipeline, @@ -57,6 +58,162 @@ def test_workspace_manager_makes_docker_mounts_container_writable(tmp_path): assert workspace.stat().st_mode & 0o777 == 0o777 assert (workspace / "torch_trace").stat().st_mode & 0o777 == 0o777 assert (workspace / "system_profile").stat().st_mode & 0o777 == 0o777 + assert (workspace / "targeted_trace").stat().st_mode & 0o777 == 0o777 + + +def test_lm_eval_quality_receipt_preserves_task_metrics(tmp_path): + eval_dir = tmp_path / "lm_eval" / "model" + eval_dir.mkdir(parents=True) + (eval_dir / "results_2026.json").write_text( + json.dumps( + { + "results": { + "gsm8k": { + "exact_match,strict-match": 0.812, + "exact_match_stderr,strict-match": 0.01, + "alias": "gsm8k", + } + } + } + ), + encoding="utf-8", + ) + (eval_dir / "samples_gsm8k.jsonl").write_text("{}\n", encoding="utf-8") + + gate = parse_lm_eval_quality(tmp_path, requested=True) + + assert gate["passed"] is True + assert gate["status"] == "passed" + assert gate["tasks"]["gsm8k"]["primary_metric"] == ( + "exact_match,strict-match" + ) + assert gate["tasks"]["gsm8k"]["value"] == pytest.approx(0.812) + assert gate["tasks"]["gsm8k"]["metrics"] == { + "exact_match,strict-match": 0.812 + } + assert "lm_eval/model/results_2026.json" in gate["artifacts"] + assert gate["result_artifact_receipts"][0]["sha256"] + assert gate["result_artifact_receipts"][0]["size_bytes"] > 0 + + +def test_lm_eval_quality_requested_missing_fails_explicitly(tmp_path): + gate = parse_lm_eval_quality(tmp_path, requested=True) + + assert gate["passed"] is False + assert gate["status"] == "missing" + assert gate["evidence_present"] is False + assert "no lm_eval/results*.json" in gate["errors"][0] + + +@pytest.mark.parametrize( + "payload", + [ + '{"results":{"gsm8k":{"acc,none":NaN}}}', + '{"results":{"gsm8k":{"acc,none":0.5,"acc,none":0.6}}}', + ], +) +def test_lm_eval_quality_rejects_noncanonical_numeric_evidence(tmp_path, payload): + eval_dir = tmp_path / "lm_eval" + eval_dir.mkdir() + (eval_dir / "results_bad.json").write_text(payload, encoding="utf-8") + + gate = parse_lm_eval_quality(tmp_path, requested=True) + + assert gate["passed"] is False + assert gate["status"] == "invalid" + assert gate["errors"] + + +def test_lm_eval_artifact_shell_helper_copies_before_upstream_move(tmp_path): + source = tmp_path / "upstream_eval" + workspace = tmp_path / "workspace" + source.mkdir() + workspace.mkdir() + (source / "results_fixture.json").write_text( + '{"results":{"gsm8k":{"acc,none":0.5}}}\n', + encoding="utf-8", + ) + helper = ( + Path(__file__).parents[1] + / "Magpie" + / "scripts" + / "benchmark" + / "magpie_bench_remote_compat.sh" + ) + proc = subprocess.run( + [ + "bash", + "-c", + 'source "$1"; magpie_preserve_lm_eval_artifacts', + "bash", + str(helper), + ], + env={ + "PATH": "/usr/bin:/bin", + "EVAL_RESULT_DIR": str(source), + "RESULT_DIR": str(workspace), + "MAGPIE_INFERENCEX_ROOT": str(tmp_path / "empty_inferencex"), + }, + capture_output=True, + text=True, + ) + + assert proc.returncode == 0, proc.stderr + assert (workspace / "lm_eval" / "results_fixture.json").is_file() + + +def test_lm_eval_fallback_ignores_stale_inferencex_results(tmp_path): + inferencex = tmp_path / "InferenceX" + workspace = tmp_path / "workspace" + inferencex.mkdir() + workspace.mkdir() + stale = inferencex / "results_stale.json" + stale.write_text("{}\n", encoding="utf-8") + helper = ( + Path(__file__).parents[1] + / "Magpie" + / "scripts" + / "benchmark" + / "magpie_bench_remote_compat.sh" + ) + proc = subprocess.run( + [ + "bash", + "-c", + 'source "$1"; magpie_mark_lm_eval_start; ' + 'touch "$MAGPIE_INFERENCEX_ROOT/results_fresh.json"; ' + "magpie_preserve_lm_eval_artifacts", + "bash", + str(helper), + ], + env={ + "PATH": "/usr/bin:/bin", + "RESULT_DIR": str(workspace), + "MAGPIE_INFERENCEX_ROOT": str(inferencex), + }, + capture_output=True, + text=True, + ) + + assert proc.returncode == 0, proc.stderr + assert not (workspace / "lm_eval" / stale.name).exists() + assert (workspace / "lm_eval" / "results_fresh.json").exists() + + +def test_qwen_acceptance_config_requests_diagnostic_accuracy_evidence(): + config_path = ( + Path(__file__).parents[1] + / "examples" + / "benchmarks" + / "benchmark_vllm_qwen3_next_80b_fp8.yaml" + ) + raw = yaml.safe_load(config_path.read_text(encoding="utf-8"))["benchmark"] + config = BenchmarkConfig.from_dict(raw) + + assert config.run_kind == "diagnostic" + assert config.reward_eligible is False + assert str(config.envs["RUN_EVAL"]).lower() == "true" + assert config.envs["MAGPIE_EVAL_TASKS"] == "gsm8k" def test_benchmark_mode_only_requests_container_writable_workspace_for_docker( diff --git a/tests/test_targeted_trace.py b/tests/test_targeted_trace.py new file mode 100644 index 0000000..d69725d --- /dev/null +++ b/tests/test_targeted_trace.py @@ -0,0 +1,588 @@ +"""Hermetic semantic tests for TargetedKernelTrace.""" + +from __future__ import annotations + +import gzip +import json +from pathlib import Path + +import pytest + +from Magpie.modes.benchmark.config import ( + BenchmarkConfig, + ProfilerConfig, + TorchProfilerConfig, +) +from Magpie.modes.benchmark.targeted_trace import run_targeted_trace_analysis +from Magpie.targeted_trace import ( + RuntimeEvidence, + TargetSpec, + TargetedTraceConfig, + TargetedTraceRecord, + TargetedTraceRecorder, + TraceContext, + TraceIdentity, + TraceValidationError, + adapt_torch_profiler_traces, + iter_trace_events, + postprocess_trace_dir, + validate_shard, +) +from Magpie.targeted_trace.sampling import should_sample, stable_key +from Magpie.targeted_trace.schema import TargetedTraceManifest +from Magpie.targeted_trace.schema import canonical_json +from Magpie.targeted_trace.writer import TraceShardWriter, default_shard_path, write_manifest + + +FIXTURE = Path(__file__).parent / "fixtures" / "targeted_trace" / "torch_trace.json" + + +class FixtureTensor: + """Torch-free tensor metadata fixture.""" + + shape = (4, 8) + dtype = "torch.float16" + device = "cuda:0" + layout = "torch.strided" + requires_grad = False + + def stride(self): + return (8, 1) + + +def make_record(run_id: str, *, rank: int = 0, pid: int = 10, key: str = "event"): + return TargetedTraceRecord( + kind="torch_profiler_kernel", + stable_event_key=key, + identity=TraceIdentity(run_id=run_id, target_id="target"), + context=TraceContext(framework="vllm", rank=rank, pid=pid), + runtime=RuntimeEvidence(gpu_symbol="kernel", duration_us=4.0), + ) + + +def test_sampling_is_reproducible_and_seeded(): + key = stable_key({"target": "moe", "shape": [4, 128]}, occurrence=3) + decisions = [should_sample("run-a", key, 0.5) for _ in range(20)] + assert len(set(decisions)) == 1 + assert should_sample("run-a", key, 0.0) is False + assert should_sample("run-a", key, 1.0) is True + assert stable_key({"shape": [4, 128], "target": "moe"}, occurrence=3) == key + + +def test_shard_round_trip_checksum_sentinel_and_cap(tmp_path): + path = default_shard_path(tmp_path, rank=0, pid=10) + writer = TraceShardWriter( + path, + run_id="run-1", + rank=0, + pid=10, + run_seed="seed", + max_records=1, + ) + assert writer.submit(make_record("run-1", key="one")) is True + assert writer.submit(make_record("run-1", key="two")) is False + receipt = writer.close() + + assert receipt.complete is True + assert receipt.counters.to_dict() == { + "seen": 2, + "sampled": 2, + "written": 1, + "dropped": 1, + "dropped_by_reason": {"cap": 1}, + } + validated = validate_shard(path, expected_receipt=receipt) + assert validated.valid is True + assert validated.event_count == 1 + + +def test_sampling_drop_is_loss_accounted(tmp_path): + path = default_shard_path(tmp_path, rank=0, pid=10) + writer = TraceShardWriter( + path, + run_id="run-sampling", + rank=0, + pid=10, + run_seed="seed", + sample_rate=0.0, + ) + assert writer.submit(make_record("run-sampling")) is False + receipt = writer.close() + + assert receipt.counters.to_dict() == { + "seen": 1, + "sampled": 0, + "written": 0, + "dropped": 1, + "dropped_by_reason": {"sampling": 1}, + } + assert validate_shard(path).valid is True + + +def test_writer_refuses_to_overwrite_existing_shard(tmp_path): + path = default_shard_path(tmp_path, rank=0, pid=10) + writer = TraceShardWriter( + path, run_id="first", rank=0, pid=10, run_seed="seed" + ) + writer.close() + + with pytest.raises(FileExistsError): + TraceShardWriter( + path, run_id="second", rank=0, pid=10, run_seed="seed" + ) + + +def test_corrupt_tail_is_reported_not_silently_skipped(tmp_path): + path = default_shard_path(tmp_path, rank=0, pid=10) + writer = TraceShardWriter( + path, run_id="run-1", rank=0, pid=10, run_seed="seed" + ) + writer.submit(make_record("run-1")) + writer.close() + raw = path.read_bytes() + path.write_bytes(raw[:-20]) + + validated = validate_shard(path) + assert validated.valid is False + assert validated.complete is False + assert any("corrupt JSON tail" in issue for issue in validated.issues) + assert any("missing end sentinel" in issue for issue in validated.issues) + + +def test_checksum_tampering_is_reported(tmp_path): + path = default_shard_path(tmp_path, rank=0, pid=10) + writer = TraceShardWriter( + path, run_id="run-1", rank=0, pid=10, run_seed="seed" + ) + writer.submit(make_record("run-1")) + writer.close() + lines = path.read_text(encoding="utf-8").splitlines() + event = json.loads(lines[1]) + event["payload"]["runtime"]["gpu_symbol"] = "tampered" + lines[1] = canonical_json(event) + path.write_text("\n".join(lines) + "\n", encoding="utf-8") + + validated = validate_shard(path) + assert validated.valid is False + assert any("checksum mismatch" in issue for issue in validated.issues) + + +def test_nonfinite_envelope_is_reported_without_crashing(tmp_path): + path = default_shard_path(tmp_path, rank=0, pid=10) + writer = TraceShardWriter( + path, run_id="run-nan", rank=0, pid=10, run_seed="seed" + ) + writer.submit(make_record("run-nan")) + writer.close() + lines = path.read_text(encoding="utf-8").splitlines() + event = json.loads(lines[1]) + event["payload"]["runtime"]["duration_us"] = float("nan") + lines[1] = json.dumps(event, allow_nan=True) + path.write_text("\n".join(lines) + "\n", encoding="utf-8") + + validated = validate_shard(path) + + assert validated.valid is False + assert any("non-finite JSON" in issue for issue in validated.issues) + + +def test_runtime_capture_records_triton_tensor_grid_meta_and_source(tmp_path): + source = tmp_path / "kernel.py" + source.write_text("def launch():\n pass\n", encoding="utf-8") + recorder = TargetedTraceRecorder( + tmp_path / "trace", + run_id="runtime-1", + run_seed="seed", + framework="vllm", + rank=2, + pid=22, + ) + assert recorder.record_triton_launch( + target_id="fused_moe", + kernel_name="fused_moe_kernel", + args=(FixtureTensor(),), + positional_names=("x",), + kwargs={"BLOCK_SIZE": 256, "num_warps": 8}, + constexpr_names=("BLOCK_SIZE",), + meta_names=("num_warps",), + grid=(32, 1, 1), + source_path=str(source), + source_line=1, + source_function="launch", + ) + receipt = recorder.close() + + records = [] + validated = validate_shard(Path(receipt.path), on_event=records.append) + assert validated.valid is True + record = records[0] + assert record.semantics.tensors[0].shape == (4, 8) + assert record.semantics.tensors[0].stride == (8, 1) + assert record.semantics.python_grid["items"] == [32, 1, 1] + assert record.semantics.constexpr == {"BLOCK_SIZE": 256} + assert record.semantics.meta == {"num_warps": 8} + assert record.semantics.source.sha256 + summary = postprocess_trace_dir(tmp_path / "trace") + assert summary["valid"] is True + assert summary["events"]["by_target"] == {"fused_moe": 1} + + +def test_runtime_recorders_merge_rank_shards_into_one_manifest(tmp_path): + output = tmp_path / "trace" + for rank, pid in ((0, 20), (1, 21)): + recorder = TargetedTraceRecorder( + output, + run_id="multi-rank", + run_seed="seed", + framework="vllm", + rank=rank, + pid=pid, + world_size=2, + ) + recorder.record_python_hip_launch( + target_id="aiter.hip_op", + kernel_name="hip_op", + args=(FixtureTensor(),), + ) + recorder.close() + + manifest = TargetedTraceManifest.from_dict( + json.loads((output / "manifest.json").read_text(encoding="utf-8")) + ) + assert len(manifest.shards) == 2 + assert manifest.coverage.written == 2 + assert {receipt.rank for receipt in manifest.shards} == {0, 1} + assert postprocess_trace_dir(output)["valid"] is True + + +def test_runtime_capture_records_python_hip_grid_meta_scalars_and_source(tmp_path): + source = tmp_path / "hip_wrapper.py" + source.write_text("def launch():\n pass\n", encoding="utf-8") + recorder = TargetedTraceRecorder( + tmp_path / "trace", + run_id="hip-runtime", + run_seed="seed", + framework="vllm", + rank=0, + pid=23, + ) + recorder.record_python_hip_launch( + target_id="aiter.hip_op", + kernel_name="hip_op", + args=(FixtureTensor(), 7), + positional_names=("x", "split_k"), + kwargs={"algorithm": 3}, + grid=(64, 1, 1), + meta_names=("algorithm",), + constexpr_names=("split_k",), + source_path=str(source), + source_line=1, + source_function="launch", + runtime=RuntimeEvidence( + gpu_symbol="hip_op_kernel", + grid=(64, 1, 1), + block=(256, 1, 1), + ), + ) + receipt = recorder.close() + + records = [] + assert validate_shard(Path(receipt.path), on_event=records.append).valid is True + record = records[0] + assert record.semantics.tensors[0].shape == (4, 8) + assert record.semantics.tensors[0].stride == (8, 1) + assert record.semantics.python_grid["items"] == [64, 1, 1] + assert record.semantics.constexpr == {"split_k": 7} + assert record.semantics.meta == {"algorithm": 3} + assert record.semantics.source.sha256 + assert record.runtime.grid == (64, 1, 1) + assert record.runtime.block == (256, 1, 1) + + +def test_torch_profiler_stream_adapter_and_manifest(tmp_path): + config = TargetedTraceConfig( + enabled=True, + run_seed="fixture-seed", + targets=[ + TargetSpec( + target_id="aiter.fused_moe", + name_patterns=("*fused_moe*",), + package="aiter", + source={"path": "aiter/moe.py", "line": 42}, + ) + ], + ) + output = tmp_path / "targeted" + manifest = adapt_torch_profiler_traces( + [FIXTURE], + output, + config=config, + run_id="fixture-run", + framework="vllm", + framework_version="0.19.1", + image="example/image:tag", + ) + + assert manifest.reward_eligible is False + assert manifest.pass_kind == "diagnostic" + assert manifest.coverage.written == 1 + records = [] + receipt = manifest.shards[0] + validated = validate_shard( + Path(receipt.path), expected_receipt=receipt, on_event=records.append + ) + assert validated.valid is True + record = records[0] + assert record.context.rank == 1 + assert record.context.stage == "decode" + assert record.context.execution_mode == "graph" + assert record.runtime.grid == (32, 2, 1) + assert record.runtime.block == (256, 1, 1) + assert record.runtime.stream == "7" + assert record.runtime.correlation_id == "99" + assert record.semantics.tensors[0].shape == (4, 128) + assert record.semantics.tensors[0].stride == (128, 1) + assert record.semantics.source.path == "aiter/moe.py" + + summary = postprocess_trace_dir(output, output_path=output / "summary.json") + assert summary["valid"] is True + assert summary["streaming"] is True + assert summary["events"]["by_target"] == {"aiter.fused_moe": 1} + + +def test_benchmark_adapter_materializes_bounded_valid_evidence(tmp_path): + targeted = TargetedTraceConfig( + enabled=True, + run_seed="benchmark-seed", + targets=[ + TargetSpec(target_id="moe", name_patterns=("*fused_moe*",)) + ], + ) + config = BenchmarkConfig( + framework="vllm", + model="model", + profiler=ProfilerConfig( + torch_profiler=TorchProfilerConfig(enabled=True), + targeted_trace=targeted, + ), + ) + + result = run_targeted_trace_analysis( + config=config, + trace_files=[FIXTURE], + workspace=tmp_path, + run_id="benchmark-run", + resolved_image="example/image@sha256:fixture", + ) + + assert result["valid"] is True + assert result["reward_eligible"] is False + assert result["coverage"] == { + "seen": 1, + "sampled": 1, + "written": 1, + "dropped": 0, + "dropped_by_reason": {}, + } + assert Path(result["manifest_path"]).is_file() + assert Path(result["summary_path"]).is_file() + + +def test_adapter_is_byte_stable_for_same_trace_seed_and_run_id(tmp_path): + config = TargetedTraceConfig( + enabled=True, + run_seed="same-seed", + targets=[TargetSpec(target_id="moe", name_patterns=("*fused_moe*",))], + ) + outputs = [] + for name in ("first", "second"): + output = tmp_path / name + manifest = adapt_torch_profiler_traces( + [FIXTURE], + output, + config=config, + run_id="same-run", + framework="vllm", + ) + outputs.append(Path(manifest.shards[0].path).read_bytes()) + + assert outputs[0] == outputs[1] + + +def test_adapter_surfaces_invalid_input_trace_as_acquisition_failure(tmp_path): + invalid = tmp_path / "not_a_trace.json" + invalid.write_text('{"metadata": {}}\n', encoding="utf-8") + config = TargetedTraceConfig( + enabled=True, + targets=[TargetSpec(target_id="moe", name_patterns=("*fused_moe*",))], + ) + output = tmp_path / "output" + adapt_torch_profiler_traces( + [FIXTURE, invalid], + output, + config=config, + run_id="invalid-input", + framework="vllm", + ) + + summary = postprocess_trace_dir(output) + assert summary["valid"] is False + assert summary["integrity_failures_by_reason"]["acquisition"] == 1 + + +def test_torch_trace_iterator_supports_gzip(tmp_path): + compressed = tmp_path / "rank_1_trace.json.gz" + with gzip.open(compressed, "wb") as stream: + stream.write(FIXTURE.read_bytes()) + events = list(iter_trace_events(compressed, chunk_size=37)) + assert [event["name"] for event in events] == [ + "cpu_only_event", + "triton_fused_moe_kernel", + "unselected_kernel", + ] + + +def test_adapter_ignores_capture_warmup_when_real_trace_exists(tmp_path): + capture_dir = tmp_path / "capture_traces" + real_dir = tmp_path / "rank_1" + capture_dir.mkdir() + real_dir.mkdir() + warmup = capture_dir / "warmup.json" + real = real_dir / "profile.json" + warmup.write_bytes(FIXTURE.read_bytes()) + real.write_bytes(FIXTURE.read_bytes()) + config = TargetedTraceConfig( + enabled=True, + targets=[TargetSpec(target_id="moe", name_patterns=("*fused_moe*",))], + ) + + manifest = adapt_torch_profiler_traces( + [warmup, real], + tmp_path / "output", + config=config, + run_id="no-warmup", + framework="vllm", + ) + + assert manifest.coverage.written == 1 + assert manifest.provenance["input_traces"] == [str(real)] + + +def test_postprocess_streams_shards_without_path_read_text(tmp_path, monkeypatch): + output = tmp_path / "targeted" + path = default_shard_path(output, rank=0, pid=10) + writer = TraceShardWriter( + path, run_id="run-stream", rank=0, pid=10, run_seed="seed" + ) + for index in range(1000): + writer.submit(make_record("run-stream", key=f"event-{index}")) + receipt = writer.close() + manifest = TargetedTraceManifest( + run_id="run-stream", + acquisition_backend="fixture", + targets=({"target_id": "target"},), + shards=(receipt,), + ) + write_manifest(output / "manifest.json", manifest) + + original = Path.read_text + + def guarded_read_text(self, *args, **kwargs): + if self.suffix == ".jsonl": + raise AssertionError("streaming postprocess must not read_text a shard") + return original(self, *args, **kwargs) + + monkeypatch.setattr(Path, "read_text", guarded_read_text) + summary = postprocess_trace_dir(output) + assert summary["valid"] is True + assert summary["coverage"]["written"] == 1000 + + +def test_postprocess_rejects_shard_not_bound_by_manifest(tmp_path): + output = tmp_path / "targeted" + first = TraceShardWriter( + default_shard_path(output, rank=0, pid=10), + run_id="bound-run", + rank=0, + pid=10, + run_seed="seed", + ) + first.submit(make_record("bound-run", pid=10)) + receipt = first.close() + extra = TraceShardWriter( + default_shard_path(output, rank=1, pid=11), + run_id="bound-run", + rank=1, + pid=11, + run_seed="seed", + ) + extra.submit(make_record("bound-run", rank=1, pid=11)) + extra.close() + write_manifest( + output / "manifest.json", + TargetedTraceManifest( + run_id="bound-run", + acquisition_backend="fixture", + targets=({"target_id": "target"},), + shards=(receipt,), + ), + ) + + summary = postprocess_trace_dir(output) + + assert summary["valid"] is False + assert any("undeclared trace shard" in issue for issue in summary["issues"]) + + +def test_unsupported_manifest_version_fails_fast(): + with pytest.raises(TraceValidationError, match="unsupported schema"): + TargetedTraceManifest.from_dict( + { + "schema_name": "magpie.targeted-kernel-trace", + "schema_version": "2.0.0", + } + ) + + +def test_measurement_and_diagnostic_lanes_are_explicit(): + measurement = BenchmarkConfig( + framework="vllm", + model="model", + run_kind="measurement", + profiler=ProfilerConfig( + torch_profiler=TorchProfilerConfig(enabled=False) + ), + ) + assert measurement.reward_eligible is True + + diagnostic = BenchmarkConfig( + framework="vllm", + model="model", + profiler=ProfilerConfig(), + ) + assert diagnostic.run_kind == "diagnostic" + assert diagnostic.reward_eligible is False + + with pytest.raises(ValueError, match="measurement.*requires"): + BenchmarkConfig( + framework="vllm", + model="model", + run_kind="measurement", + profiler=ProfilerConfig(), + ) + + +def test_targeted_trace_requires_torch_profiler(): + targeted = TargetedTraceConfig( + enabled=True, + targets=[TargetSpec(target_id="kernel", name_patterns=("*kernel*",))], + ) + with pytest.raises(ValueError, match="requires.*torch_profiler"): + BenchmarkConfig( + framework="vllm", + model="model", + profiler=ProfilerConfig( + torch_profiler=TorchProfilerConfig(enabled=False), + targeted_trace=targeted, + ), + ) From 36a843561dc30d45620bb9aeacee2d0a2b01f760 Mon Sep 17 00:00:00 2001 From: Vincent Ouyang Date: Fri, 7 Aug 2026 11:56:30 +0000 Subject: [PATCH 02/11] Lock Qwen benchmark runtime provenance --- Magpie/modes/benchmark/__init__.py | 3 +- Magpie/modes/benchmark/benchmarker.py | 306 ++++++++++- Magpie/modes/benchmark/config.py | 64 +++ Magpie/modes/benchmark/inferencex_runtime.py | 211 ++++++++ Magpie/modes/benchmark/lm_eval_runtime.py | 412 ++++++++++++++ Magpie/modes/benchmark/model_revision.py | 207 +++++++ Magpie/modes/benchmark/result.py | 44 ++ Magpie/scripts/benchmark/atom_mi300x.sh | 1 + Magpie/scripts/benchmark/atom_mi355x.sh | 1 + Magpie/scripts/benchmark/lm_eval_runtime.sh | 250 +++++++++ .../benchmark/magpie_bench_remote_compat.sh | 2 + Magpie/scripts/benchmark/sglang_mi300x.sh | 1 + Magpie/scripts/benchmark/sglang_mi355x.sh | 1 + Magpie/scripts/benchmark/vllm_mi300x.sh | 1 + Magpie/scripts/benchmark/vllm_mi355x.sh | 95 +++- Magpie/tools/amd_kernel_finder/indexer.py | 30 +- docs/how-to/benchmarking/benchmark.md | 57 +- docs/reference/benchmark-config.md | 41 ++ .../benchmark_vllm_qwen3_next_80b_fp8.yaml | 6 + tests/test_inferencex_runtime.py | 193 +++++++ tests/test_kernel_index_repo_detection.py | 100 ++++ tests/test_lm_eval_runtime.py | 510 ++++++++++++++++++ tests/test_model_revision_contract.py | 388 +++++++++++++ 23 files changed, 2882 insertions(+), 42 deletions(-) create mode 100644 Magpie/modes/benchmark/inferencex_runtime.py create mode 100644 Magpie/modes/benchmark/lm_eval_runtime.py create mode 100644 Magpie/modes/benchmark/model_revision.py create mode 100644 Magpie/scripts/benchmark/lm_eval_runtime.sh create mode 100644 tests/test_inferencex_runtime.py create mode 100644 tests/test_kernel_index_repo_detection.py create mode 100644 tests/test_lm_eval_runtime.py create mode 100644 tests/test_model_revision_contract.py diff --git a/Magpie/modes/benchmark/__init__.py b/Magpie/modes/benchmark/__init__.py index 28ca44a..a87d766 100644 --- a/Magpie/modes/benchmark/__init__.py +++ b/Magpie/modes/benchmark/__init__.py @@ -15,6 +15,7 @@ from .config import ( BenchmarkConfig, DEFAULT_SHARED_STORAGE_PATH, + LmEvalRuntimeConfig, ProfilerConfig, TorchProfilerConfig, SystemProfilerConfig, @@ -32,6 +33,7 @@ "BenchmarkMode", "BenchmarkConfig", "DEFAULT_SHARED_STORAGE_PATH", + "LmEvalRuntimeConfig", "BenchmarkResult", "ProfilerConfig", "TorchProfilerConfig", @@ -45,4 +47,3 @@ "InferenceXManager", "ensure_inferencex_available", ] - diff --git a/Magpie/modes/benchmark/benchmarker.py b/Magpie/modes/benchmark/benchmarker.py index 15e903e..4fb6659 100644 --- a/Magpie/modes/benchmark/benchmarker.py +++ b/Magpie/modes/benchmark/benchmarker.py @@ -32,6 +32,15 @@ from .config import BenchmarkConfig from .image_selector import ImageSelector from .inferencex import ensure_inferencex_available +from .inferencex_runtime import materialize_inferencex_runtime +from .lm_eval_runtime import ( + LmEvalRuntime, + collect_lm_eval_runtime_evidence, + invalid_runtime_evidence, + snapshot_runtime_manifest, + validate_lm_eval_runtime, +) +from .model_revision import collect_model_revision_evidence from .quality import parse_lm_eval_quality from .result import BenchmarkResult, LatencyMetrics, ResultParser, ThroughputMetrics from .targeted_trace import run_targeted_trace_analysis @@ -61,6 +70,10 @@ def _env_truthy(value: Any) -> bool: } ) +# These scripts source ``lm_eval_runtime.sh`` after InferenceX's mutable +# benchmark library. RUN_EVAL fails closed for every other script. +LM_EVAL_LOCKED_BUILTIN_SCRIPTS = MAGPIE_BUILTIN_SCRIPTS + class BenchmarkMode: """ @@ -95,6 +108,10 @@ def __init__( ) self._task_id: Optional[str] = None self._resolved_docker_image: Optional[str] = None + self._inferencex_source_path: Optional[str] = None + self._inferencex_runtime_receipt: Optional[Dict[str, Any]] = None + self._lm_eval_runtime: Optional[LmEvalRuntime] = None + self._lm_eval_runtime_evidence: Optional[Dict[str, Any]] = None def run(self, task_id: Optional[str] = None) -> BenchmarkResult: """ @@ -113,13 +130,37 @@ def run(self, task_id: Optional[str] = None) -> BenchmarkResult: # Ray mode: delegate to remote cluster and return when complete if self.config.is_ray: + if self._lm_eval_requested(): + result = BenchmarkResult(success=False) + result.lm_eval_runtime_receipt = invalid_runtime_evidence( + self.config.lm_eval_runtime, + "RUN_EVAL=true is unsupported in Ray mode because the " + "locked evaluator runtime cannot yet be mounted and attested", + status="unsupported", + ) + result.errors.append( + "Ray benchmark refused: RUN_EVAL=true requires the local or " + "Docker locked evaluator path" + ) + return result return self._execute_ray_benchmark() - # 0. Ensure InferenceX is available (auto-clone if needed) + self._inferencex_runtime_receipt = None + self._lm_eval_runtime = None + self._lm_eval_runtime_evidence = None + + # 0. Resolve the caller's InferenceX source checkout. On repeated runs + # of one BenchmarkMode instance, never treat the prior run's disposable + # runtime tree as the new source. try: - self.config.inferencex_path = ensure_inferencex_available( - self.config.inferencex_path + configured_source = ( + self._inferencex_source_path or self.config.inferencex_path + ) + source_path = ensure_inferencex_available( + configured_source ) + self._inferencex_source_path = str(Path(source_path).resolve()) + self.config.inferencex_path = self._inferencex_source_path except RuntimeError as e: result = BenchmarkResult() result.success = False @@ -129,8 +170,69 @@ def run(self, task_id: Optional[str] = None) -> BenchmarkResult: "git clone https://github.com/SemiAnalysisAI/InferenceX.git" ) return result - - # 0b. Optionally pick idle GPU(s) and pin VISIBLE_DEVICES. + + # 1. Create the result workspace, then export the exact InferenceX HEAD + # into a private runtime tree. Magpie and TraceLens may modify only this + # tree; the dependency/source checkout is never a write target. + workspace = self.workspace_mgr.create(self.config.to_dict()) + try: + inferencex_runtime = materialize_inferencex_runtime( + Path(self._inferencex_source_path), + workspace, + ) + except (OSError, RuntimeError) as e: + return self._workspace_failure( + workspace, + start_time, + f"Failed to materialize InferenceX runtime: {e}", + ) + self.config.inferencex_path = str(inferencex_runtime.root) + self._inferencex_runtime_receipt = dict(inferencex_runtime.receipt) + + # A requested quality evaluation must use the exact caller-built + # runtime. Validate all bytes and preserve its manifest before any + # benchmark process or container is started. + runtime_config = self.config.lm_eval_runtime + eval_requested = self._lm_eval_requested() + if eval_requested and runtime_config is None: + self._lm_eval_runtime_evidence = invalid_runtime_evidence( + None, + "RUN_EVAL=true requires benchmark.lm_eval_runtime", + ) + return self._workspace_failure( + workspace, + start_time, + "lm-eval runtime preflight failed: RUN_EVAL=true requires " + "benchmark.lm_eval_runtime", + ) + if runtime_config is not None: + try: + self._lm_eval_runtime = validate_lm_eval_runtime(runtime_config) + runtime_identity = self._lm_eval_runtime.identity + if ( + runtime_identity.get("inferencex_commit") + != inferencex_runtime.receipt.get("source_commit") + or runtime_identity.get("inferencex_tree") + != inferencex_runtime.receipt.get("source_tree") + ): + raise ValueError( + "lm-eval runtime InferenceX commit/tree does not match " + "the materialized benchmark checkout" + ) + snapshot_runtime_manifest(self._lm_eval_runtime, workspace) + except (OSError, TypeError, ValueError) as e: + self._lm_eval_runtime_evidence = invalid_runtime_evidence( + runtime_config, + str(e), + ) + return self._workspace_failure( + workspace, + start_time, + f"lm-eval runtime preflight failed: {e}", + ) + self.workspace_mgr._save_config_snapshot(workspace, self.config.to_dict()) + + # 2. Optionally pick idle GPU(s) and pin VISIBLE_DEVICES. # When server_lifecycle will reuse an existing HTTP server, skip # find_idle_gpus so ROCR/HIP/CUDA_VISIBLE_* are not reshuffled versus # the already-running server's physical devices. @@ -149,29 +251,19 @@ def run(self, task_id: Optional[str] = None) -> BenchmarkResult: else: self._apply_gpu_selection() except RuntimeError as e: - result = BenchmarkResult() - result.success = False - result.errors.append(str(e)) - return result + return self._workspace_failure(workspace, start_time, str(e)) + self.workspace_mgr._save_config_snapshot(workspace, self.config.to_dict()) - # 1. Copy Magpie generic scripts to InferenceX/benchmarks/ - self._prepare_benchmark_scripts() - - # 2. Determine runner type from GPU - runner_type = self._get_runner_type() - - # 3. Find and validate benchmark script BEFORE container starts + # 3. Add Magpie scripts to the private runtime, then select the runner + # and benchmark entrypoint before any container/process launch. try: + self._prepare_benchmark_scripts() + runner_type = self._get_runner_type() benchmark_script = self._get_benchmark_script(runner_type) + self._validate_lm_eval_benchmark_script(benchmark_script) logger.info(f"Selected benchmark script: {benchmark_script}") - except FileNotFoundError as e: - result = BenchmarkResult() - result.success = False - result.errors.append(str(e)) - return result - - # 4. Create workspace - workspace = self.workspace_mgr.create(self.config.to_dict()) + except (OSError, RuntimeError, ValueError) as e: + return self._workspace_failure(workspace, start_time, str(e)) # 4a. For Docker benchmarks, TraceLens inference can derive a patched # framework runtime image from supported official vLLM/SGLang images. @@ -209,6 +301,10 @@ def run(self, task_id: Optional[str] = None) -> BenchmarkResult: profiling_enabled=self.config.profiler.torch_profiler.enabled, run_kind=self.config.run_kind, reward_eligible=self.config.reward_eligible, + inferencex_runtime_receipt=self._inferencex_runtime_receipt, + lm_eval_runtime_receipt=self._collect_lm_eval_evidence( + workspace + ), ) result.errors.append(f"TraceLens runtime image setup failed: {e}") result.tracelens_analysis = { @@ -251,6 +347,10 @@ def run(self, task_id: Optional[str] = None) -> BenchmarkResult: profiling_enabled=self.config.profiler.torch_profiler.enabled, run_kind=self.config.run_kind, reward_eligible=self.config.reward_eligible, + inferencex_runtime_receipt=self._inferencex_runtime_receipt, + lm_eval_runtime_receipt=self._collect_lm_eval_evidence( + workspace + ), ) result.errors.append(f"TraceLens preprocess failed: {e}") result.tracelens_analysis = { @@ -338,6 +438,28 @@ def run(self, task_id: Optional[str] = None) -> BenchmarkResult: result.profiling_enabled = self.config.profiler.torch_profiler.enabled result.run_kind = self.config.run_kind result.reward_eligible = self.config.reward_eligible + result.inferencex_runtime_receipt = self._inferencex_runtime_receipt + result.lm_eval_runtime_receipt = self._collect_lm_eval_evidence(workspace) + runtime_evidence = result.lm_eval_runtime_receipt + if runtime_evidence["requested"] and not runtime_evidence["verified"]: + result.success = False + result.errors.append( + "lm-eval runtime evidence gate failed: " + f"{runtime_evidence.get('errors', [])}" + ) + + result.model_revision_receipt = collect_model_revision_evidence( + workspace, + model=self.config.model, + requested_revision=self.config.envs.get("MODEL_REVISION"), + ) + revision_evidence = result.model_revision_receipt + if revision_evidence["requested"] and not revision_evidence["verified"]: + result.success = False + result.errors.append( + "Model revision evidence gate failed: " + f"{revision_evidence.get('errors', [])}" + ) # Add GPU monitor stats if gpu_monitor_stats is not None: @@ -532,6 +654,60 @@ def run(self, task_id: Optional[str] = None) -> BenchmarkResult: logger.info(f"Benchmark completed in {result.execution_time:.2f}s") return result + + def _lm_eval_requested(self) -> bool: + return _env_truthy(self.config.envs.get("RUN_EVAL", False)) + + def _validate_lm_eval_benchmark_script(self, benchmark_script: str) -> None: + """Reject evaluator paths that do not activate the locked runtime.""" + + if not self._lm_eval_requested(): + return + path = Path(benchmark_script) + if ( + path.parent.as_posix() != "benchmarks" + or path.name not in LM_EVAL_LOCKED_BUILTIN_SCRIPTS + ): + raise RuntimeError( + "RUN_EVAL=true requires one of Magpie's locked evaluator " + "benchmark scripts; native and custom InferenceX scripts are " + "unsupported" + ) + + def _collect_lm_eval_evidence(self, workspace: Path) -> Dict[str, Any]: + if self._lm_eval_runtime_evidence is None: + self._lm_eval_runtime_evidence = collect_lm_eval_runtime_evidence( + workspace, + requested=self._lm_eval_requested(), + config=self.config.lm_eval_runtime, + execution_mode=self.config.run_mode, + ) + return dict(self._lm_eval_runtime_evidence) + + def _workspace_failure( + self, + workspace: Path, + start_time: float, + message: str, + ) -> BenchmarkResult: + """Persist a failure that occurs after the run workspace exists.""" + + result = BenchmarkResult( + success=False, + framework=self.config.framework, + model=self.config.model, + workspace_dir=str(workspace), + execution_time=time.time() - start_time, + profiling_enabled=self.config.profiler.torch_profiler.enabled, + run_kind=self.config.run_kind, + reward_eligible=self.config.reward_eligible, + inferencex_runtime_receipt=self._inferencex_runtime_receipt, + lm_eval_runtime_receipt=self._collect_lm_eval_evidence(workspace), + ) + result.errors.append(message) + self.workspace_mgr.save_report(result.to_dict()) + self.workspace_mgr.save_summary(result.get_summary()) + return result def _apply_gpu_selection(self) -> None: """Resolve idle GPU(s) and inject the selection into config.envs. @@ -623,16 +799,27 @@ def _get_runner_type(self) -> str: def _prepare_benchmark_scripts(self) -> None: """ - Copy Magpie generic benchmark scripts to InferenceX/benchmarks/. + Copy Magpie generic scripts to the run-scoped InferenceX runtime. This allows using Magpie's generic scripts while still leveraging - InferenceX's benchmark_lib.sh and other utilities. + InferenceX's benchmark_lib.sh and other utilities. ``run()`` first + materializes the exact source commit into the benchmark workspace, so + this method must never target the caller's source checkout. Always overwrites to keep scripts in sync with Magpie source. """ # Magpie scripts location: Magpie/scripts/benchmark/ magpie_scripts = Path(__file__).parent.parent.parent / "scripts" / "benchmark" - target_dir = Path(self.config.inferencex_path) / "benchmarks" + target_root = Path(self.config.inferencex_path).resolve() + if ( + self._inferencex_source_path is not None + and target_root == Path(self._inferencex_source_path).resolve() + ): + raise RuntimeError( + "refusing to install Magpie scripts into the InferenceX source " + "checkout" + ) + target_dir = target_root / "benchmarks" if not magpie_scripts.exists(): logger.debug(f"Magpie scripts directory not found: {magpie_scripts}") @@ -745,6 +932,17 @@ def _build_docker_command( if os.path.exists(inferencex_path): cmd.extend(["-v", f"{inferencex_path}:/opt/InferenceX"]) + # The evaluator is caller-built and content-addressed. Mount the whole + # root read-only so the in-container helper can independently validate + # both its manifest and site-packages tree. + if self._lm_eval_runtime is not None: + cmd.extend( + [ + "-v", + f"{self._lm_eval_runtime.root}:/opt/apex/lm-eval-runtime:ro", + ] + ) + # Model directory mount — if the model path is a local directory, mount it # so the container can access the weights (e.g. /mnt/dcgpuval/datasets/...) model_path = self.config.model @@ -760,6 +958,20 @@ def _build_docker_command( env_vars["RESULT_FILENAME"] = "inferencex_result" env_vars["RESULT_DIR"] = "/workspace" env_vars["RUNNER_TYPE"] = runner_type + if self._lm_eval_runtime is not None: + env_vars.update( + { + "MAGPIE_LM_EVAL_RUNTIME_ROOT": "/opt/apex/lm-eval-runtime", + "MAGPIE_LM_EVAL_RUNTIME_SHA256": ( + self._lm_eval_runtime.runtime_sha256 + ), + "MAGPIE_LM_EVAL_RUNTIME_RECEIPT": ( + "/workspace/lm_eval_runtime_receipt.json" + ), + "MAGPIE_LM_EVAL_EXECUTION_MODE": "docker", + "MAGPIE_LM_EVAL_REQUIRE_READONLY_MOUNT": "1", + } + ) # torch_profiler environment (matches official InferenceX: PROFILE=1) if self.config.profiler.torch_profiler.enabled: @@ -858,6 +1070,22 @@ def _build_local_command( env_vars["RESULT_DIR"] = str(workspace) env_vars["RUNNER_TYPE"] = runner_type env_vars["MAGPIE_RUN_PHASE"] = phase + if self._lm_eval_runtime is not None: + env_vars.update( + { + "MAGPIE_LM_EVAL_RUNTIME_ROOT": str( + self._lm_eval_runtime.root + ), + "MAGPIE_LM_EVAL_RUNTIME_SHA256": ( + self._lm_eval_runtime.runtime_sha256 + ), + "MAGPIE_LM_EVAL_RUNTIME_RECEIPT": str( + workspace / "lm_eval_runtime_receipt.json" + ), + "MAGPIE_LM_EVAL_EXECUTION_MODE": "local", + "MAGPIE_LM_EVAL_REQUIRE_READONLY_MOUNT": "0", + } + ) if phase == "server" and server_pid_file is not None: env_vars["MAGPIE_SERVER_PID_FILE"] = str(server_pid_file) @@ -1188,7 +1416,13 @@ def _desired_reuse_server_meta(self, port: int) -> Dict[str, Any]: upper = { str(k).upper(): str(v) for k, v in (self.config.envs or {}).items() } - ix_path = str(Path(self.config.inferencex_path).resolve()) + identity_path = self._inferencex_source_path or self.config.inferencex_path + ix_path = str(Path(identity_path).resolve()) + ix_commit = "" + if self._inferencex_runtime_receipt is not None: + ix_commit = str( + self._inferencex_runtime_receipt.get("source_commit") or "" + ) fw = self.config.framework.lower() @@ -1217,6 +1451,13 @@ def _desired_reuse_server_meta(self, port: int) -> Dict[str, Any]: "extra_atom_args": extras_atom, "max_model_len": str(max_ml), "inferencex_path": ix_path, + "inferencex_commit": ix_commit, + "model_revision": upper.get("MODEL_REVISION", ""), + "lm_eval_runtime_sha256": ( + self.config.lm_eval_runtime.sha256 + if self.config.lm_eval_runtime is not None + else "" + ), } def _reuse_meta_mismatch( @@ -1240,6 +1481,9 @@ def _reuse_meta_mismatch( ("extra_atom_args", "extra_atom_args"), ("max_model_len", "max_model_len"), ("inferencex_path", "inferencex_path"), + ("inferencex_commit", "inferencex_commit"), + ("model_revision", "model_revision"), + ("lm_eval_runtime_sha256", "lm_eval_runtime_sha256"), ) diffs = [] @@ -1875,6 +2119,12 @@ def _build_ray_benchmark_task(self) -> Tuple[Optional["Task"], Optional[str]]: """Build the ``Task`` for a Ray benchmark; sets ``self._task_id`` if unset.""" from ...core.task import ModeConfig, ModeType, Task + if self._lm_eval_requested(): + return ( + None, + "RUN_EVAL=true is unsupported in Ray mode because the locked " + "evaluator runtime cannot yet be mounted and attested", + ) if self.config.ray_config is None: return None, "ray_config is required when run_mode='ray'" diff --git a/Magpie/modes/benchmark/config.py b/Magpie/modes/benchmark/config.py index e990b40..70382dc 100644 --- a/Magpie/modes/benchmark/config.py +++ b/Magpie/modes/benchmark/config.py @@ -14,6 +14,46 @@ from ...targeted_trace.config import TargetedTraceConfig +@dataclass(frozen=True) +class LmEvalRuntimeConfig: + """Hash-locked, read-only evaluator runtime supplied by the caller. + + ``path`` names the host runtime root, ``sha256`` commits to its canonical + identity/file manifest, and ``identity`` is compared exactly with the + signed manifest. Magpie consumes this runtime; it never constructs or + updates it. + """ + + path: str + sha256: str + identity: Dict[str, Any] + + def __post_init__(self) -> None: + if not isinstance(self.path, str) or not self.path.strip(): + raise ValueError("lm_eval_runtime.path must be a non-empty string") + if not isinstance(self.sha256, str) or not self.sha256.strip(): + raise ValueError("lm_eval_runtime.sha256 must be a non-empty string") + if not isinstance(self.identity, dict) or not self.identity: + raise ValueError("lm_eval_runtime.identity must be a non-empty mapping") + + def to_dict(self) -> Dict[str, Any]: + return { + "path": self.path, + "sha256": self.sha256, + "identity": dict(self.identity), + } + + @classmethod + def from_dict(cls, data: Dict[str, Any]) -> "LmEvalRuntimeConfig": + if not isinstance(data, dict): + raise ValueError("lm_eval_runtime must be a mapping") + return cls( + path=data.get("path", ""), + sha256=data.get("sha256", ""), + identity=data.get("identity", {}), + ) + + class BenchmarkFramework(Enum): """Supported benchmark frameworks.""" @@ -686,6 +726,7 @@ class BenchmarkConfig: timeout_seconds: Benchmark timeout inferencex_path: Path to InferenceX installation hf_cache_path: HuggingFace cache directory + lm_eval_runtime: Optional caller-built, hash-locked evaluator runtime runner_type: Hardware runner type for InferenceX (e.g., "mi300x", "h100") server_lifecycle: Optional persisted-server settings (local-only) """ @@ -717,6 +758,7 @@ class BenchmarkConfig: # 3) ~/.cache/magpie/InferenceX inferencex_path: str = "" hf_cache_path: Optional[str] = None + lm_eval_runtime: Optional[LmEvalRuntimeConfig] = None # Gap analysis gap_analysis: GapAnalysisConfig = field(default_factory=GapAnalysisConfig) @@ -827,6 +869,18 @@ def __post_init__(self): self.server_lifecycle ) + if isinstance(self.lm_eval_runtime, dict): + self.lm_eval_runtime = LmEvalRuntimeConfig.from_dict( + self.lm_eval_runtime + ) + elif self.lm_eval_runtime is not None and not isinstance( + self.lm_eval_runtime, + LmEvalRuntimeConfig, + ): + raise ValueError( + "lm_eval_runtime must be a mapping or LmEvalRuntimeConfig" + ) + if self.is_server_lifecycle: if self.run_mode != "local": raise ValueError( @@ -928,6 +982,11 @@ def to_dict(self) -> Dict[str, Any]: "timeout_seconds": self.timeout_seconds, "inferencex_path": self.inferencex_path, "hf_cache_path": self.hf_cache_path, + "lm_eval_runtime": ( + self.lm_eval_runtime.to_dict() + if self.lm_eval_runtime is not None + else None + ), "runner_type": self.runner_type, "benchmark_script": self.benchmark_script, "gpu_selection": self.gpu_selection.to_dict(), @@ -990,6 +1049,11 @@ def from_dict(cls, data: Dict[str, Any]) -> "BenchmarkConfig": or "" ), hf_cache_path=data.get("hf_cache_path"), + lm_eval_runtime=( + LmEvalRuntimeConfig.from_dict(data["lm_eval_runtime"]) + if data.get("lm_eval_runtime") is not None + else None + ), runner_type=data.get("runner_type"), benchmark_script=data.get("benchmark_script"), gpu_selection=GpuSelectionConfig.from_dict(data.get("gpu_selection") or {}), diff --git a/Magpie/modes/benchmark/inferencex_runtime.py b/Magpie/modes/benchmark/inferencex_runtime.py new file mode 100644 index 0000000..3fe1c9f --- /dev/null +++ b/Magpie/modes/benchmark/inferencex_runtime.py @@ -0,0 +1,211 @@ +"""Create run-scoped InferenceX trees without mutating the source checkout.""" + +from __future__ import annotations + +import hashlib +import json +import os +import re +import shutil +import subprocess +import uuid +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Dict, Optional, Sequence + + +INFERENCEX_RUNTIME_RECEIPT_FILENAME = "inferencex_runtime_receipt.json" +INFERENCEX_RUNTIME_RECEIPT_SCHEMA = "magpie.inferencex-runtime-receipt/v1" +_COMMIT_RE = re.compile(r"^[0-9a-f]{40}$") + + +@dataclass(frozen=True) +class InferenceXRuntime: + """A run-scoped InferenceX tree and its source/materialization receipt.""" + + source_root: Path + root: Path + receipt: Dict[str, Any] + + +def _run_git( + source_root: Path, + args: Sequence[str], + *, + env: Optional[Dict[str, str]] = None, +) -> str: + command = ["git", "-C", str(source_root), *args] + try: + completed = subprocess.run( + command, + capture_output=True, + text=True, + timeout=120, + env=env, + check=False, + ) + except (OSError, subprocess.TimeoutExpired) as exc: + raise RuntimeError(f"could not run {' '.join(command[:3])}: {exc}") from exc + if completed.returncode != 0: + detail = (completed.stderr or completed.stdout or "git command failed").strip() + raise RuntimeError(f"git {' '.join(args)} failed: {detail[-1000:]}") + return completed.stdout.strip() + + +def _git_root(source_root: Path) -> Optional[Path]: + try: + top_level = _run_git(source_root, ["rev-parse", "--show-toplevel"]) + except RuntimeError: + return None + resolved = Path(top_level).resolve() + return resolved if resolved == source_root else None + + +def _git_status(source_root: Path) -> str: + return _run_git( + source_root, + ["status", "--porcelain=v1", "--untracked-files=all"], + ) + + +def _status_sha256(status: str) -> str: + return hashlib.sha256(status.encode("utf-8")).hexdigest() + + +def _checkout_commit( + source_root: Path, + runtime_root: Path, + workspace: Path, + commit: str, +) -> None: + """Export ``commit`` through a private temporary Git index. + + This reads blobs from the source object database but never changes its + working tree, index, refs, or worktree metadata. Unlike copying the source + filesystem, staged and unstaged changes cannot leak into the runtime tree. + """ + + index_path = workspace / f".inferencex-index-{uuid.uuid4().hex}" + index_lock = Path(f"{index_path}.lock") + runtime_root.mkdir(parents=True) + git_env = os.environ.copy() + git_env["GIT_INDEX_FILE"] = str(index_path) + try: + _run_git(source_root, ["read-tree", commit], env=git_env) + prefix = f"{runtime_root}{os.sep}" + _run_git( + source_root, + ["checkout-index", "--all", f"--prefix={prefix}"], + env=git_env, + ) + finally: + for temporary in (index_path, index_lock): + try: + temporary.unlink() + except FileNotFoundError: + pass + + +def _write_receipt(workspace: Path, payload: Dict[str, Any]) -> None: + receipt_path = workspace / INFERENCEX_RUNTIME_RECEIPT_FILENAME + temporary = workspace / f".{receipt_path.name}.{uuid.uuid4().hex}.tmp" + try: + with temporary.open("x", encoding="utf-8") as stream: + json.dump(payload, stream, indent=2, sort_keys=True) + stream.write("\n") + stream.flush() + os.fsync(stream.fileno()) + os.replace(temporary, receipt_path) + finally: + try: + temporary.unlink() + except FileNotFoundError: + pass + + +def materialize_inferencex_runtime( + source_root: Path, + workspace: Path, +) -> InferenceXRuntime: + """Materialize a private InferenceX tree under a benchmark workspace. + + Git repositories are exported from the exact ``HEAD`` commit through a + private index. Non-Git directories retain a compatibility path using a + filesystem copy, clearly marked as unpinned in the receipt. + """ + + source_root = Path(source_root).resolve() + workspace = Path(workspace).resolve() + runtime_root = workspace / "inferencex_runtime" + if not source_root.is_dir(): + raise RuntimeError(f"InferenceX source is not a directory: {source_root}") + if not (source_root / "benchmarks").is_dir(): + raise RuntimeError( + f"InferenceX source has no benchmarks directory: {source_root}" + ) + if workspace == source_root or source_root in workspace.parents: + raise RuntimeError("benchmark workspace must be outside the InferenceX source") + if runtime_root.exists(): + raise RuntimeError(f"InferenceX runtime already exists: {runtime_root}") + + git_root = _git_root(source_root) + if git_root is None and (source_root / ".git").exists(): + raise RuntimeError( + "InferenceX has Git metadata but its repository root/HEAD could not " + "be resolved" + ) + if git_root is not None: + commit = _run_git(source_root, ["rev-parse", "--verify", "HEAD^{commit}"]) + if not _COMMIT_RE.fullmatch(commit): + raise RuntimeError(f"InferenceX HEAD is not an exact commit: {commit!r}") + tree = _run_git(source_root, ["rev-parse", "--verify", "HEAD^{tree}"]) + if not _COMMIT_RE.fullmatch(tree): + raise RuntimeError(f"InferenceX HEAD tree is not exact: {tree!r}") + status_before = _git_status(source_root) + _checkout_commit(source_root, runtime_root, workspace, commit) + status_after = _git_status(source_root) + if status_after != status_before: + raise RuntimeError( + "InferenceX source status changed while materializing runtime" + ) + method = "git_private_index_checkout" + source_is_git = True + status_digest = _status_sha256(status_before) + source_clean = not status_before + else: + shutil.copytree( + source_root, + runtime_root, + symlinks=True, + ignore=shutil.ignore_patterns(".git"), + ) + commit = None + tree = None + method = "filesystem_copy" + source_is_git = False + status_digest = None + source_clean = None + + if not (runtime_root / "benchmarks").is_dir(): + raise RuntimeError( + "materialized InferenceX runtime has no benchmarks directory" + ) + + receipt: Dict[str, Any] = { + "schema": INFERENCEX_RUNTIME_RECEIPT_SCHEMA, + "source_root": str(source_root), + "source_is_git": source_is_git, + "source_commit": commit, + "source_tree": tree, + "source_clean": source_clean, + "source_status_sha256": status_digest, + "source_status_unchanged": True, + "runtime_path": runtime_root.name, + "materialization_method": method, + } + _write_receipt(workspace, receipt) + return InferenceXRuntime( + source_root=source_root, + root=runtime_root, + receipt=receipt, + ) diff --git a/Magpie/modes/benchmark/lm_eval_runtime.py b/Magpie/modes/benchmark/lm_eval_runtime.py new file mode 100644 index 0000000..a606e40 --- /dev/null +++ b/Magpie/modes/benchmark/lm_eval_runtime.py @@ -0,0 +1,412 @@ +"""Validate and attest caller-supplied, hash-locked lm-eval runtimes.""" + +from __future__ import annotations + +import hashlib +import json +import os +import re +import stat +import uuid +from dataclasses import dataclass +from pathlib import Path, PurePosixPath +from typing import Any, Dict, List, Mapping, Sequence, Tuple + +from .config import LmEvalRuntimeConfig + + +LM_EVAL_MANIFEST_FILENAME = "lm_eval_runtime_manifest.json" +LM_EVAL_MANIFEST_SCHEMA = "apex.lm-eval-runtime/v1" +LM_EVAL_RECEIPT_FILENAME = "lm_eval_runtime_receipt.json" +LM_EVAL_RECEIPT_SCHEMA = "magpie.lm-eval-runtime-receipt/v1" +LM_EVAL_EVIDENCE_SCHEMA = "magpie.lm-eval-runtime-evidence/v1" +MAX_MANIFEST_SIZE_BYTES = 64 * 1024 * 1024 +MAX_RECEIPT_SIZE_BYTES = 256 * 1024 + +_HEX40 = re.compile(r"^[0-9a-f]{40}$") +_SHA256 = re.compile(r"^[0-9a-f]{64}$") +_IMAGE_SHA256 = re.compile(r"^sha256:[0-9a-f]{64}$") +_REPO_DIGEST = re.compile(r"^(?:[^\s@]+@)?sha256:[0-9a-f]{64}$") +_PYTHON_ABI = re.compile(r"^[a-z0-9_]+-[a-z0-9_]+$") +_MANIFEST_KEYS = frozenset( + {"schema", "runtime_sha256", "site_packages", "identity", "files"} +) +_FILE_KEYS = frozenset({"path", "size_bytes", "mode", "sha256"}) +_IDENTITY_KEYS = frozenset( + { + "lm_eval_commit", + "lm_eval_tree", + "lm_eval_version", + "python_abi", + "base_image_id", + "base_image_repo_digest", + "inferencex_commit", + "inferencex_tree", + } +) +_RECEIPT_KEYS = frozenset( + { + "schema", + "runtime_sha256", + "identity", + "manifest_sha256", + "site_packages", + "python_abi", + "lm_eval_version", + "lm_eval_module", + "execution_mode", + "read_only_mount", + "verified", + } +) + + +@dataclass(frozen=True) +class LmEvalRuntime: + """A fully verified host runtime ready for local or read-only Docker use.""" + + root: Path + site_packages: Path + runtime_sha256: str + identity: Dict[str, Any] + manifest_bytes: bytes + manifest_sha256: str + + +def _reject_json_constant(value: str) -> None: + raise ValueError(f"non-finite JSON number is forbidden: {value}") + + +def _reject_duplicate_keys(pairs: List[Tuple[str, Any]]) -> Dict[str, Any]: + value: Dict[str, Any] = {} + for key, item in pairs: + if key in value: + raise ValueError(f"duplicate JSON object key: {key}") + value[key] = item + return value + + +def _read_json_file( + path: Path, + *, + limit: int, + require_readonly: bool = False, +) -> tuple[Mapping[str, Any], bytes]: + info = path.lstat() + if not stat.S_ISREG(info.st_mode) or info.st_nlink != 1: + raise ValueError(f"{path.name} must be a regular file with nlink=1") + if require_readonly and info.st_mode & 0o222: + raise ValueError(f"{path.name} must not have writable permission bits") + if info.st_size <= 0 or info.st_size > limit: + raise ValueError(f"{path.name} has invalid size {info.st_size}") + raw = path.read_bytes() + payload = json.loads( + raw.decode("utf-8"), + parse_constant=_reject_json_constant, + object_pairs_hook=_reject_duplicate_keys, + ) + if not isinstance(payload, Mapping): + raise ValueError(f"{path.name} root must be a JSON object") + return payload, raw + + +def _sha256_bytes(value: bytes) -> str: + return hashlib.sha256(value).hexdigest() + + +def _sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _canonical_sha256(identity: Mapping[str, Any], files: Sequence[Any]) -> str: + encoded = json.dumps( + {"identity": identity, "files": files}, + ensure_ascii=False, + separators=(",", ":"), + sort_keys=True, + ).encode("utf-8") + return _sha256_bytes(encoded) + + +def _validate_identity(identity: Any) -> Dict[str, Any]: + if not isinstance(identity, Mapping): + raise ValueError("manifest identity must be a JSON object") + missing = sorted(_IDENTITY_KEYS - set(identity)) + if missing: + raise ValueError(f"manifest identity is missing required keys: {missing}") + result = dict(identity) + for key in ("lm_eval_commit", "lm_eval_tree", "inferencex_commit", "inferencex_tree"): + if not isinstance(result[key], str) or not _HEX40.fullmatch(result[key]): + raise ValueError(f"identity.{key} must be an exact lowercase 40-hex id") + version = result["lm_eval_version"] + if not isinstance(version, str) or not version.strip(): + raise ValueError("identity.lm_eval_version must be a non-empty string") + abi = result["python_abi"] + if not isinstance(abi, str) or not _PYTHON_ABI.fullmatch(abi): + raise ValueError("identity.python_abi must be sys.implementation.cache_tag") + if not isinstance(result["base_image_id"], str) or not _IMAGE_SHA256.fullmatch( + result["base_image_id"] + ): + raise ValueError("identity.base_image_id must be an immutable image ID") + repo_digest = result["base_image_repo_digest"] + if not isinstance(repo_digest, str) or not _REPO_DIGEST.fullmatch(repo_digest): + raise ValueError("identity.base_image_repo_digest must contain a sha256 digest") + return result + + +def _validate_relative_path(value: Any) -> str: + if not isinstance(value, str) or not value: + raise ValueError("file record path must be a non-empty string") + path = PurePosixPath(value) + if path.is_absolute() or value != path.as_posix() or ".." in path.parts: + raise ValueError(f"file record path is not canonical and relative: {value!r}") + return value + + +def _validate_manifest_records(value: Any) -> List[Dict[str, Any]]: + if not isinstance(value, list) or not value: + raise ValueError("manifest files must be a non-empty list") + records: List[Dict[str, Any]] = [] + paths: List[str] = [] + for item in value: + if not isinstance(item, Mapping) or set(item) != _FILE_KEYS: + raise ValueError("each file record must have path/size_bytes/mode/sha256") + path = _validate_relative_path(item.get("path")) + size = item.get("size_bytes") + mode = item.get("mode") + digest = item.get("sha256") + if not isinstance(size, int) or isinstance(size, bool) or size < 0: + raise ValueError(f"invalid size_bytes for {path}") + if not isinstance(mode, int) or isinstance(mode, bool) or not 0 <= mode <= 0o7777: + raise ValueError(f"invalid mode for {path}") + if mode & 0o222: + raise ValueError(f"file record is writable: {path}") + if not isinstance(digest, str) or not _SHA256.fullmatch(digest): + raise ValueError(f"invalid sha256 for {path}") + paths.append(path) + records.append(dict(item)) + if paths != sorted(paths) or len(paths) != len(set(paths)): + raise ValueError("manifest file records must be unique and path-sorted") + return records + + +def _verify_site_packages(site_packages: Path, records: Sequence[Mapping[str, Any]]) -> None: + actual_files: List[str] = [] + for path in sorted(site_packages.rglob("*"), key=lambda item: item.as_posix()): + info = path.lstat() + relative = path.relative_to(site_packages).as_posix() + if stat.S_ISLNK(info.st_mode): + raise ValueError(f"runtime symlink is forbidden: {relative}") + if stat.S_ISDIR(info.st_mode): + if info.st_mode & 0o222: + raise ValueError(f"runtime directory is writable: {relative}") + continue + if not stat.S_ISREG(info.st_mode) or info.st_nlink != 1: + raise ValueError(f"runtime entry must be a regular nlink=1 file: {relative}") + actual_files.append(relative) + expected_files = [str(item["path"]) for item in records] + if actual_files != expected_files: + raise ValueError("site-packages files do not exactly match the manifest") + for item in records: + path = site_packages / str(item["path"]) + info = path.lstat() + if info.st_size != item["size_bytes"]: + raise ValueError(f"size mismatch for {item['path']}") + if stat.S_IMODE(info.st_mode) != item["mode"]: + raise ValueError(f"mode mismatch for {item['path']}") + if _sha256_file(path) != item["sha256"]: + raise ValueError(f"content digest mismatch for {item['path']}") + + +def validate_lm_eval_runtime(config: LmEvalRuntimeConfig) -> LmEvalRuntime: + """Fail closed unless the configured runtime exactly matches its manifest.""" + + root = Path(config.path) + if not root.is_absolute(): + raise ValueError("lm_eval_runtime.path must be absolute") + root_info = root.lstat() + if not stat.S_ISDIR(root_info.st_mode) or root.is_symlink(): + raise ValueError("lm_eval_runtime.path must be a real directory") + if root_info.st_mode & 0o222: + raise ValueError("lm_eval_runtime root must not be writable") + entries = {path.name for path in root.iterdir()} + if entries != {LM_EVAL_MANIFEST_FILENAME, "site-packages"}: + raise ValueError("lm_eval_runtime root must contain only manifest and site-packages") + + site_packages = root / "site-packages" + site_info = site_packages.lstat() + if not stat.S_ISDIR(site_info.st_mode) or site_packages.is_symlink(): + raise ValueError("lm_eval_runtime site-packages must be a real directory") + if site_info.st_mode & 0o222: + raise ValueError("lm_eval_runtime site-packages must not be writable") + + manifest_path = root / LM_EVAL_MANIFEST_FILENAME + manifest, manifest_bytes = _read_json_file( + manifest_path, + limit=MAX_MANIFEST_SIZE_BYTES, + require_readonly=True, + ) + if set(manifest) != _MANIFEST_KEYS: + raise ValueError("lm_eval runtime manifest keys do not match the v1 contract") + if manifest.get("schema") != LM_EVAL_MANIFEST_SCHEMA: + raise ValueError(f"unsupported lm_eval runtime schema: {manifest.get('schema')!r}") + if manifest.get("site_packages") != "site-packages": + raise ValueError("manifest site_packages must be exactly 'site-packages'") + identity = _validate_identity(manifest.get("identity")) + if identity != config.identity: + raise ValueError("manifest identity does not exactly match benchmark config") + records = _validate_manifest_records(manifest.get("files")) + computed = _canonical_sha256(identity, records) + declared = manifest.get("runtime_sha256") + if not isinstance(declared, str) or not _SHA256.fullmatch(declared): + raise ValueError("manifest runtime_sha256 is invalid") + if computed != declared or declared != config.sha256: + raise ValueError("lm_eval runtime digest does not match manifest/config") + _verify_site_packages(site_packages, records) + return LmEvalRuntime( + root=root.resolve(), + site_packages=site_packages.resolve(), + runtime_sha256=declared, + identity=identity, + manifest_bytes=manifest_bytes, + manifest_sha256=_sha256_bytes(manifest_bytes), + ) + + +def snapshot_runtime_manifest(runtime: LmEvalRuntime, workspace: Path) -> Path: + """Atomically preserve the exact consumed manifest in the run workspace.""" + + destination = Path(workspace) / LM_EVAL_MANIFEST_FILENAME + temporary = destination.with_name(f".{destination.name}.{uuid.uuid4().hex}.tmp") + try: + with temporary.open("xb") as stream: + stream.write(runtime.manifest_bytes) + stream.flush() + os.fsync(stream.fileno()) + os.replace(temporary, destination) + finally: + temporary.unlink(missing_ok=True) + return destination + + +def _base_evidence( + config: LmEvalRuntimeConfig | None, + status: str, + *, + requested: bool, +) -> Dict[str, Any]: + return { + "schema": LM_EVAL_EVIDENCE_SCHEMA, + "requested": requested, + "status": status, + "verified": status == "verified", + "evidence_present": False, + "runtime_sha256": config.sha256 if config else None, + "identity": dict(config.identity) if config else None, + "mount_mode": None, + "manifest_artifact": None, + "receipt_artifact": None, + "errors": [], + } + + +def invalid_runtime_evidence( + config: LmEvalRuntimeConfig | None, + error: str, + *, + status: str = "invalid", +) -> Dict[str, Any]: + evidence = _base_evidence(config, status, requested=True) + evidence["errors"] = [error] + return evidence + + +def collect_lm_eval_runtime_evidence( + workspace: Path, + *, + requested: bool, + config: LmEvalRuntimeConfig | None, + execution_mode: str, +) -> Dict[str, Any]: + """Validate the in-run receipt and preserved manifest for report evidence.""" + + if not requested: + return _base_evidence(config, "not_requested", requested=False) + if config is None: + return invalid_runtime_evidence( + None, + "RUN_EVAL=true requires benchmark.lm_eval_runtime", + ) + evidence = _base_evidence(config, "missing", requested=True) + receipt_path = Path(workspace) / LM_EVAL_RECEIPT_FILENAME + manifest_path = Path(workspace) / LM_EVAL_MANIFEST_FILENAME + try: + payload, receipt_bytes = _read_json_file( + receipt_path, + limit=MAX_RECEIPT_SIZE_BYTES, + ) + manifest, manifest_bytes = _read_json_file( + manifest_path, + limit=MAX_MANIFEST_SIZE_BYTES, + ) + if set(payload) != _RECEIPT_KEYS: + raise ValueError("runtime receipt keys do not match the v1 contract") + if payload.get("schema") != LM_EVAL_RECEIPT_SCHEMA: + raise ValueError("unsupported lm_eval runtime receipt schema") + if payload.get("verified") is not True: + raise ValueError("runtime receipt verified must be true") + if payload.get("runtime_sha256") != config.sha256: + raise ValueError("runtime receipt digest does not match benchmark config") + if payload.get("identity") != config.identity: + raise ValueError("runtime receipt identity does not match benchmark config") + if payload.get("manifest_sha256") != _sha256_bytes(manifest_bytes): + raise ValueError("runtime receipt manifest digest does not match artifact") + if manifest.get("runtime_sha256") != config.sha256: + raise ValueError("manifest artifact runtime digest does not match config") + if manifest.get("identity") != config.identity: + raise ValueError("manifest artifact identity does not match config") + if payload.get("site_packages") != "site-packages": + raise ValueError("runtime receipt site_packages is invalid") + if payload.get("python_abi") != config.identity.get("python_abi"): + raise ValueError("actual Python ABI does not match runtime identity") + if payload.get("lm_eval_version") != config.identity.get("lm_eval_version"): + raise ValueError("actual lm_eval version does not match runtime identity") + if payload.get("execution_mode") != execution_mode: + raise ValueError("runtime receipt execution mode does not match benchmark") + read_only_mount = payload.get("read_only_mount") + if not isinstance(read_only_mount, bool): + raise ValueError("runtime receipt read_only_mount must be boolean") + if execution_mode == "docker" and read_only_mount is not True: + raise ValueError("Docker evaluator runtime was not mounted read-only") + module = payload.get("lm_eval_module") + if not isinstance(module, str) or not module.startswith("site-packages/lm_eval/"): + raise ValueError("lm_eval module did not import from the supplied runtime") + manifest_info = manifest_path.stat() + receipt_info = receipt_path.stat() + except (OSError, UnicodeError, json.JSONDecodeError, ValueError) as exc: + return invalid_runtime_evidence(config, str(exc)) + + evidence.update( + { + "status": "verified", + "verified": True, + "evidence_present": True, + "mount_mode": "read_only" if execution_mode == "docker" else "local", + "manifest_artifact": { + "path": LM_EVAL_MANIFEST_FILENAME, + "size_bytes": manifest_info.st_size, + "sha256": _sha256_bytes(manifest_bytes), + }, + "receipt_artifact": { + "path": LM_EVAL_RECEIPT_FILENAME, + "size_bytes": receipt_info.st_size, + "sha256": _sha256_bytes(receipt_bytes), + }, + "errors": [], + } + ) + return evidence diff --git a/Magpie/modes/benchmark/model_revision.py b/Magpie/modes/benchmark/model_revision.py new file mode 100644 index 0000000..ba65da4 --- /dev/null +++ b/Magpie/modes/benchmark/model_revision.py @@ -0,0 +1,207 @@ +"""Validate model revision evidence emitted by serving benchmark scripts.""" + +from __future__ import annotations + +import hashlib +import json +import re +from pathlib import Path +from typing import Any, Dict, List, Mapping, Tuple + + +MODEL_REVISION_RECEIPT_FILENAME = "model_revision_receipt.json" +MODEL_REVISION_RECEIPT_SCHEMA = "magpie.model-revision-receipt/v1" +MODEL_REVISION_EVIDENCE_SCHEMA = "magpie.model-revision-evidence/v1" +MAX_RECEIPT_SIZE_BYTES = 64 * 1024 + +_EXACT_COMMIT_RE = re.compile(r"^[0-9a-f]{40}$") +_RECEIPT_KEYS = frozenset( + { + "schema", + "model", + "requested_revision", + "resolved_revision", + "snapshot_path", + "verified", + } +) + + +def _reject_json_constant(value: str) -> None: + raise ValueError(f"non-finite JSON number is forbidden: {value}") + + +def _reject_duplicate_keys(pairs: List[Tuple[str, Any]]) -> Dict[str, Any]: + result: Dict[str, Any] = {} + for key, value in pairs: + if key in result: + raise ValueError(f"duplicate JSON object key: {key}") + result[key] = value + return result + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _evidence( + *, + model: str, + requested_revision: str | None, + status: str, + error: str | None = None, +) -> Dict[str, Any]: + requested = requested_revision is not None + return { + "schema": MODEL_REVISION_EVIDENCE_SCHEMA, + "requested": requested, + "status": status, + "verified": status == "verified", + "evidence_present": False, + "model": model, + "requested_revision": requested_revision, + "resolved_revision": None, + "snapshot_path": None, + "receipt_artifact": None, + "errors": [error] if error else [], + } + + +def _validate_payload( + payload: Mapping[str, Any], + *, + model: str, + requested_revision: str, +) -> str | None: + keys = set(payload) + if keys != _RECEIPT_KEYS: + missing = sorted(_RECEIPT_KEYS - keys) + unexpected = sorted(keys - _RECEIPT_KEYS) + return ( + "receipt keys do not match the v1 contract " + f"(missing={missing}, unexpected={unexpected})" + ) + if payload.get("schema") != MODEL_REVISION_RECEIPT_SCHEMA: + return f"unsupported receipt schema: {payload.get('schema')!r}" + if payload.get("model") != model: + return ( + "receipt model does not match the benchmark config: " + f"{payload.get('model')!r} != {model!r}" + ) + if payload.get("requested_revision") != requested_revision: + return "receipt requested_revision does not match MODEL_REVISION" + + resolved_revision = payload.get("resolved_revision") + if not isinstance(resolved_revision, str) or not _EXACT_COMMIT_RE.fullmatch( + resolved_revision + ): + return "receipt resolved_revision is not an exact lowercase 40-hex commit" + if resolved_revision != requested_revision: + return "resolved model revision does not match MODEL_REVISION" + + snapshot_path = payload.get("snapshot_path") + if not isinstance(snapshot_path, str) or not snapshot_path.strip(): + return "receipt snapshot_path must be a non-empty string" + if not Path(snapshot_path).is_absolute(): + return "receipt snapshot_path must be absolute" + if Path(snapshot_path).name != resolved_revision: + return "receipt snapshot_path does not name the resolved revision" + if payload.get("verified") is not True: + return "receipt verified must be true" + return None + + +def collect_model_revision_evidence( + workspace: Path, + *, + model: str, + requested_revision: Any, +) -> Dict[str, Any]: + """Collect a bounded, validated model revision receipt from ``workspace``. + + No revision request is a supported legacy mode and returns ``not_requested``. + Once ``MODEL_REVISION`` is supplied, every malformed or absent receipt is a + failed evidence gate; callers must not treat agent text or server logs as a + substitute. + """ + + if requested_revision in (None, ""): + return _evidence( + model=model, + requested_revision=None, + status="not_requested", + ) + + requested = str(requested_revision).strip() + if not _EXACT_COMMIT_RE.fullmatch(requested): + return _evidence( + model=model, + requested_revision=requested, + status="invalid", + error="MODEL_REVISION must be an exact lowercase 40-hex commit", + ) + + evidence = _evidence( + model=model, + requested_revision=requested, + status="missing", + error=( + f"{MODEL_REVISION_RECEIPT_FILENAME} is missing for requested " + "MODEL_REVISION" + ), + ) + receipt_path = Path(workspace) / MODEL_REVISION_RECEIPT_FILENAME + try: + if receipt_path.is_symlink(): + raise ValueError("receipt must be a regular file, not a symlink") + if not receipt_path.is_file(): + return evidence + size_bytes = receipt_path.stat().st_size + if size_bytes <= 0 or size_bytes > MAX_RECEIPT_SIZE_BYTES: + raise ValueError( + "receipt size must be between 1 and " + f"{MAX_RECEIPT_SIZE_BYTES} bytes, got {size_bytes}" + ) + payload = json.loads( + receipt_path.read_text(encoding="utf-8"), + parse_constant=_reject_json_constant, + object_pairs_hook=_reject_duplicate_keys, + ) + if not isinstance(payload, Mapping): + raise ValueError("receipt root must be a JSON object") + validation_error = _validate_payload( + payload, + model=model, + requested_revision=requested, + ) + if validation_error: + raise ValueError(validation_error) + receipt_sha256 = _sha256(receipt_path) + except (OSError, UnicodeError, json.JSONDecodeError, ValueError) as exc: + return _evidence( + model=model, + requested_revision=requested, + status="invalid", + error=str(exc), + ) + + evidence.update( + { + "status": "verified", + "verified": True, + "evidence_present": True, + "resolved_revision": payload["resolved_revision"], + "snapshot_path": payload["snapshot_path"], + "receipt_artifact": { + "path": MODEL_REVISION_RECEIPT_FILENAME, + "size_bytes": size_bytes, + "sha256": receipt_sha256, + }, + "errors": [], + } + ) + return evidence diff --git a/Magpie/modes/benchmark/result.py b/Magpie/modes/benchmark/result.py index 973da8a..f02a6ad 100644 --- a/Magpie/modes/benchmark/result.py +++ b/Magpie/modes/benchmark/result.py @@ -156,6 +156,20 @@ class BenchmarkResult: profiling_enabled: bool = False run_kind: str = "" reward_eligible: bool = False + + # Requested/resolved Hugging Face model revision evidence. The benchmark + # controller validates the script-emitted workspace receipt before exposing + # it here; unpinned legacy runs use status="not_requested". + model_revision_receipt: Optional[Dict[str, Any]] = None + + # Run-scoped InferenceX source/materialization evidence. The runtime tree + # is disposable and may be patched; the configured source checkout remains + # untouched at its recorded commit/status. + inferencex_runtime_receipt: Optional[Dict[str, Any]] = None + + # In-container proof that the exact caller-supplied, read-only lm-eval + # runtime was validated and imported. Required whenever RUN_EVAL=true. + lm_eval_runtime_receipt: Optional[Dict[str, Any]] = None # Errors errors: List[str] = field(default_factory=list) @@ -192,6 +206,9 @@ def to_dict(self) -> Dict[str, Any]: "profiling_enabled": self.profiling_enabled, "run_kind": self.run_kind, "reward_eligible": self.reward_eligible, + "model_revision_receipt": self.model_revision_receipt, + "inferencex_runtime_receipt": self.inferencex_runtime_receipt, + "lm_eval_runtime_receipt": self.lm_eval_runtime_receipt, "errors": self.errors, } # Scriptable (server-less) extras — e.g. xDiT diffusion. Only emit when @@ -238,6 +255,33 @@ def get_summary(self) -> str: f" ITL (mean/p99): {self.latency.itl_mean:.2f}ms / {self.latency.itl_p99:.2f}ms", f" E2EL (mean/p99): {self.latency.e2el_mean:.2f}ms / {self.latency.e2el_p99:.2f}ms", ]) + + if self.model_revision_receipt is not None: + revision = self.model_revision_receipt + lines.extend( + [ + "", + "Model revision evidence:", + f" Status: {revision.get('status', 'unknown')}", + " Requested: " + f"{revision.get('requested_revision') or 'not requested'}", + " Resolved: " + f"{revision.get('resolved_revision') or 'not verified'}", + ] + ) + + if self.lm_eval_runtime_receipt is not None: + runtime = self.lm_eval_runtime_receipt + lines.extend( + [ + "", + "lm-eval runtime evidence:", + f" Status: {runtime.get('status', 'unknown')}", + " Digest: " + f"{runtime.get('runtime_sha256') or 'not supplied'}", + f" Mount: {runtime.get('mount_mode') or 'not activated'}", + ] + ) if self.top_bottlenecks: lines.extend([ diff --git a/Magpie/scripts/benchmark/atom_mi300x.sh b/Magpie/scripts/benchmark/atom_mi300x.sh index 954080d..09837aa 100644 --- a/Magpie/scripts/benchmark/atom_mi300x.sh +++ b/Magpie/scripts/benchmark/atom_mi300x.sh @@ -20,6 +20,7 @@ # launch, no server-side cleanup, no SERVER_PID monitoring). source "$(dirname "$0")/benchmark_lib.sh" +source "$(dirname "$0")/lm_eval_runtime.sh" || exit $? source "$(dirname "$0")/server_cleanup.sh" # shellcheck source=magpie_bench_remote_compat.sh [[ -f "$(dirname "$0")/magpie_bench_remote_compat.sh" ]] && source "$(dirname "$0")/magpie_bench_remote_compat.sh" diff --git a/Magpie/scripts/benchmark/atom_mi355x.sh b/Magpie/scripts/benchmark/atom_mi355x.sh index 51cfc80..649e8a2 100644 --- a/Magpie/scripts/benchmark/atom_mi355x.sh +++ b/Magpie/scripts/benchmark/atom_mi355x.sh @@ -16,6 +16,7 @@ # atom_mi300x.sh for the full contract. source "$(dirname "$0")/benchmark_lib.sh" +source "$(dirname "$0")/lm_eval_runtime.sh" || exit $? source "$(dirname "$0")/server_cleanup.sh" # shellcheck source=magpie_bench_remote_compat.sh [[ -f "$(dirname "$0")/magpie_bench_remote_compat.sh" ]] && source "$(dirname "$0")/magpie_bench_remote_compat.sh" diff --git a/Magpie/scripts/benchmark/lm_eval_runtime.sh b/Magpie/scripts/benchmark/lm_eval_runtime.sh new file mode 100644 index 0000000..3f591dd --- /dev/null +++ b/Magpie/scripts/benchmark/lm_eval_runtime.sh @@ -0,0 +1,250 @@ +#!/usr/bin/env bash +############################################################################### +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### + +# Replace InferenceX's mutable evaluator setup with a caller-provided, +# hash-locked runtime. This helper performs only local reads and imports. + +magpie_activate_lm_eval_runtime() { + local runtime_root="${MAGPIE_LM_EVAL_RUNTIME_ROOT:-}" + local expected_sha256="${MAGPIE_LM_EVAL_RUNTIME_SHA256:-}" + local receipt_path="${MAGPIE_LM_EVAL_RUNTIME_RECEIPT:-}" + local execution_mode="${MAGPIE_LM_EVAL_EXECUTION_MODE:-}" + local require_readonly_mount="${MAGPIE_LM_EVAL_REQUIRE_READONLY_MOUNT:-}" + + if [[ -z "$runtime_root" || -z "$expected_sha256" || -z "$receipt_path" \ + || ( "$execution_mode" != "docker" && "$execution_mode" != "local" ) \ + || ( "$require_readonly_mount" != "0" && "$require_readonly_mount" != "1" ) ]]; then + echo "ERROR: hash-locked lm-eval runtime environment is incomplete." >&2 + return 41 + fi + + export PYTHONDONTWRITEBYTECODE=1 + export PYTHONPATH="${runtime_root}/site-packages${PYTHONPATH:+:${PYTHONPATH}}" + + python3 - "$runtime_root" "$expected_sha256" "$receipt_path" \ + "$execution_mode" "$require_readonly_mount" <<'PY' +import hashlib +import importlib.metadata +import json +import os +import re +import stat +import sys +import uuid +from pathlib import Path, PurePosixPath + +root = Path(sys.argv[1]) +expected = sys.argv[2] +receipt = Path(sys.argv[3]) +execution_mode = sys.argv[4] +require_readonly_mount = sys.argv[5] == "1" +manifest_path = root / "lm_eval_runtime_manifest.json" +sha256_re = re.compile(r"^[0-9a-f]{64}$") + + +def reject_constant(value): + raise ValueError(f"non-finite JSON number is forbidden: {value}") + + +def reject_duplicates(pairs): + result = {} + for key, value in pairs: + if key in result: + raise ValueError(f"duplicate JSON object key: {key}") + result[key] = value + return result + + +def file_sha256(path): + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def require_readonly_directory(path, label): + info = path.lstat() + if not stat.S_ISDIR(info.st_mode) or path.is_symlink(): + raise ValueError(f"{label} must be a real directory") + if info.st_mode & 0o222: + raise ValueError(f"{label} must not have writable permission bits") + + +if not sha256_re.fullmatch(expected): + raise SystemExit("ERROR: expected lm-eval runtime digest is invalid") +require_readonly_directory(root, "runtime root") +read_only_mount = bool(os.statvfs(root).f_flag & os.ST_RDONLY) +if require_readonly_mount and not read_only_mount: + raise SystemExit("ERROR: evaluator runtime bind mount is not read-only") +if {item.name for item in root.iterdir()} != { + "lm_eval_runtime_manifest.json", + "site-packages", +}: + raise SystemExit("ERROR: runtime root contains unexpected entries") +require_readonly_directory(root / "site-packages", "site-packages") + +manifest_info = manifest_path.lstat() +if ( + not stat.S_ISREG(manifest_info.st_mode) + or manifest_info.st_nlink != 1 + or manifest_info.st_mode & 0o222 + or manifest_info.st_size <= 0 + or manifest_info.st_size > 64 * 1024 * 1024 +): + raise SystemExit("ERROR: runtime manifest is not a bounded read-only nlink=1 file") +manifest_bytes = manifest_path.read_bytes() +manifest = json.loads( + manifest_bytes.decode("utf-8"), + parse_constant=reject_constant, + object_pairs_hook=reject_duplicates, +) +if not isinstance(manifest, dict) or set(manifest) != { + "schema", + "runtime_sha256", + "site_packages", + "identity", + "files", +}: + raise SystemExit("ERROR: runtime manifest keys do not match the v1 contract") +if manifest["schema"] != "apex.lm-eval-runtime/v1": + raise SystemExit("ERROR: unsupported lm-eval runtime manifest schema") +if manifest["site_packages"] != "site-packages": + raise SystemExit("ERROR: runtime manifest site_packages is invalid") +identity = manifest["identity"] +files = manifest["files"] +if not isinstance(identity, dict) or not isinstance(files, list) or not files: + raise SystemExit("ERROR: runtime identity/files are invalid") + +canonical = json.dumps( + {"identity": identity, "files": files}, + ensure_ascii=False, + separators=(",", ":"), + sort_keys=True, +).encode("utf-8") +computed = hashlib.sha256(canonical).hexdigest() +if computed != expected or manifest["runtime_sha256"] != expected: + raise SystemExit("ERROR: runtime manifest digest does not match expected digest") + +site_packages = root / "site-packages" +actual_paths = [] +for path in sorted(site_packages.rglob("*"), key=lambda item: item.as_posix()): + info = path.lstat() + relative = path.relative_to(site_packages).as_posix() + if stat.S_ISLNK(info.st_mode): + raise SystemExit(f"ERROR: runtime symlink is forbidden: {relative}") + if stat.S_ISDIR(info.st_mode): + if info.st_mode & 0o222: + raise SystemExit(f"ERROR: runtime directory is writable: {relative}") + continue + if not stat.S_ISREG(info.st_mode) or info.st_nlink != 1: + raise SystemExit(f"ERROR: runtime entry is not a regular nlink=1 file: {relative}") + actual_paths.append(relative) + +expected_paths = [] +for item in files: + if not isinstance(item, dict) or set(item) != { + "path", + "size_bytes", + "mode", + "sha256", + }: + raise SystemExit("ERROR: invalid runtime file record") + raw_path = item["path"] + pure_path = PurePosixPath(raw_path) if isinstance(raw_path, str) else None + if ( + pure_path is None + or not raw_path + or pure_path.is_absolute() + or pure_path.as_posix() != raw_path + or ".." in pure_path.parts + ): + raise SystemExit("ERROR: runtime file path is not canonical and relative") + expected_paths.append(raw_path) +if expected_paths != sorted(expected_paths) or len(expected_paths) != len(set(expected_paths)): + raise SystemExit("ERROR: runtime file records are not unique and sorted") +if actual_paths != expected_paths: + raise SystemExit("ERROR: runtime files do not exactly match the manifest") + +for item in files: + path = site_packages / item["path"] + info = path.lstat() + if ( + not isinstance(item["size_bytes"], int) + or isinstance(item["size_bytes"], bool) + or item["size_bytes"] < 0 + or info.st_size != item["size_bytes"] + ): + raise SystemExit(f"ERROR: runtime file size mismatch: {item['path']}") + if ( + not isinstance(item["mode"], int) + or isinstance(item["mode"], bool) + or item["mode"] & 0o222 + or stat.S_IMODE(info.st_mode) != item["mode"] + ): + raise SystemExit(f"ERROR: runtime file mode mismatch: {item['path']}") + if not isinstance(item["sha256"], str) or not sha256_re.fullmatch(item["sha256"]): + raise SystemExit(f"ERROR: runtime file digest is invalid: {item['path']}") + if file_sha256(path) != item["sha256"]: + raise SystemExit(f"ERROR: runtime file content mismatch: {item['path']}") + +actual_abi = sys.implementation.cache_tag +expected_abi = identity.get("python_abi") +if actual_abi != expected_abi: + raise SystemExit(f"ERROR: Python ABI mismatch: {actual_abi} != {expected_abi}") +actual_version = importlib.metadata.version("lm_eval") +expected_version = identity.get("lm_eval_version") +if actual_version != expected_version: + raise SystemExit(f"ERROR: lm_eval version mismatch: {actual_version} != {expected_version}") + +import lm_eval + +module_path = Path(lm_eval.__file__).resolve(strict=True) +site_root = site_packages.resolve(strict=True) +try: + module_relative = module_path.relative_to(site_root).as_posix() +except ValueError as exc: + raise SystemExit("ERROR: lm_eval imported outside the supplied runtime") from exc +if not module_relative.startswith("lm_eval/"): + raise SystemExit("ERROR: imported lm_eval module path is invalid") + +payload = { + "schema": "magpie.lm-eval-runtime-receipt/v1", + "runtime_sha256": expected, + "identity": identity, + "manifest_sha256": hashlib.sha256(manifest_bytes).hexdigest(), + "site_packages": "site-packages", + "python_abi": actual_abi, + "lm_eval_version": actual_version, + "lm_eval_module": f"site-packages/{module_relative}", + "execution_mode": execution_mode, + "read_only_mount": read_only_mount, + "verified": True, +} +receipt.parent.mkdir(parents=True, exist_ok=True) +temporary = receipt.with_name(f".{receipt.name}.{uuid.uuid4().hex}.tmp") +try: + with temporary.open("x", encoding="utf-8") as stream: + json.dump(payload, stream, indent=2, sort_keys=True) + stream.write("\n") + stream.flush() + os.fsync(stream.fileno()) + os.replace(temporary, receipt) +finally: + temporary.unlink(missing_ok=True) +PY +} + +# InferenceX calls this hook immediately before importing/running lm_eval. +# Replacing it removes every mutable or network-backed dependency path. +_install_lm_eval_deps() { + local status + magpie_activate_lm_eval_runtime && return 0 + status=$? + echo "ERROR: refusing to run lm-eval without a verified locked runtime." >&2 + exit "$status" +} diff --git a/Magpie/scripts/benchmark/magpie_bench_remote_compat.sh b/Magpie/scripts/benchmark/magpie_bench_remote_compat.sh index 2b4b619..7b0eaa0 100644 --- a/Magpie/scripts/benchmark/magpie_bench_remote_compat.sh +++ b/Magpie/scripts/benchmark/magpie_bench_remote_compat.sh @@ -120,6 +120,8 @@ magpie_run_eval_remote_direct() { return 1 fi + _install_lm_eval_deps || return $? + local py="${MAGPIE_EVAL_PYTHON:-python3}" local result_dir="${RESULT_DIR:-${WORKSPACE_DIR:-/workspace}}" local out_dir="${result_dir%/}/lm_eval" diff --git a/Magpie/scripts/benchmark/sglang_mi300x.sh b/Magpie/scripts/benchmark/sglang_mi300x.sh index 5c70105..ea90f87 100644 --- a/Magpie/scripts/benchmark/sglang_mi300x.sh +++ b/Magpie/scripts/benchmark/sglang_mi300x.sh @@ -14,6 +14,7 @@ # default behaviour of launching a local server. source "$(dirname "$0")/benchmark_lib.sh" +source "$(dirname "$0")/lm_eval_runtime.sh" || exit $? source "$(dirname "$0")/server_cleanup.sh" # shellcheck source=magpie_bench_remote_compat.sh [[ -f "$(dirname "$0")/magpie_bench_remote_compat.sh" ]] && source "$(dirname "$0")/magpie_bench_remote_compat.sh" diff --git a/Magpie/scripts/benchmark/sglang_mi355x.sh b/Magpie/scripts/benchmark/sglang_mi355x.sh index 2e64d69..3c0cc23 100644 --- a/Magpie/scripts/benchmark/sglang_mi355x.sh +++ b/Magpie/scripts/benchmark/sglang_mi355x.sh @@ -15,6 +15,7 @@ # launch). See sglang_mi300x.sh for the full contract. source "$(dirname "$0")/benchmark_lib.sh" +source "$(dirname "$0")/lm_eval_runtime.sh" || exit $? source "$(dirname "$0")/server_cleanup.sh" # shellcheck source=magpie_bench_remote_compat.sh [[ -f "$(dirname "$0")/magpie_bench_remote_compat.sh" ]] && source "$(dirname "$0")/magpie_bench_remote_compat.sh" diff --git a/Magpie/scripts/benchmark/vllm_mi300x.sh b/Magpie/scripts/benchmark/vllm_mi300x.sh index 1bf78b1..36d9d39 100644 --- a/Magpie/scripts/benchmark/vllm_mi300x.sh +++ b/Magpie/scripts/benchmark/vllm_mi300x.sh @@ -19,6 +19,7 @@ # default behaviour of launching a local server. source "$(dirname "$0")/benchmark_lib.sh" +source "$(dirname "$0")/lm_eval_runtime.sh" || exit $? source "$(dirname "$0")/server_cleanup.sh" # shellcheck source=magpie_bench_remote_compat.sh [[ -f "$(dirname "$0")/magpie_bench_remote_compat.sh" ]] && source "$(dirname "$0")/magpie_bench_remote_compat.sh" diff --git a/Magpie/scripts/benchmark/vllm_mi355x.sh b/Magpie/scripts/benchmark/vllm_mi355x.sh index c4e6f41..17837fc 100644 --- a/Magpie/scripts/benchmark/vllm_mi355x.sh +++ b/Magpie/scripts/benchmark/vllm_mi355x.sh @@ -16,6 +16,7 @@ # launch). See vllm_mi300x.sh for the full contract. source "$(dirname "$0")/benchmark_lib.sh" +source "$(dirname "$0")/lm_eval_runtime.sh" || exit $? source "$(dirname "$0")/server_cleanup.sh" # shellcheck source=magpie_bench_remote_compat.sh [[ -f "$(dirname "$0")/magpie_bench_remote_compat.sh" ]] && source "$(dirname "$0")/magpie_bench_remote_compat.sh" @@ -42,12 +43,97 @@ fi MAX_MODEL_LEN=${MAX_MODEL_LEN:-4096} GPU_MEMORY_UTILIZATION=${GPU_MEMORY_UTILIZATION:-0.95} +WORKSPACE_DIR=${RESULT_DIR:-/workspace} +MODEL_REVISION_RECEIPT="$WORKSPACE_DIR/model_revision_receipt.json" + +mkdir -p "$WORKSPACE_DIR" + +MODEL_REVISION_ARGS=() +if [[ -n "${MODEL_REVISION:-}" ]]; then + if [[ ! "$MODEL_REVISION" =~ ^[0-9a-f]{40}$ ]]; then + echo "ERROR: MODEL_REVISION must be an exact lowercase 40-hex commit." >&2 + exit 4 + fi + MODEL_REVISION_ARGS+=(--revision "$MODEL_REVISION") +fi if [[ -n "$SLURM_JOB_ID" ]]; then echo "JOB $SLURM_JOB_ID running on $SLURMD_NODENAME" fi -if [[ "$PHASE" != "client" ]]; then +if [[ "$PHASE" != "client" && -n "${MODEL_REVISION:-}" ]]; then + if ! rm -f "$MODEL_REVISION_RECEIPT"; then + echo "ERROR: could not clear stale model revision receipt." >&2 + exit 4 + fi + MODEL_SNAPSHOT_PATH="$( + hf download "$MODEL" --revision "$MODEL_REVISION" --format quiet + )" + download_status=$? + if [[ $download_status -ne 0 ]]; then + echo "ERROR: hf download failed for MODEL_REVISION=$MODEL_REVISION." >&2 + exit "$download_status" + fi + + python3 - "$MODEL" "$MODEL_REVISION" "$MODEL_SNAPSHOT_PATH" \ + "$MODEL_REVISION_RECEIPT" <<'PY' +import json +import os +import re +import sys +from pathlib import Path + +model, requested, raw_snapshot, raw_receipt = sys.argv[1:] +commit_re = re.compile(r"^[0-9a-f]{40}$") +snapshot = Path(raw_snapshot).resolve(strict=True) +if not snapshot.is_dir(): + raise SystemExit(f"ERROR: hf download did not resolve to a directory: {snapshot}") +resolved = snapshot.name +if not commit_re.fullmatch(resolved): + raise SystemExit( + "ERROR: resolved Hugging Face snapshot is not an exact 40-hex commit: " + f"{resolved!r}" + ) +if resolved != requested: + raise SystemExit( + "ERROR: resolved Hugging Face snapshot does not match MODEL_REVISION: " + f"{resolved} != {requested}" + ) + +receipt = Path(raw_receipt) +receipt.parent.mkdir(parents=True, exist_ok=True) +temporary = receipt.with_name(f".{receipt.name}.{os.getpid()}.tmp") +payload = { + "schema": "magpie.model-revision-receipt/v1", + "model": model, + "requested_revision": requested, + "resolved_revision": resolved, + "snapshot_path": str(snapshot), + "verified": True, +} +try: + with temporary.open("x", encoding="utf-8") as stream: + json.dump(payload, stream, indent=2, sort_keys=True) + stream.write("\n") + stream.flush() + os.fsync(stream.fileno()) + os.replace(temporary, receipt) +except Exception: + temporary.unlink(missing_ok=True) + raise +PY + receipt_status=$? + if [[ $receipt_status -ne 0 ]]; then + echo "ERROR: failed to verify or persist model revision receipt." >&2 + exit "$receipt_status" + fi +elif [[ "$PHASE" != "client" ]]; then + # Keep legacy unpinned benchmarks working, but never leave a stale receipt + # that a report consumer could mistake for evidence for the current run. + if ! rm -f "$MODEL_REVISION_RECEIPT"; then + echo "ERROR: could not clear stale model revision receipt." >&2 + exit 4 + fi hf download "$MODEL" 2>/dev/null || true fi @@ -67,7 +153,6 @@ fi # vLLM optimizations for MI355X export VLLM_ROCM_USE_AITER=${VLLM_ROCM_USE_AITER:-1} -WORKSPACE_DIR=${RESULT_DIR:-/workspace} SERVER_LOG=${SERVER_LOG:-$WORKSPACE_DIR/server.log} PORT=${PORT:-8888} @@ -86,7 +171,8 @@ fi set -x if [[ "$PHASE" == "server" || "$PHASE" == "all" ]]; then - setsid vllm serve $MODEL --port $PORT \ + setsid vllm serve "$MODEL" --port "$PORT" \ + "${MODEL_REVISION_ARGS[@]}" \ --tensor-parallel-size=$TP \ --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" \ --max-model-len $MAX_MODEL_LEN \ @@ -148,7 +234,8 @@ if [[ "$PHASE" != "server" && "${RUN_EVAL,,}" = "true" ]]; then fi else magpie_mark_lm_eval_start || exit $? - run_eval --framework lm-eval --port "$PORT" --concurrent-requests $CONC || exit $? + EVAL_CONCURRENT_REQUESTS="${EVAL_CONCURRENT_REQUESTS:-$CONC}" \ + run_eval --framework lm-eval --port "$PORT" || exit $? magpie_preserve_lm_eval_artifacts || exit $? append_lm_eval_summary magpie_preserve_lm_eval_artifacts || exit $? diff --git a/Magpie/tools/amd_kernel_finder/indexer.py b/Magpie/tools/amd_kernel_finder/indexer.py index 440558a..0031034 100644 --- a/Magpie/tools/amd_kernel_finder/indexer.py +++ b/Magpie/tools/amd_kernel_finder/indexer.py @@ -10,13 +10,15 @@ replacing hardcoded mappings with dynamically discovered kernel locations. """ +import hashlib import json import logging import re -from dataclasses import dataclass, asdict +from dataclasses import asdict, dataclass from pathlib import Path from typing import Dict, List, Optional -import hashlib + +from .repo_config import detect_repo_type logger = logging.getLogger(__name__) @@ -48,6 +50,8 @@ class KernelIndex: Scans repositories for kernel definitions and builds a searchable index. Supports caching to avoid rescanning. """ + + CACHE_SCHEMA_VERSION = 2 KERNEL_PATTERNS = { "triton_jit": [ @@ -114,6 +118,10 @@ def _load_cache(self, cache_file: Path, repo_path: str) -> bool: try: with open(cache_file, 'r') as f: data = json.load(f) + + if data.get("schema_version") != self.CACHE_SCHEMA_VERSION: + logger.info(f"Cache schema is stale for {repo_path}") + return False cached_mtime = data.get("mtime", 0) current_mtime = Path(repo_path).stat().st_mtime @@ -138,6 +146,7 @@ def _save_cache(self, cache_file: Path, repo_path: str) -> None: } data = { + "schema_version": self.CACHE_SCHEMA_VERSION, "mtime": Path(repo_path).stat().st_mtime, "definitions": repo_defs, } @@ -200,12 +209,17 @@ def _scan_file(self, file_path: Path, pattern: str, kind: str, logger.debug(f"Error scanning {file_path}: {e}") def _detect_repo_name(self, repo_path: Path) -> str: - if (repo_path / "projects" / "composablekernel").exists(): - return "rocm-libraries" - if (repo_path / "python" / "triton").exists(): - return "triton" - if (repo_path / "vllm").exists() and (repo_path / "csrc").exists(): - return "vllm" + # Repository checkouts are frequently materialized under versioned or + # evidence-specific directory names (for example + # ``aiter-v0.1.10.post2``). Use the shared known-repository structural + # detector so emitted repo variables stay canonical (``$AITER_DIR``), + # independent of the checkout directory name. + detected = detect_repo_type(str(repo_path)) + if detected: + return detected + + # Preserve the looser legacy PyTorch fixture/layout detection. All + # other known repositories are identified through repo_config. if (repo_path / "aten").exists(): return "pytorch" return repo_path.name diff --git a/docs/how-to/benchmarking/benchmark.md b/docs/how-to/benchmarking/benchmark.md index f058012..f603ccb 100644 --- a/docs/how-to/benchmarking/benchmark.md +++ b/docs/how-to/benchmarking/benchmark.md @@ -37,6 +37,9 @@ profiling-options - **ROCm-compatible GPU with sufficient VRAM**: the example configs target AMD Instinct™ GPUs (MI300X/MI355X). DeepSeek-R1 requires 8 GPUs at fp8; smaller models need less. Magpie [selects idle GPUs automatically](automatic-gpu.md). - **HuggingFace token**: required for gated models. Set `HF_TOKEN` in your environment before running. - **InferenceX**: cloned automatically on first run; no manual install needed. +- **Pinned lm-eval runtime for accuracy runs**: `RUN_EVAL=true` requires a + caller-built `benchmark.lm_eval_runtime`. Magpie never resolves or installs + evaluator packages during a benchmark. ### Commands @@ -74,6 +77,11 @@ results/benchmark_vllm_/ ├── container_stdout.log # Container stdout ├── container_stderr.log # Container stderr ├── inferencex_result.json # Raw InferenceX output +├── inferencex_runtime_receipt.json # Exact source/runtime identity +├── inferencex_runtime/ # Private run-scoped InferenceX tree +├── model_revision_receipt.json # Requested/resolved HF snapshot (when pinned) +├── lm_eval_runtime_manifest.json # Preserved content/identity manifest +├── lm_eval_runtime_receipt.json # In-container ABI/import verification ├── lm_eval/ # Preserved serving accuracy artifacts (RUN_EVAL=true) ├── torch_trace/ # Raw torch profiler traces │ ├── *-rank-0.*.pt.trace.json.gz @@ -114,7 +122,54 @@ Every report declares `run_kind` and `reward_eligible`. A `run_kind: measurement` run rejects heavy profilers; diagnostic runs and all TargetedKernelTrace artifacts have `reward_eligible: false`. When `RUN_EVAL=true`, raw lm-eval files remain under `lm_eval/` and `quality_gate` exposes each task's -primary metric. Missing or invalid requested accuracy evidence fails the benchmark. +primary metric. The same run must provide this nested configuration: + +```yaml +benchmark: + envs: + RUN_EVAL: "true" + lm_eval_runtime: + path: /absolute/path/to/content-addressed/runtime + sha256: <64-lowercase-hex-runtime-digest> + identity: + lm_eval_commit: <40-hex-commit> + lm_eval_tree: <40-hex-tree> + lm_eval_version: 0.4.9.2 + python_abi: cpython-312 + base_image_id: sha256:<64-hex-image-id> + base_image_repo_digest: image/name@sha256:<64-hex-repo-digest> + inferencex_commit: <40-hex-commit> + inferencex_tree: <40-hex-tree> +``` + +The runtime root contains only `lm_eval_runtime_manifest.json` and +`site-packages/`. Magpie validates the exact identity, sorted file manifest, +permissions, link counts, and every file digest on the host. Docker runs mount +the root at `/opt/apex/lm-eval-runtime:ro`; the benchmark helper independently +recomputes the digest, checks the actual Python ABI and `lm_eval` version, and +proves that `lm_eval` imported from that mount. The validated report field +`lm_eval_runtime_receipt` binds the full identity and runtime digest to hashes +of the preserved manifest and receipt. The runtime's InferenceX commit/tree +must match the materialized benchmark checkout. `RUN_EVAL=true` currently +supports local or Docker execution through Magpie's built-in vLLM, SGLang, or +Atom MI300X/MI355X scripts; Ray and native/custom scripts fail closed. +Missing runtime, mutation, ABI/version mismatch, or missing receipt fails the +benchmark. There is no package-manager or network fallback. + +The MI355X vLLM script accepts `envs.MODEL_REVISION` as an exact lowercase +40-hex Hugging Face commit. When set, both `hf download` and `vllm serve` are +pinned to it. Magpie then validates `model_revision_receipt.json` and exposes a +bounded `model_revision_receipt` section in the report. A missing, malformed, +or mismatched requested receipt fails the benchmark. Without `MODEL_REVISION`, +the report uses `status: not_requested`; consumers requiring reproducible model +provenance must reject that status. + +Magpie never installs its benchmark scripts into the configured InferenceX +checkout. For a Git checkout it exports the exact `HEAD` commit into the +workspace through a private Git index, records the commit and unchanged source +status in `inferencex_runtime_receipt.json`, and modifies only that disposable +tree. A non-Git InferenceX directory uses a compatibility filesystem copy and +is explicitly marked unpinned in the receipt. Targeted trace selection uses portable symbol glob patterns under `profiler.targeted_trace.targets`; it does not depend on a fixed container-image diff --git a/docs/reference/benchmark-config.md b/docs/reference/benchmark-config.md index 3396aba..a9a7bc3 100644 --- a/docs/reference/benchmark-config.md +++ b/docs/reference/benchmark-config.md @@ -113,6 +113,20 @@ benchmark: # Paths inferencex_path: /path/to/InferenceX # InferenceX installation hf_cache_path: null # HuggingFace cache directory + # Required whenever envs.RUN_EVAL is true. This is a consumer contract: + # Magpie does not build, update, or repair this runtime. + lm_eval_runtime: + path: /absolute/path/to/content-addressed/runtime + sha256: <64-lowercase-hex-runtime-digest> + identity: + lm_eval_commit: <40-hex-commit> + lm_eval_tree: <40-hex-tree> + lm_eval_version: 0.4.9.2 + python_abi: cpython-312 # sys.implementation.cache_tag in target image + base_image_id: sha256:<64-hex-image-id> + base_image_repo_digest: image/name@sha256:<64-hex-repo-digest> + inferencex_commit: <40-hex-commit> + inferencex_tree: <40-hex-tree> # InferenceX specific runner_type: mi300x # Hardware runner type @@ -134,6 +148,33 @@ Pass these variables under `benchmark.envs:` to control request shape, concurren | `GPU_MEM_UTIL` | GPU memory utilization | 0.95 | | `ENABLE_PROFILE` | Enable torch profiler | "false" | | `EXTRA_VLLM_ARGS` | Additional arguments passed to `vllm serve` | "" | +| `RUN_EVAL` | Run the accuracy gate; requires `benchmark.lm_eval_runtime` | "false" | + +### Hash-locked lm-eval runtime + +When `RUN_EVAL=true`, Magpie fails before launching the workload unless +`benchmark.lm_eval_runtime` validates exactly. Its root must contain only a +read-only `lm_eval_runtime_manifest.json` and `site-packages/`, with no +symlinks, hardlinks, writable entries, unlisted files, or digest mismatch. The +manifest uses schema `apex.lm-eval-runtime/v1`; its `runtime_sha256` is SHA-256 +over compact, key-sorted UTF-8 JSON containing exactly `identity` and `files`. +Each path-sorted file record has `path`, `size_bytes`, integer `mode`, and +`sha256`. + +For Docker, Magpie mounts the runtime root read-only and the benchmark helper +revalidates its bytes inside the actual image before importing `lm_eval`. + +The runtime's InferenceX commit and tree must match the checkout materialized +for the benchmark. Accuracy evaluation is currently supported only for local +or Docker runs using Magpie's built-in vLLM/SGLang/Atom MI300X or MI355X +scripts. Ray mode and native/custom InferenceX scripts fail closed when +`RUN_EVAL=true` because they cannot yet attest this locked runtime. +`identity.base_image_id` and `base_image_repo_digest` describe the compatible +parent image; they are not required to equal a derived candidate image ID. +Image-parent lineage is a responsibility of the caller that produced that +candidate. Local runs perform the same content, ABI, version, and import-root +checks without a container mount. Magpie never installs evaluator dependencies +or falls back to the network. ## Examples diff --git a/examples/benchmarks/benchmark_vllm_qwen3_next_80b_fp8.yaml b/examples/benchmarks/benchmark_vllm_qwen3_next_80b_fp8.yaml index 82f0fb5..af54dc3 100644 --- a/examples/benchmarks/benchmark_vllm_qwen3_next_80b_fp8.yaml +++ b/examples/benchmarks/benchmark_vllm_qwen3_next_80b_fp8.yaml @@ -34,6 +34,12 @@ benchmark: # task/metric summary in benchmark_report.json. RUN_EVAL: "true" MAGPIE_EVAL_TASKS: gsm8k + + # Required by RUN_EVAL=true. Apex injects the exact content-addressed + # path/sha256/identity into each resolved measurement/diagnostic/replay + # view. Direct Magpie callers must provide the same nested contract; see + # docs/reference/benchmark-config.md. It is intentionally absent from this + # source workload config because no machine-local CAS path is portable. profiler: torch_profiler: diff --git a/tests/test_inferencex_runtime.py b/tests/test_inferencex_runtime.py new file mode 100644 index 0000000..552e2c2 --- /dev/null +++ b/tests/test_inferencex_runtime.py @@ -0,0 +1,193 @@ +import json +import subprocess +from pathlib import Path + +import pytest + +from Magpie.modes.benchmark.benchmarker import BenchmarkMode +from Magpie.modes.benchmark.config import BenchmarkConfig +from Magpie.modes.benchmark.inferencex_runtime import ( + INFERENCEX_RUNTIME_RECEIPT_FILENAME, + INFERENCEX_RUNTIME_RECEIPT_SCHEMA, + materialize_inferencex_runtime, +) +from Magpie.modes.benchmark.result import BenchmarkResult + + +def _git(repo: Path, *args: str) -> str: + completed = subprocess.run( + ["git", "-C", str(repo), *args], + text=True, + capture_output=True, + check=True, + ) + return completed.stdout.strip() + + +def _create_inferencex_repo(tmp_path: Path) -> Path: + repo = tmp_path / "InferenceX" + benchmarks = repo / "benchmarks" + benchmarks.mkdir(parents=True) + (benchmarks / "benchmark_lib.sh").write_text( + "run_benchmark_serving() { return 0; }\n", + encoding="utf-8", + ) + _git(repo, "init", "-q") + _git(repo, "config", "user.email", "test@example.com") + _git(repo, "config", "user.name", "Test User") + _git(repo, "add", ".") + _git(repo, "commit", "-qm", "initial") + return repo + + +def test_materialize_inferencex_runtime_exports_commit_without_source_writes( + tmp_path, +): + source = _create_inferencex_repo(tmp_path) + workspace = tmp_path / "workspace" + workspace.mkdir() + status_before = _git(source, "status", "--porcelain=v1", "--untracked-files=all") + commit = _git(source, "rev-parse", "HEAD") + tree = _git(source, "rev-parse", "HEAD^{tree}") + + runtime = materialize_inferencex_runtime(source, workspace) + + status_after = _git(source, "status", "--porcelain=v1", "--untracked-files=all") + assert status_after == status_before == "" + assert runtime.root == workspace / "inferencex_runtime" + assert not (runtime.root / ".git").exists() + assert runtime.receipt == json.loads( + (workspace / INFERENCEX_RUNTIME_RECEIPT_FILENAME).read_text( + encoding="utf-8" + ) + ) + assert runtime.receipt["schema"] == INFERENCEX_RUNTIME_RECEIPT_SCHEMA + assert runtime.receipt["source_commit"] == commit + assert runtime.receipt["source_tree"] == tree + assert runtime.receipt["source_clean"] is True + assert runtime.receipt["source_status_unchanged"] is True + assert runtime.receipt["materialization_method"] == ( + "git_private_index_checkout" + ) + + runtime_lib = runtime.root / "benchmarks/benchmark_lib.sh" + runtime_lib.write_text("runtime-only change\n", encoding="utf-8") + assert (source / "benchmarks/benchmark_lib.sh").read_text( + encoding="utf-8" + ).startswith("run_benchmark_serving") + assert _git(source, "status", "--porcelain=v1", "--untracked-files=all") == "" + + +def test_materialize_inferencex_runtime_ignores_dirty_and_untracked_source( + tmp_path, +): + source = _create_inferencex_repo(tmp_path) + tracked = source / "benchmarks/benchmark_lib.sh" + committed_content = tracked.read_text(encoding="utf-8") + tracked.write_text("unstaged source change\n", encoding="utf-8") + (source / "benchmarks/untracked.sh").write_text( + "untracked\n", encoding="utf-8" + ) + status_before = _git(source, "status", "--porcelain=v1", "--untracked-files=all") + workspace = tmp_path / "workspace" + workspace.mkdir() + + runtime = materialize_inferencex_runtime(source, workspace) + + assert (runtime.root / "benchmarks/benchmark_lib.sh").read_text( + encoding="utf-8" + ) == committed_content + assert not (runtime.root / "benchmarks/untracked.sh").exists() + assert runtime.receipt["source_clean"] is False + assert _git( + source, "status", "--porcelain=v1", "--untracked-files=all" + ) == status_before + + +def test_benchmark_mode_adds_scripts_only_to_disposable_runtime( + tmp_path, + monkeypatch, +): + source = _create_inferencex_repo(tmp_path) + source_status = _git( + source, "status", "--porcelain=v1", "--untracked-files=all" + ) + config = BenchmarkConfig( + framework="vllm", + model="example/model", + run_mode="local", + run_kind="measurement", + envs={"TP": 1}, + profiler={ + "torch_profiler": {"enabled": False}, + "gpu_monitor": {"enabled": False}, + }, + gpu_selection={"auto": False}, + inferencex_path=str(source), + runner_type="mi355x", + benchmark_script="vllm_mi355x.sh", + ) + mode = BenchmarkMode(config, output_dir=str(tmp_path / "results")) + monkeypatch.setattr(mode, "_cleanup_server_processes", lambda framework: None) + + def execute(local_cmd, local_env, workspace): + (workspace / "inferencex_result.json").write_text( + json.dumps( + { + "request_throughput": 1.0, + "output_throughput": 10.0, + "completed": 1, + } + ), + encoding="utf-8", + ) + return BenchmarkResult(success=True), "", "" + + monkeypatch.setattr(mode, "_execute_local_benchmark", execute) + + result = mode.run(task_id="non-mutating-inferencex") + + assert result.success is True + assert _git( + source, "status", "--porcelain=v1", "--untracked-files=all" + ) == source_status + assert not (source / "benchmarks/vllm_mi355x.sh").exists() + runtime_root = Path(result.workspace_dir) / "inferencex_runtime" + assert (runtime_root / "benchmarks/vllm_mi355x.sh").is_file() + assert config.inferencex_path == str(runtime_root) + assert result.inferencex_runtime_receipt["source_commit"] == _git( + source, "rev-parse", "HEAD" + ) + assert result.inferencex_runtime_receipt["source_tree"] == _git( + source, "rev-parse", "HEAD^{tree}" + ) + report = json.loads( + (Path(result.workspace_dir) / "benchmark_report.json").read_text( + encoding="utf-8" + ) + ) + assert report["inferencex_runtime_receipt"] == ( + result.inferencex_runtime_receipt + ) + + +def test_prepare_benchmark_scripts_refuses_source_checkout(tmp_path): + source = _create_inferencex_repo(tmp_path) + config = BenchmarkConfig( + framework="vllm", + model="example/model", + run_mode="local", + inferencex_path=str(source), + ) + mode = BenchmarkMode(config, output_dir=str(tmp_path / "results")) + mode._inferencex_source_path = str(source.resolve()) + source_status = _git( + source, "status", "--porcelain=v1", "--untracked-files=all" + ) + + with pytest.raises(RuntimeError, match="refusing to install"): + mode._prepare_benchmark_scripts() + + assert _git( + source, "status", "--porcelain=v1", "--untracked-files=all" + ) == source_status diff --git a/tests/test_kernel_index_repo_detection.py b/tests/test_kernel_index_repo_detection.py new file mode 100644 index 0000000..16de2f3 --- /dev/null +++ b/tests/test_kernel_index_repo_detection.py @@ -0,0 +1,100 @@ +import json +from pathlib import Path + +from Magpie.tools.amd_kernel_finder.indexer import KernelIndex + + +def _write_triton_kernel(root: Path, relative_path: str, name: str) -> None: + path = root / relative_path + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + f"""import triton + +@triton.jit +def {name}(x): + return x +""", + encoding="utf-8", + ) + + +def test_kernel_index_canonicalizes_versioned_aiter_root(tmp_path): + repo = tmp_path / "aiter-v0.1.10.post2" + (repo / "aiter/ops").mkdir(parents=True) + (repo / "csrc/kernels").mkdir(parents=True) + _write_triton_kernel( + repo, + "aiter/ops/triton/versioned_kernel.py", + "versioned_aiter_kernel", + ) + index = KernelIndex(cache_dir=str(tmp_path / "cache")) + + index.build([str(repo)], force_rebuild=True) + definition = index.lookup("versioned_aiter_kernel.kd") + + assert definition is not None + assert definition.repo_name == "aiter" + assert definition.file_path == "aiter/ops/triton/versioned_kernel.py" + repo_var = f"${definition.repo_name.upper().replace('-', '_')}_DIR" + assert f"{repo_var}/{definition.file_path}" == ( + "$AITER_DIR/aiter/ops/triton/versioned_kernel.py" + ) + + +def test_kernel_index_keeps_versioned_vllm_root_canonical(tmp_path): + repo = tmp_path / "vllm-v0.19.1" + (repo / "vllm").mkdir(parents=True) + (repo / "csrc").mkdir() + _write_triton_kernel( + repo, + "vllm/lora/ops/triton_ops/versioned_kernel.py", + "versioned_vllm_kernel", + ) + index = KernelIndex(cache_dir=str(tmp_path / "cache")) + + index.build([str(repo)], force_rebuild=True) + definition = index.lookup("versioned_vllm_kernel.kd") + + assert definition is not None + assert definition.repo_name == "vllm" + assert definition.file_path == ( + "vllm/lora/ops/triton_ops/versioned_kernel.py" + ) + repo_var = f"${definition.repo_name.upper().replace('-', '_')}_DIR" + assert repo_var == "$VLLM_DIR" + + +def test_kernel_index_rebuilds_pre_canonicalization_cache(tmp_path): + repo = tmp_path / "aiter-v0.1.10.post2" + (repo / "aiter/ops").mkdir(parents=True) + (repo / "csrc/kernels").mkdir(parents=True) + _write_triton_kernel(repo, "aiter/ops/cached.py", "cached_aiter_kernel") + index = KernelIndex(cache_dir=str(tmp_path / "cache")) + cache_file = index._get_cache_file(str(repo)) + cache_file.write_text( + json.dumps( + { + "mtime": repo.stat().st_mtime, + "definitions": { + "stale": { + "name": "cached_aiter_kernel", + "file_path": "aiter/ops/cached.py", + "repo_name": "aiter-v0.1.10.post2", + "repo_path": str(repo), + "kind": "triton_jit", + "line_number": 1, + "symbol": "stale", + } + }, + } + ), + encoding="utf-8", + ) + + index.build([str(repo)]) + definition = index.lookup("cached_aiter_kernel.kd") + + assert definition is not None + assert definition.repo_name == "aiter" + refreshed = json.loads(cache_file.read_text(encoding="utf-8")) + assert refreshed["schema_version"] == KernelIndex.CACHE_SCHEMA_VERSION diff --git a/tests/test_lm_eval_runtime.py b/tests/test_lm_eval_runtime.py new file mode 100644 index 0000000..b29c74a --- /dev/null +++ b/tests/test_lm_eval_runtime.py @@ -0,0 +1,510 @@ +import hashlib +import json +import os +import shutil +import stat +import subprocess +import sys +from pathlib import Path + +import pytest + +from Magpie.modes.benchmark.benchmarker import BenchmarkMode +from Magpie.modes.benchmark.config import BenchmarkConfig, LmEvalRuntimeConfig +from Magpie.modes.benchmark.lm_eval_runtime import ( + LM_EVAL_EVIDENCE_SCHEMA, + LM_EVAL_MANIFEST_FILENAME, + LM_EVAL_MANIFEST_SCHEMA, + LM_EVAL_RECEIPT_FILENAME, + collect_lm_eval_runtime_evidence, + snapshot_runtime_manifest, + validate_lm_eval_runtime, +) +from Magpie.modes.benchmark.result import BenchmarkResult +from Magpie.utils.gpu import GPUVendor + + +REPO_ROOT = Path(__file__).resolve().parents[1] +HELPER = REPO_ROOT / "Magpie/scripts/benchmark/lm_eval_runtime.sh" +BASE_IMAGE = ( + "vllm/vllm-openai-rocm@sha256:" + "c3457ab4702a5bd665b06d7ba57e6105fe98adc4f5b3d4afcf98ec45551988e0" +) + + +def _file_sha256(path: Path) -> str: + return hashlib.sha256(path.read_bytes()).hexdigest() + + +def _runtime_sha256(identity, files) -> str: + canonical = json.dumps( + {"identity": identity, "files": files}, + ensure_ascii=False, + separators=(",", ":"), + sort_keys=True, + ).encode("utf-8") + return hashlib.sha256(canonical).hexdigest() + + +def _build_runtime( + tmp_path: Path, + *, + python_abi: str | None = None, +) -> tuple[Path, LmEvalRuntimeConfig]: + root = tmp_path / "runtime" + site_packages = root / "site-packages" + package = site_packages / "lm_eval" + dist_info = site_packages / "lm_eval-0.4.9.2.dist-info" + package.mkdir(parents=True) + dist_info.mkdir() + (package / "__init__.py").write_text("VALUE = 'locked'\n", encoding="utf-8") + (dist_info / "METADATA").write_text( + "Metadata-Version: 2.1\nName: lm_eval\nVersion: 0.4.9.2\n", + encoding="utf-8", + ) + + for path in (package / "__init__.py", dist_info / "METADATA"): + path.chmod(0o444) + for path in (package, dist_info, site_packages): + path.chmod(0o555) + + files = [] + for path in sorted(site_packages.rglob("*"), key=lambda item: item.as_posix()): + if path.is_file(): + files.append( + { + "path": path.relative_to(site_packages).as_posix(), + "size_bytes": path.stat().st_size, + "mode": stat.S_IMODE(path.stat().st_mode), + "sha256": _file_sha256(path), + } + ) + identity = { + "lm_eval_commit": "a" * 40, + "lm_eval_tree": "b" * 40, + "lm_eval_version": "0.4.9.2", + "python_abi": python_abi or sys.implementation.cache_tag, + "base_image_id": "sha256:" + "c" * 64, + "base_image_repo_digest": "example/image@sha256:" + "d" * 64, + "inferencex_commit": "e" * 40, + "inferencex_tree": "f" * 40, + "lock_sha256": "1" * 64, + } + runtime_sha256 = _runtime_sha256(identity, files) + manifest = { + "schema": LM_EVAL_MANIFEST_SCHEMA, + "runtime_sha256": runtime_sha256, + "site_packages": "site-packages", + "identity": identity, + "files": files, + } + manifest_path = root / LM_EVAL_MANIFEST_FILENAME + manifest_path.write_text( + json.dumps(manifest, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + manifest_path.chmod(0o444) + root.chmod(0o555) + return root, LmEvalRuntimeConfig( + path=str(root), + sha256=runtime_sha256, + identity=identity, + ) + + +def _run_helper(root: Path, config: LmEvalRuntimeConfig, workspace: Path): + workspace.mkdir(exist_ok=True) + env = os.environ.copy() + env.update( + { + "MAGPIE_LM_EVAL_RUNTIME_ROOT": str(root), + "MAGPIE_LM_EVAL_RUNTIME_SHA256": config.sha256, + "MAGPIE_LM_EVAL_RUNTIME_RECEIPT": str( + workspace / LM_EVAL_RECEIPT_FILENAME + ), + "MAGPIE_LM_EVAL_EXECUTION_MODE": "local", + "MAGPIE_LM_EVAL_REQUIRE_READONLY_MOUNT": "0", + } + ) + return subprocess.run( + [ + "bash", + "-c", + 'source "$1" && _install_lm_eval_deps', + "bash", + str(HELPER), + ], + env=env, + text=True, + capture_output=True, + timeout=30, + check=False, + ) + + +def _git(repo: Path, *args: str) -> str: + completed = subprocess.run( + ["git", "-C", str(repo), *args], + text=True, + capture_output=True, + check=True, + ) + return completed.stdout.strip() + + +def _inferencex_repo(tmp_path: Path) -> Path: + repo = tmp_path / "InferenceX" + (repo / "benchmarks").mkdir(parents=True) + (repo / "benchmarks/benchmark_lib.sh").write_text( + "run_benchmark_serving() { return 0; }\n", + encoding="utf-8", + ) + _git(repo, "init", "-q") + _git(repo, "config", "user.email", "test@example.com") + _git(repo, "config", "user.name", "Test User") + _git(repo, "add", ".") + _git(repo, "commit", "-qm", "initial") + return repo + + +def test_config_round_trip_preserves_nested_runtime(tmp_path): + _, runtime_config = _build_runtime(tmp_path) + config = BenchmarkConfig( + framework="vllm", + model="example/model", + lm_eval_runtime=runtime_config, + ) + + restored = BenchmarkConfig.from_dict(config.to_dict()) + + assert restored.lm_eval_runtime == runtime_config + assert restored.to_dict()["lm_eval_runtime"] == runtime_config.to_dict() + + +def test_validate_runtime_and_helper_emit_bound_evidence(tmp_path): + root, runtime_config = _build_runtime(tmp_path) + workspace = tmp_path / "workspace" + workspace.mkdir() + runtime = validate_lm_eval_runtime(runtime_config) + snapshot_runtime_manifest(runtime, workspace) + + completed = _run_helper(root, runtime_config, workspace) + + assert completed.returncode == 0, completed.stderr + evidence = collect_lm_eval_runtime_evidence( + workspace, + requested=True, + config=runtime_config, + execution_mode="local", + ) + assert evidence["schema"] == LM_EVAL_EVIDENCE_SCHEMA + assert evidence["status"] == "verified" + assert evidence["runtime_sha256"] == runtime_config.sha256 + assert evidence["identity"] == runtime_config.identity + assert evidence["mount_mode"] == "local" + assert evidence["manifest_artifact"]["path"] == LM_EVAL_MANIFEST_FILENAME + assert evidence["receipt_artifact"]["path"] == LM_EVAL_RECEIPT_FILENAME + + +def test_validate_runtime_rejects_byte_tampering(tmp_path): + root, runtime_config = _build_runtime(tmp_path) + target = root / "site-packages/lm_eval/__init__.py" + target.chmod(0o644) + target.write_text("VALUE = 'tampered'\n", encoding="utf-8") + target.chmod(0o444) + + with pytest.raises(ValueError, match="size mismatch|content digest mismatch"): + validate_lm_eval_runtime(runtime_config) + + +def test_validate_runtime_rejects_symlink_and_hardlink(tmp_path): + root, runtime_config = _build_runtime(tmp_path) + site_packages = root / "site-packages" + root.chmod(0o755) + site_packages.chmod(0o755) + source = site_packages / "lm_eval/__init__.py" + hardlink = site_packages / "hardlink.py" + os.link(source, hardlink) + hardlink.chmod(0o444) + site_packages.chmod(0o555) + root.chmod(0o555) + + with pytest.raises(ValueError, match="nlink=1|exactly match"): + validate_lm_eval_runtime(runtime_config) + + +def test_run_eval_without_runtime_fails_before_launch(tmp_path): + inferencex = tmp_path / "InferenceX" + (inferencex / "benchmarks").mkdir(parents=True) + config = BenchmarkConfig( + framework="vllm", + model="example/model", + run_mode="local", + run_kind="measurement", + envs={"TP": 1, "RUN_EVAL": "true"}, + profiler={ + "torch_profiler": {"enabled": False}, + "gpu_monitor": {"enabled": False}, + }, + gpu_selection={"auto": False}, + inferencex_path=str(inferencex), + ) + mode = BenchmarkMode(config, output_dir=str(tmp_path / "results")) + + result = mode.run(task_id="missing-evaluator") + + assert result.success is False + assert result.lm_eval_runtime_receipt["status"] == "invalid" + assert any("RUN_EVAL=true requires" in item for item in result.errors) + + +def test_run_eval_ray_fails_closed_before_remote_dispatch(tmp_path, monkeypatch): + config = BenchmarkConfig( + framework="vllm", + model="example/model", + run_mode="ray", + envs={"RUN_EVAL": "true"}, + ) + mode = BenchmarkMode(config, output_dir=str(tmp_path / "results")) + monkeypatch.setattr( + mode, + "_execute_ray_benchmark", + lambda: pytest.fail("Ray dispatch must not run"), + ) + + result = mode.run(task_id="ray-locked-evaluator") + + assert result.success is False + assert result.lm_eval_runtime_receipt["status"] == "unsupported" + assert any("Ray benchmark refused" in item for item in result.errors) + task, error = mode._build_ray_benchmark_task() + assert task is None + assert "unsupported in Ray mode" in error + + +def test_run_eval_rejects_native_or_custom_benchmark_script(tmp_path): + config = BenchmarkConfig( + framework="vllm", + model="example/model", + run_mode="local", + envs={"RUN_EVAL": "true"}, + ) + mode = BenchmarkMode(config, output_dir=str(tmp_path / "results")) + + with pytest.raises(RuntimeError, match="native and custom"): + mode._validate_lm_eval_benchmark_script( + "benchmarks/single_node/gptoss_fp8_mi355x.sh" + ) + mode._validate_lm_eval_benchmark_script("benchmarks/vllm_mi355x.sh") + + +def test_runtime_inferencex_identity_mismatch_fails_before_launch(tmp_path): + inferencex = _inferencex_repo(tmp_path) + _, runtime_config = _build_runtime(tmp_path) + config = BenchmarkConfig( + framework="vllm", + model="example/model", + run_mode="local", + run_kind="measurement", + envs={"TP": 1, "RUN_EVAL": "true"}, + profiler={"torch_profiler": {"enabled": False}}, + gpu_selection={"auto": False}, + inferencex_path=str(inferencex), + benchmark_script="vllm_mi355x.sh", + lm_eval_runtime=runtime_config, + ) + mode = BenchmarkMode(config, output_dir=str(tmp_path / "results")) + + result = mode.run(task_id="mismatched-inferencex") + + assert result.success is False + assert result.lm_eval_runtime_receipt["status"] == "invalid" + assert any("commit/tree does not match" in item for item in result.errors) + + +def test_docker_command_mounts_runtime_read_only(tmp_path, monkeypatch): + _, runtime_config = _build_runtime(tmp_path) + config = BenchmarkConfig( + framework="vllm", + model="example/model", + run_mode="docker", + run_kind="measurement", + envs={"TP": 1, "RUN_EVAL": "true"}, + profiler={"torch_profiler": {"enabled": False}}, + gpu_selection={"auto": False}, + inferencex_path=str(tmp_path / "InferenceX"), + benchmark_script="vllm_mi355x.sh", + lm_eval_runtime=runtime_config, + ) + (tmp_path / "InferenceX").mkdir() + mode = BenchmarkMode(config, output_dir=str(tmp_path / "results")) + mode._task_id = "read-only-runtime" + mode._lm_eval_runtime = validate_lm_eval_runtime(runtime_config) + monkeypatch.setattr( + "Magpie.modes.benchmark.benchmarker.detect_gpu", + lambda: (GPUVendor.UNKNOWN, ""), + ) + monkeypatch.setattr( + mode, + "_get_benchmark_script", + lambda runner_type: "benchmarks/vllm_mi355x.sh", + ) + + command = mode._build_docker_command( + "example/image@sha256:" + "a" * 64, + tmp_path / "workspace", + "mi355x", + ) + + mount = f"{Path(runtime_config.path).resolve()}:/opt/apex/lm-eval-runtime:ro" + assert mount in command + assert "MAGPIE_LM_EVAL_RUNTIME_ROOT=/opt/apex/lm-eval-runtime" in command + assert f"MAGPIE_LM_EVAL_RUNTIME_SHA256={runtime_config.sha256}" in command + assert "MAGPIE_LM_EVAL_EXECUTION_MODE=docker" in command + assert "MAGPIE_LM_EVAL_REQUIRE_READONLY_MOUNT=1" in command + + +def test_helper_contains_no_mutable_or_network_install_path(): + source = HELPER.read_text(encoding="utf-8") + + assert "_install_lm_eval_deps()" in source + for forbidden in ("pip install", "git+https://", "curl ", "wget "): + assert forbidden not in source + + +def test_helper_failure_terminates_ignoring_upstream_caller(): + env = os.environ.copy() + for name in ( + "MAGPIE_LM_EVAL_RUNTIME_ROOT", + "MAGPIE_LM_EVAL_RUNTIME_SHA256", + "MAGPIE_LM_EVAL_RUNTIME_RECEIPT", + "MAGPIE_LM_EVAL_EXECUTION_MODE", + "MAGPIE_LM_EVAL_REQUIRE_READONLY_MOUNT", + ): + env.pop(name, None) + + completed = subprocess.run( + [ + "bash", + "-c", + 'source "$1"; _install_lm_eval_deps; echo UNSAFE_CONTINUATION', + "bash", + str(HELPER), + ], + env=env, + text=True, + capture_output=True, + timeout=10, + check=False, + ) + + assert completed.returncode != 0 + assert "UNSAFE_CONTINUATION" not in completed.stdout + assert "refusing to run" in completed.stderr + + +def _base_image_is_local() -> bool: + if shutil.which("docker") is None: + return False + inspected = subprocess.run( + ["docker", "image", "inspect", BASE_IMAGE], + text=True, + capture_output=True, + timeout=10, + check=False, + ) + return inspected.returncode == 0 + + +@pytest.mark.skipif( + not _base_image_is_local(), + reason="pinned vLLM base image is not present locally", +) +def test_helper_import_smoke_is_offline_and_mount_is_read_only(tmp_path): + root, runtime_config = _build_runtime(tmp_path, python_abi="cpython-312") + workspace = tmp_path / "workspace" + workspace.mkdir() + helper_copy = tmp_path / HELPER.name + shutil.copy2(HELPER, helper_copy) + snapshot_runtime_manifest(validate_lm_eval_runtime(runtime_config), workspace) + + completed = subprocess.run( + [ + "docker", + "run", + "--rm", + "--network=none", + "-v", + f"{root}:/opt/apex/lm-eval-runtime:ro", + "-v", + f"{helper_copy}:/opt/apex/lm_eval_runtime.sh:ro", + "-v", + f"{workspace}:/workspace", + "-e", + "MAGPIE_LM_EVAL_RUNTIME_ROOT=/opt/apex/lm-eval-runtime", + "-e", + f"MAGPIE_LM_EVAL_RUNTIME_SHA256={runtime_config.sha256}", + "-e", + "MAGPIE_LM_EVAL_RUNTIME_RECEIPT=/workspace/lm_eval_runtime_receipt.json", + "-e", + "MAGPIE_LM_EVAL_EXECUTION_MODE=docker", + "-e", + "MAGPIE_LM_EVAL_REQUIRE_READONLY_MOUNT=1", + "--entrypoint", + "bash", + BASE_IMAGE, + "-c", + "source /opt/apex/lm_eval_runtime.sh && _install_lm_eval_deps", + ], + text=True, + capture_output=True, + timeout=60, + check=False, + ) + + assert completed.returncode == 0, completed.stderr + evidence = collect_lm_eval_runtime_evidence( + workspace, + requested=True, + config=runtime_config, + execution_mode="docker", + ) + assert evidence["status"] == "verified" + assert evidence["mount_mode"] == "read_only" + + +def test_tampered_runtime_receipt_is_not_reportable(tmp_path): + root, runtime_config = _build_runtime(tmp_path) + workspace = tmp_path / "workspace" + workspace.mkdir() + snapshot_runtime_manifest(validate_lm_eval_runtime(runtime_config), workspace) + completed = _run_helper(root, runtime_config, workspace) + assert completed.returncode == 0, completed.stderr + receipt_path = workspace / LM_EVAL_RECEIPT_FILENAME + receipt = json.loads(receipt_path.read_text(encoding="utf-8")) + receipt["runtime_sha256"] = "0" * 64 + receipt_path.write_text(json.dumps(receipt), encoding="utf-8") + + evidence = collect_lm_eval_runtime_evidence( + workspace, + requested=True, + config=runtime_config, + execution_mode="local", + ) + + assert evidence["verified"] is False + assert evidence["status"] == "invalid" + assert "receipt digest" in evidence["errors"][0] + + +def test_benchmark_result_serializes_runtime_evidence(): + evidence = { + "schema": LM_EVAL_EVIDENCE_SCHEMA, + "status": "verified", + "runtime_sha256": "a" * 64, + "mount_mode": "read_only", + } + result = BenchmarkResult(lm_eval_runtime_receipt=evidence) + + assert result.to_dict()["lm_eval_runtime_receipt"] == evidence + assert "lm-eval runtime evidence:" in result.get_summary() diff --git a/tests/test_model_revision_contract.py b/tests/test_model_revision_contract.py new file mode 100644 index 0000000..f6962dc --- /dev/null +++ b/tests/test_model_revision_contract.py @@ -0,0 +1,388 @@ +import json +import os +import shutil +import subprocess +from pathlib import Path + +import pytest + +from Magpie.modes.benchmark.benchmarker import BenchmarkMode +from Magpie.modes.benchmark.config import BenchmarkConfig +from Magpie.modes.benchmark.model_revision import ( + MODEL_REVISION_EVIDENCE_SCHEMA, + MODEL_REVISION_RECEIPT_SCHEMA, + collect_model_revision_evidence, +) +from Magpie.modes.benchmark.result import BenchmarkResult + + +REPO_ROOT = Path(__file__).resolve().parents[1] +VLLM_MI355X_SCRIPT = REPO_ROOT / "Magpie/scripts/benchmark/vllm_mi355x.sh" +LM_EVAL_RUNTIME_HELPER = REPO_ROOT / "Magpie/scripts/benchmark/lm_eval_runtime.sh" +MODEL = "example-org/example-model" +REVISION = "a" * 40 + + +def _write_executable(path: Path, content: str) -> None: + path.write_text(content, encoding="utf-8") + path.chmod(0o755) + + +def _script_sandbox(tmp_path: Path) -> tuple[Path, dict[str, str]]: + scripts = tmp_path / "benchmarks" + commands = tmp_path / "bin" + workspace = tmp_path / "workspace" + snapshots = tmp_path / "snapshots" + scripts.mkdir() + commands.mkdir() + workspace.mkdir() + snapshots.mkdir() + script = scripts / VLLM_MI355X_SCRIPT.name + shutil.copy2(VLLM_MI355X_SCRIPT, script) + shutil.copy2(LM_EVAL_RUNTIME_HELPER, scripts / LM_EVAL_RUNTIME_HELPER.name) + + (scripts / "benchmark_lib.sh").write_text( + """ +check_env_vars() { + local name + for name in "$@"; do + [[ -n "${!name:-}" ]] || return 1 + done +} +wait_for_server_ready() { + local attempt + for attempt in $(seq 1 100); do + [[ -s "${VLLM_ARGS_LOG:-/nonexistent}" ]] && return 0 + sleep 0.01 + done + return 1 +} +run_benchmark_serving() { return 0; } +magpie_mark_lm_eval_start() { return 0; } +run_eval() { + printf '%s|%s\n' "${EVAL_CONCURRENT_REQUESTS:-}" "$*" > "$EVAL_LOG" +} +magpie_preserve_lm_eval_artifacts() { return 0; } +append_lm_eval_summary() { return 0; } +""", + encoding="utf-8", + ) + (scripts / "server_cleanup.sh").write_text( + "magpie_stop_benchmark_server_stack() { return 0; }\n", + encoding="utf-8", + ) + + _write_executable( + commands / "hf", + """#!/usr/bin/env bash +printf '%s\n' "$@" > "$HF_ARGS_LOG" +[[ "${HF_FAIL:-0}" == "1" ]] && exit 9 +resolved="${HF_RESOLVED_REVISION:-$MODEL_REVISION}" +snapshot="$HF_SNAPSHOT_ROOT/$resolved" +mkdir -p "$snapshot" +printf '%s\n' "$snapshot" +""", + ) + _write_executable( + commands / "vllm", + """#!/usr/bin/env bash +printf '%s\n' "$@" > "$VLLM_ARGS_LOG" +""", + ) + _write_executable( + commands / "setsid", + """#!/usr/bin/env bash +exec "$@" +""", + ) + _write_executable(commands / "rocm-smi", "#!/usr/bin/env bash\nexit 1\n") + + env = os.environ.copy() + env.update( + { + "PATH": f"{commands}:{env['PATH']}", + "MODEL": MODEL, + "MODEL_REVISION": REVISION, + "TP": "1", + "RESULT_DIR": str(workspace), + "SERVER_LOG": str(workspace / "server.log"), + "MAGPIE_RUN_PHASE": "server", + "MAGPIE_SERVER_PID_FILE": str(workspace / "server.pid"), + "HF_ARGS_LOG": str(tmp_path / "hf.args"), + "HF_SNAPSHOT_ROOT": str(snapshots), + "VLLM_ARGS_LOG": str(tmp_path / "vllm.args"), + "EVAL_LOG": str(tmp_path / "eval.args"), + "SLURM_JOB_ID": "", + "ROCR_VISIBLE_DEVICES": "", + "HIP_VISIBLE_DEVICES": "", + } + ) + return script, env + + +def _run_script(script: Path, env: dict[str, str]) -> subprocess.CompletedProcess[str]: + return subprocess.run( + ["bash", str(script)], + env=env, + text=True, + capture_output=True, + timeout=10, + check=False, + ) + + +def _receipt_payload(**overrides): + payload = { + "schema": MODEL_REVISION_RECEIPT_SCHEMA, + "model": MODEL, + "requested_revision": REVISION, + "resolved_revision": REVISION, + "snapshot_path": f"/cache/snapshots/{REVISION}", + "verified": True, + } + payload.update(overrides) + return payload + + +def test_vllm_mi355x_binds_exact_revision_and_writes_receipt(tmp_path): + script, env = _script_sandbox(tmp_path) + + completed = _run_script(script, env) + + assert completed.returncode == 0, completed.stderr + assert (tmp_path / "hf.args").read_text(encoding="utf-8").splitlines() == [ + "download", + MODEL, + "--revision", + REVISION, + "--format", + "quiet", + ] + vllm_args = (tmp_path / "vllm.args").read_text(encoding="utf-8").splitlines() + assert vllm_args[:2] == ["serve", MODEL] + revision_index = vllm_args.index("--revision") + assert vllm_args[revision_index + 1] == REVISION + + receipt = json.loads( + (tmp_path / "workspace/model_revision_receipt.json").read_text( + encoding="utf-8" + ) + ) + assert receipt == _receipt_payload( + snapshot_path=str((tmp_path / "snapshots" / REVISION).resolve()) + ) + + +@pytest.mark.parametrize( + "revision", + ["a" * 39, "A" * 40, "a" * 40 + "0", "main"], +) +def test_vllm_mi355x_rejects_non_exact_revision_before_download( + tmp_path, revision +): + script, env = _script_sandbox(tmp_path) + env["MODEL_REVISION"] = revision + + completed = _run_script(script, env) + + assert completed.returncode == 4 + assert "exact lowercase 40-hex" in completed.stderr + assert not (tmp_path / "hf.args").exists() + assert not (tmp_path / "vllm.args").exists() + + +def test_vllm_mi355x_fails_closed_when_download_fails(tmp_path): + script, env = _script_sandbox(tmp_path) + env["HF_FAIL"] = "1" + + completed = _run_script(script, env) + + assert completed.returncode == 9 + assert not (tmp_path / "workspace/model_revision_receipt.json").exists() + assert not (tmp_path / "vllm.args").exists() + + +def test_vllm_mi355x_fails_closed_when_resolved_snapshot_differs(tmp_path): + script, env = _script_sandbox(tmp_path) + env["HF_RESOLVED_REVISION"] = "b" * 40 + + completed = _run_script(script, env) + + assert completed.returncode != 0 + assert "does not match MODEL_REVISION" in completed.stderr + assert not (tmp_path / "workspace/model_revision_receipt.json").exists() + assert not (tmp_path / "vllm.args").exists() + + +def test_vllm_mi355x_passes_eval_concurrency_through_environment(tmp_path): + script, env = _script_sandbox(tmp_path) + env.update( + { + "MAGPIE_RUN_PHASE": "client", + "CONC": "16", + "ISL": "128", + "OSL": "64", + "RANDOM_RANGE_RATIO": "1", + "RESULT_FILENAME": "inferencex_result", + "RUN_EVAL": "true", + } + ) + env.pop("MODEL_REVISION") + + completed = _run_script(script, env) + + assert completed.returncode == 0, completed.stderr + assert (tmp_path / "eval.args").read_text(encoding="utf-8").strip() == ( + "16|--framework lm-eval --port 8888" + ) + + +def test_model_revision_evidence_verifies_workspace_receipt(tmp_path): + receipt_path = tmp_path / "model_revision_receipt.json" + receipt_path.write_text(json.dumps(_receipt_payload()), encoding="utf-8") + + evidence = collect_model_revision_evidence( + tmp_path, + model=MODEL, + requested_revision=REVISION, + ) + + assert evidence["schema"] == MODEL_REVISION_EVIDENCE_SCHEMA + assert evidence["status"] == "verified" + assert evidence["verified"] is True + assert evidence["resolved_revision"] == REVISION + assert evidence["receipt_artifact"]["path"] == receipt_path.name + assert len(evidence["receipt_artifact"]["sha256"]) == 64 + + +def test_model_revision_evidence_missing_requested_receipt_fails_closed(tmp_path): + evidence = collect_model_revision_evidence( + tmp_path, + model=MODEL, + requested_revision=REVISION, + ) + + assert evidence["status"] == "missing" + assert evidence["verified"] is False + assert evidence["errors"] + + +@pytest.mark.parametrize( + "overrides", + [ + {"model": "wrong/model"}, + {"resolved_revision": "b" * 40}, + {"snapshot_path": "/cache/snapshots/not-the-revision"}, + {"verified": False}, + {"unexpected": "field"}, + ], +) +def test_model_revision_evidence_rejects_tampered_receipt(tmp_path, overrides): + (tmp_path / "model_revision_receipt.json").write_text( + json.dumps(_receipt_payload(**overrides)), + encoding="utf-8", + ) + + evidence = collect_model_revision_evidence( + tmp_path, + model=MODEL, + requested_revision=REVISION, + ) + + assert evidence["status"] == "invalid" + assert evidence["verified"] is False + assert evidence["errors"] + + +def test_benchmark_result_serializes_model_revision_evidence(): + evidence = collect_model_revision_evidence( + Path("/nonexistent"), + model=MODEL, + requested_revision=None, + ) + result = BenchmarkResult(model_revision_receipt=evidence) + + report = result.to_dict() + + assert report["model_revision_receipt"]["status"] == "not_requested" + assert "Model revision evidence:" in result.get_summary() + + +@pytest.mark.parametrize( + ("write_receipt", "expected_success", "expected_status"), + [(True, True, "verified"), (False, False, "missing")], +) +def test_benchmark_mode_report_enforces_requested_revision_receipt( + tmp_path, + monkeypatch, + write_receipt, + expected_success, + expected_status, +): + inferencex = tmp_path / "InferenceX" + (inferencex / "benchmarks").mkdir(parents=True) + config = BenchmarkConfig( + framework="vllm", + model=MODEL, + run_mode="local", + run_kind="measurement", + envs={"MODEL_REVISION": REVISION, "TP": 1}, + profiler={ + "torch_profiler": {"enabled": False}, + "gpu_monitor": {"enabled": False}, + }, + gpu_selection={"auto": False}, + inferencex_path=str(inferencex), + ) + mode = BenchmarkMode(config, output_dir=str(tmp_path / "results")) + + monkeypatch.setattr( + "Magpie.modes.benchmark.benchmarker.ensure_inferencex_available", + lambda path: str(inferencex), + ) + monkeypatch.setattr(mode, "_prepare_benchmark_scripts", lambda: None) + monkeypatch.setattr(mode, "_get_runner_type", lambda: "mi355x") + monkeypatch.setattr( + mode, + "_get_benchmark_script", + lambda runner_type: "benchmarks/vllm_mi355x.sh", + ) + monkeypatch.setattr( + mode, + "_build_local_command", + lambda workspace, runner_type: (["true"], {}), + ) + monkeypatch.setattr(mode, "_cleanup_server_processes", lambda framework: None) + + def execute(local_cmd, local_env, workspace): + (workspace / "inferencex_result.json").write_text( + json.dumps( + { + "request_throughput": 1.0, + "output_throughput": 10.0, + "completed": 1, + } + ), + encoding="utf-8", + ) + if write_receipt: + (workspace / "model_revision_receipt.json").write_text( + json.dumps(_receipt_payload()), + encoding="utf-8", + ) + return BenchmarkResult(success=True), "", "" + + monkeypatch.setattr(mode, "_execute_local_benchmark", execute) + + result = mode.run(task_id="revision-contract") + + assert result.success is expected_success + assert result.model_revision_receipt["status"] == expected_status + report_path = Path(result.workspace_dir) / "benchmark_report.json" + report = json.loads(report_path.read_text(encoding="utf-8")) + assert report["model_revision_receipt"]["status"] == expected_status + if not expected_success: + assert any( + "Model revision evidence gate failed" in error + for error in result.errors + ) From 5a7fb05c01019c2a2c6f2ecff671c45be22f82d9 Mon Sep 17 00:00:00 2001 From: Vincent Ouyang Date: Fri, 7 Aug 2026 12:52:02 +0000 Subject: [PATCH 03/11] Harden TraceLens vLLM image derivation --- Magpie/modes/benchmark/tracelens_runtime.py | 135 +++- .../modes/benchmark/tracelens_vllm_image.py | 753 ++++++++++++++++++ docs/how-to/benchmarking/profiling-options.md | 27 +- docs/reference/benchmark-config.md | 6 + docs/reference/release-notes.md | 3 + tests/test_benchmark_support.py | 82 +- tests/test_tracelens_vllm_image.py | 287 +++++++ 7 files changed, 1259 insertions(+), 34 deletions(-) create mode 100644 Magpie/modes/benchmark/tracelens_vllm_image.py create mode 100644 tests/test_tracelens_vllm_image.py diff --git a/Magpie/modes/benchmark/tracelens_runtime.py b/Magpie/modes/benchmark/tracelens_runtime.py index bdd54c6..1dd7529 100644 --- a/Magpie/modes/benchmark/tracelens_runtime.py +++ b/Magpie/modes/benchmark/tracelens_runtime.py @@ -7,8 +7,8 @@ TraceLens-ready benchmark runtime image preparation. TraceLens inference mode needs framework runtime patches for some vLLM/SGLang -versions. This module uses the public TraceLens workflow build scripts to derive -patched Docker images from supported official runtime images. +versions. SGLang uses the public TraceLens workflow build script; vLLM uses a +minimal pinned-wheel builder with content-verified image reuse. """ from __future__ import annotations @@ -25,6 +25,12 @@ from typing import Any, Dict, Optional from .config import BenchmarkConfig +from .tracelens_vllm_image import ( + VllmTraceLensIdentity, + build_vllm_tracelens_image, + resolve_vllm_tracelens_identity, + validate_vllm_tracelens_image, +) logger = logging.getLogger(__name__) @@ -680,7 +686,9 @@ def prepare_tracelens_runtime_image( public_image = base_image patch_version: Optional[str] = None installed_version: Optional[str] = None + grpcio_version: Optional[str] = None tracelens_repo: Optional[Path] = None + vllm_identity: Optional[VllmTraceLensIdentity] = None if not base_is_ready: tracelens_repo = resolve_tracelens_repo_path(tl_config.tracelens_repo_path) @@ -713,6 +721,25 @@ def prepare_tracelens_runtime_image( raise RuntimeError( _vllm_patch_error(base_image, tracelens_repo, installed_version) ) + if installed_version is None: + raise RuntimeError( + "TraceLens vLLM auto patching requires the exact installed " + f"vLLM version from base image {base_image!r}." + ) + grpcio_version = docker_image_package_version(base_image, "grpcio") + if grpcio_version is None: + raise RuntimeError( + "TraceLens vLLM auto patching requires the exact installed " + f"grpcio version from base image {base_image!r}." + ) + vllm_identity = resolve_vllm_tracelens_identity( + base_image=base_image, + vllm_version=installed_version, + grpcio_version=grpcio_version, + tracelens_repo=tracelens_repo, + patch_version=patch_version, + ) + result.update(vllm_identity.metadata()) else: patch_version = "unknown" @@ -749,45 +776,101 @@ def prepare_tracelens_runtime_image( "package" if installed_version else "image" ) - if docker_image_exists(derived_image) and not tl_config.runtime_patch_force_rebuild: - result["reason"] = "derived image already exists" - return result + public_reusable = False + if not tl_config.runtime_patch_force_rebuild: + if vllm_identity is not None and docker_image_exists(public_image): + validation = validate_vllm_tracelens_image( + public_image, + vllm_identity, + ) + result["public_runtime_validation"] = validation + public_reusable = bool(validation.get("valid")) + if public_reusable: + result["public_runtime_image_id"] = validation.get("image_id") + result["dependency_wheels"] = validation.get( + "dependency_wheels", [] + ) + result["dependency_wheel_manifest_sha256"] = validation.get( + "dependency_wheel_manifest_sha256" + ) + else: + result["stale_image_rejected"] = True + result["stale_image_rejection_reason"] = validation.get("reason") + + if vllm_identity is not None and public_reusable: + if not extension: + result["reason"] = "validated derived image already exists" + return result + if docker_image_exists(derived_image): + result["reason"] = "derived extension image already exists" + return result + elif vllm_identity is None and docker_image_exists(derived_image): + # Preserve the existing SGLang and ready-image extension behavior. + result["reason"] = "derived image already exists" + return result public_built = False if not base_is_ready and ( tl_config.runtime_patch_force_rebuild or not docker_image_exists(public_image) + or (vllm_identity is not None and not public_reusable) ): assert tracelens_repo is not None assert patch_version is not None - cmd = _build_command( - config=config, - base_image=base_image, - runner_type=runner_type, - derived_image=public_image, - tracelens_repo=tracelens_repo, - patch_version=patch_version, - ) - result["command"] = cmd logger.info( "Building TraceLens-ready %s image from %s as %s", config.framework, base_image, public_image, ) - proc = subprocess.run( - cmd, - cwd=str(tracelens_repo), - text=True, - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - ) - if proc.returncode != 0: - tail = (proc.stdout or "")[-4000:] - raise RuntimeError( - "TraceLens runtime image build failed with exit code " - f"{proc.returncode}. Command: {' '.join(cmd)}\n{tail}" + if config.framework == "vllm": + assert vllm_identity is not None + build_metadata = build_vllm_tracelens_image( + identity=vllm_identity, + tracelens_repo=tracelens_repo, + derived_image=public_image, ) + result.update( + { + "command": build_metadata["command"], + "source_wheel_command": build_metadata[ + "source_wheel_command" + ], + "requirements_download_command": build_metadata[ + "requirements_download_command" + ], + "public_runtime_image_id": build_metadata["image_id"], + "public_runtime_labels": build_metadata["image_labels"], + "dependency_wheels": build_metadata["dependency_wheels"], + "dependency_wheel_manifest_sha256": build_metadata[ + "dependency_wheel_manifest_sha256" + ], + "public_runtime_validation": build_metadata["validation"], + } + ) + else: + cmd = _build_command( + config=config, + base_image=base_image, + runner_type=runner_type, + derived_image=public_image, + tracelens_repo=tracelens_repo, + patch_version=patch_version, + ) + result["command"] = cmd + proc = subprocess.run( + cmd, + cwd=str(tracelens_repo), + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + ) + if proc.returncode != 0: + tail = (proc.stdout or "")[-4000:] + raise RuntimeError( + "TraceLens runtime image build failed with exit code " + f"{proc.returncode}. Command: {' '.join(cmd)}\n{tail}" + ) public_built = True extension_built = False diff --git a/Magpie/modes/benchmark/tracelens_vllm_image.py b/Magpie/modes/benchmark/tracelens_vllm_image.py new file mode 100644 index 0000000..329d7d4 --- /dev/null +++ b/Magpie/modes/benchmark/tracelens_vllm_image.py @@ -0,0 +1,753 @@ +############################################################################### +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### +"""Deterministic, minimal TraceLens image derivation for vLLM. + +The upstream TraceLens installer intentionally serves many workflows and pulls +large, loosely resolved dependencies (including xprof). Magpie's vLLM +inference path needs a much smaller surface: the trace splitter, the CSV report +generator, Matplotlib (eagerly imported by the architecture helper), and the +vLLM instrumentation patch. This module builds exactly that surface without +changing packages already present in the vLLM base image. +""" + +from __future__ import annotations + +import hashlib +import json +import os +import re +import shutil +import subprocess +import tarfile +import tempfile +from dataclasses import dataclass, field +from pathlib import Path, PurePosixPath +from typing import Any, Dict, Mapping, Optional, Sequence + + +VLLM_TRACELENS_REQUIREMENTS = ( + ("contourpy", "1.3.3"), + ("cycler", "0.12.1"), + ("fonttools", "4.63.0"), + ("kiwisolver", "1.5.0"), + ("matplotlib", "3.11.1"), + ("pyparsing", "3.3.2"), +) +VLLM_TRACELENS_FORBIDDEN = ("xprof", "gcsfs", "grpcio-status") +VLLM_TRACELENS_SCHEMA = "magpie.tracelens-vllm-runtime/v1" +LABEL_PREFIX = "io.magpie.tracelens" +LABEL_SCHEMA = f"{LABEL_PREFIX}.schema" +LABEL_BASE_ID = f"{LABEL_PREFIX}.base-image-id" +LABEL_BASE_LOCATOR = f"{LABEL_PREFIX}.base-image-locator" +LABEL_VLLM_VERSION = f"{LABEL_PREFIX}.vllm-version" +LABEL_GRPCIO_VERSION = f"{LABEL_PREFIX}.grpcio-version" +LABEL_SOURCE_COMMIT = f"{LABEL_PREFIX}.source-commit" +LABEL_SOURCE_TREE = f"{LABEL_PREFIX}.source-tree" +LABEL_PATCH_VERSION = f"{LABEL_PREFIX}.patch-version" +LABEL_PATCH_SHA256 = f"{LABEL_PREFIX}.patch-sha256" +LABEL_WHEEL_MANIFEST = f"{LABEL_PREFIX}.wheel-manifest" +LABEL_WHEEL_MANIFEST_SHA256 = f"{LABEL_PREFIX}.wheel-manifest-sha256" +LABEL_DEPENDENCY_POLICY = f"{LABEL_PREFIX}.dependency-policy" +DEPENDENCY_POLICY = "minimal-pinned-wheels-no-deps" + +_SOURCE_PATHS = ( + "LICENSE", + "MANIFEST.in", + "README.md", + "setup.py", + "TraceLens", +) +_HASH_RE = re.compile(r"^[0-9a-f]{64}$") + + +@dataclass(frozen=True) +class VllmTraceLensIdentity: + """Immutable inputs that define a derived TraceLens vLLM image.""" + + base_image: str + base_image_id: str + base_image_locator: str + vllm_version: str + grpcio_version: str + source_commit: str + source_tree: str + patch_version: str + patch_path: str + patch_sha256: str + patch_bytes: bytes = field(repr=False) + + def labels(self) -> Dict[str, str]: + return { + LABEL_SCHEMA: VLLM_TRACELENS_SCHEMA, + LABEL_BASE_ID: self.base_image_id, + LABEL_BASE_LOCATOR: self.base_image_locator, + LABEL_VLLM_VERSION: self.vllm_version, + LABEL_GRPCIO_VERSION: self.grpcio_version, + LABEL_SOURCE_COMMIT: self.source_commit, + LABEL_SOURCE_TREE: self.source_tree, + LABEL_PATCH_VERSION: self.patch_version, + LABEL_PATCH_SHA256: self.patch_sha256, + LABEL_DEPENDENCY_POLICY: DEPENDENCY_POLICY, + } + + def metadata(self) -> Dict[str, Any]: + return { + "base_image_id": self.base_image_id, + "base_image_locator": self.base_image_locator, + "runtime_package_version": self.vllm_version, + "base_grpcio_version": self.grpcio_version, + "tracelens_source_commit": self.source_commit, + "tracelens_source_tree": self.source_tree, + "tracelens_patch_path": self.patch_path, + "tracelens_patch_sha256": self.patch_sha256, + "dependency_policy": DEPENDENCY_POLICY, + } + + +def _sha256_bytes(value: bytes) -> str: + return hashlib.sha256(value).hexdigest() + + +def _sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _completed_output(proc: subprocess.CompletedProcess[Any]) -> str: + stdout = proc.stdout or "" + stderr = proc.stderr or "" + if isinstance(stdout, bytes): + stdout = stdout.decode("utf-8", errors="replace") + if isinstance(stderr, bytes): + stderr = stderr.decode("utf-8", errors="replace") + return (stderr or stdout).strip() + + +def _git_text(repo: Path, *args: str) -> str: + proc = subprocess.run( + ["git", "-C", str(repo), *args], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + check=False, + ) + if proc.returncode != 0: + raise RuntimeError( + f"Could not inspect TraceLens source identity: {_completed_output(proc)}" + ) + return (proc.stdout or "").strip() + + +def _git_bytes(repo: Path, *args: str) -> bytes: + proc = subprocess.run( + ["git", "-C", str(repo), *args], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + check=False, + ) + if proc.returncode != 0: + raise RuntimeError( + f"Could not read committed TraceLens content: {_completed_output(proc)}" + ) + return bytes(proc.stdout or b"") + + +def docker_image_record(image: str) -> Optional[Dict[str, Any]]: + """Return Docker's inspect record for ``image``, or None when unavailable.""" + try: + proc = subprocess.run( + ["docker", "image", "inspect", image], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + timeout=120, + check=False, + ) + except (OSError, subprocess.TimeoutExpired): + return None + if proc.returncode != 0: + return None + try: + records = json.loads(proc.stdout or "[]") + except json.JSONDecodeError: + return None + if not isinstance(records, list) or len(records) != 1: + return None + return records[0] if isinstance(records[0], dict) else None + + +def docker_image_id(image: str) -> Optional[str]: + record = docker_image_record(image) + image_id = record.get("Id") if record else None + return image_id if isinstance(image_id, str) and image_id else None + + +def resolve_vllm_tracelens_identity( + *, + base_image: str, + vllm_version: str, + grpcio_version: str, + tracelens_repo: Path, + patch_version: str, +) -> VllmTraceLensIdentity: + """Resolve source, patch, and immutable base-image identity.""" + base_record = docker_image_record(base_image) + base_id = base_record.get("Id") if base_record else None + if not isinstance(base_id, str) or not base_id: + raise RuntimeError(f"Could not resolve Docker image ID for {base_image!r}") + repo_digests = base_record.get("RepoDigests") or [] + base_locator = ( + repo_digests[0] + if isinstance(repo_digests, list) + and repo_digests + and isinstance(repo_digests[0], str) + else base_image + ) + + source_commit = _git_text(tracelens_repo, "rev-parse", "HEAD") + source_tree = _git_text(tracelens_repo, "rev-parse", "HEAD^{tree}") + minor = patch_version.removeprefix("v") + if not minor.isdigit(): + raise RuntimeError(f"Invalid TraceLens vLLM patch version: {patch_version!r}") + patch_path = ( + "examples/custom_workflows/inference_analysis/vllm_patches/" + f"config_vllm_v0.{minor}.0.patch" + ) + patch_bytes = _git_bytes(tracelens_repo, "show", f"{source_commit}:{patch_path}") + if not patch_bytes: + raise RuntimeError(f"Committed TraceLens patch is empty: {patch_path}") + + return VllmTraceLensIdentity( + base_image=base_image, + base_image_id=base_id, + base_image_locator=base_locator, + vllm_version=vllm_version, + grpcio_version=grpcio_version, + source_commit=source_commit, + source_tree=source_tree, + patch_version=patch_version, + patch_path=patch_path, + patch_sha256=_sha256_bytes(patch_bytes), + patch_bytes=patch_bytes, + ) + + +def _canonical_json(value: Any) -> str: + return json.dumps(value, sort_keys=True, separators=(",", ":")) + + +def _wheel_distribution(path: Path) -> str: + return path.name.split("-", 1)[0].replace("_", "-").lower() + + +def _wheel_manifest(wheelhouse: Path) -> list[Dict[str, str]]: + expected_versions = dict(VLLM_TRACELENS_REQUIREMENTS) + records = [] + for wheel in sorted(wheelhouse.glob("*.whl")): + distribution = _wheel_distribution(wheel) + version = expected_versions.get(distribution, "source-commit") + records.append( + { + "distribution": distribution, + "filename": wheel.name, + "version": version, + "sha256": _sha256_file(wheel), + } + ) + expected = set(expected_versions) | {"tracelens"} + found = {record["distribution"] for record in records} + if found != expected: + raise RuntimeError( + "TraceLens wheelhouse is incomplete or contains unexpected wheels: " + f"expected={sorted(expected)}, found={sorted(found)}" + ) + return records + + +def _safe_extract_tar(archive_path: Path, destination: Path) -> None: + with tarfile.open(archive_path, "r") as archive: + for member in archive.getmembers(): + path = PurePosixPath(member.name) + if path.is_absolute() or ".." in path.parts: + raise RuntimeError(f"Unsafe path in git archive: {member.name!r}") + if member.issym() or member.islnk(): + raise RuntimeError( + f"Links are not allowed in the TraceLens build archive: " + f"{member.name!r}" + ) + archive.extractall(destination) + + +def _stage_committed_source( + identity: VllmTraceLensIdentity, + tracelens_repo: Path, + destination: Path, +) -> None: + archive_path = destination.parent / "tracelens-source.tar" + cmd = [ + "git", + "-C", + str(tracelens_repo), + "archive", + "--format=tar", + "--output", + str(archive_path), + identity.source_commit, + "--", + *_SOURCE_PATHS, + ] + proc = subprocess.run( + cmd, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + check=False, + ) + if proc.returncode != 0: + raise RuntimeError( + f"Could not stage committed TraceLens source: {_completed_output(proc)}" + ) + destination.mkdir(parents=True, exist_ok=True) + _safe_extract_tar(archive_path, destination) + + +def _mounted(path: Path, container_path: str, read_only: bool = False) -> str: + suffix = ":ro" if read_only else "" + return f"{path.resolve()}:{container_path}{suffix}" + + +def _build_source_wheel( + identity: VllmTraceLensIdentity, + source_dir: Path, + wheelhouse: Path, +) -> list[str]: + cmd = [ + "docker", + "run", + "--rm", + "--network", + "none", + "--user", + f"{os.getuid()}:{os.getgid()}", + "-e", + "HOME=/tmp", + "--entrypoint", + "python3", + "-v", + _mounted(source_dir, "/src"), + "-v", + _mounted(wheelhouse, "/wheels"), + identity.base_image_id, + "-m", + "pip", + "wheel", + "--disable-pip-version-check", + "--no-deps", + "--no-build-isolation", + "--wheel-dir", + "/wheels", + "/src", + ] + proc = subprocess.run( + cmd, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + check=False, + ) + if proc.returncode != 0: + raise RuntimeError( + "Could not build the pinned TraceLens wheel in the base runtime: " + f"{(proc.stdout or '')[-4000:]}" + ) + return cmd + + +def _download_requirement_wheels( + identity: VllmTraceLensIdentity, + wheelhouse: Path, +) -> list[str]: + requirements = [f"{name}=={version}" for name, version in VLLM_TRACELENS_REQUIREMENTS] + cmd = [ + "docker", + "run", + "--rm", + "--entrypoint", + "python3", + "-v", + _mounted(wheelhouse, "/wheels"), + identity.base_image_id, + "-m", + "pip", + "download", + "--disable-pip-version-check", + "--dest", + "/wheels", + "--only-binary=:all:", + "--no-deps", + *requirements, + ] + proc = subprocess.run( + cmd, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + check=False, + ) + if proc.returncode != 0: + raise RuntimeError( + "Could not download pinned TraceLens diagnostic wheels: " + f"{(proc.stdout or '')[-4000:]}" + ) + return cmd + + +def _dockerfile_labels(labels: Mapping[str, str]) -> str: + return "\n".join( + f"LABEL {key}={json.dumps(value)}" for key, value in sorted(labels.items()) + ) + + +def _verification_script() -> str: + requirements = _canonical_json(dict(VLLM_TRACELENS_REQUIREMENTS)) + forbidden = _canonical_json(list(VLLM_TRACELENS_FORBIDDEN)) + return f"""\ +import importlib.metadata as metadata +import json +import pathlib +import py_compile + +expected = json.loads({requirements!r}) +forbidden = json.loads({forbidden!r}) +actual = {{name: metadata.version(name) for name in expected}} +assert actual == expected, (actual, expected) +for name in forbidden: + try: + metadata.version(name) + except metadata.PackageNotFoundError: + continue + raise AssertionError(f"forbidden package installed: {{name}}") + +import TraceLens +from TraceLens.Agent.Analysis.utils.arch_utils import list_platforms +from TraceLens.Reporting import generate_perf_report_pytorch_inference +from TraceLens.TraceUtils import split_inference_trace_annotation +import vllm +from vllm.config.profiler import ProfilerConfig + +fields = ProfilerConfig.__dataclass_fields__ +assert "capture_torch_profiler_dir" in fields +assert "detailed_trace_annotation" in fields +root = pathlib.Path(vllm.__file__).resolve().parent +files = [ + root / "config/profiler.py", + root / "v1/worker/gpu_model_runner.py", + root / "v1/worker/gpu_worker.py", +] +for path in files: + py_compile.compile(str(path), doraise=True) +runner = files[1].read_text(encoding="utf-8") +worker = files[2].read_text(encoding="utf-8") +assert "capture_torch_profiler_dir" in runner and "profiler=profiler" in runner +assert "detailed_trace_annotation" in worker and "c_sqsk" in worker +print(json.dumps({{ + "grpcio_version": metadata.version("grpcio"), + "platforms": list_platforms(), + "tracelens_version": metadata.version("TraceLens"), + "vllm_version": metadata.version("vllm"), +}}, sort_keys=True)) +""" + + +def _write_build_context( + context: Path, + identity: VllmTraceLensIdentity, + wheel_manifest: Sequence[Mapping[str, str]], +) -> Dict[str, str]: + wheel_manifest_json = _canonical_json(list(wheel_manifest)) + labels = identity.labels() + labels.update( + { + LABEL_WHEEL_MANIFEST: wheel_manifest_json, + LABEL_WHEEL_MANIFEST_SHA256: _sha256_bytes( + wheel_manifest_json.encode("utf-8") + ), + } + ) + identity_document = { + "schema": VLLM_TRACELENS_SCHEMA, + **identity.metadata(), + "wheels": list(wheel_manifest), + } + (context / "identity.json").write_text( + json.dumps(identity_document, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + (context / "verify.py").write_text(_verification_script(), encoding="utf-8") + (context / "patch.diff").write_bytes(identity.patch_bytes) + dockerfile = f"""\ +FROM {identity.base_image_locator} +{_dockerfile_labels(labels)} +COPY wheels/ /tmp/tracelens-wheels/ +COPY patch.diff /tmp/tracelens-vllm.patch +COPY verify.py /tmp/verify-tracelens-runtime.py +COPY identity.json /opt/magpie/tracelens-runtime-identity.json +RUN python3 -m pip install --disable-pip-version-check --no-index --no-deps \\ + /tmp/tracelens-wheels/*.whl \\ + && VLLM_PACKAGE="$(python3 -c 'import os,vllm; print(os.path.dirname(vllm.__file__))')" \\ + && cd "$(dirname "$VLLM_PACKAGE")" \\ + && git apply --check /tmp/tracelens-vllm.patch \\ + && git apply /tmp/tracelens-vllm.patch \\ + && python3 /tmp/verify-tracelens-runtime.py \\ + && rm -rf /tmp/tracelens-wheels /tmp/tracelens-vllm.patch /tmp/verify-tracelens-runtime.py +WORKDIR /workspace +""" + (context / "Dockerfile").write_text(dockerfile, encoding="utf-8") + return labels + + +def build_vllm_tracelens_image( + *, + identity: VllmTraceLensIdentity, + tracelens_repo: Path, + derived_image: str, +) -> Dict[str, Any]: + """Build and verify a minimal TraceLens image from immutable inputs.""" + with tempfile.TemporaryDirectory(prefix="magpie-tracelens-vllm-") as temp_dir: + context = Path(temp_dir) + source_dir = context / "source" + wheelhouse = context / "wheels" + wheelhouse.mkdir() + _stage_committed_source(identity, tracelens_repo, source_dir) + source_wheel_command = _build_source_wheel(identity, source_dir, wheelhouse) + download_command = _download_requirement_wheels(identity, wheelhouse) + wheel_manifest = _wheel_manifest(wheelhouse) + shutil.rmtree(source_dir) + labels = _write_build_context(context, identity, wheel_manifest) + archive = context / "tracelens-source.tar" + if archive.exists(): + archive.unlink() + cmd = [ + "docker", + "build", + "--network", + "none", + "--no-cache", + "--provenance=false", + "-t", + derived_image, + str(context), + ] + proc = subprocess.run( + cmd, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + check=False, + ) + if proc.returncode != 0: + raise RuntimeError( + "TraceLens vLLM runtime image build failed with exit code " + f"{proc.returncode}. Image: {derived_image}\n{(proc.stdout or '')[-4000:]}" + ) + + validation = validate_vllm_tracelens_image(derived_image, identity) + if not validation["valid"]: + raise RuntimeError( + "Built TraceLens vLLM image failed identity validation: " + f"{validation['reason']}" + ) + record = docker_image_record(derived_image) or {} + return { + "command": cmd[:-1] + [""], + "source_wheel_command": source_wheel_command[:-1] + [""], + "requirements_download_command": download_command, + "image_id": record.get("Id"), + "image_labels": labels, + "dependency_wheels": list(wheel_manifest), + "dependency_wheel_manifest_sha256": labels[ + LABEL_WHEEL_MANIFEST_SHA256 + ], + "validation": validation, + } + + +def _validate_wheel_labels(labels: Mapping[str, Any]) -> Optional[str]: + manifest_text = labels.get(LABEL_WHEEL_MANIFEST) + manifest_sha = labels.get(LABEL_WHEEL_MANIFEST_SHA256) + if not isinstance(manifest_text, str) or not isinstance(manifest_sha, str): + return "missing wheel-manifest labels" + if _sha256_bytes(manifest_text.encode("utf-8")) != manifest_sha: + return "wheel-manifest label digest mismatch" + try: + manifest = json.loads(manifest_text) + except json.JSONDecodeError: + return "wheel-manifest label is not valid JSON" + if not isinstance(manifest, list): + return "wheel-manifest label is not a list" + expected = dict(VLLM_TRACELENS_REQUIREMENTS) + found: Dict[str, str] = {} + for item in manifest: + if not isinstance(item, dict): + return "wheel-manifest entry is not an object" + distribution = item.get("distribution") + digest = item.get("sha256") + version = item.get("version") + if not isinstance(distribution, str) or not isinstance(digest, str): + return "wheel-manifest entry is incomplete" + if not _HASH_RE.fullmatch(digest): + return f"invalid wheel digest for {distribution}" + if distribution != "tracelens": + found[distribution] = str(version) + if found != expected: + return f"pinned wheel versions mismatch: expected={expected}, found={found}" + if not any(item.get("distribution") == "tracelens" for item in manifest): + return "TraceLens source wheel is absent from wheel-manifest" + return None + + +def _runtime_probe(image: str) -> subprocess.CompletedProcess[str]: + return subprocess.run( + [ + "docker", + "run", + "--rm", + "--network", + "none", + "--entrypoint", + "python3", + image, + "-c", + _verification_script(), + ], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + timeout=120, + check=False, + ) + + +def _patch_reverse_probe( + image: str, + patch_bytes: bytes, +) -> subprocess.CompletedProcess[bytes]: + script = ( + "VLLM_PACKAGE=$(python3 -c 'import os,vllm; " + "print(os.path.dirname(vllm.__file__))'); " + 'cd "$(dirname "$VLLM_PACKAGE")"; ' + "git apply --reverse --check -" + ) + return subprocess.run( + [ + "docker", + "run", + "--rm", + "--network", + "none", + "--entrypoint", + "bash", + "-i", + image, + "-lc", + script, + ], + input=patch_bytes, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + timeout=120, + check=False, + ) + + +def validate_vllm_tracelens_image( + image: str, + identity: VllmTraceLensIdentity, +) -> Dict[str, Any]: + """Validate labels, base-layer ancestry, packages, imports, and patch markers.""" + record = docker_image_record(image) + if not record: + return {"valid": False, "reason": "image is not locally inspectable"} + labels = (record.get("Config") or {}).get("Labels") or {} + if not isinstance(labels, dict): + return {"valid": False, "reason": "image labels are unavailable"} + mismatches = { + key: {"expected": value, "actual": labels.get(key)} + for key, value in identity.labels().items() + if labels.get(key) != value + } + if mismatches: + return { + "valid": False, + "reason": "identity label mismatch", + "label_mismatches": mismatches, + } + wheel_error = _validate_wheel_labels(labels) + if wheel_error: + return {"valid": False, "reason": wheel_error} + + base_record = docker_image_record(identity.base_image_id) + base_layers = ((base_record or {}).get("RootFS") or {}).get("Layers") or [] + image_layers = (record.get("RootFS") or {}).get("Layers") or [] + if not base_layers or image_layers[: len(base_layers)] != base_layers: + return {"valid": False, "reason": "base image layer ancestry mismatch"} + + try: + probe = _runtime_probe(image) + except (OSError, subprocess.TimeoutExpired) as exc: + return {"valid": False, "reason": f"runtime probe failed: {exc}"} + if probe.returncode != 0: + return { + "valid": False, + "reason": "runtime package/import/patch probe failed", + "probe_output": _completed_output(probe)[-2000:], + } + lines = (probe.stdout or "").strip().splitlines() + try: + probe_result = json.loads(lines[-1]) if lines else {} + except json.JSONDecodeError: + return {"valid": False, "reason": "runtime probe returned invalid JSON"} + if probe_result.get("vllm_version") != identity.vllm_version: + return {"valid": False, "reason": "vLLM version changed in derived image"} + if probe_result.get("grpcio_version") != identity.grpcio_version: + return {"valid": False, "reason": "grpcio version changed in derived image"} + try: + patch_probe = _patch_reverse_probe(image, identity.patch_bytes) + except (OSError, subprocess.TimeoutExpired) as exc: + return {"valid": False, "reason": f"exact patch probe failed: {exc}"} + if patch_probe.returncode != 0: + return { + "valid": False, + "reason": "installed vLLM files do not match the committed patch", + "probe_output": _completed_output(patch_probe)[-2000:], + } + return { + "valid": True, + "reason": ( + "identity, ancestry, packages, imports, and exact patch verified" + ), + "image_id": record.get("Id"), + "runtime_probe": probe_result, + "dependency_wheels": json.loads(labels[LABEL_WHEEL_MANIFEST]), + "dependency_wheel_manifest_sha256": labels[ + LABEL_WHEEL_MANIFEST_SHA256 + ], + } + + +__all__ = [ + "VLLM_TRACELENS_FORBIDDEN", + "VLLM_TRACELENS_REQUIREMENTS", + "VllmTraceLensIdentity", + "build_vllm_tracelens_image", + "docker_image_id", + "docker_image_record", + "resolve_vllm_tracelens_identity", + "validate_vllm_tracelens_image", +] diff --git a/docs/how-to/benchmarking/profiling-options.md b/docs/how-to/benchmarking/profiling-options.md index e380626..d26bc81 100644 --- a/docs/how-to/benchmarking/profiling-options.md +++ b/docs/how-to/benchmarking/profiling-options.md @@ -43,8 +43,25 @@ legacy direct PyTorch report flow. For Docker benchmarks, `auto_patch_runtime` defaults to `true`. When TraceLens inference mode is enabled and the selected runtime image is not already TraceLens-ready, Magpie builds a derived image from supported official -vLLM/SGLang tags using the public TraceLens workflow scripts. The derived image -is tagged locally as `magpie-tracelens-:...` and reused on later runs. +vLLM/SGLang tags. SGLang continues to use the public TraceLens workflow script. +For vLLM, Magpie archives the selected committed TraceLens tree, builds its +wheel without dependencies in the exact base image, and installs only pinned +Matplotlib diagnostic wheels with `--no-deps`. It deliberately does not install +`xprof`, `gcsfs`, or `grpcio-status`, and verifies that the base image's `vllm` +and `grpcio` versions remain unchanged. The final Docker build is network-free. +The derived image is tagged locally as `magpie-tracelens-:...` and +reused on later runs only after validation. + +The vLLM image carries OCI labels and a runtime identity document containing the +base image ID, TraceLens source commit and tree, patch hash, pinned wheel hashes, +and preserved package versions. A same-name local image with missing or stale +identity, changed base ancestry, forbidden packages, import failures, or patch +marker failures is rejected and rebuilt. These fields are also returned in the +benchmark runtime metadata. TraceLens' upstream wheel metadata still declares +features outside Magpie's CSV diagnostic path, so a whole-environment +`pip check` can report intentionally omitted packages; Magpie instead validates +the exact splitter/report/import path it executes. + If no TraceLens source path is configured, Magpie shallow clones the official TraceLens `main` branch to `$XDG_CACHE_HOME/magpie/TraceLens` (or `~/.cache/magpie/TraceLens`) and reuses @@ -107,6 +124,12 @@ Magpie logs a warning and continues without architecture-specific roofline data. An explicit `tracelens.gpu_arch_config` takes priority and is passed as `--gpu_arch_json_path`. +In particular, the public TraceLens commit may not yet bundle `MI355X` even +though the vLLM instrumentation patch supports an MI355X workload. In that +case the trace, stage splitting, timing tables, kernel attribution, and CSV +reports remain valid; only architecture-dependent roofline columns are omitted. +Provide a reviewed `gpu_arch_config` or extension wheel to enable those columns. + For SGLang, TraceLens inference mode automatically adds `--enable-profile-cuda-graph`. It also adds `--enable-shape-discovery-for-cuda-graph-profile` when the configured Docker diff --git a/docs/reference/benchmark-config.md b/docs/reference/benchmark-config.md index a9a7bc3..0655c3a 100644 --- a/docs/reference/benchmark-config.md +++ b/docs/reference/benchmark-config.md @@ -275,6 +275,12 @@ when the candidate is supported; otherwise it warns and continues without architecture-specific roofline data. An explicit `gpu_arch_config` takes priority and is passed as `--gpu_arch_json_path`. +For example, if the selected public TraceLens commit does not bundle an +`MI355X.json`, an MI355X run still produces valid profiler traces, stage splits, +kernel attribution, and CSV timing reports. Only architecture-specific roofline +fields are unavailable. A reviewed `gpu_arch_config` or `TL_EXTENSION` wheel can +add that platform explicitly. + ### SGLang benchmark Basic SGLang benchmark with torch profiler enabled: diff --git a/docs/reference/release-notes.md b/docs/reference/release-notes.md index 259a8b7..fd6686f 100644 --- a/docs/reference/release-notes.md +++ b/docs/reference/release-notes.md @@ -24,6 +24,9 @@ initial beta. - TraceLens runtime patch selection is now based on the framework version installed in the runtime image and the patches available in the selected TraceLens checkout, instead of a hard-coded version list. +- Derived TraceLens vLLM images now use committed source plus a minimal pinned + wheel set, preserve the base vLLM/grpcio versions, record OCI identity labels, + and reject stale same-name images before reuse. - Trace processing now selects non-empty inference traces, improving report generation when profiler output contains empty trace files. - The compatibility matrix and the benchmarking, profiling, Ray, MCP, and diff --git a/tests/test_benchmark_support.py b/tests/test_benchmark_support.py index 053d82d..59404e3 100644 --- a/tests/test_benchmark_support.py +++ b/tests/test_benchmark_support.py @@ -3,6 +3,7 @@ import json import subprocess import zipfile +from dataclasses import replace from pathlib import Path import pytest @@ -44,10 +45,30 @@ resolve_tracelens_repo_path, runner_type_to_gpu_type, ) +from Magpie.modes.benchmark.tracelens_vllm_image import VllmTraceLensIdentity from Magpie.modes.benchmark.workspace import WorkspaceManager from Magpie.utils.gpu import GPUVendor +def _fake_vllm_tracelens_identity( + base_image="internal/vllm-rocm:latest", + vllm_version="0.22.0+rocm722", +): + return VllmTraceLensIdentity( + base_image=base_image, + base_image_id="sha256:" + "1" * 64, + base_image_locator=base_image, + vllm_version=vllm_version, + grpcio_version="1.78.0", + source_commit="2" * 40, + source_tree="3" * 40, + patch_version="v22", + patch_path="vllm_patches/config_vllm_v0.22.0.patch", + patch_sha256="4" * 64, + patch_bytes=b"patch", + ) + + def test_workspace_manager_makes_docker_mounts_container_writable(tmp_path): workspace = WorkspaceManager( base_dir=str(tmp_path), @@ -803,7 +824,25 @@ def test_prepare_tracelens_runtime_image_prefers_installed_package_version( ) monkeypatch.setattr( "Magpie.modes.benchmark.tracelens_runtime.docker_image_package_version", - lambda _image, package: "0.22.0+rocm722" if package == "vllm" else None, + lambda _image, package: { + "vllm": "0.22.0+rocm722", + "grpcio": "1.78.0", + }.get(package), + ) + identity = _fake_vllm_tracelens_identity() + monkeypatch.setattr( + "Magpie.modes.benchmark.tracelens_runtime.resolve_vllm_tracelens_identity", + lambda **_kwargs: identity, + ) + monkeypatch.setattr( + "Magpie.modes.benchmark.tracelens_runtime.validate_vllm_tracelens_image", + lambda _image, _identity: { + "valid": True, + "reason": "verified", + "image_id": "sha256:" + "5" * 64, + "dependency_wheels": [], + "dependency_wheel_manifest_sha256": "6" * 64, + }, ) cfg = BenchmarkConfig.from_dict( @@ -825,6 +864,8 @@ def test_prepare_tracelens_runtime_image_prefers_installed_package_version( assert result["patch_version"] == "v22" assert result["patch_version_source"] == "package" assert result["runtime_package_version"] == "0.22.0+rocm722" + assert result["base_grpcio_version"] == "1.78.0" + assert result["reason"] == "validated derived image already exists" assert result["image"].startswith("magpie-tracelens-vllm:v22-mi355x-") @@ -862,7 +903,37 @@ def test_prepare_tracelens_runtime_image_builds_extension_overlay( ) monkeypatch.setattr( "Magpie.modes.benchmark.tracelens_runtime.docker_image_package_version", - lambda _image, _package: None, + lambda _image, package: { + "vllm": "0.21.0+rocm", + "grpcio": "1.78.0", + }.get(package), + ) + + identity = _fake_vllm_tracelens_identity( + base_image="vllm/vllm-openai-rocm:v0.21.0", + vllm_version="0.21.0+rocm", + ) + identity = replace( + identity, + patch_version="v21", + patch_path="vllm_patches/config_vllm_v0.21.0.patch", + ) + monkeypatch.setattr( + "Magpie.modes.benchmark.tracelens_runtime.resolve_vllm_tracelens_identity", + lambda **_kwargs: identity, + ) + monkeypatch.setattr( + "Magpie.modes.benchmark.tracelens_runtime.build_vllm_tracelens_image", + lambda **_kwargs: { + "command": ["docker", "build", ""], + "source_wheel_command": ["docker", "run", ""], + "requirements_download_command": ["docker", "run", "pip", "download"], + "image_id": "sha256:" + "7" * 64, + "image_labels": identity.labels(), + "dependency_wheels": [], + "dependency_wheel_manifest_sha256": "8" * 64, + "validation": {"valid": True, "reason": "verified"}, + }, ) calls = [] @@ -912,12 +983,11 @@ def fake_run(cmd, **kwargs): f"-ext-{result['extension_wheel_sha256'][:12]}" ) assert cfg.envs["TL_EXTENSION"] == "ExistingExtension:TraceLens_Ext" - assert calls[0]["cmd"][0] == "bash" - assert calls[1]["cmd"][:2] == ["docker", "build"] - assert "ENV TL_EXTENSION=TraceLens_Ext" in calls[1]["dockerfile"] + assert calls[0]["cmd"][:2] == ["docker", "build"] + assert "ENV TL_EXTENSION=TraceLens_Ext" in calls[0]["dockerfile"] assert ( "TraceLens_Ext-0.1.0.dev20260529+gacb7fbc6-py3-none-any.whl" - in calls[1]["context_files"] + in calls[0]["context_files"] ) diff --git a/tests/test_tracelens_vllm_image.py b/tests/test_tracelens_vllm_image.py new file mode 100644 index 0000000..785d787 --- /dev/null +++ b/tests/test_tracelens_vllm_image.py @@ -0,0 +1,287 @@ +import hashlib +import json +import subprocess + +from Magpie.modes.benchmark.config import BenchmarkConfig +from Magpie.modes.benchmark.tracelens_runtime import ( + prepare_tracelens_runtime_image, +) +from Magpie.modes.benchmark.tracelens_vllm_image import ( + LABEL_WHEEL_MANIFEST, + LABEL_WHEEL_MANIFEST_SHA256, + VLLM_TRACELENS_FORBIDDEN, + VLLM_TRACELENS_REQUIREMENTS, + VllmTraceLensIdentity, + _verification_script, + _write_build_context, + resolve_vllm_tracelens_identity, + validate_vllm_tracelens_image, +) + + +def _identity(): + return VllmTraceLensIdentity( + base_image="vllm/vllm-openai-rocm:v0.19.1", + base_image_id="sha256:" + "1" * 64, + base_image_locator=( + "vllm/vllm-openai-rocm@sha256:" + "a" * 64 + ), + vllm_version="0.19.1+rocm721", + grpcio_version="1.78.0", + source_commit="2" * 40, + source_tree="3" * 40, + patch_version="v19", + patch_path=( + "examples/custom_workflows/inference_analysis/vllm_patches/" + "config_vllm_v0.19.0.patch" + ), + patch_sha256="4" * 64, + patch_bytes=b"diff --git a/vllm/example.py b/vllm/example.py\n", + ) + + +def _wheel_manifest(): + records = [ + { + "distribution": name, + "filename": f"{name}-{version}-py3-none-any.whl", + "version": version, + "sha256": hashlib.sha256(name.encode()).hexdigest(), + } + for name, version in VLLM_TRACELENS_REQUIREMENTS + ] + records.append( + { + "distribution": "tracelens", + "filename": "tracelens-0.1.0-py3-none-any.whl", + "version": "source-commit", + "sha256": hashlib.sha256(b"tracelens").hexdigest(), + } + ) + return records + + +def _labels(identity): + labels = identity.labels() + manifest = json.dumps(_wheel_manifest(), sort_keys=True, separators=(",", ":")) + labels[LABEL_WHEEL_MANIFEST] = manifest + labels[LABEL_WHEEL_MANIFEST_SHA256] = hashlib.sha256( + manifest.encode() + ).hexdigest() + return labels + + +def test_resolve_identity_uses_image_digest_and_committed_patch(monkeypatch, tmp_path): + patch = b"committed patch bytes\n" + monkeypatch.setattr( + "Magpie.modes.benchmark.tracelens_vllm_image.docker_image_record", + lambda _image: { + "Id": "sha256:" + "1" * 64, + "RepoDigests": ["registry/vllm@sha256:" + "a" * 64], + }, + ) + monkeypatch.setattr( + "Magpie.modes.benchmark.tracelens_vllm_image._git_text", + lambda _repo, *args: "2" * 40 if args[-1] == "HEAD" else "3" * 40, + ) + monkeypatch.setattr( + "Magpie.modes.benchmark.tracelens_vllm_image._git_bytes", + lambda _repo, *_args: patch, + ) + + identity = resolve_vllm_tracelens_identity( + base_image="registry/vllm:v0.19.1", + vllm_version="0.19.1+rocm721", + grpcio_version="1.78.0", + tracelens_repo=tmp_path, + patch_version="v19", + ) + + assert identity.base_image_locator == "registry/vllm@sha256:" + "a" * 64 + assert identity.source_commit == "2" * 40 + assert identity.source_tree == "3" * 40 + assert identity.patch_sha256 == hashlib.sha256(patch).hexdigest() + labels = identity.labels() + assert labels["io.magpie.tracelens.base-image-id"] == identity.base_image_id + assert labels["io.magpie.tracelens.source-tree"] == identity.source_tree + + +def test_build_context_is_offline_minimal_and_identity_labeled(tmp_path): + identity = _identity() + labels = _write_build_context(tmp_path, identity, _wheel_manifest()) + + dockerfile = (tmp_path / "Dockerfile").read_text(encoding="utf-8") + verifier = (tmp_path / "verify.py").read_text(encoding="utf-8") + identity_document = json.loads( + (tmp_path / "identity.json").read_text(encoding="utf-8") + ) + + assert dockerfile.startswith(f"FROM {identity.base_image_locator}\n") + assert "--no-index --no-deps" in dockerfile + assert "git apply --check" in dockerfile + assert "pip install --upgrade" not in dockerfile + assert labels[LABEL_WHEEL_MANIFEST_SHA256] + assert identity_document["tracelens_source_commit"] == identity.source_commit + assert identity_document["tracelens_source_tree"] == identity.source_tree + assert identity_document["tracelens_patch_sha256"] == identity.patch_sha256 + assert "metadata.version(\"grpcio\")" in verifier + assert "metadata.version(\"vllm\")" in verifier + for package in VLLM_TRACELENS_FORBIDDEN: + assert package in verifier + + +def test_existing_image_validation_preserves_vllm_grpc_and_excludes_pollution( + monkeypatch, +): + identity = _identity() + base_record = { + "Id": identity.base_image_id, + "RootFS": {"Layers": ["base-a", "base-b"]}, + "Config": {"Labels": {}}, + } + image_record = { + "Id": "sha256:" + "5" * 64, + "RootFS": {"Layers": ["base-a", "base-b", "derived"]}, + "Config": {"Labels": _labels(identity)}, + } + + monkeypatch.setattr( + "Magpie.modes.benchmark.tracelens_vllm_image.docker_image_record", + lambda image: base_record if image == identity.base_image_id else image_record, + ) + + seen = {} + + def fake_probe(image): + seen["image"] = image + seen["script"] = _verification_script() + return subprocess.CompletedProcess( + ["docker", "run"], + 0, + stdout=json.dumps( + { + "grpcio_version": "1.78.0", + "platforms": ["MI300X", "MI325X"], + "tracelens_version": "0.1.0", + "vllm_version": "0.19.1+rocm721", + } + ) + + "\n", + stderr="", + ) + + monkeypatch.setattr( + "Magpie.modes.benchmark.tracelens_vllm_image._runtime_probe", + fake_probe, + ) + monkeypatch.setattr( + "Magpie.modes.benchmark.tracelens_vllm_image._patch_reverse_probe", + lambda _image, _patch: subprocess.CompletedProcess( + ["docker", "run"], 0, stdout=b"", stderr=b"" + ), + ) + + validation = validate_vllm_tracelens_image("derived:v19", identity) + + assert validation["valid"] is True + assert validation["runtime_probe"]["grpcio_version"] == "1.78.0" + assert validation["runtime_probe"]["vllm_version"] == "0.19.1+rocm721" + for package in VLLM_TRACELENS_FORBIDDEN: + assert package in seen["script"] + + +def test_existing_image_validation_rejects_stale_identity_without_probe(monkeypatch): + identity = _identity() + stale_labels = _labels(identity) + stale_labels["io.magpie.tracelens.source-commit"] = "9" * 40 + monkeypatch.setattr( + "Magpie.modes.benchmark.tracelens_vllm_image.docker_image_record", + lambda _image: { + "Id": "sha256:" + "5" * 64, + "RootFS": {"Layers": ["base-a", "derived"]}, + "Config": {"Labels": stale_labels}, + }, + ) + monkeypatch.setattr( + "Magpie.modes.benchmark.tracelens_vllm_image._runtime_probe", + lambda _image: (_ for _ in ()).throw(AssertionError("probe must not run")), + ) + + validation = validate_vllm_tracelens_image("derived:v19", identity) + + assert validation["valid"] is False + assert validation["reason"] == "identity label mismatch" + assert "io.magpie.tracelens.source-commit" in validation["label_mismatches"] + + +def test_prepare_rebuilds_stale_vllm_tag(monkeypatch, tmp_path): + identity = _identity() + workflow = tmp_path / "examples/custom_workflows/inference_analysis" + patch_dir = workflow / "vllm_patches" + patch_dir.mkdir(parents=True) + (workflow / "build_docker_vllm.sh").write_text( + "case ${VLLM_VERSION} in\n v19) ;;\nesac\n", + encoding="utf-8", + ) + (patch_dir / "config_vllm_v0.19.0.patch").write_text( + "patch", + encoding="utf-8", + ) + monkeypatch.setenv("TRACELENS_REPO_PATH", str(tmp_path)) + monkeypatch.setattr( + "Magpie.modes.benchmark.tracelens_runtime.docker_image_exists", + lambda _image: True, + ) + monkeypatch.setattr( + "Magpie.modes.benchmark.tracelens_runtime.docker_image_package_version", + lambda _image, package: { + "vllm": identity.vllm_version, + "grpcio": identity.grpcio_version, + }.get(package), + ) + monkeypatch.setattr( + "Magpie.modes.benchmark.tracelens_runtime.resolve_vllm_tracelens_identity", + lambda **_kwargs: identity, + ) + monkeypatch.setattr( + "Magpie.modes.benchmark.tracelens_runtime.validate_vllm_tracelens_image", + lambda _image, _identity: { + "valid": False, + "reason": "identity label mismatch", + }, + ) + builds = [] + + def fake_build(**kwargs): + builds.append(kwargs) + return { + "command": ["docker", "build", ""], + "source_wheel_command": ["docker", "run", ""], + "requirements_download_command": ["docker", "run", "pip", "download"], + "image_id": "sha256:" + "6" * 64, + "image_labels": identity.labels(), + "dependency_wheels": _wheel_manifest(), + "dependency_wheel_manifest_sha256": "7" * 64, + "validation": {"valid": True, "reason": "verified after rebuild"}, + } + + monkeypatch.setattr( + "Magpie.modes.benchmark.tracelens_runtime.build_vllm_tracelens_image", + fake_build, + ) + cfg = BenchmarkConfig.from_dict( + { + "framework": "vllm", + "model": "demo", + "docker_image": identity.base_image, + "profiler": {"tracelens": {"enabled": True}}, + } + ) + + result = prepare_tracelens_runtime_image(cfg, identity.base_image, "mi355x") + + assert len(builds) == 1 + assert result["stale_image_rejected"] is True + assert result["stale_image_rejection_reason"] == "identity label mismatch" + assert result["built"] is True + assert result["public_runtime_validation"]["valid"] is True From 0c2a1b786d84b2e8050aed3c4d120acca1e85dae Mon Sep 17 00:00:00 2001 From: Vincent Ouyang Date: Fri, 7 Aug 2026 22:46:01 +0000 Subject: [PATCH 04/11] Bind Qwen evaluator budgets and raw quality evidence --- Magpie/modes/benchmark/quality.py | 70 +++++++++++++++++-- Magpie/scripts/benchmark/lm_eval_runtime.sh | 60 +++++++++++++++++ Magpie/scripts/benchmark/vllm_mi355x.sh | 9 ++- docs/how-to/benchmarking/benchmark.md | 10 ++- tests/test_benchmark_support.py | 10 +++ tests/test_lm_eval_runtime.py | 74 +++++++++++++++++++++ 6 files changed, 223 insertions(+), 10 deletions(-) diff --git a/Magpie/modes/benchmark/quality.py b/Magpie/modes/benchmark/quality.py index e75cb3c..90cdf7b 100644 --- a/Magpie/modes/benchmark/quality.py +++ b/Magpie/modes/benchmark/quality.py @@ -34,6 +34,21 @@ def _file_sha256(path: Path) -> str: return digest.hexdigest() +def _canonical_digest(value: Mapping[str, Any]) -> str: + encoded = json.dumps( + value, ensure_ascii=False, separators=(",", ":"), sort_keys=True + ).encode("utf-8") + return hashlib.sha256(encoded).hexdigest() + + +def _artifact_receipt(path: Path, workspace: Path) -> Dict[str, Any]: + return { + "path": str(path.relative_to(workspace)), + "size_bytes": path.stat().st_size, + "sha256": _file_sha256(path), + } + + def _numeric_metrics(data: Mapping[str, Any]) -> Dict[str, float]: metrics: Dict[str, float] = {} for name, value in data.items(): @@ -92,18 +107,32 @@ def parse_lm_eval_quality(workspace: Path, *, requested: bool) -> Dict[str, Any] str(path.relative_to(workspace)) for path in artifact_files[:MAX_REPORTED_ARTIFACTS] ] + artifact_receipts: List[Dict[str, Any]] = [] result_receipts: List[Dict[str, Any]] = [] + sample_receipts: List[Dict[str, Any]] = [] + for path in artifact_files[:MAX_REPORTED_ARTIFACTS]: + try: + receipt = _artifact_receipt(path, workspace) + artifact_receipts.append(receipt) + if path.name.startswith("samples") and path.suffix == ".jsonl": + sample_receipts.append(receipt) + except OSError: + pass for path in result_files[:MAX_REPORTED_ARTIFACTS]: try: - result_receipts.append( - { - "path": str(path.relative_to(workspace)), - "size_bytes": path.stat().st_size, - "sha256": _file_sha256(path), - } - ) + result_receipts.append(_artifact_receipt(path, workspace)) except OSError: pass + sample_set_digest = ( + _canonical_digest( + { + "schema": "magpie.lm-eval-sample-set/v1", + "artifacts": sample_receipts, + } + ) + if sample_receipts + else None + ) if not result_files: status = "missing" if requested else "not_requested" @@ -123,6 +152,10 @@ def parse_lm_eval_quality(workspace: Path, *, requested: bool) -> Dict[str, Any] "artifact_count": len(artifact_files), "artifacts_truncated": len(artifact_files) > MAX_REPORTED_ARTIFACTS, "result_artifact_receipts": result_receipts, + "artifact_receipts": artifact_receipts, + "sample_artifact_receipts": sample_receipts, + "sample_set_digest": sample_set_digest, + "outcome_digest": None, "result_artifact_count": len(result_files), "result_artifacts_truncated": ( len(result_files) > MAX_REPORTED_ARTIFACTS @@ -174,6 +207,23 @@ def parse_lm_eval_quality(workspace: Path, *, requested: bool) -> Dict[str, Any] passed = bool(tasks) and not errors task_items = sorted(tasks.items()) reported_tasks = dict(task_items[:MAX_REPORTED_TASKS]) + primary_outcomes = { + task: { + "metric": value["primary_metric"], + "value": value["value"], + "source": value["source"], + } + for task, value in reported_tasks.items() + } + outcome_digest = _canonical_digest( + { + "schema": "magpie.lm-eval-outcomes/v1", + "primary_metric_policy": list(PRIMARY_METRICS), + "outcomes": primary_outcomes, + "result_artifacts": result_receipts, + "sample_set_digest": sample_set_digest, + } + ) return { "kind": "lm_eval", "requested": requested, @@ -181,12 +231,18 @@ def parse_lm_eval_quality(workspace: Path, *, requested: bool) -> Dict[str, Any] "passed": passed, "evidence_present": bool(tasks), "tasks": reported_tasks, + "primary_metric_policy": list(PRIMARY_METRICS), + "primary_outcomes": primary_outcomes, + "outcome_digest": outcome_digest, + "sample_set_digest": sample_set_digest, "task_count": len(tasks), "tasks_truncated": len(tasks) > MAX_REPORTED_TASKS, "artifacts": relative_artifacts, "artifact_count": len(artifact_files), "artifacts_truncated": len(artifact_files) > MAX_REPORTED_ARTIFACTS, "result_artifact_receipts": result_receipts, + "artifact_receipts": artifact_receipts, + "sample_artifact_receipts": sample_receipts, "result_artifact_count": len(result_files), "result_artifacts_truncated": ( len(result_files) > MAX_REPORTED_ARTIFACTS diff --git a/Magpie/scripts/benchmark/lm_eval_runtime.sh b/Magpie/scripts/benchmark/lm_eval_runtime.sh index 3f591dd..97abaa8 100644 --- a/Magpie/scripts/benchmark/lm_eval_runtime.sh +++ b/Magpie/scripts/benchmark/lm_eval_runtime.sh @@ -248,3 +248,63 @@ _install_lm_eval_deps() { echo "ERROR: refusing to run lm-eval without a verified locked runtime." >&2 exit "$status" } + +# Run lm-eval from a caller-bound evaluator policy instead of inheriting +# InferenceX's coupled context/output-budget heuristic. The serving process +# continues to use MAX_MODEL_LEN; these two values govern only lm-eval's +# request admission and generated output budget. +magpie_run_lm_eval() { + local port="${PORT:-8888}" + local results_dir="${EVAL_RESULT_DIR:-${RESULT_DIR:-/workspace}/lm_eval}" + local max_length="${MAGPIE_EVAL_MAX_LENGTH:-}" + local max_gen_tokens="${MAGPIE_EVAL_MAX_GEN_TOKENS:-}" + local tasks="${MAGPIE_EVAL_TASKS:-gsm8k}" + local concurrent_requests="${EVAL_CONCURRENT_REQUESTS:-${CONC:-8}}" + local batch_size="${MAGPIE_EVAL_BATCH_SIZE:-auto}" + local python="${MAGPIE_EVAL_PYTHON:-python3}" + + while [[ $# -gt 0 ]]; do + case "$1" in + --port) port="$2"; shift 2 ;; + --results-dir) results_dir="$2"; shift 2 ;; + *) echo "ERROR: unsupported Magpie lm-eval argument: $1" >&2; return 2 ;; + esac + done + if [[ ! "$max_length" =~ ^[1-9][0-9]*$ ]] \ + || [[ ! "$max_gen_tokens" =~ ^[1-9][0-9]*$ ]] \ + || (( max_gen_tokens >= max_length )); then + echo "ERROR: evaluator policy requires positive MAGPIE_EVAL_MAX_LENGTH and MAGPIE_EVAL_MAX_GEN_TOKENS with output < context." >&2 + return 42 + fi + if [[ -z "${MAGPIE_EVAL_POLICY_ID:-}" || -z "${MAGPIE_EVAL_PRIMARY_METRIC:-}" ]]; then + echo "ERROR: evaluator policy identity and primary metric are required." >&2 + return 42 + fi + + _install_lm_eval_deps || return $? + mkdir -p "$results_dir" || return $? + export EVAL_RESULT_DIR="$results_dir" + export OPENAI_API_KEY="${OPENAI_API_KEY:-EMPTY}" + local model_name="${MODEL_NAME:-${MODEL:-}}" + local base_url="http://0.0.0.0:${port}/v1/chat/completions" + local model_args="model=${model_name},base_url=${base_url},api_key=${OPENAI_API_KEY},eos_string=,max_retries=5,num_concurrent=${concurrent_requests},timeout=1800,tokenized_requests=False,max_length=${max_length}" + local gen_kwargs="max_tokens=${max_gen_tokens},temperature=0,top_p=1" + local -a command=( + "$python" -m lm_eval + --model local-chat-completions + --apply_chat_template + --tasks "$tasks" + --output_path "$results_dir" + --log_samples + --batch_size "$batch_size" + --model_args "$model_args" + --gen_kwargs "$gen_kwargs" + ) + if [[ -n "${MAGPIE_EVAL_LIMIT:-}" ]]; then + command+=(--limit "$MAGPIE_EVAL_LIMIT") + fi + printf '[magpie] evaluator policy=%s primary=%s max_length=%s max_gen_tokens=%s\n' \ + "$MAGPIE_EVAL_POLICY_ID" "$MAGPIE_EVAL_PRIMARY_METRIC" \ + "$max_length" "$max_gen_tokens" >&2 + "${command[@]}" +} diff --git a/Magpie/scripts/benchmark/vllm_mi355x.sh b/Magpie/scripts/benchmark/vllm_mi355x.sh index 17837fc..b57b53d 100644 --- a/Magpie/scripts/benchmark/vllm_mi355x.sh +++ b/Magpie/scripts/benchmark/vllm_mi355x.sh @@ -234,8 +234,13 @@ if [[ "$PHASE" != "server" && "${RUN_EVAL,,}" = "true" ]]; then fi else magpie_mark_lm_eval_start || exit $? - EVAL_CONCURRENT_REQUESTS="${EVAL_CONCURRENT_REQUESTS:-$CONC}" \ - run_eval --framework lm-eval --port "$PORT" || exit $? + if [[ -n "${MAGPIE_EVAL_POLICY_ID:-}" ]]; then + EVAL_CONCURRENT_REQUESTS="${EVAL_CONCURRENT_REQUESTS:-$CONC}" \ + magpie_run_lm_eval --port "$PORT" || exit $? + else + EVAL_CONCURRENT_REQUESTS="${EVAL_CONCURRENT_REQUESTS:-$CONC}" \ + run_eval --framework lm-eval --port "$PORT" || exit $? + fi magpie_preserve_lm_eval_artifacts || exit $? append_lm_eval_summary magpie_preserve_lm_eval_artifacts || exit $? diff --git a/docs/how-to/benchmarking/benchmark.md b/docs/how-to/benchmarking/benchmark.md index f603ccb..8033e0f 100644 --- a/docs/how-to/benchmarking/benchmark.md +++ b/docs/how-to/benchmarking/benchmark.md @@ -122,7 +122,15 @@ Every report declares `run_kind` and `reward_eligible`. A `run_kind: measurement` run rejects heavy profilers; diagnostic runs and all TargetedKernelTrace artifacts have `reward_eligible: false`. When `RUN_EVAL=true`, raw lm-eval files remain under `lm_eval/` and `quality_gate` exposes each task's -primary metric. The same run must provide this nested configuration: +strictly ordered primary metric, a content-bound `outcome_digest`, the raw +artifact receipts, and a `sample_set_digest`. The same run must provide the +nested runtime configuration shown below. + +Apex's reviewed Qwen view also +sets `MAGPIE_EVAL_MAX_LENGTH=2248` and `MAGPIE_EVAL_MAX_GEN_TOKENS=480`: +the former is evaluator request admission, while the latter is the independent +generation budget. `MAX_MODEL_LEN` remains the serving context limit. The +locked Magpie helper constructs this argv; it does not patch InferenceX. ```yaml benchmark: diff --git a/tests/test_benchmark_support.py b/tests/test_benchmark_support.py index 59404e3..b5e15d4 100644 --- a/tests/test_benchmark_support.py +++ b/tests/test_benchmark_support.py @@ -115,6 +115,16 @@ def test_lm_eval_quality_receipt_preserves_task_metrics(tmp_path): assert "lm_eval/model/results_2026.json" in gate["artifacts"] assert gate["result_artifact_receipts"][0]["sha256"] assert gate["result_artifact_receipts"][0]["size_bytes"] > 0 + assert gate["primary_outcomes"]["gsm8k"] == { + "metric": "exact_match,strict-match", + "value": pytest.approx(0.812), + "source": "lm_eval/model/results_2026.json", + } + assert len(gate["outcome_digest"]) == 64 + assert len(gate["sample_set_digest"]) == 64 + assert gate["sample_artifact_receipts"][0]["path"].endswith( + "samples_gsm8k.jsonl" + ) def test_lm_eval_quality_requested_missing_fails_explicitly(tmp_path): diff --git a/tests/test_lm_eval_runtime.py b/tests/test_lm_eval_runtime.py index b29c74a..aae5d7d 100644 --- a/tests/test_lm_eval_runtime.py +++ b/tests/test_lm_eval_runtime.py @@ -372,6 +372,80 @@ def test_helper_contains_no_mutable_or_network_install_path(): assert forbidden not in source +def test_owned_evaluator_splits_context_and_output_budget(tmp_path): + fake_bin = tmp_path / "bin" + fake_bin.mkdir() + argv_path = tmp_path / "argv.txt" + fake_python = fake_bin / "python3" + fake_python.write_text( + "#!/usr/bin/env bash\nprintf '%s\\n' \"$@\" > \"$ARGV_PATH\"\n", + encoding="utf-8", + ) + fake_python.chmod(0o755) + env = os.environ.copy() + env.update( + { + "PATH": f"{fake_bin}:{env['PATH']}", + "ARGV_PATH": str(argv_path), + "MODEL": "Qwen/example", + "RESULT_DIR": str(tmp_path / "workspace"), + "MAGPIE_EVAL_POLICY_ID": "qwen3-next-gsm8k-v1", + "MAGPIE_EVAL_PRIMARY_METRIC": "exact_match,strict-match", + "MAGPIE_EVAL_MAX_LENGTH": "2248", + "MAGPIE_EVAL_MAX_GEN_TOKENS": "480", + } + ) + completed = subprocess.run( + [ + "bash", + "-c", + 'source "$1"; _install_lm_eval_deps() { return 0; }; ' + "magpie_run_lm_eval --port 8888", + "bash", + str(HELPER), + ], + env=env, + text=True, + capture_output=True, + timeout=10, + check=False, + ) + + assert completed.returncode == 0, completed.stderr + argv = argv_path.read_text(encoding="utf-8") + assert "max_length=2248" in argv + assert "max_tokens=480" in argv + assert "max_tokens=1124" not in argv + assert "--log_samples" in argv + + +@pytest.mark.parametrize( + ("max_length", "max_tokens"), + (("", "480"), ("2248", ""), ("2248", "2248"), ("bad", "480")), +) +def test_owned_evaluator_rejects_invalid_budget(max_length, max_tokens, tmp_path): + env = os.environ.copy() + env.update( + { + "MODEL": "Qwen/example", + "RESULT_DIR": str(tmp_path), + "MAGPIE_EVAL_POLICY_ID": "qwen3-next-gsm8k-v1", + "MAGPIE_EVAL_PRIMARY_METRIC": "exact_match,strict-match", + "MAGPIE_EVAL_MAX_LENGTH": max_length, + "MAGPIE_EVAL_MAX_GEN_TOKENS": max_tokens, + } + ) + completed = subprocess.run( + ["bash", "-c", 'source "$1"; magpie_run_lm_eval', "bash", str(HELPER)], + env=env, + text=True, + capture_output=True, + timeout=10, + check=False, + ) + assert completed.returncode == 42 + + def test_helper_failure_terminates_ignoring_upstream_caller(): env = os.environ.copy() for name in ( From ee0f1bbf50702ddc309102e012435224e4caefe9 Mon Sep 17 00:00:00 2001 From: Vincent Ouyang Date: Fri, 7 Aug 2026 23:29:06 +0000 Subject: [PATCH 05/11] Bind lm-eval to locked InferenceX tasks --- Magpie/scripts/benchmark/lm_eval_runtime.sh | 6 ++++++ Magpie/scripts/benchmark/magpie_bench_remote_compat.sh | 6 ++++++ tests/test_lm_eval_runtime.py | 5 +++++ 3 files changed, 17 insertions(+) diff --git a/Magpie/scripts/benchmark/lm_eval_runtime.sh b/Magpie/scripts/benchmark/lm_eval_runtime.sh index 97abaa8..c2afa26 100644 --- a/Magpie/scripts/benchmark/lm_eval_runtime.sh +++ b/Magpie/scripts/benchmark/lm_eval_runtime.sh @@ -259,6 +259,7 @@ magpie_run_lm_eval() { local max_length="${MAGPIE_EVAL_MAX_LENGTH:-}" local max_gen_tokens="${MAGPIE_EVAL_MAX_GEN_TOKENS:-}" local tasks="${MAGPIE_EVAL_TASKS:-gsm8k}" + local include_path="${MAGPIE_EVAL_INCLUDE_PATH:-${MAGPIE_INFERENCEX_ROOT:-$(pwd)}/utils/evals}" local concurrent_requests="${EVAL_CONCURRENT_REQUESTS:-${CONC:-8}}" local batch_size="${MAGPIE_EVAL_BATCH_SIZE:-auto}" local python="${MAGPIE_EVAL_PYTHON:-python3}" @@ -280,6 +281,10 @@ magpie_run_lm_eval() { echo "ERROR: evaluator policy identity and primary metric are required." >&2 return 42 fi + if [[ ! -d "$include_path" ]]; then + echo "ERROR: evaluator include path is unavailable: $include_path" >&2 + return 42 + fi _install_lm_eval_deps || return $? mkdir -p "$results_dir" || return $? @@ -293,6 +298,7 @@ magpie_run_lm_eval() { "$python" -m lm_eval --model local-chat-completions --apply_chat_template + --include_path "$include_path" --tasks "$tasks" --output_path "$results_dir" --log_samples diff --git a/Magpie/scripts/benchmark/magpie_bench_remote_compat.sh b/Magpie/scripts/benchmark/magpie_bench_remote_compat.sh index 7b0eaa0..b378318 100644 --- a/Magpie/scripts/benchmark/magpie_bench_remote_compat.sh +++ b/Magpie/scripts/benchmark/magpie_bench_remote_compat.sh @@ -131,6 +131,7 @@ magpie_run_eval_remote_direct() { } local tasks="${MAGPIE_EVAL_TASKS:-gsm8k}" + local include_path="${MAGPIE_EVAL_INCLUDE_PATH:-${MAGPIE_INFERENCEX_ROOT:-$(pwd)}/utils/evals}" local batch_size="${MAGPIE_EVAL_BATCH_SIZE:-auto}" local conc="${CONC:-8}" @@ -147,9 +148,14 @@ magpie_run_eval_remote_direct() { # string prompts instead. Absent env => byte-for-byte the previous behaviour. local base_url="${BENCHMARK_BASE_URL%/}/v1/completions" local model_args="model=${MODEL},base_url=${base_url},num_concurrent=${conc},tokenizer_backend=huggingface,trust_remote_code=true${MAGPIE_EVAL_TOKENIZED_REQUESTS:+,tokenized_requests=${MAGPIE_EVAL_TOKENIZED_REQUESTS}}" + if [[ ! -d "$include_path" ]]; then + echo "[magpie_bench_remote_compat] ERROR evaluator include path is unavailable: $include_path" >&2 + return 1 + fi local -a cmd=( "$py" -m lm_eval --model local-completions + --include_path "$include_path" --tasks "$tasks" --model_args "$model_args" --batch_size "$batch_size" diff --git a/tests/test_lm_eval_runtime.py b/tests/test_lm_eval_runtime.py index aae5d7d..9e0d8ce 100644 --- a/tests/test_lm_eval_runtime.py +++ b/tests/test_lm_eval_runtime.py @@ -382,6 +382,8 @@ def test_owned_evaluator_splits_context_and_output_budget(tmp_path): encoding="utf-8", ) fake_python.chmod(0o755) + include_path = tmp_path / "InferenceX" / "utils" / "evals" + include_path.mkdir(parents=True) env = os.environ.copy() env.update( { @@ -393,6 +395,7 @@ def test_owned_evaluator_splits_context_and_output_budget(tmp_path): "MAGPIE_EVAL_PRIMARY_METRIC": "exact_match,strict-match", "MAGPIE_EVAL_MAX_LENGTH": "2248", "MAGPIE_EVAL_MAX_GEN_TOKENS": "480", + "MAGPIE_EVAL_INCLUDE_PATH": str(include_path), } ) completed = subprocess.run( @@ -417,6 +420,8 @@ def test_owned_evaluator_splits_context_and_output_budget(tmp_path): assert "max_tokens=480" in argv assert "max_tokens=1124" not in argv assert "--log_samples" in argv + assert "--include_path" in argv + assert str(include_path) in argv @pytest.mark.parametrize( From ae29d4724ccae8aaff2df2a6961db740efe40f51 Mon Sep 17 00:00:00 2001 From: Vincent Ouyang Date: Sat, 8 Aug 2026 00:30:19 +0000 Subject: [PATCH 06/11] Harden kernel source and trace evidence resolution --- Magpie/modes/benchmark/targeted_trace.py | 1 + Magpie/targeted_trace/README.md | 14 + Magpie/targeted_trace/postprocess.py | 77 +++++ Magpie/targeted_trace/torch_profiler.py | 9 +- Magpie/tools/amd_kernel_finder/finder.py | 83 ++++- Magpie/tools/amd_kernel_finder/indexer.py | 149 +++++++-- Magpie/tools/amd_kernel_finder/models.py | 44 +++ Magpie/tools/amd_kernel_finder/parser.py | 10 +- Magpie/tools/amd_kernel_finder/searcher.py | 358 ++++++++++++--------- docs/how-to/kernel-source-finder.md | 18 +- tests/test_kernel_source_hardening.py | 287 +++++++++++++++++ tests/test_targeted_trace.py | 83 +++++ 12 files changed, 939 insertions(+), 194 deletions(-) create mode 100644 tests/test_kernel_source_hardening.py diff --git a/Magpie/modes/benchmark/targeted_trace.py b/Magpie/modes/benchmark/targeted_trace.py index b393757..372bf57 100644 --- a/Magpie/modes/benchmark/targeted_trace.py +++ b/Magpie/modes/benchmark/targeted_trace.py @@ -46,6 +46,7 @@ def run_targeted_trace_analysis( "manifest_path": str(targeted_dir / "manifest.json"), "summary_path": str(summary_path), "coverage": coverage, + "evidence_quality": summary["evidence_quality"], "events": summary["events"], "integrity_failures_by_reason": summary[ "integrity_failures_by_reason" diff --git a/Magpie/targeted_trace/README.md b/Magpie/targeted_trace/README.md index 6278ad8..5b0de5d 100644 --- a/Magpie/targeted_trace/README.md +++ b/Magpie/targeted_trace/README.md @@ -29,6 +29,13 @@ and per-shard file/checksum receipts. Unsupported schema versions fail fast. Postprocessing reads one JSONL line at a time and reports corrupt or missing tails; it never silently skips them. +`coverage` is loss accounting for serialized records, not a claim that the shard +has representative launch semantics. `summary.json.evidence_quality` separately +reports the record coverage fraction, missing phase/source/grid/shape fields, +correlation availability, and whether semantic coverage can be claimed. A cap or +sampling drop makes that claim false even when every retained shard is +checksum-valid. + ## Evidence fidelity The explicit `TargetedTraceRecorder` API captures Python-visible Triton and HIP @@ -41,6 +48,13 @@ rank/stage/graph context, correlation IDs, and any tensor metadata present in th trace. Missing fields remain null/empty with warnings; symbol/count/order is not treated as a globally stable CPU-to-GPU join. +Postprocessing never synthesizes a cross-event join. If phase, launch source, +grid, tensor shapes, or the required Torch-profiler correlation key is absent, +the artifact remains valid diagnostic data but its evidence quality is +`resolution_status: unresolved`, `evidence_class: diagnostic_only`, and +`semantic_coverage_claimed: false`. Consumers must not turn such records into a +patchable source binding or evaluator evidence. + Sampling uses only `{run_seed, stable_event_key}` through SHA-256. Python's process-randomized `hash()` and mutable PRNG state are not used. diff --git a/Magpie/targeted_trace/postprocess.py b/Magpie/targeted_trace/postprocess.py index 6764cfb..fefeb14 100644 --- a/Magpie/targeted_trace/postprocess.py +++ b/Magpie/targeted_trace/postprocess.py @@ -237,8 +237,18 @@ def postprocess_trace_dir( "by_kind": {}, "by_rank": {}, } + semantic_missing = { + "phase": 0, + "source": 0, + "grid": 0, + "shape": 0, + "correlation": 0, + } + complete_semantic_records = 0 + torch_profiler_records = 0 def observe(record: TargetedTraceRecord) -> None: + nonlocal complete_semantic_records, torch_profiler_records target_id = record.identity.target_id aggregates["by_target"][target_id] = ( aggregates["by_target"].get(target_id, 0) + 1 @@ -249,6 +259,24 @@ def observe(record: TargetedTraceRecord) -> None: rank = str(record.context.rank) aggregates["by_rank"][rank] = aggregates["by_rank"].get(rank, 0) + 1 + missing = [] + if not record.context.stage or record.context.stage.lower() == "unknown": + missing.append("phase") + if record.semantics.source is None: + missing.append("source") + if record.runtime.grid is None and record.semantics.python_grid is None: + missing.append("grid") + if not record.semantics.tensors: + missing.append("shape") + if record.kind == "torch_profiler_kernel": + torch_profiler_records += 1 + if record.runtime.correlation_id is None: + missing.append("correlation") + for field_name in missing: + semantic_missing[field_name] += 1 + if not missing: + complete_semantic_records += 1 + shard_paths = sorted((trace_dir / "shards").glob("*.jsonl")) if not shard_paths: shard_paths = sorted(trace_dir.glob("*.jsonl")) @@ -303,6 +331,54 @@ def observe(record: TargetedTraceRecord) -> None: reason = "other" integrity_failures[reason] = integrity_failures.get(reason, 0) + 1 + seen = int(coverage["seen"]) + written = int(coverage["written"]) + dropped = int(coverage["dropped"]) + record_coverage_fraction = written / seen if seen else 0.0 + lossless_record_coverage = seen > 0 and dropped == 0 and written == seen + complete_semantic_coverage = ( + written > 0 and complete_semantic_records == written + ) + semantic_coverage_claimed = ( + not issues and lossless_record_coverage and complete_semantic_coverage + ) + unresolved_reasons = [] + if not seen: + unresolved_reasons.append("no_records") + if dropped: + unresolved_reasons.extend( + f"dropped:{reason}" + for reason in sorted(coverage["dropped_by_reason"]) + ) + unresolved_reasons.extend( + f"missing:{field_name}" + for field_name, count in semantic_missing.items() + if count + ) + if issues: + unresolved_reasons.append("integrity_validation_failed") + + evidence_quality = { + "evidence_class": "diagnostic_only", + "resolution_status": ( + "resolved" if semantic_coverage_claimed else "unresolved" + ), + "semantic_coverage_claimed": semantic_coverage_claimed, + "record_coverage_fraction": record_coverage_fraction, + "lossless_record_coverage": lossless_record_coverage, + "records_evaluated": written, + "records_with_complete_semantics": complete_semantic_records, + "missing_by_field": semantic_missing, + # The postprocessor never synthesizes a CPU/source-to-GPU join. A + # correlation ID merely reports that a future trusted consumer has a + # join key available. + "cross_event_join": "not_performed", + "join_eligible_records": ( + torch_profiler_records - semantic_missing["correlation"] + ), + "unresolved_reasons": unresolved_reasons, + } + summary: Dict[str, Any] = { "schema_name": manifest.schema_name if manifest else None, "schema_version": manifest.schema_version if manifest else None, @@ -310,6 +386,7 @@ def observe(record: TargetedTraceRecord) -> None: "valid": not issues and all(item.valid for item in validations), "streaming": True, "coverage": coverage, + "evidence_quality": evidence_quality, "events": aggregates, "integrity_failures_by_reason": dict(sorted(integrity_failures.items())), "shards": [item.to_dict() for item in validations], diff --git a/Magpie/targeted_trace/torch_profiler.py b/Magpie/targeted_trace/torch_profiler.py index baffb45..c78e216 100644 --- a/Magpie/targeted_trace/torch_profiler.py +++ b/Magpie/targeted_trace/torch_profiler.py @@ -304,12 +304,13 @@ def get_writer(rank: int, pid: int) -> TraceShardWriter: ) graph_id = _first(args, ["graph id", "graph_id", "cuda graph id"]) execution_mode = "graph" if graph_id is not None else "unknown" + stage = _stage(event, trace_path) for target in matches: base_parts = { "target_id": target.target_id, "symbol": symbol, "rank": rank, - "stage": _stage(event, trace_path), + "stage": stage, "grid": list(runtime.grid) if runtime.grid else None, "block": list(runtime.block) if runtime.block else None, "tensors": [ @@ -329,6 +330,10 @@ def get_writer(rank: int, pid: int) -> TraceShardWriter: warnings = list(tensor_warnings) if target.source is None: warnings.append("torch_profiler_missing_launch_source") + if stage == "unknown": + warnings.append("torch_profiler_missing_phase") + if runtime.grid is None: + warnings.append("torch_profiler_missing_launch_grid") if runtime.correlation_id is None: warnings.append("torch_profiler_missing_runtime_correlation") try: @@ -351,7 +356,7 @@ def get_writer(rank: int, pid: int) -> TraceShardWriter: framework_version=framework_version, rank=rank, pid=pid, - stage=_stage(event, trace_path), + stage=stage, execution_mode=execution_mode, graph_id=( str(graph_id) if graph_id is not None else None diff --git a/Magpie/tools/amd_kernel_finder/finder.py b/Magpie/tools/amd_kernel_finder/finder.py index 259323f..7f63a71 100644 --- a/Magpie/tools/amd_kernel_finder/finder.py +++ b/Magpie/tools/amd_kernel_finder/finder.py @@ -22,7 +22,7 @@ from .models import KernelKind, KernelSourceInfo from .parser import KernelNameParser from .searcher import KernelSourceSearcher -from .repo_config import GITHUB_URL_TEMPLATES, RepoDiscovery +from .repo_config import GITHUB_URL_TEMPLATES, RepoDiscovery, detect_repo_type from .indexer import KernelIndex from .repo_manager import RepoManager @@ -136,15 +136,35 @@ def search(self, kernel_name: str) -> KernelSourceInfo: parsed = self.parser.parse(kernel_name) category = self.parser.classify_category(kernel_name) - # Skip index for kernel types where it's unreliable (CK_TILE, HIP_CPP) - # These have complex mangled names that index doesn't handle well - skip_index_kinds = {KernelKind.CK_TILE, KernelKind.HIP_CPP, KernelKind.TENSILE_GEMM} + # The index contains source-level identifiers. Mangled/generated and + # unknown names must go through their kind-aware searchers instead of + # risking a cross-repository prefix guess. + skip_index_kinds = { + KernelKind.CK_TILE, + KernelKind.HIP_CPP, + KernelKind.TENSILE_GEMM, + KernelKind.AITER, + KernelKind.UNKNOWN, + KernelKind.ANNOTATION, + } # Try index lookup first for fast results (for supported kernel types) source_match = None if self.use_index and self.index and parsed.kind not in skip_index_kinds: - index_result = self.index.lookup(kernel_name) - if index_result: + expected_kinds = {"triton_jit"} + expected_repos = { + repo_name + for repo_name in (detect_repo_type(path) for path in self.repos) + if repo_name + } + if parsed.kind == KernelKind.INDUCTOR: + expected_repos &= {"pytorch"} + index_result = self.index.lookup( + kernel_name, + expected_repo_names=expected_repos, + expected_kinds=expected_kinds, + ) + if index_result and self._indexed_source_is_trusted(index_result): from .models import SourceMatch source_match = SourceMatch( file_path=index_result.file_path, @@ -156,9 +176,23 @@ def search(self, kernel_name: str) -> KernelSourceInfo: # Fall back to searcher if index miss or skipped if not source_match: source_match = self.searcher.search_source(parsed) + + source_resolved = bool(source_match and source_match.is_resolved) + resolution_status = ( + source_match.resolution_status if source_match else "unresolved" + ) + resolution_error = ( + source_match.error if source_match and source_match.error + else "source_not_found" if not source_resolved else "" + ) + resolved_source = source_match if source_resolved else None # Search for test - test_match = self.searcher.search_test(parsed, source_match) + test_match = ( + self.searcher.search_test(parsed, resolved_source) + if source_resolved + else None + ) # Search for PyTorch eager baseline reference. We hand over the # already-computed test_match + category so the searcher does not @@ -166,22 +200,22 @@ def search(self, kernel_name: str) -> KernelSourceInfo: # and scans for `run_torch` / `ref_*` / `torch_*` / etc. by # convention. No per-kernel symbol tables involved. baseline_ref = self.searcher.search_baseline_ref( - parsed, source_match, category=category, test_match=test_match, - ) + parsed, resolved_source, category=category, test_match=test_match, + ) if source_resolved else None # Search for canonical Triton implementation reference (independent of # the eager baseline -- a kernel can have both, neither, or only one). # Discovery is also convention-driven: category -> triton-kernels dir # -> ripgrep for `@triton.jit`. triton_ref = self.searcher.search_triton_ref( - parsed, source_match, category=category, - ) + parsed, resolved_source, category=category, + ) if source_resolved else None # Build upstream URL upstream_url = "" - if source_match and source_match.repo_name in GITHUB_URL_TEMPLATES: - upstream_url = GITHUB_URL_TEMPLATES[source_match.repo_name].format( - path=source_match.file_path + if resolved_source and resolved_source.repo_name in GITHUB_URL_TEMPLATES: + upstream_url = GITHUB_URL_TEMPLATES[resolved_source.repo_name].format( + path=resolved_source.file_path ) # Build notes with more details @@ -190,12 +224,15 @@ def search(self, kernel_name: str) -> KernelSourceInfo: notes = f"{notes}; baseline: {baseline_ref.notes}" if notes else f"baseline: {baseline_ref.notes}" if triton_ref and triton_ref.notes: notes = f"{notes}; triton: {triton_ref.notes}" if notes else f"triton: {triton_ref.notes}" + if resolution_error: + marker = f"source_resolution={resolution_status}:{resolution_error}" + notes = f"{notes}; {marker}" if notes else marker return KernelSourceInfo( kind=parsed.kind.value, category=category.value, - source_repo=source_match.repo_name if source_match else "", - source_file=source_match.display_path if source_match else "", + source_repo=resolved_source.repo_name if resolved_source else "", + source_file=resolved_source.display_path if resolved_source else "", upstream_url=upstream_url, test_file=test_match.display_path if test_match else "", test_cmd=test_match.test_cmd if test_match else "", @@ -205,7 +242,21 @@ def search(self, kernel_name: str) -> KernelSourceInfo: triton_ref_file=triton_ref.display_path if triton_ref else "", triton_ref_symbol=triton_ref.ref_symbol if triton_ref else "", notes=notes, + source_resolution=resolution_status, + source_error=resolution_error, ) + + def _indexed_source_is_trusted(self, definition) -> bool: + """Verify a cached definition still belongs to a supplied source root.""" + + try: + definition_root = Path(definition.repo_path).resolve() + trusted_roots = {Path(path).resolve() for path in self.repos} + source_path = (definition_root / definition.file_path).resolve() + source_path.relative_to(definition_root) + except (OSError, ValueError): + return False + return definition_root in trusted_roots and source_path.is_file() def _build_notes(self, parsed) -> str: """Build notes from parsed information.""" diff --git a/Magpie/tools/amd_kernel_finder/indexer.py b/Magpie/tools/amd_kernel_finder/indexer.py index 0031034..53cb7a0 100644 --- a/Magpie/tools/amd_kernel_finder/indexer.py +++ b/Magpie/tools/amd_kernel_finder/indexer.py @@ -16,7 +16,7 @@ import re from dataclasses import asdict, dataclass from pathlib import Path -from typing import Dict, List, Optional +from typing import Collection, Dict, List, Optional from .repo_config import detect_repo_type @@ -232,30 +232,143 @@ def _build_name_index(self) -> None: self.name_to_keys[name] = [] self.name_to_keys[name].append(key) - def lookup(self, kernel_name: str) -> Optional[KernelDefinition]: - """Look up a kernel by name.""" + def lookup( + self, + kernel_name: str, + *, + expected_repo_names: Optional[Collection[str]] = None, + expected_kinds: Optional[Collection[str]] = None, + ) -> Optional[KernelDefinition]: + """Look up a kernel without guessing across provenance boundaries. + + ``expected_repo_names`` and ``expected_kinds`` are supplied by the + parser/finder boundary. They keep an identical symbol in an unrelated + repository or language from being accepted as source evidence. Any + empty, mangled, ambiguous, or otherwise unparseable name fails closed. + """ function_name = self._extract_function_name(kernel_name) - - if function_name in self.name_to_keys: - keys = self.name_to_keys[function_name] - if keys: - return self.index[keys[0]] - + if not function_name: + logger.debug("Kernel index rejected unparseable name: %r", kernel_name) + return None + + repo_names = self._normalize_filter(expected_repo_names) + kinds = self._normalize_filter(expected_kinds) + + exact = self._compatible_definitions( + self.name_to_keys.get(function_name, []), + repo_names=repo_names, + kinds=kinds, + ) + if len(exact) == 1: + return exact[0] + if len(exact) > 1: + logger.warning( + "Kernel index rejected ambiguous exact match for %r (%d candidates)", + kernel_name, + len(exact), + ) + return None + + prefix_keys: List[str] = [] for name, keys in self.name_to_keys.items(): - if function_name.startswith(name) or name.startswith(function_name): - return self.index[keys[0]] - + if self._is_token_boundary_prefix(function_name, name): + prefix_keys.extend(keys) + prefix_matches = self._compatible_definitions( + prefix_keys, + repo_names=repo_names, + kinds=kinds, + ) + if len(prefix_matches) == 1: + return prefix_matches[0] + if prefix_matches: + logger.warning( + "Kernel index rejected ambiguous prefix match for %r (%d candidates)", + kernel_name, + len(prefix_matches), + ) return None - - def _extract_function_name(self, kernel_name: str) -> str: - name = kernel_name.replace(".kd", "") + + @staticmethod + def _normalize_filter( + values: Optional[Collection[str]], + ) -> Optional[frozenset[str]]: + if values is None: + return None + normalized = frozenset( + str(value).strip() for value in values if str(value).strip() + ) + return normalized + + def _compatible_definitions( + self, + keys: Collection[str], + *, + repo_names: Optional[frozenset[str]], + kinds: Optional[frozenset[str]], + ) -> List[KernelDefinition]: + matches: List[KernelDefinition] = [] + seen = set() + for key in keys: + definition = self.index.get(key) + if definition is None: + continue + if repo_names is not None and definition.repo_name not in repo_names: + continue + if kinds is not None and definition.kind not in kinds: + continue + identity = ( + definition.repo_path, + definition.file_path, + definition.name, + definition.kind, + ) + if identity in seen: + continue + seen.add(identity) + matches.append(definition) + return matches + + @staticmethod + def _is_token_boundary_prefix(query: str, indexed_name: str) -> bool: + """Return true only for a unique underscore-delimited extension. + + The previous symmetric ``startswith`` accepted an empty query and + returned whichever definition happened to be inserted first. It also + let short incidental prefixes win. Configured Triton symbols use + underscore-delimited suffixes, so that is the only inexact relation we + retain. + """ + + if not query or not indexed_name or query == indexed_name: + return False + return query.startswith(f"{indexed_name}_") + + def _extract_function_name(self, kernel_name: str) -> Optional[str]: + if not isinstance(kernel_name, str): + return None + name = kernel_name.strip() + if not name: + return None + if name.endswith(".k.d"): + name = name[:-4] + elif name.endswith(".kd"): + name = name[:-3] + name = re.sub(r"\s*\[clone \.kd\]\s*$", "", name).strip() + # The index contains source-level identifiers, not a demangler. An + # Itanium symbol must be resolved by the kind-aware searcher instead + # of being truncated to an empty string and prefix-matched. + if name.startswith("_Z"): + return None + if not re.fullmatch(r"[A-Za-z_]\w*", name): + return None parts = name.split("_") - + for i, part in enumerate(parts): if part and (part[0].isupper() or part.isdigit() or part in ("bf16", "fp16", "fp32", "int8")): - return "_".join(parts[:i]) - + candidate = "_".join(parts[:i]) + return candidate or None + return name def get_all_definitions(self, kind: str = None) -> List[KernelDefinition]: diff --git a/Magpie/tools/amd_kernel_finder/models.py b/Magpie/tools/amd_kernel_finder/models.py index 771042a..0534799 100644 --- a/Magpie/tools/amd_kernel_finder/models.py +++ b/Magpie/tools/amd_kernel_finder/models.py @@ -86,6 +86,12 @@ class KernelSourceInfo: # Additional context notes: str = "" + + # Machine-readable source-resolution outcome. A missing source is never + # represented by a plausible-looking placeholder path: consumers must + # inspect this status before treating the row as patchable provenance. + source_resolution: str = "unresolved" + source_error: str = "" def to_list(self) -> List[str]: """Convert to list of values for CSV output.""" @@ -103,6 +109,8 @@ def to_list(self) -> List[str]: self.triton_ref_file, self.triton_ref_symbol, self.notes, + self.source_resolution, + self.source_error, ] @staticmethod @@ -122,6 +130,8 @@ def csv_headers() -> List[str]: "triton_ref_file", "triton_ref_symbol", "notes", + "source_resolution", + "source_error", ] @@ -148,10 +158,44 @@ class SourceMatch: line_number: Optional[int] = None repo_name: str = "" repo_var: str = "" # e.g., $TRITON_DIR + resolution_status: str = "resolved" + error: str = "" + + def __post_init__(self) -> None: + allowed = {"resolved", "unresolved", "unsupported"} + if self.resolution_status not in allowed: + raise ValueError( + f"invalid source resolution status: {self.resolution_status}" + ) + if self.resolution_status == "resolved" and not self.file_path: + raise ValueError("resolved source match requires a file path") + if self.resolution_status != "resolved" and self.file_path: + raise ValueError("unresolved source match cannot carry a file path") + if self.resolution_status != "resolved" and not self.error: + raise ValueError("unresolved source match requires an error code") + if self.resolution_status == "resolved" and self.error: + raise ValueError("resolved source match cannot carry an error") + + @property + def is_resolved(self) -> bool: + return self.resolution_status == "resolved" + + @classmethod + def unresolved( + cls, + error: str, + *, + status: str = "unresolved", + ) -> "SourceMatch": + """Return a structured fail-closed result with no source path.""" + + return cls(file_path="", resolution_status=status, error=error) @property def display_path(self) -> str: """Return path with repo variable prefix.""" + if not self.file_path: + return "" if self.repo_var: return f"{self.repo_var}/{self.file_path}" return self.file_path diff --git a/Magpie/tools/amd_kernel_finder/parser.py b/Magpie/tools/amd_kernel_finder/parser.py index f6c3dec..f01902a 100644 --- a/Magpie/tools/amd_kernel_finder/parser.py +++ b/Magpie/tools/amd_kernel_finder/parser.py @@ -19,7 +19,10 @@ class KernelNameParser: # Patterns for classification TRITON_PATTERN = re.compile(r'^[_a-zA-Z][\w]*\.kd$|\.k\.d$') TENSILE_PATTERN = re.compile(r'^Cijk_') - CK_PATTERN = re.compile(r'^_ZN7ck_tile|ck_tile::') + CK_PATTERN = re.compile( + r'^_ZN(?:7ck_tile|2ck)|(?:^|[\s:<])(?:ck_tile|ck)::|' + r'Gridwise(?:Gemm|MoeGemm)|kernel_(?:gemm_xdl|moe_gemm)' + ) ATEN_PATTERN = re.compile(r'void at::native::') INDUCTOR_PATTERN = re.compile(r'triton_\w+_fused_') HIPBLASLT_PATTERN = re.compile(r'wvSplitK|wvSpltK|DeviceGemmWmma') @@ -232,6 +235,11 @@ def _parse_ck_tile(self, name: str) -> ParsedKernelName: (r'Rmsnorm2dFwd', 'Rmsnorm2dFwd'), (r'Fmha', 'Fmha'), (r'Softmax', 'Softmax'), + (r'MoeSorting', 'MoeSorting'), + (r'GridwiseMoeGemm', 'MoeGemm'), + (r'kernel_moe_gemm', 'MoeGemm'), + (r'GridwiseGemm', 'Gemm'), + (r'kernel_gemm_xdl', 'Gemm'), (r'Gemm', 'Gemm'), ] diff --git a/Magpie/tools/amd_kernel_finder/searcher.py b/Magpie/tools/amd_kernel_finder/searcher.py index f900f4e..e05df81 100644 --- a/Magpie/tools/amd_kernel_finder/searcher.py +++ b/Magpie/tools/amd_kernel_finder/searcher.py @@ -820,7 +820,7 @@ def search_test(self, parsed: ParsedKernelName, source: Optional[SourceMatch] = elif parsed.kind == KernelKind.TENSILE_GEMM: return self._search_tensile_test(parsed) elif parsed.kind == KernelKind.CK_TILE: - return self._search_ck_test(parsed) + return self._search_ck_test(parsed, source) elif parsed.kind == KernelKind.ATEN_NATIVE: return self._search_aten_test(parsed) elif parsed.kind == KernelKind.HIP_CPP: @@ -1227,78 +1227,72 @@ def _run_ripgrep(self, pattern: str, search_path: str, return [] def _search_triton_source(self, parsed: ParsedKernelName) -> Optional[SourceMatch]: - """Search for Triton JIT kernel source.""" - function_name = parsed.function_name - - # Known kernel mappings for common kernels - # $TRITON_KERNELS_DIR = triton/python/triton_kernels/ - known_mappings = { - "_matmul_ogs": ("triton_kernels/matmul_details/_matmul.py", "$TRITON_KERNELS_DIR"), - "_matmul": ("triton_kernels/matmul_details/_matmul.py", "$TRITON_KERNELS_DIR"), - "_reduce": ("triton_kernels/reduce.py", "$TRITON_KERNELS_DIR"), - "kernel_unified_attention": ("vllm/v1/attention/ops/triton_unified_attention.py", "$VLLM_DIR"), - "_topk_forward": ("triton_kernels/topk_details/_topk_forward.py", "$TRITON_KERNELS_DIR"), - "_topk_backward": ("triton_kernels/topk_details/_topk_backward.py", "$TRITON_KERNELS_DIR"), - "_bitmatrix_metadata": ("triton_kernels/tensor_details/", "$TRITON_KERNELS_DIR"), - "_ragged_tensor_metadata": ("triton_kernels/tensor_details/", "$TRITON_KERNELS_DIR"), - "_sum_bitmatrix_rows": ("triton_kernels/tensor_details/", "$TRITON_KERNELS_DIR"), - "_fused_add_rmsnorm": ("triton_kernels/swiglu_details/", "$TRITON_KERNELS_DIR"), - "_swiglu": ("triton_kernels/swiglu_details/", "$TRITON_KERNELS_DIR"), - "_compaction": ("triton_kernels/compaction_details/", "$TRITON_KERNELS_DIR"), - } - - # Check known mappings first - for key, (path, repo_var) in known_mappings.items(): - if key in function_name: - return SourceMatch( - file_path=path, - symbol=function_name, - repo_name="triton_kernels", - repo_var=repo_var, - ) - - # Search patterns - patterns = [ - f"def {function_name}", - f"@triton.jit.*\\n.*def {function_name}", - f'def {function_name}\\(', - ] - - # Search in triton repos - triton_path = self._repo_var_map.get("$TRITON_DIR") - if triton_path: - for pattern in patterns: - files = self._run_ripgrep(pattern, triton_path, ["py"]) - if files: - rel_path = os.path.relpath(files[0], triton_path) - return SourceMatch( - file_path=rel_path, - symbol=function_name, - repo_name="triton", - repo_var="$TRITON_DIR", + """Search trusted roots for an exact Triton function definition. + + Repository roots come only from the caller-supplied ``repos`` list. + In particular, locked vLLM and AITER roots are searched explicitly; + ambient checkout guesses and placeholder paths are forbidden. + """ + + function_name = parsed.function_name.strip() + if not re.fullmatch(r"[A-Za-z_]\w*", function_name): + return SourceMatch.unresolved("triton_function_name_unparseable") + + pattern = rf"^\s*def\s+{re.escape(function_name)}\s*\(" + roots = ( + ("$VLLM_DIR", "vllm"), + ("$AITER_DIR", "aiter"), + ("$TRITON_DIR", "triton"), + ("$ROCM_LIBRARIES_DIR", "rocm-libraries"), + ) + matches = [] + for repo_var, repo_name in roots: + root_value = self._repo_var_map.get(repo_var) + if not root_value: + continue + root = Path(root_value).resolve() + files = self._search_files( + pattern, + str(root), + ["py"], + max_results=20, + ) + for file_name in files: + candidate = Path(file_name).resolve() + try: + relative = candidate.relative_to(root) + except ValueError: + logger.warning( + "Rejected Triton source outside trusted root %s: %s", + root, + candidate, ) - - # Search in rocm-libraries (triton_kernels might be there) - rocm_libs = self._repo_var_map.get("$ROCM_LIBRARIES_DIR") - if rocm_libs: - for pattern in patterns: - files = self._run_ripgrep(pattern, rocm_libs, ["py"]) - if files: - rel_path = os.path.relpath(files[0], rocm_libs) - return SourceMatch( - file_path=rel_path, + continue + if not candidate.is_file(): + continue + matches.append( + SourceMatch( + file_path=relative.as_posix(), symbol=function_name, - repo_name="rocm-libraries", - repo_var="$ROCM_LIBRARIES_DIR", + repo_name=repo_name, + repo_var=repo_var, ) - - # Default fallback for triton kernels - return SourceMatch( - file_path="(search in triton_kernels or vllm)", - symbol=function_name, - repo_name="triton", - repo_var="$TRITON_DIR", - ) + ) + + unique = { + (match.repo_var, match.file_path, match.symbol): match + for match in matches + } + if len(unique) == 1: + return next(iter(unique.values())) + if len(unique) > 1: + logger.warning( + "Rejected ambiguous Triton source for %s: %s", + function_name, + sorted(f"{item.repo_var}/{item.file_path}" for item in unique.values()), + ) + return SourceMatch.unresolved("triton_source_ambiguous") + return SourceMatch.unresolved("triton_source_not_found") def _search_tensile_source(self, parsed: ParsedKernelName) -> Optional[SourceMatch]: """Search for Tensile GEMM source (logic YAML files).""" @@ -1319,16 +1313,52 @@ def _search_tensile_source(self, parsed: ParsedKernelName) -> Optional[SourceMat return None def _search_ck_source(self, parsed: ParsedKernelName) -> Optional[SourceMatch]: - """Search for Composable Kernel source.""" + """Search a materialized, caller-supplied Composable Kernel tree.""" + + roots = [] + provenance_errors = [] rocm_libs = self._repo_var_map.get("$ROCM_LIBRARIES_DIR") - if not rocm_libs: - return None - - ck_path = Path(rocm_libs) / "projects/composablekernel" - if not ck_path.exists(): - return None - - # Map operation name to directory and kernel file + if rocm_libs: + ck_root = Path(rocm_libs) / "projects/composablekernel" + if self._is_materialized_ck_root(ck_root): + roots.append( + ( + ck_root, + "projects/composablekernel", + "rocm-libraries", + "$ROCM_LIBRARIES_DIR", + ) + ) + else: + provenance_errors.append("ck_source_tree_unavailable") + + aiter_root_value = self._repo_var_map.get("$AITER_DIR") + if aiter_root_value: + aiter_root = Path(aiter_root_value) + submodule_rel = Path("3rdparty/composable_kernel") + ck_root = aiter_root / submodule_rel + if self._is_materialized_ck_root(ck_root): + roots.append( + ( + ck_root, + submodule_rel.as_posix(), + "aiter", + "$AITER_DIR", + ) + ) + elif self._declares_submodule(aiter_root, submodule_rel): + provenance_errors.append("ck_submodule_not_materialized") + + if not roots: + error = ( + "ck_submodule_not_materialized" + if "ck_submodule_not_materialized" in provenance_errors + else "ck_source_tree_unavailable" + ) + return SourceMatch.unresolved(error, status="unsupported") + + # Map operation name to an exact source file. A generic CK directory + # is not source provenance and is therefore never emitted as resolved. op_name = parsed.function_name.lower() op_info = { "rmsnorm2dfwd": ("rmsnorm2d", "kernel/rmsnorm2d_fwd_kernel.hpp"), @@ -1340,34 +1370,51 @@ def _search_ck_source(self, parsed: ParsedKernelName) -> Optional[SourceMatch]: "moe": ("moe_sorting_topk", "kernel/moe_sorting_kernel.hpp"), } + matches = [] for op_key, (op_dir, kernel_file) in op_info.items(): - if op_key in op_name: - # Try specific kernel file first - kernel_path = f"projects/composablekernel/include/ck_tile/ops/{op_dir}/{kernel_file}" - if (Path(rocm_libs) / kernel_path).exists(): - return SourceMatch( - file_path=kernel_path, + if op_key not in op_name: + continue + relative_to_ck = Path("include/ck_tile/ops") / op_dir / kernel_file + for root, output_prefix, repo_name, repo_var in roots: + candidate = root / relative_to_ck + if not candidate.is_file(): + continue + matches.append( + SourceMatch( + file_path=(Path(output_prefix) / relative_to_ck).as_posix(), symbol=f"ck_tile::{op_dir}_kernel", - repo_name="rocm-libraries", - repo_var="$ROCM_LIBRARIES_DIR", - ) - # Fall back to directory - op_path = f"projects/composablekernel/include/ck_tile/ops/{op_dir}/" - if (Path(rocm_libs) / op_path).exists(): - return SourceMatch( - file_path=op_path, - symbol=parsed.function_name, - repo_name="rocm-libraries", - repo_var="$ROCM_LIBRARIES_DIR", + repo_name=repo_name, + repo_var=repo_var, ) - - # Generic CK search - return SourceMatch( - file_path="projects/composablekernel/include/ck_tile/ops/", - symbol=parsed.function_name, - repo_name="rocm-libraries", - repo_var="$ROCM_LIBRARIES_DIR", + ) + + if len(matches) == 1: + return matches[0] + if len(matches) > 1: + return SourceMatch.unresolved("ck_source_ambiguous") + return SourceMatch.unresolved("ck_source_not_found") + + @staticmethod + def _is_materialized_ck_root(path: Path) -> bool: + return ( + (path / "include/ck_tile/ops").is_dir() + or (path / "include/ck/tensor_operation").is_dir() ) + + @staticmethod + def _declares_submodule(repo_root: Path, relative_path: Path) -> bool: + gitmodules = repo_root / ".gitmodules" + if not gitmodules.is_file(): + return False + try: + contents = gitmodules.read_text(encoding="utf-8", errors="ignore") + except OSError: + return False + path_pattern = re.compile( + rf"^\s*path\s*=\s*{re.escape(relative_path.as_posix())}\s*$", + re.MULTILINE, + ) + return path_pattern.search(contents) is not None def _search_aten_source(self, parsed: ParsedKernelName) -> Optional[SourceMatch]: """Search for ATen native kernel source.""" @@ -1574,70 +1621,75 @@ def _search_tensile_test(self, parsed: ParsedKernelName) -> Optional[TestMatch]: ) return None - def _search_ck_test(self, parsed: ParsedKernelName) -> Optional[TestMatch]: - """Search for CK tile tests.""" + def _search_ck_test( + self, + parsed: ParsedKernelName, + source: Optional[SourceMatch], + ) -> Optional[TestMatch]: + """Return a CK test only when its materialized path exists.""" + + if source is None or not source.is_resolved: + return None + if source.repo_var == "$AITER_DIR": + relative_root = Path("3rdparty/composable_kernel") + command_root = "$AITER_DIR/3rdparty/composable_kernel" + elif source.repo_var == "$ROCM_LIBRARIES_DIR": + relative_root = Path("projects/composablekernel") + command_root = "$ROCM_LIBRARIES_DIR/projects/composablekernel" + else: + return None + + repo_root_value = self._repo_var_map.get(source.repo_var) + if not repo_root_value: + return None + op_name = parsed.function_name.lower() original_name = parsed.original_name.lower() - - # Map operation to example/test directory - # CK examples are at: projects/composablekernel/example/ck_tile/ + example = None + target = None + args = "" if "rmsnorm" in op_name or "rmsnorm" in original_name: - return TestMatch( - test_file="projects/composablekernel/example/ck_tile/10_rmsnorm2d/", - test_cmd="cd $ROCM_LIBRARIES_DIR/projects/composablekernel/build && cmake --build . -j --target tile_example_rmsnorm2d_fwd && ./bin/tile_example_rmsnorm2d_fwd -m 1024 -n 2048", - repo_var="$ROCM_LIBRARIES_DIR", - ) + example = "10_rmsnorm2d" + target = "tile_example_rmsnorm2d_fwd" + args = " -m 1024 -n 2048" elif "fmha" in op_name or "fmha" in original_name: - return TestMatch( - test_file="projects/composablekernel/example/ck_tile/01_fmha/", - test_cmd="cd $ROCM_LIBRARIES_DIR/projects/composablekernel/build && cmake --build . -j --target tile_example_fmha_fwd && ./bin/tile_example_fmha_fwd", - repo_var="$ROCM_LIBRARIES_DIR", - ) + example = "01_fmha" + target = "tile_example_fmha_fwd" elif "layernorm" in op_name or "layernorm" in original_name: - return TestMatch( - test_file="projects/composablekernel/example/ck_tile/02_layernorm2d/", - test_cmd="cd $ROCM_LIBRARIES_DIR/projects/composablekernel/build && cmake --build . -j --target tile_example_layernorm2d_fwd && ./bin/tile_example_layernorm2d_fwd", - repo_var="$ROCM_LIBRARIES_DIR", - ) + example = "02_layernorm2d" + target = "tile_example_layernorm2d_fwd" elif "gemm" in op_name or "gemm" in original_name: - return TestMatch( - test_file="projects/composablekernel/example/ck_tile/03_gemm/", - test_cmd="cd $ROCM_LIBRARIES_DIR/projects/composablekernel/build && cmake --build . -j --target tile_example_gemm && ./bin/tile_example_gemm", - repo_var="$ROCM_LIBRARIES_DIR", - ) + example = "03_gemm" + target = "tile_example_gemm" elif "topk" in op_name or "softmax" in op_name: - return TestMatch( - test_file="projects/composablekernel/example/ck_tile/09_topk_softmax/", - test_cmd="cd $ROCM_LIBRARIES_DIR/projects/composablekernel/build && cmake --build . -j --target tile_example_topk_softmax && ./bin/tile_example_topk_softmax", - repo_var="$ROCM_LIBRARIES_DIR", - ) - # MoE sorting + MoE FlatMM (top-bottleneck on MI355X gpt-oss/MoE traces). - # Added in refrence_torch follow-up; previously the CK searcher returned - # None for any MoeSorting* / MoeFlatmm* kernel and so the gap_analysis - # CSV had no test entry for ~37%+ of GPU time on MoE workloads. + example = "09_topk_softmax" + target = "tile_example_topk_softmax" elif "moesorting" in op_name or "moe_sorting" in op_name \ or "moesorting" in original_name or "moe_sorting" in original_name: - return TestMatch( - test_file="projects/composablekernel/example/ck_tile/13_moe_sorting/", - test_cmd="cd $ROCM_LIBRARIES_DIR/projects/composablekernel/build && cmake --build . -j --target tile_example_moe_sorting && ./bin/tile_example_moe_sorting", - repo_var="$ROCM_LIBRARIES_DIR", - ) + example = "13_moe_sorting" + target = "tile_example_moe_sorting" elif "moeflatmm" in op_name or "moe_flatmm" in op_name \ or "moeflatmm" in original_name or "moe_flatmm" in original_name \ or "flatmm" in original_name: - return TestMatch( - test_file="projects/composablekernel/example/ck_tile/18_flatmm/", - test_cmd="cd $ROCM_LIBRARIES_DIR/projects/composablekernel/build && cmake --build . -j --target tile_example_flatmm && ./bin/tile_example_flatmm", - repo_var="$ROCM_LIBRARIES_DIR", - ) + example = "18_flatmm" + target = "tile_example_flatmm" elif "fused_moe" in op_name or "fused_moe" in original_name: - return TestMatch( - test_file="projects/composablekernel/example/ck_tile/15_fused_moe/", - test_cmd="cd $ROCM_LIBRARIES_DIR/projects/composablekernel/build && cmake --build . -j --target tile_example_fused_moe && ./bin/tile_example_fused_moe", - repo_var="$ROCM_LIBRARIES_DIR", - ) + example = "15_fused_moe" + target = "tile_example_fused_moe" - return None + if not example or not target: + return None + test_relative = relative_root / "example/ck_tile" / example + if not (Path(repo_root_value) / test_relative).is_dir(): + return None + return TestMatch( + test_file=f"{test_relative.as_posix()}/", + test_cmd=( + f"cd {command_root}/build && cmake --build . -j --target " + f"{target} && ./bin/{target}{args}" + ), + repo_var=source.repo_var, + ) def _search_aten_test(self, parsed: ParsedKernelName) -> Optional[TestMatch]: """Search for ATen native tests.""" diff --git a/docs/how-to/kernel-source-finder.md b/docs/how-to/kernel-source-finder.md index 7f15f52..d57ce82 100644 --- a/docs/how-to/kernel-source-finder.md +++ b/docs/how-to/kernel-source-finder.md @@ -7,7 +7,7 @@ myst: # Find kernel sources with Magpie -When gap analysis identifies the GPU kernels dominating your benchmark runtime, the kernel source finder maps those mangled kernel names back to their human-readable source files and runnable test commands. It clones the relevant upstream repositories automatically, parses the kernel name to determine its type and origin, and writes source file paths, GitHub URLs, and test commands directly into the gap analysis CSV. Use this feature to quickly locate the code behind a bottleneck kernel and reproduce it in isolation. +When gap analysis identifies the GPU kernels dominating your benchmark runtime, the kernel source finder maps those mangled kernel names back to their human-readable source files and runnable test commands. It parses the kernel name, searches caller-supplied or explicitly cloned repository roots, and writes source file paths, GitHub URLs, and test commands directly into the gap analysis CSV. Source resolution is fail-closed: a missing or ambiguous definition leaves the source fields empty and records a machine-readable error instead of emitting a plausible placeholder. ## Pipeline overview @@ -93,16 +93,22 @@ ParsedKernelName( The searcher looks up source files using: - **ripgrep**: Fast regex search across repos -- **Static mappings**: Known paths for Tensile, CK Tile examples +- **Verified mappings**: Known paths are emitted only when the file exists in a supplied root - **Kernel index**: Pre-built index for faster lookups +For reproducible benchmark runs, set `gap_analysis.kernel_source_repos` to the exact locked vLLM and AITER roots. Triton source lookup searches these roots explicitly, along with supplied Triton and ROCm Libraries roots. It does not consult ambient `VLLM_DIR` or `AITER_DIR` values and does not fall back to a text placeholder. + +The index accepts only a unique source-level identifier compatible with the parser's repository and kernel-kind constraints. Empty or mangled names, cross-repository matches, and ambiguous prefix matches remain unresolved. + +Composable Kernel source embedded in AITER is usable only when `3rdparty/composable_kernel` is materialized. If the gitlink is present but its source tree is absent, the row reports `source_resolution=unsupported` with `source_error=ck_submodule_not_materialized`; a generic CK directory is never treated as an exact source file. + ### Step 4: Generate output Results are written to `gap_analysis.csv`: ```text -Name,Calls,Self CUDA total (us),...,kind,category,source_repo,source_file,upstream_url,test_file,test_cmd,notes -_matmul_ogs_NNT_bf16.kd,24552,5631747.87,...,triton_jit,gemm,triton_kernels,$TRITON_KERNELS_DIR/matmul_details/_matmul.py,https://github.com/...,$TRITON_KERNELS_DIR/tests/test_matmul.py,cd $TRITON_KERNELS_DIR && pytest tests/test_matmul.py -v,dtype=bf16 +Name,Calls,Self CUDA total (us),...,kind,category,source_repo,source_file,...,notes,source_resolution,source_error +kernel_paged_attention_2d.kd,12288,1679000.00,...,triton_jit,attention,vllm,$VLLM_DIR/vllm/v1/attention/ops/chunked_prefill_paged_decode.py,...,,resolved, ``` ## Usage @@ -132,6 +138,10 @@ The following fields are added to `gap_analysis.csv` when kernel source finding | `test_file` | Path to test file | | `test_cmd` | Command to run tests | | `notes` | Additional info (dtype, tile sizes, etc.) | +| `source_resolution` | `resolved`, `unresolved`, or `unsupported` | +| `source_error` | Stable fail-closed reason such as `triton_source_not_found`, `triton_source_ambiguous`, or `ck_submodule_not_materialized` | + +Only `source_resolution=resolved` is patchable source evidence. Empty source fields are intentional for the other statuses. ### Path variables diff --git a/tests/test_kernel_source_hardening.py b/tests/test_kernel_source_hardening.py new file mode 100644 index 0000000..0124455 --- /dev/null +++ b/tests/test_kernel_source_hardening.py @@ -0,0 +1,287 @@ +from pathlib import Path + +from Magpie.tools.amd_kernel_finder.finder import KernelSourceFinder +from Magpie.tools.amd_kernel_finder.indexer import KernelDefinition, KernelIndex +from Magpie.tools.amd_kernel_finder.models import KernelKind, KernelSourceInfo +from Magpie.tools.amd_kernel_finder.parser import KernelNameParser + + +def _definition( + root: Path, + *, + name: str, + repo_name: str, + kind: str, + relative: str, +) -> KernelDefinition: + path = root / relative + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(f"def {name}():\n pass\n", encoding="utf-8") + return KernelDefinition( + name=name, + file_path=relative, + repo_name=repo_name, + repo_path=str(root), + kind=kind, + ) + + +def _load_definitions(index: KernelIndex, definitions: list[KernelDefinition]) -> None: + for offset, definition in enumerate(definitions): + index.index[f"definition-{offset}"] = definition + index._build_name_index() + + +def _write_triton_source(root: Path, relative: str, function_name: str) -> None: + path = root / relative + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + "import triton\n\n" + "@triton.jit\n" + f"def {function_name}(x):\n" + " return x\n", + encoding="utf-8", + ) + + +def _make_vllm_root(tmp_path: Path) -> Path: + root = tmp_path / "vllm-v0.19.1" + (root / "vllm").mkdir(parents=True) + (root / "csrc").mkdir() + return root + + +def _make_aiter_root(tmp_path: Path) -> Path: + root = tmp_path / "aiter-v0.1.10.post2" + (root / "aiter/ops").mkdir(parents=True) + (root / "csrc/kernels").mkdir(parents=True) + return root + + +def test_index_rejects_empty_or_mangled_name_instead_of_first_prefix(tmp_path): + index = KernelIndex(cache_dir=str(tmp_path / "cache")) + definition = _definition( + tmp_path / "vllm", + name="unrelated_kernel", + repo_name="vllm", + kind="triton_jit", + relative="kernel.py", + ) + _load_definitions(index, [definition]) + + mangled = ( + "_ZN5aiter37dynamic_per_group_scaled_quant_kernel" + "IDF16bDB8_Li32EEEvPT0_PfPKT_PKfiliibPKii.kd" + ) + assert index.lookup(mangled) is None + assert index.lookup("") is None + assert index.lookup("void aiter::kernel()") is None + + +def test_source_resolution_fields_preserve_csv_alignment(): + info = KernelSourceInfo( + kind="triton_jit", + notes="diagnostic", + source_resolution="unresolved", + source_error="triton_source_not_found", + ) + + row = dict(zip(info.csv_headers(), info.to_list(), strict=True)) + + assert row["notes"] == "diagnostic" + assert row["source_resolution"] == "unresolved" + assert row["source_error"] == "triton_source_not_found" + + +def test_index_enforces_repo_kind_and_ambiguity_constraints(tmp_path): + shared = "shared_kernel" + vllm = _definition( + tmp_path / "vllm", + name=shared, + repo_name="vllm", + kind="triton_jit", + relative="vllm_kernel.py", + ) + aiter = _definition( + tmp_path / "aiter", + name=shared, + repo_name="aiter", + kind="hip_cpp", + relative="aiter_kernel.cu", + ) + index = KernelIndex(cache_dir=str(tmp_path / "cache")) + _load_definitions(index, [vllm, aiter]) + + assert index.lookup(f"{shared}.kd") is None + assert index.lookup( + f"{shared}.kd", + expected_repo_names={"vllm"}, + expected_kinds={"triton_jit"}, + ) == vllm + assert index.lookup( + f"{shared}.kd", + expected_repo_names={"aiter"}, + expected_kinds={"triton_jit"}, + ) is None + + +def test_index_rejects_ambiguous_token_boundary_prefix(tmp_path): + root = tmp_path / "vllm" + shorter = _definition( + root, + name="paged_kernel", + repo_name="vllm", + kind="triton_jit", + relative="short.py", + ) + longer = _definition( + root, + name="paged_kernel_decode", + repo_name="vllm", + kind="triton_jit", + relative="long.py", + ) + index = KernelIndex(cache_dir=str(tmp_path / "cache")) + _load_definitions(index, [shorter, longer]) + + assert index.lookup( + "paged_kernel_decode_config_64.kd", + expected_repo_names={"vllm"}, + expected_kinds={"triton_jit"}, + ) is None + + +def test_triton_search_uses_explicit_vllm_and_aiter_roots(tmp_path): + vllm = _make_vllm_root(tmp_path) + aiter = _make_aiter_root(tmp_path) + _write_triton_source( + vllm, + "vllm/v1/attention/ops/chunked_prefill_paged_decode.py", + "kernel_paged_attention_2d", + ) + _write_triton_source( + aiter, + "aiter/ops/triton/_triton_kernels/quant/aiter_quant.py", + "aiter_quant_kernel", + ) + finder = KernelSourceFinder( + repos=[str(vllm), str(aiter)], + auto_clone=False, + use_index=False, + auto_install_ripgrep=False, + ) + + page = finder.search("kernel_paged_attention_2d.kd") + quant = finder.search("aiter_quant_kernel.kd") + + assert page.source_resolution == "resolved" + assert page.source_file == ( + "$VLLM_DIR/vllm/v1/attention/ops/" + "chunked_prefill_paged_decode.py" + ) + assert quant.source_resolution == "resolved" + assert quant.source_file == ( + "$AITER_DIR/aiter/ops/triton/_triton_kernels/quant/aiter_quant.py" + ) + + +def test_triton_cache_miss_emits_unresolved_without_placeholder( + tmp_path, + monkeypatch, +): + vllm = _make_vllm_root(tmp_path) + rogue = _make_vllm_root(tmp_path / "rogue") + _write_triton_source(rogue, "vllm/rogue.py", "missing_kernel") + monkeypatch.setenv("VLLM_DIR", str(rogue)) + finder = KernelSourceFinder( + repos=[str(vllm)], + auto_clone=False, + use_index=False, + auto_install_ripgrep=False, + ) + + result = finder.search("missing_kernel.kd") + + assert result.source_resolution == "unresolved" + assert result.source_error == "triton_source_not_found" + assert result.source_file == "" + assert result.source_repo == "" + assert "search in" not in result.notes + + +def test_duplicate_triton_definition_across_trusted_roots_fails_closed(tmp_path): + vllm = _make_vllm_root(tmp_path) + aiter = _make_aiter_root(tmp_path) + _write_triton_source(vllm, "vllm/shared.py", "shared_kernel") + _write_triton_source(aiter, "aiter/ops/shared.py", "shared_kernel") + finder = KernelSourceFinder( + repos=[str(vllm), str(aiter)], + auto_clone=False, + use_index=False, + auto_install_ripgrep=False, + ) + + result = finder.search("shared_kernel.kd") + + assert result.source_resolution == "unresolved" + assert result.source_error == "triton_source_ambiguous" + assert result.source_file == "" + + +def test_unmaterialized_aiter_ck_submodule_is_explicitly_unsupported(tmp_path): + aiter = _make_aiter_root(tmp_path) + (aiter / ".gitmodules").write_text( + "[submodule \"3rdparty/composable_kernel\"]\n" + " path = 3rdparty/composable_kernel\n" + " url = https://github.com/ROCm/composable_kernel.git\n", + encoding="utf-8", + ) + (aiter / "3rdparty/composable_kernel").mkdir(parents=True) + finder = KernelSourceFinder( + repos=[str(aiter)], + auto_clone=False, + use_index=False, + auto_install_ripgrep=False, + ) + kernel_name = ( + "kernel_gemm_xdl_cshuffle_v3<" + "GridwiseGemmMultiD_ABScale>.kd" + ) + + assert KernelNameParser().parse(kernel_name).kind == KernelKind.CK_TILE + result = finder.search(kernel_name) + + assert result.source_resolution == "unsupported" + assert result.source_error == "ck_submodule_not_materialized" + assert result.source_file == "" + assert result.source_repo == "" + assert result.test_file == "" + assert "source_resolution=unsupported" in result.notes + + +def test_materialized_aiter_ck_source_is_bound_to_exact_file(tmp_path): + aiter = _make_aiter_root(tmp_path) + ck_source = ( + aiter + / "3rdparty/composable_kernel/include/ck_tile/ops/gemm/kernel/" + "gemm_kernel.hpp" + ) + ck_source.parent.mkdir(parents=True) + ck_source.write_text("// fixture\n", encoding="utf-8") + finder = KernelSourceFinder( + repos=[str(aiter)], + auto_clone=False, + use_index=False, + auto_install_ripgrep=False, + ) + + result = finder.search( + "kernel_gemm_xdl_cshuffle_v3.kd" + ) + + assert result.source_resolution == "resolved" + assert result.source_error == "" + assert result.source_file == ( + "$AITER_DIR/3rdparty/composable_kernel/" + "include/ck_tile/ops/gemm/kernel/gemm_kernel.hpp" + ) diff --git a/tests/test_targeted_trace.py b/tests/test_targeted_trace.py index d69725d..79cb8f8 100644 --- a/tests/test_targeted_trace.py +++ b/tests/test_targeted_trace.py @@ -348,6 +348,87 @@ def test_torch_profiler_stream_adapter_and_manifest(tmp_path): assert summary["valid"] is True assert summary["streaming"] is True assert summary["events"]["by_target"] == {"aiter.fused_moe": 1} + assert summary["evidence_quality"] == { + "evidence_class": "diagnostic_only", + "resolution_status": "resolved", + "semantic_coverage_claimed": True, + "record_coverage_fraction": 1.0, + "lossless_record_coverage": True, + "records_evaluated": 1, + "records_with_complete_semantics": 1, + "missing_by_field": { + "phase": 0, + "source": 0, + "grid": 0, + "shape": 0, + "correlation": 0, + }, + "cross_event_join": "not_performed", + "join_eligible_records": 1, + "unresolved_reasons": [], + } + + +def test_missing_targeted_semantics_are_diagnostic_only_and_unresolved(tmp_path): + config = TargetedTraceConfig( + enabled=True, + targets=[TargetSpec(target_id="moe", name_patterns=("*fused_moe*",))], + ) + output = tmp_path / "targeted" + manifest = adapt_torch_profiler_traces( + [FIXTURE], + output, + config=config, + run_id="missing-source", + framework="vllm", + ) + records = [] + assert validate_shard( + Path(manifest.shards[0].path), on_event=records.append + ).valid is True + + assert "torch_profiler_missing_launch_source" in records[0].warnings + summary = postprocess_trace_dir(output) + quality = summary["evidence_quality"] + assert summary["valid"] is True + assert quality["evidence_class"] == "diagnostic_only" + assert quality["resolution_status"] == "unresolved" + assert quality["semantic_coverage_claimed"] is False + assert quality["missing_by_field"]["source"] == 1 + assert "missing:source" in quality["unresolved_reasons"] + + +def test_capped_targeted_trace_does_not_claim_semantic_coverage(tmp_path): + output = tmp_path / "targeted" + writer = TraceShardWriter( + default_shard_path(output, rank=0, pid=10), + run_id="capped", + rank=0, + pid=10, + run_seed="seed", + max_records=1, + ) + assert writer.submit(make_record("capped", key="first")) is True + assert writer.submit(make_record("capped", key="second")) is False + receipt = writer.close() + write_manifest( + output / "manifest.json", + TargetedTraceManifest( + run_id="capped", + acquisition_backend="fixture", + targets=({"target_id": "target"},), + shards=(receipt,), + ), + ) + + summary = postprocess_trace_dir(output) + quality = summary["evidence_quality"] + assert summary["valid"] is True + assert quality["record_coverage_fraction"] == 0.5 + assert quality["lossless_record_coverage"] is False + assert quality["semantic_coverage_claimed"] is False + assert quality["cross_event_join"] == "not_performed" + assert "dropped:cap" in quality["unresolved_reasons"] def test_benchmark_adapter_materializes_bounded_valid_evidence(tmp_path): @@ -384,6 +465,8 @@ def test_benchmark_adapter_materializes_bounded_valid_evidence(tmp_path): "dropped": 0, "dropped_by_reason": {}, } + assert result["evidence_quality"]["evidence_class"] == "diagnostic_only" + assert result["evidence_quality"]["resolution_status"] == "unresolved" assert Path(result["manifest_path"]).is_file() assert Path(result["summary_path"]).is_file() From 9d41f0c3b790011b44a7333f2da283a41b7fe86f Mon Sep 17 00:00:00 2001 From: Vincent Ouyang Date: Sat, 8 Aug 2026 02:50:39 +0000 Subject: [PATCH 07/11] Protect shared-host benchmark containers --- Magpie/modes/benchmark/benchmarker.py | 60 ++++++++++--- docs/reference/troubleshooting.md | 3 +- tests/test_benchmark_support.py | 121 ++++++++++++++++++++++++++ 3 files changed, 172 insertions(+), 12 deletions(-) diff --git a/Magpie/modes/benchmark/benchmarker.py b/Magpie/modes/benchmark/benchmarker.py index 4fb6659..b4f150b 100644 --- a/Magpie/modes/benchmark/benchmarker.py +++ b/Magpie/modes/benchmark/benchmarker.py @@ -11,6 +11,7 @@ import json import logging +import math import os import shutil import signal @@ -54,6 +55,10 @@ logger = logging.getLogger(__name__) +DOCKER_STOP_PROTECTION_ENV = "MAGPIE_PROTECT_BENCHMARK_CONTAINER" +DOCKER_STOP_PROTECTION_SIGNAL = "SIGWINCH" +DOCKER_STOP_PROTECTION_GRACE_SECONDS = 60 + def _env_truthy(value: Any) -> bool: return str(value).strip().lower() in {"1", "true", "yes", "on"} @@ -112,6 +117,7 @@ def __init__( self._inferencex_runtime_receipt: Optional[Dict[str, Any]] = None self._lm_eval_runtime: Optional[LmEvalRuntime] = None self._lm_eval_runtime_evidence: Optional[Dict[str, Any]] = None + self._docker_stop_protection_active = False def run(self, task_id: Optional[str] = None) -> BenchmarkResult: """ @@ -905,6 +911,28 @@ def _build_docker_command( "--ipc=host", "--shm-size=16g", "--network=host", "--name", f"magpie-benchmark-{self._task_id}", ] + + # Shared benchmark hosts sometimes have external launchers that issue a + # blanket, graceful ``docker stop`` before starting their own workload. + # Opt-in protection makes that signal inert for this run. Magpie still + # owns timeout/cancellation cleanup through an exact ``docker kill``. + self._docker_stop_protection_active = ( + self._docker_stop_protection_requested() + ) + if self._docker_stop_protection_active: + stop_timeout = max( + DOCKER_STOP_PROTECTION_GRACE_SECONDS, + math.ceil(self.config.timeout_seconds) + + DOCKER_STOP_PROTECTION_GRACE_SECONDS, + ) + cmd.extend( + [ + "--stop-signal", + DOCKER_STOP_PROTECTION_SIGNAL, + "--stop-timeout", + str(stop_timeout), + ] + ) # Add GPU-specific flags if vendor == GPUVendor.AMD: @@ -1878,13 +1906,10 @@ def _execute_benchmark(self, cmd: List[str], workspace: Path) -> tuple: self._save_logs(workspace, stdout, stderr) - # Try to stop the container + # Terminate only this run's container. When stop protection is + # enabled, a graceful stop is intentionally inert, so use kill. try: - subprocess.run( - ["docker", "stop", f"magpie-benchmark-{self._task_id}"], - capture_output=True, - timeout=30, - ) + self._terminate_docker_benchmark_container() except Exception: pass @@ -1893,6 +1918,23 @@ def _execute_benchmark(self, cmd: List[str], workspace: Path) -> tuple: logger.exception(f"Benchmark execution failed: {e}") return result, stdout, stderr + + @staticmethod + def _docker_stop_protection_requested() -> bool: + """Whether shared-host graceful-stop isolation is explicitly enabled.""" + return _env_truthy(os.environ.get(DOCKER_STOP_PROTECTION_ENV, "")) + + def _terminate_docker_benchmark_container(self) -> None: + """Terminate the exact container owned by this benchmark invocation.""" + if not self._task_id: + return + action = "kill" if self._docker_stop_protection_active else "stop" + subprocess.run( + ["docker", action, f"magpie-benchmark-{self._task_id}"], + capture_output=True, + timeout=30, + check=False, + ) def _fix_workspace_ownership(self, workspace: Path) -> None: """chown workspace files back to the invoking user after docker run. @@ -2311,10 +2353,6 @@ def cleanup(self) -> None: self._cleanup_server_processes(self.config.framework) elif self._task_id: try: - subprocess.run( - ["docker", "stop", f"magpie-benchmark-{self._task_id}"], - capture_output=True, - timeout=30, - ) + self._terminate_docker_benchmark_container() except Exception: pass diff --git a/docs/reference/troubleshooting.md b/docs/reference/troubleshooting.md index c7a5b3a..7b9a4f5 100644 --- a/docs/reference/troubleshooting.md +++ b/docs/reference/troubleshooting.md @@ -46,6 +46,7 @@ The following errors are frequently reported in benchmark mode. | `Required TraceLens inference CLI command(s) not found on PATH` | Applies to `run_mode: local` or classic host post-processing. TraceLens auto-installs on first run. If issues persist, run: `pip install git+https://github.com/AMD-AGI/TraceLens.git`. If `TL_EXTENSION=TraceLens_NDA` is set, install the matching internal extension package. For `run_mode: docker`, commands are resolved from the runtime image. | | Timeout during model loading | Large models (for example, DeepSeek-R1) might need longer timeouts. Set `timeout_seconds: 7200` in your benchmark config. | | `gpu_selection.auto failed: ...` | Not enough idle GPUs on the host. Free a GPU, lower `gpu_selection.min_free_memory_gb`, narrow `gpu_selection.candidates`, or pin manually using `envs.ROCR_VISIBLE_DEVICES` (AMD) or `envs.CUDA_VISIBLE_DEVICES` (NVIDIA). See [Automatic GPU selection in Magpie's benchmark mode](../how-to/benchmarking/automatic-gpu.md). | +| Another shared-host launcher sends `docker stop` to every running container | Prefer host-level scheduling or exclusive workers. When that is temporarily unavailable, set `MAGPIE_PROTECT_BENCHMARK_CONTAINER=true`. Magpie then gives only its benchmark container an ignored graceful-stop signal and a stop timeout derived from `timeout_seconds`; Magpie timeout and `cleanup()` still terminate that exact container with `docker kill`. This is opt-in because an operator must use `docker kill` rather than `docker stop` to interrupt a protected run. | ### Debug mode @@ -64,4 +65,4 @@ python -m Magpie benchmark --benchmark-config config.yaml --verbose | Analyze fails on worker: missing sources | `${CK_HOME}` or paths not on worker or NFS; build artifacts not present on worker. | | Worker import errors for Magpie | Set `install_magpie: true` or bake Magpie into the worker image; check `runtime_env` pip logs. | | Benchmark TP / Ray backend wrong | Inspect `_configure_tp_isolation` logs; set `EXTRA_VLLM_ARGS` / `EXTRA_SGLANG_ARGS` explicitly. | -| Empty GPU visibility in child | Should be fixed by `_clear_hidden_gpus`; if not, inspect env in InferenceX subprocess. | \ No newline at end of file +| Empty GPU visibility in child | Should be fixed by `_clear_hidden_gpus`; if not, inspect env in InferenceX subprocess. | diff --git a/tests/test_benchmark_support.py b/tests/test_benchmark_support.py index b5e15d4..0424609 100644 --- a/tests/test_benchmark_support.py +++ b/tests/test_benchmark_support.py @@ -263,6 +263,127 @@ def test_benchmark_mode_only_requests_container_writable_workspace_for_docker( assert local_mode.workspace_mgr.container_writable is False +def test_benchmark_container_stop_protection_is_explicit_and_bounded( + tmp_path, + monkeypatch, +): + config = BenchmarkConfig( + framework="vllm", + model="demo", + run_mode="docker", + timeout_seconds=125.8, + gpu_selection={"auto": False}, + ) + mode = BenchmarkMode(config, output_dir=str(tmp_path / "results")) + mode._task_id = "protected" + monkeypatch.setattr( + "Magpie.modes.benchmark.benchmarker.detect_gpu", + lambda: (GPUVendor.UNKNOWN, ""), + ) + monkeypatch.setattr( + mode, + "_get_benchmark_script", + lambda _runner_type: "benchmarks/vllm_mi355x.sh", + ) + + monkeypatch.delenv("MAGPIE_PROTECT_BENCHMARK_CONTAINER", raising=False) + unprotected = mode._build_docker_command( + "example/image:fixed", + tmp_path / "workspace", + "mi355x", + ) + assert "--stop-signal" not in unprotected + assert "--stop-timeout" not in unprotected + + monkeypatch.setenv("MAGPIE_PROTECT_BENCHMARK_CONTAINER", "true") + protected = mode._build_docker_command( + "example/image:fixed", + tmp_path / "workspace", + "mi355x", + ) + signal_index = protected.index("--stop-signal") + timeout_index = protected.index("--stop-timeout") + assert protected[signal_index + 1] == "SIGWINCH" + assert protected[timeout_index + 1] == "186" + assert mode._docker_stop_protection_active is True + monkeypatch.setenv("MAGPIE_PROTECT_BENCHMARK_CONTAINER", "false") + assert mode._docker_stop_protection_active is True + + +@pytest.mark.parametrize( + ("protection", "expected_action"), + (("true", "kill"), ("false", "stop")), +) +def test_benchmark_container_cleanup_uses_exact_owned_name( + tmp_path, + monkeypatch, + protection, + expected_action, +): + mode = BenchmarkMode( + BenchmarkConfig(framework="vllm", model="demo", run_mode="docker"), + output_dir=str(tmp_path / "results"), + ) + mode._task_id = "owned-task" + monkeypatch.setenv("MAGPIE_PROTECT_BENCHMARK_CONTAINER", protection) + mode._docker_stop_protection_active = protection == "true" + calls = [] + + def fake_run(command, **kwargs): + calls.append((command, kwargs)) + return subprocess.CompletedProcess(command, 0) + + monkeypatch.setattr( + "Magpie.modes.benchmark.benchmarker.subprocess.run", + fake_run, + ) + + mode.cleanup() + + assert calls == [ + ( + ["docker", expected_action, "magpie-benchmark-owned-task"], + {"capture_output": True, "timeout": 30, "check": False}, + ) + ] + + +def test_benchmark_timeout_kills_latched_protected_container( + tmp_path, + monkeypatch, +): + mode = BenchmarkMode( + BenchmarkConfig(framework="vllm", model="demo", run_mode="docker"), + output_dir=str(tmp_path / "results"), + ) + mode._task_id = "timed-out-task" + mode._docker_stop_protection_active = True + calls = [] + + def fake_run(command, **kwargs): + calls.append((command, kwargs)) + if command == ["docker", "run", "example"]: + raise subprocess.TimeoutExpired(command, timeout=1) + return subprocess.CompletedProcess(command, 0) + + monkeypatch.setattr( + "Magpie.modes.benchmark.benchmarker.subprocess.run", + fake_run, + ) + + result, _stdout, _stderr = mode._execute_benchmark( + ["docker", "run", "example"], + tmp_path, + ) + + assert result.success is False + assert result.errors == ["Benchmark timed out after 3600.0s"] + assert calls[-1] == ( + ["docker", "kill", "magpie-benchmark-timed-out-task"], + {"capture_output": True, "timeout": 30, "check": False}, + ) + + def test_benchmark_server_lifecycle_requires_local_runtime(): with pytest.raises(ValueError, match="server_lifecycle"): BenchmarkConfig( From eb08aaa1c356eb20cb47e0baac54e94f042ca9e3 Mon Sep 17 00:00:00 2001 From: Vincent Ouyang Date: Sat, 8 Aug 2026 05:57:10 +0000 Subject: [PATCH 08/11] Bind TraceLens builds to local image IDs --- .../modes/benchmark/tracelens_vllm_image.py | 255 +++++++++++++--- docs/how-to/benchmarking/profiling-options.md | 7 + tests/test_tracelens_vllm_image.py | 283 +++++++++++++++++- 3 files changed, 497 insertions(+), 48 deletions(-) diff --git a/Magpie/modes/benchmark/tracelens_vllm_image.py b/Magpie/modes/benchmark/tracelens_vllm_image.py index 329d7d4..f7437c3 100644 --- a/Magpie/modes/benchmark/tracelens_vllm_image.py +++ b/Magpie/modes/benchmark/tracelens_vllm_image.py @@ -19,6 +19,7 @@ import json import os import re +import secrets import shutil import subprocess import tarfile @@ -27,7 +28,6 @@ from pathlib import Path, PurePosixPath from typing import Any, Dict, Mapping, Optional, Sequence - VLLM_TRACELENS_REQUIREMENTS = ( ("contourpy", "1.3.3"), ("cycler", "0.12.1"), @@ -61,6 +61,9 @@ "TraceLens", ) _HASH_RE = re.compile(r"^[0-9a-f]{64}$") +_IMAGE_ID_RE = re.compile(r"^sha256:([0-9a-f]{64})$") +_REPO_DIGEST_RE = re.compile(r"^[a-z0-9][a-z0-9._:/-]*@sha256:[0-9a-f]{64}$") +_LOCAL_BASE_REPOSITORY = "localhost/magpie-tracelens-vllm-base" @dataclass(frozen=True) @@ -107,6 +110,16 @@ def metadata(self) -> Dict[str, Any]: } +@dataclass(frozen=True) +class _BuildBaseReference: + """A build-only Docker reference bound to one locally inspectable image ID.""" + + locator: str + image_id: str + kind: str + owns_temporary_tag: bool = False + + def _sha256_bytes(value: bytes) -> str: return hashlib.sha256(value).hexdigest() @@ -188,6 +201,153 @@ def docker_image_id(image: str) -> Optional[str]: return image_id if isinstance(image_id, str) and image_id else None +def _require_expected_image_id(image: str, expected_id: str, *, role: str) -> None: + actual_id = docker_image_id(image) + if actual_id != expected_id: + raise RuntimeError( + f"{role} is not bound to the expected local Docker image ID: " + f"reference={image!r}, expected={expected_id!r}, actual={actual_id!r}" + ) + + +def _temporary_local_base_tag( + image_id: str, + *, + nonce: Optional[str] = None, +) -> str: + match = _IMAGE_ID_RE.fullmatch(image_id) + if not match: + raise RuntimeError(f"Invalid local Docker image ID: {image_id!r}") + unique = nonce or secrets.token_hex(16) + if not re.fullmatch(r"[0-9a-f]{32}", unique): + raise RuntimeError(f"Invalid local Docker tag nonce: {unique!r}") + return f"{_LOCAL_BASE_REPOSITORY}:sha256-{match.group(1)}-{unique}" + + +def _docker_tag_image(image_id: str, tag: str) -> None: + try: + proc = subprocess.run( + ["docker", "image", "tag", image_id, tag], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + timeout=120, + check=False, + ) + except (OSError, subprocess.TimeoutExpired) as exc: + raise RuntimeError( + f"Could not create local Docker base tag {tag!r}: {exc}" + ) from exc + if proc.returncode != 0: + raise RuntimeError( + f"Could not create local Docker base tag {tag!r}: " + f"{_completed_output(proc)}" + ) + + +def _docker_remove_tag(tag: str) -> None: + try: + proc = subprocess.run( + ["docker", "image", "rm", tag], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + timeout=120, + check=False, + ) + except (OSError, subprocess.TimeoutExpired) as exc: + raise RuntimeError( + f"Could not remove local Docker base tag {tag!r}: {exc}" + ) from exc + if proc.returncode != 0: + raise RuntimeError( + f"Could not remove local Docker base tag {tag!r}: " + f"{_completed_output(proc)}" + ) + + +def _acquire_build_base_reference( + identity: VllmTraceLensIdentity, +) -> _BuildBaseReference: + """Resolve an exact build base without treating an image ID as a tag.""" + expected_id = identity.base_image_id + if not _IMAGE_ID_RE.fullmatch(expected_id): + raise RuntimeError(f"Invalid local Docker image ID: {expected_id!r}") + _require_expected_image_id(expected_id, expected_id, role="base image ID") + + if _REPO_DIGEST_RE.fullmatch(identity.base_image_locator): + _require_expected_image_id( + identity.base_image_locator, + expected_id, + role="repository-digest build base", + ) + return _BuildBaseReference( + locator=identity.base_image_locator, + image_id=expected_id, + kind="repository-digest", + ) + + tag = _temporary_local_base_tag(expected_id) + existing_id = docker_image_id(tag) + if existing_id is not None: + raise RuntimeError( + "Unique local TraceLens base tag already exists; refusing to reuse " + "a tag this build does not own: " + f"tag={tag!r}, expected={expected_id!r}, actual={existing_id!r}" + ) + + _docker_tag_image(expected_id, tag) + try: + _require_expected_image_id(tag, expected_id, role="temporary build base tag") + except RuntimeError: + if docker_image_id(tag) == expected_id: + _docker_remove_tag(tag) + raise + return _BuildBaseReference( + locator=tag, + image_id=expected_id, + kind="temporary-local-tag", + owns_temporary_tag=True, + ) + + +def _verify_build_base_reference(reference: _BuildBaseReference) -> None: + if reference.kind == "repository-digest": + _require_expected_image_id( + reference.locator, + reference.image_id, + role="repository-digest build base", + ) + else: + _require_expected_image_id( + reference.locator, + reference.image_id, + role="local build base tag", + ) + _require_expected_image_id( + reference.image_id, + reference.image_id, + role="base image ID", + ) + + +def _release_build_base_reference(reference: _BuildBaseReference) -> None: + if not reference.owns_temporary_tag: + return + _require_expected_image_id( + reference.locator, + reference.image_id, + role="owned temporary build base tag", + ) + _docker_remove_tag(reference.locator) + remaining_id = docker_image_id(reference.locator) + if remaining_id is not None: + raise RuntimeError( + "Owned temporary build base tag still exists after cleanup: " + f"tag={reference.locator!r}, actual={remaining_id!r}" + ) + + def resolve_vllm_tracelens_identity( *, base_image: str, @@ -199,7 +359,7 @@ def resolve_vllm_tracelens_identity( """Resolve source, patch, and immutable base-image identity.""" base_record = docker_image_record(base_image) base_id = base_record.get("Id") if base_record else None - if not isinstance(base_id, str) or not base_id: + if not isinstance(base_id, str) or not _IMAGE_ID_RE.fullmatch(base_id): raise RuntimeError(f"Could not resolve Docker image ID for {base_image!r}") repo_digests = base_record.get("RepoDigests") or [] base_locator = ( @@ -373,7 +533,9 @@ def _download_requirement_wheels( identity: VllmTraceLensIdentity, wheelhouse: Path, ) -> list[str]: - requirements = [f"{name}=={version}" for name, version in VLLM_TRACELENS_REQUIREMENTS] + requirements = [ + f"{name}=={version}" for name, version in VLLM_TRACELENS_REQUIREMENTS + ] cmd = [ "docker", "run", @@ -469,6 +631,8 @@ def _write_build_context( context: Path, identity: VllmTraceLensIdentity, wheel_manifest: Sequence[Mapping[str, str]], + *, + build_base_locator: str, ) -> Dict[str, str]: wheel_manifest_json = _canonical_json(list(wheel_manifest)) labels = identity.labels() @@ -492,7 +656,7 @@ def _write_build_context( (context / "verify.py").write_text(_verification_script(), encoding="utf-8") (context / "patch.diff").write_bytes(identity.patch_bytes) dockerfile = f"""\ -FROM {identity.base_image_locator} +FROM {build_base_locator} {_dockerfile_labels(labels)} COPY wheels/ /tmp/tracelens-wheels/ COPY patch.diff /tmp/tracelens-vllm.patch @@ -529,33 +693,46 @@ def build_vllm_tracelens_image( download_command = _download_requirement_wheels(identity, wheelhouse) wheel_manifest = _wheel_manifest(wheelhouse) shutil.rmtree(source_dir) - labels = _write_build_context(context, identity, wheel_manifest) - archive = context / "tracelens-source.tar" - if archive.exists(): - archive.unlink() - cmd = [ - "docker", - "build", - "--network", - "none", - "--no-cache", - "--provenance=false", - "-t", - derived_image, - str(context), - ] - proc = subprocess.run( - cmd, - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - text=True, - check=False, - ) - if proc.returncode != 0: - raise RuntimeError( - "TraceLens vLLM runtime image build failed with exit code " - f"{proc.returncode}. Image: {derived_image}\n{(proc.stdout or '')[-4000:]}" + base_reference = _acquire_build_base_reference(identity) + try: + labels = _write_build_context( + context, + identity, + wheel_manifest, + build_base_locator=base_reference.locator, + ) + archive = context / "tracelens-source.tar" + if archive.exists(): + archive.unlink() + _verify_build_base_reference(base_reference) + cmd = [ + "docker", + "build", + "--network", + "none", + "--pull=false", + "--no-cache", + "--provenance=false", + "-t", + derived_image, + str(context), + ] + proc = subprocess.run( + cmd, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + check=False, ) + if proc.returncode != 0: + raise RuntimeError( + "TraceLens vLLM runtime image build failed with exit code " + f"{proc.returncode}. Image: {derived_image}\n" + f"{(proc.stdout or '')[-4000:]}" + ) + _verify_build_base_reference(base_reference) + finally: + _release_build_base_reference(base_reference) validation = validate_vllm_tracelens_image(derived_image, identity) if not validation["valid"]: @@ -568,12 +745,16 @@ def build_vllm_tracelens_image( "command": cmd[:-1] + [""], "source_wheel_command": source_wheel_command[:-1] + [""], "requirements_download_command": download_command, + "base_binding": { + "image_id": base_reference.image_id, + "provenance_locator": identity.base_image_locator, + "build_reference_kind": base_reference.kind, + "temporary_tag_removed": base_reference.owns_temporary_tag, + }, "image_id": record.get("Id"), "image_labels": labels, "dependency_wheels": list(wheel_manifest), - "dependency_wheel_manifest_sha256": labels[ - LABEL_WHEEL_MANIFEST_SHA256 - ], + "dependency_wheel_manifest_sha256": labels[LABEL_WHEEL_MANIFEST_SHA256], "validation": validation, } @@ -729,15 +910,11 @@ def validate_vllm_tracelens_image( } return { "valid": True, - "reason": ( - "identity, ancestry, packages, imports, and exact patch verified" - ), + "reason": ("identity, ancestry, packages, imports, and exact patch verified"), "image_id": record.get("Id"), "runtime_probe": probe_result, "dependency_wheels": json.loads(labels[LABEL_WHEEL_MANIFEST]), - "dependency_wheel_manifest_sha256": labels[ - LABEL_WHEEL_MANIFEST_SHA256 - ], + "dependency_wheel_manifest_sha256": labels[LABEL_WHEEL_MANIFEST_SHA256], } diff --git a/docs/how-to/benchmarking/profiling-options.md b/docs/how-to/benchmarking/profiling-options.md index d26bc81..b1c1296 100644 --- a/docs/how-to/benchmarking/profiling-options.md +++ b/docs/how-to/benchmarking/profiling-options.md @@ -52,6 +52,13 @@ and `grpcio` versions remain unchanged. The final Docker build is network-free. The derived image is tagged locally as `magpie-tracelens-:...` and reused on later runs only after validation. +When the selected base is a locally derived image without a repository digest, +Magpie binds its exact image ID to a unique owned `localhost/` build-only tag +instead of writing `FROM sha256:...` (which BuildKit interprets as a registry +tag). The final build disables pulls, verifies that binding before and after +the build, and removes only a tag created by that build. Repository-digest +parents keep using their immutable digest directly. + The vLLM image carries OCI labels and a runtime identity document containing the base image ID, TraceLens source commit and tree, patch hash, pinned wheel hashes, and preserved package versions. A same-name local image with missing or stale diff --git a/tests/test_tracelens_vllm_image.py b/tests/test_tracelens_vllm_image.py index 785d787..1efa8f5 100644 --- a/tests/test_tracelens_vllm_image.py +++ b/tests/test_tracelens_vllm_image.py @@ -1,6 +1,8 @@ import hashlib import json import subprocess +from dataclasses import replace +from pathlib import Path from Magpie.modes.benchmark.config import BenchmarkConfig from Magpie.modes.benchmark.tracelens_runtime import ( @@ -12,8 +14,13 @@ VLLM_TRACELENS_FORBIDDEN, VLLM_TRACELENS_REQUIREMENTS, VllmTraceLensIdentity, + _acquire_build_base_reference, + _release_build_base_reference, + _temporary_local_base_tag, + _verify_build_base_reference, _verification_script, _write_build_context, + build_vllm_tracelens_image, resolve_vllm_tracelens_identity, validate_vllm_tracelens_image, ) @@ -23,9 +30,7 @@ def _identity(): return VllmTraceLensIdentity( base_image="vllm/vllm-openai-rocm:v0.19.1", base_image_id="sha256:" + "1" * 64, - base_image_locator=( - "vllm/vllm-openai-rocm@sha256:" + "a" * 64 - ), + base_image_locator=("vllm/vllm-openai-rocm@sha256:" + "a" * 64), vllm_version="0.19.1+rocm721", grpcio_version="1.78.0", source_commit="2" * 40, @@ -65,9 +70,7 @@ def _labels(identity): labels = identity.labels() manifest = json.dumps(_wheel_manifest(), sort_keys=True, separators=(",", ":")) labels[LABEL_WHEEL_MANIFEST] = manifest - labels[LABEL_WHEEL_MANIFEST_SHA256] = hashlib.sha256( - manifest.encode() - ).hexdigest() + labels[LABEL_WHEEL_MANIFEST_SHA256] = hashlib.sha256(manifest.encode()).hexdigest() return labels @@ -108,7 +111,12 @@ def test_resolve_identity_uses_image_digest_and_committed_patch(monkeypatch, tmp def test_build_context_is_offline_minimal_and_identity_labeled(tmp_path): identity = _identity() - labels = _write_build_context(tmp_path, identity, _wheel_manifest()) + labels = _write_build_context( + tmp_path, + identity, + _wheel_manifest(), + build_base_locator=identity.base_image_locator, + ) dockerfile = (tmp_path / "Dockerfile").read_text(encoding="utf-8") verifier = (tmp_path / "verify.py").read_text(encoding="utf-8") @@ -124,12 +132,269 @@ def test_build_context_is_offline_minimal_and_identity_labeled(tmp_path): assert identity_document["tracelens_source_commit"] == identity.source_commit assert identity_document["tracelens_source_tree"] == identity.source_tree assert identity_document["tracelens_patch_sha256"] == identity.patch_sha256 - assert "metadata.version(\"grpcio\")" in verifier - assert "metadata.version(\"vllm\")" in verifier + assert 'metadata.version("grpcio")' in verifier + assert 'metadata.version("vllm")' in verifier for package in VLLM_TRACELENS_FORBIDDEN: assert package in verifier +def test_local_image_id_build_uses_bound_local_tag_without_registry_fallback( + monkeypatch, + tmp_path, +): + base_id = "sha256:" + "b" * 64 + derived_id = "sha256:" + "c" * 64 + derived_image = "magpie-tracelens-vllm:test-local-base" + identity = replace( + _identity(), + base_image=base_id, + base_image_id=base_id, + base_image_locator=base_id, + ) + nonce = "e" * 32 + local_tag = _temporary_local_base_tag(base_id, nonce=nonce) + records = { + base_id: {"Id": base_id}, + derived_image: {"Id": derived_id}, + } + commands = [] + captured = {} + + monkeypatch.setattr( + "Magpie.modes.benchmark.tracelens_vllm_image.docker_image_record", + lambda image: records.get(image), + ) + monkeypatch.setattr( + "Magpie.modes.benchmark.tracelens_vllm_image.secrets.token_hex", + lambda _size: nonce, + ) + + def fake_run(cmd, **_kwargs): + commands.append(list(cmd)) + if cmd[:3] == ["docker", "image", "tag"]: + assert cmd[3:] == [base_id, local_tag] + records[local_tag] = {"Id": base_id} + elif cmd[:3] == ["docker", "image", "rm"]: + assert cmd[3:] == [local_tag] + records.pop(local_tag) + elif cmd[:2] == ["docker", "build"]: + context = cmd[-1] + captured["dockerfile"] = (Path(context) / "Dockerfile").read_text( + encoding="utf-8" + ) + else: + raise AssertionError(f"unexpected command: {cmd}") + return subprocess.CompletedProcess(cmd, 0, stdout="", stderr="") + + def fake_stage(_identity, _repo, destination): + destination.mkdir(parents=True) + + monkeypatch.setattr( + "Magpie.modes.benchmark.tracelens_vllm_image.subprocess.run", + fake_run, + ) + monkeypatch.setattr( + "Magpie.modes.benchmark.tracelens_vllm_image._stage_committed_source", + fake_stage, + ) + monkeypatch.setattr( + "Magpie.modes.benchmark.tracelens_vllm_image._build_source_wheel", + lambda *_args: ["docker", "run", "source-wheel"], + ) + monkeypatch.setattr( + "Magpie.modes.benchmark.tracelens_vllm_image._download_requirement_wheels", + lambda *_args: ["docker", "run", "requirements"], + ) + monkeypatch.setattr( + "Magpie.modes.benchmark.tracelens_vllm_image._wheel_manifest", + lambda _wheelhouse: _wheel_manifest(), + ) + monkeypatch.setattr( + "Magpie.modes.benchmark.tracelens_vllm_image.validate_vllm_tracelens_image", + lambda _image, _identity: {"valid": True, "reason": "verified"}, + ) + + result = build_vllm_tracelens_image( + identity=identity, + tracelens_repo=tmp_path, + derived_image=derived_image, + ) + + build_command = next(cmd for cmd in commands if cmd[:2] == ["docker", "build"]) + assert "--pull=false" in build_command + assert build_command[build_command.index("--network") + 1] == "none" + assert captured["dockerfile"].startswith(f"FROM {local_tag}\n") + assert f"FROM {base_id}\n" not in captured["dockerfile"] + assert local_tag.startswith("localhost/") + assert local_tag not in records + assert result["base_binding"] == { + "image_id": base_id, + "provenance_locator": base_id, + "build_reference_kind": "temporary-local-tag", + "temporary_tag_removed": True, + } + + +def test_local_build_rejects_reserved_tag_bound_to_different_image(monkeypatch): + base_id = "sha256:" + "b" * 64 + wrong_id = "sha256:" + "d" * 64 + identity = replace( + _identity(), + base_image=base_id, + base_image_id=base_id, + base_image_locator=base_id, + ) + nonce = "e" * 32 + local_tag = _temporary_local_base_tag(base_id, nonce=nonce) + + monkeypatch.setattr( + "Magpie.modes.benchmark.tracelens_vllm_image.docker_image_id", + lambda image: wrong_id if image == local_tag else base_id, + ) + monkeypatch.setattr( + "Magpie.modes.benchmark.tracelens_vllm_image.secrets.token_hex", + lambda _size: nonce, + ) + monkeypatch.setattr( + "Magpie.modes.benchmark.tracelens_vllm_image.subprocess.run", + lambda *_args, **_kwargs: (_ for _ in ()).throw( + AssertionError("mismatched tag must fail before Docker mutation") + ), + ) + + try: + _acquire_build_base_reference(identity) + except RuntimeError as exc: + assert "refusing to reuse" in str(exc) + assert wrong_id in str(exc) + else: + raise AssertionError("mismatched reserved tag was accepted") + + +def test_repository_digest_build_rejects_locator_id_mismatch(monkeypatch): + identity = _identity() + wrong_id = "sha256:" + "d" * 64 + + monkeypatch.setattr( + "Magpie.modes.benchmark.tracelens_vllm_image.docker_image_id", + lambda image: ( + wrong_id if image == identity.base_image_locator else identity.base_image_id + ), + ) + monkeypatch.setattr( + "Magpie.modes.benchmark.tracelens_vllm_image.subprocess.run", + lambda *_args, **_kwargs: (_ for _ in ()).throw( + AssertionError("repository-digest mismatch must not mutate Docker") + ), + ) + + try: + _acquire_build_base_reference(identity) + except RuntimeError as exc: + assert "repository-digest build base" in str(exc) + assert wrong_id in str(exc) + else: + raise AssertionError("mismatched repository digest was accepted") + + +def test_local_build_detects_post_build_retag_without_removing_foreign_tag( + monkeypatch, +): + base_id = "sha256:" + "b" * 64 + wrong_id = "sha256:" + "d" * 64 + nonce = "e" * 32 + identity = replace( + _identity(), + base_image=base_id, + base_image_id=base_id, + base_image_locator=base_id, + ) + local_tag = _temporary_local_base_tag(base_id, nonce=nonce) + records = {base_id: base_id} + commands = [] + + monkeypatch.setattr( + "Magpie.modes.benchmark.tracelens_vllm_image.docker_image_id", + lambda image: records.get(image), + ) + monkeypatch.setattr( + "Magpie.modes.benchmark.tracelens_vllm_image.secrets.token_hex", + lambda _size: nonce, + ) + + def fake_run(cmd, **_kwargs): + commands.append(list(cmd)) + assert cmd[:3] == ["docker", "image", "tag"] + records[local_tag] = base_id + return subprocess.CompletedProcess(cmd, 0, stdout="", stderr="") + + monkeypatch.setattr( + "Magpie.modes.benchmark.tracelens_vllm_image.subprocess.run", + fake_run, + ) + reference = _acquire_build_base_reference(identity) + records[local_tag] = wrong_id + + try: + _verify_build_base_reference(reference) + except RuntimeError as exc: + assert "local build base tag" in str(exc) + assert wrong_id in str(exc) + else: + raise AssertionError("post-build retag was accepted") + + try: + _release_build_base_reference(reference) + except RuntimeError as exc: + assert "owned temporary build base tag" in str(exc) + else: + raise AssertionError("cleanup removed a tag whose ownership was lost") + assert all(cmd[:3] != ["docker", "image", "rm"] for cmd in commands) + + +def test_concurrent_local_builds_acquire_distinct_owned_tags(monkeypatch): + base_id = "sha256:" + "b" * 64 + nonces = iter(("e" * 32, "f" * 32)) + identity = replace( + _identity(), + base_image=base_id, + base_image_id=base_id, + base_image_locator=base_id, + ) + records = {base_id: base_id} + + monkeypatch.setattr( + "Magpie.modes.benchmark.tracelens_vllm_image.docker_image_id", + lambda image: records.get(image), + ) + monkeypatch.setattr( + "Magpie.modes.benchmark.tracelens_vllm_image.secrets.token_hex", + lambda _size: next(nonces), + ) + + def fake_run(cmd, **_kwargs): + if cmd[:3] == ["docker", "image", "tag"]: + records[cmd[4]] = base_id + elif cmd[:3] == ["docker", "image", "rm"]: + records.pop(cmd[3]) + else: + raise AssertionError(f"unexpected command: {cmd}") + return subprocess.CompletedProcess(cmd, 0, stdout="", stderr="") + + monkeypatch.setattr( + "Magpie.modes.benchmark.tracelens_vllm_image.subprocess.run", + fake_run, + ) + + first = _acquire_build_base_reference(identity) + second = _acquire_build_base_reference(identity) + assert first.locator != second.locator + assert first.owns_temporary_tag and second.owns_temporary_tag + _release_build_base_reference(first) + assert second.locator in records + _release_build_base_reference(second) + + def test_existing_image_validation_preserves_vllm_grpc_and_excludes_pollution( monkeypatch, ): From fa39dde9e18b8f50762ffbe9f5caaa2c32a9b748 Mon Sep 17 00:00:00 2001 From: Vincent Ouyang Date: Sat, 8 Aug 2026 06:04:54 +0000 Subject: [PATCH 09/11] Preserve unnamed TraceLens base images --- Magpie/modes/benchmark/tracelens_runtime.py | 1 + .../modes/benchmark/tracelens_vllm_image.py | 58 +++++++++++++++ docs/how-to/benchmarking/profiling-options.md | 12 ++-- tests/test_benchmark_support.py | 8 +++ tests/test_tracelens_vllm_image.py | 70 +++++++++++++++++-- 5 files changed, 140 insertions(+), 9 deletions(-) diff --git a/Magpie/modes/benchmark/tracelens_runtime.py b/Magpie/modes/benchmark/tracelens_runtime.py index 1dd7529..92cf2f6 100644 --- a/Magpie/modes/benchmark/tracelens_runtime.py +++ b/Magpie/modes/benchmark/tracelens_runtime.py @@ -839,6 +839,7 @@ def prepare_tracelens_runtime_image( "requirements_download_command": build_metadata[ "requirements_download_command" ], + "base_binding": build_metadata["base_binding"], "public_runtime_image_id": build_metadata["image_id"], "public_runtime_labels": build_metadata["image_labels"], "dependency_wheels": build_metadata["dependency_wheels"], diff --git a/Magpie/modes/benchmark/tracelens_vllm_image.py b/Magpie/modes/benchmark/tracelens_vllm_image.py index f7437c3..fc3e0a3 100644 --- a/Magpie/modes/benchmark/tracelens_vllm_image.py +++ b/Magpie/modes/benchmark/tracelens_vllm_image.py @@ -118,6 +118,8 @@ class _BuildBaseReference: image_id: str kind: str owns_temporary_tag: bool = False + retained_locator: Optional[str] = None + retained_locator_created: bool = False def _sha256_bytes(value: bytes) -> str: @@ -224,6 +226,14 @@ def _temporary_local_base_tag( return f"{_LOCAL_BASE_REPOSITORY}:sha256-{match.group(1)}-{unique}" +def _retained_local_base_tag(image_id: str) -> str: + """Return the stable local name that keeps an unnamed parent inspectable.""" + match = _IMAGE_ID_RE.fullmatch(image_id) + if not match: + raise RuntimeError(f"Invalid local Docker image ID: {image_id!r}") + return f"{_LOCAL_BASE_REPOSITORY}:sha256-{match.group(1)}" + + def _docker_tag_image(image_id: str, tag: str) -> None: try: proc = subprocess.run( @@ -287,6 +297,24 @@ def _acquire_build_base_reference( kind="repository-digest", ) + retained_tag = _retained_local_base_tag(expected_id) + retained_id = docker_image_id(retained_tag) + retained_created = False + if retained_id is None: + _docker_tag_image(expected_id, retained_tag) + _require_expected_image_id( + retained_tag, + expected_id, + role="retained local base tag", + ) + retained_created = True + elif retained_id != expected_id: + raise RuntimeError( + "Content-addressed local TraceLens base tag resolves to a different " + "image ID: " + f"tag={retained_tag!r}, expected={expected_id!r}, actual={retained_id!r}" + ) + tag = _temporary_local_base_tag(expected_id) existing_id = docker_image_id(tag) if existing_id is not None: @@ -308,6 +336,8 @@ def _acquire_build_base_reference( image_id=expected_id, kind="temporary-local-tag", owns_temporary_tag=True, + retained_locator=retained_tag, + retained_locator_created=retained_created, ) @@ -324,6 +354,13 @@ def _verify_build_base_reference(reference: _BuildBaseReference) -> None: reference.image_id, role="local build base tag", ) + if reference.retained_locator is None: + raise RuntimeError("Local build base is missing its retained locator") + _require_expected_image_id( + reference.retained_locator, + reference.image_id, + role="retained local base tag", + ) _require_expected_image_id( reference.image_id, reference.image_id, @@ -334,6 +371,13 @@ def _verify_build_base_reference(reference: _BuildBaseReference) -> None: def _release_build_base_reference(reference: _BuildBaseReference) -> None: if not reference.owns_temporary_tag: return + if reference.retained_locator is None: + raise RuntimeError("Owned temporary base tag has no retained locator") + _require_expected_image_id( + reference.retained_locator, + reference.image_id, + role="retained local base tag", + ) _require_expected_image_id( reference.locator, reference.image_id, @@ -346,6 +390,16 @@ def _release_build_base_reference(reference: _BuildBaseReference) -> None: "Owned temporary build base tag still exists after cleanup: " f"tag={reference.locator!r}, actual={remaining_id!r}" ) + _require_expected_image_id( + reference.retained_locator, + reference.image_id, + role="retained local base tag", + ) + _require_expected_image_id( + reference.image_id, + reference.image_id, + role="base image ID after temporary-tag cleanup", + ) def resolve_vllm_tracelens_identity( @@ -750,6 +804,10 @@ def build_vllm_tracelens_image( "provenance_locator": identity.base_image_locator, "build_reference_kind": base_reference.kind, "temporary_tag_removed": base_reference.owns_temporary_tag, + "retained_local_reference": base_reference.retained_locator, + "retained_local_reference_created": ( + base_reference.retained_locator_created + ), }, "image_id": record.get("Id"), "image_labels": labels, diff --git a/docs/how-to/benchmarking/profiling-options.md b/docs/how-to/benchmarking/profiling-options.md index b1c1296..6b985db 100644 --- a/docs/how-to/benchmarking/profiling-options.md +++ b/docs/how-to/benchmarking/profiling-options.md @@ -53,11 +53,13 @@ The derived image is tagged locally as `magpie-tracelens-:...` and reused on later runs only after validation. When the selected base is a locally derived image without a repository digest, -Magpie binds its exact image ID to a unique owned `localhost/` build-only tag -instead of writing `FROM sha256:...` (which BuildKit interprets as a registry -tag). The final build disables pulls, verifies that binding before and after -the build, and removes only a tag created by that build. Repository-digest -parents keep using their immutable digest directly. +Magpie first gives the exact image ID a stable, content-addressed `localhost/` +retention tag, then binds it to a unique owned build-only tag instead of writing +`FROM sha256:...` (which BuildKit interprets as a registry tag). The retention +tag keeps a formerly unnamed parent inspectable after Docker removes the unique +tag. The final build disables pulls, verifies both bindings before and after the +build, and removes only the unique tag created by that build. Repository-digest +parents keep using their immutable digest directly and need no retention tag. The vLLM image carries OCI labels and a runtime identity document containing the base image ID, TraceLens source commit and tree, patch hash, pinned wheel hashes, diff --git a/tests/test_benchmark_support.py b/tests/test_benchmark_support.py index 0424609..a957ead 100644 --- a/tests/test_benchmark_support.py +++ b/tests/test_benchmark_support.py @@ -1059,6 +1059,14 @@ def test_prepare_tracelens_runtime_image_builds_extension_overlay( "command": ["docker", "build", ""], "source_wheel_command": ["docker", "run", ""], "requirements_download_command": ["docker", "run", "pip", "download"], + "base_binding": { + "image_id": identity.base_image_id, + "provenance_locator": identity.base_image_locator, + "build_reference_kind": "repository-digest", + "temporary_tag_removed": False, + "retained_local_reference": None, + "retained_local_reference_created": False, + }, "image_id": "sha256:" + "7" * 64, "image_labels": identity.labels(), "dependency_wheels": [], diff --git a/tests/test_tracelens_vllm_image.py b/tests/test_tracelens_vllm_image.py index 1efa8f5..ac22c70 100644 --- a/tests/test_tracelens_vllm_image.py +++ b/tests/test_tracelens_vllm_image.py @@ -16,6 +16,7 @@ VllmTraceLensIdentity, _acquire_build_base_reference, _release_build_base_reference, + _retained_local_base_tag, _temporary_local_base_tag, _verify_build_base_reference, _verification_script, @@ -153,6 +154,7 @@ def test_local_image_id_build_uses_bound_local_tag_without_registry_fallback( ) nonce = "e" * 32 local_tag = _temporary_local_base_tag(base_id, nonce=nonce) + retained_tag = _retained_local_base_tag(base_id) records = { base_id: {"Id": base_id}, derived_image: {"Id": derived_id}, @@ -172,8 +174,9 @@ def test_local_image_id_build_uses_bound_local_tag_without_registry_fallback( def fake_run(cmd, **_kwargs): commands.append(list(cmd)) if cmd[:3] == ["docker", "image", "tag"]: - assert cmd[3:] == [base_id, local_tag] - records[local_tag] = {"Id": base_id} + assert cmd[3] == base_id + assert cmd[4] in (retained_tag, local_tag) + records[cmd[4]] = {"Id": base_id} elif cmd[:3] == ["docker", "image", "rm"]: assert cmd[3:] == [local_tag] records.pop(local_tag) @@ -227,11 +230,15 @@ def fake_stage(_identity, _repo, destination): assert f"FROM {base_id}\n" not in captured["dockerfile"] assert local_tag.startswith("localhost/") assert local_tag not in records + assert records[retained_tag] == {"Id": base_id} + assert records[base_id] == {"Id": base_id} assert result["base_binding"] == { "image_id": base_id, "provenance_locator": base_id, "build_reference_kind": "temporary-local-tag", "temporary_tag_removed": True, + "retained_local_reference": retained_tag, + "retained_local_reference_created": True, } @@ -246,10 +253,17 @@ def test_local_build_rejects_reserved_tag_bound_to_different_image(monkeypatch): ) nonce = "e" * 32 local_tag = _temporary_local_base_tag(base_id, nonce=nonce) + retained_tag = _retained_local_base_tag(base_id) monkeypatch.setattr( "Magpie.modes.benchmark.tracelens_vllm_image.docker_image_id", - lambda image: wrong_id if image == local_tag else base_id, + lambda image: ( + wrong_id + if image == local_tag + else base_id + if image in (base_id, retained_tag) + else None + ), ) monkeypatch.setattr( "Magpie.modes.benchmark.tracelens_vllm_image.secrets.token_hex", @@ -271,6 +285,37 @@ def test_local_build_rejects_reserved_tag_bound_to_different_image(monkeypatch): raise AssertionError("mismatched reserved tag was accepted") +def test_local_build_rejects_retained_tag_bound_to_different_image(monkeypatch): + base_id = "sha256:" + "b" * 64 + wrong_id = "sha256:" + "d" * 64 + identity = replace( + _identity(), + base_image=base_id, + base_image_id=base_id, + base_image_locator=base_id, + ) + retained_tag = _retained_local_base_tag(base_id) + + monkeypatch.setattr( + "Magpie.modes.benchmark.tracelens_vllm_image.docker_image_id", + lambda image: wrong_id if image == retained_tag else base_id, + ) + monkeypatch.setattr( + "Magpie.modes.benchmark.tracelens_vllm_image.subprocess.run", + lambda *_args, **_kwargs: (_ for _ in ()).throw( + AssertionError("mismatched retention tag must fail before mutation") + ), + ) + + try: + _acquire_build_base_reference(identity) + except RuntimeError as exc: + assert "Content-addressed local TraceLens base tag" in str(exc) + assert wrong_id in str(exc) + else: + raise AssertionError("mismatched retained tag was accepted") + + def test_repository_digest_build_rejects_locator_id_mismatch(monkeypatch): identity = _identity() wrong_id = "sha256:" + "d" * 64 @@ -310,6 +355,7 @@ def test_local_build_detects_post_build_retag_without_removing_foreign_tag( base_image_locator=base_id, ) local_tag = _temporary_local_base_tag(base_id, nonce=nonce) + retained_tag = _retained_local_base_tag(base_id) records = {base_id: base_id} commands = [] @@ -325,7 +371,7 @@ def test_local_build_detects_post_build_retag_without_removing_foreign_tag( def fake_run(cmd, **_kwargs): commands.append(list(cmd)) assert cmd[:3] == ["docker", "image", "tag"] - records[local_tag] = base_id + records[cmd[4]] = base_id return subprocess.CompletedProcess(cmd, 0, stdout="", stderr="") monkeypatch.setattr( @@ -350,6 +396,7 @@ def fake_run(cmd, **_kwargs): else: raise AssertionError("cleanup removed a tag whose ownership was lost") assert all(cmd[:3] != ["docker", "image", "rm"] for cmd in commands) + assert records[retained_tag] == base_id def test_concurrent_local_builds_acquire_distinct_owned_tags(monkeypatch): @@ -362,6 +409,7 @@ def test_concurrent_local_builds_acquire_distinct_owned_tags(monkeypatch): base_image_locator=base_id, ) records = {base_id: base_id} + retained_tag = _retained_local_base_tag(base_id) monkeypatch.setattr( "Magpie.modes.benchmark.tracelens_vllm_image.docker_image_id", @@ -391,8 +439,12 @@ def fake_run(cmd, **_kwargs): assert first.locator != second.locator assert first.owns_temporary_tag and second.owns_temporary_tag _release_build_base_reference(first) + assert records[retained_tag] == base_id + assert records[base_id] == base_id assert second.locator in records _release_build_base_reference(second) + assert records[retained_tag] == base_id + assert records[base_id] == base_id def test_existing_image_validation_preserves_vllm_grpc_and_excludes_pollution( @@ -523,6 +575,14 @@ def fake_build(**kwargs): "command": ["docker", "build", ""], "source_wheel_command": ["docker", "run", ""], "requirements_download_command": ["docker", "run", "pip", "download"], + "base_binding": { + "image_id": identity.base_image_id, + "provenance_locator": identity.base_image_locator, + "build_reference_kind": "repository-digest", + "temporary_tag_removed": False, + "retained_local_reference": None, + "retained_local_reference_created": False, + }, "image_id": "sha256:" + "6" * 64, "image_labels": identity.labels(), "dependency_wheels": _wheel_manifest(), @@ -549,4 +609,6 @@ def fake_build(**kwargs): assert result["stale_image_rejected"] is True assert result["stale_image_rejection_reason"] == "identity label mismatch" assert result["built"] is True + assert result["base_binding"]["image_id"] == identity.base_image_id + assert result["base_binding"]["build_reference_kind"] == "repository-digest" assert result["public_runtime_validation"]["valid"] is True From 4773cd1468f56fefec35e6e01df9c85755b7a265 Mon Sep 17 00:00:00 2001 From: Vincent Ouyang Date: Sat, 8 Aug 2026 08:20:53 +0000 Subject: [PATCH 10/11] Bind Docker benchmarks to immutable runtime evidence --- Magpie/main.py | 31 +- Magpie/modes/benchmark/benchmarker.py | 168 +++++++- Magpie/modes/benchmark/result.py | 19 + Magpie/modes/benchmark/serving_runtime.py | 306 +++++++++++++++ docs/conceptual/benchmarking-architecture.md | 3 +- docs/how-to/benchmarking/benchmark.md | 15 +- tests/test_benchmark_support.py | 35 +- tests/test_serving_runtime_receipt.py | 390 +++++++++++++++++++ 8 files changed, 952 insertions(+), 15 deletions(-) create mode 100644 Magpie/modes/benchmark/serving_runtime.py create mode 100644 tests/test_serving_runtime_receipt.py diff --git a/Magpie/main.py b/Magpie/main.py index 88d5b09..36ee781 100644 --- a/Magpie/main.py +++ b/Magpie/main.py @@ -930,8 +930,29 @@ def load_benchmark_config(benchmark_config_path: Path) -> Dict[str, Any]: Returns: Dictionary with benchmark configuration """ - data = load_yaml(benchmark_config_path) - return data.get("benchmark", {}) + config, _input_sha256 = load_benchmark_config_with_sha256( + benchmark_config_path + ) + return config + + +def load_benchmark_config_with_sha256( + benchmark_config_path: Path, +) -> tuple[Dict[str, Any], str]: + """Parse and hash the same exact benchmark YAML bytes once.""" + + from .modes.benchmark.serving_runtime import sha256_bytes + + if not benchmark_config_path.exists(): + return {}, "" + raw = benchmark_config_path.read_bytes() + loaded = yaml.safe_load(raw.decode("utf-8")) or {} + if not isinstance(loaded, dict): + return {}, sha256_bytes(raw) + benchmark = loaded.get("benchmark", {}) + if not isinstance(benchmark, dict): + return {}, sha256_bytes(raw) + return benchmark, sha256_bytes(raw) def run_gap_analysis_standalone(args) -> int: @@ -1032,10 +1053,13 @@ def run_benchmark(args, config: Dict[str, Any]) -> int: # Build benchmark config benchmark_cfg = {} + input_config_sha256 = None if args.benchmark_config: # Load from config file - benchmark_cfg = load_benchmark_config(args.benchmark_config) + benchmark_cfg, input_config_sha256 = load_benchmark_config_with_sha256( + args.benchmark_config + ) if not benchmark_cfg: logger.error(f"No benchmark config found in {args.benchmark_config}") return 1 @@ -1103,6 +1127,7 @@ def run_benchmark(args, config: Dict[str, Any]) -> int: benchmarker = BenchmarkMode( config=benchmark_config, output_dir=str(args.output_dir), + input_config_sha256=input_config_sha256, ) result = benchmarker.run() diff --git a/Magpie/modes/benchmark/benchmarker.py b/Magpie/modes/benchmark/benchmarker.py index b4f150b..14f81bc 100644 --- a/Magpie/modes/benchmark/benchmarker.py +++ b/Magpie/modes/benchmark/benchmarker.py @@ -44,6 +44,14 @@ from .model_revision import collect_model_revision_evidence from .quality import parse_lm_eval_quality from .result import BenchmarkResult, LatencyMetrics, ResultParser, ThroughputMetrics +from .serving_runtime import ( + finalize_serving_runtime_receipt, + pending_serving_runtime_receipt, + resolve_docker_image_id, + unresolved_serving_runtime_receipt, + validate_prepared_command, + write_serving_runtime_receipt, +) from .targeted_trace import run_targeted_trace_analysis from .tracelens import TraceLensAnalyzer from .tracelens_inference import ( @@ -95,6 +103,7 @@ def __init__( config: BenchmarkConfig, image_config_path: Optional[str] = None, output_dir: str = "./results", + input_config_sha256: Optional[str] = None, ): """ Initialize benchmark mode. @@ -103,6 +112,8 @@ def __init__( config: Benchmark configuration image_config_path: Path to benchmark_images.yaml output_dir: Base directory for results + input_config_sha256: SHA-256 of the exact input benchmark YAML + bytes. Docker execution requires this provenance binding. """ self.config = config self.image_selector = ImageSelector(image_config_path) @@ -113,6 +124,10 @@ def __init__( ) self._task_id: Optional[str] = None self._resolved_docker_image: Optional[str] = None + self._requested_docker_image: Optional[str] = None + self._input_config_sha256 = str(input_config_sha256 or "") + self._serving_runtime_receipt: Optional[Dict[str, Any]] = None + self._serving_runtime_workspace: Optional[Path] = None self._inferencex_source_path: Optional[str] = None self._inferencex_runtime_receipt: Optional[Dict[str, Any]] = None self._lm_eval_runtime: Optional[LmEvalRuntime] = None @@ -154,6 +169,10 @@ def run(self, task_id: Optional[str] = None) -> BenchmarkResult: self._inferencex_runtime_receipt = None self._lm_eval_runtime = None self._lm_eval_runtime_evidence = None + self._requested_docker_image = None + self._resolved_docker_image = None + self._serving_runtime_receipt = None + self._serving_runtime_workspace = None # 0. Resolve the caller's InferenceX source checkout. On repeated runs # of one BenchmarkMode instance, never treat the prior run's disposable @@ -181,6 +200,7 @@ def run(self, task_id: Optional[str] = None) -> BenchmarkResult: # into a private runtime tree. Magpie and TraceLens may modify only this # tree; the dependency/source checkout is never a write target. workspace = self.workspace_mgr.create(self.config.to_dict()) + self._serving_runtime_workspace = workspace try: inferencex_runtime = materialize_inferencex_runtime( Path(self._inferencex_source_path), @@ -420,15 +440,63 @@ def run(self, task_id: Optional[str] = None) -> BenchmarkResult: if not self.config.is_server_lifecycle: self._cleanup_server_processes(self.config.framework) else: - docker_image = self._select_image() - self._resolved_docker_image = docker_image + requested_image = self._select_image() + self._requested_docker_image = requested_image + container_name = f"magpie-benchmark-{self._task_id}" + resolved_image_id, resolution_errors = resolve_docker_image_id( + requested_image + ) + if resolution_errors: + self._serving_runtime_receipt = ( + unresolved_serving_runtime_receipt( + input_config_sha256=self._input_config_sha256, + requested_image=requested_image, + container_name=container_name, + errors=resolution_errors, + ) + ) + self._persist_serving_runtime_receipt() + if gpu_monitor is not None: + gpu_monitor.stop() + return self._workspace_failure( + workspace, + start_time, + "Docker serving runtime preflight failed: immutable image " + "identity could not be resolved", + ) + self._resolved_docker_image = resolved_image_id docker_cmd = self._build_docker_command( - docker_image=docker_image, + docker_image=resolved_image_id, workspace=workspace, runner_type=runner_type, ) - logger.info(f"Running benchmark in container with image: {docker_image}") - logger.debug(f"Docker command: {' '.join(docker_cmd)}") + self._serving_runtime_receipt = pending_serving_runtime_receipt( + execution_mode="docker", + input_config_sha256=self._input_config_sha256, + requested_image=requested_image, + resolved_image_id=resolved_image_id, + container_name=container_name, + docker_argv=docker_cmd, + ) + self._persist_serving_runtime_receipt() + if self._serving_runtime_receipt["errors"]: + if gpu_monitor is not None: + gpu_monitor.stop() + return self._workspace_failure( + workspace, + start_time, + "Docker serving runtime preflight failed: config or command " + "binding is incomplete", + ) + logger.info( + "Running benchmark in container: requested=%s resolved=%s", + requested_image, + resolved_image_id, + ) + logger.debug( + "Docker command bound by SHA-256: %s", + self._serving_runtime_receipt["docker_argv_sha256"], + ) result, stdout, stderr = self._execute_benchmark(docker_cmd, workspace) # 7b. Stop GPU monitor and collect stats @@ -446,6 +514,7 @@ def run(self, task_id: Optional[str] = None) -> BenchmarkResult: result.reward_eligible = self.config.reward_eligible result.inferencex_runtime_receipt = self._inferencex_runtime_receipt result.lm_eval_runtime_receipt = self._collect_lm_eval_evidence(workspace) + result.serving_runtime_receipt = self._copy_serving_runtime_receipt() runtime_evidence = result.lm_eval_runtime_receipt if runtime_evidence["requested"] and not runtime_evidence["verified"]: result.success = False @@ -690,6 +759,39 @@ def _collect_lm_eval_evidence(self, workspace: Path) -> Dict[str, Any]: ) return dict(self._lm_eval_runtime_evidence) + def _copy_serving_runtime_receipt(self) -> Optional[Dict[str, Any]]: + if self._serving_runtime_receipt is None: + return None + receipt = dict(self._serving_runtime_receipt) + receipt["errors"] = list(receipt.get("errors", [])) + return receipt + + def _persist_serving_runtime_receipt(self) -> None: + if ( + self._serving_runtime_receipt is None + or self._serving_runtime_workspace is None + ): + return + write_serving_runtime_receipt( + self._serving_runtime_workspace, + self._serving_runtime_receipt, + ) + + def _finish_serving_runtime_receipt( + self, + *, + process_succeeded: bool, + process_error: str = "", + ) -> None: + if self._serving_runtime_receipt is None: + return + self._serving_runtime_receipt = finalize_serving_runtime_receipt( + self._serving_runtime_receipt, + process_succeeded=process_succeeded, + process_error=process_error, + ) + self._persist_serving_runtime_receipt() + def _workspace_failure( self, workspace: Path, @@ -709,6 +811,7 @@ def _workspace_failure( reward_eligible=self.config.reward_eligible, inferencex_runtime_receipt=self._inferencex_runtime_receipt, lm_eval_runtime_receipt=self._collect_lm_eval_evidence(workspace), + serving_runtime_receipt=self._copy_serving_runtime_receipt(), ) result.errors.append(message) self.workspace_mgr.save_report(result.to_dict()) @@ -896,7 +999,7 @@ def _build_docker_command( Build Docker run command. Args: - docker_image: Docker image to use + docker_image: Immutable ``sha256:...`` Docker image ID to use workspace: Workspace directory path runner_type: InferenceX runner type @@ -1857,6 +1960,32 @@ def _execute_benchmark(self, cmd: List[str], workspace: Path) -> tuple: result = BenchmarkResult() stdout = "" stderr = "" + + if self._serving_runtime_receipt is None: + self._serving_runtime_receipt = unresolved_serving_runtime_receipt( + input_config_sha256=self._input_config_sha256, + requested_image=self._requested_docker_image or "", + container_name=f"magpie-benchmark-{self._task_id}", + errors=("serving runtime receipt was not prepared",), + ) + self._persist_serving_runtime_receipt() + + binding_errors = validate_prepared_command( + self._serving_runtime_receipt, + cmd, + ) + if binding_errors: + self._serving_runtime_receipt = finalize_serving_runtime_receipt( + self._serving_runtime_receipt, + process_succeeded=False, + process_error="; ".join(binding_errors), + ) + self._persist_serving_runtime_receipt() + result.errors.append( + "Docker serving runtime command binding failed before launch" + ) + result.serving_runtime_receipt = self._copy_serving_runtime_receipt() + return result, stdout, stderr try: # Run Docker command @@ -1881,9 +2010,17 @@ def _execute_benchmark(self, cmd: List[str], workspace: Path) -> tuple: if process.returncode == 0: result.success = True + self._finish_serving_runtime_receipt(process_succeeded=True) logger.info("Benchmark completed successfully") else: result.success = False + self._finish_serving_runtime_receipt( + process_succeeded=False, + process_error=( + "benchmark process exited with code " + f"{process.returncode}" + ), + ) result.errors.append(f"Docker command failed with code {process.returncode}") if stderr: result.errors.append(f"stderr: {stderr[:1000]}") @@ -1895,6 +2032,10 @@ def _execute_benchmark(self, cmd: List[str], workspace: Path) -> tuple: logger.debug(f"stdout (last 500 chars): {stdout[-500:]}") except subprocess.TimeoutExpired as e: + self._finish_serving_runtime_receipt( + process_succeeded=False, + process_error="benchmark process timed out", + ) result.errors.append(f"Benchmark timed out after {self.config.timeout_seconds}s") logger.error("Benchmark timed out") @@ -1914,8 +2055,17 @@ def _execute_benchmark(self, cmd: List[str], workspace: Path) -> tuple: pass except Exception as e: + self._finish_serving_runtime_receipt( + process_succeeded=False, + process_error=( + "benchmark process execution failed " + f"({type(e).__name__})" + ), + ) result.errors.append(f"Benchmark execution error: {str(e)}") logger.exception(f"Benchmark execution failed: {e}") + + result.serving_runtime_receipt = self._copy_serving_runtime_receipt() return result, stdout, stderr @@ -1954,7 +2104,11 @@ def _fix_workspace_ownership(self, workspace: Path) -> None: if uid == 0: return # nothing to fix - image = self.config.docker_image or "busybox" + image = ( + self._resolved_docker_image + or self.config.docker_image + or "busybox" + ) try: subprocess.run( [ diff --git a/Magpie/modes/benchmark/result.py b/Magpie/modes/benchmark/result.py index f02a6ad..dff91eb 100644 --- a/Magpie/modes/benchmark/result.py +++ b/Magpie/modes/benchmark/result.py @@ -170,6 +170,10 @@ class BenchmarkResult: # In-container proof that the exact caller-supplied, read-only lm-eval # runtime was validated and imported. Required whenever RUN_EVAL=true. lm_eval_runtime_receipt: Optional[Dict[str, Any]] = None + + # End-to-end binding from the input benchmark bytes to the exact immutable + # Docker image ID, owned container name, hashed argv, and process outcome. + serving_runtime_receipt: Optional[Dict[str, Any]] = None # Errors errors: List[str] = field(default_factory=list) @@ -209,6 +213,7 @@ def to_dict(self) -> Dict[str, Any]: "model_revision_receipt": self.model_revision_receipt, "inferencex_runtime_receipt": self.inferencex_runtime_receipt, "lm_eval_runtime_receipt": self.lm_eval_runtime_receipt, + "serving_runtime_receipt": self.serving_runtime_receipt, "errors": self.errors, } # Scriptable (server-less) extras — e.g. xDiT diffusion. Only emit when @@ -282,6 +287,20 @@ def get_summary(self) -> str: f" Mount: {runtime.get('mount_mode') or 'not activated'}", ] ) + + if self.serving_runtime_receipt is not None: + serving = self.serving_runtime_receipt + lines.extend( + [ + "", + "Serving runtime evidence:", + f" Verified: {serving.get('verified', False)}", + " Image ID: " + f"{serving.get('resolved_image_id') or 'not resolved'}", + " Config SHA-256: " + f"{serving.get('input_config_sha256') or 'not supplied'}", + ] + ) if self.top_bottlenecks: lines.extend([ diff --git a/Magpie/modes/benchmark/serving_runtime.py b/Magpie/modes/benchmark/serving_runtime.py new file mode 100644 index 0000000..7b5b9ab --- /dev/null +++ b/Magpie/modes/benchmark/serving_runtime.py @@ -0,0 +1,306 @@ +############################################################################### +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. +# +# See LICENSE for license information. +############################################################################### +"""Evidence binding a Docker serving benchmark to its immutable runtime.""" + +from __future__ import annotations + +import hashlib +import json +import re +import subprocess +from pathlib import Path +from typing import Any, Iterable, Mapping, Sequence, Tuple + +SERVING_RUNTIME_SCHEMA = "magpie.serving-runtime-receipt/v1" +SERVING_RUNTIME_RECEIPT = "serving_runtime_receipt.json" +SERVING_RUNTIME_KEYS = ( + "schema", + "execution_mode", + "input_config_sha256", + "requested_image", + "resolved_image_id", + "container_name", + "docker_argv_sha256", + "process_succeeded", + "verified", + "errors", +) + +_SHA256_RE = re.compile(r"^[0-9a-f]{64}$") +_IMAGE_ID_RE = re.compile(r"^sha256:[0-9a-f]{64}$") +_MAX_ERRORS = 8 +_MAX_ERROR_LENGTH = 240 + + +def sha256_bytes(content: bytes) -> str: + """Return the lowercase SHA-256 digest of exact input bytes.""" + + return hashlib.sha256(content).hexdigest() + + +def canonical_docker_argv_sha256(argv: Sequence[str]) -> str: + """Hash an unambiguous JSON encoding of an argv vector. + + Only the digest is persisted. In particular, environment values such as a + Hugging Face token that may be present in the process argv are never copied + into the receipt. + """ + + encoded = json.dumps( + [str(item) for item in argv], + ensure_ascii=False, + separators=(",", ":"), + ).encode("utf-8") + return sha256_bytes(encoded) + + +def resolve_docker_image_id(requested_image: str) -> Tuple[str, Tuple[str, ...]]: + """Resolve a tag or digest reference to the local immutable Docker image ID. + + A raw ``sha256:<64 hex>`` ID is already immutable and is used directly. + Other references are resolved with one fixed, non-shell Docker invocation. + No Docker stderr is retained in the receipt. + """ + + requested = str(requested_image).strip() + if _IMAGE_ID_RE.fullmatch(requested): + return requested, () + if not requested: + return "", ("requested Docker image is empty",) + + command = [ + "docker", + "image", + "inspect", + "--format", + "{{.Id}}", + "--", + requested, + ] + try: + completed = subprocess.run( + command, + capture_output=True, + text=True, + timeout=30, + check=False, + ) + except subprocess.TimeoutExpired: + return "", ("Docker image inspection timed out",) + except OSError as exc: + return "", ( + f"Docker image inspection could not start ({type(exc).__name__})", + ) + + if completed.returncode != 0: + return "", ( + f"Docker image inspection failed with code {completed.returncode}", + ) + lines = [line.strip() for line in (completed.stdout or "").splitlines()] + identities = [line for line in lines if line] + if len(identities) != 1 or not _IMAGE_ID_RE.fullmatch(identities[0]): + return "", ("Docker image inspection returned an invalid image ID",) + return identities[0], () + + +def pending_serving_runtime_receipt( + *, + execution_mode: str, + input_config_sha256: str, + requested_image: str, + resolved_image_id: str, + container_name: str, + docker_argv: Sequence[str], + prior_errors: Iterable[str] = (), +) -> dict[str, Any]: + """Build the pre-execution receipt and validate its command bindings.""" + + errors = list(prior_errors) + config_digest = str(input_config_sha256 or "") + image_id = str(resolved_image_id or "") + argv_digest = canonical_docker_argv_sha256(docker_argv) if docker_argv else "" + + if not _SHA256_RE.fullmatch(config_digest): + errors.append("input config SHA-256 is missing or invalid") + if not _IMAGE_ID_RE.fullmatch(image_id): + errors.append("resolved Docker image ID is missing or invalid") + errors.extend( + docker_command_binding_errors( + docker_argv, + expected_container_name=container_name, + expected_image_id=image_id, + ) + ) + return _receipt( + execution_mode=execution_mode, + input_config_sha256=config_digest, + requested_image=requested_image, + resolved_image_id=image_id, + container_name=container_name, + docker_argv_sha256=argv_digest, + process_succeeded=False, + verified=False, + errors=errors, + ) + + +def unresolved_serving_runtime_receipt( + *, + input_config_sha256: str, + requested_image: str, + container_name: str, + errors: Iterable[str], +) -> dict[str, Any]: + """Build a receipt when an immutable image could not be resolved.""" + + combined = list(errors) + if not _SHA256_RE.fullmatch(str(input_config_sha256 or "")): + combined.append("input config SHA-256 is missing or invalid") + return _receipt( + execution_mode="docker", + input_config_sha256=str(input_config_sha256 or ""), + requested_image=requested_image, + resolved_image_id="", + container_name=container_name, + docker_argv_sha256="", + process_succeeded=False, + verified=False, + errors=combined, + ) + + +def validate_prepared_command( + receipt: Mapping[str, Any], + docker_argv: Sequence[str], +) -> Tuple[str, ...]: + """Reject a command that differs from the prepared runtime receipt.""" + + errors = [] + if tuple(receipt.keys()) != SERVING_RUNTIME_KEYS: + errors.append("serving runtime receipt has an invalid shape") + if receipt.get("schema") != SERVING_RUNTIME_SCHEMA: + errors.append("serving runtime receipt schema is invalid") + if receipt.get("execution_mode") != "docker": + errors.append("serving runtime execution mode is invalid") + expected_digest = receipt.get("docker_argv_sha256") + if expected_digest != canonical_docker_argv_sha256(docker_argv): + errors.append("Docker argv does not match its prepared digest") + errors.extend( + docker_command_binding_errors( + docker_argv, + expected_container_name=str(receipt.get("container_name", "")), + expected_image_id=str(receipt.get("resolved_image_id", "")), + ) + ) + errors.extend(str(item) for item in receipt.get("errors", [])) + return tuple(_bounded_errors(errors)) + + +def finalize_serving_runtime_receipt( + receipt: Mapping[str, Any], + *, + process_succeeded: bool, + process_error: str = "", +) -> dict[str, Any]: + """Finalize process status while preserving the prepared identity fields.""" + + errors = list(receipt.get("errors", [])) + if process_error: + errors.append(process_error) + bounded = _bounded_errors(errors) + succeeded = bool(process_succeeded) + return _receipt( + execution_mode=str(receipt.get("execution_mode", "docker")), + input_config_sha256=str(receipt.get("input_config_sha256", "")), + requested_image=str(receipt.get("requested_image", "")), + resolved_image_id=str(receipt.get("resolved_image_id", "")), + container_name=str(receipt.get("container_name", "")), + docker_argv_sha256=str(receipt.get("docker_argv_sha256", "")), + process_succeeded=succeeded, + verified=succeeded and not bounded, + errors=bounded, + ) + + +def write_serving_runtime_receipt( + workspace: Path, + receipt: Mapping[str, Any], +) -> None: + """Atomically persist the bounded receipt beside the benchmark report.""" + + destination = Path(workspace) / SERVING_RUNTIME_RECEIPT + temporary = destination.with_suffix(".json.tmp") + temporary.write_text( + json.dumps(dict(receipt), indent=2, sort_keys=False) + "\n", + encoding="utf-8", + ) + temporary.replace(destination) + + +def docker_command_binding_errors( + docker_argv: Sequence[str], + *, + expected_container_name: str, + expected_image_id: str, +) -> Tuple[str, ...]: + """Validate the exact container-name and image slots in ``docker run``.""" + + argv = [str(item) for item in docker_argv] + errors = [] + if argv[:2] != ["docker", "run"]: + errors.append("Docker argv is not a docker run command") + + name_positions = [index for index, item in enumerate(argv) if item == "--name"] + if len(name_positions) != 1 or name_positions[0] + 1 >= len(argv): + errors.append("Docker argv does not contain one container name") + elif argv[name_positions[0] + 1] != expected_container_name: + errors.append("Docker argv container name does not match the receipt") + + entry_positions = [ + index for index, item in enumerate(argv) if item == "--entrypoint" + ] + if len(entry_positions) != 1 or entry_positions[0] + 2 >= len(argv): + errors.append("Docker argv does not contain one image slot") + elif argv[entry_positions[0] + 2] != expected_image_id: + errors.append("Docker argv image does not match the resolved image ID") + return tuple(_bounded_errors(errors)) + + +def _receipt( + *, + execution_mode: str, + input_config_sha256: str, + requested_image: str, + resolved_image_id: str, + container_name: str, + docker_argv_sha256: str, + process_succeeded: bool, + verified: bool, + errors: Iterable[str], +) -> dict[str, Any]: + return { + "schema": SERVING_RUNTIME_SCHEMA, + "execution_mode": execution_mode, + "input_config_sha256": input_config_sha256, + "requested_image": requested_image, + "resolved_image_id": resolved_image_id, + "container_name": container_name, + "docker_argv_sha256": docker_argv_sha256, + "process_succeeded": process_succeeded, + "verified": verified, + "errors": _bounded_errors(errors), + } + + +def _bounded_errors(errors: Iterable[str]) -> list[str]: + bounded = [] + for error in errors: + text = " ".join(str(error).split())[:_MAX_ERROR_LENGTH] + if text and text not in bounded: + bounded.append(text) + if len(bounded) == _MAX_ERRORS: + break + return bounded diff --git a/docs/conceptual/benchmarking-architecture.md b/docs/conceptual/benchmarking-architecture.md index aa94455..2979674 100644 --- a/docs/conceptual/benchmarking-architecture.md +++ b/docs/conceptual/benchmarking-architecture.md @@ -25,13 +25,14 @@ Benchmark mode consists of the following Python modules. | `TraceLensInferencePipeline` | `tracelens_inference.py` | Inference-aware TraceLens split/report flow and simple roofline summaries | | `GapAnalyzer` | `gap_analysis.py` | Kernel bottleneck analysis | | `BenchmarkResult` | `result.py` | Result data structures | +| `Serving runtime evidence` | `serving_runtime.py` | Fail-closed binding from input config bytes to the immutable Docker runtime | ### Execution flow Each benchmark run proceeds through the following stages. 1. **Configuration Loading**: Parse YAML config into `BenchmarkConfig` -2. **Runtime Setup**: For `run_mode: docker`, prepare a container with InferenceX; for `local`, use the host environment +2. **Runtime Setup**: For `run_mode: docker`, hash the exact input YAML, resolve the requested image to an immutable Docker image ID, hash the exact container argv, and prepare a serving-runtime receipt; for `local`, use the host environment 3. **Server Launch**: Start vLLM/SGLang server (in container or on host per `run_mode`) 4. **Client Execution**: Run benchmark client with profiling enabled 5. **Trace Collection**: Torch profiler traces saved to workspace diff --git a/docs/how-to/benchmarking/benchmark.md b/docs/how-to/benchmarking/benchmark.md index 8033e0f..04c3aac 100644 --- a/docs/how-to/benchmarking/benchmark.md +++ b/docs/how-to/benchmarking/benchmark.md @@ -61,8 +61,8 @@ python -m Magpie benchmark --trace-dir results/benchmark_vllm_/ # SGLang benchmark python -m Magpie benchmark --benchmark-config examples/benchmarks/benchmark_sglang_dsr1.yaml -# Ad-hoc CLI without a YAML file (framework + model; optional torch profiler) -python -m Magpie benchmark vllm --model deepseek-ai/DeepSeek-R1-0528 --torch-profiler +# Ad-hoc local CLI without a YAML file (Docker evidence requires YAML input) +python -m Magpie benchmark vllm --model deepseek-ai/DeepSeek-R1-0528 --run-mode local --torch-profiler ``` ## Output structure @@ -78,6 +78,7 @@ results/benchmark_vllm_/ ├── container_stderr.log # Container stderr ├── inferencex_result.json # Raw InferenceX output ├── inferencex_runtime_receipt.json # Exact source/runtime identity +├── serving_runtime_receipt.json # Input YAML -> immutable Docker process proof ├── inferencex_runtime/ # Private run-scoped InferenceX tree ├── model_revision_receipt.json # Requested/resolved HF snapshot (when pinned) ├── lm_eval_runtime_manifest.json # Preserved content/identity manifest @@ -118,6 +119,16 @@ machine-readable `params_json` for matched TraceLens `param:*` metadata. The primary summary file is **`benchmark_report.json`**, written to the run workspace directory. It aggregates throughput, latency, and optional `gap_analysis` and `tracelens_analysis` sections. +For Docker runs, `serving_runtime_receipt` uses schema +`magpie.serving-runtime-receipt/v1`. It binds the SHA-256 of the exact +`--benchmark-config` bytes to the requested image, the locally resolved +immutable `sha256:...` image ID, the exact owned container name, and a SHA-256 +of the canonical Docker argv. The argv itself is not persisted, so values such +as `HF_TOKEN` are not copied into evidence. Magpie executes the resolved image +ID rather than the mutable tag. Image-inspection failure, a missing input +digest, or any command-binding mismatch fails before container launch; +`verified` becomes true only after the bound process exits successfully. + Every report declares `run_kind` and `reward_eligible`. A `run_kind: measurement` run rejects heavy profilers; diagnostic runs and all TargetedKernelTrace artifacts have `reward_eligible: false`. When `RUN_EVAL=true`, diff --git a/tests/test_benchmark_support.py b/tests/test_benchmark_support.py index a957ead..09d0625 100644 --- a/tests/test_benchmark_support.py +++ b/tests/test_benchmark_support.py @@ -21,6 +21,9 @@ ) from Magpie.modes.benchmark.image_selector import ImageSelector from Magpie.modes.benchmark.result import BenchmarkResult, ResultParser +from Magpie.modes.benchmark.serving_runtime import ( + pending_serving_runtime_receipt, +) from Magpie.modes.benchmark.quality import parse_lm_eval_quality from Magpie.modes.benchmark.tracelens_inference import ( SGLANG_SHAPE_DISCOVERY_FLAG, @@ -355,14 +358,37 @@ def test_benchmark_timeout_kills_latched_protected_container( mode = BenchmarkMode( BenchmarkConfig(framework="vllm", model="demo", run_mode="docker"), output_dir=str(tmp_path / "results"), + input_config_sha256="a" * 64, ) mode._task_id = "timed-out-task" mode._docker_stop_protection_active = True + image_id = "sha256:" + "b" * 64 + command = [ + "docker", + "run", + "--name", + "magpie-benchmark-timed-out-task", + "--entrypoint", + "/bin/bash", + image_id, + "-c", + "true", + ] + mode._requested_docker_image = "example/image:fixed" + mode._resolved_docker_image = image_id + mode._serving_runtime_receipt = pending_serving_runtime_receipt( + execution_mode="docker", + input_config_sha256="a" * 64, + requested_image="example/image:fixed", + resolved_image_id=image_id, + container_name="magpie-benchmark-timed-out-task", + docker_argv=command, + ) calls = [] def fake_run(command, **kwargs): calls.append((command, kwargs)) - if command == ["docker", "run", "example"]: + if command[:2] == ["docker", "run"]: raise subprocess.TimeoutExpired(command, timeout=1) return subprocess.CompletedProcess(command, 0) @@ -372,12 +398,17 @@ def fake_run(command, **kwargs): ) result, _stdout, _stderr = mode._execute_benchmark( - ["docker", "run", "example"], + command, tmp_path, ) assert result.success is False assert result.errors == ["Benchmark timed out after 3600.0s"] + assert result.serving_runtime_receipt["process_succeeded"] is False + assert result.serving_runtime_receipt["verified"] is False + assert result.serving_runtime_receipt["errors"] == [ + "benchmark process timed out" + ] assert calls[-1] == ( ["docker", "kill", "magpie-benchmark-timed-out-task"], {"capture_output": True, "timeout": 30, "check": False}, diff --git a/tests/test_serving_runtime_receipt.py b/tests/test_serving_runtime_receipt.py new file mode 100644 index 0000000..ee8c483 --- /dev/null +++ b/tests/test_serving_runtime_receipt.py @@ -0,0 +1,390 @@ +import json +import subprocess +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from Magpie.main import load_benchmark_config_with_sha256, run_benchmark +from Magpie.modes.benchmark.benchmarker import BenchmarkMode +from Magpie.modes.benchmark.config import BenchmarkConfig +from Magpie.modes.benchmark.result import BenchmarkResult +from Magpie.modes.benchmark.serving_runtime import ( + SERVING_RUNTIME_KEYS, + SERVING_RUNTIME_SCHEMA, + canonical_docker_argv_sha256, + pending_serving_runtime_receipt, + resolve_docker_image_id, + sha256_bytes, +) + + +def _docker_command(container_name: str, image_id: str) -> list[str]: + return [ + "docker", + "run", + "--rm", + "--name", + container_name, + "--entrypoint", + "/bin/bash", + image_id, + "-c", + "true", + ] + + +def test_cli_hashes_the_exact_yaml_bytes_it_parses(tmp_path): + config_path = tmp_path / "benchmark.yaml" + raw = ( + b"# byte identity matters\r\n" + b"benchmark:\r\n" + b" framework: vllm\r\n" + b" model: example/model\r\n" + ) + config_path.write_bytes(raw) + + config, digest = load_benchmark_config_with_sha256(config_path) + + assert config == {"framework": "vllm", "model": "example/model"} + assert digest == sha256_bytes(raw) + + +def test_cli_passes_raw_config_digest_to_benchmark_mode(tmp_path, monkeypatch): + config_path = tmp_path / "benchmark.yaml" + raw = ( + b"benchmark:\n" + b" framework: vllm\n" + b" model: example/model\n" + b" run_mode: local\n" + ) + config_path.write_bytes(raw) + captured = {} + + class FakeBenchmarkMode: + def __init__(self, *, config, output_dir, input_config_sha256): + captured["config"] = config + captured["output_dir"] = output_dir + captured["input_config_sha256"] = input_config_sha256 + + def run(self): + return BenchmarkResult(success=True) + + monkeypatch.setattr( + "Magpie.modes.benchmark.BenchmarkMode", + FakeBenchmarkMode, + ) + args = SimpleNamespace( + trace_dir=None, + benchmark_config=config_path, + framework=None, + model=None, + run_mode=None, + run_kind=None, + output_dir=tmp_path / "results", + ) + + assert run_benchmark(args, {}) == 0 + assert captured["input_config_sha256"] == sha256_bytes(raw) + assert captured["config"].model == "example/model" + + +def test_resolve_tag_uses_fixed_inspect_argv(monkeypatch): + image_id = "sha256:" + "1" * 64 + calls = [] + + def fake_run(command, **kwargs): + calls.append((command, kwargs)) + return subprocess.CompletedProcess(command, 0, image_id + "\n", "") + + monkeypatch.setattr( + "Magpie.modes.benchmark.serving_runtime.subprocess.run", + fake_run, + ) + + resolved, errors = resolve_docker_image_id("example/vllm:fixed") + + assert resolved == image_id + assert errors == () + assert calls == [ + ( + [ + "docker", + "image", + "inspect", + "--format", + "{{.Id}}", + "--", + "example/vllm:fixed", + ], + { + "capture_output": True, + "text": True, + "timeout": 30, + "check": False, + }, + ) + ] + + +def test_raw_image_id_needs_no_mutable_lookup(monkeypatch): + image_id = "sha256:" + "2" * 64 + monkeypatch.setattr( + "Magpie.modes.benchmark.serving_runtime.subprocess.run", + lambda *args, **kwargs: pytest.fail("raw image ID must not be inspected"), + ) + + assert resolve_docker_image_id(image_id) == (image_id, ()) + + +def test_image_inspect_failure_is_bounded_and_does_not_copy_stderr(monkeypatch): + secret = "hf_secret_that_must_not_be_retained" + + def fake_run(command, **kwargs): + return subprocess.CompletedProcess(command, 125, "", secret) + + monkeypatch.setattr( + "Magpie.modes.benchmark.serving_runtime.subprocess.run", + fake_run, + ) + + resolved, errors = resolve_docker_image_id("missing/image:fixed") + + assert resolved == "" + assert errors == ("Docker image inspection failed with code 125",) + assert secret not in json.dumps(errors) + + +def test_pending_receipt_rejects_image_slot_mismatch(): + expected = "sha256:" + "3" * 64 + wrong = "sha256:" + "4" * 64 + command = _docker_command("magpie-benchmark-case", wrong) + + receipt = pending_serving_runtime_receipt( + execution_mode="docker", + input_config_sha256="5" * 64, + requested_image="example/vllm:fixed", + resolved_image_id=expected, + container_name="magpie-benchmark-case", + docker_argv=command, + ) + + assert receipt["verified"] is False + assert receipt["process_succeeded"] is False + assert receipt["docker_argv_sha256"] == canonical_docker_argv_sha256( + command + ) + assert receipt["errors"] == [ + "Docker argv image does not match the resolved image ID" + ] + + +def test_success_receipt_binds_command_without_persisting_token( + tmp_path, + monkeypatch, +): + input_digest = "6" * 64 + image_id = "sha256:" + "7" * 64 + container_name = "magpie-benchmark-success" + secret = "hf_a_unique_secret_value" + command = _docker_command(container_name, image_id) + command[5:5] = ["-e", f"HF_TOKEN={secret}"] + mode = BenchmarkMode( + BenchmarkConfig(framework="vllm", model="demo", run_mode="docker"), + output_dir=str(tmp_path / "results"), + input_config_sha256=input_digest, + ) + mode._task_id = "success" + mode._requested_docker_image = "example/vllm:fixed" + mode._resolved_docker_image = image_id + mode._serving_runtime_workspace = tmp_path + mode._serving_runtime_receipt = pending_serving_runtime_receipt( + execution_mode="docker", + input_config_sha256=input_digest, + requested_image="example/vllm:fixed", + resolved_image_id=image_id, + container_name=container_name, + docker_argv=command, + ) + monkeypatch.setattr(mode, "_fix_workspace_ownership", lambda workspace: None) + monkeypatch.setattr( + "Magpie.modes.benchmark.benchmarker.subprocess.run", + lambda argv, **kwargs: subprocess.CompletedProcess(argv, 0, "ok", ""), + ) + + result, _stdout, _stderr = mode._execute_benchmark(command, tmp_path) + + receipt = result.serving_runtime_receipt + assert result.success is True + assert tuple(receipt) == SERVING_RUNTIME_KEYS + assert receipt == json.loads( + (tmp_path / "serving_runtime_receipt.json").read_text(encoding="utf-8") + ) + assert receipt["schema"] == SERVING_RUNTIME_SCHEMA + assert receipt["input_config_sha256"] == input_digest + assert receipt["resolved_image_id"] == image_id + assert receipt["container_name"] == container_name + assert receipt["docker_argv_sha256"] == canonical_docker_argv_sha256( + command + ) + assert receipt["process_succeeded"] is True + assert receipt["verified"] is True + assert receipt["errors"] == [] + assert secret not in json.dumps(receipt) + assert secret not in (tmp_path / "serving_runtime_receipt.json").read_text( + encoding="utf-8" + ) + assert BenchmarkResult(serving_runtime_receipt=receipt).to_dict()[ + "serving_runtime_receipt" + ] == receipt + + +def test_benchmark_run_reports_resolved_immutable_docker_runtime( + tmp_path, + monkeypatch, +): + source = tmp_path / "InferenceX" + source.mkdir() + requested = "example/vllm:fixed" + image_id = "sha256:" + "a" * 64 + input_digest = "b" * 64 + config = BenchmarkConfig( + framework="vllm", + model="example/model", + run_mode="docker", + run_kind="measurement", + docker_image=requested, + inferencex_path=str(source), + gpu_selection={"auto": False}, + profiler={ + "torch_profiler": {"enabled": False}, + "gpu_monitor": {"enabled": False}, + }, + ) + mode = BenchmarkMode( + config, + output_dir=str(tmp_path / "results"), + input_config_sha256=input_digest, + ) + monkeypatch.setattr( + "Magpie.modes.benchmark.benchmarker.ensure_inferencex_available", + lambda path: str(source), + ) + monkeypatch.setattr( + "Magpie.modes.benchmark.benchmarker.materialize_inferencex_runtime", + lambda source_path, workspace: SimpleNamespace( + root=source_path, + receipt={"source_commit": "c" * 40, "source_tree": "d" * 40}, + ), + ) + monkeypatch.setattr(mode, "_prepare_benchmark_scripts", lambda: None) + monkeypatch.setattr(mode, "_get_runner_type", lambda: "mi355x") + monkeypatch.setattr( + mode, + "_get_benchmark_script", + lambda runner_type: "benchmarks/vllm_mi355x.sh", + ) + monkeypatch.setattr( + "Magpie.modes.benchmark.benchmarker.detect_gpu", + lambda: ("unknown", ""), + ) + main_commands = [] + + def fake_benchmark_run(command, **kwargs): + if command[:3] == ["docker", "image", "inspect"]: + return subprocess.CompletedProcess( + command, + 0, + image_id + "\n", + "", + ) + main_commands.append(command) + workspace = mode.workspace_mgr.workspace_path + if "--name" in command: + (workspace / "inferencex_result.json").write_text( + json.dumps( + { + "request_throughput": 1.0, + "output_throughput": 10.0, + "completed": 1, + } + ), + encoding="utf-8", + ) + return subprocess.CompletedProcess(command, 0, "ok", "") + + monkeypatch.setattr( + "Magpie.modes.benchmark.benchmarker.subprocess.run", + fake_benchmark_run, + ) + + result = mode.run(task_id="full-success") + + assert result.success is True + receipt = result.serving_runtime_receipt + assert receipt["verified"] is True + assert receipt["requested_image"] == requested + assert receipt["resolved_image_id"] == image_id + assert receipt["input_config_sha256"] == input_digest + benchmark_command = next( + command for command in main_commands if "--name" in command + ) + entrypoint = benchmark_command.index("--entrypoint") + assert benchmark_command[entrypoint + 2] == image_id + report = json.loads( + ( + tmp_path + / "results" + / Path(result.workspace_dir).name + / "benchmark_report.json" + ).read_text(encoding="utf-8") + ) + assert report["serving_runtime_receipt"] == receipt + + +def test_command_digest_mismatch_fails_before_process_launch( + tmp_path, + monkeypatch, +): + image_id = "sha256:" + "8" * 64 + container_name = "magpie-benchmark-command-swap" + prepared_command = _docker_command(container_name, image_id) + changed_command = list(prepared_command) + changed_command[-1] = "false" + mode = BenchmarkMode( + BenchmarkConfig(framework="vllm", model="demo", run_mode="docker"), + output_dir=str(tmp_path / "results"), + input_config_sha256="9" * 64, + ) + mode._task_id = "command-swap" + mode._requested_docker_image = "example/vllm:fixed" + mode._resolved_docker_image = image_id + mode._serving_runtime_workspace = tmp_path + mode._serving_runtime_receipt = pending_serving_runtime_receipt( + execution_mode="docker", + input_config_sha256="9" * 64, + requested_image="example/vllm:fixed", + resolved_image_id=image_id, + container_name=container_name, + docker_argv=prepared_command, + ) + monkeypatch.setattr( + "Magpie.modes.benchmark.benchmarker.subprocess.run", + lambda *args, **kwargs: pytest.fail("mismatched command must not launch"), + ) + + result, _stdout, _stderr = mode._execute_benchmark( + changed_command, + tmp_path, + ) + + assert result.success is False + assert result.errors == [ + "Docker serving runtime command binding failed before launch" + ] + assert result.serving_runtime_receipt["process_succeeded"] is False + assert result.serving_runtime_receipt["verified"] is False + assert any( + "prepared digest" in error + for error in result.serving_runtime_receipt["errors"] + ) From 210513b31b2f3607920be4000d37fc51f14c5711 Mon Sep 17 00:00:00 2001 From: Vincent Ouyang Date: Sat, 8 Aug 2026 10:49:37 +0000 Subject: [PATCH 11/11] Bind serving receipts to TraceLens image lineage --- Magpie/modes/benchmark/benchmarker.py | 47 ++- Magpie/modes/benchmark/result.py | 11 +- Magpie/modes/benchmark/serving_runtime.py | 354 +++++++++++++++++- Magpie/modes/benchmark/tracelens_runtime.py | 2 + .../modes/benchmark/tracelens_vllm_image.py | 3 +- docs/conceptual/benchmarking-architecture.md | 2 +- docs/how-to/benchmarking/benchmark.md | 24 +- docs/how-to/benchmarking/profiling-options.md | 10 +- tests/test_benchmark_support.py | 3 + tests/test_serving_runtime_receipt.py | 239 +++++++++++- tests/test_tracelens_vllm_image.py | 32 ++ 11 files changed, 698 insertions(+), 29 deletions(-) diff --git a/Magpie/modes/benchmark/benchmarker.py b/Magpie/modes/benchmark/benchmarker.py index 14f81bc..c575f2d 100644 --- a/Magpie/modes/benchmark/benchmarker.py +++ b/Magpie/modes/benchmark/benchmarker.py @@ -123,8 +123,12 @@ def __init__( container_writable=config.run_mode == "docker", ) self._task_id: Optional[str] = None + self._configured_docker_image: Optional[str] = None + self._input_docker_image: Optional[str] = None + self._input_docker_image_id: Optional[str] = None self._resolved_docker_image: Optional[str] = None self._requested_docker_image: Optional[str] = None + self._tracelens_runtime_result: Optional[Dict[str, Any]] = None self._input_config_sha256 = str(input_config_sha256 or "") self._serving_runtime_receipt: Optional[Dict[str, Any]] = None self._serving_runtime_workspace: Optional[Path] = None @@ -169,8 +173,11 @@ def run(self, task_id: Optional[str] = None) -> BenchmarkResult: self._inferencex_runtime_receipt = None self._lm_eval_runtime = None self._lm_eval_runtime_evidence = None + self._input_docker_image = None + self._input_docker_image_id = None self._requested_docker_image = None self._resolved_docker_image = None + self._tracelens_runtime_result = None self._serving_runtime_receipt = None self._serving_runtime_workspace = None @@ -294,12 +301,18 @@ def run(self, task_id: Optional[str] = None) -> BenchmarkResult: # 4a. For Docker benchmarks, TraceLens inference can derive a patched # framework runtime image from supported official vLLM/SGLang images. tracelens_runtime_result: Optional[Dict[str, Any]] = None + if self.config.run_mode == "docker": + if self._configured_docker_image is None: + self._configured_docker_image = self._select_image() + self._input_docker_image = self._configured_docker_image + self.config.docker_image = self._configured_docker_image if ( is_tracelens_inference_enabled(self.config) and self.config.run_mode == "docker" ): try: - base_image = self._select_image() + base_image = self._input_docker_image + assert base_image is not None tracelens_runtime_result = prepare_tracelens_runtime_image( config=self.config, base_image=base_image, @@ -308,6 +321,7 @@ def run(self, task_id: Optional[str] = None) -> BenchmarkResult: self.config.docker_image = str( tracelens_runtime_result.get("image") or base_image ) + self._tracelens_runtime_result = dict(tracelens_runtime_result) logger.info( "TraceLens runtime image: %s (%s)", self.config.docker_image, @@ -442,16 +456,29 @@ def run(self, task_id: Optional[str] = None) -> BenchmarkResult: else: requested_image = self._select_image() self._requested_docker_image = requested_image + input_image = self._input_docker_image or requested_image container_name = f"magpie-benchmark-{self._task_id}" - resolved_image_id, resolution_errors = resolve_docker_image_id( - requested_image - ) + input_image_id, input_errors = resolve_docker_image_id(input_image) + self._input_docker_image_id = input_image_id or None + if requested_image == input_image: + resolved_image_id = input_image_id + runtime_errors: tuple[str, ...] = () + else: + resolved_image_id, runtime_errors = resolve_docker_image_id( + requested_image + ) + resolution_errors = tuple(input_errors) + tuple(runtime_errors) if resolution_errors: self._serving_runtime_receipt = ( unresolved_serving_runtime_receipt( input_config_sha256=self._input_config_sha256, + framework=self.config.framework, + input_image=input_image, + input_image_id=input_image_id, requested_image=requested_image, + resolved_image_id=resolved_image_id, container_name=container_name, + tracelens_runtime=self._tracelens_runtime_result, errors=resolution_errors, ) ) @@ -473,10 +500,14 @@ def run(self, task_id: Optional[str] = None) -> BenchmarkResult: self._serving_runtime_receipt = pending_serving_runtime_receipt( execution_mode="docker", input_config_sha256=self._input_config_sha256, + framework=self.config.framework, + input_image=input_image, + input_image_id=input_image_id, requested_image=requested_image, resolved_image_id=resolved_image_id, container_name=container_name, docker_argv=docker_cmd, + tracelens_runtime=self._tracelens_runtime_result, ) self._persist_serving_runtime_receipt() if self._serving_runtime_receipt["errors"]: @@ -763,6 +794,9 @@ def _copy_serving_runtime_receipt(self) -> Optional[Dict[str, Any]]: if self._serving_runtime_receipt is None: return None receipt = dict(self._serving_runtime_receipt) + derivation = receipt.get("image_derivation") + if isinstance(derivation, dict): + receipt["image_derivation"] = dict(derivation) receipt["errors"] = list(receipt.get("errors", [])) return receipt @@ -1964,8 +1998,13 @@ def _execute_benchmark(self, cmd: List[str], workspace: Path) -> tuple: if self._serving_runtime_receipt is None: self._serving_runtime_receipt = unresolved_serving_runtime_receipt( input_config_sha256=self._input_config_sha256, + framework=self.config.framework, + input_image=self._input_docker_image or "", + input_image_id=self._input_docker_image_id or "", requested_image=self._requested_docker_image or "", + resolved_image_id=self._resolved_docker_image or "", container_name=f"magpie-benchmark-{self._task_id}", + tracelens_runtime=self._tracelens_runtime_result, errors=("serving runtime receipt was not prepared",), ) self._persist_serving_runtime_receipt() diff --git a/Magpie/modes/benchmark/result.py b/Magpie/modes/benchmark/result.py index dff91eb..ea6495a 100644 --- a/Magpie/modes/benchmark/result.py +++ b/Magpie/modes/benchmark/result.py @@ -171,8 +171,9 @@ class BenchmarkResult: # runtime was validated and imported. Required whenever RUN_EVAL=true. lm_eval_runtime_receipt: Optional[Dict[str, Any]] = None - # End-to-end binding from the input benchmark bytes to the exact immutable - # Docker image ID, owned container name, hashed argv, and process outcome. + # End-to-end binding from the input benchmark bytes and configured image + # through any validated TraceLens derivation to the exact immutable runtime + # image ID, owned container name, hashed argv, and process outcome. serving_runtime_receipt: Optional[Dict[str, Any]] = None # Errors @@ -295,8 +296,12 @@ def get_summary(self) -> str: "", "Serving runtime evidence:", f" Verified: {serving.get('verified', False)}", - " Image ID: " + " Input image ID: " + f"{serving.get('input_image_id') or 'not resolved'}", + " Runtime image ID: " f"{serving.get('resolved_image_id') or 'not resolved'}", + " Derivation: " + f"{(serving.get('image_derivation') or {}).get('kind', 'unknown')}", " Config SHA-256: " f"{serving.get('input_config_sha256') or 'not supplied'}", ] diff --git a/Magpie/modes/benchmark/serving_runtime.py b/Magpie/modes/benchmark/serving_runtime.py index 7b5b9ab..d18cde5 100644 --- a/Magpie/modes/benchmark/serving_runtime.py +++ b/Magpie/modes/benchmark/serving_runtime.py @@ -11,26 +11,57 @@ import json import re import subprocess -from pathlib import Path -from typing import Any, Iterable, Mapping, Sequence, Tuple +from pathlib import Path, PurePosixPath +from typing import Any, Iterable, Mapping, Optional, Sequence, Tuple -SERVING_RUNTIME_SCHEMA = "magpie.serving-runtime-receipt/v1" +SERVING_RUNTIME_SCHEMA = "magpie.serving-runtime-receipt/v2" SERVING_RUNTIME_RECEIPT = "serving_runtime_receipt.json" SERVING_RUNTIME_KEYS = ( "schema", "execution_mode", "input_config_sha256", + "input_image", + "input_image_id", "requested_image", "resolved_image_id", + "image_derivation", "container_name", "docker_argv_sha256", "process_succeeded", "verified", "errors", ) +SERVING_IMAGE_DERIVATION_KEYS = ( + "kind", + "framework", + "runtime_schema", + "base_image", + "base_image_id", + "base_image_locator", + "derived_image", + "derived_image_id", + "tracelens_source_commit", + "tracelens_source_tree", + "patch_version", + "patch_path", + "patch_sha256", + "dependency_wheel_manifest_sha256", + "validator", + "verified", +) _SHA256_RE = re.compile(r"^[0-9a-f]{64}$") _IMAGE_ID_RE = re.compile(r"^sha256:[0-9a-f]{64}$") +_GIT_OBJECT_RE = re.compile(r"^[0-9a-f]{40}$") +_PATCH_VERSION_RE = re.compile(r"^v[0-9]+$") +_PATCH_PATH_RE = re.compile( + r"^examples/custom_workflows/inference_analysis/vllm_patches/" + r"config_vllm_v0\.([0-9]+)\.0\.patch$" +) +_REPO_DIGEST_RE = re.compile(r"^[a-z0-9][a-z0-9._:/-]*@sha256:[0-9a-f]{64}$") +_TRACELENS_VLLM_RUNTIME_SCHEMA = "magpie.tracelens-vllm-runtime/v1" +_DIRECT_VALIDATOR = "docker-image-id" +_TRACELENS_VALIDATOR = "vllm-tracelens-runtime-validation/v1" _MAX_ERRORS = 8 _MAX_ERROR_LENGTH = 240 @@ -106,27 +137,285 @@ def resolve_docker_image_id(requested_image: str) -> Tuple[str, Tuple[str, ...]] return identities[0], () +def image_derivation_receipt( + *, + framework: str, + input_image: str, + input_image_id: str, + requested_image: str, + resolved_image_id: str, + tracelens_runtime: Optional[Mapping[str, Any]] = None, +) -> Tuple[dict[str, Any], Tuple[str, ...]]: + """Bind the frozen input image to the exact image selected for execution.""" + + input_ref = str(input_image or "").strip() + input_id = str(input_image_id or "").strip() + runtime_ref = str(requested_image or "").strip() + runtime_id = str(resolved_image_id or "").strip() + framework_name = str(framework or "").strip().lower() + if input_ref == runtime_ref: + derivation = _direct_derivation( + framework=framework_name, + image=input_ref, + image_id=input_id, + ) + else: + derivation = _tracelens_derivation( + framework=framework_name, + input_image=input_ref, + input_image_id=input_id, + requested_image=runtime_ref, + resolved_image_id=runtime_id, + runtime=tracelens_runtime, + ) + errors = _image_derivation_errors( + derivation, + input_image=input_ref, + input_image_id=input_id, + requested_image=runtime_ref, + resolved_image_id=runtime_id, + ) + if errors and derivation.get("verified") is True: + derivation = dict(derivation) + derivation["verified"] = False + errors = _image_derivation_errors( + derivation, + input_image=input_ref, + input_image_id=input_id, + requested_image=runtime_ref, + resolved_image_id=runtime_id, + ) + return derivation, errors + + +def _direct_derivation( + *, + framework: str, + image: str, + image_id: str, +) -> dict[str, Any]: + return _ordered_derivation( + { + "kind": "direct", + "framework": framework, + "runtime_schema": None, + "base_image": image, + "base_image_id": image_id, + "base_image_locator": image, + "derived_image": image, + "derived_image_id": image_id, + "tracelens_source_commit": None, + "tracelens_source_tree": None, + "patch_version": None, + "patch_path": None, + "patch_sha256": None, + "dependency_wheel_manifest_sha256": None, + "validator": _DIRECT_VALIDATOR, + "verified": True, + } + ) + + +def _tracelens_derivation( + *, + framework: str, + input_image: str, + input_image_id: str, + requested_image: str, + resolved_image_id: str, + runtime: Optional[Mapping[str, Any]], +) -> dict[str, Any]: + data = runtime if isinstance(runtime, Mapping) else {} + validation = data.get("public_runtime_validation") + validation = validation if isinstance(validation, Mapping) else {} + metadata = { + "kind": "tracelens-derived", + "framework": framework, + "runtime_schema": _string(data.get("runtime_schema")), + "base_image": input_image, + "base_image_id": input_image_id, + "base_image_locator": _string(data.get("base_image_locator")), + "derived_image": requested_image, + "derived_image_id": resolved_image_id, + "tracelens_source_commit": _string(data.get("tracelens_source_commit")), + "tracelens_source_tree": _string(data.get("tracelens_source_tree")), + "patch_version": _string(data.get("patch_version")), + "patch_path": _string(data.get("tracelens_patch_path")), + "patch_sha256": _string(data.get("tracelens_patch_sha256")), + "dependency_wheel_manifest_sha256": _string( + data.get("dependency_wheel_manifest_sha256") + ), + "validator": _TRACELENS_VALIDATOR, + "verified": False, + } + runtime_matches = ( + data.get("enabled") is True + and data.get("framework") == framework + and data.get("base_image") == input_image + and data.get("base_image_id") == input_image_id + and data.get("image") == requested_image + and data.get("public_runtime_image") == requested_image + and data.get("public_runtime_image_id") == resolved_image_id + and validation.get("valid") is True + and validation.get("image_id") == resolved_image_id + ) + metadata["verified"] = bool(runtime_matches) + return _ordered_derivation(metadata) + + +def _ordered_derivation(value: Mapping[str, Any]) -> dict[str, Any]: + return {key: value.get(key) for key in SERVING_IMAGE_DERIVATION_KEYS} + + +def _image_derivation_errors( + value: object, + *, + input_image: str, + input_image_id: str, + requested_image: str, + resolved_image_id: str, +) -> Tuple[str, ...]: + if not isinstance(value, Mapping): + return ("serving image derivation is missing",) + errors = [] + if tuple(value.keys()) != SERVING_IMAGE_DERIVATION_KEYS: + errors.append("serving image derivation has an invalid shape") + if not _IMAGE_ID_RE.fullmatch(input_image_id): + errors.append("input Docker image ID is missing or invalid") + if not _IMAGE_ID_RE.fullmatch(resolved_image_id): + errors.append("resolved Docker image ID is missing or invalid") + if not input_image or not requested_image: + errors.append("serving image references are missing") + if value.get("framework") not in {"vllm", "sglang", "atom"}: + errors.append("serving image derivation framework is invalid") + if ( + value.get("base_image") != input_image + or value.get("base_image_id") != input_image_id + or value.get("derived_image") != requested_image + or value.get("derived_image_id") != resolved_image_id + ): + errors.append("serving image derivation does not match its receipt") + + if value.get("kind") == "direct": + errors.extend( + _direct_derivation_errors( + value, + input_image=input_image, + input_image_id=input_image_id, + requested_image=requested_image, + resolved_image_id=resolved_image_id, + ) + ) + elif value.get("kind") == "tracelens-derived": + errors.extend(_tracelens_derivation_errors(value)) + else: + errors.append("serving image derivation kind is invalid") + if value.get("verified") is not True: + errors.append("serving image derivation is not verified") + return tuple(_bounded_errors(errors)) + + +def _direct_derivation_errors( + value: Mapping[str, Any], + *, + input_image: str, + input_image_id: str, + requested_image: str, + resolved_image_id: str, +) -> list[str]: + errors = [] + nullable = SERVING_IMAGE_DERIVATION_KEYS[2:3] + SERVING_IMAGE_DERIVATION_KEYS[8:14] + if any(value.get(key) is not None for key in nullable): + errors.append("direct image derivation carries TraceLens identity") + if ( + input_image != requested_image + or input_image_id != resolved_image_id + or value.get("base_image_locator") != input_image + or value.get("validator") != _DIRECT_VALIDATOR + ): + errors.append("direct image derivation changed the configured image") + return errors + + +def _tracelens_derivation_errors(value: Mapping[str, Any]) -> list[str]: + errors = [] + patch_path = value.get("patch_path") + patch_match = None + if isinstance(patch_path, str) and patch_path: + parsed = PurePosixPath(patch_path) + patch_match = ( + _PATCH_PATH_RE.fullmatch(patch_path) + if not parsed.is_absolute() and ".." not in parsed.parts + else None + ) + if value.get("framework") != "vllm": + errors.append("verified TraceLens derivation currently requires vLLM") + if value.get("runtime_schema") != _TRACELENS_VLLM_RUNTIME_SCHEMA: + errors.append("TraceLens runtime schema is invalid") + base_locator = value.get("base_image_locator") + base_id = value.get("base_image_id") + locator_valid = bool( + isinstance(base_locator, str) + and ( + (_IMAGE_ID_RE.fullmatch(base_locator) and base_locator == base_id) + or _REPO_DIGEST_RE.fullmatch(base_locator) + ) + ) + if not locator_valid: + errors.append("TraceLens base image locator is missing") + if not _GIT_OBJECT_RE.fullmatch(str(value.get("tracelens_source_commit") or "")): + errors.append("TraceLens source commit is invalid") + if not _GIT_OBJECT_RE.fullmatch(str(value.get("tracelens_source_tree") or "")): + errors.append("TraceLens source tree is invalid") + patch_version = str(value.get("patch_version") or "") + if not _PATCH_VERSION_RE.fullmatch(patch_version): + errors.append("TraceLens patch version is invalid") + if patch_match is None or patch_version != f"v{int(patch_match.group(1))}": + errors.append("TraceLens patch path is invalid") + if not _SHA256_RE.fullmatch(str(value.get("patch_sha256") or "")): + errors.append("TraceLens patch SHA-256 is invalid") + if not _SHA256_RE.fullmatch( + str(value.get("dependency_wheel_manifest_sha256") or "") + ): + errors.append("TraceLens wheel manifest SHA-256 is invalid") + if value.get("validator") != _TRACELENS_VALIDATOR: + errors.append("TraceLens image validator is invalid") + return errors + + def pending_serving_runtime_receipt( *, execution_mode: str, input_config_sha256: str, + framework: str, + input_image: str, + input_image_id: str, requested_image: str, resolved_image_id: str, container_name: str, docker_argv: Sequence[str], + tracelens_runtime: Optional[Mapping[str, Any]] = None, prior_errors: Iterable[str] = (), ) -> dict[str, Any]: """Build the pre-execution receipt and validate its command bindings.""" errors = list(prior_errors) config_digest = str(input_config_sha256 or "") + input_id = str(input_image_id or "") image_id = str(resolved_image_id or "") argv_digest = canonical_docker_argv_sha256(docker_argv) if docker_argv else "" + derivation, derivation_errors = image_derivation_receipt( + framework=framework, + input_image=input_image, + input_image_id=input_id, + requested_image=requested_image, + resolved_image_id=image_id, + tracelens_runtime=tracelens_runtime, + ) if not _SHA256_RE.fullmatch(config_digest): errors.append("input config SHA-256 is missing or invalid") - if not _IMAGE_ID_RE.fullmatch(image_id): - errors.append("resolved Docker image ID is missing or invalid") + errors.extend(derivation_errors) errors.extend( docker_command_binding_errors( docker_argv, @@ -137,8 +426,11 @@ def pending_serving_runtime_receipt( return _receipt( execution_mode=execution_mode, input_config_sha256=config_digest, + input_image=input_image, + input_image_id=input_id, requested_image=requested_image, resolved_image_id=image_id, + image_derivation=derivation, container_name=container_name, docker_argv_sha256=argv_digest, process_succeeded=False, @@ -150,8 +442,13 @@ def pending_serving_runtime_receipt( def unresolved_serving_runtime_receipt( *, input_config_sha256: str, + framework: str, + input_image: str, + input_image_id: str, requested_image: str, + resolved_image_id: str, container_name: str, + tracelens_runtime: Optional[Mapping[str, Any]] = None, errors: Iterable[str], ) -> dict[str, Any]: """Build a receipt when an immutable image could not be resolved.""" @@ -159,11 +456,23 @@ def unresolved_serving_runtime_receipt( combined = list(errors) if not _SHA256_RE.fullmatch(str(input_config_sha256 or "")): combined.append("input config SHA-256 is missing or invalid") + derivation, derivation_errors = image_derivation_receipt( + framework=framework, + input_image=input_image, + input_image_id=input_image_id, + requested_image=requested_image, + resolved_image_id=resolved_image_id, + tracelens_runtime=tracelens_runtime, + ) + combined.extend(derivation_errors) return _receipt( execution_mode="docker", input_config_sha256=str(input_config_sha256 or ""), + input_image=input_image, + input_image_id=input_image_id, requested_image=requested_image, - resolved_image_id="", + resolved_image_id=resolved_image_id, + image_derivation=derivation, container_name=container_name, docker_argv_sha256="", process_succeeded=False, @@ -195,6 +504,15 @@ def validate_prepared_command( expected_image_id=str(receipt.get("resolved_image_id", "")), ) ) + errors.extend( + _image_derivation_errors( + receipt.get("image_derivation"), + input_image=str(receipt.get("input_image", "")), + input_image_id=str(receipt.get("input_image_id", "")), + requested_image=str(receipt.get("requested_image", "")), + resolved_image_id=str(receipt.get("resolved_image_id", "")), + ) + ) errors.extend(str(item) for item in receipt.get("errors", [])) return tuple(_bounded_errors(errors)) @@ -210,13 +528,25 @@ def finalize_serving_runtime_receipt( errors = list(receipt.get("errors", [])) if process_error: errors.append(process_error) + errors.extend( + _image_derivation_errors( + receipt.get("image_derivation"), + input_image=str(receipt.get("input_image", "")), + input_image_id=str(receipt.get("input_image_id", "")), + requested_image=str(receipt.get("requested_image", "")), + resolved_image_id=str(receipt.get("resolved_image_id", "")), + ) + ) bounded = _bounded_errors(errors) succeeded = bool(process_succeeded) return _receipt( execution_mode=str(receipt.get("execution_mode", "docker")), input_config_sha256=str(receipt.get("input_config_sha256", "")), + input_image=str(receipt.get("input_image", "")), + input_image_id=str(receipt.get("input_image_id", "")), requested_image=str(receipt.get("requested_image", "")), resolved_image_id=str(receipt.get("resolved_image_id", "")), + image_derivation=receipt.get("image_derivation"), container_name=str(receipt.get("container_name", "")), docker_argv_sha256=str(receipt.get("docker_argv_sha256", "")), process_succeeded=succeeded, @@ -273,8 +603,11 @@ def _receipt( *, execution_mode: str, input_config_sha256: str, + input_image: str, + input_image_id: str, requested_image: str, resolved_image_id: str, + image_derivation: object, container_name: str, docker_argv_sha256: str, process_succeeded: bool, @@ -285,8 +618,13 @@ def _receipt( "schema": SERVING_RUNTIME_SCHEMA, "execution_mode": execution_mode, "input_config_sha256": input_config_sha256, + "input_image": input_image, + "input_image_id": input_image_id, "requested_image": requested_image, "resolved_image_id": resolved_image_id, + "image_derivation": _ordered_derivation( + image_derivation if isinstance(image_derivation, Mapping) else {} + ), "container_name": container_name, "docker_argv_sha256": docker_argv_sha256, "process_succeeded": process_succeeded, @@ -304,3 +642,7 @@ def _bounded_errors(errors: Iterable[str]) -> list[str]: if len(bounded) == _MAX_ERRORS: break return bounded + + +def _string(value: object) -> Optional[str]: + return value if isinstance(value, str) and value else None diff --git a/Magpie/modes/benchmark/tracelens_runtime.py b/Magpie/modes/benchmark/tracelens_runtime.py index 92cf2f6..6de1e2f 100644 --- a/Magpie/modes/benchmark/tracelens_runtime.py +++ b/Magpie/modes/benchmark/tracelens_runtime.py @@ -26,6 +26,7 @@ from .config import BenchmarkConfig from .tracelens_vllm_image import ( + VLLM_TRACELENS_SCHEMA, VllmTraceLensIdentity, build_vllm_tracelens_image, resolve_vllm_tracelens_identity, @@ -740,6 +741,7 @@ def prepare_tracelens_runtime_image( patch_version=patch_version, ) result.update(vllm_identity.metadata()) + result["runtime_schema"] = VLLM_TRACELENS_SCHEMA else: patch_version = "unknown" diff --git a/Magpie/modes/benchmark/tracelens_vllm_image.py b/Magpie/modes/benchmark/tracelens_vllm_image.py index fc3e0a3..8db3e18 100644 --- a/Magpie/modes/benchmark/tracelens_vllm_image.py +++ b/Magpie/modes/benchmark/tracelens_vllm_image.py @@ -421,7 +421,7 @@ def resolve_vllm_tracelens_identity( if isinstance(repo_digests, list) and repo_digests and isinstance(repo_digests[0], str) - else base_image + else base_id ) source_commit = _git_text(tracelens_repo, "rev-parse", "HEAD") @@ -979,6 +979,7 @@ def validate_vllm_tracelens_image( __all__ = [ "VLLM_TRACELENS_FORBIDDEN", "VLLM_TRACELENS_REQUIREMENTS", + "VLLM_TRACELENS_SCHEMA", "VllmTraceLensIdentity", "build_vllm_tracelens_image", "docker_image_id", diff --git a/docs/conceptual/benchmarking-architecture.md b/docs/conceptual/benchmarking-architecture.md index 2979674..029419b 100644 --- a/docs/conceptual/benchmarking-architecture.md +++ b/docs/conceptual/benchmarking-architecture.md @@ -32,7 +32,7 @@ Benchmark mode consists of the following Python modules. Each benchmark run proceeds through the following stages. 1. **Configuration Loading**: Parse YAML config into `BenchmarkConfig` -2. **Runtime Setup**: For `run_mode: docker`, hash the exact input YAML, resolve the requested image to an immutable Docker image ID, hash the exact container argv, and prepare a serving-runtime receipt; for `local`, use the host environment +2. **Runtime Setup**: For `run_mode: docker`, hash the exact input YAML, resolve the configured image to an immutable pre-derivation ID, bind any validated TraceLens vLLM derived image to that input, resolve the actual runtime image ID, hash the exact container argv, and prepare a serving-runtime receipt; for `local`, use the host environment 3. **Server Launch**: Start vLLM/SGLang server (in container or on host per `run_mode`) 4. **Client Execution**: Run benchmark client with profiling enabled 5. **Trace Collection**: Torch profiler traces saved to workspace diff --git a/docs/how-to/benchmarking/benchmark.md b/docs/how-to/benchmarking/benchmark.md index 04c3aac..85acb2c 100644 --- a/docs/how-to/benchmarking/benchmark.md +++ b/docs/how-to/benchmarking/benchmark.md @@ -120,14 +120,22 @@ machine-readable `params_json` for matched TraceLens `param:*` metadata. The primary summary file is **`benchmark_report.json`**, written to the run workspace directory. It aggregates throughput, latency, and optional `gap_analysis` and `tracelens_analysis` sections. For Docker runs, `serving_runtime_receipt` uses schema -`magpie.serving-runtime-receipt/v1`. It binds the SHA-256 of the exact -`--benchmark-config` bytes to the requested image, the locally resolved -immutable `sha256:...` image ID, the exact owned container name, and a SHA-256 -of the canonical Docker argv. The argv itself is not persisted, so values such -as `HF_TOKEN` are not copied into evidence. Magpie executes the resolved image -ID rather than the mutable tag. Image-inspection failure, a missing input -digest, or any command-binding mismatch fails before container launch; -`verified` becomes true only after the bound process exits successfully. +`magpie.serving-runtime-receipt/v2`. It binds the SHA-256 of the exact +`--benchmark-config` bytes to `input_image` and its immutable pre-derivation +`input_image_id`, then to the `requested_image` actually selected for execution +and its `resolved_image_id`. `image_derivation.kind=direct` requires both +references and IDs to remain equal. A TraceLens vLLM auto-patch instead records +`kind=tracelens-derived` plus the validated base/derived IDs, runtime schema, +TraceLens commit/tree, patch version/path/hash, and dependency-wheel manifest +hash. Missing or inconsistent lineage fails before container launch. + +The receipt also binds the exact owned container name and a SHA-256 of the +canonical Docker argv. The argv itself is not persisted, so values such as +`HF_TOKEN` are not copied into evidence. Magpie executes the resolved image ID +rather than the mutable tag. Image-inspection failure, a missing input digest, +an unverified derivation, or any command-binding mismatch fails before launch; +top-level `verified` becomes true only after the bound process exits +successfully. Every report declares `run_kind` and `reward_eligible`. A `run_kind: measurement` run rejects heavy profilers; diagnostic runs and all diff --git a/docs/how-to/benchmarking/profiling-options.md b/docs/how-to/benchmarking/profiling-options.md index 6b985db..f7b1e93 100644 --- a/docs/how-to/benchmarking/profiling-options.md +++ b/docs/how-to/benchmarking/profiling-options.md @@ -66,11 +66,19 @@ base image ID, TraceLens source commit and tree, patch hash, pinned wheel hashes and preserved package versions. A same-name local image with missing or stale identity, changed base ancestry, forbidden packages, import failures, or patch marker failures is rejected and rebuilt. These fields are also returned in the -benchmark runtime metadata. TraceLens' upstream wheel metadata still declares +benchmark runtime metadata and copied into the serving-runtime v2 derivation +only after the derived image ID and validation result agree. TraceLens' upstream +wheel metadata still declares features outside Magpie's CSV diagnostic path, so a whole-environment `pip check` can report intentionally omitted packages; Magpie instead validates the exact splitter/report/import path it executes. +Serving-runtime v2 currently certifies an automatically derived TraceLens image +only for the validated vLLM builder above. For SGLang, provide an already +TraceLens-ready image so the serving receipt can bind it through the direct +immutable-image path; the public SGLang auto-build is not accepted as verified +serving lineage. + If no TraceLens source path is configured, Magpie shallow clones the official TraceLens `main` branch to `$XDG_CACHE_HOME/magpie/TraceLens` (or `~/.cache/magpie/TraceLens`) and reuses diff --git a/tests/test_benchmark_support.py b/tests/test_benchmark_support.py index 09d0625..d74a4e7 100644 --- a/tests/test_benchmark_support.py +++ b/tests/test_benchmark_support.py @@ -379,6 +379,9 @@ def test_benchmark_timeout_kills_latched_protected_container( mode._serving_runtime_receipt = pending_serving_runtime_receipt( execution_mode="docker", input_config_sha256="a" * 64, + framework="vllm", + input_image="example/image:fixed", + input_image_id=image_id, requested_image="example/image:fixed", resolved_image_id=image_id, container_name="magpie-benchmark-timed-out-task", diff --git a/tests/test_serving_runtime_receipt.py b/tests/test_serving_runtime_receipt.py index ee8c483..9544825 100644 --- a/tests/test_serving_runtime_receipt.py +++ b/tests/test_serving_runtime_receipt.py @@ -34,6 +34,39 @@ def _docker_command(container_name: str, image_id: str) -> list[str]: ] +def _tracelens_runtime( + *, + base_image: str, + base_image_id: str, + derived_image: str, + derived_image_id: str, +) -> dict[str, object]: + return { + "enabled": True, + "framework": "vllm", + "runtime_schema": "magpie.tracelens-vllm-runtime/v1", + "base_image": base_image, + "base_image_id": base_image_id, + "base_image_locator": base_image_id, + "image": derived_image, + "public_runtime_image": derived_image, + "public_runtime_image_id": derived_image_id, + "tracelens_source_commit": "1" * 40, + "tracelens_source_tree": "2" * 40, + "patch_version": "v19", + "tracelens_patch_path": ( + "examples/custom_workflows/inference_analysis/vllm_patches/" + "config_vllm_v0.19.0.patch" + ), + "tracelens_patch_sha256": "3" * 64, + "dependency_wheel_manifest_sha256": "4" * 64, + "public_runtime_validation": { + "valid": True, + "image_id": derived_image_id, + }, + } + + def test_cli_hashes_the_exact_yaml_bytes_it_parses(tmp_path): config_path = tmp_path / "benchmark.yaml" raw = ( @@ -163,6 +196,9 @@ def test_pending_receipt_rejects_image_slot_mismatch(): receipt = pending_serving_runtime_receipt( execution_mode="docker", input_config_sha256="5" * 64, + framework="vllm", + input_image="example/vllm:fixed", + input_image_id=expected, requested_image="example/vllm:fixed", resolved_image_id=expected, container_name="magpie-benchmark-case", @@ -179,6 +215,126 @@ def test_pending_receipt_rejects_image_slot_mismatch(): ] +def test_tracelens_receipt_binds_input_image_to_validated_derived_runtime(): + input_image = "sha256:" + "a" * 64 + derived_image = "magpie-tracelens-vllm:v19-candidate" + derived_image_id = "sha256:" + "b" * 64 + container_name = "magpie-benchmark-tracelens" + command = _docker_command(container_name, derived_image_id) + runtime = _tracelens_runtime( + base_image=input_image, + base_image_id=input_image, + derived_image=derived_image, + derived_image_id=derived_image_id, + ) + + receipt = pending_serving_runtime_receipt( + execution_mode="docker", + input_config_sha256="5" * 64, + framework="vllm", + input_image=input_image, + input_image_id=input_image, + requested_image=derived_image, + resolved_image_id=derived_image_id, + container_name=container_name, + docker_argv=command, + tracelens_runtime=runtime, + ) + + assert receipt["errors"] == [] + assert receipt["input_image"] == input_image + assert receipt["input_image_id"] == input_image + assert receipt["requested_image"] == derived_image + assert receipt["resolved_image_id"] == derived_image_id + assert receipt["image_derivation"] == { + "kind": "tracelens-derived", + "framework": "vllm", + "runtime_schema": "magpie.tracelens-vllm-runtime/v1", + "base_image": input_image, + "base_image_id": input_image, + "base_image_locator": input_image, + "derived_image": derived_image, + "derived_image_id": derived_image_id, + "tracelens_source_commit": "1" * 40, + "tracelens_source_tree": "2" * 40, + "patch_version": "v19", + "patch_path": ( + "examples/custom_workflows/inference_analysis/vllm_patches/" + "config_vllm_v0.19.0.patch" + ), + "patch_sha256": "3" * 64, + "dependency_wheel_manifest_sha256": "4" * 64, + "validator": "vllm-tracelens-runtime-validation/v1", + "verified": True, + } + + +@pytest.mark.parametrize( + ("mutation", "expected_error"), + [ + ( + lambda runtime: runtime["public_runtime_validation"].update(valid=False), + "serving image derivation is not verified", + ), + ( + lambda runtime: runtime.update(base_image_id="sha256:" + "c" * 64), + "serving image derivation is not verified", + ), + ( + lambda runtime: runtime.update(dependency_wheel_manifest_sha256="bad"), + "TraceLens wheel manifest SHA-256 is invalid", + ), + ( + lambda runtime: runtime.update(base_image_locator="mutable:tag"), + "TraceLens base image locator is missing", + ), + ( + lambda runtime: runtime.update( + tracelens_patch_path=( + "examples/custom_workflows/inference_analysis/vllm_patches/" + "config_vllm_v0.20.0.patch" + ) + ), + "TraceLens patch path is invalid", + ), + ], +) +def test_tracelens_receipt_rejects_unverified_or_malformed_lineage( + mutation, + expected_error, +): + input_image = "sha256:" + "a" * 64 + derived_image = "magpie-tracelens-vllm:v19-candidate" + derived_image_id = "sha256:" + "b" * 64 + runtime = _tracelens_runtime( + base_image=input_image, + base_image_id=input_image, + derived_image=derived_image, + derived_image_id=derived_image_id, + ) + mutation(runtime) + + receipt = pending_serving_runtime_receipt( + execution_mode="docker", + input_config_sha256="5" * 64, + framework="vllm", + input_image=input_image, + input_image_id=input_image, + requested_image=derived_image, + resolved_image_id=derived_image_id, + container_name="magpie-benchmark-tracelens-invalid", + docker_argv=_docker_command( + "magpie-benchmark-tracelens-invalid", + derived_image_id, + ), + tracelens_runtime=runtime, + ) + + assert receipt["verified"] is False + assert receipt["image_derivation"]["verified"] is False + assert expected_error in receipt["errors"] + + def test_success_receipt_binds_command_without_persisting_token( tmp_path, monkeypatch, @@ -201,6 +357,9 @@ def test_success_receipt_binds_command_without_persisting_token( mode._serving_runtime_receipt = pending_serving_runtime_receipt( execution_mode="docker", input_config_sha256=input_digest, + framework="vllm", + input_image="example/vllm:fixed", + input_image_id=image_id, requested_image="example/vllm:fixed", resolved_image_id=image_id, container_name=container_name, @@ -222,7 +381,28 @@ def test_success_receipt_binds_command_without_persisting_token( ) assert receipt["schema"] == SERVING_RUNTIME_SCHEMA assert receipt["input_config_sha256"] == input_digest + assert receipt["input_image"] == "example/vllm:fixed" + assert receipt["input_image_id"] == image_id + assert receipt["requested_image"] == "example/vllm:fixed" assert receipt["resolved_image_id"] == image_id + assert receipt["image_derivation"] == { + "kind": "direct", + "framework": "vllm", + "runtime_schema": None, + "base_image": "example/vllm:fixed", + "base_image_id": image_id, + "base_image_locator": "example/vllm:fixed", + "derived_image": "example/vllm:fixed", + "derived_image_id": image_id, + "tracelens_source_commit": None, + "tracelens_source_tree": None, + "patch_version": None, + "patch_path": None, + "patch_sha256": None, + "dependency_wheel_manifest_sha256": None, + "validator": "docker-image-id", + "verified": True, + } assert receipt["container_name"] == container_name assert receipt["docker_argv_sha256"] == canonical_docker_argv_sha256( command @@ -239,26 +419,31 @@ def test_success_receipt_binds_command_without_persisting_token( ] == receipt +@pytest.mark.parametrize("tracelens_derived", [False, True]) def test_benchmark_run_reports_resolved_immutable_docker_runtime( tmp_path, monkeypatch, + tracelens_derived, ): source = tmp_path / "InferenceX" source.mkdir() requested = "example/vllm:fixed" image_id = "sha256:" + "a" * 64 + derived_image = "magpie-tracelens-vllm:v19-test" + derived_image_id = "sha256:" + "e" * 64 input_digest = "b" * 64 config = BenchmarkConfig( framework="vllm", model="example/model", run_mode="docker", - run_kind="measurement", + run_kind="diagnostic" if tracelens_derived else "measurement", docker_image=requested, inferencex_path=str(source), gpu_selection={"auto": False}, profiler={ "torch_profiler": {"enabled": False}, "gpu_monitor": {"enabled": False}, + "tracelens": {"enabled": tracelens_derived}, }, ) mode = BenchmarkMode( @@ -288,14 +473,44 @@ def test_benchmark_run_reports_resolved_immutable_docker_runtime( "Magpie.modes.benchmark.benchmarker.detect_gpu", lambda: ("unknown", ""), ) + if tracelens_derived: + runtime = _tracelens_runtime( + base_image=requested, + base_image_id=image_id, + derived_image=derived_image, + derived_image_id=derived_image_id, + ) + monkeypatch.setattr( + "Magpie.modes.benchmark.benchmarker.prepare_tracelens_runtime_image", + lambda **_kwargs: runtime, + ) + + class FakeTraceLensPipeline: + def __init__(self, _config): + pass + + def prepare(self, _workspace): + return {"warnings": []} + + def restore(self): + return {"warnings": []} + + monkeypatch.setattr( + "Magpie.modes.benchmark.benchmarker.TraceLensInferencePipeline", + FakeTraceLensPipeline, + ) main_commands = [] def fake_benchmark_run(command, **kwargs): if command[:3] == ["docker", "image", "inspect"]: + inspected = command[-1] + inspected_id = ( + derived_image_id if inspected == derived_image else image_id + ) return subprocess.CompletedProcess( command, 0, - image_id + "\n", + inspected_id + "\n", "", ) main_commands.append(command) @@ -323,14 +538,25 @@ def fake_benchmark_run(command, **kwargs): assert result.success is True receipt = result.serving_runtime_receipt assert receipt["verified"] is True - assert receipt["requested_image"] == requested - assert receipt["resolved_image_id"] == image_id + assert receipt["input_image"] == requested + assert receipt["input_image_id"] == image_id + assert receipt["requested_image"] == ( + derived_image if tracelens_derived else requested + ) + assert receipt["resolved_image_id"] == ( + derived_image_id if tracelens_derived else image_id + ) + assert receipt["image_derivation"]["kind"] == ( + "tracelens-derived" if tracelens_derived else "direct" + ) assert receipt["input_config_sha256"] == input_digest benchmark_command = next( command for command in main_commands if "--name" in command ) entrypoint = benchmark_command.index("--entrypoint") - assert benchmark_command[entrypoint + 2] == image_id + assert benchmark_command[entrypoint + 2] == ( + derived_image_id if tracelens_derived else image_id + ) report = json.loads( ( tmp_path @@ -363,6 +589,9 @@ def test_command_digest_mismatch_fails_before_process_launch( mode._serving_runtime_receipt = pending_serving_runtime_receipt( execution_mode="docker", input_config_sha256="9" * 64, + framework="vllm", + input_image="example/vllm:fixed", + input_image_id=image_id, requested_image="example/vllm:fixed", resolved_image_id=image_id, container_name=container_name, diff --git a/tests/test_tracelens_vllm_image.py b/tests/test_tracelens_vllm_image.py index ac22c70..cc88672 100644 --- a/tests/test_tracelens_vllm_image.py +++ b/tests/test_tracelens_vllm_image.py @@ -110,6 +110,37 @@ def test_resolve_identity_uses_image_digest_and_committed_patch(monkeypatch, tmp assert labels["io.magpie.tracelens.source-tree"] == identity.source_tree +def test_resolve_identity_uses_image_id_when_local_tag_has_no_repo_digest( + monkeypatch, + tmp_path, +): + base_id = "sha256:" + "1" * 64 + monkeypatch.setattr( + "Magpie.modes.benchmark.tracelens_vllm_image.docker_image_record", + lambda _image: {"Id": base_id, "RepoDigests": []}, + ) + monkeypatch.setattr( + "Magpie.modes.benchmark.tracelens_vllm_image._git_text", + lambda _repo, *args: "2" * 40 if args[-1] == "HEAD" else "3" * 40, + ) + monkeypatch.setattr( + "Magpie.modes.benchmark.tracelens_vllm_image._git_bytes", + lambda _repo, *_args: b"patch", + ) + + identity = resolve_vllm_tracelens_identity( + base_image="local/vllm:candidate", + vllm_version="0.19.1+rocm721", + grpcio_version="1.78.0", + tracelens_repo=tmp_path, + patch_version="v19", + ) + + assert identity.base_image == "local/vllm:candidate" + assert identity.base_image_id == base_id + assert identity.base_image_locator == base_id + + def test_build_context_is_offline_minimal_and_identity_labeled(tmp_path): identity = _identity() labels = _write_build_context( @@ -609,6 +640,7 @@ def fake_build(**kwargs): assert result["stale_image_rejected"] is True assert result["stale_image_rejection_reason"] == "identity label mismatch" assert result["built"] is True + assert result["runtime_schema"] == "magpie.tracelens-vllm-runtime/v1" assert result["base_binding"]["image_id"] == identity.base_image_id assert result["base_binding"]["build_reference_kind"] == "repository-digest" assert result["public_runtime_validation"]["valid"] is True