From f6f1cacbce0e39992d04843f652c1adda373ae43 Mon Sep 17 00:00:00 2001 From: Gahow Wang Date: Sat, 11 Jul 2026 17:29:02 +0800 Subject: [PATCH 1/5] Add lightweight per-step OpProf telemetry Assisted-by: OpenAI Codex --- vllm/envs.py | 4 + vllm/v1/core/sched/scheduler.py | 28 +++ vllm/v1/opprof.py | 337 +++++++++++++++++++++++++++++ vllm/v1/worker/gpu_model_runner.py | 6 +- 4 files changed, 374 insertions(+), 1 deletion(-) create mode 100644 vllm/v1/opprof.py diff --git a/vllm/envs.py b/vllm/envs.py index 27a85bb..b3093e9 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -45,6 +45,7 @@ if TYPE_CHECKING: VLLM_LOGGING_COLOR: str = "auto" NO_COLOR: bool = False VLLM_LOG_STATS_INTERVAL: float = 10.0 + VLLM_OPPROF_DIR: str = "" VLLM_TRACE_FUNCTION: int = 0 VLLM_USE_FLASHINFER_SAMPLER: bool = True VLLM_PP_LAYER_PARTITION: str | None = None @@ -786,6 +787,9 @@ environment_variables: dict[str, Callable[[], Any]] = { if (val := float(os.getenv("VLLM_LOG_STATS_INTERVAL", "10."))) > 0.0 else 10.0 ), + # Directory for per-step OpProf JSONL telemetry. + # Empty disables OpProf. + "VLLM_OPPROF_DIR": lambda: os.getenv("VLLM_OPPROF_DIR", ""), # Trace function calls # If set to 1, vllm will trace function calls # Useful for debugging diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 90d93a1..303c562 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -7,6 +7,7 @@ from collections.abc import Iterable from dataclasses import replace from typing import Any +import vllm.envs as envs from vllm.compilation.cuda_graph import CUDAGraphStat from vllm.config import VllmConfig from vllm.distributed.ec_transfer.ec_connector.base import ( @@ -55,6 +56,7 @@ from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutp from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.metrics.perf import ModelMetrics, PerfStats from vllm.v1.metrics.stats import PrefixCacheStats, SchedulerStats +from vllm.v1.opprof import OpProfRecorder from vllm.v1.outputs import DraftTokenIds, KVConnectorOutput, ModelRunnerOutput from vllm.v1.request import Request, RequestStatus, StreamingUpdate from vllm.v1.spec_decode.dynamic.utils import build_dynamic_sd_schedule_lookup @@ -271,6 +273,12 @@ class Scheduler(SchedulerInterface): if self.connector is not None: self.connector.bind_gpu_block_pool(self.kv_cache_manager.block_pool) + self.opprof = OpProfRecorder.create( + output_dir=envs.VLLM_OPPROF_DIR, + dp_rank=self.parallel_config.data_parallel_index, + log_stats=self.log_stats, + ) + self.use_pp = self.parallel_config.pipeline_parallel_size > 1 self.use_v2_model_runner = vllm_config.use_v2_model_runner # Scheduler iteration counter. Drives the V2+PP+async decode-throttle @@ -386,6 +394,9 @@ class Scheduler(SchedulerInterface): return num_new_tokens def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: + opprof_start = ( + self.opprof.capture_start(self) if self.opprof is not None else None + ) self.current_step += 1 # NOTE(woosuk) on the scheduling algorithm: # There's no "decoding phase" nor "prefill phase" in the scheduler. @@ -1090,6 +1101,14 @@ class Scheduler(SchedulerInterface): ) scheduler_output.ec_connector_metadata = ec_meta + if self.opprof is not None: + assert opprof_start is not None + self.opprof.begin( + scheduler=self, + output=scheduler_output, + start=opprof_start, + ) + # Advance the fence only for non-empty steps (those that actually # write KV and have their output processed later in update_from_output). if self.defer_block_free and total_num_scheduled_tokens > 0: @@ -1800,6 +1819,12 @@ class Scheduler(SchedulerInterface): engine_core_outputs[0] = eco = EngineCoreOutputs() eco.scheduler_stats = stats + if self.opprof is not None: + self.opprof.finalize( + output=scheduler_output, + cudagraph_stat=cudagraph_stats, + ) + return engine_core_outputs @staticmethod @@ -2292,6 +2317,9 @@ class Scheduler(SchedulerInterface): if self.ec_connector is not None: self.ec_connector.shutdown() + if self.opprof is not None: + self.opprof.close() + logger.debug_once("[shutdown] Scheduler: complete") ######################################################################## diff --git a/vllm/v1/opprof.py b/vllm/v1/opprof.py new file mode 100644 index 0000000..f0330d0 --- /dev/null +++ b/vllm/v1/opprof.py @@ -0,0 +1,337 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import atexit +import logging +import os +import queue +import threading +import time +from bisect import bisect_left +from pathlib import Path +from typing import Any + +import msgspec + +logger = logging.getLogger(__name__) + +SCHEMA_VERSION = 1 +CONTEXT_LENGTH_EDGES = tuple(1 << exponent for exponent in range(7, 18)) +CHUNK_SIZE_EDGES = tuple(1 << exponent for exponent in range(4, 12)) +DEFAULT_QUEUE_CAPACITY = 8192 +_CLOSE_TIMEOUT_SECONDS = 1.0 +_STOP = object() +_PREFIX_FIELDS = ( # noqa: SIM905 + "requests queries hits preempted_requests preempted_queries preempted_hits" +).split() + + +ScheduleStart = tuple[int, int, dict[str, dict[str, int] | None]] + + +def classify_chunk(was_chunk: bool, end: int, target: int) -> str: + assert 0 <= end <= target + if was_chunk: + return "middle" if end < target else "final" + return "first" if end < target else "unsplit" + + +def _prefix_values(stats: Any | None) -> dict[str, int] | None: + if stats is None: + return None + return {name: int(getattr(stats, name, 0)) for name in _PREFIX_FIELDS} + + +def _prefix_snapshot(scheduler: Any) -> dict[str, dict[str, int] | None]: + return { + "local": _prefix_values(scheduler.kv_cache_manager.prefix_cache_stats), + "external": _prefix_values(scheduler.connector_prefix_cache_stats), + } + + +def _prefix_delta( + before: dict[str, int] | None, after: dict[str, int] | None +) -> dict[str, int] | None: + if before is None and after is None: + return None + before = before or dict.fromkeys(_PREFIX_FIELDS, 0) + after = after or dict.fromkeys(_PREFIX_FIELDS, 0) + delta = {name: after[name] - before[name] for name in _PREFIX_FIELDS} + assert all(value >= 0 for value in delta.values()) + return delta + + +class JSONLWriter: + def __init__( + self, + path: Path, + capacity: int = DEFAULT_QUEUE_CAPACITY, + start: bool = True, + ) -> None: + self._queue: queue.Queue[Any] = queue.Queue(capacity) + self._encoder = msgspec.json.Encoder() + self._file = path.open("xb", buffering=1 << 20) + self._thread = threading.Thread(target=self._run, daemon=True) + self._failure_lock = threading.Lock() + self._started = self._closed = False + self.failed = False + self.failure: Exception | None = None + self.encoded_records = self.written_records = 0 + self.dropped_records = self._unreported_drops = 0 + if start: + self.start() + + def start(self) -> None: + if not self._started: + self._started = True + self._thread.start() + + def _record_failure(self, error: Exception) -> None: + with self._failure_lock: + if self.failed: + return + self.failed = True + self.failure = error + logger.error("OpProf writer failed: %s", error) + + def _writer_unavailable(self) -> bool: + if self.failed: + return True + if self._started and not self._thread.is_alive(): + self._record_failure(RuntimeError("writer thread stopped unexpectedly")) + return True + return False + + def _drop(self, pending: int) -> bool: + self.dropped_records += 1 + self._unreported_drops = pending + 1 + return False + + def submit(self, record: dict[str, Any]) -> bool: + if self._closed: + raise RuntimeError("OpProf writer is closed") + pending = self._unreported_drops + record["dropped_records_before"] = pending + if self._writer_unavailable(): + return self._drop(pending) + payload = self._encoder.encode(record) + b"\n" + self.encoded_records += 1 + if self._writer_unavailable(): + return self._drop(pending) + try: + self._queue.put_nowait(payload) + except queue.Full: + return self._drop(pending) + self._unreported_drops = 0 + return True + + def _run(self) -> None: + buffered = 0 + last_flush = time.monotonic() + try: + while True: + try: + item = self._queue.get(timeout=1.0) + except queue.Empty: + self._file.flush() + buffered = 0 + last_flush = time.monotonic() + continue + try: + if item is _STOP: + break + self._file.write(item) + buffered += len(item) + self.written_records += 1 + now = time.monotonic() + if buffered >= 1 << 20 or now - last_flush >= 1.0: + self._file.flush() + buffered = 0 + last_flush = now + finally: + self._queue.task_done() + except Exception as error: + self._record_failure(error) + finally: + footer = dict( + schema=SCHEMA_VERSION, + record_type="footer", + encoded_records=self.encoded_records, + written_records=self.written_records, + dropped_records=self.dropped_records, + ) + try: + self._file.write(self._encoder.encode(footer) + b"\n") + self._file.flush() + except Exception as error: + self._record_failure(error) + finally: + try: + self._file.close() + except Exception as error: + self._record_failure(error) + + def close(self) -> None: + if self._closed: + return + self._closed = True + self.start() + if not self._thread.is_alive(): + return + try: + self._queue.put(_STOP, timeout=_CLOSE_TIMEOUT_SECONDS) + except queue.Full: + if not self._thread.is_alive(): + return + try: + self._queue.put_nowait(_STOP) + except queue.Full: + self._record_failure( + TimeoutError("timed out enqueueing writer stop sentinel") + ) + return + self._thread.join(timeout=_CLOSE_TIMEOUT_SECONDS) + if self._thread.is_alive(): + self._record_failure(TimeoutError("timed out joining writer thread")) + + +class OpProfRecorder: + def __init__(self, engine_id: str, writer: JSONLWriter) -> None: + self.engine_id = engine_id + self.writer = writer + self._next_step = 0 + self._pending: dict[int, dict[str, Any]] = {} + atexit.register(self.close) + + @classmethod + def create( + cls, output_dir: str, dp_rank: int, log_stats: bool + ) -> "OpProfRecorder | None": + if not output_dir: + return None + if not log_stats: + raise ValueError("VLLM_OPPROF_DIR requires log stats to be enabled") + directory = Path(output_dir).expanduser() + if not directory.is_absolute(): + raise ValueError("VLLM_OPPROF_DIR must be an absolute path") + directory.mkdir(parents=True, exist_ok=True) + engine_id = f"dp{dp_rank}-pid{os.getpid()}" + name = f"opprof-v{SCHEMA_VERSION}-{engine_id}-{time.time_ns()}.jsonl" + return cls(engine_id, JSONLWriter(directory / name)) + + @staticmethod + def capture_start(scheduler: Any) -> ScheduleStart: + return time.time_ns(), time.monotonic_ns(), _prefix_snapshot(scheduler) + + def begin(self, scheduler: Any, output: Any, start: ScheduleStart) -> None: + key = id(output) + assert key not in self._pending, "duplicate OpProf begin" + new_ids = {request.req_id for request in output.scheduled_new_reqs} + prefill_requests = prefill_tokens = decode_requests = decode_tokens = 0 + context_length_hist = [0] * (len(CONTEXT_LENGTH_EDGES) + 1) + chunk_size_hist = [0] * (len(CHUNK_SIZE_EDGES) + 1) + chunks: dict[str, Any] = dict.fromkeys( + ("first", "middle", "final", "unsplit"), 0 + ) + for req_id, num_tokens in output.num_scheduled_tokens.items(): + request = scheduler.requests[req_id] + end = request.num_computed_tokens + num_tokens + assert end >= 0 + context_length_hist[bisect_left(CONTEXT_LENGTH_EDGES, end)] += 1 + is_prefill = ( + req_id in new_ids + or output.scheduled_cached_reqs.is_context_phase(req_id) + ) + if is_prefill: + prefill_requests += 1 + prefill_tokens += num_tokens + assert num_tokens >= 0 + chunk_size_hist[bisect_left(CHUNK_SIZE_EDGES, num_tokens)] += 1 + target = request.num_tokens + request.num_output_placeholders + chunks[classify_chunk(request.is_prefill_chunk, end, target)] += 1 + else: + decode_requests += 1 + decode_tokens += num_tokens + prefix_after = _prefix_snapshot(scheduler) + block_pool = scheduler.kv_cache_manager.block_pool + total_blocks = block_pool.num_gpu_blocks - 1 + free_blocks = block_pool.get_num_free_blocks() + assert 0 <= free_blocks <= total_blocks + chunks["tokens"] = prefill_tokens + chunks["chunk_size_hist"] = chunk_size_hist + values = dict( + schema=SCHEMA_VERSION, + engine_id=self.engine_id, + step_index=self._next_step, + submit_wall_ns=start[0], + submit_mono_ns=start[1], + model_executed=output.total_num_scheduled_tokens > 0, + scheduled_requests=len(output.num_scheduled_tokens), + decode_batch_size=decode_requests, + prefill_requests=prefill_requests, + prefill_tokens=prefill_tokens, + decode_tokens=decode_tokens, + chunked_prefill=chunks, + context_length_hist=context_length_hist, + preemptions=len(output.preempted_req_ids or ()), + queues=dict( + running=len(scheduler.running), + waiting=len(scheduler.waiting), + deferred=len(scheduler.skipped_waiting), + ), + kv=dict( + total_blocks=total_blocks, + free_blocks=free_blocks, + used_blocks=total_blocks - free_blocks, + usage=scheduler.kv_cache_manager.usage, + ), + prefix=dict( + local=_prefix_delta(start[2]["local"], prefix_after["local"]), + external=_prefix_delta(start[2]["external"], prefix_after["external"]), + ), + ) + assert prefill_tokens + decode_tokens == output.total_num_scheduled_tokens + self._pending[key] = values + self._next_step += 1 + + def finalize(self, output: Any, cudagraph_stat: Any | None) -> bool: + try: + values = self._pending.pop(id(output)) + except KeyError: + raise AssertionError("missing or already finalized OpProf step") from None + if cudagraph_stat is None: + assert not values["model_executed"] + cudagraph = dict( + hit=False, + runtime_mode="NONE", + unpadded_tokens=0, + bucket_tokens=0, + padding_tokens=0, + ) + else: + mode = str(cudagraph_stat.runtime_mode).rsplit(".", 1)[-1] + cudagraph = dict( + hit=mode != "NONE", + runtime_mode=mode, + unpadded_tokens=cudagraph_stat.num_unpadded_tokens, + bucket_tokens=cudagraph_stat.num_padded_tokens, + padding_tokens=cudagraph_stat.num_paddings, + ) + record = dict( + values, + complete_mono_ns=time.monotonic_ns(), + cudagraph=cudagraph, + moe_expert_load=None, + dropped_records_before=0, + ) + return self.writer.submit(record) + + def close(self) -> None: + self.writer.close() + + @property + def failed(self) -> bool: + return self.writer.failed + + @property + def failure(self) -> Exception | None: + return self.writer.failure diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 74938a8..c11d773 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -437,6 +437,7 @@ class GPUModelRunner( self.scheduler_config = vllm_config.scheduler_config self.speculative_config = vllm_config.speculative_config self.observability_config = vllm_config.observability_config + self.opprof_enabled = bool(envs.VLLM_OPPROF_DIR) model_config = self.model_config cache_config = self.cache_config @@ -3917,7 +3918,10 @@ class GPUModelRunner( assert batch_descriptor.num_tokens == num_tokens_padded cudagraph_stats = None - if self.vllm_config.observability_config.cudagraph_metrics: + if ( + self.vllm_config.observability_config.cudagraph_metrics + or self.opprof_enabled + ): cudagraph_stats = CUDAGraphStat( num_unpadded_tokens=num_tokens, num_padded_tokens=batch_descriptor.num_tokens, -- 2.43.0