majiayu000 GPU Scheduling
VRAM-aware GPU scheduling for ML model orchestration.
Model Types and VRAM Configuration
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum ModelType {
Whisper = 0, // 4.5 GB
VideoMAE = 1, // 5.0 GB
CLAP = 2, // 2.0 GB
Qwen3 = 3, // 1.5 GB
DINOv3 = 4, // 1.0 GB
}
#[derive(Debug, Clone)]
pub struct VramConfig {
pub total_vram_gb: f64, // e.g., 16.0
pub safety_factor: f64, // e.g., 0.9 (use 90%)
pub whisper_vram_gb: f64,
pub videomae_vram_gb: f64,
pub clap_vram_gb: f64,
pub qwen3_vram_gb: f64,
pub dinov3_vram_gb: f64,
}
impl VramConfig {
pub fn usable_vram_gb(&self) -> f64 {
self.total_vram_gb * self.safety_factor
}
pub fn model_vram(&self, model: ModelType) -> f64 {
match model {
ModelType::Whisper => self.whisper_vram_gb,
ModelType::VideoMAE => self.videomae_vram_gb,
ModelType::CLAP => self.clap_vram_gb,
ModelType::Qwen3 => self.qwen3_vram_gb,
ModelType::DINOv3 => self.dinov3_vram_gb,
}
}
pub fn can_fit(&self, model: ModelType) -> bool {
self.model_vram(model) <= self.usable_vram_gb()
}
}
impl Default for VramConfig {
fn default() -> Self {
Self {
total_vram_gb: 16.0,
safety_factor: 0.9,
whisper_vram_gb: 4.5,
videomae_vram_gb: 5.0,
clap_vram_gb: 2.0,
qwen3_vram_gb: 1.5,
dinov3_vram_gb: 1.0,
}
}
}
Work Item with Dependencies
use std::collections::HashSet;
use std::time::{SystemTime, UNIX_EPOCH};
#[derive(Debug, Clone)]
pub struct WorkItem {
pub id: String,
pub video_id: String,
pub model_type: ModelType,
pub input_data: String,
pub priority: u32,
pub dependencies: Vec<String>,
pub submitted_at: f64,
pub resolved_deps: HashSet<String>,
}
impl WorkItem {
pub fn new(
video_id: String,
model_type: ModelType,
input_data: String,
dependencies: Option<Vec<String>>,
priority: Option<u32>,
) -> Self {
let id = format!("{}_{:?}_{}", video_id, model_type, uuid::Uuid::new_v4());
Self {
id,
video_id,
model_type,
input_data,
priority: priority.unwrap_or(0),
dependencies: dependencies.unwrap_or_default(),
submitted_at: SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_secs_f64())
.unwrap_or(0.0),
resolved_deps: HashSet::new(),
}
}
pub fn is_ready(&self) -> bool {
self.dependencies.iter()
.all(|dep| self.resolved_deps.contains(dep))
}
pub fn resolve_dependency(&mut self, dep_id: &str) {
self.resolved_deps.insert(dep_id.to_string());
}
pub fn waiting_time(&self) -> f64 {
let now = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_secs_f64())
.unwrap_or(0.0);
now - self.submitted_at
}
}
Work Queue with Priority
use std::collections::{HashMap, VecDeque};
pub struct WorkQueue {
model_type: ModelType,
ready_items: VecDeque<WorkItem>,
waiting_items: Vec<WorkItem>,
}
impl WorkQueue {
pub fn new(model_type: ModelType) -> Self {
Self {
model_type,
ready_items: VecDeque::new(),
waiting_items: Vec::new(),
}
}
pub fn push(&mut self, item: WorkItem) {
if item.is_ready() {
// Insert by priority (higher priority first)
let pos = self.ready_items
.iter()
.position(|i| i.priority < item.priority)
.unwrap_or(self.ready_items.len());
self.ready_items.insert(pos, item);
} else {
self.waiting_items.push(item);
}
}
pub fn pop(&mut self) -> Option<WorkItem> {
self.ready_items.pop_front()
}
pub fn resolve_dependency(&mut self, dep_id: &str) {
let mut newly_ready = Vec::new();
let mut i = 0;
while i < self.waiting_items.len() {
self.waiting_items[i].resolve_dependency(dep_id);
if self.waiting_items[i].is_ready() {
let item = self.waiting_items.swap_remove(i);
newly_ready.push(item);
} else {
i += 1;
}
}
for item in newly_ready {
self.push(item);
}
}
pub fn ready_count(&self) -> usize {
self.ready_items.len()
}
pub fn waiting_count(&self) -> usize {
self.waiting_items.len()
}
}
Model State Manager
use std::time::{Duration, Instant};
pub struct ModelStateManager {
current_model: Option<ModelType>,
loaded_at: Option<Instant>,
load_counts: [u32; 5],
last_load_times: [Option<Duration>; 5],
}
impl ModelStateManager {
pub fn new() -> Self {
Self {
current_model: None,
loaded_at: None,
load_counts: [0; 5],
last_load_times: [None; 5],
}
}
pub fn mark_loaded(&mut self, model: ModelType, load_time: Duration) {
self.current_model = Some(model);
self.loaded_at = Some(Instant::now());
self.load_counts[model as usize] += 1;
self.last_load_times[model as usize] = Some(load_time);
}
pub fn mark_unloaded(&mut self) {
self.current_model = None;
self.loaded_at = None;
}
pub fn current_model(&self) -> Option<ModelType> {
self.current_model
}
pub fn estimated_load_time(&self, model: ModelType) -> Duration {
self.last_load_times[model as usize].unwrap_or_else(|| {
// Default estimates
match model {
ModelType::Whisper => Duration::from_secs(5),
ModelType::VideoMAE => Duration::from_secs(8),
ModelType::CLAP => Duration::from_secs(4),
ModelType::Qwen3 => Duration::from_secs(10),
ModelType::DINOv3 => Duration::from_secs(6),
}
})
}
pub fn needs_switch(&self, target: ModelType) -> bool {
self.current_model != Some(target)
}
}
GPU Scheduler
use parking_lot::RwLock;
use std::sync::Arc;
pub struct GPUScheduler {
queues: Arc<RwLock<HashMap<ModelType, WorkQueue>>>,
state: Arc<RwLock<ModelStateManager>>,
vram_config: VramConfig,
}
impl GPUScheduler {
pub fn new(vram_config: VramConfig) -> Self {
let mut queues = HashMap::new();
for model in [
ModelType::Whisper,
ModelType::VideoMAE,
ModelType::CLAP,
ModelType::Qwen3,
ModelType::DINOv3,
] {
queues.insert(model, WorkQueue::new(model));
}
Self {
queues: Arc::new(RwLock::new(queues)),
state: Arc::new(RwLock::new(ModelStateManager::new())),
vram_config,
}
}
pub fn submit(&self, item: WorkItem) -> String {
let id = item.id.clone();
let model = item.model_type;
let mut queues = self.queues.write();
if let Some(queue) = queues.get_mut(&model) {
queue.push(item);
}
id
}
pub fn resolve_dependency(&self, dep_id: &str) {
let mut queues = self.queues.write();
for queue in queues.values_mut() {
queue.resolve_dependency(dep_id);
}
}
pub fn get_next_batch(&self, max_batch_size: usize) -> Option<(ModelType, Vec<WorkItem>)> {
let state = self.state.read();
let mut queues = self.queues.write();
// Prefer current model to minimize switches
if let Some(current) = state.current_model() {
if let Some(queue) = queues.get_mut(¤t) {
if queue.ready_count() > 0 {
let mut batch = Vec::new();
while batch.len() < max_batch_size {
if let Some(item) = queue.pop() {
batch.push(item);
} else {
break;
}
}
if !batch.is_empty() {
return Some((current, batch));
}
}
}
}
// Find queue with most ready items
let best_model = queues
.iter()
.filter(|(_, q)| q.ready_count() > 0)
.max_by_key(|(_, q)| q.ready_count())
.map(|(m, _)| *m);
if let Some(model) = best_model {
let queue = queues.get_mut(&model).unwrap();
let mut batch = Vec::new();
while batch.len() < max_batch_size {
if let Some(item) = queue.pop() {
batch.push(item);
} else {
break;
}
}
if !batch.is_empty() {
return Some((model, batch));
}
}
None
}
pub fn total_pending(&self) -> usize {
self.queues.read()
.values()
.map(|q| q.ready_count() + q.waiting_count())
.sum()
}
}
Semaphore-Based VRAM Limiting
use tokio::sync::Semaphore;
// Alternative: Use semaphore for simpler VRAM tracking
pub struct VramSemaphore {
semaphore: Semaphore,
units_per_gb: u32,
}
impl VramSemaphore {
pub fn new(total_gb: f64, units_per_gb: u32) -> Self {
let total_units = (total_gb * units_per_gb as f64) as usize;
Self {
semaphore: Semaphore::new(total_units),
units_per_gb,
}
}
pub async fn acquire(&self, vram_gb: f64) -> Result<SemaphorePermit<'_>> {
let units = (vram_gb * self.units_per_gb as f64) as u32;
self.semaphore.acquire_many(units).await
.map_err(|_| Error::VramAcquisition)
}
}
// Usage
let vram = VramSemaphore::new(16.0, 10); // 160 units total
async fn run_whisper(vram: &VramSemaphore) -> Result<()> {
let _permit = vram.acquire(4.5).await?; // 45 units
// Run model - permit released on drop
Ok(())
}
Guidelines
- Track VRAM usage per model type
- Use safety factor (90%) to avoid OOM
- Minimize model switches (prefer current model)
- Support work item dependencies
- Use priority queues for important work
- Track model load times for scheduling decisions
- Use semaphores for simpler VRAM limiting
Examples
See hercules-local-algo/src/scheduler/ for complete implementation.