iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Profiler] Add per-session profiling controls to Python and Rust frontends (#57875)

Browse Source 
Signed-off-by: czhu <conway.zhu@cohere.com>
Signed-off-by: Conway Zhu <conway.zhu@cohere.com>
Co-authored-by: aoshen02 <aoshen@inferact.ai>
This commit is contained in:
czhu-cohere
2026-10-06 05:21:19 +00:00
committed by GitHub
co-authored by aoshen02
parent 9a1f6fd4cb
commit 6fddc9b21e
25 changed files with 1497 additions and 164 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/contributing/profiling.md
+37
View File
@@ -31,6 +31,18 @@ When using `vllm bench serve`, you can enable profiling by passing the `--profil
Traces can be visualized using <https://ui.perfetto.dev/>.
Each worker writes its trace and summary table under `torch_profiler_dir`,
named after its rank, for example `rank0.<id>.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
`<hostname>_<pid>.async_llm` instead of the rank name. When a session sets
`profile_prefix`, it is prepended as `<profile_prefix>_`, so summaries from
different sessions do not overwrite each other.
!!! note
The summary table was previously written to `profiler_out_<rank>.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)
rust/src/engine-core-client/src/client.rs
+31 -3
View File
@@ -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<ClientInner>,
coordinator: Option<CoordinatorHandle>,
abort_tx: mpsc::UnboundedSender<AbortRequest>,
/// 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<bool>,
/// 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<u64>,
max_iterations: Option<u64>,
) -> 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(())
}
rust/src/engine-core-client/src/error.rs
+2
View File
@@ -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,
rust/src/server/src/error.rs
+9
View File
@@ -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 }
rust/src/server/src/routes/profile.rs
+73 -4
View File
@@ -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<String>,
delay_iterations: Option<u64>,
max_iterations: Option<u64>,
}
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<Arc<AppState>>) -> Result<StatusCode, ApiError> {
pub async fn start_profile(
State(state): State<Arc<AppState>>,
body: Bytes,
) -> Result<StatusCode, ApiError> {
let body = if body.is_empty() {
StartProfileRequest::default()
} else {
Json::<StartProfileRequest>::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<Arc<AppState>>) -> Result<StatusCo
info!("profiler stopped");
Ok(StatusCode::OK)
}
#[cfg(test)]
mod tests {
use super::valid_profile_prefix;
#[test]
fn validates_profile_prefixes() {
for prefix in ["benchmark", "sharegpt_run-1", "profile.2"] {
assert!(valid_profile_prefix(prefix), "{prefix}");
}
for prefix in ["", "../trace", "/tmp/trace", "trace name", "_trace"] {
assert!(!valid_profile_prefix(prefix), "{prefix}");
}
assert!(!valid_profile_prefix(&"a".repeat(129)));
}
}
rust/src/server/src/routes/tests.rs
+276 -24
View File
@@ -7802,27 +7802,146 @@ where
)
}
/// An empty body starts profiling without overrides, with or without a JSON
/// content type.
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn start_profile_route_sends_expected_utility_call() {
for content_type in [None, Some("application/json")] {
let (app, engine_task) = test_profiling_app_with_engine_script(|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;
})
})
.await;
let mut request = Request::builder().method("POST").uri("/start_profile");
if let Some(content_type) = content_type {
request = request.header("content-type", content_type);
}
let response = app
.clone()
.call(request.body(Body::empty()).expect("build request"))
.await
.expect("call app");
let status = response.status();
let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
assert_eq!(
status,
StatusCode::OK,
"{content_type:?}: {}",
String::from_utf8_lossy(&body)
);
assert!(body.is_empty());
engine_task.await.expect("mock engine task");
}
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn start_profile_route_forwards_session_overrides() {
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(true), Value::Nil]));
let start = vec![
Value::from(true),
Value::from("sharegpt_run-1"),
Value::from(5000),
Value::from(20),
];
let call_id = recv_profile_call(dealer, start).await;
send_outputs(push, utility_outputs(call_id, utility_none_result())).await;
})
})
.await;
let response = app
.clone()
.call(
Request::builder()
.method("POST")
.uri("/start_profile")
.header("content-type", "application/json")
.body(Body::from(
json!({
"profile_prefix": "sharegpt_run-1",
"delay_iterations": 5000,
"max_iterations": 20,
})
.to_string(),
))
.expect("build request"),
)
.await
.expect("call app");
assert_eq!(response.status(), StatusCode::OK);
engine_task.await.expect("mock engine task");
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn start_profile_route_ignores_completion_payload() {
let (app, engine_task) = test_profiling_app_with_engine_script(|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;
})
})
.await;
let response = app
.clone()
.call(
Request::builder()
.method("POST")
.uri("/start_profile")
.header("content-type", "application/json")
.body(Body::from(
json!({
"model": "Qwen/Qwen3-0.6B",
"prompt": "profile this request",
"max_tokens": 16,
"stream": false,
})
.to_string(),
))
.expect("build request"),
)
.await
.expect("call app");
assert_eq!(response.status(), StatusCode::OK);
engine_task.await.expect("mock engine task");
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn duplicate_start_profile_returns_conflict() {
let (app, engine_task) = test_profiling_app_with_engine_script(|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;
})
})
.await;
let first = app
.clone()
.call(
Request::builder()
.method("POST")
.uri("/start_profile")
.body(Body::empty())
.expect("build request"),
)
.await
.expect("call app");
let second = app
.clone()
.call(
Request::builder()
@@ -7834,10 +7953,151 @@ async fn start_profile_route_sends_expected_utility_call() {
.await
.expect("call app");
let status = response.status();
let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
assert_eq!(status, StatusCode::OK, "{}", String::from_utf8_lossy(&body));
assert!(body.is_empty());
assert_eq!(first.status(), StatusCode::OK);
assert_eq!(second.status(), StatusCode::CONFLICT);
engine_task.await.expect("mock engine task");
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
#[serial]
async fn start_profile_succeeds_again_after_stop() {
let (app, engine_task) = test_profiling_app_with_engine_script(|dealer, push| {
boxed_test_future(async move {
let expected_args = [
vec![Value::from(true), Value::Nil, Value::Nil, Value::Nil],
vec![Value::from(false), Value::Nil],
vec![Value::from(true), Value::Nil, Value::Nil, Value::Nil],
vec![Value::from(false), Value::Nil],
];
for expected in expected_args {
let call_id = recv_profile_call(dealer, expected).await;
send_outputs(push, utility_outputs(call_id, utility_none_result())).await;
}
})
})
.await;
for uri in [
"/start_profile",
"/stop_profile",
"/start_profile",
"/stop_profile",
] {
let response = app
.clone()
.call(
Request::builder()
.method("POST")
.uri(uri)
.body(Body::empty())
.expect("build request"),
)
.await
.expect("call app");
assert_eq!(response.status(), StatusCode::OK, "{uri}");
}
engine_task.await.expect("mock engine task");
}
/// Receive a `profile` utility call with the expected args and return its call id.
async fn recv_profile_call(dealer: &mut DealerSocket, expected_args: Vec<Value>) -> 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;
})
})
tests/benchmarks/test_audio_dataset.py
+47
View File
@@ -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|>"
tests/entrypoints/serve/test_profile_api_router.py
+86
View File
@@ -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()
tests/v1/engine/test_async_llm.py
+283 -9
View File
@@ -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,
tests/v1/engine/test_llm_engine.py
+33
View File
@@ -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],
tests/v1/worker/test_gpu_profiler.py
+172 -12
View File
@@ -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()
vllm/benchmarks/lib/endpoint_request_func.py
+28 -2
View File
@@ -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,
vllm/benchmarks/serve.py
+8 -25
View File
@@ -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")
vllm/config/profiler.py
+39 -7
View File
@@ -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."
vllm/engine/protocol.py
+8 -2
View File
@@ -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
vllm/entrypoints/llm.py
+17 -3
View File
@@ -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 "<prefix>_dp<X>_pp<Y>_tp<Z>".
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()
vllm/entrypoints/serve/profile/api_router.py
+21 -2
View File
@@ -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)
vllm/exceptions.py
+11
View File
@@ -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,
)
vllm/profiler/wrapper.py
+104 -26
View File
@@ -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:
vllm/v1/engine/async_llm.py
+65 -26
View File
@@ -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
vllm/v1/engine/core.py
+13 -2
View File
@@ -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
vllm/v1/engine/core_client.py
+51 -8
View File
@@ -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")
vllm/v1/engine/llm_engine.py
+30 -2
View File
@@ -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
vllm/v1/executor/abstract.py
+37 -2
View File
@@ -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,
vllm/v1/worker/gpu_worker.py
+16 -5
View File
@@ -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.")