From c263cca2a37ebd3c26435961c83bb5c46ab66053 Mon Sep 17 00:00:00 2001 From: mohammad abdul basit Date: Thu, 10 Sep 2026 15:27:40 +0000 Subject: [PATCH 1/6] code to go to patchless sglang --- .../multi_node/commands/infera.py | 13 +++ .../multi_node/scripts/launch_infera_node.py | 44 ++++++++ .../multi_node/scripts/launch_multinode.py | 46 ++++++++ .../executors/_multi_node_server_lifecycle.py | 14 ++- .../actions/executors/_server_patcher.py | 66 ++++++++++++ .../actions/executors/_workload_envs.py | 101 ++++++++++++------ .../orchestrator/actions/executors/profile.py | 19 ++-- 7 files changed, 262 insertions(+), 41 deletions(-) diff --git a/src/hyperloom/inference_optimizer/multi_node/commands/infera.py b/src/hyperloom/inference_optimizer/multi_node/commands/infera.py index 36bc709df4..2f046320c6 100644 --- a/src/hyperloom/inference_optimizer/multi_node/commands/infera.py +++ b/src/hyperloom/inference_optimizer/multi_node/commands/infera.py @@ -90,6 +90,19 @@ def _collect_forward_env() -> dict[str, str]: trace_dir = os.environ.get("HYPERLOOM_MN_PROFILE_TRACE_DIR", "").strip() if trace_dir and "SGLANG_TORCH_PROFILER_DIR" not in fwd: fwd["SGLANG_TORCH_PROFILER_DIR"] = trace_dir + # SGLang >= 0.5.18 no-patch shape discovery: the pod-side launcher + # (launch_infera_node.py) puts kernel_shape_tool on PYTHONPATH from + # TRACELENS_ROOT and honors these flags. PYTHONPATH itself is blocked from + # SSH forwarding (BLOCKED_UNTRUSTED_ENV_NAMES) and is set pod-side instead. + for _shape_key in ( + "TRACELENS_ROOT", + "TRACELENS_SHAPE_DISCOVERY", + "HYPERLOOM_SGLANG_SHAPE_MODE", + "HYPERLOOM_SGLANG_VERSION_PIN", + ): + _shape_val = os.environ.get(_shape_key, "").strip() + if _shape_val and _shape_key not in fwd: + fwd[_shape_key] = _shape_val unset_fwd = os.environ.get("HYPERLOOM_MN_UNSET_FWD_ENV", "").strip() if unset_fwd: try: diff --git a/src/hyperloom/inference_optimizer/multi_node/scripts/launch_infera_node.py b/src/hyperloom/inference_optimizer/multi_node/scripts/launch_infera_node.py index 9ad9e00c12..fbaa2e49e3 100644 --- a/src/hyperloom/inference_optimizer/multi_node/scripts/launch_infera_node.py +++ b/src/hyperloom/inference_optimizer/multi_node/scripts/launch_infera_node.py @@ -6,6 +6,7 @@ import argparse import json import os +import re import shlex import subprocess import sys @@ -411,6 +412,48 @@ def _reap_stale_engine_ports() -> None: break +# SGLang >= 0.5.18 uses the no-patch kernel_shape_tool (PYTHONPATH + +# sitecustomize + TRACELENS_SHAPE_DISCOVERY) for shape discovery. This script is +# standalone (no hyperloom import), so the gate is mirrored inline. +_KERNEL_SHAPE_TOOL_REL = ("TraceLens", "TraceUtils", "kernel_shape_tool") +_SGLANG_SITECUSTOMIZE_MIN_VERSION = (0, 5, 18) + + +def _sglang_shape_mode() -> str: + """Pod-side mirror of hyperloom's SGLang shape-mode gate (no hyperloom import).""" + override = os.environ.get("HYPERLOOM_SGLANG_SHAPE_MODE", "auto").strip().lower() + if override in {"patch", "patched"}: + return "patched" + if override == "sitecustomize": + return "sitecustomize" + version = "" + try: + import sglang # type: ignore + + version = (getattr(sglang, "__version__", "") or "").strip() + except Exception: # noqa: BLE001 + version = os.environ.get("HYPERLOOM_SGLANG_VERSION_PIN", "").strip() + m = re.match(r"^\s*v?(\d+(?:\.\d+)*)", version) + if not m: + return "patched" + vt = tuple(int(p) for p in m.group(1).split(".")) + return "sitecustomize" if vt >= _SGLANG_SITECUSTOMIZE_MIN_VERSION else "patched" + + +def _maybe_activate_kernel_shape_tool(env: dict[str, str]) -> None: + """SGLang >= 0.5.18: put the no-patch kernel_shape_tool on PYTHONPATH.""" + root = (env.get("TRACELENS_ROOT") or os.environ.get("TRACELENS_ROOT") or "").strip() + if not root or _sglang_shape_mode() != "sitecustomize": + return + tool = Path(root).joinpath(*_KERNEL_SHAPE_TOOL_REL) + if not tool.is_dir(): + _log(f"WARN kernel_shape_tool not found at {tool}; SGLang shape discovery disabled") + return + existing = (env.get("PYTHONPATH") or "").strip() + env["PYTHONPATH"] = f"{tool}{os.pathsep}{existing}" if existing else str(tool) + env.setdefault("TRACELENS_SHAPE_DISCOVERY", "1") + + def _build_sglang_cmd( a: argparse.Namespace, node_rank: int, @@ -724,6 +767,7 @@ def main() -> int: if _shared_log_dir.startswith("/") and "$" not in _shared_log_dir: log_file = Path(_shared_log_dir) / f"mn_infera_server_{advertise_host}_r{node_rank}.log" if args.framework == "sglang": + _maybe_activate_kernel_shape_tool(env) cmd = _build_sglang_cmd(args, node_rank, leader, advertise_host=advertise_host) pid = _detach_launch(cmd, log_file, pid_file, env) else: diff --git a/src/hyperloom/inference_optimizer/multi_node/scripts/launch_multinode.py b/src/hyperloom/inference_optimizer/multi_node/scripts/launch_multinode.py index cd9958e454..e5866a73d7 100755 --- a/src/hyperloom/inference_optimizer/multi_node/scripts/launch_multinode.py +++ b/src/hyperloom/inference_optimizer/multi_node/scripts/launch_multinode.py @@ -9,6 +9,7 @@ import argparse import json import os +import re import shlex import subprocess import pathlib @@ -310,6 +311,50 @@ def _probe_mec_firmware_lt_177() -> bool: return False +# SGLang >= 0.5.18 uses the no-patch kernel_shape_tool (PYTHONPATH + +# sitecustomize + TRACELENS_SHAPE_DISCOVERY) for shape discovery. This script is +# standalone (no hyperloom import), so the gate is mirrored inline. +_KERNEL_SHAPE_TOOL_REL = ("TraceLens", "TraceUtils", "kernel_shape_tool") +_SGLANG_SITECUSTOMIZE_MIN_VERSION = (0, 5, 18) + + +def _sglang_shape_mode() -> str: + """Pod-side mirror of hyperloom's SGLang shape-mode gate (no hyperloom import).""" + override = os.environ.get("HYPERLOOM_SGLANG_SHAPE_MODE", "auto").strip().lower() + if override in {"patch", "patched"}: + return "patched" + if override == "sitecustomize": + return "sitecustomize" + version = "" + try: + import sglang # type: ignore + + version = (getattr(sglang, "__version__", "") or "").strip() + except Exception: # noqa: BLE001 + version = os.environ.get("HYPERLOOM_SGLANG_VERSION_PIN", "").strip() + m = re.match(r"^\s*v?(\d+(?:\.\d+)*)", version) + if not m: + return "patched" + vt = tuple(int(p) for p in m.group(1).split(".")) + return "sitecustomize" if vt >= _SGLANG_SITECUSTOMIZE_MIN_VERSION else "patched" + + +def _maybe_activate_kernel_shape_tool(sub_env: dict[str, str]) -> None: + """SGLang >= 0.5.18: put the no-patch kernel_shape_tool on PYTHONPATH.""" + root = os.environ.get("TRACELENS_ROOT", "").strip() + if not root or _sglang_shape_mode() != "sitecustomize": + return + tool = Path(root).joinpath(*_KERNEL_SHAPE_TOOL_REL) + if not tool.is_dir(): + sys.stderr.write( + f"WARN kernel_shape_tool not found at {tool}; SGLang shape discovery disabled\n" + ) + return + existing = sub_env.get("PYTHONPATH", "").strip() + sub_env["PYTHONPATH"] = f"{tool}{os.pathsep}{existing}" if existing else str(tool) + sub_env.setdefault("TRACELENS_SHAPE_DISCOVERY", "1") + + def _subprocess_env() -> dict[str, str]: """Build the framework launcher subprocess env.""" env = dict(os.environ) @@ -443,6 +488,7 @@ def _spawn_remote( dist_init_addr = f"{head_ip}:{dist_init_port}" fw = framework.lower() if fw == "sglang": + _maybe_activate_kernel_shape_tool(sub_env) cmd = _build_sglang_cmd( model=model, tp=tp, diff --git a/src/hyperloom/orchestrator/actions/executors/_multi_node_server_lifecycle.py b/src/hyperloom/orchestrator/actions/executors/_multi_node_server_lifecycle.py index 0355dffc8a..105fa29285 100644 --- a/src/hyperloom/orchestrator/actions/executors/_multi_node_server_lifecycle.py +++ b/src/hyperloom/orchestrator/actions/executors/_multi_node_server_lifecycle.py @@ -372,12 +372,22 @@ async def restart_server_for_round( # Multi-node TraceLens SGLang patch fan-out (fail-soft). try: - from ._server_patcher import _tracelens_patch_enabled + from ._server_patcher import _tracelens_patch_enabled, resolve_sglang_shape_mode except Exception: # noqa: BLE001 _tracelens_patch_enabled_fn = lambda: True # noqa: E731 - safe default + _sglang_shape_mode_val = "patched" else: _tracelens_patch_enabled_fn = _tracelens_patch_enabled - if _tracelens_patch_enabled_fn() and (os.environ.get("TRACELENS_ROOT", "").strip()): + _sglang_shape_mode_val = resolve_sglang_shape_mode() + if _sglang_shape_mode_val == "sitecustomize": + # SGLang >= 0.5.18: shapes come from the no-patch kernel_shape_tool + # (PYTHONPATH + TRACELENS_SHAPE_DISCOVERY, wired in launch_multinode / + # infera), so the git-apply fan-out is skipped entirely. + log.info( + "restart_server_for_round: SGLang shape mode=sitecustomize; " + "skipping TraceLens patch fan-out (shapes via kernel_shape_tool)." + ) + elif _tracelens_patch_enabled_fn() and (os.environ.get("TRACELENS_ROOT", "").strip()): try: from hyperloom.inference_optimizer.multi_node.cli import cmd_apply_tracelens_patch diff --git a/src/hyperloom/orchestrator/actions/executors/_server_patcher.py b/src/hyperloom/orchestrator/actions/executors/_server_patcher.py index 17be1b9c05..edbab6c596 100644 --- a/src/hyperloom/orchestrator/actions/executors/_server_patcher.py +++ b/src/hyperloom/orchestrator/actions/executors/_server_patcher.py @@ -213,6 +213,72 @@ def ensure_sglang_patched_for_tracelens( return _ensure_patched(plan) +# SGLang shape-discovery mechanism gate. +# +# From 0.5.18 the kernel shape profiler is delivered as the no-patch TraceLens +# tool (PYTHONPATH + sitecustomize + TRACELENS_SHAPE_DISCOVERY) instead of a +# ``git apply`` of the shape-profiler patches; ``detailed_annotations`` is +# upstream by then, so dropping the patch loses only shapes (the tool restores +# them). Older versions keep the patch mechanism. +_SGLANG_SITECUSTOMIZE_MIN_VERSION: tuple[int, ...] = (0, 5, 18) +_SGLANG_SHAPE_MODE_ENV = "HYPERLOOM_SGLANG_SHAPE_MODE" +# The no-patch kernel shape tool lives in the TraceLens package under +# ``TraceLens/TraceUtils/kernel_shape_tool`` (relative to TRACELENS_ROOT). +_KERNEL_SHAPE_TOOL_REL: tuple[str, ...] = ("TraceLens", "TraceUtils", "kernel_shape_tool") + + +def sglang_shape_mode(version: str) -> str: + """Return the shape-discovery mechanism for an SGLang version. + + ``"sitecustomize"`` (>= 0.5.18) uses the no-patch TraceLens tool; + ``"patched"`` (< 0.5.18) uses the legacy ``git apply`` flow. + ``HYPERLOOM_SGLANG_SHAPE_MODE=patch|sitecustomize`` overrides the gate + (``auto`` / unset = version-based). + """ + override = os.environ.get(_SGLANG_SHAPE_MODE_ENV, "auto").strip().lower() + if override in {"patch", "patched"}: + return "patched" + if override == "sitecustomize": + return "sitecustomize" + vt = _version_tuple(version) + if vt is None: + # Unparseable version: keep the safe legacy mechanism. + return "patched" + return "sitecustomize" if vt >= _SGLANG_SITECUSTOMIZE_MIN_VERSION else "patched" + + +def kernel_shape_tool_dir(tracelens_root: Path | str | None = None) -> Path | None: + """Resolve the no-patch ``kernel_shape_tool`` dir under TRACELENS_ROOT, or None.""" + root = _resolve_tracelens_root(tracelens_root) + if root is None: + return None + tool = root.joinpath(*_KERNEL_SHAPE_TOOL_REL) + return tool if tool.is_dir() else None + + +def _detect_installed_sglang_version() -> str | None: + """Return the locally-installed SGLang version, or ``None`` if unimportable.""" + try: + import sglang # type: ignore # noqa: I001 - runtime probe + except Exception: # noqa: BLE001 + return None + return (getattr(sglang, "__version__", "") or "").strip() or None + + +def resolve_sglang_shape_mode() -> str: + """Resolve the SGLang shape mode from override -> local install -> MN version pin. + + ``HYPERLOOM_SGLANG_SHAPE_MODE`` wins; otherwise the version comes from the + locally-installed SGLang (single-node / sandbox) or, when SGLang is not + importable in the controller (multi-node), ``HYPERLOOM_SGLANG_VERSION_PIN``. + Falls back to ``"patched"`` (legacy) when the version cannot be determined. + """ + version = _detect_installed_sglang_version() or os.environ.get( + "HYPERLOOM_SGLANG_VERSION_PIN", "" + ).strip() + return sglang_shape_mode(version) + + def ensure_sglang_patched_for_ck_blockscale( kernelforge_root: Path | str | None = None, ) -> bool: diff --git a/src/hyperloom/orchestrator/actions/executors/_workload_envs.py b/src/hyperloom/orchestrator/actions/executors/_workload_envs.py index 79601ef21d..a4ecc435a9 100644 --- a/src/hyperloom/orchestrator/actions/executors/_workload_envs.py +++ b/src/hyperloom/orchestrator/actions/executors/_workload_envs.py @@ -62,6 +62,8 @@ ensure_sglang_patched_for_ck_blockscale, ensure_sglang_patched_for_tracelens, ensure_vllm_patched_for_tracelens, + kernel_shape_tool_dir, + resolve_sglang_shape_mode, ) from hyperloom.inference_optimizer.model_config_utils import ( _fp8_is_per_channel_per_token, @@ -1303,7 +1305,12 @@ def materialize_config_with_envs( # try to patch, fall back to the safe set on failure. Default-on # (HYPERLOOM_ENABLE_PATCH=0 disables); skip for atom. tracelens_patch_ok = False - patch_attempted = _tracelens_patch_enabled() and not is_atom + # SGLang >= 0.5.18 uses the no-patch kernel_shape_tool (PYTHONPATH + + # sitecustomize + TRACELENS_SHAPE_DISCOVERY) instead of git-apply; older + # versions keep the patch mechanism. + is_sglang = not is_atom and "vllm" not in fw + sglang_sitecustomize = is_sglang and resolve_sglang_shape_mode() == "sitecustomize" + patch_attempted = _tracelens_patch_enabled() and not is_atom and not sglang_sitecustomize if patch_attempted: if "vllm" in fw: tracelens_patch_ok = ensure_vllm_patched_for_tracelens() @@ -1398,39 +1405,67 @@ def materialize_config_with_envs( "imprecise.", _model, ) - # Both capture options are annotation-only and need TraceLens - # server-side support to land: without it the trace carries no - # ``kernel_shape_profiler`` events (trace-health check 5), so asking - # for them pays the capture cost for data nothing downstream reads. - # Keyed on the degraded *reason* rather than ``tracelens_patch_ok``: - # a patch that was never attempted (HYPERLOOM_ENABLE_PATCH=0) can - # still be baked into the image, and must keep the annotations. - _patch_degraded = envs.get("HYPERLOOM_PROFILE_DEGRADED_REASON") == _TRACELENS_PATCH_UNAVAILABLE - if _patch_degraded: - _shape_disc = False - extra_body["shape_discovery"] = _shape_disc - if _patch_degraded: - extra_body["detailed_annotations"] = False - else: + if sglang_sitecustomize: + # No-patch path (SGLang >= 0.5.18): shapes come from the + # kernel_shape_tool via PYTHONPATH + TRACELENS_SHAPE_DISCOVERY, + # NOT a request-body flag (unpatched SGLang has no + # ``shape_discovery`` field and would reject it), and NOT the + # --enable-shape-discovery-for-cuda-graph-profile arg (unpatched + # SGLang errors on it). ``detailed_annotations`` is upstream, so + # keep it. + extra_body.pop("shape_discovery", None) extra_body.setdefault("detailed_annotations", True) - # NOTE: this write happens before the per-task ``extra_envs`` merge, so - # an ``extra_envs`` entry for PROFILE_EXTRA_BODY can still drop - # start_step/num_steps the way ``args_mode="replace"`` used to drop - # vLLM's --profiler-config bounds. The vLLM side is re-asserted at the - # end of this function; SGLang is NOT, because deciding whether a - # non-positive num_steps means "unbounded" or "no capture" needs a - # SGLang-side answer this layer does not have. Every OOM observed so - # far was vLLM. - envs["PROFILE_EXTRA_BODY"] = _json.dumps(extra_body) - if tracelens_patch_ok and _shape_disc: - # TraceLens-patched SGLang exposes - # --enable-shape-discovery-for-cuda-graph-profile; unpatched - # SGLang errors on it. - existing_sglang = str(envs.get("EXTRA_SGLANG_ARGS", "")) - if "shape-discovery-for-cuda-graph-profile" not in existing_sglang: - envs["EXTRA_SGLANG_ARGS"] = ( - f"{existing_sglang} --enable-shape-discovery-for-cuda-graph-profile" - ).strip() + _tool_dir = kernel_shape_tool_dir() + if _shape_disc and _tool_dir is not None: + _existing_pp = str(envs.get("PYTHONPATH", "")).strip() + envs["PYTHONPATH"] = ( + f"{_tool_dir}{os.pathsep}{_existing_pp}" + if _existing_pp + else str(_tool_dir) + ) + envs["TRACELENS_SHAPE_DISCOVERY"] = "1" + else: + envs["TRACELENS_SHAPE_DISCOVERY"] = "0" + if _shape_disc and _tool_dir is None: + log.warning( + "SGLang shape mode=sitecustomize but kernel_shape_tool " + "not found under TRACELENS_ROOT; shapes will be absent " + "(set TRACELENS_ROOT to an NFS path visible to the server).", + ) + envs["PROFILE_EXTRA_BODY"] = _json.dumps(extra_body) + else: + # Legacy patched path (SGLang < 0.5.18). Both capture options are + # annotation-only and need the TraceLens git-apply patch to land: + # without it the trace carries no ``kernel_shape_profiler`` events + # (trace-health check 5). Keyed on the degraded *reason* rather + # than ``tracelens_patch_ok``: a patch that was never attempted + # (HYPERLOOM_ENABLE_PATCH=0) can still be baked into the image. + _patch_degraded = envs.get("HYPERLOOM_PROFILE_DEGRADED_REASON") == _TRACELENS_PATCH_UNAVAILABLE + if _patch_degraded: + _shape_disc = False + extra_body["shape_discovery"] = _shape_disc + if _patch_degraded: + extra_body["detailed_annotations"] = False + else: + extra_body.setdefault("detailed_annotations", True) + # NOTE: this write happens before the per-task ``extra_envs`` merge, so + # an ``extra_envs`` entry for PROFILE_EXTRA_BODY can still drop + # start_step/num_steps the way ``args_mode="replace"`` used to drop + # vLLM's --profiler-config bounds. The vLLM side is re-asserted at the + # end of this function; SGLang is NOT, because deciding whether a + # non-positive num_steps means "unbounded" or "no capture" needs a + # SGLang-side answer this layer does not have. Every OOM observed so + # far was vLLM. + envs["PROFILE_EXTRA_BODY"] = _json.dumps(extra_body) + if tracelens_patch_ok and _shape_disc: + # TraceLens-patched SGLang exposes + # --enable-shape-discovery-for-cuda-graph-profile; unpatched + # SGLang errors on it. + existing_sglang = str(envs.get("EXTRA_SGLANG_ARGS", "")) + if "shape-discovery-for-cuda-graph-profile" not in existing_sglang: + envs["EXTRA_SGLANG_ARGS"] = ( + f"{existing_sglang} --enable-shape-discovery-for-cuda-graph-profile" + ).strip() if not _is_scriptable_profile: # NUM_PROMPTS / NUM_WARMUPS are serving-request concepts; xDiT drives its diff --git a/src/hyperloom/orchestrator/actions/executors/profile.py b/src/hyperloom/orchestrator/actions/executors/profile.py index 14644f1c3e..5fbe69af06 100644 --- a/src/hyperloom/orchestrator/actions/executors/profile.py +++ b/src/hyperloom/orchestrator/actions/executors/profile.py @@ -577,20 +577,27 @@ def _note_check( skip_reason="main trace could not be sampled", ) else: + # Shape discovery lands via one of two mechanisms and leaves either + # marker: the ``sglang_profiler::`` custom-op namespace (both the legacy + # patch and the no-patch kernel_shape_tool use it) or the + # ``kernel_shape_profiler`` module frame (present when with_stack is on). + _shape_markers = ("sglang_profiler::", "kernel_shape_profiler") + _shape_present = any(m in main_text for m in _shape_markers) _note_check( CHECK_SGLANG_SHAPE_PROFILER, - status="passed" if "kernel_shape_profiler" in main_text else "failed", + status="passed" if _shape_present else "failed", sampled_file=main_traces[0].name, sampled_bytes=_TRACE_INSPECT_BYTES, ) - if "kernel_shape_profiler" not in main_text: + if not _shape_present: issues.append( f"[5] sglang main trace ({main_traces[0].name}, sampled " f"first {_TRACE_INSPECT_BYTES // 1_000_000} MB) lacks " - "kernel_shape_profiler events — shape-discovery " - "patch didn't reach the live SGLang. Verify " - "_server_patcher (PR #207) succeeded for the " - "deployed SGLang version (check log warnings)." + "kernel-shape events — shape discovery didn't reach the live " + "SGLang. For SGLang < 0.5.18 verify the _server_patcher " + "git-apply succeeded; for >= 0.5.18 verify the kernel_shape_tool " + "is on the server PYTHONPATH and TRACELENS_SHAPE_DISCOVERY=1 " + "(check log warnings)." ) if issues: From b6d9de28e0e24ffa2e2d5ca264b729088cac9ab7 Mon Sep 17 00:00:00 2001 From: mohbasit Date: Fri, 11 Sep 2026 10:54:07 -0400 Subject: [PATCH 2/6] fix(profile): make sglang shape-mode imports function-local to break cyclic import Move kernel_shape_tool_dir and resolve_sglang_shape_mode out of the module-level _server_patcher import and into materialize_config_with_envs, clearing the CodeQL module-level cyclic-import alerts on those lines. Matches the function-local pattern in _multi_node_server_lifecycle. Co-authored-by: Cursor --- .../orchestrator/actions/executors/_workload_envs.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/src/hyperloom/orchestrator/actions/executors/_workload_envs.py b/src/hyperloom/orchestrator/actions/executors/_workload_envs.py index a4ecc435a9..2a1c68dfb6 100644 --- a/src/hyperloom/orchestrator/actions/executors/_workload_envs.py +++ b/src/hyperloom/orchestrator/actions/executors/_workload_envs.py @@ -62,8 +62,6 @@ ensure_sglang_patched_for_ck_blockscale, ensure_sglang_patched_for_tracelens, ensure_vllm_patched_for_tracelens, - kernel_shape_tool_dir, - resolve_sglang_shape_mode, ) from hyperloom.inference_optimizer.model_config_utils import ( _fp8_is_per_channel_per_token, @@ -1307,7 +1305,12 @@ def materialize_config_with_envs( tracelens_patch_ok = False # SGLang >= 0.5.18 uses the no-patch kernel_shape_tool (PYTHONPATH + # sitecustomize + TRACELENS_SHAPE_DISCOVERY) instead of git-apply; older - # versions keep the patch mechanism. + # versions keep the patch mechanism. Imported function-locally (not at + # module top) so these symbols stay out of the module-level import cycle + # ``_workload_envs`` -> executors package __init__ -> baseline -> + # ``_workload_envs`` (matches ``_multi_node_server_lifecycle``). + from ._server_patcher import kernel_shape_tool_dir, resolve_sglang_shape_mode + is_sglang = not is_atom and "vllm" not in fw sglang_sitecustomize = is_sglang and resolve_sglang_shape_mode() == "sitecustomize" patch_attempted = _tracelens_patch_enabled() and not is_atom and not sglang_sitecustomize From 0e2e774a7c0a06ef5697d0781065b6c43cf489c6 Mon Sep 17 00:00:00 2001 From: mohbasit Date: Mon, 14 Sep 2026 16:11:51 -0400 Subject: [PATCH 3/6] fix lint --- .../multi_node/scripts/launch_multinode.py | 4 +--- .../orchestrator/actions/executors/_server_patcher.py | 4 +--- .../orchestrator/actions/executors/_workload_envs.py | 6 +----- 3 files changed, 3 insertions(+), 11 deletions(-) diff --git a/src/hyperloom/inference_optimizer/multi_node/scripts/launch_multinode.py b/src/hyperloom/inference_optimizer/multi_node/scripts/launch_multinode.py index e5866a73d7..7b12ec3cf9 100755 --- a/src/hyperloom/inference_optimizer/multi_node/scripts/launch_multinode.py +++ b/src/hyperloom/inference_optimizer/multi_node/scripts/launch_multinode.py @@ -346,9 +346,7 @@ def _maybe_activate_kernel_shape_tool(sub_env: dict[str, str]) -> None: return tool = Path(root).joinpath(*_KERNEL_SHAPE_TOOL_REL) if not tool.is_dir(): - sys.stderr.write( - f"WARN kernel_shape_tool not found at {tool}; SGLang shape discovery disabled\n" - ) + sys.stderr.write(f"WARN kernel_shape_tool not found at {tool}; SGLang shape discovery disabled\n") return existing = sub_env.get("PYTHONPATH", "").strip() sub_env["PYTHONPATH"] = f"{tool}{os.pathsep}{existing}" if existing else str(tool) diff --git a/src/hyperloom/orchestrator/actions/executors/_server_patcher.py b/src/hyperloom/orchestrator/actions/executors/_server_patcher.py index edbab6c596..594310c47f 100644 --- a/src/hyperloom/orchestrator/actions/executors/_server_patcher.py +++ b/src/hyperloom/orchestrator/actions/executors/_server_patcher.py @@ -273,9 +273,7 @@ def resolve_sglang_shape_mode() -> str: importable in the controller (multi-node), ``HYPERLOOM_SGLANG_VERSION_PIN``. Falls back to ``"patched"`` (legacy) when the version cannot be determined. """ - version = _detect_installed_sglang_version() or os.environ.get( - "HYPERLOOM_SGLANG_VERSION_PIN", "" - ).strip() + version = _detect_installed_sglang_version() or os.environ.get("HYPERLOOM_SGLANG_VERSION_PIN", "").strip() return sglang_shape_mode(version) diff --git a/src/hyperloom/orchestrator/actions/executors/_workload_envs.py b/src/hyperloom/orchestrator/actions/executors/_workload_envs.py index 2a1c68dfb6..ae1a46452e 100644 --- a/src/hyperloom/orchestrator/actions/executors/_workload_envs.py +++ b/src/hyperloom/orchestrator/actions/executors/_workload_envs.py @@ -1421,11 +1421,7 @@ def materialize_config_with_envs( _tool_dir = kernel_shape_tool_dir() if _shape_disc and _tool_dir is not None: _existing_pp = str(envs.get("PYTHONPATH", "")).strip() - envs["PYTHONPATH"] = ( - f"{_tool_dir}{os.pathsep}{_existing_pp}" - if _existing_pp - else str(_tool_dir) - ) + envs["PYTHONPATH"] = f"{_tool_dir}{os.pathsep}{_existing_pp}" if _existing_pp else str(_tool_dir) envs["TRACELENS_SHAPE_DISCOVERY"] = "1" else: envs["TRACELENS_SHAPE_DISCOVERY"] = "0" From 58d3693fcccf0405d88742766cb81b6212c89aee Mon Sep 17 00:00:00 2001 From: mohbasit Date: Mon, 14 Sep 2026 17:28:37 -0400 Subject: [PATCH 4/6] reduce comments --- .../multi_node/commands/infera.py | 6 +-- .../multi_node/scripts/launch_infera_node.py | 4 +- .../multi_node/scripts/launch_multinode.py | 4 +- .../executors/_multi_node_server_lifecycle.py | 4 +- .../actions/executors/_server_patcher.py | 12 ++---- .../actions/executors/_workload_envs.py | 42 ++++++------------- .../orchestrator/actions/executors/profile.py | 6 +-- 7 files changed, 22 insertions(+), 56 deletions(-) diff --git a/src/hyperloom/inference_optimizer/multi_node/commands/infera.py b/src/hyperloom/inference_optimizer/multi_node/commands/infera.py index 2f046320c6..67c5cd1d6e 100644 --- a/src/hyperloom/inference_optimizer/multi_node/commands/infera.py +++ b/src/hyperloom/inference_optimizer/multi_node/commands/infera.py @@ -90,10 +90,8 @@ def _collect_forward_env() -> dict[str, str]: trace_dir = os.environ.get("HYPERLOOM_MN_PROFILE_TRACE_DIR", "").strip() if trace_dir and "SGLANG_TORCH_PROFILER_DIR" not in fwd: fwd["SGLANG_TORCH_PROFILER_DIR"] = trace_dir - # SGLang >= 0.5.18 no-patch shape discovery: the pod-side launcher - # (launch_infera_node.py) puts kernel_shape_tool on PYTHONPATH from - # TRACELENS_ROOT and honors these flags. PYTHONPATH itself is blocked from - # SSH forwarding (BLOCKED_UNTRUSTED_ENV_NAMES) and is set pod-side instead. + # Forward no-patch shape-discovery config; the pod-side launcher sets + # PYTHONPATH itself (it is blocked from SSH forwarding). for _shape_key in ( "TRACELENS_ROOT", "TRACELENS_SHAPE_DISCOVERY", diff --git a/src/hyperloom/inference_optimizer/multi_node/scripts/launch_infera_node.py b/src/hyperloom/inference_optimizer/multi_node/scripts/launch_infera_node.py index fbaa2e49e3..5a3403f1dc 100644 --- a/src/hyperloom/inference_optimizer/multi_node/scripts/launch_infera_node.py +++ b/src/hyperloom/inference_optimizer/multi_node/scripts/launch_infera_node.py @@ -412,9 +412,7 @@ def _reap_stale_engine_ports() -> None: break -# SGLang >= 0.5.18 uses the no-patch kernel_shape_tool (PYTHONPATH + -# sitecustomize + TRACELENS_SHAPE_DISCOVERY) for shape discovery. This script is -# standalone (no hyperloom import), so the gate is mirrored inline. +# SGLang >= 0.5.18 no-patch shape tool; gate mirrored inline (no hyperloom import). _KERNEL_SHAPE_TOOL_REL = ("TraceLens", "TraceUtils", "kernel_shape_tool") _SGLANG_SITECUSTOMIZE_MIN_VERSION = (0, 5, 18) diff --git a/src/hyperloom/inference_optimizer/multi_node/scripts/launch_multinode.py b/src/hyperloom/inference_optimizer/multi_node/scripts/launch_multinode.py index 7b12ec3cf9..e03aa20c3f 100755 --- a/src/hyperloom/inference_optimizer/multi_node/scripts/launch_multinode.py +++ b/src/hyperloom/inference_optimizer/multi_node/scripts/launch_multinode.py @@ -311,9 +311,7 @@ def _probe_mec_firmware_lt_177() -> bool: return False -# SGLang >= 0.5.18 uses the no-patch kernel_shape_tool (PYTHONPATH + -# sitecustomize + TRACELENS_SHAPE_DISCOVERY) for shape discovery. This script is -# standalone (no hyperloom import), so the gate is mirrored inline. +# SGLang >= 0.5.18 no-patch shape tool; gate mirrored inline (no hyperloom import). _KERNEL_SHAPE_TOOL_REL = ("TraceLens", "TraceUtils", "kernel_shape_tool") _SGLANG_SITECUSTOMIZE_MIN_VERSION = (0, 5, 18) diff --git a/src/hyperloom/orchestrator/actions/executors/_multi_node_server_lifecycle.py b/src/hyperloom/orchestrator/actions/executors/_multi_node_server_lifecycle.py index 105fa29285..b5d26803fe 100644 --- a/src/hyperloom/orchestrator/actions/executors/_multi_node_server_lifecycle.py +++ b/src/hyperloom/orchestrator/actions/executors/_multi_node_server_lifecycle.py @@ -380,9 +380,7 @@ async def restart_server_for_round( _tracelens_patch_enabled_fn = _tracelens_patch_enabled _sglang_shape_mode_val = resolve_sglang_shape_mode() if _sglang_shape_mode_val == "sitecustomize": - # SGLang >= 0.5.18: shapes come from the no-patch kernel_shape_tool - # (PYTHONPATH + TRACELENS_SHAPE_DISCOVERY, wired in launch_multinode / - # infera), so the git-apply fan-out is skipped entirely. + # sitecustomize mode: shapes come from the no-patch tool; skip the patch fan-out. log.info( "restart_server_for_round: SGLang shape mode=sitecustomize; " "skipping TraceLens patch fan-out (shapes via kernel_shape_tool)." diff --git a/src/hyperloom/orchestrator/actions/executors/_server_patcher.py b/src/hyperloom/orchestrator/actions/executors/_server_patcher.py index 594310c47f..3e36d9e394 100644 --- a/src/hyperloom/orchestrator/actions/executors/_server_patcher.py +++ b/src/hyperloom/orchestrator/actions/executors/_server_patcher.py @@ -213,17 +213,11 @@ def ensure_sglang_patched_for_tracelens( return _ensure_patched(plan) -# SGLang shape-discovery mechanism gate. -# -# From 0.5.18 the kernel shape profiler is delivered as the no-patch TraceLens -# tool (PYTHONPATH + sitecustomize + TRACELENS_SHAPE_DISCOVERY) instead of a -# ``git apply`` of the shape-profiler patches; ``detailed_annotations`` is -# upstream by then, so dropping the patch loses only shapes (the tool restores -# them). Older versions keep the patch mechanism. +# SGLang shape-discovery gate: >= 0.5.18 uses the no-patch TraceLens tool +# (PYTHONPATH + sitecustomize + TRACELENS_SHAPE_DISCOVERY); older uses git-apply. _SGLANG_SITECUSTOMIZE_MIN_VERSION: tuple[int, ...] = (0, 5, 18) _SGLANG_SHAPE_MODE_ENV = "HYPERLOOM_SGLANG_SHAPE_MODE" -# The no-patch kernel shape tool lives in the TraceLens package under -# ``TraceLens/TraceUtils/kernel_shape_tool`` (relative to TRACELENS_ROOT). +# No-patch tool location, relative to TRACELENS_ROOT. _KERNEL_SHAPE_TOOL_REL: tuple[str, ...] = ("TraceLens", "TraceUtils", "kernel_shape_tool") diff --git a/src/hyperloom/orchestrator/actions/executors/_workload_envs.py b/src/hyperloom/orchestrator/actions/executors/_workload_envs.py index ae1a46452e..3a90466156 100644 --- a/src/hyperloom/orchestrator/actions/executors/_workload_envs.py +++ b/src/hyperloom/orchestrator/actions/executors/_workload_envs.py @@ -1303,12 +1303,8 @@ def materialize_config_with_envs( # try to patch, fall back to the safe set on failure. Default-on # (HYPERLOOM_ENABLE_PATCH=0 disables); skip for atom. tracelens_patch_ok = False - # SGLang >= 0.5.18 uses the no-patch kernel_shape_tool (PYTHONPATH + - # sitecustomize + TRACELENS_SHAPE_DISCOVERY) instead of git-apply; older - # versions keep the patch mechanism. Imported function-locally (not at - # module top) so these symbols stay out of the module-level import cycle - # ``_workload_envs`` -> executors package __init__ -> baseline -> - # ``_workload_envs`` (matches ``_multi_node_server_lifecycle``). + # Function-local import to stay out of the module-level import cycle + # (matches _multi_node_server_lifecycle). from ._server_patcher import kernel_shape_tool_dir, resolve_sglang_shape_mode is_sglang = not is_atom and "vllm" not in fw @@ -1409,13 +1405,8 @@ def materialize_config_with_envs( _model, ) if sglang_sitecustomize: - # No-patch path (SGLang >= 0.5.18): shapes come from the - # kernel_shape_tool via PYTHONPATH + TRACELENS_SHAPE_DISCOVERY, - # NOT a request-body flag (unpatched SGLang has no - # ``shape_discovery`` field and would reject it), and NOT the - # --enable-shape-discovery-for-cuda-graph-profile arg (unpatched - # SGLang errors on it). ``detailed_annotations`` is upstream, so - # keep it. + # No-patch path: shapes come from the tool via PYTHONPATH, not a + # request-body flag or CUDA-graph arg (unpatched SGLang rejects both). extra_body.pop("shape_discovery", None) extra_body.setdefault("detailed_annotations", True) _tool_dir = kernel_shape_tool_dir() @@ -1433,12 +1424,9 @@ def materialize_config_with_envs( ) envs["PROFILE_EXTRA_BODY"] = _json.dumps(extra_body) else: - # Legacy patched path (SGLang < 0.5.18). Both capture options are - # annotation-only and need the TraceLens git-apply patch to land: - # without it the trace carries no ``kernel_shape_profiler`` events - # (trace-health check 5). Keyed on the degraded *reason* rather - # than ``tracelens_patch_ok``: a patch that was never attempted - # (HYPERLOOM_ENABLE_PATCH=0) can still be baked into the image. + # Legacy patched path: capture options need the git-apply patch to + # land. Keyed on the degraded reason, not tracelens_patch_ok, since a + # patch may be baked into the image without being attempted here. _patch_degraded = envs.get("HYPERLOOM_PROFILE_DEGRADED_REASON") == _TRACELENS_PATCH_UNAVAILABLE if _patch_degraded: _shape_disc = False @@ -1447,19 +1435,13 @@ def materialize_config_with_envs( extra_body["detailed_annotations"] = False else: extra_body.setdefault("detailed_annotations", True) - # NOTE: this write happens before the per-task ``extra_envs`` merge, so - # an ``extra_envs`` entry for PROFILE_EXTRA_BODY can still drop - # start_step/num_steps the way ``args_mode="replace"`` used to drop - # vLLM's --profiler-config bounds. The vLLM side is re-asserted at the - # end of this function; SGLang is NOT, because deciding whether a - # non-positive num_steps means "unbounded" or "no capture" needs a - # SGLang-side answer this layer does not have. Every OOM observed so - # far was vLLM. + # Written before the per-task extra_envs merge, so an extra_envs + # PROFILE_EXTRA_BODY can still drop start_step/num_steps. Not + # re-asserted for SGLang (unlike vLLM): "unbounded" vs "no capture" + # for non-positive num_steps needs a SGLang-side answer. envs["PROFILE_EXTRA_BODY"] = _json.dumps(extra_body) if tracelens_patch_ok and _shape_disc: - # TraceLens-patched SGLang exposes - # --enable-shape-discovery-for-cuda-graph-profile; unpatched - # SGLang errors on it. + # Patched SGLang exposes this arg; unpatched errors on it. existing_sglang = str(envs.get("EXTRA_SGLANG_ARGS", "")) if "shape-discovery-for-cuda-graph-profile" not in existing_sglang: envs["EXTRA_SGLANG_ARGS"] = ( diff --git a/src/hyperloom/orchestrator/actions/executors/profile.py b/src/hyperloom/orchestrator/actions/executors/profile.py index 5fbe69af06..4ac6bab565 100644 --- a/src/hyperloom/orchestrator/actions/executors/profile.py +++ b/src/hyperloom/orchestrator/actions/executors/profile.py @@ -577,10 +577,8 @@ def _note_check( skip_reason="main trace could not be sampled", ) else: - # Shape discovery lands via one of two mechanisms and leaves either - # marker: the ``sglang_profiler::`` custom-op namespace (both the legacy - # patch and the no-patch kernel_shape_tool use it) or the - # ``kernel_shape_profiler`` module frame (present when with_stack is on). + # Shape markers left by either mechanism: the sglang_profiler:: op + # namespace, or the kernel_shape_profiler frame (when with_stack is on). _shape_markers = ("sglang_profiler::", "kernel_shape_profiler") _shape_present = any(m in main_text for m in _shape_markers) _note_check( From 66b7b0d6870459fa2e7cb22dbaf06849c9623d9a Mon Sep 17 00:00:00 2001 From: mohbasit Date: Mon, 14 Sep 2026 18:35:57 -0400 Subject: [PATCH 5/6] sgalng check --- src/hyperloom/orchestrator/actions/executors/_workload_envs.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/hyperloom/orchestrator/actions/executors/_workload_envs.py b/src/hyperloom/orchestrator/actions/executors/_workload_envs.py index 3a90466156..42e9f79988 100644 --- a/src/hyperloom/orchestrator/actions/executors/_workload_envs.py +++ b/src/hyperloom/orchestrator/actions/executors/_workload_envs.py @@ -1307,7 +1307,7 @@ def materialize_config_with_envs( # (matches _multi_node_server_lifecycle). from ._server_patcher import kernel_shape_tool_dir, resolve_sglang_shape_mode - is_sglang = not is_atom and "vllm" not in fw + is_sglang = "sglang" in fw sglang_sitecustomize = is_sglang and resolve_sglang_shape_mode() == "sitecustomize" patch_attempted = _tracelens_patch_enabled() and not is_atom and not sglang_sitecustomize if patch_attempted: From 27d62c08aaa741f7013a78089303477290697096 Mon Sep 17 00:00:00 2001 From: mohammad abdul basit Date: Thu, 17 Sep 2026 21:02:36 +0000 Subject: [PATCH 6/6] change TraceLens ref --- src/hyperloom/agents/kernel/scripts/install.sh | 2 +- src/hyperloom/agents/kernel/tools/tracelens_analysis.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/hyperloom/agents/kernel/scripts/install.sh b/src/hyperloom/agents/kernel/scripts/install.sh index 515fbe5fa1..dac4444c83 100644 --- a/src/hyperloom/agents/kernel/scripts/install.sh +++ b/src/hyperloom/agents/kernel/scripts/install.sh @@ -156,7 +156,7 @@ INFERENCEX_PATH="${INFERENCEX_PATH:-}" # The internal extension is used ONLY when $TRACELENS_INTERNAL_ROOT is set # (env / .env); leave it unset for the base-only report. No separate toggle. TRACELENS_REPO="https://github.com/AMD-AGI/TraceLens.git" -TRACELENS_REF="c74d4d2ca48d6fcd7e7e829b409446000fe4300f" +TRACELENS_REF="9fc0dc6487bde554c6ed314a15b61022e5ec62ea" # Operator override iff TRACELENS_ROOT points OUTSIDE the pod-local default. # The persistent kernel-agent env re-exports the resolved default path, so a # presence-only check (${VAR:+1}) would misclassify it as an override and skip diff --git a/src/hyperloom/agents/kernel/tools/tracelens_analysis.py b/src/hyperloom/agents/kernel/tools/tracelens_analysis.py index fa8bb1ed76..338ed6450c 100755 --- a/src/hyperloom/agents/kernel/tools/tracelens_analysis.py +++ b/src/hyperloom/agents/kernel/tools/tracelens_analysis.py @@ -5685,7 +5685,7 @@ def run_command( # Defaults kept in sync with src/hyperloom/agents/kernel/scripts/install.sh (TRACELENS_REPO / # TRACELENS_REF). Overridable via env so a run can pin its own SHA. _TRACELENS_REPO_DEFAULT = "https://github.com/AMD-AGI/TraceLens.git" -_TRACELENS_REF_DEFAULT = "c74d4d2ca48d6fcd7e7e829b409446000fe4300f" +_TRACELENS_REF_DEFAULT = "9fc0dc6487bde554c6ed314a15b61022e5ec62ea" def _default_tracelens_root() -> Path: