[rust/rqd] Add OOM prevention logic with kill frame selection (#2064)

Implement a new OOM (Out-Of-Memory) prevention module that proactively
kills frames when memory usage exceeds configured thresholds, preventing
system-wide OOM killer from terminating the RQD process.

When memory usage exceeds `memory_oom_margin_percentage` (default 96%),
the system calculates a target memory level (5% below the threshold) and
selects frames to kill to reach that target safely.

Frame Selection Algorithm -------------------------

Frames are sorted using a multi-criteria scoring system that balances
three factors:

1. **Memory Impact (weight: 10x)** 
- Measures absolute memory savings: consumed_memory /
total_memory_consumed
- Prioritizes frames that will free the most system memory - Normalized
to [0,1] range by dividing by the maximum memory impact

3. **Overboard Rate (weight: 7x)** 
- Measures relative excess: (consumed - limit) / limit 
- Targets frames most aggressively exceeding their soft limits 
- A frame using 10GB with a 1GB limit (900% over) has higher overboard
rate than one using 60GB with a 50GB limit (20% over) - Normalized to
[0,1] range by dividing by the maximum overboard rate

5. **Duration Rate (weight: 12x - highest)** 
- Prefers killing more recent frames: (max_duration - frame_duration) /
max_duration
- Minimizes wasted compute by preserving long-running frames 
- Inverted scale: newer frames score higher - Normalized to [0,1] range

Each frame receives a composite score calculated as: score =
(memory_impact × 10) + (overboard_rate × 7) + (duration_rate × 12)

All metrics are normalized before weighting to ensure each factor
contributes meaningfully regardless of absolute values. This prevents a
frame with very high memory usage from completely dominating the score,
allowing the algorithm to balance all three criteria.

The algorithm sorts frames by descending score and kills them
iteratively until enough memory is freed to reach the target level. This
conservative approach avoids unnecessary termination while effectively
preventing OOM.

The weighting scheme reflects production priorities: 
- Duration (12x): Preserve investment in long-running frames 
- Memory Impact (10x): Maximize immediate memory relief 
- Overboard Rate (7x): Discourage limit violations

Implementation includes comprehensive test coverage validating: 
- Threshold triggering behavior 
- Target memory calculation 
- Single and multi-frame selection 
- Each scoring criterion independently
- Normalized scoring with edge cases 
- Stop-when-sufficient-memory-freed logic

Author: DiegoTavares

rust/crates/rqd/src/config/mod.rs

@@ -89,6 +89,7 @@ pub struct MachineConfig {
     #[serde(with = "humantime_serde")]
     pub nimby_start_retry_interval: Duration,
     pub nimby_display_xauthority_path: String,
+    pub memory_oom_margin_percentage: u32,
 }
 
 impl Default for MachineConfig {
@@ -111,6 +112,7 @@ impl Default for MachineConfig {
             nimby_display_file_path: None,
             nimby_start_retry_interval: Duration::from_secs(60 * 5), // 5 min
             nimby_display_xauthority_path: "/home/{username}/Xauthority".to_string(),
+            memory_oom_margin_percentage: 96,
         }
     }
 }

rust/crates/rqd/src/frame/running_frame.rs

@@ -24,6 +24,7 @@ use tokio::io::AsyncReadExt;
 use tokio::{io::AsyncBufReadExt, task::JoinHandle};
 use tracing::{error, info, trace, warn};
 
+use crate::system::OOM_REASON_MSG;
 use crate::{frame::frame_cmd::FrameCmdBuilder, system::manager::ProcessStats};
 
 use serde::{Deserialize, Serialize};
@@ -179,13 +180,69 @@ impl RunningFrame {
         }
     }
 
