Skip to content

GitLab

Unverified Commit d870d4b0 authored by Yan Ru Pei's avatar Yan Ru Pei Committed by GitHub
Browse files

chore: mocker states do not need arc mutex (#3883) 


Signed-off-by: default avatarPeaBrane <yanrpei@gmail.com>
parent ee275cbf
Show whitespace changes
Inline Side-by-side
Showing with 71 additions and 145 deletions
lib/llm/src/mocker/scheduler.rs
View file @ d870d4b0
......@@ -38,8 +38,7 @@ use crate::tokens::BlockHash;
use crate::tokens::blocks::UniqueBlock;
use std::collections::HashMap;
use std::collections::VecDeque;
use std::sync::Arc;
use tokio::sync::{Mutex, mpsc};
use tokio::sync::mpsc;
use tokio::time::Duration;
use tokio_util::sync::CancellationToken;
use uuid::Uuid;
......@@ -238,8 +237,6 @@ impl SchedulerState {
/// Manages scheduling of requests using KvManager resources
#[derive(Clone)]
pub struct Scheduler {
state: Arc<Mutex<SchedulerState>>,
kv_manager: Arc<Mutex<KvManager>>,
request_tx: mpsc::UnboundedSender<DirectRequest>,
metrics_rx: tokio::sync::watch::Receiver<ForwardPassMetrics>,
}
......@@ -253,8 +250,6 @@ impl Scheduler {
kv_events_tx: Option<mpsc::UnboundedSender<KvCacheEventData>>,
cancellation_token: Option<CancellationToken>,
) -> Self {
let state = Arc::new(Mutex::new(SchedulerState::new(args.max_num_batched_tokens)));
// Create internal channel for KV events only if needed
let (block_resp_tx, mut block_resp_rx) = if kv_events_tx.is_some() {
let (tx, rx) = mpsc::unbounded_channel::<MoveBlockResponse>();
......@@ -263,13 +258,6 @@ impl Scheduler {
(None, None)
};
let kv_manager = Arc::new(Mutex::new(KvManager::new_with_sender(
args.num_gpu_blocks,
args.block_size,
block_resp_tx,
)));
let hit_rates = Arc::new(Mutex::new(VecDeque::with_capacity(1000)));
// Assert speedup_ratio is greater than 0
assert!(
args.speedup_ratio > 0.0,
......@@ -284,30 +272,26 @@ impl Scheduler {
let (metrics_tx, metrics_rx) =
tokio::sync::watch::channel::<ForwardPassMetrics>(initial_metrics);
// Create a clone for the background task
let state_clone = state.clone();
let kv_manager_clone = kv_manager.clone();
let output_tx_clone = output_tx.clone();
let cancel_token_clone = cancellation_token.unwrap_or_default().clone();
// Spawn main background task with cancellation token
tokio::spawn(async move {
// Create state and kv_manager as local variables owned by this task
let mut state = SchedulerState::new(args.max_num_batched_tokens);
let mut kv_manager =
KvManager::new_with_sender(args.num_gpu_blocks, args.block_size, block_resp_tx);
let mut hit_rates = VecDeque::with_capacity(1000);
let mut should_schedule = true;
loop {
{
let state_guard = state_clone.lock().await;
// Enqueue new request, blocks until at least one is received, so no redundant work is done
// TODO: clean this up? double lock acquisition is ugly, but needed to not hold the lock forever
if state_guard.is_empty() {
drop(state_guard);
if state.is_empty() {
let Some(request) = request_rx.recv().await else {
tracing::warn!("request sender is dropped");
break;
};
let mut state_guard = state_clone.lock().await;
state_guard.receive(request);
state.receive(request);
}
}
......@@ -316,7 +300,6 @@ impl Scheduler {
// Enqueue new request
Some(request) = request_rx.recv() => {
let mut state = state_clone.lock().await;
state.receive(request);
}
......@@ -327,20 +310,17 @@ impl Scheduler {
continue;
}
let mut state_guard = state_clone.lock().await;
let kv_manager_guard = kv_manager_clone.lock().await;
// Process DirectRequests, converting them to ActiveSequence and scheduling them until we can't
// schedule anymore.
let mut current_blocks = kv_manager_guard.num_active_blocks();
let mut current_tokens = state_guard.active_tokens + state_guard.waiting_tokens;
let mut current_seqs = state_guard.num_active_requests();
let mut current_blocks = kv_manager.num_active_blocks();
let mut current_tokens = state.active_tokens + state.waiting_tokens;
let mut current_seqs = state.num_active_requests();
while let Some((uuid, request)) = state_guard.next() {
while let Some((uuid, request)) = state.next() {
let active_sequence = get_active_sequence(request, args.block_size, args.enable_prefix_caching);
// Update predictive budgets
let prefill_cost = kv_manager_guard.get_prefill_cost(&active_sequence);
let prefill_cost = kv_manager.get_prefill_cost(&active_sequence);
let total_tokens = active_sequence.len();
// this is conservative, assumes no cache hit so never over-schedules
let new_blocks = (total_tokens as u32).div_ceil(args.block_size as u32) as usize;
......@@ -351,7 +331,7 @@ impl Scheduler {
current_seqs += 1;
// Check various budgets to see if possible to schedule
let under_block_budget = current_blocks as f64 <= (1. - args.watermark) * kv_manager_guard.max_capacity() as f64;
let under_block_budget = current_blocks as f64 <= (1. - args.watermark) * kv_manager.max_capacity() as f64;
// If chunked prefill is enabled, we can be under token budget when scheduling
let comparison_tokens = if args.enable_chunked_prefill {current_tokens - new_tokens} else {current_tokens};
let under_token_budget = args.max_num_batched_tokens.is_none_or(|limit| comparison_tokens <= limit);
......@@ -359,21 +339,18 @@ impl Scheduler {
// Cannot schedule, put first in line instead
if !(under_block_budget && under_token_budget && under_seq_budget) {
state_guard.first_in_line(uuid, Request::Active(active_sequence));
state.first_in_line(uuid, Request::Active(active_sequence));
break;
}
// Compute and store hit rate
let hit_rate = if !active_sequence.is_empty() { 1.0 - (new_tokens as f32 / active_sequence.len() as f32) } else { 0.0 };
{
let mut hit_rates_guard = hit_rates.lock().await;
hit_rates_guard.push_back(hit_rate);
if hit_rates_guard.len() > 1000 {
hit_rates_guard.pop_front();
}
hit_rates.push_back(hit_rate);
if hit_rates.len() > 1000 {
hit_rates.pop_front();
}
state_guard.move_to_prefill(uuid, active_sequence, prefill_cost);
state.move_to_prefill(uuid, active_sequence, prefill_cost);
should_schedule = false;
}
}
......@@ -385,11 +362,8 @@ impl Scheduler {
}
// Simulates prefill + decode
let mut state_guard = state_clone.lock().await;
let mut kv_manager_guard = kv_manager_clone.lock().await;
// Base time needed for decoding using active percentage and quadratic formula
let active_perc = kv_manager_guard.get_active_perc();
let active_perc = kv_manager.get_active_perc();
let decoding_time = -5.47 * active_perc.powi(2) + 43.88 * active_perc + 19.44;
let mut total_time = Duration::from_secs_f64(decoding_time / 1000.0);
......@@ -399,17 +373,15 @@ impl Scheduler {
maybe_creation_signal,
block_hashes,
is_full_prefill,
)) = state_guard.try_prefill()
)) = state.try_prefill()
{
// NOTE: Prefill cost/time is always incremented for new blocks, even if they
// could be cached by other requests in the same batch. This matches vLLM behavior.
total_time += Duration::from_secs_f64(prefill_compute / 1000.0);
if let Some(creation_signal) = maybe_creation_signal {
if !process_signals(
&mut kv_manager_guard,
std::slice::from_ref(&creation_signal),
) {
if !process_signals(&mut kv_manager, std::slice::from_ref(&creation_signal))
{
panic!("Block allocation for prefilling cannot fail.");
}
......@@ -428,36 +400,25 @@ impl Scheduler {
}
}
state_guard.reset_active_tokens();
{
let hit_rates_guard = hit_rates.lock().await;
let metrics = get_fwd_pass_metrics(
&state_guard,
&kv_manager_guard,
&hit_rates_guard,
dp_rank,
);
let _ = metrics_tx.send(metrics);
}
state.reset_active_tokens();
// Process decoding
let uuids: Vec<Uuid> = state_guard.decode.keys().cloned().collect();
let uuids: Vec<Uuid> = state.decode.keys().cloned().collect();
if !uuids.is_empty() {
should_schedule = true
};
for uuid in uuids {
let Some(sequence) = state_guard.run(uuid) else {
let Some(sequence) = state.run(uuid) else {
continue;
};
let signals = sequence.generate();
// Process all signals with the KvManager
// Handling of preemption on failure
if !process_signals(&mut kv_manager_guard, &signals) {
if !process_signals(&mut kv_manager, &signals) {
sequence.pop(); // revert the failed generation op
for signal in state_guard.preempt() {
kv_manager_guard.process(&signal);
for signal in state.preempt() {
kv_manager.process(&signal);
}
continue;
}
......@@ -477,7 +438,7 @@ impl Scheduler {
let mut send_failed = false;
if should_output {
send_failed = output_tx_clone.as_ref().is_some_and(|tx| {
send_failed = output_tx.as_ref().is_some_and(|tx| {
tx.send(OutputSignal {
uuid,
completed: is_complete,
......@@ -488,30 +449,23 @@ impl Scheduler {
if send_failed {
for signal in &sequence.free_signal() {
kv_manager_guard.process(signal);
kv_manager.process(signal);
}
}
{
let hit_rates_guard = hit_rates.lock().await;
let metrics = get_fwd_pass_metrics(
&state_guard,
&kv_manager_guard,
&hit_rates_guard,
dp_rank,
);
let _ = metrics_tx.send(metrics);
}
if send_failed || is_complete {
state_guard.complete(&uuid);
state.complete(&uuid);
continue;
}
}
// Send metrics once per forward pass (after all prefill and decode processing)
{
let metrics = get_fwd_pass_metrics(&state, &kv_manager, &hit_rates, dp_rank);
let _ = metrics_tx.send(metrics);
}
// Sleep once for the adjusted duration
drop(kv_manager_guard);
drop(state_guard);
let adjusted_time =
Duration::from_secs_f64(total_time.as_secs_f64() / args.speedup_ratio);
if adjusted_time.as_millis() > 0 {
......@@ -521,8 +475,6 @@ impl Scheduler {
});
Self {
state,
kv_manager,
request_tx,
metrics_rx,
}
......@@ -537,31 +489,6 @@ impl Scheduler {
self.request_tx.clone()
}
pub async fn waiting_count(&self) -> usize {
let state = self.state.lock().await;
state.waiting.len()
}
pub async fn running_count(&self) -> usize {
let state = self.state.lock().await;
state.decode.len()
}
pub async fn waiting_tokens(&self) -> usize {
let state = self.state.lock().await;
state.waiting_tokens
}
pub async fn active_tokens(&self) -> usize {
let state = self.state.lock().await;
state.active_tokens
}
pub async fn kv_usage_perc(&self) -> f64 {
let kv_manager = self.kv_manager.lock().await;
kv_manager.current_capacity_perc()
}
/// Get a watch receiver for forward pass metrics
pub fn metrics_receiver(&self) -> tokio::sync::watch::Receiver<ForwardPassMetrics> {
self.metrics_rx.clone()
......@@ -648,12 +575,9 @@ fn get_active_sequence(
/// This validation is important because in normal operation, the only legitimate failure
/// case should be when trying to acquire a new generation block - any other failures would
/// indicate an unexpected state in the system.
fn process_signals(
kv_manager_guard: &mut tokio::sync::MutexGuard<'_, KvManager>,
signals: &[MoveBlock],
) -> bool {
fn process_signals(kv_manager: &mut KvManager, signals: &[MoveBlock]) -> bool {
for signal in signals {
if kv_manager_guard.process(signal) {
if kv_manager.process(signal) {
continue;
}
......@@ -666,7 +590,7 @@ fn process_signals(
// Verify the signal contains exactly one block
let num_blocks = blocks.len();
let num_active_blocks = kv_manager_guard.num_active_blocks();
let num_active_blocks = kv_manager.num_active_blocks();
if num_blocks != 1 {
panic!(
"Failed signal is Invalid. Tried to create (prefill) {num_blocks} blocks on top of {num_active_blocks} active blocks."
......@@ -691,6 +615,30 @@ mod tests {
use std::time::Duration;
use tokio::time::interval;
/// Helper function to verify that the scheduler is idle (no active or waiting requests/resources)
fn assert_scheduler_idle(metrics: &ForwardPassMetrics) {
assert_eq!(
metrics.worker_stats.request_active_slots, 0,
"Expected 0 active slots, got {}",
metrics.worker_stats.request_active_slots
);
assert_eq!(
metrics.worker_stats.num_requests_waiting, 0,
"Expected 0 waiting requests, got {}",
metrics.worker_stats.num_requests_waiting
);
assert_eq!(
metrics.kv_stats.kv_active_blocks, 0,
"Expected 0 active blocks, got {}",
metrics.kv_stats.kv_active_blocks
);
assert_eq!(
metrics.kv_stats.gpu_cache_usage_perc, 0.0,
"Expected 0% GPU cache usage, got {}",
metrics.kv_stats.gpu_cache_usage_perc
);
}
#[rstest]
#[case::case_1(false, false, false)]
#[case::case_2(false, true, false)]
......@@ -820,17 +768,11 @@ mod tests {
"Received {received_tokens} tokens but expected exactly {expected_tokens}"
);
let active_tokens = scheduler.active_tokens().await;
assert!(
active_tokens == 0,
"Scheduler still have {active_tokens} active tokens but expected 0"
);
// Wait a bit for final metrics update to propagate
tokio::time::sleep(Duration::from_millis(100)).await;
let waiting_tokens = scheduler.waiting_tokens().await;
assert!(
waiting_tokens == 0,
"Scheduler still have {waiting_tokens} waiting tokens but expected 0"
);
let metrics = scheduler.metrics_receiver().borrow().clone();
assert_scheduler_idle(&metrics);
}
#[tokio::test]
......@@ -913,12 +855,7 @@ mod tests {
// Verify forward pass metrics
let metrics = metrics_rx.borrow().clone();
assert_eq!(
metrics.worker_stats.num_requests_waiting, 0,
"Expected no waiting requests, got {}",
metrics.worker_stats.num_requests_waiting
);
assert_scheduler_idle(&metrics);
assert!(
metrics.kv_stats.gpu_prefix_cache_hit_rate > 0.8,
"Expected cache hit rate > 0.8, got {}",
......@@ -983,17 +920,6 @@ mod tests {
let metrics_rx = scheduler.metrics_receiver();
let metrics = metrics_rx.borrow().clone();
assert_eq!(
metrics.kv_stats.gpu_cache_usage_perc,
0.0,
"Expected GPU cache usage to be 0%, got {}%",
metrics.kv_stats.gpu_cache_usage_perc * 100.0
);
assert_eq!(
metrics.kv_stats.kv_active_blocks, 0,
"Expected 0 active blocks, got {}",
metrics.kv_stats.kv_active_blocks
);
assert_scheduler_idle(&metrics);
}
}
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment