diff --git a/docs/contributing/profiling.md b/docs/contributing/profiling.md index 86bd68acef..5e5e270263 100644 --- a/docs/contributing/profiling.md +++ b/docs/contributing/profiling.md @@ -31,6 +31,18 @@ When using `vllm bench serve`, you can enable profiling by passing the `--profil Traces can be visualized using . +Each worker writes its trace and summary table under `torch_profiler_dir`, +named after its rank, for example `rank0..pt.trace.json.gz` and +`rank0.profiler_out.txt`. Depending on the parallel configuration, the rank +name can also include the parallel ranks, such as +`dp0_pp0_tp0_dcp0_ep0_rank0`. The AsyncLLM CPU trace uses +`_.async_llm` instead of the rank name. When a session sets +`profile_prefix`, it is prepended as `_`, so summaries from +different sessions do not overwrite each other. + +!!! note + The summary table was previously written to `profiler_out_.txt`. + !!! tip You can directly call bench module without installing vLLM using `python -m vllm.entrypoints.cli.main bench`. @@ -88,6 +100,31 @@ curl -X POST http://localhost:8000/v1/chat/completions \ $ curl -X POST http://localhost:8000/stop_profile ``` +`/start_profile` also accepts an optional JSON body with settings for that +profiling session. The Python and Rust API frontends accept the same fields: + +```shell +curl -X POST http://localhost:8000/start_profile \ + -H "Content-Type: application/json" \ + -d '{ + "profile_prefix": "sharegpt_run_1", + "delay_iterations": 5000, + "max_iterations": 20 + }' +``` + +`profile_prefix` is prepended to the generated trace and summary file names. It +must be 1-128 ASCII characters, start with a letter or number, and contain only +letters, numbers, `.`, `_`, or `-`. `delay_iterations` skips worker iterations +before collection starts, and `max_iterations` limits the number of collected +worker iterations (`0` means no limit). Both must be non-negative JSON integers; +values of the wrong type, such as `"5"` or `5.0`, are rejected rather than +converted. Unknown fields are ignored. These iteration +bounds do not apply to the AsyncLLM frontend CPU profiler, which records from +`/start_profile` until `/stop_profile`. Always call `/stop_profile`, including +after the worker reaches `max_iterations`; starting another session before +stopping the current one returns HTTP 409. + ## Profile with Triton Proton [Proton](https://github.com/triton-lang/triton/tree/main/third_party/proton) diff --git a/rust/src/engine-core-client/src/client.rs b/rust/src/engine-core-client/src/client.rs index 0514292b90..36a4fbffa5 100644 --- a/rust/src/engine-core-client/src/client.rs +++ b/rust/src/engine-core-client/src/client.rs @@ -9,7 +9,7 @@ use futures::future::join_all; use itertools::Itertools; use serde::Serialize; use serde_json::Value as JsonValue; -use tokio::sync::mpsc; +use tokio::sync::{Mutex, mpsc}; use tokio_util::task::AbortOnDropHandle; use tracing::{debug, info, trace}; @@ -283,6 +283,9 @@ pub struct EngineCoreClient { inner: Arc, coordinator: Option, abort_tx: mpsc::UnboundedSender, + /// Whether a profiling session is awaiting an explicit stop. Held across each + /// start/stop utility call so a stale stop reply cannot clear a newer session. + profile_active: Mutex, /// Runtime used to send messages to the engine and drive all background tasks. runtime: BackgroundShutdownRuntime, @@ -431,6 +434,7 @@ impl EngineCoreClient { inner, coordinator, abort_tx, + profile_active: Mutex::new(false), runtime, output_task, dispatcher_task, @@ -963,14 +967,38 @@ impl EngineCoreClient { } /// Start profiling the engine. - pub async fn start_profile(&self, profile_prefix: Option<&str>) -> Result<()> { - self.call_utility::<(), _>("profile", (true, profile_prefix)).await?; + pub async fn start_profile( + &self, + profile_prefix: Option<&str>, + delay_iterations: Option, + max_iterations: Option, + ) -> Result<()> { + let mut active = self.profile_active.lock().await; + if *active { + return Err(Error::ProfileAlreadyActive); + } + + // A cancelled start may still reach the engine, so only a reported + // failure clears the session; otherwise a stop is needed. + *active = true; + if let Err(error) = self + .call_utility::<(), _>( + "profile", + (true, profile_prefix, delay_iterations, max_iterations), + ) + .await + { + *active = false; + return Err(error); + } Ok(()) } /// Stop profiling the engine. pub async fn stop_profile(&self, profile_prefix: Option<&str>) -> Result<()> { + let mut active = self.profile_active.lock().await; self.call_utility::<(), _>("profile", (false, profile_prefix)).await?; + *active = false; Ok(()) } diff --git a/rust/src/engine-core-client/src/error.rs b/rust/src/engine-core-client/src/error.rs index d43952f023..d55f2cbbcb 100644 --- a/rust/src/engine-core-client/src/error.rs +++ b/rust/src/engine-core-client/src/error.rs @@ -83,6 +83,8 @@ pub enum Error { ClientClosed { message: String }, #[error("request output stream for `{request_id}` closed unexpectedly")] RequestStreamClosed { request_id: String }, + #[error("profiling session is already active")] + ProfileAlreadyActive, #[error("utility call `{method}` failed (call_id={call_id}): {message}")] UtilityCallFailed { method: String, diff --git a/rust/src/server/src/error.rs b/rust/src/server/src/error.rs index 5f7adc8fbb..77f7295b56 100644 --- a/rust/src/server/src/error.rs +++ b/rust/src/server/src/error.rs @@ -22,6 +22,8 @@ pub enum ApiError { ModelNotFound { model: String }, /// The request body could not be parsed as valid JSON. JsonParseError { message: String }, + /// The requested operation conflicts with current server state. + Conflict { message: String }, /// An unexpected internal failure happened before streaming started. ServerError { message: String }, } @@ -32,6 +34,7 @@ impl ApiError { match self { Self::InvalidRequest { .. } => StatusCode::BAD_REQUEST, Self::ModelNotFound { .. } => StatusCode::NOT_FOUND, + Self::Conflict { .. } => StatusCode::CONFLICT, Self::ServerError { .. } => StatusCode::INTERNAL_SERVER_ERROR, Self::JsonParseError { .. } => StatusCode::BAD_REQUEST, } @@ -65,6 +68,12 @@ impl ApiError { param: None, code: Some("json_parse_error".to_string()), }, + Self::Conflict { message } => ErrorDetail { + message: message.clone(), + error_type: "conflict_error".to_string(), + param: None, + code: Some("conflict_error".to_string()), + }, }; ErrorResponse { error } diff --git a/rust/src/server/src/routes/profile.rs b/rust/src/server/src/routes/profile.rs index de793a7b3b..763a4f8c1b 100644 --- a/rust/src/server/src/routes/profile.rs +++ b/rust/src/server/src/routes/profile.rs @@ -3,22 +3,75 @@ use std::sync::Arc; +use axum::Json; +use axum::body::Bytes; use axum::extract::State; use axum::http::StatusCode; +use serde::Deserialize; use tracing::info; +use vllm_engine_core_client::Error as EngineCoreClientError; use crate::error::ApiError; use crate::state::AppState; use crate::utils::utility_call_error; +#[derive(Debug, Default, Deserialize)] +pub(crate) struct StartProfileRequest { + profile_prefix: Option, + delay_iterations: Option, + max_iterations: Option, +} + +fn valid_profile_prefix(prefix: &str) -> bool { + let bytes = prefix.as_bytes(); + (1..=128).contains(&bytes.len()) + && bytes[0].is_ascii_alphanumeric() + && bytes[1..] + .iter() + .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'.' | b'_' | b'-')) +} + /// Start profiling the engine. -pub async fn start_profile(State(state): State>) -> Result { +pub async fn start_profile( + State(state): State>, + body: Bytes, +) -> Result { + let body = if body.is_empty() { + StartProfileRequest::default() + } else { + Json::::from_bytes(&body)?.0 + }; + if body + .profile_prefix + .as_deref() + .is_some_and(|prefix| !valid_profile_prefix(prefix)) + { + return Err(ApiError::invalid_request( + "profile_prefix must be 1-128 ASCII characters, start with a letter or \ + number, and contain only letters, numbers, '.', '_', or '-'" + .to_string(), + Some("profile_prefix"), + )); + } info!("starting profiler"); - state + if let Err(error) = state .engine_core_client() - .start_profile(None) + .start_profile( + body.profile_prefix.as_deref(), + body.delay_iterations, + body.max_iterations, + ) .await - .map_err(|error| utility_call_error("start_profile", error))?; + { + if matches!(error, EngineCoreClientError::ProfileAlreadyActive) { + return Err(ApiError::conflict( + "A profiling session is already active. Call /stop_profile before \ + starting another session." + .to_string(), + )); + } + return Err(utility_call_error("start_profile", error)); + } info!("profiler started"); Ok(StatusCode::OK) } @@ -34,3 +87,19 @@ pub async fn stop_profile(State(state): State>) -> Result) -> u64 { + let utility = recv_engine_message(dealer).await; + assert_eq!(utility[0].as_ref(), &[0x03]); + let payload = decode_value(&utility[1]).expect("decode utility payload"); + let array = payload.as_array().expect("utility payload array"); + assert_eq!(array[2], Value::from("profile")); + assert_eq!(array[3], Value::Array(expected_args)); + array[1].as_u64().expect("call id") +} + +/// A start issued while a stop is in flight waits for that stop instead of +/// racing it, so a stale stop reply cannot clear the new session. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn start_profile_waits_for_inflight_stop() { + let (stop_received_tx, stop_received_rx) = tokio::sync::oneshot::channel(); + let (release_stop_tx, release_stop_rx) = tokio::sync::oneshot::channel::<()>(); + let (app, engine_task) = test_profiling_app_with_engine_script(move |dealer, push| { + boxed_test_future(async move { + let start = || vec![Value::from(true), Value::Nil, Value::Nil, Value::Nil]; + + let call_id = recv_profile_call(dealer, start()).await; + send_outputs(push, utility_outputs(call_id, utility_none_result())).await; + + let call_id = recv_profile_call(dealer, vec![Value::from(false), Value::Nil]).await; + stop_received_tx.send(()).expect("notify stop received"); + release_stop_rx.await.expect("release stop"); + send_outputs(push, utility_outputs(call_id, utility_none_result())).await; + + let call_id = recv_profile_call(dealer, start()).await; + send_outputs(push, utility_outputs(call_id, utility_none_result())).await; + }) + }) + .await; + + let post = |uri: &'static str| { + let mut app = app.clone(); + async move { + app.call( + Request::builder() + .method("POST") + .uri(uri) + .body(Body::empty()) + .expect("build request"), + ) + .await + .expect("call app") + .status() + } + }; + + let timeout = Duration::from_secs(2); + assert_eq!(post("/start_profile").await, StatusCode::OK); + let stop = tokio::spawn(post("/stop_profile")); + tokio::time::timeout(timeout, stop_received_rx) + .await + .expect("stop dispatch timeout") + .expect("stop received"); + + // Nothing before the profile lock awaits, so one poll reaches it; an + // unserialized client would instead reject against the stale session. + let mut restart = std::pin::pin!(post("/start_profile")); + assert!(futures::poll!(restart.as_mut()).is_pending()); + + release_stop_tx.send(()).expect("release stop"); + let stop = tokio::time::timeout(timeout, stop).await.expect("stop timeout"); + assert_eq!(stop.expect("stop task"), StatusCode::OK); + let restart = tokio::time::timeout(timeout, restart).await.expect("restart timeout"); + assert_eq!(restart, StatusCode::OK); + tokio::time::timeout(timeout, engine_task) + .await + .expect("mock engine timeout") + .expect("mock engine task"); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn start_profile_route_rejects_invalid_prefix() { + let (app, engine_task) = + test_profiling_app_with_engine_script(|_, _| boxed_test_future(async move {})).await; + + let response = app + .clone() + .call( + Request::builder() + .method("POST") + .uri("/start_profile") + .header("content-type", "application/json") + .body(Body::from( + json!({"profile_prefix": "../trace"}).to_string(), + )) + .expect("build request"), + ) + .await + .expect("call app"); + + assert_eq!(response.status(), StatusCode::BAD_REQUEST); engine_task.await.expect("mock engine task"); } @@ -7846,16 +8106,8 @@ async fn start_profile_route_sends_expected_utility_call() { async fn stop_profile_route_sends_expected_utility_call() { let (app, engine_task) = test_profiling_app_with_engine_script(|dealer, push| { boxed_test_future(async move { - let utility = recv_engine_message(dealer).await; - assert_eq!(utility[0].as_ref(), &[0x03]); - - let payload = decode_value(&utility[1]).expect("decode utility payload"); - let array = payload.as_array().expect("utility payload array"); - let call_id = array[1].as_u64().expect("call id"); - - assert_eq!(array[2], Value::from("profile")); - assert_eq!(array[3], Value::Array(vec![Value::from(false), Value::Nil])); - + let stop = vec![Value::from(false), Value::Nil]; + let call_id = recv_profile_call(dealer, stop).await; send_outputs(push, utility_outputs(call_id, utility_none_result())).await; }) }) diff --git a/tests/benchmarks/test_audio_dataset.py b/tests/benchmarks/test_audio_dataset.py index e81c0cd173..54a165dc61 100644 --- a/tests/benchmarks/test_audio_dataset.py +++ b/tests/benchmarks/test_audio_dataset.py @@ -98,6 +98,17 @@ class _FakeSession: return _FakeResponse() +class _FakeProfileSession: + def __init__(self) -> None: + self.url: str | None = None + self.headers: dict[str, str] | None = None + + def post(self, *, url: str, headers: dict[str, str]): + self.url = url + self.headers = headers + return _FakeResponse() + + def test_asr_dataset_sample_handles_local_audio_paths(tmp_path: Path) -> None: audio_path = tmp_path / "earnings.wav" _write_wav(audio_path, duration_s=0.1) @@ -222,6 +233,42 @@ def test_async_request_openai_audio_handles_decoded_audio_arrays( assert output.generated_text == "hello" +def test_audio_request_rejects_profile_endpoint() -> None: + request_input = RequestFuncInput( + prompt="", + api_url="http://localhost:8000/start_profile", + prompt_len=1, + output_len=32, + model="openai/whisper-large-v3", + ) + + with pytest.raises(ValueError, match="OpenAI Audio API"): + asyncio.run( + request_func_module.async_request_openai_audio( + request_input, _FakeProfileSession() + ) + ) + + +@pytest.mark.parametrize("endpoint", ["start_profile", "stop_profile"]) +def test_audio_benchmark_profile_request_has_no_inference_payload( + endpoint: str, +) -> None: + session = _FakeProfileSession() + + output = asyncio.run( + request_func_module.async_request_profile( + f"http://localhost:8000/{endpoint}", + session, + extra_headers={"X-Test-Header": "profile"}, + ) + ) + + assert output.success is True + assert session.url == f"http://localhost:8000/{endpoint}" + assert session.headers == {"X-Test-Header": "profile"} + + _COHERE_ASR_PROMPT = ( "<|startofcontext|><|startoftranscript|>" "<|emo:undefined|><|en|><|en|><|pnc|><|noitn|>" diff --git a/tests/entrypoints/serve/test_profile_api_router.py b/tests/entrypoints/serve/test_profile_api_router.py new file mode 100644 index 0000000000..c5ceddd1bc --- /dev/null +++ b/tests/entrypoints/serve/test_profile_api_router.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from unittest.mock import AsyncMock + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from vllm.entrypoints.serve.profile.api_router import router + + +@pytest.fixture +def engine_client() -> AsyncMock: + return AsyncMock() + + +@pytest.fixture +def client(engine_client: AsyncMock) -> TestClient: + app = FastAPI() + app.state.engine_client = engine_client + app.include_router(router) + return TestClient(app) + + +def test_start_profile_without_body_is_backwards_compatible( + client: TestClient, engine_client: AsyncMock +) -> None: + response = client.post("/start_profile") + + assert response.status_code == 200 + engine_client.start_profile.assert_awaited_once_with() + + +def test_start_profile_ignores_completion_payload( + client: TestClient, engine_client: AsyncMock +) -> None: + response = client.post( + "/start_profile", + json={ + "model": "Qwen/Qwen3-0.6B", + "prompt": "profile this request", + "max_tokens": 16, + "stream": False, + }, + ) + + assert response.status_code == 200 + engine_client.start_profile.assert_awaited_once_with() + + +def test_start_profile_forwards_session_overrides( + client: TestClient, engine_client: AsyncMock +) -> None: + response = client.post( + "/start_profile", + json={ + "profile_prefix": "sharegpt_run-1", + "delay_iterations": 5000, + "max_iterations": 20, + }, + ) + + assert response.status_code == 200 + engine_client.start_profile.assert_awaited_once_with( + profile_prefix="sharegpt_run-1", + delay_iterations=5000, + max_iterations=20, + ) + + +@pytest.mark.parametrize( + "body", + [ + {"profile_prefix": 1}, + {"delay_iterations": "1"}, + {"max_iterations": "20"}, + ], +) +def test_start_profile_rejects_wrongly_typed_overrides( + client: TestClient, engine_client: AsyncMock, body: dict +) -> None: + response = client.post("/start_profile", json=body) + + assert response.status_code == 422 + engine_client.start_profile.assert_not_awaited() diff --git a/tests/v1/engine/test_async_llm.py b/tests/v1/engine/test_async_llm.py index dc3cf1aa00..3ac687c918 100644 --- a/tests/v1/engine/test_async_llm.py +++ b/tests/v1/engine/test_async_llm.py @@ -2,8 +2,10 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import asyncio +import threading import time from contextlib import ExitStack +from http import HTTPStatus from unittest.mock import AsyncMock, MagicMock, call import pytest @@ -11,7 +13,7 @@ import pytest import vllm.v1.engine.async_llm as async_llm_module from vllm import SamplingParams from vllm.assets.image import ImageAsset -from vllm.config import LoggingConfig, VllmConfig +from vllm.config import LoggingConfig, ProfilerConfig, VllmConfig from vllm.engine.arg_utils import AsyncEngineArgs from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionRequest, @@ -20,10 +22,11 @@ from vllm.entrypoints.openai.chat_completion.protocol import ( from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat from vllm.entrypoints.openai.models.protocol import BaseModelPath from vllm.entrypoints.openai.models.serving import OpenAIServingModels -from vllm.exceptions import VLLMValidationError +from vllm.exceptions import ProfilerAlreadyActiveError, VLLMValidationError from vllm.inputs import PromptType from vllm.outputs import RequestOutput from vllm.platforms import current_platform +from vllm.profiler.wrapper import WorkerProfiler from vllm.sampling_params import RequestOutputKind from vllm.utils.torch_utils import set_default_torch_num_threads from vllm.v1.engine.async_llm import AsyncLLM @@ -87,13 +90,19 @@ def _mock_async_llm_dependencies(monkeypatch: pytest.MonkeyPatch): return engine_core -def test_cuda_profiler_requests_reach_engine_core(monkeypatch: pytest.MonkeyPatch): +def _mock_vllm_config() -> MagicMock: vllm_config = MagicMock() vllm_config.logging_config = LoggingConfig() vllm_config.observability_config.otlp_traces_endpoint = None vllm_config.scheduler_config.stream_interval = 1 + return vllm_config + + +def test_cuda_profiler_requests_reach_engine_core(monkeypatch: pytest.MonkeyPatch): + vllm_config = _mock_vllm_config() vllm_config.profiler_config.profiler = "cuda" vllm_config.profiler_config.ignore_frontend = False + vllm_config.profiler_config.should_profile_frontend = False engine_core = _mock_async_llm_dependencies(monkeypatch) engine = AsyncLLM(vllm_config, MagicMock(), log_stats=False) @@ -105,22 +114,22 @@ def test_cuda_profiler_requests_reach_engine_core(monkeypatch: pytest.MonkeyPatc asyncio.run(profile()) assert engine.profiler is None - engine_core.profile_async.assert_has_awaits([call(True, None), call(False)]) + engine_core.profile_async.assert_has_awaits( + [call(True, None, None, None), call(False)] + ) def test_cuda_only_torch_profiler_skips_frontend_cpu_trace( monkeypatch: pytest.MonkeyPatch, ): - vllm_config = MagicMock() - vllm_config.logging_config = LoggingConfig() - vllm_config.observability_config.otlp_traces_endpoint = None - vllm_config.scheduler_config.stream_interval = 1 + vllm_config = _mock_vllm_config() vllm_config.profiler_config.profiler = "torch" vllm_config.profiler_config.ignore_frontend = False vllm_config.profiler_config.torch_profiler_activities = ["CUDA"] + vllm_config.profiler_config.should_profile_frontend = False _mock_async_llm_dependencies(monkeypatch) profiler = MagicMock() - monkeypatch.setattr(async_llm_module, "TorchProfilerWrapper", profiler) + monkeypatch.setattr(async_llm_module, "create_frontend_profiler", profiler) engine = AsyncLLM(vllm_config, MagicMock(), log_stats=False) @@ -128,6 +137,271 @@ def test_cuda_only_torch_profiler_skips_frontend_cpu_trace( profiler.assert_not_called() +def test_profile_forwards_overrides_rejects_duplicate_and_allows_restart( + monkeypatch: pytest.MonkeyPatch, +): + vllm_config = _mock_vllm_config() + vllm_config.profiler_config.should_profile_frontend = False + engine_core = _mock_async_llm_dependencies(monkeypatch) + + engine = AsyncLLM(vllm_config, MagicMock(), log_stats=False) + + async def profile(): + for kwargs in ( + {"profile_prefix": "../trace"}, + {"delay_iterations": -1}, + {"max_iterations": -1}, + ): + with pytest.raises(ValueError): + await engine.start_profile(**kwargs) + await engine.start_profile( + "session", + delay_iterations=5, + max_iterations=2, + ) + with pytest.raises(ProfilerAlreadyActiveError) as exc_info: + await engine.start_profile("duplicate") + assert exc_info.value.http_status == HTTPStatus.CONFLICT + await engine.stop_profile() + await engine.start_profile("next-session") + await engine.stop_profile() + + asyncio.run(profile()) + + engine_core.profile_async.assert_has_awaits( + [ + call(True, "session", 5, 2), + call(False), + call(True, "next-session", None, None), + call(False), + ] + ) + + +def test_stop_profile_waits_for_inflight_start(monkeypatch: pytest.MonkeyPatch): + vllm_config = _mock_vllm_config() + engine_core = _mock_async_llm_dependencies(monkeypatch) + events: list[str] = [] + frontend_profiler = MagicMock() + frontend_profiler.start.side_effect = lambda: events.append("frontend_start") + frontend_profiler.stop.side_effect = lambda: events.append("frontend_stop") + + engine = AsyncLLM( + vllm_config, + MagicMock(), + log_stats=False, + profiler=frontend_profiler, + ) + + async def profile(): + start_rpc_entered = asyncio.Event() + release_start_rpc = asyncio.Event() + + async def profile_async(is_start, *args): + if is_start: + start_rpc_entered.set() + await release_start_rpc.wait() + events.append("engine_start_done") + else: + events.append("engine_stop") + + engine_core.profile_async.side_effect = profile_async + + start_task = asyncio.create_task(engine.start_profile()) + await start_rpc_entered.wait() + + stop_task = asyncio.create_task(engine.stop_profile()) + await asyncio.sleep(0) + release_start_rpc.set() + await asyncio.wait_for(asyncio.gather(start_task, stop_task), timeout=5) + + asyncio.run(profile()) + + assert events[:2] == ["engine_start_done", "frontend_start"] + assert sorted(events[2:]) == ["engine_stop", "frontend_stop"] + assert engine._profile_session_active is False + + +@pytest.mark.parametrize("failure", ["error", "cancel"]) +def test_interrupted_engine_profile_start_skips_frontend_profiler( + monkeypatch: pytest.MonkeyPatch, failure: str +): + """The frontend trace starts only after the engine start succeeds. A + cancelled start may still reach the engine, so it keeps the session active + until /stop_profile.""" + vllm_config = _mock_vllm_config() + engine_core = _mock_async_llm_dependencies(monkeypatch) + frontend_profiler = MagicMock() + engine = AsyncLLM( + vllm_config, + MagicMock(), + log_stats=False, + profiler=frontend_profiler, + ) + + async def profile(): + start_rpc_entered = asyncio.Event() + + async def profile_async(is_start, *args): + if not is_start: + return + start_rpc_entered.set() + if failure == "error": + raise RuntimeError("engine profiler failed to start") + await asyncio.Event().wait() + + engine_core.profile_async.side_effect = profile_async + + start_task = asyncio.create_task(engine.start_profile()) + await start_rpc_entered.wait() + if failure == "cancel": + start_task.cancel() + expected = RuntimeError if failure == "error" else asyncio.CancelledError + with pytest.raises(expected): + await start_task + + assert engine._profile_session_active is (failure == "cancel") + await engine.stop_profile() + assert engine._profile_session_active is False + + asyncio.run(profile()) + + frontend_profiler.start.assert_not_called() + + +class _BlockingStartProfiler(WorkerProfiler): + def __init__(self) -> None: + self.in_start = threading.Event() + self.release_start = threading.Event() + self.stop_call_count = 0 + super().__init__( + ProfilerConfig(profiler="torch", torch_profiler_dir="/tmp/mock") + ) + + def _start(self) -> None: + self.in_start.set() + if not self.release_start.wait(timeout=5): + raise TimeoutError("start was never released") + + def _stop(self) -> None: + self.stop_call_count += 1 + + +def test_stop_after_cancelled_frontend_start_stops_frontend_profiler( + monkeypatch: pytest.MonkeyPatch, +): + """Cancelling start_profile while the frontend profiler is starting keeps + the session active, and the next stop_profile still stops the backend.""" + vllm_config = _mock_vllm_config() + engine_core = _mock_async_llm_dependencies(monkeypatch) + frontend_profiler = _BlockingStartProfiler() + engine = AsyncLLM( + vllm_config, + MagicMock(), + log_stats=False, + profiler=frontend_profiler, + ) + + async def profile(): + try: + start_task = asyncio.create_task(engine.start_profile()) + assert await asyncio.to_thread(frontend_profiler.in_start.wait, 5) + start_task.cancel() + with pytest.raises(asyncio.CancelledError): + await start_task + assert engine._profile_session_active is True + + stop_task = asyncio.create_task(engine.stop_profile()) + done, _ = await asyncio.wait({stop_task}, timeout=0.1) + assert not done, "stop_profile did not wait for the in-flight start" + frontend_profiler.release_start.set() + await asyncio.wait_for(stop_task, timeout=5) + finally: + frontend_profiler.release_start.set() + + asyncio.run(profile()) + + assert frontend_profiler.stop_call_count == 1 + assert frontend_profiler._active is False + assert frontend_profiler._running is False + assert engine._profile_session_active is False + assert engine_core.profile_async.await_args_list == [ + call(True, None, None, None), + call(False), + ] + + +def test_multi_client_profile_uses_idempotent_engine_requests( + monkeypatch: pytest.MonkeyPatch, tmp_path +): + vllm_config = _mock_vllm_config() + vllm_config.profiler_config = ProfilerConfig( + profiler="torch", + torch_profiler_dir=str(tmp_path), + torch_profiler_activities=["CPU"], + ) + engine_core = _mock_async_llm_dependencies(monkeypatch) + profiler = MagicMock() + monkeypatch.setattr(async_llm_module, "create_frontend_profiler", profiler) + engines = [ + AsyncLLM( + vllm_config, + MagicMock(), + log_stats=False, + client_count=2, + client_index=client_index, + ) + for client_index in range(2) + ] + + assert all(engine.profiler is None for engine in engines) + profiler.assert_not_called() + + async def profile(): + await engines[0].start_profile() + await engines[1].start_profile() + await engines[1].stop_profile() + await engines[0].start_profile() + + asyncio.run(profile()) + + engine_core.profile_async.assert_has_awaits( + [ + call(True, None, None, None), + call(True, None, None, None), + call(False), + call(True, None, None, None), + ] + ) + + +def test_frontend_profiler_ignores_worker_iteration_bounds( + monkeypatch: pytest.MonkeyPatch, tmp_path +): + vllm_config = _mock_vllm_config() + vllm_config.profiler_config = ProfilerConfig( + profiler="torch", + torch_profiler_dir=str(tmp_path), + torch_profiler_activities=["CPU", "CUDA"], + delay_iterations=5, + max_iterations=2, + wait_iterations=3, + warmup_iterations=1, + ) + _mock_async_llm_dependencies(monkeypatch) + profiler = MagicMock() + monkeypatch.setattr("vllm.profiler.wrapper.TorchProfilerWrapper", profiler) + + engine = AsyncLLM(vllm_config, MagicMock(), log_stats=False) + + frontend_config = profiler.call_args.args[0] + assert frontend_config.delay_iterations == 0 + assert frontend_config.max_iterations == 0 + assert frontend_config.wait_iterations == 0 + assert frontend_config.warmup_iterations == 0 + assert engine.profiler is profiler.return_value + + async def generate( engine: AsyncLLM, request_id: str, diff --git a/tests/v1/engine/test_llm_engine.py b/tests/v1/engine/test_llm_engine.py index b09f9405d7..69241889dc 100644 --- a/tests/v1/engine/test_llm_engine.py +++ b/tests/v1/engine/test_llm_engine.py @@ -7,6 +7,7 @@ from unittest.mock import MagicMock, call import pytest from vllm import LLM +from vllm.exceptions import ProfilerAlreadyActiveError from vllm.sampling_params import SamplingParams, StructuredOutputsParams from vllm.v1.engine.llm_engine import LLMEngine from vllm.v1.metrics.reader import Counter, Gauge, Histogram, Metric, Vector @@ -52,6 +53,38 @@ def test_step_invalidates_mm_processor_cache_misses(): ] +def test_duplicate_profile_start_rejected_before_engine_core_dispatch(): + engine = object.__new__(LLMEngine) + engine.engine_core = MagicMock() + engine._profile_session_active = False + + engine.start_profile("first") + + with pytest.raises(ProfilerAlreadyActiveError): + engine.start_profile("duplicate") + + engine.engine_core.profile.assert_called_once_with(True, "first", None, None) + + +@pytest.mark.parametrize( + "kwargs", + [ + {"profile_prefix": "../trace"}, + {"delay_iterations": -1}, + {"max_iterations": -1}, + ], +) +def test_invalid_profile_controls_rejected_before_engine_core_dispatch(kwargs): + engine = object.__new__(LLMEngine) + engine.engine_core = MagicMock() + engine._profile_session_active = False + + with pytest.raises(ValueError): + engine.start_profile(**kwargs) + + engine.engine_core.profile.assert_not_called() + + def _vllm_model( apc: bool, vllm_runner: type[VllmRunner], diff --git a/tests/v1/worker/test_gpu_profiler.py b/tests/v1/worker/test_gpu_profiler.py index e140803ef0..e272ed5106 100644 --- a/tests/v1/worker/test_gpu_profiler.py +++ b/tests/v1/worker/test_gpu_profiler.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import json import os +import threading from contextlib import nullcontext from types import SimpleNamespace from unittest.mock import MagicMock, Mock, call, patch @@ -27,6 +28,7 @@ from vllm.profiler.wrapper import ( validate_worker_profiler_config, ) from vllm.v1.core.sched.output import CachedRequestData +from vllm.v1.executor.abstract import Executor from vllm.v1.worker.gpu_model_runner import GPUModelRunner from vllm.v1.worker.gpu_worker import Worker from vllm.v1.worker.xpu_worker import XPUWorker @@ -159,13 +161,33 @@ def test_torch_profiler_device_summary(tmp_path, capsys, activities, dump_device wrapper.start() wrapper.stop() - summary = tmp_path / "profiler_out_0.txt" + summary = tmp_path / "worker.profiler_out.txt" assert summary.exists() == dump_device_time assert ("device times" in capsys.readouterr().out) == dump_device_time if dump_device_time: assert summary.read_text() == "device times\n" +def test_torch_profiler_summary_follows_output_name(tmp_path): + config = ProfilerConfig( + profiler="torch", + torch_profiler_dir=str(tmp_path), + ) + with patch("vllm.profiler.wrapper.torch.profiler.profile") as profile: + table = profile.return_value.key_averages.return_value.table + table.side_effect = ["first run", "second run"] + wrapper = TorchProfilerWrapper( + config, worker_name="initial", local_rank=1, activities=["CUDA"] + ) + for output_name in ["first_rank1", "second_rank1"]: + wrapper.set_output_name(output_name) + wrapper.start() + wrapper.stop() + + assert (tmp_path / "first_rank1.profiler_out.txt").read_text() == "first run\n" + assert (tmp_path / "second_rank1.profiler_out.txt").read_text() == "second run\n" + + @pytest.mark.parametrize( "activities", [[], ["CPU", "CPU"], ["INVALID"]], @@ -184,6 +206,30 @@ def test_torch_profiler_activities_require_torch_profiler(): ProfilerConfig(torch_profiler_activities=["CPU"]) +@pytest.mark.parametrize( + ("profiler", "activities", "ignore_frontend", "expected"), + [ + ("torch", None, False, True), + ("torch", ["CPU"], False, True), + ("torch", ["CPU", "CUDA"], False, True), + ("torch", ["CUDA"], False, False), + ("torch", ["CPU"], True, False), + ("cuda", None, False, False), + ], +) +def test_should_profile_frontend( + profiler, activities, ignore_frontend, expected, tmp_path +): + config = ProfilerConfig( + profiler=profiler, + torch_profiler_dir=str(tmp_path) if profiler == "torch" else "", + torch_profiler_activities=activities, + ignore_frontend=ignore_frontend, + ) + + assert config.should_profile_frontend is expected + + @pytest.mark.parametrize( ("device_type", "activities"), [ @@ -270,16 +316,81 @@ def test_worker_reuses_torch_wrapper_across_profile_rounds(worker_type): patch("vllm.distributed.utils.get_worker_rank_suffix", return_value="rank0"), patch("vllm.profiler.wrapper.TorchProfilerWrapper") as wrapper, ): - worker.profile() + worker.profile( + profile_prefix="first", + delay_iterations=5, + max_iterations=2, + ) worker.profile(is_start=False) - worker.profile() + worker.profile(profile_prefix="second") assert worker.profiler is wrapper.return_value wrapper.assert_called_once() - assert wrapper.return_value.start.call_count == 2 + assert wrapper.return_value.set_output_name.call_args_list == [ + call("first_rank0"), + call("second_rank0"), + ] + assert wrapper.return_value.start.call_args_list == [ + call(delay_iterations=5, max_iterations=2), + call(delay_iterations=None, max_iterations=None), + ] wrapper.return_value.stop.assert_called_once_with() +def test_executor_profile_only_forwards_supplied_override_kwargs(): + executor = MagicMock() + executor.vllm_config.parallel_config.worker_cls = "vllm.v1.worker.gpu_worker.Worker" + + Executor.profile(executor, True, "default") + Executor.profile( + executor, + True, + "bounded", + delay_iterations=5, + max_iterations=2, + ) + + assert executor.collective_rpc.call_args_list == [ + call("profile", args=(True, "default"), kwargs=None), + call( + "profile", + args=(True, "bounded"), + kwargs={"delay_iterations": 5, "max_iterations": 2}, + ), + ] + + +class _LegacyProfileWorker: + def profile(self, is_start: bool = True, profile_prefix: str | None = None): + pass + + +def test_executor_rejects_overrides_unsupported_by_worker(): + """Out-of-tree workers with the old ``profile`` signature are rejected + before dispatch, so a failed collective_rpc cannot leave replies queued.""" + executor = MagicMock() + executor.vllm_config.parallel_config.worker_cls = f"{__name__}._LegacyProfileWorker" + + with pytest.raises(ValueError, match="max_iterations"): + Executor.profile(executor, True, None, max_iterations=2) + Executor.profile(executor, True, "legacy") + + executor.collective_rpc.assert_called_once_with( + "profile", args=(True, "legacy"), kwargs=None + ) + + +def test_worker_rejects_invalid_profile_prefix(): + worker = object.__new__(Worker) + worker.profiler = None + worker.profiler_config = ProfilerConfig( + profiler="torch", torch_profiler_dir="/tmp/mock" + ) + + with pytest.raises(ValueError, match="profile_prefix"): + worker.profile(profile_prefix="../trace") + + def test_immediate_start_stop(default_profiler_config): """Test standard start without delay.""" profiler = ConcreteWorkerProfiler(default_profiler_config) @@ -294,6 +405,33 @@ def test_immediate_start_stop(default_profiler_config): assert profiler.stop_call_count == 1 +def test_stop_waits_for_inflight_start(default_profiler_config): + """A stop() from another thread during a slow start() must still stop the + backend instead of leaving it running with the session marked inactive.""" + profiler = ConcreteWorkerProfiler(default_profiler_config) + in_start = threading.Event() + release_start = threading.Event() + + def slow_start(): + in_start.set() + release_start.wait() + + with patch.object(profiler, "_start", side_effect=slow_start): + start_thread = threading.Thread(target=profiler.start) + start_thread.start() + assert in_start.wait(timeout=5) + stop_thread = threading.Thread(target=profiler.stop) + stop_thread.start() + stop_thread.join(timeout=0.1) + release_start.set() + start_thread.join(timeout=5) + stop_thread.join(timeout=5) + + assert profiler._active is False + assert profiler._running is False + assert profiler.stop_call_count == 1 + + def test_delayed_start(default_profiler_config): """Test that profiler waits for N steps before actually starting.""" default_profiler_config.delay_iterations = 2 @@ -400,19 +538,37 @@ def test_delayed_start_and_max_iters(default_profiler_config): assert profiler.stop_call_count == 1 -def test_idempotency(default_profiler_config): - """Test that calling start/stop multiple times doesn't break logic.""" +def test_duplicate_start_is_ignored(default_profiler_config): profiler = ConcreteWorkerProfiler(default_profiler_config) - # Double Start profiler.start() profiler.start() - assert profiler.start_call_count == 1 # Should only start once + assert profiler.start_call_count == 1 - # Double Stop profiler.stop() profiler.stop() - assert profiler.stop_call_count == 1 # Should only stop once + assert profiler.stop_call_count == 1 + + +def test_session_bounds_override_and_reset_to_defaults(): + profiler = ConcreteWorkerProfiler( + ProfilerConfig( + profiler="torch", + torch_profiler_dir="/tmp/mock", + delay_iterations=3, + max_iterations=4, + ) + ) + + profiler.start(delay_iterations=1, max_iterations=2) + assert profiler._delay_iters == 1 + assert profiler._max_iters == 2 + profiler.stop() + + profiler.start() + assert profiler._delay_iters == 3 + assert profiler._max_iters == 4 + profiler.stop() def test_step_inactive(default_profiler_config): @@ -1095,7 +1251,9 @@ def test_gpu_worker_creates_proton_profiler(): Worker.profile(worker) wrapper.assert_called_once_with(worker.profiler_config, worker_name="rank1") - worker.profiler.start.assert_called_once_with() + worker.profiler.start.assert_called_once_with( + delay_iterations=None, max_iterations=None + ) @_requires_cuda_for_proton @@ -1141,7 +1299,9 @@ def test_gpu_worker_reuses_cuda_graph_proton_session(): Worker.profile(worker, is_start=False) worker.profiler.set_output_name.assert_called_once_with("first_rank1") - worker.profiler.start.assert_called_once_with() + worker.profiler.start.assert_called_once_with( + delay_iterations=None, max_iterations=None + ) worker.profiler.stop.assert_called_once_with() diff --git a/vllm/benchmarks/lib/endpoint_request_func.py b/vllm/benchmarks/lib/endpoint_request_func.py index 14ff6c1ab1..778a4e50d3 100644 --- a/vllm/benchmarks/lib/endpoint_request_func.py +++ b/vllm/benchmarks/lib/endpoint_request_func.py @@ -122,8 +122,6 @@ def _validate_api_url( if isinstance(expected_suffixes, str): expected_suffixes = {expected_suffixes} - expected_suffixes = {*expected_suffixes, "profile"} - if not api_url.endswith(tuple(expected_suffixes)): raise ValueError(f"{api_name} URL must end with one of: {expected_suffixes}.") @@ -158,6 +156,34 @@ def _get_headers(content_type: str | None = None) -> dict[str, str]: return headers +async def async_request_profile( + api_url: str, + session: aiohttp.ClientSession, + extra_headers: dict | None = None, +) -> RequestFuncOutput: + _validate_api_url( + api_url, + "Profiling API", + {"start_profile", "stop_profile"}, + ) + + headers = _get_headers() + if extra_headers: + headers |= extra_headers + + output = RequestFuncOutput() + try: + async with session.post(url=api_url, headers=headers) as response: + output.success = response.status == 200 + if not output.success: + output.error = response.reason or "" + except Exception: + exc_info = sys.exc_info() + output.error = "".join(traceback.format_exception(*exc_info)) + + return output + + async def async_request_openai_completions( request_func_input: RequestFuncInput, session: aiohttp.ClientSession, diff --git a/vllm/benchmarks/serve.py b/vllm/benchmarks/serve.py index 19130716f4..b551361f55 100644 --- a/vllm/benchmarks/serve.py +++ b/vllm/benchmarks/serve.py @@ -49,6 +49,7 @@ from vllm.benchmarks.lib.endpoint_request_func import ( POOLING_BACKENDS, RequestFuncInput, RequestFuncOutput, + async_request_profile, ) from vllm.benchmarks.lib.ready_checker import wait_for_endpoint from vllm.benchmarks.lib.utils import ( @@ -936,22 +937,10 @@ async def benchmark( if profile: print("Starting profiler...") - profile_input = RequestFuncInput( - model=model_id, - model_name=model_name, - prompt=test_prompt, - api_url=base_url + "/start_profile", - prompt_len=test_prompt_len, - output_len=test_output_len, - logprobs=logprobs, - multi_modal_content=test_mm_content, - ignore_eos=ignore_eos, + profile_output = await async_request_profile( + base_url + "/start_profile", + session, extra_headers=extra_headers, - extra_body=test_extra_body, - chat_messages=test_chat_messages, - ) - profile_output = await request_func( - request_func_input=profile_input, session=session ) if profile_output.success: print("Profiler started") @@ -1469,16 +1458,10 @@ async def benchmark( if profile: print("Stopping profiler...") - profile_input = RequestFuncInput( - model=model_id, - prompt=test_prompt, - api_url=base_url + "/stop_profile", - prompt_len=test_prompt_len, - output_len=test_output_len, - logprobs=logprobs, - ) - profile_output = await request_func( - request_func_input=profile_input, session=session + profile_output = await async_request_profile( + base_url + "/stop_profile", + session, + extra_headers=extra_headers, ) if profile_output.success: print("Profiler stopped") diff --git a/vllm/config/profiler.py b/vllm/config/profiler.py index 3ade4bceed..7547546d14 100644 --- a/vllm/config/profiler.py +++ b/vllm/config/profiler.py @@ -22,6 +22,34 @@ ProtonHook = Literal["triton"] ProtonOutputFormat = Literal["hatchet", "hatchet_msgpack", "chrome_trace"] +def validate_profile_prefix(prefix: str | None) -> str | None: + """Validate a trace filename prefix supplied for one profiling session.""" + if prefix is None: + return None + if ( + not 1 <= len(prefix) <= 128 + or not prefix.isascii() + or not prefix[0].isalnum() + or any(not (char.isalnum() or char in "._-") for char in prefix[1:]) + ): + raise ValueError( + "profile_prefix must be 1-128 ASCII characters, start with a letter " + "or number, and contain only letters, numbers, '.', '_', or '-'" + ) + return prefix + + +def validate_profile_iteration_bounds( + delay_iterations: int | None, + max_iterations: int | None, +) -> None: + """Validate iteration bounds supplied for one profiling session.""" + if delay_iterations is not None and delay_iterations < 0: + raise ValueError("delay_iterations must be greater than or equal to 0") + if max_iterations is not None and max_iterations < 0: + raise ValueError("max_iterations must be greater than or equal to 0") + + def _is_uri_path(path: str) -> bool: """Check if path is a URI (scheme://...), excluding Windows drive letters. @@ -136,6 +164,16 @@ class ProfilerConfig: Defaults to 0, meaning no limit. """ + @property + def should_profile_frontend(self) -> bool: + """Whether the frontend CPU profiler should be enabled.""" + activities = self.torch_profiler_activities + return ( + self.profiler == "torch" + and not self.ignore_frontend + and (activities is None or "CPU" in activities) + ) + warmup_iterations: int = Field(default=0, ge=0) """Number of warmup iterations for PyTorch profiler schedule. During warmup, the profiler runs but data is discarded. This helps reduce @@ -188,13 +226,7 @@ class ProfilerConfig: raise ValueError( "torch_profiler_activities must not contain duplicates" ) - records_cpu_activity = activities is None or "CPU" in activities - if ( - self.profiler == "torch" - and has_delay_or_limit - and not self.ignore_frontend - and records_cpu_activity - ): + if has_delay_or_limit and self.should_profile_frontend: logger.warning_once( "Using 'torch' profiler with delay_iterations or max_iterations " "while ignore_frontend is False may result in high overhead." diff --git a/vllm/engine/protocol.py b/vllm/engine/protocol.py index a34981d651..2bb128ba07 100644 --- a/vllm/engine/protocol.py +++ b/vllm/engine/protocol.py @@ -162,8 +162,14 @@ class EngineClient(ABC): ... @abstractmethod - async def start_profile(self) -> None: - """Start profiling the engine.""" + async def start_profile( + self, + profile_prefix: str | None = None, + *, + delay_iterations: int | None = None, + max_iterations: int | None = None, + ) -> None: + """Start profiling with optional per-session overrides.""" ... @abstractmethod diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py index 3de7bab133..45b6aede58 100644 --- a/vllm/entrypoints/llm.py +++ b/vllm/entrypoints/llm.py @@ -785,16 +785,30 @@ class LLM(BeamSearchOfflineMixin, PoolingOfflineMixin, OfflineInferenceMixin): mm_processor_kwargs=mm_processor_kwargs, ) - def start_profile(self, profile_prefix: str | None = None) -> None: - """Start profiling with optional custom trace prefix. + def start_profile( + self, + profile_prefix: str | None = None, + *, + delay_iterations: int | None = None, + max_iterations: int | None = None, + ) -> None: + """Start profiling with optional per-session overrides. Args: profile_prefix: Optional prefix for the trace file names. If provided, trace files will be named as "_dp_pp_tp". If not provided, default naming will be used. + delay_iterations: Optional number of worker iterations to skip before + profiling starts. + max_iterations: Optional maximum number of worker iterations to profile. + Zero means no limit. """ - self.llm_engine.start_profile(profile_prefix) + self.llm_engine.start_profile( + profile_prefix, + delay_iterations=delay_iterations, + max_iterations=max_iterations, + ) def stop_profile(self) -> None: self.llm_engine.stop_profile() diff --git a/vllm/entrypoints/serve/profile/api_router.py b/vllm/entrypoints/serve/profile/api_router.py index eeed6b45ef..2051da2b1d 100644 --- a/vllm/entrypoints/serve/profile/api_router.py +++ b/vllm/entrypoints/serve/profile/api_router.py @@ -4,6 +4,7 @@ from fastapi import APIRouter, FastAPI, Request from fastapi.responses import Response +from pydantic import BaseModel, ConfigDict from vllm.config import ProfilerConfig from vllm.engine.protocol import EngineClient @@ -14,14 +15,32 @@ logger = init_logger(__name__) router = APIRouter() +class StartProfileRequest(BaseModel): + """Optional overrides for one profiling session. + + Values are validated by the engine client. + """ + + model_config = ConfigDict(extra="ignore", strict=True) + + profile_prefix: str | None = None + delay_iterations: int | None = None + max_iterations: int | None = None + + def engine_client(request: Request) -> EngineClient: return request.app.state.engine_client @router.post("/start_profile") -async def start_profile(raw_request: Request): +async def start_profile( + raw_request: Request, profile_request: StartProfileRequest | None = None +): logger.info("Starting profiler...") - await engine_client(raw_request).start_profile() + # Only pass supplied overrides so clients without them keep the + # zero-argument call. + overrides = profile_request.model_dump(exclude_none=True) if profile_request else {} + await engine_client(raw_request).start_profile(**overrides) logger.info("Profiler started.") return Response(status_code=200) diff --git a/vllm/exceptions.py b/vllm/exceptions.py index 9bfd4634bc..11ca13fe0b 100644 --- a/vllm/exceptions.py +++ b/vllm/exceptions.py @@ -181,3 +181,14 @@ class MaxQueuedTokensError(GracefulHTTPError): "Please try again later or on a different instance.", HTTPStatus.SERVICE_UNAVAILABLE, ) + + +class ProfilerAlreadyActiveError(GracefulHTTPError): + """Raised when a profiling session is already active.""" + + def __init__(self): + super().__init__( + "A profiling session is already active. Call /stop_profile before " + "starting another session.", + HTTPStatus.CONFLICT, + ) diff --git a/vllm/profiler/wrapper.py b/vllm/profiler/wrapper.py index bd9b6c231d..ebe7d90c3b 100644 --- a/vllm/profiler/wrapper.py +++ b/vllm/profiler/wrapper.py @@ -5,6 +5,7 @@ import importlib import inspect import json import os +import threading from abc import ABC, abstractmethod from collections.abc import Callable, Iterator, Sequence from contextlib import AbstractContextManager, contextmanager, nullcontext, suppress @@ -17,7 +18,13 @@ from typing_extensions import override import vllm.version from vllm.config import ProfilerConfig -from vllm.config.profiler import ProfilerKind, TorchProfilerActivity, _is_uri_path +from vllm.config.profiler import ( + ProfilerKind, + TorchProfilerActivity, + _is_uri_path, + validate_profile_iteration_bounds, +) +from vllm.config.utils import replace from vllm.logger import init_logger from vllm.platforms import current_platform @@ -28,14 +35,16 @@ _TRITON_PROTON_3_7_VERSION = Version("3.7.0") class WorkerProfiler(ABC): def __init__(self, profiler_config: ProfilerConfig) -> None: - self._delay_iters = profiler_config.delay_iterations + self._default_delay_iters = profiler_config.delay_iterations + self._delay_iters = self._default_delay_iters if self._delay_iters > 0: logger.info_once( "GPU profiling will start " f"{self._delay_iters} steps after start_profile." ) - self._max_iters = profiler_config.max_iterations + self._default_max_iters = profiler_config.max_iterations + self._max_iters = self._default_max_iters if self._max_iters > 0: logger.info_once( "GPU profiling will stop " @@ -43,6 +52,9 @@ class WorkerProfiler(ABC): "or when stop_profile is received." ) + # Serializes start()/stop(), which AsyncLLM may call from different threads. + self._lifecycle_lock = threading.Lock() + # Track when the profiler gets triggered by start_profile self._active_iteration_count = 0 self._active = False @@ -88,17 +100,32 @@ class WorkerProfiler(ABC): logger.warning("Failed to stop profiler: %s", e) self._running = False # Always mark as not running, assume stop worked - def start(self) -> None: + def start( + self, + *, + delay_iterations: int | None = None, + max_iterations: int | None = None, + ) -> None: """Attempt to start the profiler, accounting for delayed starts.""" - if self._active: - logger.debug( - "start_profile received when profiler is already active. " - "Ignoring request." + with self._lifecycle_lock: + if self._active: + logger.debug( + "start_profile received when profiler is already active. " + "Ignoring request." + ) + return + validate_profile_iteration_bounds(delay_iterations, max_iterations) + self._delay_iters = ( + self._default_delay_iters + if delay_iterations is None + else delay_iterations ) - return - self._active = True - if self._delay_iters == 0: - self._call_start() + self._max_iters = ( + self._default_max_iters if max_iterations is None else max_iterations + ) + self._active = True + if self._delay_iters == 0: + self._call_start() def step(self) -> None: """Update the profiler state at each worker step, @@ -148,17 +175,19 @@ class WorkerProfiler(ABC): def stop(self) -> None: """Attempt to stop the profiler, accounting for overlapped calls.""" - if not self._active: - logger.debug( - "stop_profile received when profiler is not active. Ignoring request." - ) - return - self._active = False - self._active_iteration_count = 0 - self._profiling_for_iters = 0 + with self._lifecycle_lock: + if not self._active: + logger.debug( + "stop_profile received when profiler is not active. " + "Ignoring request." + ) + return + self._active = False + self._active_iteration_count = 0 + self._profiling_for_iters = 0 - if self._running: - self._call_stop() + if self._running: + self._call_stop() def shutdown(self) -> None: """Ensure profiler is stopped when shutting down.""" @@ -175,6 +204,10 @@ class WorkerProfiler(ABC): """Observe graph creation for backends that attribute replay activity.""" return nullcontext() + def set_output_name(self, worker_name: str) -> None: + """Set the output name for the next profiling session.""" + return None + def annotate_context_manager(self, name: str): """Return a context manager to annotate profiler traces.""" return nullcontext() @@ -202,6 +235,10 @@ class TorchProfilerWrapper(WorkerProfiler): self.local_rank = local_rank self.profiler_config = profiler_config torch_profiler_trace_dir = profiler_config.torch_profiler_dir + self._torch_profiler_trace_dir = torch_profiler_trace_dir + self._torch_profiler_use_gzip = profiler_config.torch_profiler_use_gzip + self._worker_name = worker_name + self._custom_trace_handler = on_trace_ready is not None if local_rank in (None, 0): logger.info_once( "Torch profiling enabled. Traces will be saved to: %s", @@ -276,6 +313,21 @@ class TorchProfilerWrapper(WorkerProfiler): self._warmup_steps_remaining = self._initial_warmup_steps_remaining self._version_metadata_added = False + @override + def set_output_name(self, worker_name: str) -> None: + if self._active: + return + self._worker_name = worker_name + if self._custom_trace_handler: + return + self._profiler_kwargs["on_trace_ready"] = ( + torch.profiler.tensorboard_trace_handler( + self._torch_profiler_trace_dir, + worker_name=worker_name, + use_gzip=self._torch_profiler_use_gzip, + ) + ) + def _build_profiler_table( self, sort_key: str, @@ -293,13 +345,15 @@ class TorchProfilerWrapper(WorkerProfiler): row_limit=row_limit, ) - def _write_profiler_table(self, rank: int, table: str) -> None: + def _write_profiler_table(self, table: str) -> None: profiler_dir = self.profiler_config.torch_profiler_dir # Skip file write for URI paths (gs://, s3://, etc.) # as standard file I/O doesn't work with URI schemes if not _is_uri_path(profiler_dir): - profiler_out_file = f"{profiler_dir}/profiler_out_{rank}.txt" + profiler_out_file = os.path.join( + profiler_dir, f"{self._worker_name}.profiler_out.txt" + ) with open(profiler_out_file, "w") as f: print(table, file=f) @@ -344,7 +398,7 @@ class TorchProfilerWrapper(WorkerProfiler): rank = self.local_rank if self.dump_device_time_total: table = self._build_profiler_table(sort_key="self_device_time_total") - self._write_profiler_table(rank, table) + self._write_profiler_table(table) # only print profiler results on rank 0 if rank == 0: @@ -354,7 +408,7 @@ class TorchProfilerWrapper(WorkerProfiler): table = self._build_profiler_table( sort_key="self_cpu_time_total", row_limit=50 ) - self._write_profiler_table(rank, table) + self._write_profiler_table(table) # only print profiler results on rank 0 if rank == 0: @@ -494,6 +548,7 @@ class ProtonProfilerWrapper(WorkerProfiler): def has_cuda_graph_session(self) -> bool: return self._graph_session + @override def set_output_name(self, worker_name: str) -> None: """Set the next run's output name after startup graph capture.""" if self._active: @@ -701,6 +756,29 @@ def create_worker_profiler( return ProtonProfilerWrapper(profiler_config, worker_name=worker_name) +def create_frontend_profiler( + profiler_config: ProfilerConfig, *, worker_name: str +) -> TorchProfilerWrapper: + """Create the CPU-only frontend profiler, which runs for the whole session + rather than following engine iterations.""" + logger.info( + "Torch profiler enabled. AsyncLLM CPU traces will be collected under %s", + profiler_config.torch_profiler_dir, + ) + return TorchProfilerWrapper( + replace( + profiler_config, + delay_iterations=0, + max_iterations=0, + wait_iterations=0, + warmup_iterations=0, + ), + worker_name=worker_name, + local_rank=0, + activities=["CPU"], + ) + + def create_graph_capture_profiler( profiler_config: ProfilerConfig, global_rank: int ) -> WorkerProfiler | None: diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index 37025a8632..7e50aa0451 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -13,6 +13,10 @@ import vllm.envs as envs from vllm import TokensPrompt from vllm.config import VllmConfig from vllm.config.kv_events import KVEventsConfig +from vllm.config.profiler import ( + validate_profile_iteration_bounds, + validate_profile_prefix, +) from vllm.distributed.weight_transfer.base import ( WeightTransferInitRequest, WeightTransferUpdateRequest, @@ -23,6 +27,7 @@ from vllm.entrypoints.serve.elastic_ep.middleware import set_scaling_elastic_ep from vllm.exceptions import ( GracefulHTTPError, MaxQueuedTokensError, + ProfilerAlreadyActiveError, QueueOverflowError, VLLMClientError, VLLMValidationError, @@ -33,7 +38,7 @@ from vllm.lora.request import LoRARequest from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry from vllm.outputs import STREAM_FINISHED, PoolingRequestOutput, RequestOutput from vllm.pooling_params import PoolingParams -from vllm.profiler.wrapper import TorchProfilerWrapper +from vllm.profiler.wrapper import TorchProfilerWrapper, create_frontend_profiler from vllm.renderers import renderer_from_config from vllm.renderers.inputs.preprocess import extract_prompt_components from vllm.sampling_params import RequestOutputKind, SamplingParams @@ -211,24 +216,27 @@ class AsyncLLM(EngineClient): pass self.profiler = profiler - configured_activities = vllm_config.profiler_config.torch_profiler_activities + self._frontend_profiler_injected = profiler is not None + self._profile_lock = asyncio.Lock() + self._profile_session_guard_enabled = client_count == 1 + self._profile_session_active = False if ( - vllm_config.profiler_config.profiler == "torch" - and not vllm_config.profiler_config.ignore_frontend - and (configured_activities is None or "CPU" in configured_activities) + not self._frontend_profiler_injected + and vllm_config.profiler_config.should_profile_frontend ): - profiler_dir = vllm_config.profiler_config.torch_profiler_dir - logger.info( - "Torch profiler enabled. AsyncLLM CPU traces will be collected under %s", # noqa: E501 - profiler_dir, - ) - worker_name = f"{socket.gethostname()}_{os.getpid()}.async_llm" - self.profiler = TorchProfilerWrapper( - vllm_config.profiler_config, - worker_name=worker_name, - local_rank=0, - activities=["CPU"], - ) + if self._profile_session_guard_enabled: + self._frontend_profiler_worker_name = ( + f"{socket.gethostname()}_{os.getpid()}.async_llm" + ) + self.profiler = create_frontend_profiler( + vllm_config.profiler_config, + worker_name=self._frontend_profiler_worker_name, + ) + else: + logger.warning( + "Frontend CPU profiling is disabled when multiple API server " + "processes share an engine." + ) @classmethod def from_vllm_config( @@ -1081,17 +1089,48 @@ class AsyncLLM(EngineClient): if self.errored: raise self.dead_error - async def start_profile(self, profile_prefix: str | None = None) -> None: - coros = [self.engine_core.profile_async(True, profile_prefix)] - if self.profiler is not None: - coros.append(asyncio.to_thread(self.profiler.start)) - await asyncio.gather(*coros) + async def start_profile( + self, + profile_prefix: str | None = None, + *, + delay_iterations: int | None = None, + max_iterations: int | None = None, + ) -> None: + async with self._profile_lock: + if self._profile_session_guard_enabled and self._profile_session_active: + raise ProfilerAlreadyActiveError() + validate_profile_prefix(profile_prefix) + validate_profile_iteration_bounds(delay_iterations, max_iterations) + + if self._profile_session_guard_enabled: + self._profile_session_active = True + # A cancelled start may still take effect, so only a reported + # engine failure clears the session; otherwise /stop_profile is needed. + try: + await self.engine_core.profile_async( + True, + profile_prefix, + delay_iterations, + max_iterations, + ) + except Exception: + self._profile_session_active = False + raise + if self.profiler is not None: + if not self._frontend_profiler_injected: + worker_name = self._frontend_profiler_worker_name + if profile_prefix is not None: + worker_name = f"{profile_prefix}_{worker_name}" + self.profiler.set_output_name(worker_name) + await asyncio.to_thread(self.profiler.start) async def stop_profile(self) -> None: - coros = [self.engine_core.profile_async(False)] - if self.profiler is not None: - coros.append(asyncio.to_thread(self.profiler.stop)) - await asyncio.gather(*coros) + async with self._profile_lock: + coros = [self.engine_core.profile_async(False)] + if self.profiler is not None: + coros.append(asyncio.to_thread(self.profiler.stop)) + await asyncio.gather(*coros) + self._profile_session_active = False async def reset_mm_cache(self) -> None: # Join the background MM warmup first: the mm_processor_cache is not diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index cc0a8beddd..e0a003af62 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -810,8 +810,19 @@ class EngineCore: cleanup_dist_env_and_memory() logger.debug_once("[shutdown] EngineCore: local resource teardown complete") - def profile(self, is_start: bool = True, profile_prefix: str | None = None): - self.model_executor.profile(is_start, profile_prefix) + def profile( + self, + is_start: bool = True, + profile_prefix: str | None = None, + delay_iterations: int | None = None, + max_iterations: int | None = None, + ): + self.model_executor.profile( + is_start, + profile_prefix, + delay_iterations=delay_iterations, + max_iterations=max_iterations, + ) def reset_mm_cache(self): # NOTE: Since this is mainly for debugging, we don't attempt to diff --git a/vllm/v1/engine/core_client.py b/vllm/v1/engine/core_client.py index cb6705966a..88fdc1d1e7 100644 --- a/vllm/v1/engine/core_client.py +++ b/vllm/v1/engine/core_client.py @@ -181,7 +181,13 @@ class EngineCoreClient(ABC): def add_request(self, request: EngineCoreRequest) -> None: raise NotImplementedError - def profile(self, is_start: bool = True, profile_prefix: str | None = None) -> None: + def profile( + self, + is_start: bool = True, + profile_prefix: str | None = None, + delay_iterations: int | None = None, + max_iterations: int | None = None, + ) -> None: raise NotImplementedError def reset_mm_cache(self) -> None: @@ -275,7 +281,11 @@ class EngineCoreClient(ABC): raise NotImplementedError async def profile_async( - self, is_start: bool = True, profile_prefix: str | None = None + self, + is_start: bool = True, + profile_prefix: str | None = None, + delay_iterations: int | None = None, + max_iterations: int | None = None, ) -> None: raise NotImplementedError @@ -392,8 +402,19 @@ class InprocClient(EngineCoreClient): if self._finalizer.detach() is not None: self.engine_core.shutdown() - def profile(self, is_start: bool = True, profile_prefix: str | None = None) -> None: - self.engine_core.profile(is_start, profile_prefix) + def profile( + self, + is_start: bool = True, + profile_prefix: str | None = None, + delay_iterations: int | None = None, + max_iterations: int | None = None, + ) -> None: + self.engine_core.profile( + is_start, + profile_prefix, + delay_iterations=delay_iterations, + max_iterations=max_iterations, + ) def reset_mm_cache(self) -> None: self.engine_core.reset_mm_cache() @@ -1019,8 +1040,20 @@ class SyncMPClient(MPClient): if request_ids and not self.resources.engine_dead: self._send_input(EngineCoreRequestType.ABORT, request_ids) - def profile(self, is_start: bool = True, profile_prefix: str | None = None) -> None: - self.call_utility("profile", is_start, profile_prefix) + def profile( + self, + is_start: bool = True, + profile_prefix: str | None = None, + delay_iterations: int | None = None, + max_iterations: int | None = None, + ) -> None: + self.call_utility( + "profile", + is_start, + profile_prefix, + delay_iterations, + max_iterations, + ) def reset_mm_cache(self) -> None: self.call_utility("reset_mm_cache") @@ -1294,9 +1327,19 @@ class AsyncMPClient(MPClient): return await self.call_utility_consensus_async("is_scheduler_paused") async def profile_async( - self, is_start: bool = True, profile_prefix: str | None = None + self, + is_start: bool = True, + profile_prefix: str | None = None, + delay_iterations: int | None = None, + max_iterations: int | None = None, ) -> None: - await self.call_utility_async("profile", is_start, profile_prefix) + await self.call_utility_async( + "profile", + is_start, + profile_prefix, + delay_iterations, + max_iterations, + ) async def reset_mm_cache_async(self) -> None: await self.call_utility_async("reset_mm_cache") diff --git a/vllm/v1/engine/llm_engine.py b/vllm/v1/engine/llm_engine.py index 509d46703d..2dd70fc179 100644 --- a/vllm/v1/engine/llm_engine.py +++ b/vllm/v1/engine/llm_engine.py @@ -12,9 +12,14 @@ from typing_extensions import TypeVar import vllm.envs as envs from vllm.config import ParallelConfig, VllmConfig +from vllm.config.profiler import ( + validate_profile_iteration_bounds, + validate_profile_prefix, +) from vllm.distributed import stateless_destroy_torch_distributed_process_group from vllm.distributed.parallel_state import get_dp_group from vllm.engine.arg_utils import EngineArgs +from vllm.exceptions import ProfilerAlreadyActiveError from vllm.inputs import EngineInput, PromptType from vllm.logger import configure_logging_if_needed, init_logger from vllm.lora.request import LoRARequest @@ -116,6 +121,7 @@ class LLMEngine: log_stats=self.log_stats, renderer=renderer, ) + self._profile_session_active = False self.logger_manager: StatLoggerManager | None = None if self.log_stats: @@ -358,11 +364,33 @@ class LLMEngine: return processed_outputs.request_outputs - def start_profile(self, profile_prefix: str | None = None): - self.engine_core.profile(True, profile_prefix) + def start_profile( + self, + profile_prefix: str | None = None, + *, + delay_iterations: int | None = None, + max_iterations: int | None = None, + ): + if self._profile_session_active: + raise ProfilerAlreadyActiveError() + validate_profile_prefix(profile_prefix) + validate_profile_iteration_bounds(delay_iterations, max_iterations) + + self._profile_session_active = True + try: + self.engine_core.profile( + True, + profile_prefix, + delay_iterations, + max_iterations, + ) + except BaseException: + self._profile_session_active = False + raise def stop_profile(self): self.engine_core.profile(False) + self._profile_session_active = False def reset_mm_cache(self): # Join the background MM warmup first: the mm_processor_cache is not diff --git a/vllm/v1/executor/abstract.py b/vllm/v1/executor/abstract.py index f234e943b4..0af9508662 100644 --- a/vllm/v1/executor/abstract.py +++ b/vllm/v1/executor/abstract.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import inspect import time from abc import ABC, abstractmethod from collections.abc import Callable @@ -38,6 +39,18 @@ SLEEP_TAGS = frozenset(("weights", "kv_cache")) FailureCallback = Callable[[], None] +def _check_worker_profile_kwargs(worker_cls: str, kwargs: dict[str, int]) -> None: + """Raise if the worker's ``profile`` cannot accept the given overrides.""" + params = inspect.signature(resolve_obj_by_qualname(worker_cls).profile).parameters + if any(p.kind is inspect.Parameter.VAR_KEYWORD for p in params.values()): + return + if unsupported := [name for name in kwargs if name not in params]: + raise ValueError( + f"Worker class {worker_cls} does not support per-session profiling " + f"overrides: {', '.join(unsupported)}" + ) + + class Executor(ABC): """Abstract base class for vLLM executors. @@ -269,8 +282,30 @@ class Executor(ABC): output: list[DraftTokenIds] = self.collective_rpc("take_draft_token_ids") return output[0] - def profile(self, is_start: bool = True, profile_prefix: str | None = None): - self.collective_rpc("profile", args=(is_start, profile_prefix)) + def profile( + self, + is_start: bool = True, + profile_prefix: str | None = None, + *, + delay_iterations: int | None = None, + max_iterations: int | None = None, + ): + kwargs: dict[str, int] = {} + if delay_iterations is not None: + kwargs["delay_iterations"] = delay_iterations + if max_iterations is not None: + kwargs["max_iterations"] = max_iterations + if kwargs: + # Reject before dispatch: a failing collective_rpc may leave + # unread worker replies queued. + _check_worker_profile_kwargs( + self.vllm_config.parallel_config.worker_cls, kwargs + ) + self.collective_rpc( + "profile", + args=(is_start, profile_prefix), + kwargs=kwargs or None, + ) def save_sharded_state( self, diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index f9b58831e4..d4c64c5542 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -19,6 +19,7 @@ import torch.nn as nn import vllm.envs as envs from vllm.config import CUDAGraphMode, VllmConfig, set_current_vllm_config from vllm.config.compilation import CompilationMode +from vllm.config.profiler import validate_profile_prefix from vllm.device_allocator import get_mem_allocator_instance from vllm.distributed import ( ensure_model_parallel_initialized, @@ -1359,7 +1360,14 @@ class Worker(WorkerBase): def take_draft_token_ids(self) -> DraftTokenIds | None: return self.model_runner.take_draft_token_ids() - def profile(self, is_start: bool = True, profile_prefix: str | None = None): + def profile( + self, + is_start: bool = True, + profile_prefix: str | None = None, + *, + delay_iterations: int | None = None, + max_iterations: int | None = None, + ): # Check if profiling is enabled if self.profiler_config is None or self.profiler_config.profiler is None: raise RuntimeError( @@ -1370,6 +1378,8 @@ class Worker(WorkerBase): ) if is_start: + validate_profile_prefix(profile_prefix) + # Generate the trace name by combining prefix with comprehensive rank suffix from vllm.distributed.utils import get_worker_rank_suffix @@ -1381,9 +1391,6 @@ class Worker(WorkerBase): else: trace_name = rank_suffix - if self.profiler_config.profiler == "proton" and self.profiler is not None: - self.profiler.set_output_name(trace_name) - # Create the profiler wrapper only on the first start call if self.profiler is None: self.profiler = create_worker_profiler( @@ -1392,7 +1399,11 @@ class Worker(WorkerBase): local_rank=self.local_rank, ) - self.profiler.start() + self.profiler.set_output_name(trace_name) + self.profiler.start( + delay_iterations=delay_iterations, + max_iterations=max_iterations, + ) else: if self.profiler is None: logger.warning("Profiler was not started, nothing to stop.")