"""Shared helpers for unified LLM torch-profiler skill scripts.""" from __future__ import annotations import gzip import json import re import shutil import sys import tempfile import time from collections import Counter, defaultdict from dataclasses import dataclass from functools import lru_cache from pathlib import Path from typing import Callable, Dict, Iterable, List, Optional, Sequence, Tuple from urllib import request STAGE_ORDER = {"extend": 0, "prefill": 0, "decode": 1, "all": 2} FRAMEWORK_LABELS = { "auto": "auto", "sglang": "SGLang", "vllm": "vLLM", "trtllm": "TensorRT-LLM", "tokenspeed": "TokenSpeed", } TRACE_FILE_PATTERNS = ( "*.trace.json", "*.trace.json.gz", "*.pt.trace.json", "*.pt.trace.json.gz", "*.json", "*.json.gz", ) TRACE_FILE_IGNORE_NAMES = { "server_args.json", "metadata.json", "config.json", } TRACE_METADATA_NAMES = { "process_name", "thread_name", "process_sort_index", "thread_sort_index", } NON_KERNEL_TRACE_CATEGORIES = ("python_function", "cpu_op", "trace") PYTHON_SCOPE_NAME_PREFIXES = ("python/", "nn.module:") PROFILE_WORKLOAD_CHOICES = ("legacy", "prefill", "decode", "both") DEFAULT_PREFILL_INPUT_LEN = 4090 DEFAULT_PREFILL_OUTPUT_LEN = 1 DEFAULT_DECODE_INPUT_LEN = 1 DEFAULT_DECODE_OUTPUT_LEN = 2048 DEFAULT_WARMUP_STEPS = 10 @dataclass(frozen=True) class ProbePlan: prompt: str capture_max_new_tokens: int capture_requests: int warmup_max_new_tokens: int warmup_requests: int @lru_cache(maxsize=65536) def _normalize_text_cached(text: str) -> str: text = text.strip() if not text: return "" for token in (" ", "\t", "\n", "\r", "\v", "\f"): if token in text: return " ".join(text.split()) return text def normalize_text(value: object) -> str: return _normalize_text_cached(value if isinstance(value, str) else str(value)) def canonicalize_framework(value: object) -> str: lowered = normalize_text(value).lower().replace("_", "-") aliases = { "": "auto", "auto": "auto", "sglang": "sglang", "sgl": "sglang", "vllm": "vllm", "trt": "trtllm", "tllm": "trtllm", "trtllm": "trtllm", "tensorrt-llm": "trtllm", "tensorrtllm": "trtllm", "tokenspeed": "tokenspeed", "token-speed": "tokenspeed", "ts": "tokenspeed", } return aliases.get(lowered, "auto") def framework_display_name(value: object) -> str: return FRAMEWORK_LABELS.get(canonicalize_framework(value), str(value)) @lru_cache(maxsize=65536) def _normalize_repo_relative_path_cached(text: str) -> str: text = text.replace("\\", "/") lowered = text.lower() for marker, normalized_marker in ( ("python/sglang/", "python/sglang/"), ("sgl_kernel/", "sgl_kernel/"), ("vllm/", "vllm/"), ("python/tokenspeed/", "python/tokenspeed/"), ("tokenspeed/", "tokenspeed/"), ("tensorrt_llm/", "tensorrt_llm/"), ("tensorrt-llm/", "tensorrt_llm/"), ): idx = lowered.find(marker) if idx != -1: suffix = text[idx + len(marker) :].lstrip("/") return f"{normalized_marker}{suffix}".lstrip("/") idx = lowered.find("sglang/") if idx != -1: return ("python/" + text[idx:]).lstrip("/") return text.lstrip("/") def normalize_repo_relative_path(path: object) -> str: return _normalize_repo_relative_path_cached(normalize_text(path)) def contains_any_keyword(text: str, keywords: Iterable[str]) -> bool: return any(keyword in text for keyword in keywords) def coerce_optional_int(value: object) -> Optional[int]: if value in (None, "", "None"): return None if isinstance(value, int): return value if isinstance(value, float): return int(value) if value.is_integer() else None try: return int(str(value)) except (TypeError, ValueError): return None def extract_trace_events(trace: object) -> Sequence[dict]: if isinstance(trace, dict): events = trace.get("traceEvents", []) return events if isinstance(events, list) else [] if isinstance(trace, list): return trace return [] def is_trace_metadata_name(name: object) -> bool: return str(name) in TRACE_METADATA_NAMES def is_complete_duration_event(event: dict) -> bool: if event.get("ph") != "X": return False dur = event.get("dur") ts = event.get("ts") if dur is None or ts is None: return False try: return float(dur) > 0 except (TypeError, ValueError): return False def is_annotation_event(name: object, category: object) -> bool: lowered_name = normalize_text(name).lower() lowered_category = normalize_text(category).lower() return "annotation" in lowered_category or lowered_name.startswith("## call ") def is_non_kernel_trace_category(category: object) -> bool: lowered_category = normalize_text(category).lower() return any(token in lowered_category for token in NON_KERNEL_TRACE_CATEGORIES) def looks_like_python_scope_name(name: object) -> bool: lowered_name = normalize_text(name).lower() return ".py(" in lowered_name or lowered_name.startswith(PYTHON_SCOPE_NAME_PREFIXES) def has_stream_marker(args: Optional[dict]) -> bool: trace_args = args or {} return "stream" in trace_args or "cuda_stream" in trace_args def load_trace_json(path: Path) -> dict: if path.suffix == ".gz": with gzip.open(path, "rt", encoding="utf-8") as handle: return json.load(handle) with open(path, "r", encoding="utf-8") as handle: return json.load(handle) def load_server_args(path: Path) -> Optional[dict]: resolved = path.resolve() candidate_dirs: List[Path] = [] if resolved.is_file(): candidate_dirs.extend([resolved.parent, resolved.parent.parent]) else: candidate_dirs.extend([resolved, resolved.parent]) seen: set[Path] = set() for candidate_dir in candidate_dirs: if candidate_dir in seen: continue seen.add(candidate_dir) candidate = candidate_dir / "server_args.json" if candidate.exists(): with open(candidate, "r", encoding="utf-8") as handle: return json.load(handle) return None def try_get_json(url: str, timeout: float = 60.0) -> Optional[object]: try: with request.urlopen(url, timeout=timeout) as response: raw = response.read() except Exception: return None if not raw: return None try: return json.loads(raw.decode("utf-8")) except json.JSONDecodeError: return None def _flatten_chat_text_parts(value: object) -> List[str]: if value is None: return [] if isinstance(value, str): text = value.strip() return [text] if text else [] if isinstance(value, list): parts: List[str] = [] for item in value: parts.extend(_flatten_chat_text_parts(item)) return parts if isinstance(value, dict): parts: List[str] = [] text_keys = ( "text", "content", "reasoning_content", "reasoning", "output_text", ) if any(key in value for key in text_keys): for key in text_keys: parts.extend(_flatten_chat_text_parts(value.get(key))) if parts: return parts item_type = normalize_text(value.get("type")).lower() if item_type in {"text", "output_text", "input_text"}: for key in ("text", "content", "value"): parts.extend(_flatten_chat_text_parts(value.get(key))) elif item_type in {"reasoning", "thinking"}: for key in ("text", "content", "reasoning_content", "reasoning"): parts.extend(_flatten_chat_text_parts(value.get(key))) return parts return [] def flatten_chat_text(value: object) -> str: return "\n".join(_flatten_chat_text_parts(value)).strip() def extract_openai_chat_text(body: object) -> Tuple[str, str]: if not isinstance(body, dict): return "", "invalid_body" choices = body.get("choices") if not isinstance(choices, list) or not choices: fallback = flatten_chat_text(body.get("output_text")) if fallback: return fallback, "body.output_text" return "", "missing_choices" first_choice = choices[0] if not isinstance(first_choice, dict): return "", "invalid_choice" message = first_choice.get("message") if isinstance(message, dict): for key in ("content", "reasoning_content", "reasoning"): text = flatten_chat_text(message.get(key)) if text: return text, f"message.{key}" for key in ("text", "content", "reasoning_content", "reasoning"): text = flatten_chat_text(first_choice.get(key)) if text: return text, f"choice.{key}" delta = first_choice.get("delta") if isinstance(delta, dict): for key in ("content", "reasoning_content", "reasoning"): text = flatten_chat_text(delta.get(key)) if text: return text, f"delta.{key}" fallback = flatten_chat_text(body.get("output_text")) if fallback: return fallback, "body.output_text" return "", "empty" def detect_framework_from_text(text: object) -> Optional[str]: lowered = normalize_text(text).lower() if not lowered: return None if any(token in lowered for token in ("tokenspeed", "token-speed", "/ts/")): return "tokenspeed" if any( token in lowered for token in ( "tensorrt_llm", "tensorrt-llm", "trtllm", "pyexecutor", ) ): return "trtllm" if "vllm" in lowered: return "vllm" if any(token in lowered for token in ("python/sglang/", "sgl_kernel/", "sglang/")): return "sglang" return None def detect_framework_from_server_args(server_args: Optional[dict]) -> Optional[str]: if not isinstance(server_args, dict) or not server_args: return None lowered_keys = {normalize_text(key).lower() for key in server_args} text = json.dumps(server_args, sort_keys=True) if any(token in text.lower() for token in ("tokenspeed", "token-speed")): return "tokenspeed" if lowered_keys & { "attn_tp_size", "dense_tp_size", "moe_tp_size", "enable_mla_l1_5_cache", "mla_chunk_multiplier", "comm_fusion_max_num_tokens", "enable_allreduce_fusion", }: return "tokenspeed" text_hint = detect_framework_from_text(text) if text_hint: return text_hint if lowered_keys & { "attention_backend", "sampling_backend", "disable_cuda_graph", "disable_piecewise_cuda_graph", "chunked_prefill_size", "schedule_policy", }: return "sglang" return None def detect_framework_from_trace(trace: object) -> Optional[str]: text_samples: List[str] = [] for event in extract_trace_events(trace)[:256]: text_samples.extend( [ str(event.get("name", "")), str(event.get("cat", "")), str(event.get("pid", "")), ] ) trace_args = event.get("args") if isinstance(trace_args, dict): for key, value in list(trace_args.items())[:8]: text_samples.append(str(key)) if isinstance(value, str): text_samples.append(value) return detect_framework_from_text(" ".join(text_samples)) def detect_framework_from_path(path: Path) -> Optional[str]: hint = detect_framework_from_text(str(path)) if hint: return hint server_args = load_server_args(path) hint = detect_framework_from_server_args(server_args) if hint: return hint if path.is_file(): try: return detect_framework_from_trace(load_trace_json(path)) except Exception: return None trace_files = discover_trace_files(path, recursive=True, limit=3) for trace_file in trace_files: try: hint = detect_framework_from_trace(load_trace_json(trace_file)) except Exception: hint = None if hint: return hint return None def detect_framework_from_url( url: str, output_dir: Optional[str] = None ) -> Optional[str]: hint = detect_framework_from_text(output_dir or "") if hint: return hint server_info = try_get_json(url.rstrip("/") + "/server_info") if isinstance(server_info, dict) and ( "internal_states" in server_info or "tokenizer_path" in server_info or "prefill" in server_info or "decode" in server_info ): return "sglang" readiness = try_get_json(url.rstrip("/") + "/readiness", timeout=5.0) if readiness is not None: return "tokenspeed" models = try_get_json(url.rstrip("/") + "/v1/models") if isinstance(models, dict) and isinstance(models.get("data"), list): return "vllm" return None def resolve_framework( requested: object, *, input_path: Optional[Path] = None, url: Optional[str] = None, server_args: Optional[dict] = None, ) -> str: explicit = canonicalize_framework(requested) if explicit != "auto": return explicit for hint in ( detect_framework_from_server_args(server_args), detect_framework_from_path(input_path) if input_path else None, ( detect_framework_from_url(url, str(input_path) if input_path else None) if url else None ), ): if hint: return hint return "sglang" def parse_stage(path: Path) -> str: parts = [part.lower() for part in path.parts[-6:]] name = " ".join(parts) segment_path = "/" + "/".join(parts) + "/" if any(marker in name for marker in ("-extend", "-prefill", "_extend", "_prefill")): return "extend" if any(f"/{segment}/" in segment_path for segment in ("extend", "prefill")): return "extend" if any(marker in name for marker in ("-decode", "_decode")): return "decode" if "/decode/" in segment_path: return "decode" return "all" def parse_tp_rank(path: Path) -> Optional[int]: for pattern in ( r"(?:^|[_-])tp(\d+)(?:[_.-]|$)", r"TP-(\d+)", r"(?:^|[_-])rank(\d+)(?:[_.-]|$)", r"(?:^|[_-])worker(\d+)(?:[_.-]|$)", ): match = re.search(pattern, path.name, re.IGNORECASE) if match: return int(match.group(1)) return None def file_looks_like_trace(path: Path) -> bool: name = path.name.lower() if name in TRACE_FILE_IGNORE_NAMES: return False if path.is_dir(): return False if any(name.endswith(suffix) for suffix in (".trace.json", ".trace.json.gz")): return True if ".pt.trace.json" in name: return True if not any(name.endswith(suffix) for suffix in (".json", ".json.gz")): return False try: trace = load_trace_json(path) except Exception: return False if isinstance(trace, dict): return isinstance(trace.get("traceEvents"), list) if isinstance(trace, list): return bool(trace) and all(isinstance(item, dict) for item in trace[:8]) return False def discover_trace_files( path: Path, *, recursive: bool, limit: Optional[int] = None, ) -> List[Path]: if path.is_file(): return [path] if file_looks_like_trace(path) else [] candidates: List[Path] = [] seen: set[Path] = set() for pattern in TRACE_FILE_PATTERNS: iterator = path.rglob(pattern) if recursive else path.glob(pattern) for candidate in iterator: resolved = candidate.resolve() if resolved in seen: continue seen.add(resolved) candidates.append(resolved) candidates = [ candidate for candidate in candidates if candidate.exists() and file_looks_like_trace(candidate) ] candidates.sort(key=lambda item: item.stat().st_mtime) if limit is not None and limit >= 0: return candidates[-limit:] if limit else [] return candidates def newest_trace_dir(path: Path) -> Path: if path.is_file(): return path.parent direct = discover_trace_files(path, recursive=False) if direct: return path traces = discover_trace_files(path, recursive=True) trace_dirs = list({trace.parent for trace in traces}) if not trace_dirs: raise FileNotFoundError(f"No trace files found under {path}") trace_dirs.sort( key=lambda item: max( trace.stat().st_mtime for trace in traces if trace.parent == item ) ) return trace_dirs[-1] def discover_trace_targets( path: Path, all_traces: bool ) -> Tuple[List[Path], Optional[dict]]: if path.is_file(): return [path], load_server_args(path) direct_traces = discover_trace_files(path, recursive=False) recursive_traces = discover_trace_files(path, recursive=True) recursive_stages = {parse_stage(trace) for trace in recursive_traces} if ( not direct_traces and recursive_traces and any(stage != "all" for stage in recursive_stages) ): traces = recursive_traces trace_dir = path else: trace_dir = newest_trace_dir(path) traces = discover_trace_files(trace_dir, recursive=False) if not traces: raise FileNotFoundError(f"No trace files found under {trace_dir}") non_merged = [trace for trace in traces if not trace.name.startswith("merged-")] selected = non_merged or traces if not all_traces: ranks = sorted( { rank for rank in (parse_tp_rank(trace) for trace in selected) if rank is not None } ) if ranks: rank = 0 if 0 in ranks else ranks[0] selected = [trace for trace in selected if parse_tp_rank(trace) == rank] grouped: Dict[str, List[Path]] = defaultdict(list) for trace in selected: grouped[parse_stage(trace)].append(trace) selected = [ sorted(group, key=lambda item: item.stat().st_mtime)[-1] for group in grouped.values() ] selected.sort(key=lambda item: (STAGE_ORDER.get(parse_stage(item), 99), item.name)) return selected, load_server_args(trace_dir) def post_json( url: str, payload: Optional[dict] = None, timeout: float = 60.0 ) -> Optional[dict]: req = request.Request( url=url, data=(None if payload is None else json.dumps(payload).encode("utf-8")), headers={"Content-Type": "application/json"}, method="POST", ) with request.urlopen(req, timeout=timeout) as response: raw = response.read() return json.loads(raw.decode("utf-8")) if raw else None def send_probe_request( url: str, prompt: str, max_new_tokens: int, sampling_seed: int, framework: str, model: Optional[str] = None, ) -> None: framework = canonicalize_framework(framework) if framework == "sglang": payload = { "text": prompt, "sampling_params": { "sampling_seed": sampling_seed, "temperature": 0.0, "max_new_tokens": max_new_tokens, }, "stream": False, } post_json(url.rstrip("/") + "/generate", payload, timeout=300.0) return resolved_model = model or discover_openai_model(url) chat_payload = { "model": resolved_model, "messages": [{"role": "user", "content": prompt}], "temperature": 0.0, "max_tokens": max_new_tokens, "stream": False, } try: post_json(url.rstrip("/") + "/v1/chat/completions", chat_payload, timeout=300.0) return except Exception: completion_payload = { "model": resolved_model, "prompt": prompt, "temperature": 0.0, "max_tokens": max_new_tokens, "stream": False, } post_json( url.rstrip("/") + "/v1/completions", completion_payload, timeout=300.0, ) def unique_probe_prompt(prompt: str, probe_index: int) -> str: marker = f"profile_probe_{max(0, int(probe_index))}" parts = prompt.split(maxsplit=1) suffix = parts[1] if len(parts) == 2 else prompt return f"{marker} {suffix}".strip() def send_probe_requests( *, url: str, prompt: str, max_new_tokens: int, request_count: int, framework: str, model: Optional[str] = None, sampling_seed_offset: int = 0, ) -> None: request_count = max(0, int(request_count)) seed_offset = max(0, int(sampling_seed_offset)) for request_idx in range(request_count): probe_index = seed_offset + request_idx send_probe_request( url=url, prompt=unique_probe_prompt(prompt, probe_index), max_new_tokens=max_new_tokens, sampling_seed=probe_index, framework=framework, model=model, ) def synthetic_prompt(input_len: int) -> str: token_count = max(1, int(input_len)) return " ".join(["profile"] * token_count) def workload_probe( stage: str, *, prefill_input_len: int, prefill_output_len: int, decode_input_len: int, decode_output_len: int, ) -> Tuple[str, int]: if stage == "prefill": return synthetic_prompt(prefill_input_len), max(1, int(prefill_output_len)) if stage == "decode": return synthetic_prompt(decode_input_len), max(1, int(decode_output_len)) raise ValueError(f"unknown profile workload stage: {stage}") def build_probe_plan( stage: str, *, prompt: str, max_new_tokens: int, num_steps: int, probe_requests: int, warmup_steps: int, ) -> ProbePlan: active_steps = max(1, int(num_steps)) requested_probes = max(1, int(probe_requests)) warmup_steps = max(0, int(warmup_steps)) max_new_tokens = max(1, int(max_new_tokens)) if stage == "prefill": return ProbePlan( prompt=prompt, capture_max_new_tokens=max_new_tokens, capture_requests=max(requested_probes, active_steps), warmup_max_new_tokens=max_new_tokens, warmup_requests=warmup_steps, ) if stage == "decode": return ProbePlan( prompt=prompt, capture_max_new_tokens=max_new_tokens, capture_requests=requested_probes, warmup_max_new_tokens=max(1, warmup_steps), warmup_requests=1 if warmup_steps else 0, ) return ProbePlan( prompt=prompt, capture_max_new_tokens=max_new_tokens, capture_requests=requested_probes, warmup_max_new_tokens=max_new_tokens, warmup_requests=warmup_steps, ) def expand_profile_workload(profile_workload: str) -> List[str]: workload = normalize_text(profile_workload).lower() if workload not in PROFILE_WORKLOAD_CHOICES: raise ValueError( f"--profile-workload must be one of {', '.join(PROFILE_WORKLOAD_CHOICES)}" ) if workload == "both": return ["prefill", "decode"] if workload == "legacy": return ["legacy"] return [workload] def discover_openai_model(url: str) -> str: payload = try_get_json(url.rstrip("/") + "/v1/models", timeout=60.0) if not isinstance(payload, dict): raise RuntimeError(f"Could not read {url.rstrip('/')}/v1/models") data = payload.get("data") if not isinstance(data, list) or not data: raise RuntimeError(f"No models returned by {url.rstrip('/')}/v1/models") first = data[0] if isinstance(first, dict) and first.get("id"): return str(first["id"]) raise RuntimeError(f"Malformed /v1/models payload from {url.rstrip('/')}") def ensure_remote_profiler_output_path( output_dir: Optional[str], framework: str ) -> Path: if not output_dir: raise ValueError( f"{framework_display_name(framework)} live capture requires --output-dir " "to point at the server-side torch profiler trace path that is visible " "from this machine." ) output_path = Path(output_dir).expanduser().resolve() if output_path.suffix in {".json", ".gz"}: output_path.parent.mkdir(parents=True, exist_ok=True) else: output_path.mkdir(parents=True, exist_ok=True) return output_path def wait_for_profiler_artifact(path: Path, timeout_s: float = 60.0) -> Path: deadline = time.time() + timeout_s while time.time() < deadline: if path.is_file() and file_looks_like_trace(path): return path if path.exists(): trace_files = discover_trace_files(path, recursive=True) if trace_files: return newest_trace_dir(path) if path.is_dir(): child_dirs = [item for item in path.iterdir() if item.is_dir()] if child_dirs: child_dirs.sort(key=lambda item: item.stat().st_mtime) newest_child = child_dirs[-1] child_traces = discover_trace_files(newest_child, recursive=True) if child_traces: return newest_child time.sleep(0.5) return path def start_remote_profiler( url: str, framework: str, payload: Optional[dict] = None ) -> None: try: post_json(url.rstrip("/") + "/start_profile", payload=payload, timeout=60.0) except Exception as exc: if framework == "vllm": raise RuntimeError( "vLLM live torch profiling requires the server to be launched with " '--profiler-config \'{"profiler":"torch","torch_profiler_dir":"..."}\' ' "and to expose POST /start_profile." ) from exc if framework == "trtllm": raise RuntimeError( "TensorRT-LLM live torch profiling requires " "a server build that exposes POST /start_profile plus the env vars " "TLLM_PROFILE_START_STOP=- and " "TLLM_TORCH_PROFILE_TRACE=/shared/path." ) from exc if framework == "tokenspeed": raise RuntimeError( "TokenSpeed live torch profiling requires a server build that " "exposes POST /start_profile and POST /stop_profile. The helper " "passes output_dir, activities, and profile_id in the start payload." ) from exc raise def build_remote_profiler_start_payload( framework: str, output_path: Path, profile_prefix: Optional[str], stage: Optional[str], ) -> Optional[dict]: if framework != "tokenspeed": return None profile_id = profile_prefix or "triage-trace" if stage: profile_id = f"{profile_id}-{stage}" return { "output_dir": str(output_path), "activities": ["CPU", "GPU"], "with_stack": True, "record_shapes": False, "profile_id": profile_id, } def stop_remote_profiler(url: str, framework: str) -> None: try: post_json(url.rstrip("/") + "/stop_profile", timeout=300.0) except Exception as exc: raise RuntimeError( f"Failed to stop {framework_display_name(framework)} profiler via " f"{url.rstrip('/')}/stop_profile" ) from exc def run_remote_profiler( url: str, output_dir: Optional[str], framework: str, probe_plan: ProbePlan, probe_delay: float, profile_prefix: Optional[str] = None, stage: Optional[str] = None, ) -> Path: framework = canonicalize_framework(framework) output_path = ensure_remote_profiler_output_path(output_dir, framework) if stage and output_path.is_file(): raise ValueError( "--profile-workload both requires a directory output path for " f"{framework_display_name(framework)} so each stage trace can be labeled." ) before_traces = ( set(discover_trace_files(output_path, recursive=True)) if output_path.exists() else set() ) model = ( discover_openai_model(url) if framework in {"vllm", "trtllm", "tokenspeed"} else None ) if probe_plan.warmup_requests > 0: send_probe_requests( url=url, prompt=probe_plan.prompt, max_new_tokens=probe_plan.warmup_max_new_tokens, request_count=probe_plan.warmup_requests, framework=framework, model=model, ) start_payload = build_remote_profiler_start_payload( framework=framework, output_path=output_path, profile_prefix=profile_prefix, stage=stage, ) start_remote_profiler(url, framework, payload=start_payload) stop_error: Optional[BaseException] = None try: if probe_plan.capture_requests > 0: # Server-side profilers may do setup work after POST /start_profile. # A very short delay can send probes too early and miss the window. time.sleep(max(5.0, probe_delay)) send_probe_requests( url=url, prompt=probe_plan.prompt, max_new_tokens=probe_plan.capture_max_new_tokens, request_count=probe_plan.capture_requests, framework=framework, model=model, sampling_seed_offset=probe_plan.warmup_requests, ) finally: try: stop_remote_profiler(url, framework) except BaseException as exc: # pragma: no cover - preserve original failure stop_error = exc if stop_error is not None: raise stop_error artifact = wait_for_profiler_artifact(output_path) if stage and output_path.is_dir(): after_traces = set(discover_trace_files(output_path, recursive=True)) new_traces = sorted(after_traces - before_traces, key=lambda item: item.name) if new_traces: stage_dir = output_path / stage stage_dir.mkdir(parents=True, exist_ok=True) for trace in new_traces: if stage_dir in trace.parents: continue target = stage_dir / trace.name if target.exists(): target = stage_dir / f"{time.time_ns()}-{trace.name}" shutil.move(str(trace), str(target)) return stage_dir return artifact def run_sglang_profiler( url: str, output_dir: Optional[str], num_steps: int, profile_by_stage: bool, merge_profiles: bool, profile_prefix: Optional[str], probe_plan: ProbePlan, probe_delay: float, start_step: Optional[int] = None, ) -> Path: if output_dir is None: output_dir = tempfile.mkdtemp(prefix="sglang-torch-profile-") output_root = Path(output_dir).resolve() output_root.mkdir(parents=True, exist_ok=True) output_path = output_root / str(time.time()) output_path.mkdir(parents=True, exist_ok=True) server_args = try_get_json(url.rstrip("/") + "/server_info", timeout=60.0) if server_args is not None: with open(output_path / "server_args.json", "w", encoding="utf-8") as handle: json.dump(server_args, handle) payload = { "output_dir": str(output_path), "num_steps": str(num_steps), "activities": ["CPU", "GPU"], "profile_by_stage": profile_by_stage, "merge_profiles": merge_profiles, "profile_prefix": profile_prefix, } if start_step is not None: payload["start_step"] = str(start_step) if probe_plan.warmup_requests > 0: send_probe_requests( url=url, prompt=probe_plan.prompt, max_new_tokens=probe_plan.warmup_max_new_tokens, request_count=probe_plan.warmup_requests, framework="sglang", ) req = request.Request( url.rstrip("/") + "/start_profile", data=json.dumps(payload).encode("utf-8"), headers={"Content-Type": "application/json"}, ) with request.urlopen(req, timeout=300.0): pass if probe_plan.capture_requests > 0: time.sleep(max(0.0, probe_delay)) send_probe_requests( url=url, prompt=probe_plan.prompt, max_new_tokens=probe_plan.capture_max_new_tokens, request_count=probe_plan.capture_requests, framework="sglang", sampling_seed_offset=probe_plan.warmup_requests, ) try: stop_remote_profiler(url, "sglang") except RuntimeError: pass return wait_for_profiler_artifact(output_path, timeout_s=180.0) def run_profiler( url: str, output_dir: Optional[str], num_steps: int, profile_by_stage: bool, merge_profiles: bool, profile_prefix: Optional[str], probe_requests: int, probe_prompt: str, probe_max_new_tokens: Optional[int], probe_delay: float, warmup_steps: int = DEFAULT_WARMUP_STEPS, start_step: Optional[int] = None, framework: str = "auto", framework_hint_path: Optional[str] = None, profile_workload: str = "both", prefill_input_len: int = DEFAULT_PREFILL_INPUT_LEN, prefill_output_len: int = DEFAULT_PREFILL_OUTPUT_LEN, decode_input_len: int = DEFAULT_DECODE_INPUT_LEN, decode_output_len: int = DEFAULT_DECODE_OUTPUT_LEN, ) -> Path: resolved_framework = resolve_framework( framework, url=url, input_path=( Path(framework_hint_path).expanduser().resolve() if framework_hint_path else None ), ) if resolved_framework == "sglang": stages = expand_profile_workload(profile_workload) if stages != ["legacy"]: output_root = ( Path(output_dir).expanduser().resolve() if output_dir else Path(tempfile.mkdtemp(prefix="sglang-torch-profile-")) ) output_root.mkdir(parents=True, exist_ok=True) for stage in stages: prompt, max_new_tokens = workload_probe( stage, prefill_input_len=prefill_input_len, prefill_output_len=prefill_output_len, decode_input_len=decode_input_len, decode_output_len=decode_output_len, ) probe_plan = build_probe_plan( stage, prompt=prompt, max_new_tokens=max_new_tokens, num_steps=num_steps, probe_requests=probe_requests, warmup_steps=warmup_steps, ) # SGLang increments `forward_ct` before checking whether the # profiler reached its target. Ask for one extra step so the # requested stage forward is captured instead of stopping just # before it runs. stage_num_steps = max(1, int(num_steps)) + 1 run_sglang_profiler( url=url, output_dir=str(output_root / stage), num_steps=stage_num_steps, profile_by_stage=False, merge_profiles=merge_profiles, profile_prefix=( f"{profile_prefix}-{stage}" if profile_prefix else stage ), probe_plan=probe_plan, probe_delay=probe_delay, start_step=start_step, ) return output_root legacy_max_new_tokens = probe_max_new_tokens or max(64, num_steps * 8) legacy_plan = build_probe_plan( "legacy", prompt=probe_prompt, max_new_tokens=legacy_max_new_tokens, num_steps=num_steps, probe_requests=probe_requests, warmup_steps=warmup_steps, ) return run_sglang_profiler( url=url, output_dir=output_dir, num_steps=num_steps, profile_by_stage=profile_by_stage, merge_profiles=merge_profiles, profile_prefix=profile_prefix, probe_plan=legacy_plan, probe_delay=probe_delay, start_step=start_step, ) if start_step is not None: raise ValueError("--start-step is only supported for SGLang live capture.") if profile_by_stage: raise ValueError( "--profile-by-stage is only supported for SGLang live capture. " "Disable it when profiling vLLM, TensorRT-LLM, or TokenSpeed." ) if merge_profiles: raise ValueError( "--merge-profiles is only supported for SGLang live capture. " "Disable it when profiling vLLM, TensorRT-LLM, or TokenSpeed." ) if profile_prefix and resolved_framework in {"vllm", "trtllm"}: print( f"Note: {framework_display_name(resolved_framework)} ignores " "--profile-prefix on the HTTP profiler control path.", file=sys.stderr, ) stages = expand_profile_workload(profile_workload) if stages == ["legacy"]: legacy_max_new_tokens = probe_max_new_tokens or max(64, num_steps * 8) return run_remote_profiler( url=url, output_dir=output_dir, framework=resolved_framework, probe_plan=build_probe_plan( "legacy", prompt=probe_prompt, max_new_tokens=legacy_max_new_tokens, num_steps=num_steps, probe_requests=probe_requests, warmup_steps=warmup_steps, ), probe_delay=probe_delay, profile_prefix=profile_prefix, ) output_root = ensure_remote_profiler_output_path(output_dir, resolved_framework) for stage in stages: prompt, max_new_tokens = workload_probe( stage, prefill_input_len=prefill_input_len, prefill_output_len=prefill_output_len, decode_input_len=decode_input_len, decode_output_len=decode_output_len, ) run_remote_profiler( url=url, output_dir=str(output_root), framework=resolved_framework, probe_plan=build_probe_plan( stage, prompt=prompt, max_new_tokens=max_new_tokens, num_steps=num_steps, probe_requests=probe_requests, warmup_steps=warmup_steps, ), probe_delay=probe_delay, profile_prefix=profile_prefix, stage=stage, ) return output_root def select_heaviest_pid( events: Sequence[dict], event_filter: Callable[[dict], bool], pid_substring: Optional[str] = None, preferred_substrings: Iterable[str] = (), ) -> Optional[str]: durations: Counter = Counter() for event in events: if not event_filter(event): continue pid = str(event.get("pid")) if pid_substring and pid_substring not in pid: continue durations[pid] += float(event["dur"]) if not durations: return None for substring in preferred_substrings: preferred = [pid for pid in durations if substring in pid] if preferred: return max(preferred, key=lambda pid: durations[pid]) return max(durations, key=lambda pid: durations[pid])