+    #[cfg(test)]
+    pub fn init_started_for_test(
+        request: RunFrame,
+        uid: u32,
+        config: RunnerConfig,
+        cpu_list: Option<Vec<u32>>,
+        gpu_list: Option<Vec<u32>>,
+        hostname: String,
+        duration: Duration,
+    ) -> Self {
+        let instance = Self::init(request, uid, config, cpu_list, gpu_list, hostname);
+
+        {
+            let mut state = instance
+                .state
+                .write()
+                .unwrap_or_else(|err| err.into_inner());
+
+            match &mut *state {
+                FrameState::Created(created_state) => {
+                    *state = FrameState::Running(RunningState {
+                        pid: 999, // Dummy pid
+                        start_time: SystemTime::now()
+                            .checked_sub(duration)
+                            .unwrap_or(SystemTime::now()),
+                        launch_thread_handle: created_state.launch_thread_handle.take(),
+                        kill_reason: None,
+                    });
+                }
+                FrameState::Running(running_state) => warn!(
+                    "Invalid State. Frame {} has already started {:?}",
+                    instance, running_state
+                ),
+                FrameState::Finished(_) => {
+                    warn!("Invalid States. Frame {} has already finished", instance)
+                }
+                FrameState::FailedBeforeStart => {
+                    warn!("Invalid States. Frame {} failed before starting", instance)
+                }
+            }
+        } // state is dropped here
+
+        instance
+    }
+
     pub fn update_frame_stats(&self, proc_stats: ProcessStats) {
         self.frame_stats
             .write()
             .unwrap_or_else(|poisoned| poisoned.into_inner())
             .update(proc_stats);
     }
 
