Fix AgentX profiling on SGLang without shape discovery

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"],
         }

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"),

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

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

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

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)
 
 

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"}

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(

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):