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: diff --git a/src/hyperloom/inference_optimizer/multi_node/commands/infera.py b/src/hyperloom/inference_optimizer/multi_node/commands/infera.py index 36bc709df4..67c5cd1d6e 100644 --- a/src/hyperloom/inference_optimizer/multi_node/commands/infera.py +++ b/src/hyperloom/inference_optimizer/multi_node/commands/infera.py @@ -90,6 +90,17 @@ 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 + # 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", + "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..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 @@ -6,6 +6,7 @@ import argparse import json import os +import re import shlex import subprocess import sys @@ -411,6 +412,46 @@ def _reap_stale_engine_ports() -> None: break +# 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) + + +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 +765,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..e03aa20c3f 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,46 @@ def _probe_mec_firmware_lt_177() -> bool: return False +# 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) + + +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 +484,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..b5d26803fe 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,20 @@ 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": + # 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)." + ) + 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 86d37c45eb..5d884943db 100644 --- a/src/hyperloom/orchestrator/actions/executors/_server_patcher.py +++ b/src/hyperloom/orchestrator/actions/executors/_server_patcher.py @@ -253,6 +253,64 @@ def ensure_sglang_patched_for_tracelens( return _ensure_patched(plan) +# 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" +# No-patch tool location, 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 7a2fcf2a2e..13a1e791a5 100644 --- a/src/hyperloom/orchestrator/actions/executors/_workload_envs.py +++ b/src/hyperloom/orchestrator/actions/executors/_workload_envs.py @@ -1513,7 +1513,13 @@ 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 + # 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 = "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 # Written in every branch, not only the failing one. "No status" used to mean both "patched fine" and # "never tried because the image already carries it", and those two call for different reactions when a # trace later turns up without annotations. @@ -1612,39 +1618,49 @@ 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: 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) - # 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: 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 + extra_body["shape_discovery"] = _shape_disc + if _patch_degraded: + extra_body["detailed_annotations"] = False + else: + extra_body.setdefault("detailed_annotations", True) + # 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: + # 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"] = ( + 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 0f7644834e..2da0f74296 100644 --- a/src/hyperloom/orchestrator/actions/executors/profile.py +++ b/src/hyperloom/orchestrator/actions/executors/profile.py @@ -660,20 +660,25 @@ def _note_check( skip_reason="main trace could not be sampled", ) else: + # 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( 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: