A high-throughput and memory-efficient inference and serving engine for LLMs
[Rust Frontend] Add external→internal request-id map for abort() (#45137)
Signed-off-by: Sahil Singh <sahiilsiingh37@gmail.com>
Sahil Singh committed
Jun 15, 2026 at 12:21 UTC
64833f8158236a7bddeeb89efc6a3bde5d16f468
13 files changed
+447
-27
rust/Cargo.lock
+1
@@ -5804,6 +5804,7 @@ dependencies = [
5804
"enum-as-inner",
5805
"expect-test",
5806
"futures",
5807
+ "parking_lot",
5808
"rmp-serde",
5809
"serde",
5810
"serde_json",
rust/src/chat/src/lib.rs
+6
@@ -235,6 +235,12 @@ impl ChatLlm {
235
Ok(token_ids)
236
}
237
238
+ /// Abort in-flight requests by their external (user-supplied) request ids.
239
+ pub async fn abort(&self, external_ids: &[String]) -> Result<()> {
240
+ self.text.abort(external_ids).await?;
241
+ Ok(())
242
+ }
243
+
244
/// Shut down the underlying LLM client and its background tasks.
245
pub async fn shutdown(self) -> Result<()> {
246
self.text.shutdown().await?;
rust/src/engine-core-client/src/client.rs
+4
@@ -490,6 +490,10 @@ impl EngineCoreClient {
490
return Ok(());
491
}
492
493
+ // Finalize the consumer streams first, before the engine round-trip.
494
+ let all_request_ids: Vec<String> = abortable.values().flatten().cloned().collect();
495
+ self.inner.abort_requests_locally(&all_request_ids);
496
+
497
for (engine_id, request_ids) in abortable {
498
self.inner.do_abort_requests(&engine_id, &request_ids).await?;
499
}
rust/src/engine-core-client/src/client/imp.rs
+15
@@ -1,5 +1,6 @@
1
use std::collections::{BTreeMap, BTreeSet};
2
use std::sync::Arc;
3
+use std::time::{SystemTime, UNIX_EPOCH};
4
5
use arc_swap::ArcSwapOption;
6
use parking_lot::Mutex;
@@ -126,6 +127,20 @@ impl ClientInner {
127
self.request_reg.lock().finish_many(request_ids)
128
}
129
130
+ /// Finalize client-initiated aborts by pushing a terminal `Abort` output
131
+ /// down each request's stream and removing it from the registry. Returns
132
+ /// the request ids that were still active. See [`RequestRegistry::abort_many`].
133
+ pub fn abort_requests_locally<'a>(
134
+ &self,
135
+ request_ids: impl IntoIterator<Item = &'a String>,
136
+ ) -> Vec<String> {
137
+ let timestamp = SystemTime::now()
138
+ .duration_since(UNIX_EPOCH)
139
+ .map(|d| d.as_secs_f64())
140
+ .unwrap_or(0.0);
141
+ self.request_reg.lock().abort_many(request_ids, timestamp)
142
+ }
143
+
144
/// Apply one scheduler stats update for the given engine to the local
145
/// routing state. Returns `false` if the engine is unknown to the
146
/// client.
rust/src/engine-core-client/src/client/state.rs
+29
-1
@@ -9,7 +9,7 @@ use crate::client::stream::EngineCoreStreamOutput;
9
use crate::error::{Error, Result};
10
use crate::protocol::stats::SchedulerStats;
11
use crate::protocol::utility::UtilityOutput;
12
-use crate::protocol::{EngineCoreEventType, EngineCoreOutput};
12
+use crate::protocol::{EngineCoreEventType, EngineCoreFinishReason, EngineCoreOutput};
13
use crate::transport::ConnectedEngine;
14
15
pub type OutputSender = mpsc::UnboundedSender<Result<EngineCoreStreamOutput>>;
@@ -289,6 +289,34 @@ impl RequestRegistry {
289
.collect()
290
}
291
292
+ /// Finalize client-initiated aborts: remove each request and push a
293
+ /// terminal output with `finish_reason = Abort` down its stream before the
294
+ /// sender drops. Returns the request ids that were still active.
295
+ pub fn abort_many<'a>(
296
+ &mut self,
297
+ request_ids: impl IntoIterator<Item = &'a String>,
298
+ timestamp: f64,
299
+ ) -> Vec<String> {
300
+ let mut aborted = Vec::new();
301
+ for request_id in request_ids {
302
+ let Some((sender, engine_id)) = self.remove(request_id) else {
303
+ continue;
304
+ };
305
+ let output = EngineCoreStreamOutput {
306
+ engine_index: engine_id.engine_index().unwrap_or(0),
307
+ timestamp,
308
+ output: EngineCoreOutput {
309
+ request_id: request_id.clone(),
310
+ finish_reason: Some(EngineCoreFinishReason::Abort),
311
+ ..EngineCoreOutput::default()
312
+ },
313
+ };
314
+ let _ = sender.send(Ok(output));
315
+ aborted.push(request_id.clone());
316
+ }
317
+ aborted
318
+ }
319
+
320
/// Remove one request from the local registry. Returns the tracked entry if
321
/// it exists.
322
#[must_use]
rust/src/engine-core-client/src/tests/client.rs
+2
-2
@@ -1939,7 +1939,7 @@ async fn multi_engine_abort_is_grouped_and_utility_fans_out_to_all_engines() {
1939
1940
let (shutdown_tx_0, engine_task_0) = spawn_mock_engine_task(
1941
handshake_address.clone(),
1942
- b"engine-0".to_vec(),
1942
+ EngineId::from_engine_index(0).into_frame().to_vec(),
1943
|dealer, push| {
1944
Box::pin(async move {
1945
let utility = recv_engine_message(dealer).await;
@@ -1993,7 +1993,7 @@ async fn multi_engine_abort_is_grouped_and_utility_fans_out_to_all_engines() {
1993
tokio::time::sleep(Duration::from_millis(50)).await;
1994
let (shutdown_tx_1, engine_task_1) = spawn_mock_engine_task(
1995
handshake_address.clone(),
1996
- b"engine-1".to_vec(),
1996
+ EngineId::from_engine_index(1).into_frame().to_vec(),
1997
|dealer, push| {
1998
Box::pin(async move {
1999
let utility = recv_engine_message(dealer).await;
rust/src/llm/Cargo.toml
+1
@@ -11,6 +11,7 @@ test-util = []
11
easy-ext.workspace = true
12
enum-as-inner.workspace = true
13
futures.workspace = true
14
+parking_lot.workspace = true
15
serde.workspace = true
16
serde_json.workspace = true
17
thiserror.workspace = true
rust/src/llm/src/inflight.rs
new
+179
@@ -0,0 +1,179 @@
1
+//! Tracking of the external→internal request-id mapping for in-flight requests.
2
+//!
3
+//! When request-id randomization is enabled (the default), [`crate::Llm`]
4
+//! rewrites the external (user-supplied) request id into a unique internal
5
+//! engine id before reaching engine-core. Engine-core only ever knows the
6
+//! internal id, so aborting a request by its external id requires resolving it
7
+//! back to the internal id(s) first.
8
+
9
+use std::collections::HashMap;
10
+use std::sync::{Arc, Weak};
11
+
12
+use parking_lot::Mutex;
13
+
14
+/// external id → internal id → number of live guards holding that edge.
15
+type InflightMap = HashMap<String, HashMap<String, usize>>;
16
+
17
+/// Maps external (user-supplied) request ids to the set of live internal engine
18
+/// request ids they currently expand into.
19
+///
20
+/// One external id may map to multiple internal ids: duplicate external ids
21
+/// submitted concurrently each get their own randomized internal id, and an
22
+/// abort by the shared external id must reach all of them. Edges are
23
+/// refcounted: with randomization disabled the same (external, internal) pair
24
+/// can be tracked by several guards in sequence (e.g. a finished request whose
25
+/// stream is still held alongside a fresh submission reusing the id), and the
26
+/// edge must survive until the last guard drops.
27
+#[derive(Default)]
28
+pub(crate) struct InflightRequests {
29
+ map: Arc<Mutex<InflightMap>>,
30
+}
31
+
32
+impl InflightRequests {
33
+ pub(crate) fn new() -> Self {
34
+ Self::default()
35
+ }
36
+
37
+ /// Record that `internal` is now an in-flight engine request for the
38
+ /// `external` request id, returning a guard that removes the edge when the
39
+ /// request's output stream is dropped (on clean finish or cancellation).
40
+ pub(crate) fn track(&self, external: String, internal: String) -> RequestGuard {
41
+ *self
42
+ .map
43
+ .lock()
44
+ .entry(external.clone())
45
+ .or_default()
46
+ .entry(internal.clone())
47
+ .or_insert(0) += 1;
48
+ RequestGuard {
49
+ map: Arc::downgrade(&self.map),
50
+ external,
51
+ internal,
52
+ }
53
+ }
54
+
55
+ /// Resolve external request ids to the internal engine ids currently
56
+ /// in-flight for them. Unknown or already-finished ids contribute nothing.
57
+ pub(crate) fn resolve(&self, external_ids: &[String]) -> Vec<String> {
58
+ let map = self.map.lock();
59
+ external_ids
60
+ .iter()
61
+ .filter_map(|external| map.get(external))
62
+ .flat_map(|internal_ids| internal_ids.keys())
63
+ .cloned()
64
+ .collect()
65
+ }
66
+
67
+ #[cfg(test)]
68
+ fn is_empty(&self) -> bool {
69
+ self.map.lock().is_empty()
70
+ }
71
+}
72
+
73
+/// RAII guard that releases one refcount on a single external→internal edge
74
+/// when dropped, removing the edge once no live guard holds it.
75
+///
76
+/// Held by the per-request output stream, so cleanup runs whether the stream
77
+/// terminates cleanly or is cancelled. A [`Weak`] handle is used so a stream
78
+/// outliving its owning [`InflightRequests`] does not keep the map alive.
79
+pub(crate) struct RequestGuard {
80
+ map: Weak<Mutex<InflightMap>>,
81
+ external: String,
82
+ internal: String,
83
+}
84
+
85
+impl Drop for RequestGuard {
86
+ fn drop(&mut self) {
87
+ let Some(map) = self.map.upgrade() else {
88
+ return;
89
+ };
90
+ let mut map = map.lock();
91
+ if let Some(internal_ids) = map.get_mut(&self.external) {
92
+ if let Some(count) = internal_ids.get_mut(&self.internal) {
93
+ *count -= 1;
94
+ if *count == 0 {
95
+ internal_ids.remove(&self.internal);
96
+ }
97
+ }
98
+ if internal_ids.is_empty() {
99
+ map.remove(&self.external);
100
+ }
101
+ }
102
+ }
103
+}
104
+
105
+#[cfg(test)]
106
+mod tests {
107
+ use super::*;
108
+
109
+ #[test]
110
+ fn resolves_external_to_internal() {
111
+ let inflight = InflightRequests::new();
112
+ let _guard = inflight.track("ext".to_string(), "ext-abc".to_string());
113
+
114
+ assert_eq!(
115
+ inflight.resolve(&["ext".to_string()]),
116
+ vec!["ext-abc".to_string()]
117
+ );
118
+ assert!(inflight.resolve(&["unknown".to_string()]).is_empty());
119
+ }
120
+
121
+ #[test]
122
+ fn one_external_maps_to_many_internal() {
123
+ let inflight = InflightRequests::new();
124
+ let _g1 = inflight.track("dup".to_string(), "dup-1".to_string());
125
+ let _g2 = inflight.track("dup".to_string(), "dup-2".to_string());
126
+
127
+ let mut resolved = inflight.resolve(&["dup".to_string()]);
128
+ resolved.sort();
129
+ assert_eq!(resolved, vec!["dup-1".to_string(), "dup-2".to_string()]);
130
+ }
131
+
132
+ #[test]
133
+ fn dropping_guard_removes_only_its_own_edge_then_cleans_empty_key() {
134
+ let inflight = InflightRequests::new();
135
+ let g1 = inflight.track("dup".to_string(), "dup-1".to_string());
136
+ let g2 = inflight.track("dup".to_string(), "dup-2".to_string());
137
+
138
+ drop(g1);
139
+ assert_eq!(
140
+ inflight.resolve(&["dup".to_string()]),
141
+ vec!["dup-2".to_string()]
142
+ );
143
+
144
+ drop(g2);
145
+ assert!(inflight.resolve(&["dup".to_string()]).is_empty());
146
+ assert!(
147
+ inflight.is_empty(),
148
+ "empty external key must be removed, not left dangling"
149
+ );
150
+ }
151
+
152
+ #[test]
153
+ fn identical_edges_are_refcounted_across_guards() {
154
+ // With request-id randomization disabled, internal == external, so two
155
+ // tracked requests can share the exact same edge. Dropping one guard
156
+ // (e.g. a stale stream, or the error path of a rejected duplicate
157
+ // submission) must not untrack the other still-live request.
158
+ let inflight = InflightRequests::new();
159
+ let g1 = inflight.track("x".to_string(), "x".to_string());
160
+ let g2 = inflight.track("x".to_string(), "x".to_string());
161
+
162
+ drop(g1);
163
+ assert_eq!(inflight.resolve(&["x".to_string()]), vec!["x".to_string()]);
164
+
165
+ drop(g2);
166
+ assert!(inflight.resolve(&["x".to_string()]).is_empty());
167
+ assert!(inflight.is_empty());
168
+ }
169
+
170
+ #[test]
171
+ fn guard_drop_is_a_noop_after_inflight_is_gone() {
172
+ let guard = {
173
+ let inflight = InflightRequests::new();
174
+ inflight.track("ext".to_string(), "ext-abc".to_string())
175
+ };
176
+ // Dropping the guard after the owning map is gone must not panic.
177
+ drop(guard);
178
+ }
179
+}
rust/src/llm/src/lib.rs
+33
-3
@@ -2,6 +2,7 @@ use tracing::Span;
2
use vllm_engine_core_client::EngineCoreClient;
3
4
mod error;
5
+mod inflight;
6
mod log_stats;
7
mod output;
8
mod request;
@@ -15,18 +16,22 @@ pub use output::{
16
pub use request::GenerateRequest;
17
pub use vllm_engine_core_client::protocol::logprobs::{Logprobs, PositionLogprobs, TokenLogprob};
18
19
+use crate::inflight::InflightRequests;
20
use crate::log_stats::StatsLogger;
21
use crate::request_metrics::RequestMetricsTracker;
22
21
-/// Thin generate-only facade over [`EngineCoreClient`].
23
+/// Thin generate-and-abort facade over [`EngineCoreClient`].
24
///
25
/// This mirrors the narrow public shape of Python `AsyncLLM.generate()` and
26
/// `abort()`, but keeps the boundary close to raw engine-core requests and
25
-/// outputs.
27
+/// outputs. It tracks an in-flight external→internal request-id index (see
28
+/// [`InflightRequests`]) so that aborts issued against external (user-supplied)
29
+/// ids can be resolved to the internal engine ids that engine-core understands.
30
pub struct Llm {
31
client: EngineCoreClient,
32
randomize_request_id: bool,
33
stats_logger: Option<StatsLogger>,
34
+ inflight: InflightRequests,
35
}
36
37
impl Llm {
@@ -37,6 +42,7 @@ impl Llm {
42
client,
43
randomize_request_id: true,
44
stats_logger: None,
45
+ inflight: InflightRequests::new(),
46
}
47
}
48
@@ -72,9 +78,15 @@ impl Llm {
78
pub async fn generate(&self, req: GenerateRequest) -> Result<GenerateOutputStream> {
79
let prepared = req.prepare(self.randomize_request_id)?;
80
let prompt_token_ids = prepared.prompt_token_ids().into();
81
+ let external_request_id = prepared
82
+ .engine_request
83
+ .external_req_id
84
+ .clone()
85
+ .expect("prepare always sets external_req_id");
86
+ let internal_request_id = prepared.engine_request.request_id.clone();
87
88
// Record internal engine-core request ID in the current tracing span.
77
- Span::current().record("engine_request_id", &prepared.engine_request.request_id);
89
+ Span::current().record("engine_request_id", &internal_request_id);
90
91
let request_metrics = RequestMetricsTracker::new(
92
self.client.model_name().to_string(),
@@ -84,14 +96,32 @@ impl Llm {
96
1,
97
);
98
let stream = self.client.call(prepared.engine_request).await?;
99
+ let guard = self.inflight.track(external_request_id, internal_request_id);
100
101
Ok(GenerateOutputStream::new(
102
prompt_token_ids,
103
stream,
104
request_metrics,
105
+ guard,
106
))
107
}
108
109
+ /// Abort in-flight requests by their external (user-supplied) request ids.
110
+ ///
111
+ /// External ids are resolved to the internal engine ids actually known to
112
+ /// engine-core (one external id may map to several internal ids). Unknown
113
+ /// or already-finished ids resolve to nothing and are a safe no-op. The
114
+ /// tracking entries themselves are removed when the corresponding output
115
+ /// streams are dropped, not here.
116
+ pub async fn abort(&self, external_ids: &[String]) -> Result<()> {
117
+ let internal_ids = self.inflight.resolve(external_ids);
118
+ if internal_ids.is_empty() {
119
+ return Ok(());
120
+ }
121
+ self.client.abort(&internal_ids).await?;
122
+ Ok(())
123
+ }
124
+
125
/// Shut down the underlying engine-core client and its background tasks.
126
pub async fn shutdown(self) -> Result<()> {
127
self.client.shutdown().await?;
rust/src/llm/src/output.rs
+10
-2
@@ -12,6 +12,7 @@ use vllm_engine_core_client::protocol::{EngineCoreFinishReason, StopReason};
12
use vllm_engine_core_client::{AbortCause, EngineCoreOutputStream};
13
14
use crate::error::Result;
15
+use crate::inflight::RequestGuard;
16
use crate::request_metrics::{RequestMetricsTracker, current_unix_timestamp_secs};
17
18
/// Token usage metadata for one request.
@@ -195,12 +196,17 @@ impl GenerateOutput {
196
197
/// Stream of per-request generate outputs for one request.
198
///
198
-/// - A normal termination of the stream represents a clean completion of the request.
199
-/// - For errors, unexpected closes, or explicit aborts, the stream terminates with an error.
199
+/// - A normal termination of the stream represents a clean completion of the
200
+/// request, including a client-initiated abort, which yields a final output
201
+/// with `finish_reason = Abort` before the stream ends.
202
+/// - For errors or unexpected engine-side closes, the stream terminates with an error.
203
pub struct GenerateOutputStream {
204
pending_prompt_info: Option<GeneratePromptInfo>,
205
raw_stream: EngineCoreOutputStream,
206
request_metrics: RequestMetricsTracker,
207
+ /// Removes this request's external→internal tracking edge on drop. Held for
208
+ /// its `Drop` side effect only; never read directly.
209
+ _request_guard: RequestGuard,
210
}
211
212
impl GenerateOutputStream {
@@ -210,6 +216,7 @@ impl GenerateOutputStream {
216
prompt_token_ids: Arc<[u32]>,
217
raw_stream: EngineCoreOutputStream,
218
request_metrics: RequestMetricsTracker,
219
+ request_guard: RequestGuard,
220
) -> Self {
221
Self {
222
pending_prompt_info: Some(GeneratePromptInfo {
@@ -218,6 +225,7 @@ impl GenerateOutputStream {
225
}),
226
raw_stream,
227
request_metrics,
228
+ _request_guard: request_guard,
229
}
230
}
231
rust/src/llm/src/request_metrics.rs
+25
-19
@@ -98,27 +98,33 @@ impl RequestMetricsTracker {
98
self.observe_events(engine_index, events);
99
}
100
101
- if self.is_prefilling {
102
- if let Some(prefill_stats) = &output.prefill_stats {
103
- record_prompt_tokens(&self.model_name, engine_index, prefill_stats);
101
+ // Only outputs that actually carry tokens drive token-timing metrics.
102
+ // A terminal output with no new tokens (e.g. the synthesized abort
103
+ // output) must not log a stray time-to-first-token or inter-token
104
+ // sample.
105
+ if !output.new_token_ids.is_empty() {
106
+ if self.is_prefilling {
107
+ if let Some(prefill_stats) = &output.prefill_stats {
108
+ record_prompt_tokens(&self.model_name, engine_index, prefill_stats);
109
+ }
110
+ self.first_token_latency = received_at - self.arrival_time;
111
+ observe_time_to_first_token_seconds(
112
+ &self.model_name,
113
+ engine_index,
114
+ self.first_token_latency,
115
+ );
116
+ self.first_token_ts = batch_timestamp;
117
+ self.is_prefilling = false;
118
+ } else if self.last_token_ts > 0.0 {
119
+ observe_inter_token_latency_seconds(
120
+ &self.model_name,
121
+ engine_index,
122
+ batch_timestamp - self.last_token_ts,
123
+ );
124
}
105
- self.first_token_latency = received_at - self.arrival_time;
106
- observe_time_to_first_token_seconds(
107
- &self.model_name,
108
- engine_index,
109
- self.first_token_latency,
110
- );
111
- self.first_token_ts = batch_timestamp;
112
- self.is_prefilling = false;
113
- } else if self.last_token_ts > 0.0 {
114
- observe_inter_token_latency_seconds(
115
- &self.model_name,
116
- engine_index,
117
- batch_timestamp - self.last_token_ts,
118
- );
119
- }
125
121
- self.last_token_ts = batch_timestamp;
126
+ self.last_token_ts = batch_timestamp;
127
+ }
128
}
129
130
/// Emit the terminal request metrics once a finished output has been
rust/src/llm/tests/generate.rs
+136
@@ -554,6 +554,142 @@ async fn duplicate_external_request_ids_are_randomized_before_reaching_engine_co
554
llm.shutdown().await.unwrap();
555
}
556
557
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
558
+async fn abort_resolves_external_request_id_to_internal_before_reaching_engine() {
559
+ let ipc = IpcNamespace::new().unwrap();
560
+ let handshake_address = ipc.handshake_endpoint();
561
+ let engine_id = b"engine-abort".to_vec();
562
+
563
+ let (shutdown_tx, engine_task) = spawn_mock_engine_task(
564
+ handshake_address.clone(),
565
+ engine_id.clone(),
566
+ |dealer, push| {
567
+ Box::pin(async move {
568
+ let add = recv_engine_message(dealer).await;
569
+ assert_eq!(add[0].as_ref(), &[0x00]);
570
+ let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap();
571
+ assert_eq!(request.external_req_id.as_deref(), Some("req-abort"));
572
+ assert!(request.request_id.starts_with("req-abort-"));
573
+ assert_ne!(request.request_id, "req-abort");
574
+
575
+ send_outputs(
576
+ push,
577
+ EngineCoreOutputs {
578
+ outputs: vec![request_output(&request.request_id, vec![7], None)],
579
+ ..Default::default()
580
+ },
581
+ )
582
+ .await;
583
+
584
+ // The abort frame must carry the internal engine id, not the
585
+ // external "req-abort" id the caller aborted by.
586
+ let abort =
587
+ timeout(Duration::from_secs(1), recv_engine_message(dealer)).await.unwrap();
588
+ assert_eq!(abort[0].as_ref(), &[0x01]);
589
+ let aborted_ids: Vec<String> = rmp_serde::from_slice(&abort[1]).unwrap();
590
+ assert_eq!(aborted_ids, vec![request.request_id]);
591
+ })
592
+ },
593
+ );
594
+
595
+ let llm = connect_async_llm_with_ipc(handshake_address, 0, "test-model", &ipc).await;
596
+ let mut stream = llm.generate(sample_generate_request("req-abort", 4)).await.unwrap();
597
+ let internal_id = stream.request_id().to_string();
598
+ assert_ne!(internal_id, "req-abort");
599
+
600
+ assert_eq!(stream.next().await.unwrap().unwrap().token_ids, vec![7]);
601
+
602
+ // Abort by the external id; engine-core only knows the internal id.
603
+ llm.abort(&["req-abort".to_string()]).await.unwrap();
604
+
605
+ // The consumer stream is finalized locally with a clean abort terminal
606
+ // rather than hanging or surfacing as RequestStreamClosed. The engine sends
607
+ // no final output for a client abort, so this output is synthesized.
608
+ let terminal = stream.next().await.unwrap().unwrap();
609
+ assert_eq!(terminal.finish_reason, Some(FinishReason::Abort));
610
+ assert!(terminal.token_ids.is_empty());
611
+ assert!(stream.next().await.is_none());
612
+
613
+ let _ = shutdown_tx.send(());
614
+ engine_task.await.unwrap();
615
+ drop(stream);
616
+ llm.shutdown().await.unwrap();
617
+}
618
+
619
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
620
+async fn abort_by_external_id_aborts_all_internal_requests() {
621
+ let ipc = IpcNamespace::new().unwrap();
622
+ let handshake_address = ipc.handshake_endpoint();
623
+ let engine_id = b"engine-abort-many".to_vec();
624
+
625
+ let (shutdown_tx, engine_task) = spawn_mock_engine_task(
626
+ handshake_address.clone(),
627
+ engine_id.clone(),
628
+ |dealer, push| {
629
+ Box::pin(async move {
630
+ let add_1 = recv_engine_message(dealer).await;
631
+ assert_eq!(add_1[0].as_ref(), &[0x00]);
632
+ let request_1: EngineCoreRequest = rmp_serde::from_slice(&add_1[1]).unwrap();
633
+
634
+ let add_2 = recv_engine_message(dealer).await;
635
+ assert_eq!(add_2[0].as_ref(), &[0x00]);
636
+ let request_2: EngineCoreRequest = rmp_serde::from_slice(&add_2[1]).unwrap();
637
+
638
+ assert_eq!(request_1.external_req_id.as_deref(), Some("req-dup-abort"));
639
+ assert_eq!(request_2.external_req_id.as_deref(), Some("req-dup-abort"));
640
+ assert_ne!(request_1.request_id, request_2.request_id);
641
+
642
+ send_outputs(
643
+ push,
644
+ EngineCoreOutputs {
645
+ outputs: vec![
646
+ request_output(&request_1.request_id, vec![7], None),
647
+ request_output(&request_2.request_id, vec![8], None),
648
+ ],
649
+ ..Default::default()
650
+ },
651
+ )
652
+ .await;
653
+
654
+ // A single abort by the shared external id must abort both
655
+ // internal engine ids it expanded into.
656
+ let abort =
657
+ timeout(Duration::from_secs(1), recv_engine_message(dealer)).await.unwrap();
658
+ assert_eq!(abort[0].as_ref(), &[0x01]);
659
+ let mut aborted_ids: Vec<String> = rmp_serde::from_slice(&abort[1]).unwrap();
660
+ aborted_ids.sort();
661
+ let mut expected = vec![request_1.request_id, request_2.request_id];
662
+ expected.sort();
663
+ assert_eq!(aborted_ids, expected);
664
+ })
665
+ },
666
+ );
667
+
668
+ let llm = connect_async_llm_with_ipc(handshake_address, 0, "test-model", &ipc).await;
669
+ let mut stream_1 = llm.generate(sample_generate_request("req-dup-abort", 4)).await.unwrap();
670
+ let mut stream_2 = llm.generate(sample_generate_request("req-dup-abort", 4)).await.unwrap();
671
+ assert_ne!(stream_1.request_id(), stream_2.request_id());
672
+
673
+ assert_eq!(stream_1.next().await.unwrap().unwrap().token_ids, vec![7]);
674
+ assert_eq!(stream_2.next().await.unwrap().unwrap().token_ids, vec![8]);
675
+
676
+ llm.abort(&["req-dup-abort".to_string()]).await.unwrap();
677
+
678
+ // Both internal requests the external id expanded into are finalized with a
679
+ // clean abort terminal.
680
+ for stream in [&mut stream_1, &mut stream_2] {
681
+ let terminal = stream.next().await.unwrap().unwrap();
682
+ assert_eq!(terminal.finish_reason, Some(FinishReason::Abort));
683
+ assert!(stream.next().await.is_none());
684
+ }
685
+
686
+ let _ = shutdown_tx.send(());
687
+ engine_task.await.unwrap();
688
+ drop(stream_1);
689
+ drop(stream_2);
690
+ llm.shutdown().await.unwrap();
691
+}
692
+
693
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
694
async fn generate_records_request_metrics_in_prometheus_output() {
695
let ipc = IpcNamespace::new().unwrap();
rust/src/text/src/lib.rs
+6
@@ -151,6 +151,12 @@ impl TextLlm {
151
Ok((text_request, raw_stream))
152
}
153
154
+ /// Abort in-flight requests by their external (user-supplied) request ids.
155
+ pub async fn abort(&self, external_ids: &[String]) -> Result<()> {
156
+ self.llm.abort(external_ids).await?;
157
+ Ok(())
158
+ }
159
+
160
/// Shut down the underlying LLM client and its background tasks.
161
pub async fn shutdown(self) -> Result<()> {
162
self.llm.shutdown().await?;