feat(scheduler): add GPU-aware DRF scheduling with resource requirements

.gitignore

@@ -35,3 +35,6 @@ dist/
 tmp/
 session_manager/events/
 certs/
+
+# AI agent working directories
+.sisyphus/

common/Cargo.toml

@@ -32,6 +32,7 @@ serde_yaml = { workspace = true }
 rustix = { version = "1.1" , features = ["system"] }
 num_cpus = "1.17"
 bytesize = "1.3"
+nvml-wrapper = "0.10"
 
 axum = { workspace = true }
 pprof = { workspace = true }

common/src/apis/from_rpc.rs

@@ -23,6 +23,7 @@ impl From<rpc::ResourceRequirement> for ResourceRequirement {
         Self {
             cpu: req.cpu,
             memory: req.memory,
+            gpu: req.gpu,
         }
     }
 }

common/src/apis/to_rpc.rs

@@ -20,7 +20,7 @@ impl From<ResourceRequirement> for rpc::ResourceRequirement {
         Self {
             cpu: req.cpu,
             memory: req.memory,
-            gpu: 0,
+            gpu: req.gpu,
         }
     }
 }
@@ -202,6 +202,7 @@ impl From<&Session> for rpc::Session {
                 max_instances: ssn.max_instances,
                 batch_size: ssn.batch_size,
                 priority: ssn.priority,
+                resreq: None,
             }),
             status: Some(status),
         }

common/src/apis/types.rs

@@ -180,6 +180,7 @@ pub struct SessionAttributes {
     pub max_instances: Option<u32>,
     pub batch_size: u32,
     pub priority: u32,
+    pub resreq: Option<ResourceRequirement>,
 }
 
 impl Default for SessionAttributes {
@@ -193,6 +194,7 @@ impl Default for SessionAttributes {
             max_instances: None,
             batch_size: 1,
             priority: 0,
+            resreq: None,
         }
     }
 }
@@ -226,6 +228,7 @@ pub struct Session {
     pub max_instances: Option<u32>,
     pub batch_size: u32,
     pub priority: u32,
+    pub resreq: Option<ResourceRequirement>,
 }
 
 #[derive(Clone, Copy, Debug, Default, Eq, PartialEq, Hash, strum_macros::Display)]
@@ -351,6 +354,7 @@ pub struct NodeInfo {
 pub struct ResourceRequirement {
     pub cpu: u64,
     pub memory: u64,
+    pub gpu: i32,
 }
 
 #[derive(Clone, Debug, Default)]
@@ -382,6 +386,32 @@ fn totalram() -> u64 {
     system::sysinfo().totalram
 }
 
+fn detect_gpus() -> i32 {
+    if let Ok(devices) = std::env::var("CUDA_VISIBLE_DEVICES") {
+        if devices.is_empty() || devices == "-1" {
+            return 0;
+        }
+        return devices.split(',').count() as i32;
+    }
+
+    match nvml_wrapper::Nvml::init() {
+        Ok(nvml) => match nvml.device_count() {
+            Ok(count) => {
+                tracing::info!("Detected {} GPU(s) via NVML", count);
+                count as i32
+            }
+            Err(e) => {
+                tracing::warn!("NVML initialized but failed to get device count: {}", e);
+                0
+            }
+        },
+        Err(e) => {
+            tracing::debug!("NVML not available, GPU detection skipped: {}", e);
+            0
+        }
+    }
+}
+
 impl Node {
     pub fn new() -> Self {
         let name = uname();
@@ -397,7 +427,8 @@ impl Node {
     pub fn refresh(&mut self) {
         let memory = totalram();
         let cpu = num_cpus::get() as u64;
-        let capacity = ResourceRequirement { cpu, memory };
+        let gpu = detect_gpus();
+        let capacity = ResourceRequirement { cpu, memory, gpu };
         let allocatable = capacity.clone();
         let info = NodeInfo {
             arch: env::consts::ARCH.to_string(),
@@ -420,6 +451,7 @@ impl From<&String> for ResourceRequirement {
         let parts = s.split(',');
         let mut cpu = 0;
         let mut memory = 0;
+        let mut gpu = 0;
         for p in parts {
             let mut parts = p.split('=').map(|s| s.trim());
             let key = parts.next();
@@ -428,12 +460,13 @@ impl From<&String> for ResourceRequirement {
                 (Some("cpu"), Some(value)) => cpu = value.parse::<u64>().unwrap_or(0),
                 (Some("memory"), Some(value)) => memory = Self::parse_memory(value),
                 (Some("mem"), Some(value)) => memory = Self::parse_memory(value),
+                (Some("gpu"), Some(value)) => gpu = value.parse::<i32>().unwrap_or(0),
                 _ => {
                     tracing::error!("Invalid resource requirement: {s}");
                 }
             }
         }
-        Self { cpu, memory }
+        Self { cpu, memory, gpu }
     }
 }
 
@@ -442,11 +475,26 @@ impl ResourceRequirement {
         Self {
             cpu: slots as u64 * unit.cpu,
             memory: slots as u64 * unit.memory,
+            gpu: slots as i32 * unit.gpu,
         }
     }
 
     pub fn to_slots(&self, unit: &ResourceRequirement) -> u32 {
-        (self.cpu / unit.cpu).min(self.memory / unit.memory) as u32
+        let cpu_slots = self.cpu.checked_div(unit.cpu).unwrap_or(u64::MAX);
+
+        let mem_slots = self.memory.checked_div(unit.memory).unwrap_or(u64::MAX);
+
+        let gpu_slots = if unit.gpu > 0 {
+            if self.gpu < 0 {
+                0
+            } else {
+                (self.gpu as u64) / (unit.gpu as u64)
+            }
+        } else {
+            u64::MAX
+        };
+
+        cpu_slots.min(mem_slots).min(gpu_slots) as u32
     }
 
     pub(crate) fn parse_memory(s: &str) -> u64 {

docker/Dockerfile.console

@@ -1,4 +1,4 @@
-FROM rust:1.94 AS builder
+FROM rust:1.95 AS builder
 
 WORKDIR /usr/src/flame
 COPY . .

docker/Dockerfile.fem

@@ -1,4 +1,4 @@
-FROM rust:1.94 AS builder
+FROM rust:1.95 AS builder
 
 WORKDIR /usr/src/flame
 COPY . .

docker/Dockerfile.foc

@@ -1,4 +1,4 @@
-FROM rust:1.94 AS builder
+FROM rust:1.95 AS builder
 
 WORKDIR /usr/src/flame
 COPY . .

docker/Dockerfile.fsm

@@ -1,4 +1,4 @@
-FROM rust:1.94 AS builder
+FROM rust:1.95 AS builder
 
 WORKDIR /usr/src/flame
 COPY . .