fix(spurd): refresh node inventory and re-register when it changes

crates/spur-sched/src/cons_tres.rs

@@ -66,6 +66,29 @@ impl NodeAllocation {
         }
     }
 
+    /// Update total capacity from a re-discovered inventory, preserving current
+    /// allocations. CPUs/memory resize in place; GPUs are re-indexed by device
+    /// id so existing GPU allocations follow their device. A device that
+    /// vanished while allocated stops counting toward the total but its owning
+    /// job is untouched (it releases normally). Without this the agent's local
+    /// capacity stays frozen at startup and rejects launches for GPUs that only
+    /// appeared later (e.g. an SPX->CPX partition change), even after the
+    /// controller has learned about them via re-register.
+    pub fn update_capacity(&mut self, resources: &ResourceSet) {
+        self.total_cpus = resources.cpus;
+        self.allocated_cpus.resize(resources.cpus as usize, false);
+        self.total_memory_mb = resources.memory_mb;
+
+        let allocated_ids: std::collections::HashSet<u32> =
+            self.allocated_gpu_ids().into_iter().collect();
+        self.gpus = resources.gpus.clone();
+        self.gpu_allocated = self
+            .gpus
+            .iter()
+            .map(|g| allocated_ids.contains(&g.device_id))
+            .collect();
+    }
+
     /// Available (unallocated) CPU count.
     pub fn free_cpus(&self) -> u32 {
         self.allocated_cpus.iter().filter(|&&a| !a).count() as u32
@@ -326,6 +349,43 @@ mod tests {
         assert_eq!(node.free_gpus(None), 8);
     }
 
+    #[test]
+    fn test_update_capacity_grows_and_preserves_allocations() {
+        // Start with 4 GPUs, allocate two of them, then an inventory refresh
+        // reveals 8 (e.g. SPX->CPX). The new GPUs become schedulable while the
+        // two already allocated stay allocated (by device id).
+        let mut node = make_node(32, 128_000, 4, "mi300x");
+        node.allocate_for_job(1, 8, 0, &[1, 2]).unwrap();
+        assert_eq!(node.free_gpus(None), 2);
+
+        let bigger = ResourceSet {
+            cpus: 64,
+            memory_mb: 256_000,
+            gpus: (0..8u32)
+                .map(|device_id| GpuResource {
+                    device_id,
+                    gpu_type: "mi300x".into(),
+                    memory_mb: 192_000,
+                    peer_gpus: vec![],
+                    link_type: GpuLinkType::XGMI,
+                })
+                .collect(),
+            ..Default::default()
+        };
+        node.update_capacity(&bigger);
+
+        assert_eq!(node.total_cpus, 64);
+        assert_eq!(node.free_memory_mb(), 256_000 - node.allocated_memory_mb);
+        assert_eq!(
+            node.free_gpus(None),
+            6,
+            "4 new GPUs added, 2 still allocated"
+        );
+        assert_eq!(node.allocated_gpu_ids(), vec![1, 2]);
+        // CPU allocation from before is preserved (8 cores still busy).
+        assert_eq!(node.free_cpus(), 64 - 8);
+    }
+
     #[test]
     fn test_allocate_cpus() {
         let mut node = make_node(64, 256_000, 0, "");

crates/spurd/src/agent_server.rs

@@ -654,7 +654,7 @@ impl AgentService {
             hostname::get()
                 .map(|h| h.to_string_lossy().to_string())
                 .unwrap_or_else(|_| "unknown".into()),
-            &reporter.resources,
+            &reporter.resources.read().unwrap(),
         );
 
         // Load SPANK plugins from plugstack.conf if available
@@ -731,6 +731,14 @@ impl AgentService {
         self.k0s.clone()
     }
 
+    /// Shared handle to this node's local allocation, so the inventory-refresh
+    /// task can update its capacity (via `update_capacity`) when devices appear
+    /// or vanish. Without it the agent's capacity stays frozen at startup and
+    /// rejects launches for devices the controller has already re-learned.
+    pub fn allocation_handle(&self) -> Arc<Mutex<NodeAllocation>> {
+        self.allocation.clone()
+    }
+
     /// Spawn a background task to monitor running jobs and report completions.
     pub fn start_monitor(&self, controller_addr: String) {
         let running = self.running.clone();
@@ -2124,9 +2132,9 @@ impl SlurmAgent for AgentService {
         &self,
         _request: Request<()>,
     ) -> Result<Response<NodeResourcesResponse>, Status> {
-        let resources = &self.reporter.resources;
+        let resources = self.reporter.resources.read().unwrap();
         Ok(Response::new(NodeResourcesResponse {
-            total: Some(crate::reporter::resource_to_proto(resources)),
+            total: Some(crate::reporter::resource_to_proto(&resources)),
             used: Some(crate::reporter::allocations_to_proto(
                 &spur_core::resource::ResourceAllocations::default(),
             )),

crates/spurd/src/main.rs

@@ -422,6 +422,51 @@ async fn main() -> anyhow::Result<()> {
 
     agent_service.start_monitor(args.controller.clone());
 
+    // Periodically re-discover node inventory. It is otherwise frozen at startup,
+    // so a device count that changes out of band (a GPU partition-mode switch, a
+    // GPU dropping off the bus) never reaches the controller and it keeps
+    // scheduling against hardware that no longer exists. On a change: swap the
+    // shared registry (so injection uses the new set), resize the local
+    // allocation capacity (so this node accepts launches for devices that
+    // appeared), and re-register with the controller.
+    {
+        let reporter = reporter.clone();
+        let registry = registry.clone();
+        let config = config.clone();
+        let allocation = agent_service.allocation_handle();
+        tokio::spawn(async move {
+            let mut ticker = tokio::time::interval(std::time::Duration::from_secs(60));
+            ticker.tick().await; // the first tick fires immediately; skip it
+                                 // `update_resources` commits the new inventory to the reporter as it
+                                 // detects the change, so a re-register that then fails would never be
+                                 // retried (the next tick sees no further change). Track that a
+                                 // re-register is still owed and keep trying until it lands.
+            let mut reregister_pending = false;
+            loop {
+                ticker.tick().await;
+                let rebuilt = init_device_registry(config.as_ref());
+                let fresh = reporter::discover_resources(&rebuilt);
+                let changed = reporter.update_resources(fresh.clone());
+                if changed {
+                    *registry.lock().await = rebuilt;
+                    allocation.lock().await.update_capacity(&fresh);
+                    reregister_pending = true;
+                }
+                if reregister_pending {
+                    match reporter.register().await {
+                        Ok(()) => {
+                            reregister_pending = false;
+                            info!("node inventory changed; re-registered with the controller")
+                        }
+                        Err(e) => {
+                            warn!(error = %e, "re-register after inventory change failed; will retry")
+                        }
+                    }
+                }
+            }
+        });
+    }
+
     let addr = args.listen.parse()?;
     info!(%addr, "agent gRPC server listening");
 

crates/spurd/src/reporter.rs

@@ -31,7 +31,10 @@ impl<T: Send> HeldJobs for Mutex<HashMap<u32, T>> {
 pub struct NodeReporter {
     pub hostname: String,
     pub controller_addr: String,
-    pub resources: ResourceSet,
+    /// Node inventory. Behind a lock so a background refresh can update it live
+    /// (a device count that changes after startup, e.g. a GPU partition-mode
+    /// switch, must reach the controller rather than being frozen at boot).
+    pub resources: std::sync::RwLock<ResourceSet>,
     pub node_address: spur_net::NodeAddress,
     pub labels: HashMap<String, String>,
     pub free_memory_mb: AtomicU64,
@@ -64,7 +67,7 @@ impl NodeReporter {
         Self {
             hostname,
             controller_addr,
-            resources,
+            resources: std::sync::RwLock::new(resources),
             node_address,
             labels,
             free_memory_mb: AtomicU64::new(0),
@@ -94,17 +97,32 @@ impl NodeReporter {
         self.held_jobs.held_job_ids()
     }
 
+    /// Replace the reported inventory if it changed; returns whether it did, so
+    /// the caller can re-register the node with the controller.
+    pub fn update_resources(&self, fresh: ResourceSet) -> bool {
+        let mut cur = self.resources.write().unwrap();
+        if *cur != fresh {
+            *cur = fresh;
+            true
+        } else {
+            false
+        }
+    }
+
     /// Register with the controller.
     pub async fn register(&self) -> anyhow::Result<()> {
         let channel = spur_client::connect_channel(&self.controller_addr)
             .await
             .context("failed to connect to spurctld for registration")?;
         let mut client = spur_proto::controller_client(channel);
 
+        // Snapshot the inventory before the await so the lock guard is not held
+        // across it (the future must stay Send).
+        let resources = resource_to_proto(&self.resources.read().unwrap());
         let resp = client
             .register_agent(RegisterAgentRequest {
                 hostname: self.hostname.clone(),
-                resources: Some(resource_to_proto(&self.resources)),
+                resources: Some(resources),
                 version: env!("CARGO_PKG_VERSION").into(),
                 address: self.node_address.ip.clone(),
                 port: self.node_address.port as u32,
@@ -576,6 +594,44 @@ mod tests {
         assert_eq!(resources.generic.get("bandwidth:lustre"), Some(&4096));
     }
 
+    #[test]
+    fn update_resources_reports_only_real_changes() {
+        let held: Arc<dyn HeldJobs> = Arc::new(Mutex::new(HashMap::<u32, ()>::new()));
+        let reporter = NodeReporter::new(
+            "host".into(),
+            "http://controller".into(),
+            ResourceSet {
+                cpus: 8,
+                memory_mb: 1000,
+                ..Default::default()
+            },
+            spur_net::address::NodeAddress {
+                ip: "127.0.0.1".into(),
+                hostname: "host".into(),
+                port: 6818,
+                source: spur_net::address::AddressSource::Static,
+            },
+            HashMap::new(),
+            String::new(),
+            "spur0".into(),
+            held,
+        );
+
+        // Re-discovering the same inventory is not a change.
+        assert!(!reporter.update_resources(ResourceSet {
+            cpus: 8,
+            memory_mb: 1000,
+            ..Default::default()
+        }));
+        // A changed device/cpu count is, and becomes the new reported inventory.
+        assert!(reporter.update_resources(ResourceSet {
+            cpus: 16,
+            memory_mb: 1000,
+            ..Default::default()
+        }));
+        assert_eq!(reporter.resources.read().unwrap().cpus, 16);
+    }
+
     #[test]
     fn should_reregister_on_not_found() {
         assert!(should_reregister(&tonic::Status::not_found(