+    pub fn get_duration(&self) -> Duration {
+        let state = self.state.read().unwrap_or_else(|err| err.into_inner());
+
+        match *state {
+            FrameState::Created(_) => Duration::ZERO,
+            FrameState::Running(ref r) => r.start_time.elapsed().unwrap_or(Duration::ZERO),
+            FrameState::Finished(ref r) => r.start_time.elapsed().unwrap_or(Duration::ZERO),
+            FrameState::FailedBeforeStart => Duration::ZERO,
+        }
+    }
+
     pub fn get_state_copy(&self) -> FrameState {
         let state = self.state.read().unwrap_or_else(|err| err.into_inner());
 
@@ -250,13 +307,23 @@ impl RunningFrame {
         match &mut *state {
             FrameState::Created(_) => Err(miette!("Invalid State. Frame {} hasn't started", self)),
             FrameState::Running(running_state) => {
+                // Replace exit_signal to memory signal if kill_reason matches the memory check message
+                let modified_exit_signal = match &running_state.kill_reason {
+                    Some(reason) if reason.contains(OOM_REASON_MSG) => {
+                        // 33 is the error signal hardcoded on Cuebot for memory issues
+                        // (See Dispatcher.java:EXIT_STATUS_MEMORY_FAILURE)
+                        Some(33)
+                    }
+                    _ => exit_signal,
+                };
+
                 // Create a new FinishedState with the current running state values
                 let finished_state = FinishedState {
                     pid: running_state.pid,
                     start_time: running_state.start_time,
                     end_time: SystemTime::now(),
                     exit_code,
-                    exit_signal,
+                    exit_signal: modified_exit_signal,
                     kill_reason: running_state.kill_reason.clone(),
                 };
 

rust/crates/rqd/src/system/machine.rs

@@ -1,6 +1,11 @@
 use std::sync::Arc;
 
-use crate::{config::CONFIG, frame::manager, report::report_client};
+use crate::{
+    config::CONFIG,
+    frame::manager,
+    report::report_client,
+    system::oom::{self, OOM_REASON_MSG},
+};
 use async_trait::async_trait;
 use bytesize::KIB;
 use itertools::Either;
@@ -61,7 +66,7 @@ pub struct MachineMonitor {
     pub system_manager: Mutex<SystemManagerType>,
     pub core_manager: Arc<RwLock<CoreStateManager>>,
     pub running_frames_cache: Arc<RunningFrameCache>,
-    last_host_state: Arc<Mutex<Option<RenderHost>>>,
+    last_host_state: Arc<RwLock<Option<RenderHost>>>,
     interrupt: Mutex<Option<broadcast::Sender<()>>>,
     reboot_when_idle: Mutex<bool>,
     #[cfg(feature = "nimby")]
@@ -136,7 +141,7 @@ impl MachineMonitor {
             report_client,
             system_manager: Mutex::new(system_manager),
             running_frames_cache: RunningFrameCache::init(),
-            last_host_state: Arc::new(Mutex::new(None)),
+            last_host_state: Arc::new(RwLock::new(None)),
             interrupt: Mutex::new(None),
             reboot_when_idle: Mutex::new(false),
             #[cfg(feature = "nimby")]
@@ -162,7 +167,7 @@ impl MachineMonitor {
         };
 
         self.last_host_state
-            .lock()
+            .write()
             .await
             .replace(host_state.clone());
 
@@ -203,40 +208,7 @@ impl MachineMonitor {
 
                         #[cfg(feature = "nimby")]
                         if let Some(nimby) = &*self.nimby {
-                            match (nimby.is_user_active(), last_lock_state) {
-                                // Became locked
-                                (true, LockState::Open) => {
-                                    last_lock_state = LockState::NimbyLocked;
-
-                                    // Update registered state
-                                    let mut nimby_state = self.nimby_state.write().await;
-                                    *nimby_state = last_lock_state;
-                                    drop(nimby_state);
-
-                                    info!("Host became nimby locked");
-                                    self.lock_all_cores().await;
-                                    let count = manager::instance().await?.kill_all_running_frames("Host has been Nimby-Locked").await?;
-                                    info!("{count} frames killed after the machine became locked")
-                                }
-                                // Continues locked
-                                (true, LockState::NimbyLocked) => {}
-                                // Continues open
-                                (false, LockState::Open) => {}
-                                // Became unlocked
-                                (false, LockState::NimbyLocked) => {
-                                    last_lock_state = LockState::Open;
-
-                                    // Update registered state
-                                    let mut nimby_state = self.nimby_state.write().await;
-                                    *nimby_state = last_lock_state;
-                                    drop(nimby_state);
-
-                                    info!("Host became nimby unlocked");
-                                    self.unlock_all_cores().await;
-                                }
-                                // NoOp
-                                _ => ()
-                            }
+                            last_lock_state = self.handle_nimby_state_change(nimby, last_lock_state).await?;
                         }
                 }
 
@@ -245,6 +217,57 @@ impl MachineMonitor {
         Ok(())
     }
 
+    /// Handles NIMBY state changes and performs necessary actions when the host
+    /// becomes locked or unlocked based on user activity.
+    ///
+    /// Returns the new lock state after handling any transitions.
+    #[cfg(feature = "nimby")]
+    async fn handle_nimby_state_change(
+        &self,
+        nimby: &Nimby,
+        current_state: LockState,
+    ) -> Result<LockState> {
+        match (nimby.is_user_active(), current_state) {
+            // Became locked
+            (true, LockState::Open) => {
+                let new_state = LockState::NimbyLocked;
+
+                // Update registered state
+                let mut nimby_state = self.nimby_state.write().await;
+                *nimby_state = new_state;
+                drop(nimby_state);
+
+                info!("Host became nimby locked");
+                self.lock_all_cores().await;
+                let count = manager::instance()
+                    .await?
+                    .kill_all_running_frames("Host has been Nimby-Locked")
+                    .await?;
+                info!("{count} frames killed after the machine became locked");
+                Ok(new_state)
+            }
+            // Continues locked
+            (true, LockState::NimbyLocked) => Ok(current_state),
+            // Continues open
+            (false, LockState::Open) => Ok(current_state),
+            // Became unlocked
+            (false, LockState::NimbyLocked) => {
+                let new_state = LockState::Open;
+
+                // Update registered state
+                let mut nimby_state = self.nimby_state.write().await;
+                *nimby_state = new_state;
+                drop(nimby_state);
+
+                info!("Host became nimby unlocked");
+                self.unlock_all_cores().await;
+                Ok(new_state)
+            }
+            // NoOp
+            _ => Ok(current_state),
+        }
+    }
+
     #[cfg(feature = "nimby")]
     async fn start_nimby(&self, term_receiver: Receiver<()>) {
         // Start nimby monitor
@@ -306,6 +329,7 @@ impl MachineMonitor {
     async fn monitor_running_frames(&self) -> Result<()> {
         let mut finished_frames: Vec<Arc<RunningFrame>> = Vec::new();
         let mut running_frames: Vec<(Arc<RunningFrame>, RunningState)> = Vec::new();
+        let mut memory_aggressors: Vec<(Arc<RunningFrame>, u64)> = Vec::new();
 
         // Only keep running frames on the cache and store a copy of their state
         // to avoid having to deal with the state lock
@@ -340,6 +364,13 @@ impl MachineMonitor {
             };
 
             if let Some(proc_stats) = proc_stats_opt {
+                // Mark memory aggressor frames
+                if running_frame.request.soft_memory_limit > 0
+                    && proc_stats.rss as i64 > running_frame.request.soft_memory_limit
+                {
+                    memory_aggressors.push((Arc::clone(running_frame), proc_stats.max_rss));
+                }
+
                 // Update stats for running frames
                 running_frame.update_frame_stats(proc_stats);
             } else if running_frame.is_marked_for_cache_removal() {
@@ -361,6 +392,26 @@ impl MachineMonitor {
         }
 
         self.handle_finished_frames(finished_frames).await;
+        if let Some((memory_usage, total_memory)) = self.memory_usage().await {
+            let frames_to_kill =
+                oom::choose_frames_to_kill(memory_usage, total_memory, memory_aggressors);
+
+            // Attempt to kill selected frames.
+            // Logic will ignore kill errors and try again on the next iteration
+            for frame in frames_to_kill {
+                if let Ok(manager) = manager::instance().await {
+                    let kill_result = manager
+                        .kill_running_frame(&frame.frame_id, OOM_REASON_MSG.to_string())
+                        .await;
+                    if let Err(err) = kill_result {
+                        warn!(
+                            "Failed to kill frame {} when under OOM pressure. {}",
+                            frame, err
+                        )
+                    }
+                }
+            }
+        }
 
         // Sanitize dangling reservations
         // This mechanism is redundant as handle_finished_frames releases resources reserved to
@@ -387,7 +438,7 @@ impl MachineMonitor {
 
         // Avoid holding a lock while reporting back to cuebot
         let host_state_opt = {
-            let host_state_lock = self.last_host_state.lock().await;
+            let host_state_lock = self.last_host_state.read().await;
             host_state_lock.clone()
         };
 
@@ -487,6 +538,7 @@ impl MachineMonitor {
 #[async_trait]
 pub trait Machine {
     async fn hardware_state(&self) -> Option<HardwareState>;
+    async fn memory_usage(&self) -> Option<(u32, u64)>;
     async fn nimby_locked(&self) -> bool;
 
     /// Reserve CPU cores for a resource
@@ -591,15 +643,23 @@ pub trait Machine {
 impl Machine for MachineMonitor {
     async fn hardware_state(&self) -> Option<HardwareState> {
         self.last_host_state
-            .lock()
+            .read()
             .await
             .as_ref()
             .map(|hs| hs.state())
     }
 
+    async fn memory_usage(&self) -> Option<(u32, u64)> {
+        self.last_host_state.read().await.as_ref().map(|hs| {
+            let memory_percentage =
+                (((hs.total_mem - hs.free_mem) as f64 / hs.total_mem as f64) * 100.0) as u32;
+            (memory_percentage, hs.total_mem as u64)
+        })
+    }
+
     async fn nimby_locked(&self) -> bool {
         self.last_host_state
-            .lock()
+            .read()
             .await
             .as_ref()
             .map(|hs| hs.nimby_locked)
@@ -634,7 +694,7 @@ impl Machine for MachineMonitor {
     }
 
     async fn get_host_name(&self) -> String {
-        let lock = self.last_host_state.lock().await;
+        let lock = self.last_host_state.read().await;
 
         lock.as_ref()
             .map(|h| h.name.clone())
@@ -725,10 +785,11 @@ impl Machine for MachineMonitor {
         };
 
         // Store the last host_state on self
-        let mut self_host_state_lock = self.last_host_state.lock().await;
+        let mut self_host_state_lock = self.last_host_state.write().await;
         self_host_state_lock.replace(render_host.clone());
         drop(self_host_state_lock);
 
+        // Refresh list of running frames
         self.monitor_running_frames().await?;
 
         Ok(HostReport {

rust/crates/rqd/src/system/mod.rs

@@ -4,6 +4,7 @@ pub mod linux;
 pub mod machine;
 #[cfg(feature = "nimby")]
 pub mod nimby;
+mod oom;
 mod reservation;
 
 #[cfg(target_os = "macos")]
@@ -14,3 +15,5 @@ pub type ResourceId = Uuid;
 pub type CoreId = u32;
 pub type PhysId = u32;
 pub type ThreadId = u32;
+
+pub use oom::OOM_REASON_MSG;