[Frontend] Add /abort_requests to the RLHF dev API router (#47173)
Signed-off-by: aoshen02 <aoshen@inferact.ai> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
This commit is contained in:
co-authored by
Claude Opus 4.8
mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
parent
a02984ed47
commit
83762b77b0
@@ -170,6 +170,7 @@ For further details on Weight Transfer, please refer to [this page](../../traini
|
||||
- `/pause` - Pause generation (causes denial of service)
|
||||
- `/resume` - Resume generation
|
||||
- `/is_paused` - Check if generation is paused
|
||||
- `/abort_requests` - Abort in-flight requests (all in-flight, or the given `request_ids`) without pausing the scheduler
|
||||
- `/init_weight_transfer_engine` - Initialize weight transfer engine for RLHF
|
||||
- `/start_weight_update` - Prepares the inference engine for a weight update.
|
||||
- `/update_weights` - Update model weights (can alter model behavior)
|
||||
|
||||
@@ -42,6 +42,7 @@ When using the vLLM HTTP server, the same functionality is available via:
|
||||
|
||||
- `POST /pause?mode=keep` - Pause generation
|
||||
- `POST /resume` - Resume generation
|
||||
- `POST /abort_requests` - Abort in-flight requests without pausing the scheduler (send `{}` to abort all, or `{"request_ids": [...]}`)
|
||||
|
||||
!!! note "Data Parallelism"
|
||||
When using data parallelism with vLLM's **internal load balancer** (i.e. `data_parallel_backend="ray"`), pause and resume are handled automatically across all DP ranks -- a single call is sufficient. When using an **external load balancer** (i.e. multiple independent vLLM instances behind a proxy), you must send pause and resume requests to **every** engine instance individually before and after the weight update.
|
||||
|
||||
@@ -191,6 +191,7 @@ The following endpoints **do not require authentication** even when `--api-key`
|
||||
- `/pause` - Pause generation (causes denial of service)
|
||||
- `/resume` - Resume generation
|
||||
- `/is_paused` - Check if generation is paused
|
||||
- `/abort_requests` - Abort in-flight requests (causes loss of in-flight work)
|
||||
- `/scale_elastic_ep` - Trigger scaling operations
|
||||
- `/is_scaling_elastic_ep` - Check if scaling is in progress
|
||||
- `/init_weight_transfer_engine` - Initialize weight transfer engine for RLHF
|
||||
|
||||
@@ -64,6 +64,13 @@ impl InflightRequests {
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Collect the internal engine ids of every in-flight request. Used to
|
||||
/// abort all outstanding requests when no external ids are given.
|
||||
pub(crate) fn all_internal_ids(&self) -> Vec<String> {
|
||||
let map = self.map.lock();
|
||||
map.values().flat_map(|internal_ids| internal_ids.keys()).cloned().collect()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
fn is_empty(&self) -> bool {
|
||||
self.map.lock().is_empty()
|
||||
|
||||
@@ -122,7 +122,12 @@ impl Llm {
|
||||
/// tracking entries themselves are removed when the corresponding output
|
||||
/// streams are dropped, not here.
|
||||
pub async fn abort(&self, external_ids: &[String]) -> Result<()> {
|
||||
let internal_ids = self.inflight.resolve(external_ids);
|
||||
// Empty `external_ids` means abort every in-flight request.
|
||||
let internal_ids = if external_ids.is_empty() {
|
||||
self.inflight.all_internal_ids()
|
||||
} else {
|
||||
self.inflight.resolve(external_ids)
|
||||
};
|
||||
if internal_ids.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
@@ -20,12 +20,8 @@ pub async fn abort_requests(
|
||||
body: Result<Json<AbortRequestsRequest>, JsonRejection>,
|
||||
) -> Result<StatusCode, ApiError> {
|
||||
let Json(body) = body.map_err(|error| ApiError::json_parse_error(error.body_text()))?;
|
||||
let request_ids = body.request_ids.ok_or_else(|| {
|
||||
ApiError::invalid_request(
|
||||
"Missing 'request_ids' in request body".to_string(),
|
||||
Some("request_ids"),
|
||||
)
|
||||
})?;
|
||||
// Empty/missing `request_ids` aborts all in-flight requests.
|
||||
let request_ids = body.request_ids.unwrap_or_default();
|
||||
|
||||
state
|
||||
.chat
|
||||
|
||||
@@ -5276,10 +5276,12 @@ async fn abort_requests_route_returns_ok_for_well_formed_body() {
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
#[serial]
|
||||
async fn abort_requests_route_rejects_missing_request_ids() {
|
||||
async fn abort_requests_route_aborts_all_when_request_ids_missing() {
|
||||
let (app, engine_task) =
|
||||
test_admin_app_with_engine_script(|_dealer, _push| boxed_test_future(async move {})).await;
|
||||
|
||||
// Missing `request_ids` means "abort all in-flight requests"; with no
|
||||
// in-flight requests this is a no-op that still succeeds.
|
||||
let response = app
|
||||
.clone()
|
||||
.call(
|
||||
@@ -5293,11 +5295,10 @@ async fn abort_requests_route_rejects_missing_request_ids() {
|
||||
.await
|
||||
.expect("call app");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
|
||||
let status = response.status();
|
||||
let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
|
||||
let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json");
|
||||
assert_eq!(json["error"]["type"], "invalid_request_error");
|
||||
assert_eq!(json["error"]["param"], "request_ids");
|
||||
assert_eq!(status, StatusCode::OK, "{}", String::from_utf8_lossy(&body));
|
||||
assert!(body.is_empty());
|
||||
engine_task.abort_and_join().await;
|
||||
}
|
||||
|
||||
|
||||
@@ -91,6 +91,51 @@ async def resume_generation(raw_request: Request) -> JSONResponse:
|
||||
)
|
||||
|
||||
|
||||
@router.post("/abort_requests")
|
||||
async def abort_requests(raw_request: Request) -> JSONResponse:
|
||||
"""Abort in-flight requests without pausing the scheduler.
|
||||
|
||||
Empty/missing ``request_ids`` aborts all in-flight requests.
|
||||
"""
|
||||
|
||||
engine = engine_client(raw_request)
|
||||
|
||||
try:
|
||||
body = await raw_request.json()
|
||||
except json.JSONDecodeError as e:
|
||||
raise HTTPException(status_code=400, detail="Invalid JSON format") from e # noqa: B904
|
||||
|
||||
request_ids = body.get("request_ids")
|
||||
|
||||
try:
|
||||
if request_ids:
|
||||
# Body ids are external (user-supplied) request ids.
|
||||
await engine.abort(request_ids)
|
||||
else:
|
||||
# The dev RL server runs AsyncLLM; abort everything it is tracking.
|
||||
# request_states is keyed by internal ids; parent_requests holds
|
||||
# parallel-sampling parents. Abort both as internal ids.
|
||||
from vllm.v1.engine.async_llm import AsyncLLM
|
||||
|
||||
assert isinstance(engine, AsyncLLM)
|
||||
op = engine.output_processor
|
||||
request_ids = [
|
||||
*op.request_states.keys(),
|
||||
*op.parent_requests.keys(),
|
||||
]
|
||||
await engine.abort(request_ids, internal=True)
|
||||
return JSONResponse(
|
||||
content={"status": "aborted", "aborted": len(request_ids)},
|
||||
status_code=HTTPStatus.OK.value,
|
||||
)
|
||||
except Exception as err: # pragma: no cover - defensive
|
||||
logger.exception("Failed to abort requests")
|
||||
return JSONResponse(
|
||||
content={"error": f"Failed to abort requests: {err}"},
|
||||
status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value,
|
||||
)
|
||||
|
||||
|
||||
@router.get("/is_paused")
|
||||
async def is_paused(raw_request: Request) -> JSONResponse:
|
||||
"""Return the current pause status."""
|
||||
|
||||
Reference in New Issue
Block a user