diff --git a/Magpie/modes/benchmark/agentx_profile_capabilities.py b/Magpie/modes/benchmark/agentx_profile_capabilities.py index e004f12..d8920d6 100644 --- a/Magpie/modes/benchmark/agentx_profile_capabilities.py +++ b/Magpie/modes/benchmark/agentx_profile_capabilities.py @@ -80,8 +80,8 @@ def fields(cls): 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. + Shape discovery is optional for SGLang annotations. Graph support is reported + separately so callers can require it only for graph-mode captures. """ if framework not in {"sglang", "vllm"}: raise ValueError("AgentX detailed profiling supports SGLang and vLLM") @@ -113,14 +113,12 @@ def validate_profile_capabilities(framework: str, value: Any) -> dict[str, Any]: 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"): + for name in ("shape_discovery", "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, + "shape_discovery": value["shape_discovery"], "graph_capture": value["graph_capture"], "graph_shape_discovery": value["graph_shape_discovery"], } diff --git a/Magpie/modes/benchmark/agentx_profile_config.py b/Magpie/modes/benchmark/agentx_profile_config.py index 9df3b2d..ea7e98b 100644 --- a/Magpie/modes/benchmark/agentx_profile_config.py +++ b/Magpie/modes/benchmark/agentx_profile_config.py @@ -208,12 +208,9 @@ def profile_server_spec( SGLANG_GRAPH_BATCH_CAPTURE="True", ) if "--disable-cuda-graph" not in effective_flags: - if ( - not capabilities["graph_capture"] - or not capabilities["graph_shape_discovery"] - ): + if not capabilities["graph_capture"]: raise ValueError( - "SGLang detailed annotations with CUDA graphs require graph profiling and shape discovery capabilities" + "SGLang detailed annotations with CUDA graphs require graph profiling support" ) for name, flag in ( ("graph_capture", "--enable-profile-cuda-graph"), diff --git a/Magpie/modes/benchmark/agentx_profiling.py b/Magpie/modes/benchmark/agentx_profiling.py index 8b52184..8696711 100644 --- a/Magpie/modes/benchmark/agentx_profiling.py +++ b/Magpie/modes/benchmark/agentx_profiling.py @@ -631,9 +631,10 @@ def check_before_start() -> None: body.update( with_stack=True, record_shapes=True, - shape_discovery=True, **{capabilities["annotation_field"]: True}, ) + if capabilities["shape_discovery"]: + body["shape_discovery"] = True deadline = time.monotonic() + capture_timeout if ( measurement_duration_seconds is not None diff --git a/docs/how-to/benchmarking/profiling-options.md b/docs/how-to/benchmarking/profiling-options.md index 3a9655f..19719d6 100644 --- a/docs/how-to/benchmarking/profiling-options.md +++ b/docs/how-to/benchmarking/profiling-options.md @@ -103,14 +103,20 @@ Ordinary torch traces and PyTorch TraceLens reports work with diagnostics require a framework runtime instrumented for those annotations. Opt in with `detailed_annotations: true` or `--torch-profiler-detailed-annotations`; Magpie probes the actual server Python -environment before launch and rejects missing capabilities. It does not patch -the pinned InferenceX checkout. Supply the instrumented image or source overlay -before enabling this option. +environment before launch and rejects missing required capabilities. It does +not patch the pinned InferenceX checkout. Supply the instrumented image or +source overlay before enabling this option. -For SGLang, Magpie enables shape/stack recording and shape discovery, and selects +For SGLang, Magpie enables shape/stack recording and selects the supported `roofline_annotations` or older `detailed_annotations` request -field. With CUDA graphs enabled it also requires graph profiling and graph -shape-discovery support. For vLLM, it sets detailed trace annotations, +field. Shape discovery is optional: Magpie sends `shape_discovery` and adds the +graph shape-discovery flag only when the respective capability is supported. +SGLang runtimes with `detailed_annotations` but no shape discovery can still +capture annotated traces. With CUDA graphs enabled, +graph profiling support is still required; Magpie does not disable graphs to +work around missing support. The launch evidence and capture manifests retain +the actual capability flags, including absent shape-discovery support. +For vLLM, it sets detailed trace annotations, shape/stack recording, and the supported graph-capture profiler field. An eager server stays eager. Callers pass the canonical setting above; Magpie owns the framework-specific profiler flags and environment. Enhanced annotation support diff --git a/docs/reference/benchmark-config.md b/docs/reference/benchmark-config.md index 526e2d4..e785733 100644 --- a/docs/reference/benchmark-config.md +++ b/docs/reference/benchmark-config.md @@ -340,7 +340,10 @@ not enable profiling. CLI equivalents are `--torch-profiler-start-seconds` and and PyTorch TraceLens reports on stock profiler-capable frameworks. Set it to `true` for enhanced shape/roofline diagnostics only with an instrumented SGLang or vLLM runtime. Magpie probes the actual server environment and fails before -launch if the required annotation/shape capabilities are absent. It maps the +launch if required annotation capabilities are absent. SGLang request and graph +shape discovery are optional and enabled independently when supported; absent +support is recorded in launch evidence and capture manifests. Graph-mode SGLang +still requires graph profiling support. Magpie maps the setting to the supported new or old framework fields, including shape/stack recording and graph profiling when graphs are enabled; it keeps eager mode unchanged. Magpie owns these framework flags, while callers supply the diff --git a/tests/unit/modes/test_agentx_profile_capabilities.py b/tests/unit/modes/test_agentx_profile_capabilities.py index fc3dac5..57334ea 100644 --- a/tests/unit/modes/test_agentx_profile_capabilities.py +++ b/tests/unit/modes/test_agentx_profile_capabilities.py @@ -34,7 +34,9 @@ def _probe(root, framework, argv=None, **kwargs): ) -def _sglang(root, annotation="roofline_annotations", shape=True, graph=True): +def _sglang( + root, annotation="roofline_annotations", shape=True, graph=True, graph_shape=True +): _module( root, "sglang.srt.server_args", @@ -43,10 +45,14 @@ def _sglang(root, annotation="roofline_annotations", shape=True, graph=True): "@dataclass\nclass ServerArgs:\n" + ( " enable_profile_cuda_graph: bool = False\n" - " enable_shape_discovery_for_cuda_graph_profile: bool = False\n" if graph else " unrelated: bool = False\n" ) + + ( + " enable_shape_discovery_for_cuda_graph_profile: bool = False\n" + if graph and graph_shape + else "" + ) ), ) _module( @@ -104,6 +110,16 @@ def test_sglang_allows_eager_runtime_without_graph_capabilities(tmp_path): assert result["graph_shape_discovery"] is False +def test_sglang_annotations_without_shape_discovery_preserve_capabilities(tmp_path): + _sglang(tmp_path, "detailed_annotations", shape=False, graph_shape=False) + assert _probe(tmp_path, "sglang") == { + "annotation_field": "detailed_annotations", + "shape_discovery": False, + "graph_capture": True, + "graph_shape_discovery": False, + } + + @pytest.mark.parametrize( "capture", ["capture_torch_profiler_dir", "capture_torch_profiler"] ) @@ -126,20 +142,11 @@ def test_vllm_profiler_submodule_and_preferred_capture_field(tmp_path): assert _probe(tmp_path, "vllm")["capture_field"] == "capture_torch_profiler_dir" -@pytest.mark.parametrize( - "annotation,shape,missing", - [ - ("other", True, "roofline_annotations/detailed_annotations"), - ("detailed_annotations", False, "shape_discovery"), - ], -) -def test_missing_sglang_capability_reports_interpreter( - tmp_path, annotation, shape, missing -): - _sglang(tmp_path, annotation, shape) +def test_missing_sglang_annotation_support_reports_interpreter(tmp_path): + _sglang(tmp_path, "other") with pytest.raises(RuntimeError) as error: _probe(tmp_path, "sglang") - assert missing in str(error.value) + assert "roofline_annotations/detailed_annotations" in str(error.value) assert sys.executable in str(error.value) diff --git a/tests/unit/modes/test_agentx_profile_integration.py b/tests/unit/modes/test_agentx_profile_integration.py index e539988..4d12525 100644 --- a/tests/unit/modes/test_agentx_profile_integration.py +++ b/tests/unit/modes/test_agentx_profile_integration.py @@ -366,7 +366,7 @@ def test_invalid_first_capture_window_fails_before_server_start( assert not (_fixture(config) / "server.json").exists() -def _enhanced_runtime(config, tmp_path, framework, *, legacy=False): +def _enhanced_runtime(config, tmp_path, framework, *, legacy=False, shape=True): config.framework = framework config.profiler.torch_profiler.detailed_annotations = True spec = config.agentx.resolved["server-launch-spec"] @@ -376,10 +376,16 @@ def _enhanced_runtime(config, tmp_path, framework, *, legacy=False): "sglang/srt/server_args.py": ( "class ServerArgs:\n" " enable_profile_cuda_graph: bool = False\n" - " enable_shape_discovery_for_cuda_graph_profile: bool = False\n" + + ( + " enable_shape_discovery_for_cuda_graph_profile: bool = False\n" + if shape + else "" + ) ), "sglang/srt/managers/io_struct.py": ( - "class ProfileReq:\n shape_discovery: bool = False\n " + "class ProfileReq:\n" + + (" shape_discovery: bool = False\n" if shape else "") + + " " + ("detailed_annotations" if legacy else "roofline_annotations") + ": bool = False\n" ), @@ -403,13 +409,21 @@ def _enhanced_runtime(config, tmp_path, framework, *, legacy=False): spec["env"].update(PYTHONPATH=str(packages), PROFILE_TEST_ENHANCED="1") -@pytest.mark.parametrize("framework", ["sglang", "vllm"]) -@pytest.mark.parametrize("legacy", [False, True]) +@pytest.mark.parametrize( + "framework,legacy,shape", + [ + ("sglang", False, True), + ("sglang", True, True), + pytest.param("sglang", True, False, id="sglang-0.5.20-capabilities"), + ("vllm", False, True), + ("vllm", True, True), + ], +) def test_enhanced_runtime_probe_and_two_captures_keep_shared_graph_evidence( - diagnostic_config, tmp_path, framework, legacy + diagnostic_config, tmp_path, framework, legacy, shape ): config = diagnostic_config - _enhanced_runtime(config, tmp_path, framework, legacy=legacy) + _enhanced_runtime(config, tmp_path, framework, legacy=legacy, shape=shape) config.profiler.torch_profiler.num_profiles = 2 config.profiler.torch_profiler.interval_seconds = 0.05 config.envs["PROFILE_TEST_DURATION"] = "2.5" @@ -439,8 +453,26 @@ def test_enhanced_runtime_probe_and_two_captures_keep_shared_graph_evidence( assert len(starts) == 2 if framework == "sglang": field = "detailed_annotations" if legacy else "roofline_annotations" + capabilities = { + "annotation_field": field, + "shape_discovery": shape, + "graph_capture": True, + "graph_shape_discovery": shape, + } + assert receipt["torch_profiler"]["capabilities"] == capabilities + assert capture["capabilities"] == capabilities + assert all(p["capabilities"] == capabilities for p in capture["profiles"]) + argv = json.loads((_fixture(config) / "server.json").read_text())["argv"] + assert "--disable-cuda-graph" not in argv + assert "--enable-profile-cuda-graph" in argv + assert ("--enable-shape-discovery-for-cuda-graph-profile" in argv) is shape assert all(start["body"][field] is True for start in starts) - assert all(start["body"]["shape_discovery"] is True for start in starts) + assert all(start["body"]["record_shapes"] is True for start in starts) + assert all(start["body"]["with_stack"] is True for start in starts) + if shape: + assert all(start["body"]["shape_discovery"] is True for start in starts) + else: + assert all("shape_discovery" not in start["body"] for start in starts) receipt["runtime_environment"]["SGLANG_PROFILE_RECORD_SHAPES"] = "False" receipt["evidence_sha256"] = digest( {key: value for key, value in receipt.items() if key != "evidence_sha256"} diff --git a/tests/unit/modes/test_agentx_profile_launch.py b/tests/unit/modes/test_agentx_profile_launch.py index fb32ab4..98d374f 100644 --- a/tests/unit/modes/test_agentx_profile_launch.py +++ b/tests/unit/modes/test_agentx_profile_launch.py @@ -219,12 +219,16 @@ def test_invalid_start_seconds_cannot_enter_plan_or_launch(tmp_path, start): @pytest.mark.parametrize("field", ["roofline_annotations", "detailed_annotations"]) -def test_sglang_enhanced_launch_binds_detected_capabilities(tmp_path, field): +@pytest.mark.parametrize("shape", [False, True]) +@pytest.mark.parametrize("graph_shape", [False, True]) +def test_sglang_enhanced_launch_binds_detected_capabilities( + tmp_path, field, shape, graph_shape +): capabilities = { "annotation_field": field, - "shape_discovery": True, + "shape_discovery": shape, "graph_capture": True, - "graph_shape_discovery": True, + "graph_shape_discovery": graph_shape, } spec = {"framework": "sglang", "argv": ["python", "--model", "fixture"]} derived = profile_server_spec( @@ -235,7 +239,10 @@ def test_sglang_enhanced_launch_binds_detected_capabilities(tmp_path, field): capabilities=capabilities, ) assert "--enable-profile-cuda-graph" in derived["argv"] - assert "--enable-shape-discovery-for-cuda-graph-profile" in derived["argv"] + assert ( + "--enable-shape-discovery-for-cuda-graph-profile" in derived["argv"] + ) is graph_shape + assert "--disable-cuda-graph" not in derived["argv"] assert derived["env"]["SGLANG_PROFILE_RECORD_SHAPES"] == "True" assert derived["torch_profiler"]["capabilities"] == capabilities eager = profile_server_spec( @@ -247,6 +254,33 @@ def test_sglang_enhanced_launch_binds_detected_capabilities(tmp_path, field): launch_overrides={"version": 1, "append_args": ["--disable-cuda-graph"]}, ) assert "--enable-profile-cuda-graph" not in eager["argv"] + assert "--enable-shape-discovery-for-cuda-graph-profile" not in eager["argv"] + + +def test_sglang_graph_profiling_remains_required_only_for_graph_mode(tmp_path): + capabilities = { + "annotation_field": "detailed_annotations", + "shape_discovery": False, + "graph_capture": False, + "graph_shape_discovery": False, + } + spec = {"framework": "sglang", "argv": ["python", "--model", "fixture"]} + settings = {**SETTINGS, "detailed_annotations": True} + with pytest.raises(ValueError, match="require graph profiling support"): + profile_server_spec( + spec, settings, tmp_path, CAPTURE_ID, capabilities=capabilities + ) + eager = profile_server_spec( + spec, + settings, + tmp_path, + CAPTURE_ID, + capabilities=capabilities, + launch_overrides={"version": 1, "append_args": ["--disable-cuda-graph"]}, + ) + assert "--enable-profile-cuda-graph" not in eager["argv"] + assert "--enable-shape-discovery-for-cuda-graph-profile" not in eager["argv"] + assert eager["torch_profiler"]["capabilities"] == capabilities @pytest.mark.parametrize( diff --git a/tests/unit/modes/test_agentx_profiling.py b/tests/unit/modes/test_agentx_profiling.py index c70ac75..45a32f0 100644 --- a/tests/unit/modes/test_agentx_profiling.py +++ b/tests/unit/modes/test_agentx_profiling.py @@ -119,10 +119,11 @@ def test_replay_ending_before_first_delayed_capture_is_a_failure(harness): @pytest.mark.parametrize("field", ["roofline_annotations", "detailed_annotations"]) -def test_sglang_enhanced_request_uses_probed_field(harness, field): +@pytest.mark.parametrize("shape", [False, True]) +def test_sglang_enhanced_request_uses_probed_field(harness, field, shape): capabilities = { "annotation_field": field, - "shape_discovery": True, + "shape_discovery": shape, "graph_capture": True, "graph_shape_discovery": True, } @@ -131,10 +132,14 @@ def test_sglang_enhanced_request_uses_probed_field(harness, field): ) body = next(body for url, body in harness.calls if url.endswith("/start_profile")) assert body[field] is True - assert body["shape_discovery"] is True + if shape: + assert body["shape_discovery"] is True + else: + assert "shape_discovery" not in body assert body["record_shapes"] is True assert body["with_stack"] is True assert result["capabilities"] == capabilities + assert manifest(harness)["capabilities"] == capabilities def test_rank_trace_mapping_uses_validated_global_rank_with_repeated_local_tp(harness):