diff --git a/Magpie/benchmark_config.yaml.example b/Magpie/benchmark_config.yaml.example index b51f948..bcaef61 100644 --- a/Magpie/benchmark_config.yaml.example +++ b/Magpie/benchmark_config.yaml.example @@ -307,6 +307,12 @@ benchmark: # EVAL_MAX_MODEL_LEN - lm-eval local-completions context length # EVAL_MAX_GEN_TOKS - lm-eval generation budget # +# Magpie's built-in SGLang/vLLM/Atom scripts supply the upstream local-eval +# defaults EVAL_ONLY=false, IS_AGENTIC=0, and OPENAI_API_KEY=EMPTY. Explicit +# benchmark.envs values take precedence. Local runs also preserve existing +# process-environment values; Docker runs do not inherit a host API key. +# These defaults do not change AgentX recipes or custom benchmark scripts. +# # vLLM specific: # EXTRA_VLLM_ARGS - Additional vLLM server arguments # VLLM_ROCM_USE_AITER - Enable AITER for AMD ROCm (0/1) diff --git a/Magpie/config.yaml b/Magpie/config.yaml index 21f6f90..947700e 100644 --- a/Magpie/config.yaml +++ b/Magpie/config.yaml @@ -194,9 +194,10 @@ benchmark: # Path to separate image mapping configuration file image_config: "./benchmark_images.yaml" - # Path to InferenceX installation + # InferenceX repository root or its inferencex-e2e project directory # If not configured or path doesn't exist, InferenceX will be auto-cloned from: # https://github.com/SemiAnalysisAI/InferenceX.git + # New clones pin commit 408c015be4b22d14c69518643609669405507077; existing checkouts are preserved. inferencex_path: "" # Default timeout for benchmark tasks diff --git a/Magpie/main.py b/Magpie/main.py index 2ce7173..13733cf 100644 --- a/Magpie/main.py +++ b/Magpie/main.py @@ -1058,6 +1058,11 @@ def run_benchmark(args, config: Dict[str, Any]) -> int: "profiler": { "torch_profiler": { "enabled": args.torch_profiler, + "num_steps": getattr(args, "torch_profiler_steps", 20), + "num_profiles": getattr(args, "torch_profiler_count", 1), + "interval_seconds": getattr(args, "torch_profiler_interval", 200.0), + "start_seconds": getattr(args, "torch_profiler_start_seconds", 0.0), + "detailed_annotations": getattr(args, "torch_profiler_detailed_annotations", False), }, "system_profiler": { "enabled": args.system_profiler, @@ -1081,8 +1086,8 @@ def run_benchmark(args, config: Dict[str, Any]) -> int: if run_mode: benchmark_cfg["run_mode"] = run_mode - # AgentX remains an InferenceX-owned workload. This flag only asks Magpie - # to select and execute the matching InferenceX AgentX launcher/recipe. + # AgentX preserves the InferenceX workload while Magpie owns serving for + # the packaged client layout; old checkouts retain their launcher path. if getattr(args, "agentx", False): benchmark_cfg["agentx"] = { "enabled": True, @@ -1297,6 +1302,26 @@ def create_parser() -> argparse.ArgumentParser: "--torch-profiler", action="store_true", help="Enable torch profiler" ) + benchmark_parser.add_argument( + "--torch-profiler-steps", type=int, default=20, + help="Server execution steps to capture for AgentX torch profiling (default: 20)", + ) + benchmark_parser.add_argument( + "--torch-profiler-count", type=int, default=1, + help="Sequential AgentX captures in one replay (default: 1)", + ) + benchmark_parser.add_argument( + "--torch-profiler-interval", type=float, default=200.0, + help="Seconds after trace flushing before the next AgentX capture (default: 200; 0 starts immediately)", + ) + benchmark_parser.add_argument( + "--torch-profiler-start-seconds", type=float, default=0.0, + help="Delay from AgentX measurement start before the first capture (default: 0)", + ) + benchmark_parser.add_argument( + "--torch-profiler-detailed-annotations", action="store_true", + help="Require enhanced AgentX traces from a TraceLens-capable framework runtime", + ) benchmark_parser.add_argument( "--system-profiler", action="store_true", help="Enable system profiler (rocprof/ncu)" @@ -1304,7 +1329,7 @@ def create_parser() -> argparse.ArgumentParser: benchmark_parser.add_argument( "--agentx", action="store_true", - help="Run the matching InferenceX AgentX trace-replay launcher", + help="Run AgentX trace replay with the matching InferenceX recipe", ) benchmark_parser.add_argument( "--agentx-mode", @@ -1324,7 +1349,7 @@ def create_parser() -> argparse.ArgumentParser: benchmark_parser.add_argument( "--inferencex-path", type=str, default="", - help="Path to InferenceX installation (auto-cloned if not specified)" + help="InferenceX repository or inferencex-e2e directory (auto-cloned at a pinned commit)" ) benchmark_parser.add_argument( "--benchmark-script", type=str, diff --git a/Magpie/mcp/server.py b/Magpie/mcp/server.py index 124996e..5086aa8 100644 --- a/Magpie/mcp/server.py +++ b/Magpie/mcp/server.py @@ -1168,7 +1168,7 @@ async def benchmark( concurrency: int = 32, input_len: Optional[int] = None, output_len: Optional[int] = None, - torch_profiler: bool = True, + torch_profiler: Optional[bool] = None, system_profiler: bool = False, tracelens: bool = False, tracelens_export_format: str = "csv", @@ -1196,6 +1196,11 @@ async def benchmark( ray_num_nodes: int = 1, agentx: bool = False, agentx_mode: str = "canonical", + torch_profiler_steps: int = 20, + torch_profiler_count: int = 1, + torch_profiler_interval_seconds: float = 200.0, + torch_profiler_start_seconds: float = 0.0, + torch_profiler_detailed_annotations: bool = False, ) -> str: """ Run a framework-level LLM inference benchmark (vLLM, SGLang, or Atom). @@ -1203,8 +1208,8 @@ async def benchmark( Launches the inference server using InferenceX scripts, runs a benchmark client, and collects throughput/latency metrics. Ordinary serving benchmarks can optionally collect torch profiler traces, run TraceLens - analysis, and perform gap analysis. AgentX currently reports AIPerf and - server metrics but does not support Magpie profiler collection. + analysis, and perform gap analysis. AgentX can explicitly collect bounded + torch profiler traces as a diagnostic run; its metrics cannot be published. InferenceX is auto-cloned if not present. @@ -1222,7 +1227,16 @@ async def benchmark( concurrency: Request concurrency (default: 32) input_len: Fixed-sequence input length (default: 1024); ignored by AgentX output_len: Fixed-sequence output length (default: 512); ignored by AgentX - torch_profiler: Enable PyTorch profiler traces (default: True) + torch_profiler: Enable PyTorch profiler traces (default: True for ordinary + benchmarks, False for AgentX). Explicit True enables AgentX diagnostics. + torch_profiler_steps: Server execution steps to capture for AgentX (default: 20). + torch_profiler_count: Sequential profiles in one AgentX replay (default: 1). + torch_profiler_interval_seconds: Delay after trace flushing before the next + AgentX capture (default: 200). Zero starts the next capture immediately. + torch_profiler_start_seconds: Delay from AgentX measurement start before + the first capture (default: 0). + torch_profiler_detailed_annotations: Require enhanced framework annotations + and graph shape traces (default: False); requires a capable runtime. system_profiler: Enable system profiler - rocprof (AMD) or ncu (NVIDIA) (default: False) tracelens: Enable TraceLens trace analysis on host after benchmark (default: False) tracelens_export_format: TraceLens export format - "csv" or "excel" (default: "csv") @@ -1238,7 +1252,7 @@ async def benchmark( Default: ["gpu_user_annotation"] docker_image: Override automatic Docker image selection (optional) gpu_arch: Force GPU architecture, e.g. "gfx942" (auto-detected if omitted) - inferencex_path: Path to InferenceX installation (auto-cloned if empty) + inferencex_path: InferenceX repository or inferencex-e2e project directory (auto-cloned at a pinned commit if empty) hf_cache_path: HuggingFace cache directory (default: ~/.cache/huggingface) benchmark_script: Override benchmark script name (e.g., "dsr1_fp8_mi300x.sh") runner_type: Hardware runner type (e.g., "mi300x", "h100") - auto-detected if omitted @@ -1255,7 +1269,7 @@ async def benchmark( ray_multi_node: Whether the benchmark needs multiple nodes (default: False) ray_total_num_gpus: Total GPUs across nodes for multi-node (default: 8) ray_num_nodes: Number of nodes for multi-node (default: 1) - agentx: Run the matching InferenceX AgentX launcher and recipe + agentx: Run trace replay with the matching InferenceX AgentX recipe agentx_mode: "canonical" (publishable) or "fast" (validation only) Returns: @@ -1291,12 +1305,18 @@ async def benchmark( envs.update(extra_envs) profiler_cfg = { - # AgentX owns its profiling loop. The MCP's historical torch - # profiler default is true, so disable it automatically here. - "torch_profiler": {"enabled": torch_profiler and not agentx}, + "torch_profiler": { + "enabled": (not agentx) if torch_profiler is None else torch_profiler, + "num_steps": torch_profiler_steps, + "num_profiles": torch_profiler_count, + "interval_seconds": torch_profiler_interval_seconds, + "start_seconds": torch_profiler_start_seconds, + "detailed_annotations": torch_profiler_detailed_annotations, + }, "system_profiler": {"enabled": system_profiler}, "tracelens": { "enabled": tracelens, + **({"analysis_mode": "pytorch"} if agentx else {}), "export_format": tracelens_export_format, "perf_report_enabled": True, "multi_rank_report_enabled": tp > 1, diff --git a/Magpie/modes/benchmark/agentx.py b/Magpie/modes/benchmark/agentx.py index 2ab1a88..44fb465 100644 --- a/Magpie/modes/benchmark/agentx.py +++ b/Magpie/modes/benchmark/agentx.py @@ -18,7 +18,10 @@ import yaml +from .agentx_custom import resolve_custom_entry +from .agentx_recipe import CLIENT_SCRIPT, custom_server_spec, native_server_spec from .config import AgentXConfig, BenchmarkConfig +from .inferencex import resolve_inferencex_root logger = logging.getLogger(__name__) @@ -28,6 +31,10 @@ _DEFAULT_AGENTIC_DURATION_SECONDS = 3600 +class NoMatchingAgentXRecipe(ValueError): + """The registered matrix has no recipe matching the requested identity.""" + + @dataclass(frozen=True) class AgentXLaunchSpec: """One concrete single-node AgentX point resolved from InferenceX.""" @@ -35,6 +42,7 @@ class AgentXLaunchSpec: recipe: str config_file: str entry: Dict[str, Any] + server: Optional[Dict[str, Any]] = None def _recipe_fingerprint(entry: Dict[str, Any]) -> str: @@ -148,7 +156,7 @@ def _infer_recipe( if not matches: locations = ", ".join(checked) or "" - raise ValueError( + raise NoMatchingAgentXRecipe( "InferenceX has no single-node AgentX recipe matching " f"model={config.model!r}, framework={config.framework!r}, " f"precision={config.precision!r}, gpu={runner_type!r} in: " @@ -181,6 +189,24 @@ def _selector_matches(entry: Dict[str, Any], selector: Dict[str, Any]) -> bool: return True +def _requested_topology_matches(entry: Dict[str, Any], config: BenchmarkConfig) -> bool: + """A fresh native request's explicit topology narrows the matrix as well.""" + if not entry.get("srt-recipe"): + return True + names = { + "TP": "tp", + "EP": "ep", + "EP_SIZE": "ep", + "PP_SIZE": "pp", + "DCP_SIZE": "dcp-size", + "PCP_SIZE": "pcp-size", + } + for name, field in names.items(): + if name in config.envs and str(config.envs[name]) != str(entry.get(field, 1)): + return False + return True + + def _describe_point(entry: Dict[str, Any]) -> str: fields = ( "tp", @@ -409,6 +435,8 @@ def _expand_single_node_agentx_entries( component = arm.get(key, value.get(key)) if component is not None: entry[key] = component + if arm.get("srt-recipe"): + entry["srt-recipe"] = arm["srt-recipe"] entries.append(entry) return entries @@ -423,16 +451,55 @@ def resolve_agentx_recipe( if config.agentx is None: raise ValueError("AgentX configuration is missing") - inferencex_root = Path(inferencex_path).resolve() + # Resolution runs in the caller's directory, while the server and client + # run in InferenceX. Bind local paths once before reading metadata or + # constructing argv so both phases (and Docker mounts) use the same files. + model_path_keys = [key for key in config.envs if key.upper() == "MODEL_PATH"] + if len(model_path_keys) > 1: + raise ValueError("AgentX has duplicate MODEL_PATH environment settings") + if model_path_keys: + key = model_path_keys[0] + value = config.envs.pop(key) + if not isinstance(value, str) or not value.strip(): + raise ValueError("AgentX MODEL_PATH must be a nonempty local path") + config.envs["MODEL_PATH"] = str(Path(value).expanduser().resolve()) + local_model = Path(config.model).expanduser() + if local_model.is_dir(): + config.model = str(local_model.resolve()) + + inferencex_root = Path(resolve_inferencex_root(inferencex_path)) agentx = config.agentx - if not agentx.recipe: + custom_name = f"custom-{config.framework}-{runner_type}" + requested_custom = agentx.recipe == custom_name + if not agentx.recipe or requested_custom: if not runner_type: raise ValueError("runner_type is required when agentx.recipe is not set") - config_file = _infer_recipe( - config, - inferencex_root, - runner_type, - ) + try: + config_file = _infer_recipe(config, inferencex_root, runner_type) + except NoMatchingAgentXRecipe: + entry = resolve_custom_entry(config, inferencex_root, runner_type) + if entry is None: + raise + if not _selector_matches(entry, agentx.selector): + raise ValueError("Custom AgentX point does not match agentx.selector") + entry["recipe-fingerprint"] = _recipe_fingerprint(entry) + _resolve_launcher(config, inferencex_root, custom_name, entry, runner_type) + _apply_launch_entry(config, entry) + agentx.resolved = entry + return AgentXLaunchSpec( + custom_name, + ( + "magpie:custom" + if entry.get("server-launch-spec") + else "configs/agentx-launchers.json" + ), + entry, + entry.get("server-launch-spec"), + ) + if requested_custom: + raise ValueError( + "Custom AgentX identity now matches a registered recipe; select it explicitly" + ) else: config_file = _find_recipe_config(inferencex_root, agentx) recipe = agentx.recipe @@ -456,6 +523,7 @@ def resolve_agentx_recipe( if isinstance(entry, dict) and entry.get("scenario-type") == "agentic-coding" and _selector_matches(entry, agentx.selector) + and _requested_topology_matches(entry, config) ] if not entries: selector = json.dumps(agentx.selector, sort_keys=True) @@ -517,6 +585,14 @@ def resolve_agentx_recipe( entry["image"] = config.docker_image entry["recipe-fingerprint"] = _recipe_fingerprint(entry) + _resolve_launcher(config, inferencex_root, recipe, entry, runner_type) + if entry.get("server-launch-spec"): + for source in (config_file, inferencex_root / "configs/runners.yaml"): + entry["server-launch-spec"]["source_files"][str(source.resolve())] = ( + hashlib.sha256(source.read_bytes()).hexdigest() + ) + entry["recipe-fingerprint"] = _recipe_fingerprint(entry) + _apply_launch_entry(config, entry) try: config_file_label = str(config_file.relative_to(inferencex_root)) @@ -526,12 +602,113 @@ def resolve_agentx_recipe( recipe=recipe, config_file=config_file_label, entry=entry, + server=entry.get("server-launch-spec"), ) agentx.resolved = entry logger.info("Resolved AgentX recipe %s: %s", recipe, _describe_point(entry)) return spec +def _resolve_launcher( + config: BenchmarkConfig, + root: Path, + recipe: str, + entry: Dict[str, Any], + runner_type: Optional[str], +) -> None: + """Use InferenceX's declared recipe/launcher mapping, never filename guesses.""" + client = root / "benchmarks" / CLIENT_SCRIPT + if entry.get("srt-recipe") or (entry.get("custom") and client.is_file()): + if not client.is_file(): + raise ValueError( + "Native AgentX requires the official srt_agentic.sh client" + ) + if config.benchmark_script not in {None, "", CLIENT_SCRIPT}: + raise ValueError( + "Magpie-owned AgentX uses benchmark_script=srt_agentic.sh; " + "legacy server launcher paths cannot be used with a native YAML recipe" + ) + server = ( + custom_server_spec(config, root, entry) + if entry.get("custom") + else native_server_spec(config, root, entry) + ) + config.benchmark_script = CLIENT_SCRIPT + entry["benchmark-script"] = CLIENT_SCRIPT + entry["launch-overrides-version"] = 1 + entry["server-launch-spec"] = server + # Native YAML source, variant, golden acceptance, and actual server + # configuration are part of the materialized Magpie recipe identity. + entry["recipe-fingerprint"] = _recipe_fingerprint(entry) + return + if config.profiler.torch_profiler.enabled: + raise ValueError( + "AgentX torch profiling requires Magpie-managed serving and the native InferenceX client layout" + ) + manifest = root / "configs" / "agentx-launchers.json" + row = None + if manifest.is_file(): + try: + data = json.loads(manifest.read_text(encoding="utf-8")) + except (OSError, ValueError) as exc: + raise ValueError(f"Invalid AgentX launcher manifest: {manifest}") from exc + if ( + not isinstance(data, dict) + or data.get("version") != 1 + or not isinstance(data.get("recipes"), dict) + ): + raise ValueError(f"Unsupported AgentX launcher manifest: {manifest}") + entries = data.get("generic", {}) if entry.get("custom") else data["recipes"] + row = entries.get(recipe) if isinstance(entries, dict) else None + if row is not None: + if not isinstance(row, dict) or row.get("framework") != config.framework: + raise ValueError( + f"AgentX launcher manifest framework mismatch for {recipe}" + ) + identities = ( + () + if entry.get("custom") + else (("model", config.model), ("precision", config.precision)) + ) + for field, expected in identities: + if row.get(field) != expected: + raise ValueError( + f"AgentX launcher manifest {field} mismatch for {recipe}" + ) + if runner_type and row.get("runner_type") != runner_type.lower(): + raise ValueError(f"AgentX launcher manifest runner mismatch for {recipe}") + script = row.get("benchmark_script") + if not isinstance(script, str) or not script.startswith("single_node/agentic/"): + raise ValueError(f"Invalid AgentX launcher path for {recipe}") + path = (root / "benchmarks" / script).resolve() + agentic_root = (root / "benchmarks" / "single_node" / "agentic").resolve() + if not path.is_relative_to(agentic_root) or not path.is_file(): + raise ValueError( + "AgentX declared launcher is missing or escapes its directory: " + f"{script}" + ) + if config.benchmark_script and config.benchmark_script != script: + raise ValueError( + f"AgentX recipe {recipe} declares launcher {script}, " + f"not {config.benchmark_script}" + ) + config.benchmark_script = script + entry["benchmark-script"] = script + entry["launch-overrides-version"] = row.get("launch_overrides_version") + if not config.benchmark_script: + raise ValueError( + f"InferenceX has no declared AgentX launcher for {recipe}; " + "use a checkout with configs/agentx-launchers.json or pin benchmark_script" + ) + assert config.agentx is not None + if config.agentx.launch_overrides is not None and ( + row is None or row.get("launch_overrides_version") != 1 + ): + raise ValueError( + f"AgentX launcher for {recipe} does not declare launch_overrides v1 support" + ) + + def _apply_launch_entry(config: BenchmarkConfig, entry: Dict[str, Any]) -> None: assert config.agentx is not None config.framework = str(entry["framework"]).lower() @@ -585,6 +762,22 @@ def _apply_launch_entry(config: BenchmarkConfig, entry: Dict[str, Any]) -> None: "AIPERF_FAILED_REQUEST_THRESHOLD": (config.agentx.failed_request_threshold), } ) + if entry.get("custom"): + runtime_env.update( + { + "AGENTX_CUSTOM_RECIPE": "1", + "AGENTX_NATIVE_CONTEXT_LENGTH": entry["native-context-length"], + "AGENTX_MAX_MODEL_LEN": entry["max-model-len"], + "AGENTX_MODEL_CONFIG_SHA256": entry["model-config-sha256"], + "MAX_MODEL_LEN": entry["max-model-len"], + "WEKA_LOADER_OVERRIDE": entry["trace-loader"], + } + ) + if entry.get("model-revision"): + runtime_env["AGENTX_MODEL_REVISION"] = entry["model-revision"] + server = entry.get("server-launch-spec") + if server: + runtime_env.update(server["client_env"]) config.envs = runtime_env # Loading a large model plus the canonical one-hour profile can exceed the @@ -596,10 +789,11 @@ def _apply_launch_entry(config: BenchmarkConfig, entry: Dict[str, Any]) -> None: def ensure_agentx_dependencies(inferencex_path: str) -> None: """Initialize the pinned AIPerf submodule required by AgentX.""" - root = Path(inferencex_path).resolve() + root = Path(resolve_inferencex_root(inferencex_path)) requirements = root / "utils" / "agentic-benchmark" / "requirements.txt" aiperf_project = root / "utils" / "aiperf" / "pyproject.toml" - if not requirements.is_file(): + native_client = root / "benchmarks" / CLIENT_SCRIPT + if not requirements.is_file() and not native_client.is_file(): raise RuntimeError( "InferenceX checkout does not contain AgentX support: " f"missing {requirements}" @@ -607,16 +801,19 @@ def ensure_agentx_dependencies(inferencex_path: str) -> None: if aiperf_project.is_file(): return + repo_root = root.parent if root.name == "inferencex-e2e" else root + submodule = str((root / "utils" / "aiperf").relative_to(repo_root)) completed = subprocess.run( [ "git", "-C", - str(root), + str(repo_root), "submodule", "update", "--init", "--recursive", - "utils/aiperf", + "--", + submodule, ], capture_output=True, text=True, diff --git a/Magpie/modes/benchmark/agentx_custom.py b/Magpie/modes/benchmark/agentx_custom.py new file mode 100644 index 0000000..8438ffd --- /dev/null +++ b/Magpie/modes/benchmark/agentx_custom.py @@ -0,0 +1,244 @@ +"""Resolve explicit custom-model AgentX points without modifying InferenceX.""" + +from __future__ import annotations + +import hashlib +import json +import re +from pathlib import Path +from typing import TYPE_CHECKING, Any +from urllib.parse import quote +from urllib.request import Request, urlopen + +if TYPE_CHECKING: + from .config import BenchmarkConfig + + +def native_context_length(config: dict[str, Any]) -> int: + text = config.get("text_config", config.get("language_config", config)) + if not isinstance(text, dict): + raise ValueError("Model text configuration must be an object") # noqa: TRY004 + values = [ + text.get(name) + for name in ( + "max_position_embeddings", + "max_sequence_length", + "seq_length", + "n_positions", + ) + ] + confirmed = [value for value in values if type(value) is int and value > 0] + if not confirmed: + raise ValueError( + "config.json does not declare a positive native context length" + ) + return min(confirmed) + + +def _model_metadata( + config: BenchmarkConfig, *, pin_remote: bool = False +) -> tuple[dict[str, Any], str, str, str | None]: + envs = {key.upper(): value for key, value in config.envs.items()} + local = str(envs.get("MODEL_PATH") or "") + revision = None + if local: + path = Path(local).expanduser() / "config.json" + try: + contents = path.read_bytes() + except OSError as exc: + raise ValueError( + f"Custom AgentX requires readable local model metadata: {path}" + ) from exc + source = str(path.resolve()) + elif Path(config.model).is_dir(): + path = Path(config.model) / "config.json" + contents = path.read_bytes() + source = str(path.resolve()) + else: + if not re.fullmatch(r"[A-Za-z0-9_.-]+/[A-Za-z0-9_.-]+", config.model): + raise ValueError( + "Custom AgentX model must be a local directory or HuggingFace model ID" + ) + saved = (config.agentx.resolved if config.agentx else None) or {} + revision = ( + saved.get("model-revision") + if saved.get("custom") + else envs.get("AGENTX_MODEL_REVISION") + ) + if revision is not None and not re.fullmatch(r"[0-9a-f]{40}", str(revision)): + raise ValueError("Custom AgentX model revision must be an immutable commit") + source = f"https://huggingface.co/{quote(config.model, safe='/')}/resolve/{revision or 'main'}/config.json" + headers = {"Accept": "application/json"} + try: + with urlopen(Request(source, headers=headers), timeout=30) as response: + contents = response.read(8 * 1024 * 1024 + 1) + observed = getattr(response, "headers", {}).get("X-Repo-Commit") + if observed: + if revision and revision != observed: + raise ValueError("Custom model metadata revision changed") + revision = observed + except OSError as exc: + raise ValueError( + f"Cannot read public model config metadata for {config.model}; " + "gated/private models require a local MODEL_PATH" + ) from exc + if len(contents) > 8 * 1024 * 1024: + raise ValueError("Custom AgentX model config metadata exceeds 8 MiB") + if pin_remote and not re.fullmatch(r"[0-9a-f]{40}", str(revision or "")): + raise ValueError( + "Custom AgentX remote config must identify its immutable HF revision" + ) + if revision: + source = f"https://huggingface.co/{quote(config.model, safe='/')}/resolve/{revision}/config.json" + data = json.loads(contents) + if not isinstance(data, dict): + raise ValueError( + "Custom AgentX config.json must contain an object" + ) # noqa: TRY004 + return data, hashlib.sha256(contents).hexdigest(), source, revision + + +def _positive_int(value: Any, name: str) -> int: + if isinstance(value, bool) or not re.fullmatch(r"[1-9][0-9]*", str(value)): + raise ValueError(f"Custom AgentX {name} must be an explicit positive integer") + return int(value) + + +def resolve_custom_entry( + config: BenchmarkConfig, + root: Path, + runner_type: str, +) -> dict[str, Any] | None: + """Return a declared custom point, or None when no generic capability exists.""" + assert config.agentx is not None + name = f"custom-{config.framework}-{runner_type}" + manifest = root / "configs/agentx-launchers.json" + owned = (root / "benchmarks/srt_agentic.sh").is_file() + if owned: + if runner_type not in {"mi300x", "mi325x", "mi355x"}: + raise ValueError( + "Custom AgentX supports AMD MI300X, MI325X, and MI355X only" + ) + row = { + "framework": config.framework, + "runner_type": runner_type, + "launch_overrides_version": 1, + "max_gpus": 8, + } + else: + if not manifest.is_file(): + return None + data = json.loads(manifest.read_text(encoding="utf-8")) + if not isinstance(data, dict) or data.get("version") != 1: + raise ValueError("Unsupported AgentX launcher manifest") + generic = data.get("generic", {}) + if not isinstance(generic, dict): + raise ValueError( + "AgentX generic launcher manifest must be an object" + ) # noqa: TRY004 + row = generic.get(name) + if row is None: + return None + if ( + not isinstance(row, dict) + or row.get("framework") != config.framework + or row.get("runner_type") != runner_type + ): + raise ValueError(f"Generic AgentX capability identity mismatch: {name}") + if row.get("launch_overrides_version") != 1: + raise ValueError("Generic AgentX requires the verified launch extension v1") + if not config.docker_image: + raise ValueError( + "Custom AgentX requires an explicit docker_image, including local runs" + ) + if config.framework not in {"sglang", "vllm"}: + raise ValueError("Custom AgentX supports SGLang and vLLM only") + if any(character.isspace() for character in config.model): + raise ValueError( + "Custom AgentX model identifier cannot contain whitespace; use MODEL_PATH for local files" + ) + envs = {key.upper(): value for key, value in config.envs.items()} + if envs.get("SERVED_MODEL_NAME", config.model) != config.model: + raise ValueError( + "Custom AgentX SERVED_MODEL_NAME must match the model identity" + ) + tp = _positive_int(envs.get("TP"), "TP") + ep = _positive_int(envs.get("EP_SIZE", envs.get("EP")), "EP_SIZE") + if "EP" in envs and _positive_int(envs["EP"], "EP") != ep: + raise ValueError("Custom AgentX EP and EP_SIZE disagree") + conc = _positive_int(config.agentx.concurrency or envs.get("CONC"), "CONC") + if tp > int(row.get("max_gpus", 8)) or tp % ep: + raise ValueError("Custom AgentX EP must divide TP within one declared GPU node") + if config.framework == "vllm" and ep not in {1, tp}: + raise ValueError("Custom vLLM AgentX supports EP=1 or EP=TP") + for key in ("PP", "PP_SIZE", "DCP_SIZE", "PCP_SIZE", "DP_SIZE", "NNODES"): + if str(envs.get(key, 1)) != "1": + raise ValueError(f"Custom AgentX supports {key}=1 only") + for key in ("DP_ATTENTION", "DISAGG"): + if str(envs.get(key, False)).lower() not in {"false", "0", ""}: + raise ValueError(f"Custom AgentX does not support {key}") + if str(envs.get("KV_OFFLOADING", "none")) != "none" or envs.get( + "KV_OFFLOAD_BACKEND" + ): + raise ValueError("Custom AgentX currently supports KV_OFFLOADING=none only") + if str(envs.get("SPEC_DECODING", "none")) != "none": + raise ValueError("Custom AgentX does not infer speculative decoding recipes") + precision = config.precision.lower() + if precision not in { + "bf16", + "bfloat16", + "fp16", + "float16", + "fp32", + "float32", + "auto", + "fp8", + "fp4", + "mxfp4", + }: + raise ValueError(f"Unsupported custom AgentX precision: {precision}") + metadata, checksum, source, revision = _model_metadata(config, pin_remote=owned) + if precision in {"fp8", "fp4", "mxfp4"} and not metadata.get("quantization_config"): + raise ValueError( + "Quantized custom AgentX requires quantization_config in model config.json" + ) + native = native_context_length(metadata) + maximum = _positive_int(envs.get("MAX_MODEL_LEN", native), "MAX_MODEL_LEN") + if maximum > native: + raise ValueError( + f"Custom AgentX MAX_MODEL_LEN={maximum} exceeds confirmed native context {native}" + ) + config.agentx.recipe = name + config.agentx.concurrency = conc + result = { + "custom": True, + "model": config.model, + "framework": config.framework, + "precision": precision, + "image": config.docker_image, + "model-prefix": "custom-" + + re.sub(r"[^A-Za-z0-9_.-]+", "-", config.model.rsplit("/", 1)[-1]), + "runner": runner_type, + "tp": tp, + "pp": 1, + "ep": ep, + "dcp-size": 1, + "pcp-size": 1, + "dp-attn": False, + "conc": conc, + "spec-decoding": "none", + "kv-offloading": "none", + "total-cpu-dram-gb": 0, + "duration": 3600, + "run-eval": False, + "scenario-type": "agentic-coding", + "exp-name": f"{name}_tp{tp}_conc{conc}", + "native-context-length": native, + "max-model-len": maximum, + "model-config-sha256": checksum, + "model-config-source": source, + "trace-loader": "semianalysis_cc_traces_weka_062126_256k", + } + if revision: + result["model-revision"] = revision + return result diff --git a/Magpie/modes/benchmark/agentx_launch.py b/Magpie/modes/benchmark/agentx_launch.py new file mode 100644 index 0000000..37b47f4 --- /dev/null +++ b/Magpie/modes/benchmark/agentx_launch.py @@ -0,0 +1,601 @@ +"""Versioned AgentX launch requests and independently persisted launch evidence.""" + +from __future__ import annotations + +import hashlib +import json +import os +import re +import shutil +from pathlib import Path +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from .config import BenchmarkConfig + +PROTECTED_ARGS = frozenset( + { + "--model", + "--model-path", + "--profiler-config", + "--enable-profile-cuda-graph", + "--enable-shape-discovery-for-cuda-graph-profile", + "--revision", + "--code-revision", + "--tokenizer-revision", + "--max-model-len", + "--context-length", + "--served-model-name", + "--host", + "--port", + "--tp", + "--tp-size", + "--tensor-parallel-size", + "--pp", + "--pp-size", + "--pipeline-parallel-size", + "--ep", + "--ep-size", + "--expert-parallel-size", + "--dp", + "--dp-size", + "--data-parallel-size", + "--data-parallel-rank", + "--data-parallel-start-rank", + "--data-parallel-size-local", + "--data-parallel-address", + "--data-parallel-rpc-port", + "--enable-dp-attention", + "--enable-expert-parallel", + "--decode-context-parallel-size", + "--prefill-context-parallel-size", + "--dcp-size", + "--pcp-size", + "--nnodes", + "--node-rank", + "--dist-init-addr", + } +) +PROTECTED_ENV = frozenset( + { + "PROFILE", + "SGLANG_TORCH_PROFILER_DIR", + "SGLANG_PROFILE_WITH_STACK", + "SGLANG_PROFILE_RECORD_SHAPES", + "SGLANG_GRAPH_BATCH_CAPTURE", + "VLLM_TORCH_PROFILER_DIR", + "MODEL", + "MODEL_PATH", + "MAX_MODEL_LEN", + "MODEL_NAME", + "SERVED_MODEL_NAME", + "MODEL_PREFIX", + "TP", + "PP_SIZE", + "EP_SIZE", + "DP_ATTENTION", + "DCP_SIZE", + "PCP_SIZE", + "CONC", + "DURATION", + "PORT", + "RUN_EVAL", + "EVAL_ONLY", + "IS_AGENTIC", + "SCENARIO_TYPE", + "SCENARIO_SUBDIR", + "RECIPE_FINGERPRINT", + "ROCR_VISIBLE_DEVICES", + "HIP_VISIBLE_DEVICES", + "CUDA_VISIBLE_DEVICES", + } +) +PROTECTED_ENV_PREFIXES = ("AGENTX_", "AGENTIC_", "AIPERF_", "SGLANG_SIMULATE_ACC_") +FIELDS = { + "version", + "append_args", + "remove_args", + "replace_args", + "env", + "unset_env", + "source_files", + "absent_source_files", + "executable", +} + + +def digest(value: Any) -> str: + return hashlib.sha256( + json.dumps( + value, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=True, + ).encode("utf-8") + ).hexdigest() + + +def validate_overrides(value: Any) -> dict[str, Any]: + if ( + not isinstance(value, dict) + or type(value.get("version")) is not int + or value["version"] != 1 + ): + raise ValueError("AgentX launch_overrides requires version: 1") + if unknown := value.keys() - FIELDS: + raise ValueError(f"Unknown AgentX launch override fields: {sorted(unknown)}") + result = { + "version": 1, + "append_args": [], + "remove_args": [], + "replace_args": False, + "env": {}, + "unset_env": [], + "source_files": {}, + "absent_source_files": [], + "executable": None, + **value, + } + for key in ("append_args", "remove_args", "unset_env", "absent_source_files"): + if not isinstance(result[key], list) or any( + not isinstance(item, str) or "\0" in item for item in result[key] + ): + raise ValueError(f"{key} must be a list of strings without NUL bytes") + if type(result["replace_args"]) is not bool: + raise ValueError("replace_args must be a boolean") + for key in ("env", "source_files"): + if not isinstance(result[key], dict) or any( + not isinstance(name, str) + or not isinstance(item, str) + or "\0" in name + or "\0" in item + for name, item in result[key].items() + ): + raise ValueError(f"{key} must map strings to strings without NUL bytes") + if len(set(result["remove_args"])) != len(result["remove_args"]): + raise ValueError("remove_args must not contain duplicates") + for name in result["remove_args"]: + if not re.fullmatch(r"--[A-Za-z0-9][A-Za-z0-9_-]*", name): + raise ValueError(f"remove_args must contain long option names: {name!r}") + if name.replace("_", "-") in PROTECTED_ARGS: + raise ValueError(f"Cannot remove protocol option {name}") + names = set(result["env"]) | set(result["unset_env"]) + if set(result["env"]) & set(result["unset_env"]): + raise ValueError("An environment variable cannot be both set and unset") + for name in names: + if not re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", name): + raise ValueError(f"Invalid environment variable name: {name!r}") + if name in PROTECTED_ENV or name.startswith(PROTECTED_ENV_PREFIXES): + raise ValueError(f"Cannot override protocol environment variable {name}") + for name, checksum in result["source_files"].items(): + if not Path(name).is_absolute() or not re.fullmatch(r"[0-9a-f]{64}", checksum): + raise ValueError( + "source_files requires absolute paths and lowercase SHA256 values" + ) + for name in result["absent_source_files"]: + if not Path(name).is_absolute() or name in result["source_files"]: + raise ValueError("absent_source_files requires distinct absolute paths") + executable = result["executable"] + if executable is not None and ( + not isinstance(executable, str) + or "\0" in executable + or not Path(executable).is_absolute() + ): + raise ValueError("executable must be an absolute path or null") + return result + + +def option_groups(argv: list[str]) -> tuple[list[str], list[list[str]]]: + """Separate fixed executable/positionals and unambiguous long options.""" + prefix: list[str] = [] + groups: list[list[str]] = [] + for token in argv: + if token.startswith("--"): + name = token.partition("=")[0] + if not re.fullmatch(r"--[A-Za-z0-9][A-Za-z0-9_-]*", name): + raise ValueError(f"Ambiguous server option: {token!r}") + groups.append([token]) + elif groups: + # Short options are ambiguous with values and cannot safely be + # removed as part of the preceding long option. Negative numbers + # are ordinary values. Equals-style options already own a value. + if "=" in groups[-1][0] or ( + token.startswith("-") + and not re.fullmatch(r"-\d+(?:\.\d+)?(?:[eE][+-]?\d+)?", token) + ): + raise ValueError(f"Ambiguous server option value: {token!r}") + groups[-1].append(token) + else: + prefix.append(token) + return prefix, groups + + +def apply_launch_args( + argv: list[str], + overrides: dict[str, Any], + framework: str, +) -> list[str]: + """Apply the v1 argv contract without reading files or the process environment.""" + overrides = validate_overrides(overrides) + if framework not in {"sglang", "vllm"} or not argv: + raise ValueError( + "AgentX launch extension supports nonempty SGLang/vLLM commands" + ) + effective = list(argv) + # Empty requests do not parse/rewrite an existing command. This preserves + # every canonical launcher token, including options unfamiliar to v1. + if ( + overrides["append_args"] + or overrides["remove_args"] + or overrides["replace_args"] + ): + prefix, groups = option_groups(argv) + extra_prefix, extra = option_groups(overrides["append_args"]) + if extra_prefix: + raise ValueError( + "append_args cannot inject an executable or positional arguments" + ) + if any( + group[0].partition("=")[0].replace("_", "-") in PROTECTED_ARGS + for group in extra + ): + raise ValueError( + "append_args cannot change model, topology, host or port; profiling is also protected" + ) + for name in overrides["remove_args"]: + matches = [group for group in groups if group[0].partition("=")[0] == name] + if len(matches) != 1: + raise ValueError( + f"remove_args must match exactly one existing option: {name}" + ) + groups.remove(matches[0]) + if overrides["replace_args"]: + # Keep the executable/model positionals and all protocol options; + # only the optional serving configuration is replaced. + groups = [ + group + for group in groups + if group[0].partition("=")[0].replace("_", "-") in PROTECTED_ARGS + ] + effective = prefix + [token for group in groups + extra for token in group] + if overrides["executable"] is not None: + effective[0] = overrides["executable"] + return effective + + +def _validate_golden_launch(spec: dict[str, Any], effective: list[str]) -> None: + """A candidate cannot change the acceptance curve without re-resolution.""" + if not spec.get("golden_acceptance_overrides"): + return + + def values(argv: list[str]) -> dict[str, list[list[str]]]: + result: dict[str, list[list[str]]] = {} + for group in option_groups(argv)[1]: + key, equal, inline = group[0].partition("=") + result.setdefault(key, []).append(([inline] if equal else []) + group[1:]) + return result + + base, changed = values(spec["argv"]), values(effective) + if spec["framework"] == "sglang": + keys = { + "--speculative-algorithm", + "--speculative-algo", + "--speculative-num-steps", + "--speculative-dspark-block-size", + "--speculative-draft-model-path", + } + if any(base.get(key) != changed.get(key) for key in keys): + raise ValueError( + "Changing AgentX golden acceptance inputs requires a new resolved recipe" + ) + elif spec["framework"] == "vllm": + + def parameters(options): + groups = options.get("--speculative-config", []) + if not groups: + return {} + if len(groups) != 1 or len(groups[0]) != 1: + raise ValueError( + "AgentX requires exactly one speculative-config JSON object" + ) + value = json.loads(groups[0][0]) + if not isinstance(value, dict): + raise ValueError("AgentX speculative-config must be a JSON object") + return value + + old, new = parameters(base), parameters(changed) + keys = { + "method", + "model", + "num_speculative_tokens", + "rejection_sample_method", + "synthetic_acceptance_length", + } + if any(old.get(key) != new.get(key) for key in keys): + raise ValueError("AgentX golden acceptance parameters cannot be overridden") + + +def launch_environment( + config: BenchmarkConfig, + workspace: Path, + *, + docker: bool = False, +) -> dict[str, str]: + """Persist the exact normalized request and bind this run's output path.""" + result = { + "AGENTX_LAUNCH_OVERRIDES_FILE": "", + "AGENTX_LAUNCH_OVERRIDES_SHA256": "", + "AGENTX_SERVER_LAUNCH_FILE": "", + } + if config.agentx is None or config.agentx.launch_overrides is None: + return result + overrides = validate_overrides(config.agentx.launch_overrides) + if overrides != config.agentx.launch_overrides: + raise ValueError("AgentX launch overrides changed after config normalization") + workspace.mkdir(parents=True, exist_ok=True) + request = workspace / "agentx_launch_overrides.json" + request.write_text(json.dumps(overrides, indent=2) + "\n", encoding="utf-8") + # A stale receipt must never make a failed new launch appear verified. + (workspace / "agentx_server_launch.json").unlink(missing_ok=True) + runtime_workspace = Path("/workspace") if docker else workspace.resolve() + result.update( + { + "AGENTX_LAUNCH_OVERRIDES_FILE": str(runtime_workspace / request.name), + "AGENTX_LAUNCH_OVERRIDES_SHA256": digest(overrides), + "AGENTX_SERVER_LAUNCH_FILE": str( + runtime_workspace / "agentx_server_launch.json" + ), + } + ) + return result + + +def read_launch_evidence(config: BenchmarkConfig, workspace: Path) -> dict[str, Any]: + """Verify the receipt against the immutable request persisted for this run.""" + assert config.agentx is not None + overrides = validate_overrides(config.agentx.launch_overrides or {"version": 1}) + request = json.loads( + (workspace / "agentx_launch_overrides.json").read_text(encoding="utf-8") + ) + if request != overrides: + raise ValueError( + "Persisted AgentX launch request differs from the benchmark config" + ) + evidence = json.loads( + (workspace / "agentx_server_launch.json").read_text(encoding="utf-8") + ) + if not isinstance(evidence, dict): + raise ValueError( + "AgentX server launch evidence must be an object" + ) # noqa: TRY004 + payload = { + key: value for key, value in evidence.items() if key != "evidence_sha256" + } + if evidence.get("evidence_sha256") != digest(payload): + raise ValueError("AgentX server launch evidence hash mismatch") + if ( + type(evidence.get("version")) is not int + or evidence.get("version") != 1 + or evidence.get("framework") != config.framework + ): + raise ValueError("AgentX server launch evidence identity mismatch") + if evidence.get("overrides_sha256") != digest(overrides): + raise ValueError("AgentX server did not apply this run's launch overrides") + if evidence.get("source_files") != overrides["source_files"]: + raise ValueError("AgentX source file launch evidence mismatch") + if evidence.get("absent_source_files") != overrides["absent_source_files"]: + raise ValueError("AgentX deleted source file launch evidence mismatch") + names = sorted(set(overrides["env"]) | set(overrides["unset_env"])) + expected_env = {name: overrides["env"].get(name) for name in names} + if evidence.get("effective_env") != expected_env: + raise ValueError("AgentX effective server environment mismatch") + if not isinstance(evidence.get("base_env"), dict) or set( + evidence["base_env"] + ) != set(names): + raise ValueError("AgentX base server environment evidence is incomplete") + for key in ("base_argv", "effective_argv"): + argv = evidence.get(key) + if ( + not isinstance(argv, list) + or not argv + or any(not isinstance(value, str) for value in argv) + ): + raise ValueError(f"AgentX {key} evidence must contain server argv tokens") + if ( + overrides["executable"] + and evidence["effective_argv"][0] != overrides["executable"] + ): + raise ValueError("AgentX server executable override was not applied") + executable = evidence.get("resolved_executable") + if not isinstance(executable, str) or not Path(executable).is_absolute(): + raise ValueError("AgentX resolved executable evidence is missing") + runtime = evidence.get("runtime_environment") + if not isinstance(runtime, dict) or any( + not isinstance(name, str) or not isinstance(value, str) + for name, value in runtime.items() + ): + raise ValueError("AgentX runtime environment evidence is incomplete") + spec = (config.agentx.resolved or {}).get("server-launch-spec") + if spec is not None: + from .agentx_profile_config import profile_server_spec, profile_settings + + settings = profile_settings(config) + profiling = evidence.get("torch_profiler") + if settings is not None: + if not isinstance(profiling, dict): + raise ValueError("AgentX profiler launch evidence is missing") + spec = profile_server_spec( + spec, + settings, + workspace, + profiling.get("capture_id"), + capabilities=profiling.get("capabilities"), + launch_overrides=config.agentx.launch_overrides, + ) + if settings.get("detailed_annotations") and not profiling.get( + "capabilities" + ): + raise ValueError( + "AgentX detailed annotation capability evidence is missing" + ) + if profiling != spec["torch_profiler"]: + raise ValueError("AgentX profiler configuration evidence mismatch") + if ( + config.framework == "sglang" + and runtime.get("SGLANG_TORCH_PROFILER_DIR") + != spec["env"]["SGLANG_TORCH_PROFILER_DIR"] + ): + raise ValueError("AgentX profiler output environment mismatch") + if config.framework == "sglang" and settings.get("detailed_annotations"): + for name in ( + "SGLANG_PROFILE_WITH_STACK", + "SGLANG_PROFILE_RECORD_SHAPES", + "SGLANG_GRAPH_BATCH_CAPTURE", + ): + if runtime.get(name) != spec["env"][name]: + raise ValueError( + f"AgentX annotation environment mismatch: {name}" + ) + elif profiling is not None: + raise ValueError("Unexpected AgentX profiler launch evidence") + if evidence.get("owner") != "magpie": + raise ValueError("AgentX server launch evidence is not Magpie-owned") + if evidence.get("server_spec_sha256") != digest(spec): + raise ValueError("AgentX server launch specification identity mismatch") + if evidence.get("recipe_source_files") != spec.get("source_files", {}): + raise ValueError("AgentX recipe source evidence is incomplete") + if evidence["effective_argv"] != apply_launch_args( + evidence["base_argv"], overrides, config.framework + ): + raise ValueError("AgentX effective argv does not match this candidate") + _validate_golden_launch(spec, evidence["effective_argv"]) + if evidence["base_argv"] != spec.get("argv"): + raise ValueError("AgentX base argv does not match the resolved recipe") + return evidence + + +def prepare_server_launch( + argv: list[str], + environment: dict[str, str], + overrides: dict[str, Any] | None, + framework: str, + workspace: Path, + *, + server_spec: dict[str, Any] | None = None, +) -> tuple[list[str], dict[str, str], dict[str, Any]]: + """Attest a Magpie-owned server in its actual runtime, after setup. + + The caller immediately spawns the returned argv/environment. This function + must run inside the serving container for Docker jobs. The client never + invokes it, so client setup cannot overwrite the server's receipt. + """ + workspace.mkdir(parents=True, exist_ok=True) + receipt = workspace / "agentx_server_launch.json" + receipt.unlink(missing_ok=True) + request = validate_overrides(overrides or {"version": 1}) + if not isinstance(argv, list) or any( + not isinstance(token, str) or "\0" in token for token in argv + ): + raise ValueError("AgentX server argv must contain literal string tokens") + if any( + not isinstance(name, str) + or not isinstance(value, str) + or "\0" in name + or "\0" in value + for name, value in environment.items() + ): + raise ValueError("AgentX server environment must contain strings without NUL") + if server_spec is not None and server_spec.get("argv") != argv: + raise ValueError("AgentX server argv differs from its resolved specification") + effective = apply_launch_args(argv, request, framework) + if server_spec is not None: + _validate_golden_launch(server_spec, effective) + child_env = dict(environment) + touched = sorted(set(request["env"]) | set(request["unset_env"])) + base_env = {name: environment.get(name) for name in touched} + for name in request["unset_env"]: + child_env.pop(name, None) + child_env.update(request["env"]) + + sources = {} + for name, expected in request["source_files"].items(): + actual = hashlib.sha256(Path(name).read_bytes()).hexdigest() + if actual != expected: + raise ValueError(f"AgentX source changed before server launch: {name}") + sources[name] = actual + recipe_sources = {} + for name, expected in (server_spec or {}).get("source_files", {}).items(): + if not Path(name).is_absolute() or not re.fullmatch(r"[0-9a-f]{64}", expected): + raise ValueError("AgentX recipe sources require absolute paths and SHA256") + actual = hashlib.sha256(Path(name).read_bytes()).hexdigest() + if actual != expected: + raise ValueError( + f"AgentX recipe input changed before server launch: {name}" + ) + recipe_sources[name] = actual + for name in request["absent_source_files"]: + if os.path.lexists(name): + raise ValueError(f"AgentX deleted source reappeared before launch: {name}") + resolved_executable = shutil.which(effective[0], path=child_env.get("PATH", "")) + if resolved_executable is None: + raise ValueError(f"AgentX server executable not found: {effective[0]}") + evidence: dict[str, Any] = { + "version": 1, + "owner": "magpie", + "framework": framework, + "base_argv": list(argv), + "effective_argv": effective, + "base_env": base_env, + "effective_env": {name: child_env.get(name) for name in touched}, + "overrides_sha256": digest(request), + "source_files": sources, + "absent_source_files": request["absent_source_files"], + "resolved_executable": str(Path(resolved_executable).resolve()), + "runtime_environment": { + name: value + for name, value in sorted(child_env.items()) + if ( + name + in { + "PATH", + "PYTHONPATH", + "LD_LIBRARY_PATH", + "LIBRARY_PATH", + "GPU_ARCHS", + } + or name.startswith( + ( + "SGLANG_", + "VLLM_", + "AITER_", + "TRITON_", + "TORCH_", + "PYTORCH_", + "HIP_", + "ROCR_", + "HSA_", + "NCCL_", + "RCCL_", + "OMP_", + ) + ) + ) + and not any( + word in name.upper() + for word in ("TOKEN", "PASSWORD", "SECRET", "API_KEY") + ) + }, + } + if server_spec is not None: + if server_spec.get("torch_profiler"): + evidence["torch_profiler"] = server_spec["torch_profiler"] + evidence["server_spec_sha256"] = digest(server_spec) + evidence["recipe_source_files"] = recipe_sources + evidence["evidence_sha256"] = digest(evidence) + (workspace / "agentx_launch_overrides.json").write_text( + json.dumps(request, indent=2) + "\n", encoding="utf-8" + ) + temporary = receipt.with_name(receipt.name + ".tmp") + temporary.write_text(json.dumps(evidence, indent=2) + "\n", encoding="utf-8") + temporary.replace(receipt) + return effective, child_env, evidence diff --git a/Magpie/modes/benchmark/agentx_profile_capabilities.py b/Magpie/modes/benchmark/agentx_profile_capabilities.py new file mode 100644 index 0000000..e004f12 --- /dev/null +++ b/Magpie/modes/benchmark/agentx_profile_capabilities.py @@ -0,0 +1,280 @@ +############################################################################### +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. +# See LICENSE for license information. +############################################################################### +"""Inspect profiling support with the interpreter that will launch the server. + +This module and its child probe use only the standard library. Importing the +framework's configuration classes does not instantiate a model or a GPU worker. +""" + +from __future__ import annotations + +import json +import math +import os +from pathlib import Path +import re +import shlex +import shutil +import subprocess +from typing import Any + +_RESULT_PREFIX = "MAGPIE_PROFILE_CAPABILITIES=" +_PYTHON_NAME = re.compile(r"(?:python|pypy)(?:\d+(?:\.\d+)*)?(?:\.exe)?$") +_PROBE_SCRIPT = r""" +import dataclasses +import importlib +import json +import sys + +framework, script_directory = sys.argv[1:3] +if script_directory and not (getattr(sys.flags, "safe_path", False) or sys.flags.isolated): + sys.path[0] = script_directory + +def fields(cls): + names = set() + for base in cls.__mro__: + names.update(getattr(base, "__annotations__", {})) + if dataclasses.is_dataclass(cls): + names.update(field.name for field in dataclasses.fields(cls)) + for attribute in ("model_fields", "__fields__", "__struct_fields__"): + names.update(getattr(cls, attribute, {}) or {}) + return names + +if framework == "sglang": + server = importlib.import_module("sglang.srt.server_args") + requests = importlib.import_module("sglang.srt.managers.io_struct") + request_types = [getattr(requests, name) for name in + ("ProfileReqInput", "ProfileReq") if hasattr(requests, name)] + if not request_types: + raise RuntimeError("SGLang lacks ProfileReqInput and ProfileReq") + request_fields = set.intersection(*(fields(cls) for cls in request_types)) + server_fields = fields(server.ServerArgs) + value = { + "annotation_field": next((name for name in + ("roofline_annotations", "detailed_annotations") + if name in request_fields), None), + "shape_discovery": "shape_discovery" in request_fields, + "graph_capture": "enable_profile_cuda_graph" in server_fields, + "graph_shape_discovery": + "enable_shape_discovery_for_cuda_graph_profile" in server_fields, + } +else: + config = importlib.import_module("vllm.config") + cls = getattr(config, "ProfilerConfig", None) + if cls is None: + cls = importlib.import_module("vllm.config.profiler").ProfilerConfig + names = fields(cls) + value = { + "annotation_field": ("detailed_trace_annotation" + if "detailed_trace_annotation" in names else None), + "capture_field": next((name for name in + ("capture_torch_profiler_dir", "capture_torch_profiler") + if name in names), None), + } +print("MAGPIE_PROFILE_CAPABILITIES=" + json.dumps(value, sort_keys=True)) +""" + + +def validate_profile_capabilities(framework: str, value: Any) -> dict[str, Any]: + """Validate and canonicalize capabilities persisted in launch evidence. + + Graph support is reported separately because an eager candidate need not + support graph capture. The caller must require it before enabling graph flags. + """ + if framework not in {"sglang", "vllm"}: + raise ValueError("AgentX detailed profiling supports SGLang and vLLM") + if not isinstance(value, dict): + raise ValueError("AgentX profile capabilities must be an object") + expected = ( + { + "annotation_field", + "shape_discovery", + "graph_capture", + "graph_shape_discovery", + } + if framework == "sglang" + else {"annotation_field", "capture_field"} + ) + if value.keys() - expected: + raise ValueError( + f"Unknown AgentX profile capability fields: {sorted(value.keys() - expected)}" + ) + if expected - value.keys(): + raise ValueError( + f"Missing AgentX profile capability fields: {sorted(expected - value.keys())}" + ) + if framework == "sglang": + if value["annotation_field"] not in ( + "roofline_annotations", + "detailed_annotations", + ): + raise ValueError( + "SGLang lacks profile request roofline_annotations/detailed_annotations" + ) + if value["shape_discovery"] is not True: + raise ValueError("SGLang lacks profile request shape_discovery") + for name in ("graph_capture", "graph_shape_discovery"): + if type(value[name]) is not bool: + raise ValueError(f"SGLang profile capability {name} must be boolean") + return { + "annotation_field": value["annotation_field"], + "shape_discovery": True, + "graph_capture": value["graph_capture"], + "graph_shape_discovery": value["graph_shape_discovery"], + } + if value["annotation_field"] != "detailed_trace_annotation": + raise ValueError("vLLM ProfilerConfig lacks detailed_trace_annotation") + if value["capture_field"] not in ( + "capture_torch_profiler_dir", + "capture_torch_profiler", + ): + raise ValueError( + "vLLM ProfilerConfig lacks capture_torch_profiler_dir/capture_torch_profiler" + ) + return { + "annotation_field": "detailed_trace_annotation", + "capture_field": value["capture_field"], + } + + +def _resolve_executable(executable: str, env: dict[str, str], cwd: Path) -> str: + if os.sep in executable: + path = Path(executable) + path = path if path.is_absolute() else cwd / path + if not path.is_file() or not os.access(path, os.X_OK): + raise ValueError(f"Server executable is unavailable: {path}") + # Preserve virtualenv symlinks: resolving them can select another install. + return str(path.absolute()) + search_path = os.pathsep.join( + str(Path(part) if Path(part).is_absolute() else cwd / part) + for part in env.get("PATH", os.defpath).split(os.pathsep) + ) + resolved = shutil.which(executable, path=search_path) + if resolved is None: + raise ValueError(f"Server executable is unavailable on its PATH: {executable}") + return resolved + + +def _python_prefix( + argv: list[str], env: dict[str, str], cwd: Path +) -> tuple[list[str], str]: + executable = _resolve_executable(argv[0], env, cwd) + arguments = argv[1:] + script_directory = "" + if not _PYTHON_NAME.fullmatch(Path(executable).name): + with open(executable, "rb") as handle: + line = handle.readline(4096) + if not line.startswith(b"#!"): + raise ValueError(f"Server executable has no Python shebang: {executable}") + try: + shebang = shlex.split(line[2:].decode("utf-8").strip()) + except (UnicodeError, ValueError) as exc: + raise ValueError(f"Invalid server Python shebang: {executable}") from exc + if shebang and Path(shebang[0]).name == "env": + shebang = shebang[1:] + if shebang[:1] == ["-S"]: + shebang = shebang[1:] + if not shebang or not _PYTHON_NAME.fullmatch(Path(shebang[0]).name): + raise ValueError(f"Unsupported server Python shebang: {executable}") + script_directory = str(Path(executable).parent) + executable = _resolve_executable(shebang[0], env, cwd) + arguments = shebang[1:] + + prefix = [executable] + index = 0 + while index < len(arguments): + arg = arguments[index] + if arg in {"-m", "-c"}: + break + if arg == "--" or not arg.startswith("-"): + if arg == "--" and index + 1 >= len(arguments): + raise ValueError(f"Missing Python script after --: {executable}") + script = arguments[index + 1] if arg == "--" else arg + if script != "-": + path = Path(script) + script_directory = str( + (path if path.is_absolute() else cwd / path).parent + ) + break + if arg in {"-W", "-X", "--check-hash-based-pycs"}: + if index + 1 >= len(arguments): + raise ValueError(f"Incomplete Python option {arg}: {executable}") + prefix.extend(arguments[index : index + 2]) + index += 2 + continue + if re.fullmatch(r"-[bBdEHiIOPqRsSuv]+", arg) or arg.startswith(("-W", "-X")): + # -i would hang the probe after -c. Its import semantics are unchanged. + retained = arg.replace("i", "") if not arg.startswith(("-W", "-X")) else arg + if retained != "-": + prefix.append(retained) + index += 1 + continue + raise ValueError(f"Unsupported server Python option {arg!r}: {executable}") + return prefix, script_directory + + +def probe_profile_capabilities( + framework: str, + argv: list[str], + env: dict[str, str], + cwd: str | Path, + timeout: float, +) -> dict[str, Any]: + """Probe the selected server runtime; never substitute the Magpie interpreter.""" + if framework not in {"sglang", "vllm"}: + raise ValueError("AgentX detailed profiling supports SGLang and vLLM") + if ( + not isinstance(argv, list) + or not argv + or any(not isinstance(arg, str) or "\0" in arg for arg in argv) + or not argv[0] + ): + raise ValueError("AgentX profile probe requires the resolved server argv") + if type(timeout) not in (int, float) or not math.isfinite(timeout) or timeout <= 0: + raise ValueError("AgentX profile probe timeout must be positive and finite") + location = Path(cwd).absolute() + try: + prefix, script_directory = _python_prefix(argv, env, location) + except (OSError, ValueError) as exc: + raise RuntimeError( + f"Cannot select profile probe interpreter for {argv[0]!r}: {exc}" + ) from exc + label = shlex.join(prefix) + try: + result = subprocess.run( + [*prefix, "-c", _PROBE_SCRIPT, framework, script_directory], + env=env, + cwd=location, + timeout=timeout, + text=True, + capture_output=True, + check=False, + ) + except subprocess.TimeoutExpired as exc: + raise RuntimeError( + f"AgentX {framework} profile capability probe timed out using {label}" + ) from exc + except OSError as exc: + raise RuntimeError( + f"AgentX {framework} profile capability probe failed using {label}: {exc}" + ) from exc + if result.returncode: + detail = (result.stderr or result.stdout).strip()[-2000:] + raise RuntimeError( + f"AgentX {framework} profile capability probe failed using {label}: {detail}" + ) + lines = [ + line[len(_RESULT_PREFIX) :] + for line in result.stdout.splitlines() + if line.startswith(_RESULT_PREFIX) + ] + try: + if len(lines) != 1: + raise ValueError("probe did not return exactly one capability object") + return validate_profile_capabilities(framework, json.loads(lines[0])) + except (ValueError, TypeError) as exc: + raise RuntimeError( + f"AgentX {framework} detailed profiling unavailable using {label}: {exc}" + ) from exc diff --git a/Magpie/modes/benchmark/agentx_profile_config.py b/Magpie/modes/benchmark/agentx_profile_config.py new file mode 100644 index 0000000..9df3b2d --- /dev/null +++ b/Magpie/modes/benchmark/agentx_profile_config.py @@ -0,0 +1,230 @@ +############################################################################### +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. +# See LICENSE for license information. +############################################################################### +"""Derive a diagnostic launch without mutating the accepted AgentX candidate.""" + +from __future__ import annotations + +import copy +import json +import math +import re +from fractions import Fraction +from pathlib import Path +from typing import Any + + +def profile_settings(config) -> dict[str, Any] | None: + profiler = getattr(getattr(config, "profiler", None), "torch_profiler", None) + if profiler is None or not profiler.enabled: + return None + return { + "num_steps": profiler.num_steps, + "num_profiles": profiler.num_profiles, + "interval_seconds": profiler.interval_seconds, + "start_seconds": profiler.start_seconds, + "detailed_annotations": profiler.detailed_annotations, + "capture_timeout_seconds": profiler.capture_timeout_seconds, + "flush_timeout_seconds": profiler.flush_timeout_seconds, + } + + +def profile_plan( + settings: dict[str, Any], client_env: dict[str, str] +) -> dict[str, Any]: + """Bound capture starts by the replay window, without predicting step cost.""" + count = settings.get("num_profiles", 1) + if type(count) is not int or count <= 0: + raise ValueError("AgentX num_profiles must be a positive integer") + interval = settings.get("interval_seconds", 200.0) + start = settings.get("start_seconds", 0.0) + raw_duration = ( + 1200 + if client_env.get("AIPERF_EXPERIMENTAL_FAST") == "1" + else client_env.get("DURATION", 3600) + ) + try: + duration = float(raw_duration) + valid = ( + not isinstance(raw_duration, bool) + and math.isfinite(duration) + and duration > 0 + and type(interval) in (int, float) + and math.isfinite(interval) + and interval >= 0 + and type(start) in (int, float) + and math.isfinite(start) + and start >= 0 + ) + except (TypeError, ValueError, OverflowError): + valid = False + if not valid: + raise ValueError( + "AgentX profiling needs a positive finite replay duration and nonnegative finite interval/start_seconds" + ) + if start >= duration: + raise ValueError( + "AgentX profiler start_seconds must be less than the measurement duration" + ) + maximum = None + if interval > 0: + # A start at exactly the end of measurement is too late. Decimal + # rational arithmetic avoids overflow and ceil errors near multiples. + remaining = Fraction(str(duration)) - Fraction(str(start)) + ratio = remaining / Fraction(str(interval)) + maximum = -(-ratio.numerator // ratio.denominator) + return { + "requested_profiles": count, + "max_profiles": maximum, + "planned_profiles": min(count, maximum) if maximum is not None else count, + "measurement_duration_seconds": duration, + "start_seconds": float(start), + } + + +def profile_server_spec( + spec: dict[str, Any], + settings: dict[str, Any], + workspace: Path, + capture_id: str, + *, + capabilities: dict[str, Any] | None = None, + launch_overrides: dict[str, Any] | None = None, +) -> dict[str, Any]: + """Bind framework instrumentation and its output directory into launch evidence.""" + if not isinstance(capture_id, str) or not re.fullmatch(r"[0-9a-f]{32}", capture_id): + raise ValueError("AgentX profiler requires a unique capture id") + if type(settings.get("num_steps")) is not int or settings["num_steps"] <= 0: + raise ValueError("AgentX profiler num_steps must be a positive integer") + annotations = settings.get("detailed_annotations", False) + if type(annotations) is not bool: + raise ValueError("AgentX profiler detailed_annotations must be a boolean") + if capabilities is not None: + if not annotations: + raise ValueError("Unexpected AgentX annotation capabilities") + if __package__: + from .agentx_profile_capabilities import validate_profile_capabilities + else: + from agentx_profile_capabilities import validate_profile_capabilities + capabilities = validate_profile_capabilities(spec["framework"], capabilities) + if __package__: + from .agentx_launch import apply_launch_args + else: + from agentx_launch import apply_launch_args + effective_args = apply_launch_args( + spec["argv"], launch_overrides or {"version": 1}, spec["framework"] + ) + effective_flags = { + arg.split("=", 1)[0].replace("_", "-") for arg in effective_args + } + enforce_eager = False + for arg in effective_args: + name = arg.replace("_", "-") + if name in {"--enforce-eager", "--no-enforce-eager"}: + enforce_eager = name == "--enforce-eager" + count = settings.get("num_profiles", 1) + if type(count) is not int or count <= 0: + raise ValueError("AgentX profiler num_profiles must be a positive integer") + interval = settings.get("interval_seconds", 200.0) + try: + valid_interval = ( + type(interval) in (int, float) and math.isfinite(interval) and interval >= 0 + ) + except OverflowError: + valid_interval = False + if not valid_interval: + raise ValueError( + "AgentX profiler interval_seconds must be nonnegative and finite" + ) + start = settings.get("start_seconds", 0.0) + try: + valid_start = ( + type(start) in (int, float) and math.isfinite(start) and start >= 0 + ) + except OverflowError: + valid_start = False + if not valid_start: + raise ValueError("AgentX profiler start_seconds must be nonnegative and finite") + for name in ("capture_timeout_seconds", "flush_timeout_seconds"): + value = settings.get(name) + try: + valid = type(value) in (int, float) and math.isfinite(value) and value > 0 + except OverflowError: + valid = False + if not valid: + raise ValueError(f"AgentX profiler {name} must be positive and finite") + derived = copy.deepcopy(spec) + directory = workspace.resolve() / "torch_trace" / capture_id + output_directory = directory / "active" if count > 1 else directory + instrumentation = { + **settings, + "capture_id": capture_id, + "trace_dir": str(directory), + } + if capabilities is not None: + instrumentation["capabilities"] = capabilities + if any( + arg.replace("_", "-").startswith("--profiler-config") for arg in derived["argv"] + ): + raise ValueError("AgentX profiler configuration must be owned by Magpie") + if derived["framework"] == "vllm": + options = { + "profiler": "torch", + "torch_profiler_dir": str(output_directory), + "max_iterations": settings["num_steps"], + "ignore_frontend": True, + "torch_profiler_use_gzip": True, + } + if capabilities is not None: + options.update( + detailed_trace_annotation=True, + torch_profiler_record_shapes=True, + torch_profiler_with_stack=True, + ) + if not enforce_eager: + field = capabilities["capture_field"] + options[field] = ( + str(output_directory / "capture_traces") + if field == "capture_torch_profiler_dir" + else True + ) + derived["argv"] += [ + "--profiler-config", + json.dumps( + options, + sort_keys=True, + separators=(",", ":"), + ), + ] + elif derived["framework"] == "sglang": + derived.setdefault("env", {})["SGLANG_TORCH_PROFILER_DIR"] = str( + output_directory + ) + if capabilities is not None: + derived["env"].update( + SGLANG_PROFILE_WITH_STACK="True", + SGLANG_PROFILE_RECORD_SHAPES="True", + SGLANG_GRAPH_BATCH_CAPTURE="True", + ) + if "--disable-cuda-graph" not in effective_flags: + if ( + not capabilities["graph_capture"] + or not capabilities["graph_shape_discovery"] + ): + raise ValueError( + "SGLang detailed annotations with CUDA graphs require graph profiling and shape discovery capabilities" + ) + for name, flag in ( + ("graph_capture", "--enable-profile-cuda-graph"), + ( + "graph_shape_discovery", + "--enable-shape-discovery-for-cuda-graph-profile", + ), + ): + if capabilities[name] and flag not in effective_flags: + derived["argv"].append(flag) + else: + raise ValueError("AgentX torch profiling supports SGLang and vLLM") + derived["torch_profiler"] = instrumentation + return derived diff --git a/Magpie/modes/benchmark/agentx_profiling.py b/Magpie/modes/benchmark/agentx_profiling.py new file mode 100644 index 0000000..8b52184 --- /dev/null +++ b/Magpie/modes/benchmark/agentx_profiling.py @@ -0,0 +1,949 @@ +# SPDX-FileCopyrightText: 2026 Advanced Micro Devices, Inc. +# SPDX-License-Identifier: MIT + +"""Collect bounded AgentX server traces after AIPerf enters its measured phase. + +The bounded-memory trace reader is adapted from Hyperloom's MIT-licensed +``inference_optimizer/assets/agentx/aiperf_phase_gate.py``. This module has no +Magpie imports so the standalone container worker can use it directly. +""" + +from __future__ import annotations + +import gzip +import http.client +import json +import math +import re +import time +import urllib.error +import urllib.request +import zlib +from collections.abc import Callable +from pathlib import Path +from typing import Any + +_POLL_SECONDS = 0.2 +_HTTP_SECONDS = 2.0 +_STARTUP_TRACE_DIRS = {"graph_capture_profile", "capture_traces"} +_TRACE_RANK_PATTERNS = ( + re.compile(r"(?:^|[-_.])rank[-_]?(\d+)(?=[-_.]|$)", re.IGNORECASE), + re.compile(r"^r(\d+)(?=[-.])", re.IGNORECASE), +) +_LOCAL_TP_RANK_PATTERN = re.compile(r"(?:^|[-_.])tp[-_](\d+)(?=[-_.]|$)", re.IGNORECASE) + + +class ReplayFinished(RuntimeError): + """The replay ended normally before another capture could be started.""" + + capture_started = False + + +def _reject_json_constant(value: str) -> Any: + raise ValueError(f"invalid JSON constant: {value}") + + +_TRACE_READ_CHARS = 64 * 1024 +_MAX_JSON_VALUE_CHARS = 1024 * 1024 +_MAX_JSON_DEPTH = 128 +_JSON_DECODER = json.JSONDecoder(parse_constant=_reject_json_constant) +_JSON_STRING_SPECIAL = re.compile(r'["\\\x00-\x1f]') + + +class _TraceValueTooLarge(ValueError): + """Switch from whole-value decoding to field-wise streaming.""" + + +def _check_deadline(deadline: float | None) -> None: + if deadline is not None and time.monotonic() >= deadline: + raise TimeoutError("trace check exceeded its remaining budget") + + +class _TraceJSONReader: + """Read a complete trace document without retaining its event array.""" + + def __init__( + self, handle: Any, deadline: float, check_alive: Callable[[], None] + ) -> None: + self.handle = handle + self.deadline = deadline + self.check_alive = check_alive + self.buffer = "" + self.position = 0 + self.eof = False + + def refill(self) -> None: + self.check_alive() + _check_deadline(self.deadline) + self.buffer = self.buffer[self.position :] + self.position = 0 + if len(self.buffer) >= _MAX_JSON_VALUE_CHARS: + raise _TraceValueTooLarge("trace JSON value requires field-wise streaming") + chunk = self.handle.read( + min(_TRACE_READ_CHARS, _MAX_JSON_VALUE_CHARS - len(self.buffer)) + ) + self.check_alive() + _check_deadline(self.deadline) + self.eof = not chunk + self.buffer += chunk + + def peek(self) -> str: + self.check_alive() + _check_deadline(self.deadline) + while True: + while ( + self.position < len(self.buffer) + and self.buffer[self.position] in " \t\r\n" + ): + self.position += 1 + if self.position < len(self.buffer): + return self.buffer[self.position] + if self.eof: + return "" + self.refill() + + def expect(self, character: str) -> None: + if self.peek() != character: + raise ValueError(f"expected {character!r} in trace JSON") + self.position += 1 + + def value(self) -> Any: + if not self.peek(): + raise ValueError("incomplete trace JSON value") + while True: + self.check_alive() + _check_deadline(self.deadline) + try: + value, end = _JSON_DECODER.raw_decode(self.buffer, self.position) + except json.JSONDecodeError: + if self.eof: + raise + self.refill() + continue + # A number can end at a chunk boundary before its exponent arrives. + if end == len(self.buffer) and not self.eof: + self.refill() + continue + if end < len(self.buffer) and self.buffer[end] not in " \t\r\n,]}": + if ( + type(value) in (int, float) + and self.buffer[end] in ".eE" + and not self.eof + ): + self.refill() + continue + raise ValueError("invalid trace JSON value delimiter") + self.position = end + return value + + def string(self) -> str | None: + """Validate strings of any length, retaining only short metadata values.""" + self.expect('"') + parts = [] + length = 0 + while True: + self.check_alive() + _check_deadline(self.deadline) + match = _JSON_STRING_SPECIAL.search(self.buffer, self.position) + end = match.start() if match else len(self.buffer) + if length <= 64: + length += end - self.position + if length <= 64: + parts.append(self.buffer[self.position : end]) + else: + parts.clear() + self.position = end + if match is None: + if self.eof: + raise ValueError("incomplete trace JSON string") + self.refill() + continue + character = self.buffer[self.position] + self.position += 1 + if character == '"': + return "".join(parts) if length <= 64 else None + if character != "\\": + raise ValueError("unescaped control character in trace JSON string") + while len(self.buffer) - self.position < 1 and not self.eof: + self.refill() + if self.position == len(self.buffer): + raise ValueError("incomplete trace JSON escape") + escape_length = 5 if self.buffer[self.position] == "u" else 1 + while len(self.buffer) - self.position < escape_length and not self.eof: + self.refill() + escaped = self.buffer[self.position : self.position + escape_length] + decoded = _JSON_DECODER.decode('"\\' + escaped + '"') + self.position += escape_length + length += len(decoded) + if length <= 64: + parts.append(decoded) + else: + parts.clear() + + def event(self) -> bool: + try: + event = self.value() + except _TraceValueTooLarge: + if self.peek() != "{": + self.skip() + return False + fields = {} + for key in self.members(): + if key in {"cat", "ph"} and self.peek() == '"': + fields[key] = self.string() + else: + if key in {"cat", "ph"}: + fields[key] = None + self.skip() + return fields.get("cat") == "kernel" and fields.get("ph") == "X" + return ( + isinstance(event, dict) + and event.get("cat") == "kernel" + and event.get("ph") == "X" + ) + + def members(self): + self.expect("{") + if self.peek() == "}": + self.position += 1 + return + while True: + if self.peek() != '"': + raise ValueError("trace JSON object key must be a string") + key = self.string() + self.expect(":") + yield key + if self.peek() == "}": + self.position += 1 + return + self.expect(",") + + def skip(self, depth: int = 0) -> None: + if depth >= _MAX_JSON_DEPTH: + raise ValueError("trace JSON nesting exceeds the parsing limit") + character = self.peek() + if character == "{": + for _key in self.members(): + self.skip(depth + 1) + elif character == "[": + self.position += 1 + if self.peek() == "]": + self.position += 1 + return + while True: + self.skip(depth + 1) + if self.peek() == "]": + self.position += 1 + return + self.expect(",") + elif character == '"': + self.string() + else: + self.value() + + +def _trace_metadata( + path: Path, *, deadline: float, check_alive: Callable[[], None] +) -> dict[str, Any]: + rank = None + for token in (path.name, path.parent.name): + for pattern in _TRACE_RANK_PATTERNS: + match = pattern.search(token) + if match: + rank = int(match.group(1)) + break + if rank is not None: + break + global_filename_rank = rank + if rank is None: + # SGLang's TP component is local to a pipeline stage. Prefer the global + # process rank in PyTorch's distributedInfo when that metadata exists. + for token in (path.name, path.parent.name): + match = _LOCAL_TP_RANK_PATTERN.search(token) + if match: + rank = int(match.group(1)) + break + opener = gzip.open if path.suffix == ".gz" else open + has_kernel = False + seen = set() + with opener(path, "rt", encoding="utf-8") as handle: + reader = _TraceJSONReader(handle, deadline, check_alive) + for key in reader.members(): + if key in {"traceEvents", "distributedInfo"}: + if key in seen: + raise ValueError("duplicate trace metadata field") + seen.add(key) + if key == "traceEvents": + reader.expect("[") + if reader.peek() == "]": + reader.position += 1 + continue + while True: + if reader.event(): + has_kernel = True + if reader.peek() == "]": + reader.position += 1 + break + reader.expect(",") + elif key == "distributedInfo" and reader.peek() == "{": + rank_seen = False + for field in reader.members(): + if field == "rank": + header_rank = reader.value() + if ( + rank_seen + or type(header_rank) is not int + or ( + global_filename_rank is not None + and global_filename_rank != header_rank + ) + ): + raise ValueError("conflicting trace rank") + rank = header_rank + rank_seen = True + else: + reader.skip() + else: + reader.skip() + if reader.peek(): + raise ValueError("trailing content after trace JSON document") + return {"rank": rank, "has_kernel": has_kernel and "traceEvents" in seen} + + +def _trace_files(directory: Path) -> list[Path]: + return sorted( + path + for path in directory.rglob("*.trace.json*") + if path.is_file() + and not path.is_symlink() + and path.name.endswith((".trace.json", ".trace.json.gz")) + and not path.name.startswith(("graph_capture_", "merged-")) + and not (_STARTUP_TRACE_DIRS | {"trace_split"}).intersection( + path.relative_to(directory).parts + ) + ) + + +def _signature(path: Path) -> tuple[int, ...]: + stat = path.stat() + return stat.st_mtime_ns, stat.st_size, stat.st_ctime_ns, stat.st_ino + + +def _complete_traces( + directory: Path, + expected_ranks: int, + deadline: float, + check_alive: Callable[[], None], + cache: dict[Path, tuple[tuple[int, ...], dict[str, Any] | None]], +) -> list[str] | None: + files = _trace_files(directory) + ranks: set[int] = set() + valid = [] + for path in files: + check_alive() + _check_deadline(deadline) + signature = _signature(path) + cached = cache.get(path) + if cached is not None and cached[0] == signature: + metadata = cached[1] + else: + try: + metadata = _trace_metadata( + path, deadline=deadline, check_alive=check_alive + ) + except (EOFError, ValueError, UnicodeError, RecursionError, zlib.error): + metadata = None + except OSError as exc: + if isinstance(exc, TimeoutError): + raise + metadata = None + cache[path] = signature, metadata + if _signature(path) != signature or metadata is None: + return None + if not metadata["has_kernel"]: + continue + rank = metadata["rank"] + if rank is None and expected_ranks == 1: + rank = 0 + if type(rank) is not int or rank not in range(expected_ranks): + return None + # A worker may export several stage/schedule traces. They add evidence + # for that rank, but cannot substitute for another missing worker. + ranks.add(rank) + valid.append(str(path)) + if ranks == set(range(expected_ranks)) and files == _trace_files(directory): + return valid + return None + + +def _http(url: str, *, timeout: float, body: dict[str, Any] | None = None) -> bytes: + data = None if body is None else json.dumps(body).encode("utf-8") + request = urllib.request.Request( + url, + data=data, + headers={"Accept": "application/json", "Content-Type": "application/json"}, + method="GET" if body is None else "POST", + ) + opener = urllib.request.build_opener(urllib.request.ProxyHandler({})) + with opener.open(request, timeout=timeout) as response: + payload = response.read(1024 * 1024 + 1) + if len(payload) > 1024 * 1024: + raise ValueError("AgentX profiler control response exceeds 1 MiB") + return payload + + +def _wait_phase( + progress_url: str, + deadline: float, + check_alive: Callable[[], None], + client_started_ns: int | None, + measurement_duration_seconds: float | None = None, + start_seconds: float = 0.0, +) -> int: + last_error = "" + while True: + check_alive() + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError(f"AgentX profiling phase did not start: {last_error}") + try: + payload = json.loads( + _http(progress_url, timeout=min(_HTTP_SECONDS, remaining)), + parse_constant=_reject_json_constant, + ) + phases = payload.get("phases") if isinstance(payload, dict) else None + stats = phases.get("profiling") if isinstance(phases, dict) else None + if not isinstance(stats, dict): + raise ValueError( + "AIPerf progress has no profiling phase" + ) # noqa: TRY004 + start = stats.get("start_ns") + if ( + type(start) is int + and start > 0 + and client_started_ns is not None + and start < client_started_ns + ): + raise ValueError("AIPerf progress belongs to an earlier client") + if stats.get("was_cancelled") is True: + raise RuntimeError("AIPerf profiling phase was cancelled") + if type(start) is int and start > 0: + if ( + stats.get("requests_end_ns") is not None + or stats.get("sent_end_ns") is not None + ): + raise ReplayFinished( + "AIPerf profiling phase ended before capture started" + ) + if ( + measurement_duration_seconds is not None + and (time.time_ns() - start) / 1e9 >= measurement_duration_seconds + ): + raise ReplayFinished( + "AIPerf measurement window ended before capture started" + ) + if (time.time_ns() - start) / 1e9 >= start_seconds: + return start + last_error = "waiting for the requested measurement start offset" + else: + last_error = "AIPerf is still preparing or warming up" + except (OSError, ValueError, http.client.HTTPException) as exc: + last_error = str(exc) + time.sleep(min(_POLL_SECONDS, max(0, deadline - time.monotonic()))) + + +def _positive_seconds(value: Any, name: str) -> float: + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise ValueError( + f"AgentX {name} must be a positive finite number" + ) # noqa: TRY004 + try: + result = float(value) + except OverflowError as exc: + raise ValueError(f"AgentX {name} must be a positive finite number") from exc + if not math.isfinite(result) or result <= 0: + raise ValueError(f"AgentX {name} must be a positive finite number") + return result + + +def _nonnegative_seconds(value: Any, name: str) -> float: + if type(value) in (int, float) and value == 0: + return 0.0 + return _positive_seconds(value, name) + + +def _save_capture(directory: Path, capture: dict[str, Any]) -> None: + capture["recorded_at_ns"] = time.time_ns() + temporary = directory / ".capture.json.tmp" + temporary.write_text(json.dumps(capture, indent=2, sort_keys=True) + "\n") + temporary.replace(directory / "capture.json") + + +def _check_fresh_trace_directory(directory: Path) -> None: + for path in directory.rglob("*"): + if not path.is_symlink(): + if path.is_dir(): + continue + if _STARTUP_TRACE_DIRS.intersection(path.relative_to(directory).parts[:-1]): + continue + raise ValueError( + "AgentX profile capture requires a fresh unique trace directory; " + "only empty directories and startup trace directories may exist" + ) + + +def _share_graph_traces(archive: Path, root: Path) -> dict[str, str]: + """Keep startup evidence in this invocation, reachable from every window.""" + shared = {} + for name in sorted(_STARTUP_TRACE_DIRS): + source, target = archive / name, root / name + if source.is_symlink() or target.is_symlink(): + raise ValueError("AgentX graph trace directories cannot be external links") + if source.exists(): + if not source.is_dir() or target.exists(): + raise ValueError( + "AgentX graph trace directory was recreated during capture" + ) + if any(path.is_symlink() for path in source.rglob("*")): + raise ValueError("AgentX graph trace files cannot be external links") + source.rename(target) + if target.exists(): + if not target.is_dir() or any( + path.is_symlink() for path in target.rglob("*") + ): + raise ValueError( + "AgentX shared graph traces must remain within this run" + ) + source.symlink_to(Path("..") / name, target_is_directory=True) + shared[name] = str(target) + return shared + + +def capture_profile( + *, + framework: str, + server_url: str, + progress_url: str, + trace_dir: Path, + settings: dict[str, Any], + expected_ranks: int, + phase_timeout_seconds: float, + check_alive: Callable[[], None], + client_started_ns: int | None = None, + check_replay_alive: Callable[[], None] | None = None, + measurement_duration_seconds: float | None = None, +) -> dict[str, Any]: + """Wait for the measured replay, trigger bounded profiling and verify its files. + + The caller configures vLLM's worker step limit before launching the server. + ``check_alive`` must raise on cancellation, failed child processes or the + overall deadline; a successful client exit must still permit trace flushing. + ``check_replay_alive`` also rejects a successful client exit while waiting + to start profiling, when further replay work is still required. + The trace directory must be unique to this capture. Empty directories and + known server-startup graph trace directories may already exist in it. + """ + if framework not in {"sglang", "vllm"}: + raise ValueError("AgentX profiling supports SGLang and vLLM") + if type(expected_ranks) is not int or expected_ranks <= 0: + raise ValueError("AgentX expected_ranks must be a positive integer") + steps = settings.get("num_steps", 20) + if type(steps) is not int or steps <= 0: + raise ValueError("AgentX num_steps must be a positive integer") + start_seconds = _nonnegative_seconds( + settings.get("start_seconds", 0), "start_seconds" + ) + annotations = settings.get("detailed_annotations", False) + if type(annotations) is not bool: + raise ValueError("AgentX detailed_annotations must be a boolean") + capabilities = None + if annotations: + if __package__: + from .agentx_profile_capabilities import validate_profile_capabilities + else: + from agentx_profile_capabilities import validate_profile_capabilities + capabilities = validate_profile_capabilities( + framework, settings.get("capabilities") + ) + phase_timeout = _positive_seconds(phase_timeout_seconds, "phase_timeout_seconds") + capture_timeout = _positive_seconds( + settings.get("capture_timeout_seconds", 300), "capture_timeout_seconds" + ) + flush_timeout = _positive_seconds( + settings.get("flush_timeout_seconds", 1800), "flush_timeout_seconds" + ) + if measurement_duration_seconds is not None: + measurement_duration_seconds = _positive_seconds( + measurement_duration_seconds, "measurement_duration_seconds" + ) + if start_seconds >= measurement_duration_seconds: + raise ValueError( + "AgentX start_seconds leaves no measurement capture window" + ) + if client_started_ns is not None and ( + type(client_started_ns) is not int or client_started_ns <= 0 + ): + raise ValueError("AgentX client_started_ns must be a positive integer") + trace_dir = Path(trace_dir).resolve() + trace_dir.mkdir(parents=True, exist_ok=True) + _check_fresh_trace_directory(trace_dir) + capture: dict[str, Any] = { + "version": 1, + "capture_id": trace_dir.name, + "status": "failed", + "framework": framework, + "num_steps": steps, + "start_seconds": start_seconds, + "detailed_annotations": annotations, + "expected_ranks": expected_ranks, + "phase_start_ns": None, + "trace_files": [], + } + if capabilities is not None: + capture["capabilities"] = capabilities + attempted = False + failure: BaseException | None = None + + def check_before_start() -> None: + check_alive() + if check_replay_alive is not None: + check_replay_alive() + + try: + capture["phase_start_ns"] = _wait_phase( + progress_url, + time.monotonic() + phase_timeout, + check_before_start, + client_started_ns, + measurement_duration_seconds, + start_seconds, + ) + check_before_start() + body = {} + if framework == "sglang": + body = { + "num_steps": steps, + "output_dir": str(trace_dir), + "activities": ["CPU", "GPU"], + "profile_prefix": trace_dir.name, + } + if capabilities is not None: + body.update( + with_stack=True, + record_shapes=True, + shape_discovery=True, + **{capabilities["annotation_field"]: True}, + ) + deadline = time.monotonic() + capture_timeout + if ( + measurement_duration_seconds is not None + and (time.time_ns() - capture["phase_start_ns"]) / 1e9 + >= measurement_duration_seconds + ): + raise ReplayFinished( + "AIPerf measurement window ended before capture started" + ) + capture["capture_started_ns"] = time.time_ns() + attempted = True + _http( + server_url.rstrip("/") + "/start_profile", + timeout=capture_timeout, + body=body, + ) + check_alive() + if time.monotonic() >= deadline: + raise TimeoutError("AgentX profiling start exceeded capture timeout") + while not _trace_files(trace_dir): + check_alive() + if time.monotonic() >= deadline: + raise TimeoutError( + "AgentX bounded profiling produced no trace before capture timeout" + ) + time.sleep(min(_POLL_SECONDS, max(0, deadline - time.monotonic()))) + capture["flush_started_ns"] = time.time_ns() + deadline = time.monotonic() + flush_timeout + cache: dict[Path, tuple[tuple[int, ...], dict[str, Any] | None]] = {} + while True: + check_alive() + if time.monotonic() >= deadline: + raise TimeoutError( + "AgentX trace flush did not produce complete GPU traces for " + f"ranks 0..{expected_ranks - 1}" + ) + complete = _complete_traces( + trace_dir, expected_ranks, deadline, check_alive, cache + ) + if complete is not None: + rank_trace_files: dict[str, list[str]] = {} + for name in complete: + rank = cache[Path(name)][1]["rank"] + rank_trace_files.setdefault( + str(0 if rank is None else rank), [] + ).append(name) + capture.update( + status="complete", + trace_files=complete, + rank_trace_files=rank_trace_files, + ) + capture["graph_trace_dirs"] = { + name: str(trace_dir / name) + for name in sorted(_STARTUP_TRACE_DIRS) + if (trace_dir / name).is_dir() + } + return capture + time.sleep(min(_POLL_SECONDS, max(0, deadline - time.monotonic()))) + except BaseException as exc: + # Persist failure and stop an uncertain capture even on cancellation; + # the original exception still propagates to the runtime's owner. + failure = exc + if isinstance(exc, ReplayFinished): + exc.capture_started = attempted + capture["error"] = f"{type(exc).__name__}: {exc}" + raise + finally: + if failure is not None and attempted: + try: + _http( + server_url.rstrip("/") + "/stop_profile", + timeout=min(_HTTP_SECONDS, flush_timeout), + body={}, + ) + capture["cleanup"] = "stop_requested" + except (OSError, ValueError, http.client.HTTPException) as exc: + capture["cleanup"] = f"stop request: {type(exc).__name__}: {exc}" + _save_capture(trace_dir, capture) + + +def capture_profiles( + *, + framework: str, + server_url: str, + progress_url: str, + trace_dir: Path, + settings: dict[str, Any], + expected_ranks: int, + phase_timeout_seconds: float, + check_alive: Callable[[], None], + client_started_ns: int | None = None, + check_replay_alive: Callable[[], None] | None = None, + profile_plan: dict[str, Any] | None = None, +) -> dict[str, Any]: + """Collect sequential captures from one server and one measured replay. + + Single captures retain the original directory and manifest layout. Multiple + captures use ``active`` as the fixed server output directory and archive each + completed capture before the next starts. The replay deadline is shared by + all waits and intervals; trace flushing retains its separate budget. + """ + count = settings.get("num_profiles", 1) + if type(count) is not int or count <= 0: + raise ValueError("AgentX num_profiles must be a positive integer") + interval = _nonnegative_seconds( + settings.get("interval_seconds", 200), "interval_seconds" + ) + start_seconds = _nonnegative_seconds( + settings.get("start_seconds", 0), "start_seconds" + ) + plan: dict[str, Any] = { + "requested_profiles": count, + "max_profiles": None, + "planned_profiles": count, + "measurement_duration_seconds": None, + "start_seconds": start_seconds, + } + if profile_plan is not None: + if not isinstance(profile_plan, dict): + raise ValueError("AgentX profile_plan must be an object") + plan.update({key: profile_plan.get(key, value) for key, value in plan.items()}) + requested, planned, maximum = ( + plan["requested_profiles"], + plan["planned_profiles"], + plan["max_profiles"], + ) + if ( + type(requested) is not int + or requested != count + or type(planned) is not int + or not 1 <= planned <= count + or (maximum is not None and (type(maximum) is not int or maximum < planned)) + or plan["start_seconds"] != start_seconds + ): + raise ValueError("AgentX profile_plan has inconsistent capture counts") + plan["measurement_duration_seconds"] = _positive_seconds( + plan["measurement_duration_seconds"], "measurement_duration_seconds" + ) + if start_seconds >= plan["measurement_duration_seconds"]: + raise ValueError( + "AgentX start_seconds leaves no measurement capture window" + ) + planned_count = plan["planned_profiles"] + capture_args = { + "framework": framework, + "server_url": server_url, + "progress_url": progress_url, + "trace_dir": trace_dir, + "settings": settings, + "expected_ranks": expected_ranks, + "phase_timeout_seconds": phase_timeout_seconds, + "check_alive": check_alive, + "client_started_ns": client_started_ns, + "check_replay_alive": check_replay_alive, + "measurement_duration_seconds": plan["measurement_duration_seconds"], + } + if count == 1: + capture = capture_profile(**capture_args) + if profile_plan is not None: + capture.update(plan, effective_profiles=1, completed_profiles=1) + _save_capture(Path(trace_dir).resolve(), capture) + return capture + + deadline = time.monotonic() + _positive_seconds( + phase_timeout_seconds, "phase_timeout_seconds" + ) + capture_timeout = _positive_seconds( + settings.get("capture_timeout_seconds", 300), "capture_timeout_seconds" + ) + flush_timeout = _positive_seconds( + settings.get("flush_timeout_seconds", 1800), "flush_timeout_seconds" + ) + overall_deadline = deadline + capture_timeout + flush_timeout + trace_dir = Path(trace_dir).resolve() + trace_dir.mkdir(parents=True, exist_ok=True) + _check_fresh_trace_directory(trace_dir) + active = trace_dir / "active" + summary: dict[str, Any] = { + "version": 1, + "capture_id": trace_dir.name, + "status": "running", + "framework": framework, + "num_steps": settings.get("num_steps", 20), + "detailed_annotations": settings.get("detailed_annotations", False), + "expected_ranks": expected_ranks, + **plan, + "effective_profiles": planned_count, + "completed_profiles": 0, + "interval_seconds": interval, + "profiles": [], + "trace_files": [], + } + if settings.get("capabilities") is not None: + summary["capabilities"] = settings["capabilities"] + if planned_count < count: + summary["stop_reason"] = "duration_cap" + preparing = True + measurement_deadline_ns: float | None = None + + def check_replay_budget() -> float: + check_alive() + if check_replay_alive is not None: + check_replay_alive() + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError( + "AgentX replay budget ended before all profiles completed" + ) + return remaining + + try: + for index in range(1, planned_count + 1): + preparing = True + remaining = check_replay_budget() + archive = trace_dir / f"profile_{index:03d}" + if archive.exists() or archive.is_symlink(): + raise FileExistsError( + f"AgentX profile archive already exists: {archive}" + ) + capture_args.update(trace_dir=active, phase_timeout_seconds=remaining) + capture = capture_profile(**capture_args) + preparing = False + if ( + measurement_deadline_ns is None + and plan["measurement_duration_seconds"] is not None + ): + measurement_deadline_ns = ( + capture["phase_start_ns"] + + plan["measurement_duration_seconds"] * 1e9 + ) + if framework == "vllm": + # Older vLLM workers export at max_iterations but keep their + # active flag set. Reset it before another start can take effect. + remaining = min(flush_timeout, overall_deadline - time.monotonic()) + if remaining <= 0: + raise TimeoutError( + "AgentX profiler reset exceeded the overall budget" + ) + _http( + server_url.rstrip("/") + "/stop_profile", + timeout=remaining, + body={}, + ) + check_alive() + if time.monotonic() >= overall_deadline: + raise TimeoutError( + "AgentX profiler reset exceeded the overall budget" + ) + active.rename(archive) + capture["graph_trace_dirs"] = _share_graph_traces(archive, trace_dir) + capture["capture_id"] = trace_dir.name + capture["profile_index"] = index + capture["trace_dir"] = str(archive) + capture["trace_files"] = [ + str(archive / Path(path).relative_to(active)) + for path in capture["trace_files"] + ] + capture["rank_trace_files"] = { + rank: [str(archive / Path(path).relative_to(active)) for path in paths] + for rank, paths in capture["rank_trace_files"].items() + } + _save_capture(archive, capture) + summary["profiles"].append(capture) + if capture["graph_trace_dirs"]: + summary["graph_trace_dirs"] = capture["graph_trace_dirs"] + summary["completed_profiles"] = index + summary["trace_files"].extend(capture["trace_files"]) + _save_capture(trace_dir, summary) + if index < planned_count: + preparing = True + check_alive() + if ( + measurement_deadline_ns is not None + and (measurement_deadline_ns - time.time_ns()) / 1e9 <= interval + ): + summary["effective_profiles"] = index + summary["stop_reason"] = "insufficient_measurement_time" + break + interval_deadline = time.monotonic() + interval + while time.monotonic() < interval_deadline: + remaining = check_replay_budget() + time.sleep( + max( + 0, + min( + _POLL_SECONDS, + interval_deadline - time.monotonic(), + remaining, + ), + ) + ) + summary["status"] = "complete" + return summary + except BaseException as exc: + if ( + isinstance(exc, ReplayFinished) + and not exc.capture_started + and preparing + and summary["completed_profiles"] > 0 + ): + summary["status"] = "complete" + summary["effective_profiles"] = summary["completed_profiles"] + summary["stop_reason"] = "replay_finished" + summary["stop_detail"] = str(exc) + return summary + summary["status"] = "failed" + summary["error"] = f"{type(exc).__name__}: {exc}" + summary["failed_profile"] = index + if active.exists(): + summary["active_trace_dir"] = str(active) + raise + finally: + _save_capture(trace_dir, summary) diff --git a/Magpie/modes/benchmark/agentx_recipe.py b/Magpie/modes/benchmark/agentx_recipe.py new file mode 100644 index 0000000..c3c2151 --- /dev/null +++ b/Magpie/modes/benchmark/agentx_recipe.py @@ -0,0 +1,565 @@ +"""Materialize single-node native recipes without a Slurm deployment. + +InferenceX still owns point validation and golden acceptance. Magpie owns the +literal server command and its lifecycle; unsupported deployment semantics are +rejected instead of silently discarded. +""" + +from __future__ import annotations + +import copy +import hashlib +import json +import re +import subprocess +import sys +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import yaml + +if TYPE_CHECKING: + from .config import BenchmarkConfig + +CLIENT_SCRIPT = "srt_agentic.sh" + +# These upstream APIs perform no submission and import srtctl only in other, +# unused functions. A separate interpreter avoids altering Magpie's sys.path or +# reusing a module cached from a different checkout. +_VALIDATE = """ +import json, sys +sys.path.insert(0, sys.argv[1]) +from infx.srt_slurm.single_node import validate_recipe +from infx.srt_slurm.synthetic_acceptance import build_overrides +data = json.load(sys.stdin) +matches, errors = [], [] +for name, recipe in data['variants']: + try: + validate_recipe(recipe, data['environment']) + except (ValueError, KeyError, TypeError) as exc: + errors.append(f'{name}: {exc}') + else: + matches.append((name, recipe)) +if len(matches) != 1: + raise ValueError('Expected exactly one matching native recipe: ' + + ('; '.join(errors) if not matches else str([x[0] for x in matches]))) +name, recipe = matches[0] +overrides = build_overrides(recipe, data['environment']['FRAMEWORK'], data['environment']) +print(json.dumps({'variant': name, 'recipe': recipe, 'overrides': overrides})) +""" + + +def _merge(base: dict[str, Any], override: dict[str, Any]) -> dict[str, Any]: + result = copy.deepcopy(base) + for key, value in override.items(): + if isinstance(value, dict) and isinstance(result.get(key), dict): + result[key] = _merge(result[key], value) + else: + result[key] = copy.deepcopy(value) + return result + + +def _variants(raw: dict[str, Any], selector: str) -> list[tuple[str, dict[str, Any]]]: + if "base" not in raw: + if selector: + raise ValueError("Native recipe selector requires a base/override recipe") + return [("", raw)] + if any(key.startswith("zip_override_") for key in raw): + raise ValueError("Magpie does not support zipped native recipe sweeps") + if not isinstance(raw["base"], dict): + raise ValueError("Native recipe base must be a mapping") + names = sorted(key for key in raw if key.startswith("override_")) + if selector == "base" or not names: + return [("base", raw["base"])] + if selector: + if selector not in names: + raise ValueError(f"Unknown native recipe variant: {selector}") + names = [selector] + if any(not isinstance(raw[name], dict) for name in names): + raise ValueError("Native recipe overrides must be mappings") + return [(name, _merge(raw["base"], raw[name])) for name in names] + + +def _strings(value: Any, description: str) -> dict[str, str]: + if not isinstance(value, dict) or any( + not isinstance(key, str) + or not re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", key) + or not isinstance(item, (str, int, float, bool)) + for key, item in value.items() + ): + raise ValueError(f"{description} must map environment names to scalar values") + return { + key: str(item).lower() if isinstance(item, bool) else str(item) + for key, item in value.items() + } + + +def _apply_golden(recipe: dict[str, Any], overrides: list[str]) -> None: + if len(overrides) % 2: + raise ValueError("Malformed upstream golden acceptance overrides") + for operation, raw in zip(overrides[::2], overrides[1::2]): + path, _, value = raw.partition("=") + keys = path.split(".") + if keys[0] not in {"roles", "environment"}: + raise ValueError(f"Unexpected upstream golden acceptance target: {path}") + target = recipe + for key in keys[:-1]: + if not isinstance(target, dict): + raise ValueError("Native recipe roles/environment must be mappings") + target = target.setdefault(key, {}) + if not isinstance(target, dict): + raise ValueError("Native recipe roles/environment must be mappings") + if operation == "--set": + target[keys[-1]] = json.loads(value) + elif operation == "--unset": + target.pop(keys[-1], None) + else: + raise ValueError( + f"Unexpected upstream golden acceptance operation: {operation}" + ) + + +def _flags(arguments: dict[str, Any]) -> list[str]: + tokens: list[str] = [] + for name, value in arguments.items(): + if not isinstance(name, str) or not re.fullmatch(r"[a-zA-Z][\w-]*", name): + raise ValueError(f"Invalid native server option: {name!r}") + if value is None or value is False: + continue + tokens.append("--" + name) + if value is True: + continue + if isinstance(value, dict): + tokens.append(json.dumps(value, separators=(",", ":"))) + elif isinstance(value, list): + if any(not isinstance(item, (str, int, float)) for item in value): + raise ValueError(f"Unsupported native server list argument: {name}") + tokens.extend(str(item) for item in value) + elif isinstance(value, (str, int, float)): + tokens.append(str(value)) + else: + raise ValueError(f"Unsupported native server argument: {name}") + return tokens + + +def _port(config: BenchmarkConfig) -> int: + raw = config.envs.get("PORT", 8888) + if isinstance(raw, bool) or not str(raw).isdigit() or not 0 < int(raw) < 65536: + raise ValueError("AgentX PORT must be an integer between 1 and 65535") + return int(raw) + + +def _file_hashes(paths: list[Path]) -> dict[str, str]: + try: + return { + str(path.resolve()): hashlib.sha256(path.read_bytes()).hexdigest() + for path in paths + } + except OSError as exc: + raise ValueError( + f"Required AgentX protocol/recipe source is unreadable: {exc}" + ) from exc + + +def resolve_client_sources(root: Path) -> dict[str, str]: + """Bind the official client, shared shell, aggregation, and acceptance gate.""" + files = [ + root / "benchmarks/srt_agentic.sh", + root / "benchmarks/benchmark_lib.sh", + root / "benchmarks/runtime_settings.sh", + root / "runners/srt-slurm/hooks/common.sh", + root / "infx/results/agentic/process_agentic_result.py", + root / "infx/results/agentic/validate_agentic_result.py", + root / "infx/__init__.py", + root / "infx/results/__init__.py", + root / "infx/results/metadata.py", + root / "infx/results/topology.py", + ] + files += sorted((root / "infx/results/agentic").rglob("*.py")) + files += sorted((root / "infx/results/power").rglob("*.py")) + return _file_hashes(files) + + +def _client_revisions(root: Path) -> dict[str, str]: + """Record checkout and harness pins when resolving an actual Git checkout.""" + head = subprocess.run( + ["git", "-C", str(root), "rev-parse", "HEAD"], + capture_output=True, + text=True, + timeout=10, + ) + if head.returncode or not re.fullmatch(r"[0-9a-f]{40}", head.stdout.strip()): + return {} + result = {"client_revision": head.stdout.strip()} + harness = subprocess.run( + ["git", "-C", str(root), "ls-tree", "HEAD", "--", "utils/aiperf"], + capture_output=True, + text=True, + timeout=10, + ) + match = re.fullmatch( + r"160000 commit ([0-9a-f]{40})\tutils/aiperf", harness.stdout.strip() + ) + if harness.returncode == 0 and match: + result["aiperf_revision"] = match[1] + return result + + +def _client_defaults(root: Path, config: BenchmarkConfig) -> dict[str, str]: + result = {} + for line in (root / "benchmarks/runtime_settings.sh").read_text().splitlines(): + match = re.fullmatch(r"export ([A-Z_]+)='([^']*)'", line) + if match and ( + match[1].startswith(("AIPERF_", "AGENTIC_")) + or match[1] == "ENABLE_AGENTX_POWER" + ): + result[match[1]] = match[2] + required = { + "AIPERF_PYTHON_VERSION", + "AIPERF_LIVE_FAILED_REQUEST_THRESHOLD", + "AIPERF_TRACE_IDLE_GAP_CAP_SECONDS", + "AIPERF_WARMUP_REQUESTS_PER_LANE", + "AGENTIC_WARMUP_GRACE_PERIOD", + "AIPERF_DATASET_WEKA_LIVE_ASSISTANT_RESPONSES", + "AIPERF_DYNAMO_SESSION_TIMEOUT_SECONDS", + "AIPERF_HTTP_X_DYNAMO_SESSION_ID_FROM_CORRELATION_ID", + "AIPERF_UNSAFE_OVERRIDE", + "AIPERF_USE_DYNAMO_CONV_AWARE_ROUTING", + "ENABLE_AGENTX_POWER", + } + if any(not result.get(name) for name in required): + raise ValueError( + "InferenceX client runtime settings are missing literal AgentX defaults" + ) + assert config.agentx is not None + result["AIPERF_FAILED_REQUEST_THRESHOLD"] = str( + config.agentx.failed_request_threshold + ) + result["AIPERF_EXPERIMENTAL_FAST"] = "1" if config.agentx.mode == "fast" else "0" + # The upstream workflow (benchmark-tmpl.yml), not runtime_settings.sh, + # supplies this default. Enable collection without requiring valid power + # telemetry unless the caller or recipe explicitly requests that gate. + result["REQUIRE_POWER"] = "0" + result.update(_client_power_settings(result, config)) + return result + + +def _client_power_settings( + client: dict[str, str], config: BenchmarkConfig +) -> dict[str, str]: + result = {} + for name in ("ENABLE_AGENTX_POWER", "REQUIRE_POWER"): + value = str(config.envs.get(name, client[name])).lower() + if value not in {"0", "1", "false", "true", "no", "yes"}: + raise ValueError(f"AgentX {name} must be a boolean power setting") + result[name] = "1" if value in {"1", "true", "yes"} else "0" + return result + + +def _server_environment(config: BenchmarkConfig) -> dict[str, str]: + # Materialization writes replay metadata back to config.envs. Do not feed + # those fields into the server or change spec identity on snapshot reload. + if any(key.upper().startswith("SGLANG_SIMULATE_ACC_") for key in config.envs): + raise ValueError("AgentX golden acceptance environment is owned by InferenceX") + client_names = { + "MODEL", + "MODEL_PATH", + "MODEL_PREFIX", + "FRAMEWORK", + "IMAGE", + "PRECISION", + "EXP_NAME", + "TP", + "PP", + "PP_SIZE", + "EP", + "EP_SIZE", + "DCP_SIZE", + "PCP_SIZE", + "DP_ATTENTION", + "CONC", + "SPEC_DECODING", + "DISAGG", + "IS_AGENTIC", + "IS_MULTINODE", + "KV_OFFLOADING", + "KV_OFFLOAD_BACKEND", + "KV_OFFLOAD_BACKEND_METADATA", + "ROUTER_METADATA", + "KV_P2P_TRANSFER", + "TOTAL_CPU_DRAM_GB", + "DURATION", + "RUN_EVAL", + "EVAL_ONLY", + "RECIPE_FINGERPRINT", + "MAX_MODEL_LEN", + "WEKA_LOADER_OVERRIDE", + "PORT", + "GPU_COUNT", + "THINKING_MODE", + "ENABLE_AGENTX_POWER", + "REQUIRE_POWER", + } + return _strings( + { + key: value + for key, value in config.envs.items() + if key.upper() not in client_names + and not (key.upper().startswith("EXTRA_") and key.upper().endswith("_ARGS")) + and not key.upper().startswith( + ("AIPERF_", "AGENTIC_", "AGENTX_", "SCENARIO_") + ) + }, + "Explicit server environment", + ) + + +def native_server_spec( + config: BenchmarkConfig, root: Path, entry: dict[str, Any] +) -> dict[str, Any]: + """Resolve one YAML variant, preserving upstream validation and acceptance.""" + filename, _, selector = str(entry["srt-recipe"]).partition(":") + path = (root / filename).resolve() + if not path.is_relative_to(root) or not path.is_file(): + raise ValueError( + f"Native AgentX recipe is missing or outside InferenceX: {filename}" + ) + raw = yaml.safe_load(path.read_text(encoding="utf-8")) + if not isinstance(raw, dict): + raise ValueError("Native AgentX recipe must be a mapping") + environment = { + "MODEL": entry["model"], + "MODEL_PREFIX": entry["model-prefix"], + "FRAMEWORK": entry["framework"], + "PRECISION": entry["precision"], + "IMAGE": config.docker_image or entry["image"], + "TP": str(entry["tp"]), + "EP_SIZE": str(entry["ep"]), + "PP_SIZE": str(entry.get("pp", 1)), + "DCP_SIZE": str(entry.get("dcp-size", 1)), + "PCP_SIZE": str(entry.get("pcp-size", 1)), + "GPU_COUNT": str(entry["tp"] * entry.get("pp", 1) * entry.get("pcp-size", 1)), + "DP_ATTENTION": str(entry.get("dp-attn", False)).lower(), + "CONC": str(entry["conc"]), + "KV_OFFLOADING": entry["kv-offloading"], + "TOTAL_CPU_DRAM_GB": str(entry["total-cpu-dram-gb"]), + "SPEC_DECODING": entry["spec-decoding"], + "EVAL_ONLY": "false", + "IS_AGENTIC": "1", + # This is the official single-node workflow default, not an inferred AL. + "THINKING_MODE": str(config.envs.get("THINKING_MODE", "thinking_on")), + } + completed = subprocess.run( + [sys.executable, "-c", _VALIDATE, str(root)], + input=json.dumps( + {"variants": _variants(raw, selector), "environment": environment} + ), + capture_output=True, + text=True, + timeout=30, + ) + if completed.returncode: + raise ValueError( + "InferenceX native recipe validation failed: " + completed.stderr.strip() + ) + selected = json.loads(completed.stdout) + recipe = selected["recipe"] + _apply_golden(recipe, selected["overrides"]) + allowed = { + "schema", + "name", + "model", + "resources", + "frontend", + "observability", + "engine", + "health_check", + "roles", + "benchmark", + "environment", + "setup_script", + } + if set(recipe) - allowed or recipe.get("schema") != 2: + raise ValueError("Unsupported native AgentX recipe deployment fields or schema") + engine = recipe["engine"] + if isinstance(engine, dict): + if set(engine) - {"type", "connector"} or engine.get("connector") is not None: + raise ValueError("Magpie does not support native engine connectors") + engine = engine["type"] + if engine not in {"sglang", "vllm"}: + raise ValueError("Magpie-owned AgentX supports SGLang and vLLM only") + frontend = recipe["frontend"] + if ( + frontend.get("type") != engine + or frontend.get("enable_multiple_frontends", False) + or frontend.get("args") + or frontend.get("env") + ): + raise ValueError( + "Magpie-owned AgentX does not support router or multiple frontends" + ) + observability = recipe.get("observability", {}) + if observability.get("enabled") or observability.get("tachometer", {}).get( + "enabled" + ): + raise ValueError( + "Magpie-owned AgentX does not support recipe observability services" + ) + role = recipe["roles"]["agg"] + if set(role) - {"nodes", "workers", "gpus", "args", "env"}: + raise ValueError("Unsupported native AgentX worker deployment fields") + arguments = dict(role["args"]) + for name in ("host", "port", "model", "model-path"): + if name in arguments: + raise ValueError(f"Native recipe must leave runtime-owned {name} to Magpie") + env = _strings(role.get("env", {}), "Native worker environment") + env.update(_strings(recipe.get("environment", {}), "Native global environment")) + env.update(_server_environment(config)) + port = _port(config) + model_path = str(config.envs.get("MODEL_PATH") or config.model) + if engine == "sglang": + argv = ["python3", "-m", "sglang.launch_server", "--model-path", model_path] + else: + argv = [ + "python3", + "-m", + "vllm.entrypoints.openai.api_server", + "--model", + model_path, + ] + argv += ["--host", "0.0.0.0", "--port", str(port), *_flags(arguments)] + setup = [] + files = [ + path, + root / "infx/srt_slurm/single_node.py", + root / "infx/srt_slurm/synthetic_acceptance.py", + root / "infx/golden_al_distribution/__init__.py", + root / "infx/golden_al_distribution/curves.py", + ] + files += sorted((root / "infx/golden_al_distribution").glob("*.yaml")) + if recipe.get("setup_script"): + if recipe["setup_script"] != "pip-runtime-deps.sh": + raise ValueError( + "Unsupported native AgentX setup_script: " + str(recipe["setup_script"]) + ) + script = ( + root / "benchmarks/multi_node/srt-slurm-recipes/configs/pip-runtime-deps.sh" + ) + files.append(script) + setup = [["bash", str(script)]] + client = _client_defaults(root, config) + client.update( + _strings(recipe["benchmark"].get("env", {}), "Native client environment") + ) + client.update(_client_power_settings(client, config)) + client.update(environment) + client.update( + PORT=str(port), + AIPERF_SERVER_URL=f"http://127.0.0.1:{port}", + AIPERF_SERVER_METRICS_URLS=f"http://127.0.0.1:{port}/metrics", + ) + context = arguments.get("context-length" if engine == "sglang" else "max-model-len") + if context is not None: + client["AIPERF_MAX_CONTEXT_LENGTH"] = str(context) + health = recipe.get("health_check", {}) + return { + **_client_revisions(root), + "version": 1, + "framework": engine, + "argv": argv, + "env": env, + "setup_commands": setup, + "port": port, + "health_path": "/health", + "health_timeout_seconds": int(health.get("interval_seconds", 10)) + * int(health.get("max_attempts", 180)), + "model_metadata": None, + "client_env": client, + "source_files": {**_file_hashes(files), **resolve_client_sources(root)}, + "recipe_variant": selected["variant"], + "golden_acceptance_overrides": selected["overrides"], + } + + +def custom_server_spec( + config: BenchmarkConfig, root: Path, entry: dict[str, Any] +) -> dict[str, Any]: + """Explicit custom models share the official replay, without a fake recipe.""" + port = _port(config) + model = str(config.envs.get("MODEL_PATH") or config.model) + dtype = {"bf16": "bfloat16", "fp16": "float16", "fp32": "float32"}.get( + config.precision, config.precision + ) + if dtype in {"fp8", "fp4", "mxfp4"}: + dtype = "auto" + if config.framework == "sglang": + argv = [ + "python3", + "-m", + "sglang.launch_server", + "--model-path", + model, + "--tensor-parallel-size", + str(entry["tp"]), + "--expert-parallel-size", + str(entry["ep"]), + "--context-length", + str(entry["max-model-len"]), + "--enable-metrics", + ] + else: + argv = [ + "python3", + "-m", + "vllm.entrypoints.openai.api_server", + "--model", + model, + "--tensor-parallel-size", + str(entry["tp"]), + "--max-model-len", + str(entry["max-model-len"]), + ] + if entry["ep"] > 1: + argv.append("--enable-expert-parallel") + argv += [ + "--served-model-name", + config.model, + "--host", + "0.0.0.0", + "--port", + str(port), + "--trust-remote-code", + "--dtype", + dtype, + ] + if entry.get("model-revision"): + argv += ["--revision", entry["model-revision"]] + return { + **_client_revisions(root), + "version": 1, + "framework": config.framework, + "argv": argv, + "env": _server_environment(config), + "setup_commands": [], + "port": port, + "health_path": "/health", + "health_timeout_seconds": 1800, + "model_metadata": { + "source": entry["model-config-source"], + "sha256": entry["model-config-sha256"], + "native_context_length": entry["native-context-length"], + "max_model_len": entry["max-model-len"], + "revision": entry.get("model-revision"), + }, + "client_env": { + **_client_defaults(root, config), + "AIPERF_MAX_CONTEXT_LENGTH": str(entry["max-model-len"]), + "WEKA_LOADER_OVERRIDE": entry["trace-loader"], + "PORT": str(port), + "AIPERF_SERVER_URL": f"http://127.0.0.1:{port}", + "AIPERF_SERVER_METRICS_URLS": f"http://127.0.0.1:{port}/metrics", + }, + "source_files": resolve_client_sources(root), + "recipe_variant": "custom", + "golden_acceptance_overrides": [], + } diff --git a/Magpie/modes/benchmark/agentx_runtime.py b/Magpie/modes/benchmark/agentx_runtime.py new file mode 100644 index 0000000..3e432c3 --- /dev/null +++ b/Magpie/modes/benchmark/agentx_runtime.py @@ -0,0 +1,919 @@ +############################################################################### +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. +# See LICENSE for license information. +############################################################################### +"""Own one AgentX server and its client, without invoking an upstream launcher.""" + +from __future__ import annotations + +import argparse +import hashlib +import http.client +import json +import math +import os +import re +import signal +import socket +import subprocess +import time +import urllib.error +import urllib.request +import uuid +from pathlib import Path +from typing import Any +from urllib.parse import urlparse + +# Direct execution inside the serving image does not import Magpie's optional +# analysis/evaluation dependencies. Both files are mounted from the same build. +if __package__: + from .agentx_custom import native_context_length + from .agentx_launch import prepare_server_launch + from .agentx_profile_config import ( + profile_plan, + profile_server_spec, + profile_settings, + ) +else: + from agentx_custom import native_context_length + from agentx_launch import prepare_server_launch + from agentx_profile_config import ( + profile_plan, + profile_server_spec, + profile_settings, + ) + + +class RuntimeInterrupted(RuntimeError): + """A signal interrupted this point; its owned children must be reaped.""" + + +class _PublicModelRedirect(urllib.request.HTTPRedirectHandler): + def redirect_request(self, req, fp, code, msg, headers, newurl): + parsed = urlparse(newurl) + if ( + parsed.scheme != "https" + or parsed.hostname != "huggingface.co" + or parsed.port not in {None, 443} + or parsed.username + or parsed.password + ): + raise ValueError( + "AgentX model metadata redirect left the public HuggingFace host" + ) + return super().redirect_request(req, fp, code, msg, headers, newurl) + + +def _signal_interrupt(signum, _frame): + raise RuntimeInterrupted(f"AgentX interrupted by signal {signum}") + + +def _stop_owned(process: subprocess.Popen) -> None: + """Stop only the process group created by this runtime, including descendants.""" + if process.pid <= 1 or process.pid == os.getpgrp(): + raise RuntimeError("Refusing to terminate an unowned process group") + try: + os.killpg(process.pid, signal.SIGTERM) + except ProcessLookupError: + process.wait() + return + try: + process.wait(timeout=5) + except subprocess.TimeoutExpired: + pass + # The root may exit before its workers, so leader exit alone is not cleanup. + try: + os.killpg(process.pid, signal.SIGKILL) + except ProcessLookupError: + pass + process.wait(timeout=5) + + +def _spawn( + argv: list[str], *, env: dict[str, str], cwd: Path, output +) -> subprocess.Popen: + return subprocess.Popen( + argv, + env=env, + cwd=cwd, + stdout=output, + stderr=subprocess.STDOUT, + start_new_session=True, + ) + + +def _run_owned(argv, *, env, cwd, output, timeout, server=None) -> int: + process = _spawn(argv, env=env, cwd=cwd, output=output) + deadline = time.monotonic() + timeout + try: + while process.poll() is None: + if server is not None and server.poll() is not None: + raise RuntimeError( + f"AgentX server exited while the client was running ({server.returncode})" + ) + if time.monotonic() >= deadline: + raise TimeoutError(f"AgentX command exceeded its {timeout:g}s deadline") + time.sleep(0.1) + if server is not None and server.poll() is not None: + raise RuntimeError( + f"AgentX server exited before the client completed ({server.returncode})" + ) + return process.returncode + finally: + _stop_owned(process) + + +def _check_profiled_replay_result(workspace: Path, client_started_ns: int) -> None: + """AIPerf can exit zero after cancellation; require its final completion flag.""" + result_path = workspace / "aiperf_artifacts" / "profile_export_aiperf.json" + try: + if result_path.stat().st_mtime_ns < client_started_ns: + raise ValueError("AIPerf result belongs to an earlier client") + result = json.loads(result_path.read_text(encoding="utf-8")) + if ( + not isinstance(result, dict) + or type(result.get("was_cancelled")) is not bool + ): + raise ValueError("AIPerf result has no boolean was_cancelled flag") + except (OSError, ValueError) as exc: + raise RuntimeError( + f"AgentX cannot verify replay completion from {result_path}: {exc}" + ) from exc + if result["was_cancelled"]: + raise RuntimeError("AIPerf replay was cancelled") + + +def _run_profiled_client(request, spec, *, server, output, outcome) -> int: + """Bracket framework step profiling while the official replay remains unchanged.""" + if __package__: + from .agentx_profiling import ReplayFinished, capture_profiles + else: + from agentx_profiling import ReplayFinished, capture_profiles + + workspace = Path(request["workspace"]) + root = Path(request["inferencex_path"]) + trace_dir = Path(spec["torch_profiler"]["trace_dir"]) + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe: + probe.bind(("127.0.0.1", 0)) + progress_port = probe.getsockname()[1] + env = _client_environment(request, spec, root, workspace) + env.update( + AIPERF_API_SERVER_HOST="127.0.0.1", AIPERF_API_SERVER_PORT=str(progress_port) + ) + plan = profile_plan(request["profile"], env) + output.write("Magpie profiler plan: " + json.dumps(plan, sort_keys=True) + "\n") + output.flush() + started_ns = time.time_ns() + deadline = time.monotonic() + request["client_timeout"] + overall_deadline = ( + deadline + + request["profile"]["capture_timeout_seconds"] + + request["profile"]["flush_timeout_seconds"] + ) + client = _spawn(_client_command(root), env=env, cwd=root, output=output) + replay_verified = False + + def check_alive(): + nonlocal replay_verified + if time.monotonic() >= overall_deadline: + raise TimeoutError("AgentX profiling exceeded its overall deadline") + if server.poll() is not None: + raise RuntimeError("AgentX server exited during profiler capture") + return_code = client.poll() + if return_code not in (None, 0): + raise RuntimeError( + f"AgentX client failed during profiler capture ({client.returncode})" + ) + if return_code == 0 and not replay_verified: + _check_profiled_replay_result(workspace, started_ns) + replay_verified = True + + def check_replay_alive(): + check_alive() + if client.poll() == 0: + raise ReplayFinished( + "AgentX replay ended before all requested profiler captures completed" + ) + if time.monotonic() >= deadline: + raise TimeoutError("AgentX client deadline reached before the next capture") + + try: + outcome["profile_capture"] = capture_profiles( + framework=spec["framework"], + server_url=f"http://127.0.0.1:{spec['port']}", + progress_url=f"http://127.0.0.1:{progress_port}/api/progress", + trace_dir=trace_dir, + settings={ + **request["profile"], + "capabilities": spec["torch_profiler"].get("capabilities"), + }, + expected_ranks=request["profile_ranks"], + phase_timeout_seconds=max(0.001, deadline - time.monotonic()), + check_alive=check_alive, + check_replay_alive=check_replay_alive, + profile_plan=plan, + client_started_ns=started_ns, + ) + while client.poll() is None: + check_alive() + if time.monotonic() >= deadline: + raise TimeoutError( + "AgentX client exceeded its deadline after profiler capture" + ) + time.sleep(0.1) + check_alive() + return client.returncode + finally: + try: + manifest = trace_dir / "capture.json" + if manifest.is_file(): + outcome["profile_capture"] = json.loads( + manifest.read_text(encoding="utf-8") + ) + finally: + _stop_owned(client) + + +def _require_free_port(port: int) -> None: + # Do not reuse an unrelated healthy endpoint. A new server owns every point. + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe: + probe.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + try: + probe.bind(("127.0.0.1", port)) + except OSError as exc: + raise RuntimeError( + f"AgentX server port {port} is unavailable; existing servers are not reused" + ) from exc + + +def _wait_health(server, url: str, deadline: float) -> None: + opener = urllib.request.build_opener(urllib.request.ProxyHandler({})) + while time.monotonic() < deadline: + if server.poll() is not None: + raise RuntimeError( + f"AgentX server exited before becoming healthy ({server.returncode})" + ) + try: + with opener.open( + url, timeout=min(2.0, max(0.1, deadline - time.monotonic())) + ) as response: + healthy = response.status == 200 + except (OSError, urllib.error.URLError): + healthy = False + if healthy: + time.sleep(0.1) + if server.poll() is not None: + raise RuntimeError( + f"AgentX server exited during its health check ({server.returncode})" + ) + return + time.sleep(0.1) + raise TimeoutError( + "AgentX server did not become healthy before its startup deadline" + ) + + +def _verify_model_metadata(spec: dict[str, Any]) -> None: + metadata = spec.get("model_metadata") + if not metadata: + return + source = str(metadata.get("source") or "") + if source.startswith("https://"): + revision = str(metadata.get("revision") or "") + parsed = urlparse(source) + if ( + parsed.hostname != "huggingface.co" + or parsed.port not in {None, 443} + or parsed.username + or parsed.password + or parsed.query + or parsed.fragment + or not re.fullmatch(r"[a-f0-9]{40}", revision) + or not parsed.path.endswith(f"/resolve/{revision}/config.json") + ): + raise ValueError( + "Remote AgentX metadata requires an immutable public HuggingFace revision" + ) + argv = spec["argv"] + if "--revision" not in argv or argv[argv.index("--revision") + 1 :][:1] != [ + revision + ]: + raise ValueError("AgentX server revision does not match its model metadata") + opener = urllib.request.build_opener( + urllib.request.ProxyHandler({}), _PublicModelRedirect() + ) + try: + with opener.open(source, timeout=30) as response: + contents = response.read(4 * 1024 * 1024 + 1) + except urllib.error.HTTPError as exc: + if exc.code in {401, 403}: + raise ValueError( + "Gated AgentX models require an already downloaded local MODEL_PATH/config.json" + ) from exc + raise + else: + path = Path(source) + if not path.is_absolute() or not path.is_file(): + raise ValueError( + "AgentX model metadata requires its resolved local config.json" + ) + contents = path.read_bytes() + if len(contents) > 4 * 1024 * 1024: + raise ValueError("AgentX model configuration exceeds the 4 MiB metadata limit") + if hashlib.sha256(contents).hexdigest() != metadata.get("sha256"): + raise ValueError("AgentX model metadata changed after recipe resolution") + native = native_context_length(json.loads(contents)) + maximum = metadata.get("max_model_len") + if ( + native != metadata.get("native_context_length") + or type(maximum) is not int + or not 0 < maximum <= native + ): + raise ValueError( + "AgentX model context limits differ from the resolved metadata" + ) + + +def _verify_recipe_inputs(spec: dict[str, Any]) -> None: + for name, expected in spec.get("source_files", {}).items(): + if hashlib.sha256(Path(name).read_bytes()).hexdigest() != expected: + raise ValueError(f"AgentX recipe input changed before setup: {name}") + + +def _aiperf_git(root: Path, *arguments: str) -> str: + path = root / "utils/aiperf" + completed = subprocess.run( + ["git", "-c", f"safe.directory={path}", "-C", str(path), *arguments], + capture_output=True, + text=True, + timeout=10, + ) + if completed.returncode: + raise ValueError("AgentX requires the resolved AIPerf Git checkout") + return completed.stdout.strip() + + +def _verify_aiperf_revision(root: Path, spec: dict[str, Any]) -> None: + expected = spec.get("aiperf_revision") + if expected is None: + return + if not isinstance(expected, str) or not re.fullmatch(r"[0-9a-f]{40}", expected): + raise ValueError("AgentX AIPerf revision must be a complete Git commit") + if ( + Path(_aiperf_git(root, "rev-parse", "--show-toplevel")).resolve() + != (root / "utils/aiperf").resolve() + ): + raise ValueError("AgentX AIPerf submodule is not initialized") + if _aiperf_git(root, "rev-parse", "HEAD") != expected: + raise ValueError("AgentX AIPerf revision changed after recipe resolution") + + +def _validate_spec(spec: Any) -> dict[str, Any]: + if ( + not isinstance(spec, dict) + or type(spec.get("version")) is not int + or spec["version"] != 1 + ): + raise ValueError( + "AgentX runtime requires a resolved server-launch-spec version 1" + ) + if spec.get("framework") not in {"sglang", "vllm"}: + raise ValueError("Magpie-owned AgentX supports SGLang and vLLM") + setup = spec.get("setup_commands", []) + if not isinstance(setup, list): + raise ValueError("AgentX setup_commands must be a list of argv token lists") + for argv in [spec.get("argv"), *setup]: + if ( + not isinstance(argv, list) + or not argv + or any(not isinstance(token, str) or "\0" in token for token in argv) + ): + raise ValueError( + "AgentX server/setup commands must be literal argv token lists" + ) + for key in ("env", "client_env"): + mapping = spec.get(key, {}) + if not isinstance(mapping, dict) or any( + not isinstance(name, str) + or not isinstance(value, str) + or "=" in name + or "\0" in name + value + for name, value in mapping.items() + ): + raise ValueError( + f"AgentX {key} must contain literal string environment values" + ) + port = spec.get("port") + if type(port) is not int or not 1 <= port <= 65535: + raise ValueError("AgentX server port must be between 1 and 65535") + health_path = spec.get("health_path", "/health") + if ( + not isinstance(health_path, str) + or not health_path.startswith("/") + or health_path.startswith("//") + ): + raise ValueError("AgentX health_path must be a local absolute URL path") + return spec + + +def _client_command(root: Path) -> list[str]: + current = root / "benchmarks" / "srt_agentic.sh" + if current.is_file(): + return ["bash", str(current)] + if (root / "benchmarks" / "benchmark_lib.sh").is_file(): + adapter = Path(__file__).parents[2] / "scripts" / "agentx" / "client_legacy.sh" + return ["bash", str(adapter)] + raise ValueError("InferenceX project has no supported AgentX client entrypoint") + + +def _client_environment(request, spec, root, workspace) -> dict[str, str]: + env = { + key: value + for key, value in os.environ.items() + if not key.startswith(("AIPERF_", "AGENTIC_", "AGENTX_", "SRT_", "SRTCTL_")) + and key not in {"CONC_LIST", "BASH_ENV", "ENV", "SHELLOPTS", "BASHOPTS"} + } + server_only = set(spec.get("env", {})) - { + "ROCR_VISIBLE_DEVICES", + "HIP_VISIBLE_DEVICES", + "CUDA_VISIBLE_DEVICES", + } + env.update( + { + key: value + for key, value in request["client_env"].items() + if key not in server_only + } + ) + env.update(spec.get("client_env", {})) + # A candidate framework overlay must not change the measurement client. + env["PYTHONPATH"] = str(root) + env.pop("PYTHONHOME", None) + env.update( + { + "INFMAX_CONTAINER_WORKSPACE": str(root), + "RESULT_DIR": str(workspace), + "AGENTIC_OUTPUT_DIR": str(workspace), + "RESULT_FILENAME": "inferencex_result", + "SERVER_LOG": str(workspace / "server.log"), + "AIPERF_RUNTIME_DIR": str(workspace / ".agentx-runtime"), + "AIPERF_SERVER_URL": f"http://127.0.0.1:{spec['port']}", + "AIPERF_SERVER_METRICS_URLS": f"http://127.0.0.1:{spec['port']}/metrics", + "PORT": str(spec["port"]), + "IS_MULTINODE": "false", + "IS_AGENTIC": "1", + "SCENARIO_TYPE": "agentic-coding", + "EVAL_ONLY": "false", + "RUN_EVAL": "false", + } + ) + # Client adapters must not alter or delete the server receipt. + for name in tuple(env): + if name.startswith(("AGENTX_LAUNCH_", "SRT_", "SRTCTL_")) or name in { + "AGENTX_SERVER_LAUNCH_FILE", + "CONC_LIST", + }: + env.pop(name) + return env + + +def _preflight_client_environment(env: dict[str, str]) -> None: + """Validate the pinned native client's single-node replay contract. + + These are the check_env_vars inputs in srt_agentic.sh and its dependency, + replay, and power helpers at InferenceX 408c015. Keep this contract in sync + when changing the client pin; do not execute the upstream shell to check it. + """ + required = { + "RESULT_DIR", + "EVAL_ONLY", + "MODEL", + "MODEL_PREFIX", + "FRAMEWORK", + "PRECISION", + "CONC", + "RESULT_FILENAME", + "DURATION", + "PORT", + "INFMAX_CONTAINER_WORKSPACE", + "AIPERF_PYTHON_VERSION", + "AIPERF_FAILED_REQUEST_THRESHOLD", + "AIPERF_LIVE_FAILED_REQUEST_THRESHOLD", + "AIPERF_TRACE_IDLE_GAP_CAP_SECONDS", + "AGENTIC_WARMUP_GRACE_PERIOD", + "AIPERF_DATASET_WEKA_LIVE_ASSISTANT_RESPONSES", + "AIPERF_DYNAMO_SESSION_TIMEOUT_SECONDS", + "AIPERF_EXPERIMENTAL_FAST", + "AIPERF_HTTP_X_DYNAMO_SESSION_ID_FROM_CORRELATION_ID", + "AIPERF_UNSAFE_OVERRIDE", + "AIPERF_USE_DYNAMO_CONV_AWARE_ROUTING", + "AIPERF_WARMUP_REQUESTS_PER_LANE", + "ENABLE_AGENTX_POWER", + "IS_MULTINODE", + "REQUIRE_POWER", + } + if env.get("ENABLE_AGENTX_POWER", "").lower() in {"1", "true", "yes"}: + required.update({"TP", "PP_SIZE", "PCP_SIZE"}) + missing = sorted(name for name in required if not env.get(name)) + if missing: + raise ValueError( + "AgentX client environment is missing required values before server startup: " + + ", ".join(missing) + ) + + +def _execute_local(request: dict[str, Any]) -> dict[str, Any]: + started = time.monotonic() + workspace = Path(request["workspace"]) + root = Path(request["inferencex_path"]) + workspace.mkdir(parents=True, exist_ok=True) + (workspace / "agentx_server_launch.json").unlink(missing_ok=True) + server = None + outcome = {"success": False, "errors": [], "execution_time": 0.0} + try: + spec = _validate_spec(request["server_spec"]) + _verify_recipe_inputs(spec) + _verify_aiperf_revision(root, spec) + command = _client_command(root) + if "recipe_variant" in spec and Path(command[-1]).name == "srt_agentic.sh": + _preflight_client_environment( + _client_environment(request, spec, root, workspace) + ) + if request.get("profile"): + # Reserve this invocation before startup: frameworks may write + # CUDA-graph warmup traces while the service is coming online. + Path(spec["torch_profiler"]["trace_dir"]).mkdir( + parents=True, exist_ok=False + ) + _require_free_port(spec["port"]) + deadline = time.monotonic() + request["ready_timeout"] + # Acceptance simulation must come from the resolved recipe, including + # when this recipe intentionally has no simulation settings. + server_env = { + key: value + for key, value in os.environ.items() + if not key.startswith("SGLANG_SIMULATE_ACC_") + } + server_env.update(request.get("gpu_env", {})) + server_env.update(spec.get("env", {})) + for key in ("BASH_ENV", "ENV", "SHELLOPTS", "BASHOPTS"): + server_env.pop(key, None) + with (workspace / "agentx_setup.log").open("w", encoding="utf-8") as setup_log: + for setup in spec.get("setup_commands", []): + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError( + "AgentX preparation exceeded its startup deadline" + ) + rc = _run_owned( + setup, env=server_env, cwd=root, output=setup_log, timeout=remaining + ) + if rc: + raise RuntimeError( + f"AgentX server preparation failed with exit code {rc}" + ) + _verify_aiperf_revision(root, spec) + _verify_model_metadata(spec) + argv, server_env, _evidence = prepare_server_launch( + spec["argv"], + server_env, + request.get("overrides"), + spec["framework"], + workspace, + server_spec=spec, + ) + if (request.get("profile") or {}).get("detailed_annotations"): + if __package__: + from .agentx_profile_capabilities import probe_profile_capabilities + else: + from agentx_profile_capabilities import probe_profile_capabilities + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError("AgentX annotation probe exceeded startup deadline") + capabilities = probe_profile_capabilities( + spec["framework"], + argv, + server_env, + root, + remaining, + ) + spec = profile_server_spec( + _validate_spec(request.get("base_server_spec")), + request["profile"], + workspace, + spec["torch_profiler"]["capture_id"], + capabilities=capabilities, + launch_overrides=request.get("overrides"), + ) + server_env.update(spec.get("env", {})) + argv, server_env, _evidence = prepare_server_launch( + spec["argv"], + server_env, + request.get("overrides"), + spec["framework"], + workspace, + server_spec=spec, + ) + with (workspace / "server.log").open("w", encoding="utf-8") as server_log: + server = _spawn(argv, env=server_env, cwd=root, output=server_log) + _wait_health( + server, + f"http://127.0.0.1:{spec['port']}{spec.get('health_path', '/health')}", + deadline, + ) + with (workspace / "agentx_client.log").open( + "w", encoding="utf-8" + ) as client_log: + if request.get("profile"): + rc = _run_profiled_client( + request, spec, server=server, output=client_log, outcome=outcome + ) + else: + rc = _run_owned( + command, + env=_client_environment(request, spec, root, workspace), + cwd=root, + output=client_log, + timeout=request["client_timeout"], + server=server, + ) + if rc: + raise RuntimeError( + f"AgentX benchmark client failed with exit code {rc}" + ) + outcome["success"] = True + except ( + OSError, + ValueError, + RuntimeError, + TimeoutError, + subprocess.SubprocessError, + http.client.HTTPException, + ) as exc: + outcome["errors"].append(str(exc)) + finally: + if server is not None: + try: + _stop_owned(server) + except (OSError, subprocess.SubprocessError) as exc: + outcome["success"] = False + outcome["errors"].append(f"AgentX server cleanup failed: {exc}") + outcome["execution_time"] = time.monotonic() - started + return outcome + + +def _docker_command( + request, *, image: str, name: str, request_path: Path, runner_type: str +) -> list[str]: + command = [ + "docker", + "run", + "--init", + "--name", + name, + "--network=host", + "--ipc=host", + "--shm-size=16g", + ] + if runner_type.lower().startswith(("mi", "gfx", "radeon")): + command += [ + "--device=/dev/kfd", + "--device=/dev/dri", + "--group-add=video", + "--cap-add=SYS_PTRACE", + "--security-opt=seccomp=unconfined", + ] + else: + command += ["--gpus=all"] + package = Path(__file__).resolve().parents[2] + mounts = { + str(package): "ro", + request["inferencex_path"]: "rw", + request["workspace"]: "rw", + } + if request["server_spec"].get("aiperf_revision") is not None: + root = Path(request["inferencex_path"]) + _verify_aiperf_revision(root, request["server_spec"]) + # A submodule's .git file points outside the mounted project directory. + # Preserve that exact path without exposing the parent checkout. + git_directory = Path(_aiperf_git(root, "rev-parse", "--absolute-git-dir")) + mounts[str(git_directory.resolve())] = "ro" + model_path = request["client_env"].get( + "MODEL_PATH", request["client_env"].get("MODEL", "") + ) + if model_path and Path(model_path).is_dir(): + mounts[str(Path(model_path).resolve())] = "ro" + cache = request.get("hf_cache_path") + if cache and Path(cache).is_dir(): + mounts[str(Path(cache).resolve())] = "rw" + command += ["--env", f"HF_HOME={Path(cache).resolve()}"] + for source, mode in mounts.items(): + if ":" in source or "\n" in source: + raise ValueError( + "AgentX Docker mount paths cannot contain ':' or a newline" + ) + command += ["--volume", f"{source}:{source}:{mode}"] + for key, value in request.get("gpu_env", {}).items(): + command += ["--env", f"{key}={value}"] + if os.environ.get("HF_TOKEN"): + command += ["--env", "HF_TOKEN"] + command += [ + "--entrypoint", + "python3", + "--", + image, + str(Path(__file__).resolve()), + "--worker", + str(request_path), + ] + return command + + +def _execute_docker(request, *, image: str, runner_type: str) -> dict[str, Any]: + workspace = Path(request["workspace"]) + request_path = workspace / "agentx_runtime_request.json" + result_path = workspace / "agentx_runtime_result.json" + result_path.unlink(missing_ok=True) + request_path.write_text(json.dumps(request), encoding="utf-8") + name = f"magpie-agentx-{uuid.uuid4().hex}" + outcome = {"success": False, "errors": []} + started = time.monotonic() + try: + command = _docker_command( + request, + image=image, + name=name, + request_path=request_path, + runner_type=runner_type, + ) + with (workspace / "agentx_container.log").open("w", encoding="utf-8") as output: + rc = _run_owned( + command, + env=dict(os.environ), + cwd=workspace, + output=output, + timeout=( + request["ready_timeout"] + + request["client_timeout"] + + 60 + + (request.get("profile") or {}).get("capture_timeout_seconds", 0) + + (request.get("profile") or {}).get("flush_timeout_seconds", 0) + ), + ) + if result_path.is_file(): + outcome = json.loads(result_path.read_text(encoding="utf-8")) + if rc: + outcome["success"] = False + outcome["errors"].append(f"AgentX container exited with code {rc}") + elif not result_path.is_file(): + outcome["errors"].append("AgentX container produced no runtime result") + except ( + OSError, + ValueError, + RuntimeError, + TimeoutError, + subprocess.SubprocessError, + ) as exc: + outcome["errors"].append(str(exc)) + finally: + try: + removed = subprocess.run( + ["docker", "rm", "--force", name], + capture_output=True, + text=True, + timeout=30, + ) + if removed.returncode and "No such container" not in removed.stderr: + raise RuntimeError(removed.stderr.strip()) + except (OSError, RuntimeError, subprocess.SubprocessError) as exc: + outcome["success"] = False + outcome["errors"].append(f"AgentX container cleanup failed: {exc}") + outcome["execution_time"] = time.monotonic() - started + return outcome + + +def execute_agentx( + config, workspace: Path, runner_type: str, *, docker_image: str | None = None +): + """Run one fresh server/client point and return Magpie's normal execution tuple.""" + from .result import BenchmarkResult + + workspace = workspace.resolve() + workspace.mkdir(parents=True, exist_ok=True) + resolved = config.agentx.resolved if config.agentx else {} + spec = _validate_spec(resolved.get("server-launch-spec")) + if spec["framework"] != config.framework: + raise ValueError("AgentX server framework differs from the benchmark") + lifecycle = config.server_lifecycle + request = { + "workspace": str(workspace), + "inferencex_path": str(Path(config.inferencex_path).resolve()), + "server_spec": spec, + "overrides": config.agentx.launch_overrides, + "client_env": config.get_env_vars(), + "hf_cache_path": config.hf_cache_path + or os.environ.get("HF_HOME") + or str(Path.home() / ".cache/huggingface"), + "ready_timeout": float( + lifecycle.server_ready_timeout_s + if lifecycle + else spec.get("health_timeout_seconds", 2700) + ), + "client_timeout": float(config.timeout_seconds), + "gpu_env": { + key: str(value) + for key, value in config.envs.items() + if key + in {"ROCR_VISIBLE_DEVICES", "HIP_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES"} + }, + } + settings = profile_settings(config) + if settings is not None: + # Reject an impossible first capture before starting a costly server. + # The worker recomputes this plan from its final client environment. + profile_plan( + settings, + _client_environment( + request, spec, Path(request["inferencex_path"]), workspace + ), + ) + request["profile"] = settings + request["base_server_spec"] = spec + request["server_spec"] = profile_server_spec( + spec, settings, workspace, uuid.uuid4().hex + ) + ranks = 1 + for entry_name, env_name in ( + ("tp", "TP"), + ("pp", "PP_SIZE"), + ("pcp-size", "PCP_SIZE"), + ): + value = resolved.get(entry_name, config.envs.get(env_name, 1)) + if ( + not isinstance(value, (str, int)) + or isinstance(value, bool) + or int(value) <= 0 + ): + raise ValueError( + "AgentX profiling requires positive resolved GPU topology" + ) + ranks *= int(value) + request["profile_ranks"] = ranks + if any( + not math.isfinite(request[key]) or request[key] <= 0 + for key in ("ready_timeout", "client_timeout") + ): + raise ValueError( + "AgentX startup and client timeouts must be positive finite durations" + ) + if config.run_mode == "docker": + image = docker_image or config.docker_image + if not image: + raise ValueError("AgentX Docker execution requires a pinned image") + outcome = _execute_docker(request, image=image, runner_type=runner_type) + elif config.run_mode == "local": + outcome = _execute_local(request) + else: + raise ValueError("Magpie-owned AgentX supports local and Docker execution") + result = BenchmarkResult( + success=bool(outcome["success"]), + framework=config.framework, + model=config.model, + scenario="agentx", + workspace_dir=str(workspace), + execution_time=outcome["execution_time"], + errors=outcome["errors"], + ) + if settings is not None: + result.profiling_enabled = True + result.benchmark_valid = False + result.publishable = False + result.agentx_metrics = { + "diagnostic_only": True, + "profile_capture": outcome.get("profile_capture", {"status": "failed"}), + } + log = workspace / "agentx_client.log" + stdout = ( + log.read_text(encoding="utf-8", errors="replace")[-131072:] + if log.is_file() + else "" + ) + return result, stdout, "\n".join(outcome["errors"]) + + +def _worker_main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--worker", type=Path, required=True) + args = parser.parse_args() + request = json.loads(args.worker.read_text(encoding="utf-8")) + signal.signal(signal.SIGTERM, _signal_interrupt) + signal.signal(signal.SIGINT, _signal_interrupt) + outcome = _execute_local(request) + (Path(request["workspace"]) / "agentx_runtime_result.json").write_text( + json.dumps(outcome), encoding="utf-8" + ) + return 0 if outcome["success"] else 1 + + +if __name__ == "__main__": + raise SystemExit(_worker_main()) diff --git a/Magpie/modes/benchmark/benchmarker.py b/Magpie/modes/benchmark/benchmarker.py index d85d3c9..4bae871 100644 --- a/Magpie/modes/benchmark/benchmarker.py +++ b/Magpie/modes/benchmark/benchmarker.py @@ -12,6 +12,7 @@ import json import logging import os +import shlex import shutil import signal import subprocess @@ -33,6 +34,13 @@ ensure_agentx_dependencies, resolve_agentx_recipe, ) +from .agentx_launch import ( + apply_launch_args, + launch_environment, + read_launch_evidence, + validate_overrides, +) +from .agentx_runtime import execute_agentx from .config import BenchmarkConfig from .image_selector import ImageSelector from .inferencex import ensure_inferencex_available @@ -116,6 +124,49 @@ def __init__( ) self._task_id: Optional[str] = None + @property + def _managed_agentx(self) -> Optional[Dict[str, Any]]: + """The persisted recipe decides who owns the server, not script names.""" + if not self.config.is_agentx or self.config.agentx is None: + return None + return (self.config.agentx.resolved or {}).get("server-launch-spec") + + def _normalize_agentx_extra_args(self) -> None: + """Consume legacy extra-args once into the literal candidate contract.""" + if not self._managed_agentx: + return + key = f"EXTRA_{self.config.framework.upper()}_ARGS" + present = [ + name + for name in self.config.envs + if name.upper().startswith("EXTRA_") and name.upper().endswith("_ARGS") + ] + if any(name.upper() != key for name in present): + raise ValueError( + f"Magpie-owned AgentX only accepts {key} for this framework" + ) + if not present: + return + if len(present) != 1: + raise ValueError(f"Duplicate AgentX {key} environment settings") + name = present[0] + raw = self.config.envs[name] + if not isinstance(raw, str): + raise ValueError(f"{key} must be a shell-quoted argument string") + overrides = validate_overrides( + self.config.agentx.launch_overrides or {"version": 1} + ) + overrides["append_args"] = [*overrides["append_args"], *shlex.split(raw)] + # Reject protocol changes before dependency setup; the runtime repeats + # these checks immediately before starting the server. + overrides = validate_overrides(overrides) + apply_launch_args( + self._managed_agentx["argv"], overrides, self.config.framework + ) + self.config.agentx.launch_overrides = overrides + self.config.envs.pop(name) + self._managed_agentx.get("client_env", {}).pop(name, None) + def run(self, task_id: Optional[str] = None) -> BenchmarkResult: """ Run benchmark. @@ -154,8 +205,8 @@ def run(self, task_id: Optional[str] = None) -> BenchmarkResult: ) return result - # AgentX execution stays inside InferenceX. Magpie selects the existing - # launcher/recipe and forwards its resolved environment. + # Resolve one workload. New recipes separate Magpie-owned serving from + # the upstream client; explicit legacy checkouts retain their launcher. runner_type: Optional[str] = None if self.config.is_agentx: try: @@ -165,7 +216,13 @@ def run(self, task_id: Optional[str] = None) -> BenchmarkResult: self.config.inferencex_path, runner_type=runner_type, ) + self._normalize_agentx_extra_args() ensure_agentx_dependencies(self.config.inferencex_path) + if self.config.is_server_lifecycle and not self._managed_agentx: + raise ValueError( + "Legacy AgentX launchers cannot use server_lifecycle; " + "use the pinned InferenceX client layout for Magpie-managed serving" + ) self.workspace_mgr.framework = self.config.framework except (OSError, RuntimeError, ValueError) as e: result = BenchmarkResult( @@ -182,6 +239,7 @@ def run(self, task_id: Optional[str] = None) -> BenchmarkResult: # the already-running server's physical devices. lifecycle_server_healthy = ( self.config.is_server_lifecycle + and not self.config.is_agentx and self._reuse_http_healthy(self._reuse_benchmark_port()) ) skip_idle_gpu_for_reuse = bool(lifecycle_server_healthy) @@ -210,7 +268,8 @@ def run(self, task_id: Optional[str] = None) -> BenchmarkResult: return result # 1. Copy Magpie generic scripts to InferenceX/benchmarks/ - self._prepare_benchmark_scripts() + if not self._managed_agentx: + self._prepare_benchmark_scripts() # 2. Determine runner type from GPU runner_type = runner_type or self._get_runner_type() @@ -344,7 +403,14 @@ def run(self, task_id: Optional[str] = None) -> BenchmarkResult: gpu_monitor = None # 5-7. Build and execute (Docker or local) - if self.config.is_local: + if self._managed_agentx: + result, stdout, stderr = execute_agentx( + self.config, + workspace, + runner_type, + docker_image=(None if self.config.is_local else self._select_image()), + ) + elif self.config.is_local: symlink = self._create_workspace_symlink(workspace) try: if self.config.is_server_lifecycle: @@ -399,6 +465,8 @@ def run(self, task_id: Optional[str] = None) -> BenchmarkResult: gpu_monitor_stats = gpu_monitor.stop() logger.info(f"GPU monitor stopped: {gpu_monitor_stats.sample_count} samples collected") + profile_capture = (result.agentx_metrics or {}).get("profile_capture") + # 8. Collect results result.workspace_dir = str(workspace) result.execution_time = time.time() - start_time @@ -526,11 +594,69 @@ def run(self, task_id: Optional[str] = None) -> BenchmarkResult: "Benchmark produced no valid throughput/latency metrics" ) - # Parse torch trace if available - if self.config.profiler.torch_profiler.enabled: + if self.config.is_agentx and ( + self._managed_agentx or self.config.agentx.launch_overrides is not None + ): + try: + evidence = read_launch_evidence(self.config, workspace) + if result.agentx_metrics is None: + result.agentx_metrics = {} + result.agentx_metrics["server_launch"] = evidence + if self._managed_agentx: + result.agentx_metrics["execution_owner"] = "magpie" + metadata = self._managed_agentx.get("model_metadata") + if metadata: + result.agentx_metrics["model_context"] = metadata + except (OSError, ValueError) as exc: + result.success = False + result.benchmark_valid = False + result.publishable = False + result.errors.append(f"AgentX server launch evidence rejected: {exc}") + + if self.config.is_agentx and self.config.profiler.torch_profiler.enabled: + # A successful diagnostic capture is never canonical KEEP evidence, + # even when the upstream replay's own validity checks pass. + result.benchmark_valid = False + result.publishable = False + if result.agentx_metrics is None: + result.agentx_metrics = {} + result.agentx_metrics["diagnostic_only"] = True + result.agentx_metrics["profile_capture"] = profile_capture or { + "status": "failed" + } + if not profile_capture or profile_capture.get("status") != "complete": + result.success = False + result.errors.append("AgentX torch profiler capture did not complete") + + repeated_profile = ( + self.config.is_agentx + and profile_capture + and isinstance(profile_capture.get("profiles"), list) + ) + if repeated_profile and result.success: + self._analyze_agentx_profile_series(result, profile_capture, workspace) + + # Parse each repeated capture separately so ranks from different time + # windows are never combined into a single TraceLens collective report. + if self.config.profiler.torch_profiler.enabled and not repeated_profile: torch_trace_dir = workspace / "torch_trace" + if ( + self.config.is_agentx + and profile_capture + and profile_capture.get("capture_id") + ): + # Analyze only this invocation, never a prior diagnostic capture + # left in a reused workspace. + torch_trace_dir /= profile_capture["capture_id"] # 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 = ( + [ + path for path in torch_trace_dir.rglob("*") + if path.is_file() and path.name != "capture.json" + and path.name.endswith((".json", ".json.gz")) + ] + if torch_trace_dir.is_dir() else [] + ) has_traces = len(trace_files) > 0 if not result.success or not has_traces: @@ -587,6 +713,44 @@ def run(self, task_id: Optional[str] = None) -> BenchmarkResult: return result + def _analyze_agentx_profile_series( + self, result: BenchmarkResult, capture: Dict[str, Any], workspace: Path + ) -> None: + analyses = [] + tracelens_profiles = [] + for profile in capture["profiles"]: + if profile.get("status") != "complete": + continue + trace_dir = ( + workspace + / "torch_trace" + / capture["capture_id"] + / f"profile_{profile['profile_index']:03d}" + ) + kernels = ResultParser.parse_torch_trace(trace_dir) + analyses.append( + { + "profile_index": profile["profile_index"], + "trace_dir": str(trace_dir), + "kernel_summary": [kernel.to_dict() for kernel in kernels], + "top_bottlenecks": [kernel.name for kernel in kernels[:10]], + } + ) + # The existing top-level fields show the latest capture; all + # individual windows remain available in profile_analyses. + result.kernel_summary = kernels + result.top_bottlenecks = analyses[-1]["top_bottlenecks"] + if self.config.profiler.tracelens.enabled: + output_dir = workspace / f"profile_{profile['profile_index']:03d}" + output_dir.mkdir(parents=True, exist_ok=True) + analysis = self._run_tracelens_analysis(trace_dir, output_dir) + tracelens_profiles.append( + {"profile_index": profile["profile_index"], "analysis": analysis} + ) + result.agentx_metrics["profile_analyses"] = analyses + if self.config.profiler.tracelens.enabled: + result.tracelens_analysis = {"enabled": True, "profiles": tracelens_profiles} + def _apply_gpu_selection(self) -> None: """Resolve idle GPU(s) and inject the selection into config.envs. @@ -861,7 +1025,9 @@ def _build_docker_command( cmd.extend(["-v", f"{workspace}:/workspace"]) # Environment variables + benchmark_script = self._get_benchmark_script(runner_type) env_vars = self.config.get_env_vars() + self._set_generic_eval_defaults(env_vars, benchmark_script) env_vars["RESULT_FILENAME"] = "inferencex_result" env_vars["RESULT_DIR"] = "/workspace" env_vars["RUNNER_TYPE"] = runner_type @@ -870,6 +1036,7 @@ def _build_docker_command( env_vars["INFMAX_CONTAINER_WORKSPACE"] = "/opt/InferenceX" env_vars["AGENTIC_OUTPUT_DIR"] = "/workspace" env_vars["AIPERF_RUNTIME_DIR"] = "/tmp/inferencex-agentx" + env_vars.update(launch_environment(self.config, workspace, docker=True)) if phase == "server": env_vars["MAGPIE_SERVER_PID_FILE"] = "/workspace/reuse_server_spawn.pid" env_vars["MAGPIE_KEEP_CONTAINER_ALIVE"] = "1" @@ -881,7 +1048,6 @@ def _build_docker_command( env_vars["SGLANG_TORCH_PROFILER_DIR"] = "/workspace/torch_trace" env_vars["ATOM_TORCH_PROFILER_DIR"] = "/workspace/torch_trace" - # HuggingFace token from environment hf_token = os.environ.get("HF_TOKEN", "") if hf_token: @@ -897,13 +1063,10 @@ def _build_docker_command( cmd.extend(["--entrypoint", "/bin/bash"]) cmd.append(docker_image) - # Build the benchmark command - benchmark_script = self._get_benchmark_script(runner_type) - # With --entrypoint /bin/bash, pass -c as first arg cmd.extend([ "-c", - f"cd /opt/InferenceX && bash {benchmark_script}" + f"cd /opt/InferenceX && bash {shlex.quote(benchmark_script)}" ]) return cmd @@ -937,6 +1100,21 @@ def _remove_workspace_symlink(symlink: Optional[Path]) -> None: except OSError as e: logger.warning(f"Could not remove symlink {symlink}: {e}") + def _set_generic_eval_defaults( + self, env: dict[str, str], benchmark_script: str + ) -> None: + """Supply workflow-owned inputs required by InferenceX's local eval. + + Generic scripts use the upstream eval functions without the upstream + workflow. Preserve explicit values, keep custom scripts and AgentX + unchanged, and never fetch a host credential for a Docker benchmark. + """ + if self.config.is_agentx or Path(benchmark_script).name not in MAGPIE_BUILTIN_SCRIPTS: + return + env.setdefault("EVAL_ONLY", "false") + env.setdefault("IS_AGENTIC", "0") + env.setdefault("OPENAI_API_KEY", "EMPTY") + def _build_local_command( self, workspace: Path, @@ -961,7 +1139,7 @@ def _build_local_command( cmd = [ "bash", "-c", - f"cd {inferencex_path} && bash {benchmark_script}", + f"cd {shlex.quote(inferencex_path)} && bash {shlex.quote(benchmark_script)}", ] env = os.environ.copy() @@ -977,6 +1155,7 @@ def _build_local_command( env_vars["AIPERF_RUNTIME_DIR"] = str( workspace / ".agentx-runtime" ) + env_vars.update(launch_environment(self.config, workspace)) if phase == "server" and server_pid_file is not None: env_vars["MAGPIE_SERVER_PID_FILE"] = str(server_pid_file) @@ -997,6 +1176,7 @@ def _build_local_command( for key, value in env_vars.items(): env[key] = str(value) + self._set_generic_eval_defaults(env, benchmark_script) return cmd, env def _execute_local_benchmark_with_reuse( @@ -2310,6 +2490,13 @@ def _run_tracelens_analysis( # Get number of ranks from TP config num_ranks = int(self.config.envs.get("TP", 8)) + if self.config.is_agentx: + resolved = self.config.agentx.resolved + num_ranks = 1 + for key, env in ( + ("tp", "TP"), ("pp", "PP_SIZE"), ("pcp-size", "PCP_SIZE") + ): + num_ranks *= int(resolved.get(key, self.config.envs.get(env, 1))) results = analyzer.analyze( trace_dir=torch_trace_dir, diff --git a/Magpie/modes/benchmark/config.py b/Magpie/modes/benchmark/config.py index 63c0476..7efb8e5 100644 --- a/Magpie/modes/benchmark/config.py +++ b/Magpie/modes/benchmark/config.py @@ -8,10 +8,13 @@ """ import logging +import math from dataclasses import dataclass, field from enum import Enum from typing import Any, Dict, List, Optional +from .agentx_launch import validate_overrides + logger = logging.getLogger(__name__) @@ -62,18 +65,76 @@ class TorchProfilerConfig: Attributes: enabled: Whether torch_profiler is enabled (default: True) + num_steps: Framework steps captured in each AgentX diagnostic profile. + capture_timeout_seconds: Maximum wait to start capture and see its first trace. + flush_timeout_seconds: Maximum wait for all AgentX rank traces to finish. + num_profiles: Number of sequential captures in the same AgentX replay. + interval_seconds: Delay after complete trace flushing before the next capture. + start_seconds: Delay from measurement start before the first capture. + detailed_annotations: Require framework annotations and graph shape traces. """ enabled: bool = True + num_steps: int = 20 + capture_timeout_seconds: float = 300.0 + flush_timeout_seconds: float = 1800.0 + num_profiles: int = 1 + interval_seconds: float = 200.0 + start_seconds: float = 0.0 + detailed_annotations: bool = False + + def __post_init__(self): + """Validate bounded AgentX diagnostic capture settings.""" + if type(self.detailed_annotations) is not bool: + raise ValueError("torch_profiler.detailed_annotations must be a boolean") + for name in ("num_steps", "num_profiles"): + value = getattr(self, name) + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise ValueError(f"torch_profiler.{name} must be a positive integer") + for name in ( + "capture_timeout_seconds", "flush_timeout_seconds", "interval_seconds", + "start_seconds", + ): + value = getattr(self, name) + allow_zero = name in {"interval_seconds", "start_seconds"} + requirement = "non-negative" if allow_zero else "positive" + error = f"torch_profiler.{name} must be a finite {requirement} number" + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise ValueError(error) + try: + value = float(value) + except OverflowError as exc: + raise ValueError(error) from exc + if not math.isfinite(value) or value < 0 or (value == 0 and not allow_zero): + raise ValueError(error) + setattr(self, name, value) def to_dict(self) -> Dict[str, Any]: """Convert to dictionary.""" - return {"enabled": self.enabled} + return { + "enabled": self.enabled, + "num_steps": self.num_steps, + "capture_timeout_seconds": self.capture_timeout_seconds, + "flush_timeout_seconds": self.flush_timeout_seconds, + "num_profiles": self.num_profiles, + "interval_seconds": self.interval_seconds, + "start_seconds": self.start_seconds, + "detailed_annotations": self.detailed_annotations, + } @classmethod def from_dict(cls, data: Dict[str, Any]) -> "TorchProfilerConfig": """Create from dictionary.""" - return cls(enabled=data.get("enabled", True)) + return cls( + enabled=data.get("enabled", True), + num_steps=data.get("num_steps", 20), + capture_timeout_seconds=data.get("capture_timeout_seconds", 300.0), + flush_timeout_seconds=data.get("flush_timeout_seconds", 1800.0), + num_profiles=data.get("num_profiles", 1), + interval_seconds=data.get("interval_seconds", 200.0), + start_seconds=data.get("start_seconds", 0.0), + detailed_annotations=data.get("detailed_annotations", False), + ) @dataclass @@ -668,7 +729,7 @@ class AgentXConfig: """Configuration for an InferenceX AgentX trace replay. The workload switch is simply ``agentx: enable``; ``docker_image`` and - ``benchmark_script`` remain explicit top-level benchmark pins. ``recipe`` + ``benchmark_script`` can override pins resolved from InferenceX. ``recipe`` and ``selector`` are escape hatches for configurations that cannot be selected unambiguously from model, framework, precision, GPU, and concurrency. @@ -681,6 +742,7 @@ class AgentXConfig: config_file: Optional[str] = None selector: Dict[str, Any] = field(default_factory=dict) failed_request_threshold: float = 0.10 + launch_overrides: Optional[Dict[str, Any]] = None resolved: Optional[Dict[str, Any]] = field(default=None, repr=False) def __post_init__(self) -> None: @@ -693,6 +755,8 @@ def __post_init__(self) -> None: raise ValueError("agentx.mode must be 'canonical' or 'fast'") if not 0.0 <= self.failed_request_threshold <= 1.0: raise ValueError("agentx.failed_request_threshold must be between 0 and 1") + if self.launch_overrides is not None: + self.launch_overrides = validate_overrides(self.launch_overrides) def to_dict(self) -> Dict[str, Any]: result: Dict[str, Any] = { @@ -710,6 +774,8 @@ def to_dict(self) -> Dict[str, Any]: result["selector"] = self.selector if self.resolved is not None: result["resolved"] = self.resolved + if self.launch_overrides is not None: + result["launch_overrides"] = self.launch_overrides return result @classmethod @@ -737,6 +803,7 @@ def from_dict(cls, data: Dict[str, Any]) -> "AgentXConfig": config_file=data.get("config_file"), selector=dict(data.get("selector") or {}), failed_request_threshold=float(data.get("failed_request_threshold", 0.10)), + launch_overrides=data.get("launch_overrides"), ) @classmethod @@ -793,7 +860,7 @@ class BenchmarkConfig: envs: Dict[str, Any] = field(default_factory=dict) # Profiler configuration - profiler: ProfilerConfig = field(default_factory=ProfilerConfig) + profiler: Optional[ProfilerConfig] = None # Docker/execution settings docker_image: Optional[str] = None @@ -886,9 +953,15 @@ def __post_init__(self): "RANDOM_RANGE_RATIO": 0.5, } - # Convert profiler dict to ProfilerConfig if needed - if isinstance(self.profiler, dict): - self.profiler = ProfilerConfig.from_dict(self.profiler) + # Omitted AgentX sections cannot implicitly enable diagnostics. + if self.profiler is None or isinstance(self.profiler, dict): + profiler_data = dict(self.profiler or {}) + if self.is_agentx: + for name in ("torch_profiler", "gpu_monitor"): + section = dict(profiler_data.get(name) or {}) + section.setdefault("enabled", False) + profiler_data[name] = section + self.profiler = ProfilerConfig.from_dict(profiler_data) # Convert gap_analysis dict to GapAnalysisConfig if needed if isinstance(self.gap_analysis, dict): @@ -917,25 +990,19 @@ def __post_init__(self): "Magpie AgentX v1 supports run_mode='local' or 'docker'; " "Ray and multi-node execution are not supported yet" ) - if self.is_server_lifecycle: - raise ValueError("AgentX cannot be combined with server_lifecycle") - if not self.benchmark_script: - raise ValueError( - "AgentX requires benchmark_script so the InferenceX " - "launcher is explicitly pinned" - ) - if self.run_mode == "docker" and not self.docker_image: + if self.is_server_lifecycle and ( + not self.server_lifecycle.cleanup or self.server_lifecycle.force_reuse + ): raise ValueError( - "Docker AgentX requires docker_image so the runtime is " - "explicitly pinned" + "AgentX cannot be combined with server_lifecycle reuse: " + "set cleanup=true and force_reuse=false for a fresh server per point" ) incompatible = [] - if self.profiler.torch_profiler.enabled: - incompatible.append("torch_profiler") if self.profiler.system_profiler.enabled: incompatible.append("system_profiler") - if self.profiler.tracelens.enabled: - incompatible.append("tracelens") + tracelens = self.profiler.tracelens + if tracelens.enabled and tracelens.is_inference_mode: + incompatible.append("tracelens.analysis_mode=inference") if self.gap_analysis.enabled: incompatible.append("gap_analysis") if incompatible: @@ -943,6 +1010,12 @@ def __post_init__(self): "AgentX v1 does not support Magpie profiling options: " + ", ".join(incompatible) ) + if tracelens.enabled and not self.profiler.torch_profiler.enabled: + raise ValueError( + "AgentX TraceLens analysis requires " + "profiler.torch_profiler.enabled=true and " + "profiler.tracelens.analysis_mode=pytorch" + ) if self.is_server_lifecycle: if self.run_mode not in {"local", "docker"}: @@ -1079,15 +1152,7 @@ def from_dict(cls, data: Dict[str, Any]) -> "BenchmarkConfig": agentx_data = legacy_scenario agentx = AgentXConfig.from_value(agentx_data) - profiler_data = dict(data.get("profiler", {}) or {}) - if agentx is not None and agentx.enabled: - profiler_data.setdefault("torch_profiler", {"enabled": False}) - profiler_data.setdefault("gpu_monitor", {"enabled": False}) - profiler = ( - ProfilerConfig.from_dict(profiler_data) - if profiler_data - else ProfilerConfig() - ) + profiler = data.get("profiler") gap_data = data.get("gap_analysis", {}) gap_analysis = ( diff --git a/Magpie/modes/benchmark/inferencex.py b/Magpie/modes/benchmark/inferencex.py index 905b478..7608f25 100644 --- a/Magpie/modes/benchmark/inferencex.py +++ b/Magpie/modes/benchmark/inferencex.py @@ -14,12 +14,12 @@ import os import subprocess from pathlib import Path -from typing import Optional logger = logging.getLogger(__name__) # InferenceX repository configuration INFERENCEX_REPO_URL = "https://github.com/SemiAnalysisAI/InferenceX.git" +INFERENCEX_DEFAULT_REF = "408c015be4b22d14c69518643609669405507077" # Default directory resolution order (first writable candidate wins): _PROJECT_ROOT = Path(__file__).parent.parent.parent.parent INFERENCEX_DEFAULT_DIR = str(_PROJECT_ROOT / "InferenceX") @@ -31,6 +31,41 @@ } +def resolve_inferencex_root(path: str) -> str: + """Return a supported InferenceX project root without altering the checkout. + + A benchmarks directory alone is not an installation: Magpie may have copied + its own scripts there. Require the upstream library and client source too. + """ + root = Path(path).expanduser().resolve() + candidates = [] + for candidate in (root, root / "inferencex-e2e"): + if not (candidate / "benchmarks" / "benchmark_lib.sh").is_file(): + continue + legacy = candidate / "utils" / "bench_serving" / "benchmark_serving.py" + packaged = candidate / "infx" / "bench_serving" / "benchmark_serving.py" + if legacy.is_file() or ( + packaged.is_file() and (candidate / "pyproject.toml").is_file() + ): + candidates.append(candidate) + if len(candidates) != 1: + raise RuntimeError( + f"Expected one supported InferenceX project at {root}; provide the " + "repository root or its inferencex-e2e directory containing " + "benchmarks/benchmark_lib.sh and the upstream benchmark client" + ) + return str(candidates[0]) + + +def resolve_benchmark_serving(path: str) -> Path: + """Locate the supported client source for profiling and runtime adapters.""" + root = Path(resolve_inferencex_root(path)) + packaged = root / "infx" / "bench_serving" / "benchmark_serving.py" + if packaged.is_file() and (root / "pyproject.toml").is_file(): + return packaged + return root / "utils" / "bench_serving" / "benchmark_serving.py" + + def _resolve_default_inferencex_dir() -> str: """Pick the first writable default location for InferenceX. @@ -46,8 +81,11 @@ def _resolve_default_inferencex_dir() -> str: candidates = [ INFERENCEX_DEFAULT_DIR, - str(Path(os.environ.get("XDG_CACHE_HOME", str(Path.home() / ".cache"))) - / "magpie" / "InferenceX"), + str( + Path(os.environ.get("XDG_CACHE_HOME", str(Path.home() / ".cache"))) + / "magpie" + / "InferenceX" + ), "/root/workspace/InferenceX", ] for candidate in candidates: @@ -68,21 +106,21 @@ def _resolve_default_inferencex_dir() -> str: class InferenceXManager: """ Manages InferenceX repository for benchmark execution. - + Handles: - Validation of existing installation - Automatic cloning if not present - Path resolution for placeholder values """ - + def __init__( self, repo_url: str = INFERENCEX_REPO_URL, - default_dir: Optional[str] = None, + default_dir: str | None = None, ): """ Initialize InferenceX manager. - + Args: repo_url: Git repository URL for InferenceX default_dir: Default directory to clone into. If None, resolves @@ -90,114 +128,127 @@ def __init__( """ self.repo_url = repo_url self.default_dir = default_dir or _resolve_default_inferencex_dir() - - def ensure_available(self, configured_path: Optional[str] = None) -> str: + + def ensure_available(self, configured_path: str | None = None) -> str: """ Ensure InferenceX is available at the specified path. - + If the path doesn't exist or is a placeholder value, automatically clone InferenceX from the repository. - + Args: configured_path: Configured path to InferenceX (may be None or placeholder) - + Returns: Valid path to InferenceX installation - + Raises: RuntimeError: If unable to clone or validate InferenceX """ # Determine if path is a placeholder or not configured is_placeholder = self._is_placeholder(configured_path) - - # Check if configured path exists - if not is_placeholder and configured_path and os.path.exists(configured_path): - logger.debug(f"InferenceX found at: {configured_path}") - return configured_path - + # Determine target directory - target_dir = self.default_dir if is_placeholder or not configured_path else configured_path - + target_dir = ( + self.default_dir + if is_placeholder or not configured_path + else configured_path + ) + target_dir = str(Path(target_dir).expanduser().resolve()) + # Check if already exists at target directory if os.path.exists(target_dir): - if self._validate_installation(target_dir): - logger.info(f"InferenceX already exists at: {target_dir}") - return target_dir - else: - logger.warning(f"Directory exists but doesn't appear to be InferenceX: {target_dir}") - # Still return it, let downstream code handle validation - return target_dir - + project_root = resolve_inferencex_root(target_dir) + logger.info(f"InferenceX already exists at: {project_root}") + return project_root + # Clone the repository return self._clone_repository(target_dir) - - def _is_placeholder(self, path: Optional[str]) -> bool: + + def _is_placeholder(self, path: str | None) -> bool: """Check if the path is a placeholder value.""" if path is None: return True return path in PLACEHOLDER_VALUES - + def _validate_installation(self, path: str) -> bool: """ Validate that the path contains a valid InferenceX installation. - + Args: path: Path to validate - + Returns: True if valid InferenceX installation """ - required_paths = [ - "benchmarks", # InferenceX benchmark scripts directory - ] - - for required in required_paths: - if not os.path.exists(os.path.join(path, required)): - return False - + try: + resolve_inferencex_root(path) + except (OSError, RuntimeError): + return False return True - + def _clone_repository(self, target_dir: str) -> str: """ Clone InferenceX repository. - + Args: target_dir: Directory to clone into - + Returns: Path to cloned repository - + Raises: RuntimeError: If clone fails """ logger.info(f"InferenceX not found. Cloning from {self.repo_url}...") logger.info(f"Clone destination: {target_dir}") - + try: # Ensure parent directory exists parent_dir = os.path.dirname(target_dir) if parent_dir and not os.path.exists(parent_dir): os.makedirs(parent_dir, exist_ok=True) - - # Clone the repository - result = subprocess.run( - ["git", "clone", self.repo_url, target_dir], - capture_output=True, - text=True, - timeout=300, # 5 minute timeout - ) - - if result.returncode != 0: - error_msg = result.stderr.strip() if result.stderr else "Unknown error" - logger.error(f"Failed to clone InferenceX: {error_msg}") - raise RuntimeError(f"git clone failed: {error_msg}") - + + commands = [ + ["git", "clone", "--no-checkout", self.repo_url, target_dir], + [ + "git", + "-C", + target_dir, + "checkout", + "--detach", + INFERENCEX_DEFAULT_REF, + ], + [ + "git", + "-C", + target_dir, + "submodule", + "update", + "--init", + "--recursive", + "--", + "inferencex-e2e/utils/aiperf", + ], + ] + for command in commands: + result = subprocess.run( + command, capture_output=True, text=True, timeout=300, check=False + ) + if result.returncode != 0: + error_msg = ( + result.stderr.strip() if result.stderr else "Unknown error" + ) + raise RuntimeError( + f"InferenceX {' '.join(command)} failed: {error_msg}" + ) + logger.info(f"Successfully cloned InferenceX to: {target_dir}") - return target_dir - + return resolve_inferencex_root(target_dir) + except subprocess.TimeoutExpired: - logger.error("InferenceX clone timed out after 5 minutes") - raise RuntimeError("git clone timed out after 5 minutes") + logger.error("InferenceX installation command timed out after 5 minutes") + raise RuntimeError("InferenceX git command timed out after 5 minutes") except FileNotFoundError: logger.error("git command not found. Please install git.") raise RuntimeError("git is not installed. Please install git first.") @@ -207,7 +258,7 @@ def _clone_repository(self, target_dir: str) -> str: # Module-level instance for convenience -_manager: Optional[InferenceXManager] = None +_manager: InferenceXManager | None = None def get_manager() -> InferenceXManager: @@ -218,15 +269,14 @@ def get_manager() -> InferenceXManager: return _manager -def ensure_inferencex_available(configured_path: Optional[str] = None) -> str: +def ensure_inferencex_available(configured_path: str | None = None) -> str: """ Convenience function to ensure InferenceX is available. - + Args: configured_path: Configured path to InferenceX - + Returns: Valid path to InferenceX installation """ return get_manager().ensure_available(configured_path) - diff --git a/Magpie/modes/benchmark/result.py b/Magpie/modes/benchmark/result.py index df5d52a..1aedadc 100644 --- a/Magpie/modes/benchmark/result.py +++ b/Magpie/modes/benchmark/result.py @@ -565,6 +565,25 @@ def _parse_agentx_data( data.get("num_requests_total", accounting.get("records_total")) ) errors = _as_non_negative_int(accounting.get("records_error_dropped")) + warmup = _as_non_negative_int(accounting.get("records_warmup_dropped")) + if "records_dropped_total" in accounting: + # InferenceX counts warmup and error records independently: a + # failed warmup is in both counters. The dropped total is their + # union, so removing all warmup rows leaves only measured errors. + # Older aggregates lack that union; keep their conservative error + # count rather than guessing how many errors occurred in warmup. + dropped = accounting["records_dropped_total"] + if not ( + type(dropped) is int + and 0 <= warmup <= dropped + and 0 <= dropped - warmup <= errors <= dropped + and successful + dropped == total + ): + result.benchmark_valid = False + result.publishable = False + result.errors.append("AgentX request accounting is inconsistent") + return result + errors = dropped - warmup completed = successful + errors error_rate = errors / completed if completed else 1.0 @@ -613,6 +632,26 @@ def _parse_agentx_data( "AgentX result is missing a positive duration or total throughput" ) + if data.get("custom_recipe"): + native_context = data.get("native_context_length") + max_context = data.get("max_model_len") + metadata_hash = data.get("model_config_sha256") + if not ( + data["custom_recipe"] is True + and type(native_context) is int + and type(max_context) is int + and 0 < max_context <= native_context + and isinstance(metadata_hash, str) + and len(metadata_hash) == 64 + and all( + character in "0123456789abcdef" for character in metadata_hash + ) + ): + valid = False + result.errors.append( + "Custom AgentX result has invalid model/context metadata" + ) + result.success = valid result.benchmark_valid = valid result.publishable = valid and mode == "canonical" and fingerprint_valid @@ -626,9 +665,7 @@ def _parse_agentx_data( "profiled_total": completed, "successful": successful, "errors": errors, - "warmup_dropped": _as_non_negative_int( - accounting.get("records_warmup_dropped") - ), + "warmup_dropped": warmup, "error_rate": error_rate, "threshold": failed_request_threshold, }, @@ -664,6 +701,10 @@ def _parse_agentx_data( "allocated_cpu_dram_gb", "router", "kv_p2p_transfer", + "custom_recipe", + "native_context_length", + "max_model_len", + "model_config_sha256", ) if key in data }, @@ -702,7 +743,9 @@ def parse_torch_trace(trace_dir: Path) -> List[KernelMetrics]: candidate_files = [ f for f in trace_files - if "capture_traces" not in (part.lower() for part in f.parts) + if not {"capture_traces", "graph_capture_profile"}.intersection( + part.lower() for part in f.parts + ) ] if not candidate_files: candidate_files = trace_files diff --git a/Magpie/modes/benchmark/tracelens.py b/Magpie/modes/benchmark/tracelens.py index 056b0ca..f36e7f2 100644 --- a/Magpie/modes/benchmark/tracelens.py +++ b/Magpie/modes/benchmark/tracelens.py @@ -15,9 +15,11 @@ """ import logging +import re import shutil import subprocess import sys +import tempfile from pathlib import Path from typing import Dict, List, Optional, Any @@ -198,7 +200,10 @@ def analyze( results["errors"].append(rank0_result["error"]) # 2. Multi-rank collective analysis (TraceLens_generate_multi_rank_collective_report_pytorch) - if self.config.multi_rank_report_enabled and len(trace_files) >= num_ranks and num_ranks > 1: + if self.config.multi_rank_report_enabled and num_ranks > 1 and ( + len(trace_files) >= num_ranks + or any(re.search(r"-TP-\d+", path.name) for path in trace_files) + ): logger.info("Running TraceLens_generate_multi_rank_collective_report_pytorch...") collective_csv_dir = output_dir / "tracelens_collective_csvs" if use_csv else None @@ -306,7 +311,11 @@ def _find_trace_files(self, trace_dir: Path) -> List[Path]: trace_files = [] for pattern in ["*.json.gz", "*.json"]: - trace_files.extend(trace_dir.rglob(pattern)) + trace_files.extend( + path for path in trace_dir.rglob(pattern) + if path.name != "capture.json" + and not {"capture_traces", "graph_capture_profile"}.intersection(path.parts) + ) # Filter out async_llm traces (main process, not worker) worker_traces = [ @@ -434,6 +443,29 @@ def _run_multi_rank_collective( # Find trace pattern - TraceLens expects pattern like "rank*_trace.json.gz" # or we can use a pattern that matches our trace files trace_files = self._find_trace_files(trace_dir) + + if any(re.search(r"-TP-\d+", path.name) for path in trace_files): + # TraceLens's pattern accepts one rank placeholder, not a glob. + # SGLang includes independent TP and EP ids (and a timestamp), so + # use validated TP ranks without renaming or modifying the traces. + try: + ranks = self._sglang_trace_ranks(trace_files, num_ranks) + report_dir = output_csv_dir or output_xlsx.parent + report_dir.mkdir(parents=True, exist_ok=True) + with tempfile.TemporaryDirectory( + prefix="tracelens-sglang-ranks-", dir=report_dir + ) as temporary: + rank_dir = Path(temporary) + for rank, path in ranks.items(): + suffix = ".json.gz" if path.name.endswith(".gz") else ".json" + (rank_dir / f"rank-{rank}{suffix}").symlink_to(path.resolve()) + return self._run_multi_rank_collective( + rank_dir, output_csv_dir, output_xlsx, num_ranks + ) + except (OSError, ValueError) as exc: + result["error"] = f"SGLang multi-rank trace mapping failed: {exc}" + logger.error(result["error"]) + return result if len(trace_files) < num_ranks: logger.warning( @@ -505,6 +537,37 @@ def _run_multi_rank_collective( return result + @staticmethod + def _sglang_trace_ranks( + trace_files: list[Path], num_ranks: int + ) -> dict[int, Path]: + """Map one single-node TP capture; PP/DP need explicit topology support.""" + ranks = {} + families = set() + for path in trace_files: + if re.search(r"-(?:DP|PP)-", path.name): + raise ValueError(f"PP/DP trace topology is unsupported: {path.name}") + match = re.fullmatch( + r"(.+)-TP-(\d+)(?:-EP-\d+)?(\.trace\.json(?:\.gz)?)", path.name + ) + if match is None: + raise ValueError( + f"Expected one TP-only SGLang capture (optional EP); " + f"PP/DP or mixed filenames are unsupported: {path.name}" + ) + rank = int(match[2]) + if rank in ranks: + raise ValueError(f"Duplicate TP rank {rank}; select one capture/stage") + ranks[rank] = path + families.add((path.parent, match[1], match[3])) + if len(families) != 1: + raise ValueError("Mixed SGLang capture/stage files; select one capture/stage") + if set(ranks) != set(range(num_ranks)): + raise ValueError( + f"Expected TP ranks 0..{num_ranks - 1}; found {sorted(ranks)}" + ) + return ranks + def _detect_trace_pattern( self, trace_dir: Path, @@ -528,8 +591,6 @@ def _detect_trace_pattern( if not trace_files: return None - import re - # Operate on the path relative to trace_dir so a rank_/ directory # component is part of the pattern instead of being dropped. try: diff --git a/Magpie/modes/benchmark/tracelens_inference.py b/Magpie/modes/benchmark/tracelens_inference.py index d2fbe2c..80ae0fa 100644 --- a/Magpie/modes/benchmark/tracelens_inference.py +++ b/Magpie/modes/benchmark/tracelens_inference.py @@ -31,6 +31,7 @@ from ...utils.gpu import detect_gpu from .config import BenchmarkConfig +from .inferencex import resolve_benchmark_serving, resolve_inferencex_root from .tracelens import ensure_tracelens_installed from .tracelens_runtime import ATOM_DETAILED_ANNOTATION_SCRIPT, docker_image_probe @@ -386,16 +387,20 @@ def prepare( if "NUM_PROMPTS" not in envs: envs["NUM_PROMPTS"] = str(int(float(envs.get("CONC", 32))) * 10) - # ATOM ships its own runner and does not source InferenceX benchmark_lib.sh. - if self.config.framework != "atom": - self._patch_benchmark_lib(result) - - if self.config.framework == "vllm": - self._prepare_vllm_env(envs, max_iters, delay_iters) - elif self.config.framework == "sglang": - self._prepare_sglang_env(envs, max_iters, delay_iters, result) - elif self.config.framework == "atom": - self._prepare_atom_env(envs, result, runtime) + try: + # ATOM ships its own runner and does not source InferenceX benchmark_lib.sh. + if self.config.framework != "atom": + self._patch_benchmark_lib(result) + + if self.config.framework == "vllm": + self._prepare_vllm_env(envs, max_iters, delay_iters) + elif self.config.framework == "sglang": + self._prepare_sglang_env(envs, max_iters, delay_iters, result) + elif self.config.framework == "atom": + self._prepare_atom_env(envs, result, runtime) + except Exception: + self.restore() + raise result["env_updates"] = { key: envs.get(key) @@ -773,23 +778,21 @@ def _should_enable_sglang_shape_discovery(self) -> bool: ) or (self.config.is_local and host_sglang_supports_shape_discovery()) def _patch_benchmark_lib(self, result: Dict[str, Any]) -> None: - path = Path(self.config.inferencex_path) / "benchmarks" / "benchmark_lib.sh" - if not path.exists(): - warning = f"benchmark_lib.sh not found, skipping TraceLens preprocess patch: {path}" - logger.warning(warning) - result["warnings"].append(warning) - return + path = ( + Path(resolve_inferencex_root(self.config.inferencex_path)) + / "benchmarks" + / "benchmark_lib.sh" + ) text = path.read_text(encoding="utf-8") old = 'num_prompts="$max_concurrency"' new = 'num_prompts="$num_prompts"' if new in text: return - if old not in text: - warning = f"benchmark_lib.sh profile num_prompts anchor not found: {path}" - logger.warning(warning) - result["warnings"].append(warning) - return + if text.count(old) != 1: + raise RuntimeError( + f"Unsupported TraceLens benchmark_lib.sh profile anchor: {path}" + ) self._backup_file(path) path.write_text(text.replace(old, new, 1), encoding="utf-8") @@ -802,49 +805,27 @@ def _patch_sglang_benchmark_serving( delay_iters: int, result: Dict[str, Any], ) -> None: - path = ( - Path(self.config.inferencex_path) - / "utils" - / "bench_serving" - / "benchmark_serving.py" - ) - if not path.exists(): - warning = f"benchmark_serving.py not found, skipping SGLang TraceLens patch: {path}" - logger.warning(warning) - result["warnings"].append(warning) - return - + path = resolve_benchmark_serving(self.config.inferencex_path) text = path.read_text(encoding="utf-8") - common_anchor = '"num_steps": 1, "merge_profiles": True, "profile_by_stage": True' - common_repl = ( - '"shape_discovery": True, "detailed_annotations": True, ' - '"num_steps": 1, "merge_profiles": True, "profile_by_stage": True' + # Legacy clients keep this dict on one line; packaged infx clients use + # multiple lines. Match the known request fields, not source formatting. + common_anchor = ( + r'"num_steps"\s*:\s*1,\s*"merge_profiles"\s*:\s*True,' + r'\s*"profile_by_stage"\s*:\s*True' ) steady_repl = ( + '"shape_discovery": True, "detailed_annotations": True, ' f'"start_step": {delay_iters}, "num_steps": {max_iters}, ' '"merge_profiles": False, "profile_by_stage": False' ) - - if steady_repl in text and "detailed_annotations" in text: - return - - if "detailed_annotations" not in text: - if common_anchor not in text: - warning = f"SGLang benchmark_serving.py common patch anchor not found: {path}" - logger.warning(warning) - result["warnings"].append(warning) - return - self._backup_file(path) - text = text.replace(common_anchor, common_repl, 1) - - if common_anchor not in text: - warning = f"SGLang benchmark_serving.py steady-state patch anchor not found: {path}" - logger.warning(warning) - result["warnings"].append(warning) + if steady_repl in text: return - + if len(re.findall(common_anchor, text)) != 1: + raise RuntimeError( + f"Unsupported SGLang TraceLens profile request anchor: {path}" + ) self._backup_file(path) - text = text.replace(common_anchor, steady_repl, 1) + text = re.sub(common_anchor, lambda _match: steady_repl, text, count=1) path.write_text(text, encoding="utf-8") result["patched_files"].append(str(path)) logger.info("Patched SGLang benchmark_serving.py for TraceLens inference profiling") diff --git a/Magpie/modes/benchmark/workspace.py b/Magpie/modes/benchmark/workspace.py index ae45f60..d1613ad 100644 --- a/Magpie/modes/benchmark/workspace.py +++ b/Magpie/modes/benchmark/workspace.py @@ -201,7 +201,8 @@ def collect_results(self) -> Dict[str, Any]: torch_trace_dir = self._workspace_path / "torch_trace" if torch_trace_dir.exists(): results["torch_trace_files"] = [ - str(f) for f in torch_trace_dir.iterdir() if f.is_file() + str(f) for f in torch_trace_dir.rglob("*") + if f.is_file() and f.name != "capture.json" ] # List system profile files diff --git a/Magpie/scripts/agentx/client_legacy.sh b/Magpie/scripts/agentx/client_legacy.sh new file mode 100644 index 0000000..359555f --- /dev/null +++ b/Magpie/scripts/agentx/client_legacy.sh @@ -0,0 +1,15 @@ +#!/usr/bin/env bash +# Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. +# See LICENSE for license information. +# Explicit adapter for the pre-srt client library. Never launches a server. +set -eo pipefail +: "${INFMAX_CONTAINER_WORKSPACE:?}" "${AIPERF_SERVER_URL:?}" "${RESULT_DIR:?}" +source "$INFMAX_CONTAINER_WORKSPACE/benchmarks/benchmark_lib.sh" +for function in resolve_trace_source install_agentic_deps build_replay_cmd run_agentic_replay_and_write_outputs; do + declare -F "$function" >/dev/null || { echo "Unsupported legacy AgentX client library: $function missing" >&2; exit 2; } +done +resolve_trace_source +install_agentic_deps +mkdir -p "$RESULT_DIR" +build_replay_cmd "$RESULT_DIR" +run_agentic_replay_and_write_outputs "$RESULT_DIR" diff --git a/Magpie/scripts/benchmark/atom_mi300x.sh b/Magpie/scripts/benchmark/atom_mi300x.sh index 9e22259..9969f40 100644 --- a/Magpie/scripts/benchmark/atom_mi300x.sh +++ b/Magpie/scripts/benchmark/atom_mi300x.sh @@ -122,11 +122,17 @@ if [[ -n "${SERVER_PID:-}" ]]; then fi if [[ "$PHASE" == "client" || "$PHASE" == "all" ]]; then + # Client-only tokenizer selection; retain Hyperloom's existing environment alias. + CLIENT_TOKENIZER_MODE="${MAGPIE_CLIENT_TOKENIZER_MODE:-${HYPERLOOM_CLIENT_TOKENIZER_MODE:-}}" + CLIENT_TOKENIZER_ARGS=() + if [[ -n "$CLIENT_TOKENIZER_MODE" ]]; then + CLIENT_TOKENIZER_ARGS+=(--tokenizer-mode "$CLIENT_TOKENIZER_MODE") + fi if [[ -n "${BENCHMARK_BASE_URL:-}" ]]; then SERVER_MONITOR_ARGS=() - magpie_run_benchmark_serving_remote_direct trust || exit $? + magpie_run_benchmark_serving_remote_direct trust "${CLIENT_TOKENIZER_ARGS[@]}" || exit $? else - run_benchmark_serving \ + magpie_run_benchmark_serving \ --model "$MODEL" \ --port "$PORT" \ --backend vllm \ @@ -138,6 +144,7 @@ if [[ "$PHASE" == "client" || "$PHASE" == "all" ]]; then --result-filename "$RESULT_FILENAME" \ --result-dir "$WORKSPACE_DIR/" \ "${SERVER_MONITOR_ARGS[@]}" \ + "${CLIENT_TOKENIZER_ARGS[@]}" \ --trust-remote-code || exit $? fi fi diff --git a/Magpie/scripts/benchmark/atom_mi355x.sh b/Magpie/scripts/benchmark/atom_mi355x.sh index c88143c..d8d4bec 100644 --- a/Magpie/scripts/benchmark/atom_mi355x.sh +++ b/Magpie/scripts/benchmark/atom_mi355x.sh @@ -118,11 +118,17 @@ if [[ -n "${SERVER_PID:-}" ]]; then fi if [[ "$PHASE" == "client" || "$PHASE" == "all" ]]; then + # Client-only tokenizer selection; retain Hyperloom's existing environment alias. + CLIENT_TOKENIZER_MODE="${MAGPIE_CLIENT_TOKENIZER_MODE:-${HYPERLOOM_CLIENT_TOKENIZER_MODE:-}}" + CLIENT_TOKENIZER_ARGS=() + if [[ -n "$CLIENT_TOKENIZER_MODE" ]]; then + CLIENT_TOKENIZER_ARGS+=(--tokenizer-mode "$CLIENT_TOKENIZER_MODE") + fi if [[ -n "${BENCHMARK_BASE_URL:-}" ]]; then SERVER_MONITOR_ARGS=() - magpie_run_benchmark_serving_remote_direct trust || exit $? + magpie_run_benchmark_serving_remote_direct trust "${CLIENT_TOKENIZER_ARGS[@]}" || exit $? else - run_benchmark_serving \ + magpie_run_benchmark_serving \ --model "$MODEL" \ --port "$PORT" \ --backend vllm \ @@ -134,6 +140,7 @@ if [[ "$PHASE" == "client" || "$PHASE" == "all" ]]; then --result-filename "$RESULT_FILENAME" \ --result-dir "$WORKSPACE_DIR/" \ "${SERVER_MONITOR_ARGS[@]}" \ + "${CLIENT_TOKENIZER_ARGS[@]}" \ --trust-remote-code || exit $? fi fi diff --git a/Magpie/scripts/benchmark/magpie_bench_remote_compat.sh b/Magpie/scripts/benchmark/magpie_bench_remote_compat.sh index 9fc5816..b967464 100644 --- a/Magpie/scripts/benchmark/magpie_bench_remote_compat.sh +++ b/Magpie/scripts/benchmark/magpie_bench_remote_compat.sh @@ -1,59 +1,166 @@ #!/usr/bin/env bash ############################################################################### -# Remote benchmark compat (Magpie). -# -# InferenceX benchmarks/benchmark_lib.sh defines run_benchmark_serving() which -# parses a fixed set of flags. Older (and some current) trees reject -# --base-url at the bash layer even though utils/bench_serving/benchmark_serving.py -# accepts --base-url for OpenAI-compatible servers (SGLang/vLLM HTTP). -# -# When BENCHMARK_BASE_URL is set, Magpie *mi*.sh scripts call -# magpie_run_benchmark_serving_remote_direct() instead of passing --base-url into -# run_benchmark_serving(). -# -# Working directory must be the InferenceX repository root (Magpie benchmarker -# runs: cd && bash benchmarks/