[rust/rqd] Add more protections for accountability issues (#1768)

Fix core accountability issues caused by frames that failed before
having a chance to start not having their cores released

Author: DiegoTavares

rust/crates/opencue-proto/src/lib.rs

@@ -139,7 +139,10 @@ impl CoreDetail {
     ///
     /// * `Ok(())` if cores were reserved successfully
     /// * `Err(String)` if trying to reserve more cores than are available
-    pub fn reserve(&mut self, core_count_with_multiplier: usize) -> Result<(), String> {
+    pub fn register_reservation(
+        &mut self,
+        core_count_with_multiplier: usize,
+    ) -> Result<(), String> {
         if self.idle_cores - (core_count_with_multiplier as i32) < 0 {
             Err(format!(
                 "Tried to reserve {} out of {} cores available",
@@ -162,7 +165,7 @@ impl CoreDetail {
     ///
     /// * `Ok(())` if cores were released successfully
     /// * `Err(String)` if trying to release more cores than are currently reserved
-    pub fn release(&mut self, core_count_with_multiplier: u32) -> Result<(), String> {
+    pub fn register_release(&mut self, core_count_with_multiplier: u32) -> Result<(), String> {
         if self.booked_cores < core_count_with_multiplier as i32 {
             Err(format!(
                 "Tried to release {} out of {} cores reserved",

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

@@ -71,10 +71,10 @@ impl FrameManager {
             .environment
             .get("CUE_THREADABLE")
             .is_some_and(|v| v == "1");
-        let cpu_request = run_frame.num_cores as u32 / self.config.machine.core_multiplier;
+        let num_cores = (run_frame.num_cores as u32).div_ceil(self.config.machine.core_multiplier);
         let thread_ids = self
             .machine
-            .reserve_cores(cpu_request as usize, run_frame.resource_id(), hyperthreaded)
+            .reserve_cores(num_cores as usize, run_frame.resource_id(), hyperthreaded)
             .await
             .map_err(|err| {
                 FrameManagerError::Aborted(format!(
@@ -90,11 +90,7 @@ impl FrameManager {
                 let reserved_res = self.machine.reserve_gpus(run_frame.num_gpus as u32).await;
                 if reserved_res.is_err() {
                     // Release cores reserved on the last step
-                    if let Some(thread_ids) = &thread_ids {
-                        self.machine.release_threads(thread_ids).await;
-                    } else {
-                        self.machine.release_cores(cpu_request).await;
-                    }
+                    self.machine.release_cores(num_cores, &thread_ids).await;
                 }
                 Some(reserved_res.map_err(|err| {
                     FrameManagerError::Aborted(format!(
@@ -109,16 +105,18 @@ impl FrameManager {
             run_frame,
             uid,
             self.config.runner.clone(),
-            thread_ids,
+            thread_ids.clone(),
             gpu_list,
             self.machine.get_host_name().await,
         ));
 
         if self.config.runner.run_on_docker {
             self.spawn_docker_frame(running_frame, false);
-        } else {
-            self.spawn_running_frame(running_frame, false);
+        } else if self.spawn_running_frame(running_frame, false).is_err() {
+            // Release cores reserved on the last step
+            self.machine.release_cores(num_cores, &thread_ids).await;
         }
+
         Ok(())
     }
 
@@ -169,8 +167,8 @@ impl FrameManager {
                             .await
                             .map(Some),
                         None => {
-                            let num_cores = running_frame.request.num_cores as u32
-                                / self.config.machine.core_multiplier;
+                            let num_cores = (running_frame.request.num_cores as u32)
+                                .div_ceil(self.config.machine.core_multiplier);
                             self.machine
                                 .reserve_cores(
                                     num_cores as usize,
@@ -182,10 +180,14 @@ impl FrameManager {
                     } {
                         errors.push(err.to_string());
                     }
+
+                    let num_cores = (running_frame.request.num_cores as u32)
+                        .div_ceil(self.config.machine.core_multiplier);
+                    let thread_ids = &running_frame.thread_ids.clone();
                     if self.config.runner.run_on_docker {
                         todo!("Recovering frames when running on docker is not yet supported")
-                    } else {
-                        self.spawn_running_frame(running_frame, true)
+                    } else if self.spawn_running_frame(running_frame, true).is_err() {
+                        self.machine.release_cores(num_cores, thread_ids).await;
                     }
                 }
                 Err(err) => {
@@ -204,16 +206,24 @@ impl FrameManager {
         }
     }
 
-    fn spawn_running_frame(&self, running_frame: Arc<RunningFrame>, recover_mode: bool) {
+    fn spawn_running_frame(
+        &self,
+        running_frame: Arc<RunningFrame>,
+        recover_mode: bool,
+    ) -> Result<()> {
         self.machine.add_running_frame(Arc::clone(&running_frame));
         let running_frame_ref: Arc<RunningFrame> = Arc::clone(&running_frame);
         let thread_handle = tokio::spawn(async move { running_frame.run(recover_mode).await });
-        if let Err(err) = running_frame_ref.update_launch_thread_handle(thread_handle) {
-            warn!(
-                "Failed to update thread handle for frame {}. {}",
-                running_frame_ref, err
-            );
-        }
+        running_frame_ref
+            .update_launch_thread_handle(thread_handle)
+            .map_err(|err| {
+                warn!(
+                    "Failed to update thread handle for frame {}. {}",
+                    running_frame_ref, err
+                );
+                err
+            })?;
+        Ok(())
     }
 
     #[cfg(feature = "containerized_frames")]

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

@@ -370,15 +370,13 @@ impl MachineMonitor {
                 info!("Sending frame complete report: {}", frame);
 
                 // Release resources
-                if let Some(threads) = &frame.thread_ids {
-                    self.release_threads(threads).await;
-                } else {
-                    // Ensure the division rounds up if num_cores is not a multiple of
-                    // core_multiplier
-                    let num_cores_to_release = (frame.request.num_cores as u32)
-                        .div_ceil(self.maching_config.core_multiplier);
-                    self.release_cores(num_cores_to_release).await;
-                }
+                // Ensure the division rounds up if num_cores is not a multiple of
+                // core_multiplier
+                let num_cores_to_release =
+                    (frame.request.num_cores as u32).div_ceil(self.maching_config.core_multiplier);
+
+                self.release_cores(num_cores_to_release, &frame.thread_ids)
+                    .await;
 
                 // Send complete report
                 if let Err(err) = self
@@ -471,19 +469,14 @@ pub trait Machine {
     /// Vector of successfully reserved CPU core IDs
     async fn reserve_cores_by_id(&self, thread_ids: &[u32], resource_id: Uuid) -> Result<Vec<u32>>;
 
-    /// Release specific threads
-    ///
-    /// # Arguments
-    ///
-    /// * `threads` - Vector of thread IDs to release
-    async fn release_threads(&self, thread_ids: &[u32]);
-
     /// Releases a specified number of CPU cores
     ///
     /// # Arguments
     ///
     /// * `num_cores` - The number of cores to release
-    async fn release_cores(&self, num_cores: u32);
+    /// * `thread_ids` - List of threads to be released, None means the reservation didn't
+    ///   allocate specific threads
+    async fn release_cores(&self, num_cores: u32, thread_ids: &Option<Vec<u32>>);
 
     /// Reserve GPU units
     ///
@@ -585,11 +578,15 @@ impl Machine for MachineMonitor {
         // Record reservation to be reported to cuebot
         if cores_result.is_ok() {
             let mut core_state = self.core_state.lock().await;
-            debug!("Before: {:?}", *core_state);
-            core_state
-                .reserve(num_cores * self.maching_config.core_multiplier as usize)
-                .map_err(|err| miette!(err))?;
-            debug!("After: {:?}", *core_state);
+            debug!("Reservation: Before: {:?}", *core_state);
+            if let Err(err) = core_state
+                .register_reservation(num_cores * self.maching_config.core_multiplier as usize)
+                .map_err(|err| miette!(err))
+            {
+                error!("Accountability Error: {err}");
+                Err(err)?;
+            }
+            debug!("Reservation: After: {:?}", *core_state);
         }
         cores_result
     }
@@ -605,52 +602,46 @@ impl Machine for MachineMonitor {
 
         // Record reservation to be reported to cuebot
         let mut core_state = self.core_state.lock().await;
-        core_state
-            .reserve(thread_ids.len() * self.maching_config.core_multiplier as usize)
-            .map_err(|err| miette!(err))?;
+        debug!("Reservation: Before: {:?}", *core_state);
+        if let Err(err) = core_state
+            .register_reservation(thread_ids.len() * self.maching_config.core_multiplier as usize)
+            .map_err(|err| miette!(err))
+        {
+            error!("Accountability Error: {err}");
+            Err(err)?;
+        }
+        debug!("Reservation: After: {:?}", *core_state);
 
         Ok(thread_ids)
     }
 
-    async fn release_threads(&self, thread_ids: &[u32]) {
-        let mut released_cores = HashSet::new();
-        {
-            let mut system = self.system_manager.lock().await;
-            for thread_id in thread_ids {
-                match system.release_core_by_thread(thread_id) {
-                    Ok((phys_id, core_id)) => {
-                        released_cores.insert((phys_id, core_id));
-                    }
-                    Err(err) => match err {
-                        ReservationError::ReservationNotFound(_) => {
-                            // NoOp. When releasing a thread, the entire core might be released,
-                            // threfore misses are expected
-                        }
-                        _ => {
-                            error!("Failed to release proc {thread_id}. Unexpected error")
+    async fn release_cores(&self, num_cores: u32, thread_ids: &Option<Vec<u32>>) {
+        if let Some(thread_ids) = thread_ids {
+            let mut released_cores = HashSet::new();
+            {
+                let mut system = self.system_manager.lock().await;
+                for thread_id in thread_ids {
+                    match system.release_core_by_thread(thread_id) {
+                        Ok((phys_id, core_id)) => {
+                            released_cores.insert((phys_id, core_id));
                         }
-                    },
+                        Err(err) => match err {
+                            ReservationError::ReservationNotFound(_) => {
+                                // NoOp. When releasing a thread, the entire core might be released,
+                                // threfore misses are expected
+                            }
+                            _ => {
+                                error!("Failed to release proc {thread_id}. Unexpected error")
+                            }
+                        },
+                    }
                 }
             }
         }
         // Record reservation to be reported to cuebot
         let mut core_state = self.core_state.lock().await;
         if let Err(err) = core_state
-            .release(released_cores.len() as u32 * self.maching_config.core_multiplier)
-            .map_err(|err| miette!(err))
-        {
-            error!(
-                "Accountability error. Failed to release the requested number of cores. {}",
-                err
-            )
-        };
-    }
-
-    async fn release_cores(&self, num_cores: u32) {
-        // Record reservation to be reported to cuebot
-        let mut core_state = self.core_state.lock().await;
-        if let Err(err) = core_state
-            .release(num_cores * self.maching_config.core_multiplier)
+            .register_release(num_cores * self.maching_config.core_multiplier)
             .map_err(|err| miette!(err))
         {
             error!(