[rust/rqd] Add option to measure proc memory using PSS (#2089)

Overview

Add a configuration option to RQD to choose between RSS (Resident Set
Size) and PSS (Proportional Set Size) for measuring process memory usage
on Linux systems.

**Current State:** 
- RQD measures memory using RSS from `/proc/[pid]/statm` 
- RSS counts shared libraries fully for each process, leading to
inflated memory reports
- PSS divides shared memory proportionally among processes, providing
more accurate measurements

**Key Differences:** - **RSS**: Total physical memory including full
count of shared libraries (can be misleading)
- **PSS**: Proportional share - shared pages divided by number of
processes sharing them (more accurate)
- **Source**: `/proc/[pid]/smaps_rollup` (Linux kernel 4.14+) provides
pre-summed PSS values

Author: DiegoTavares

rust/Cargo.toml

@@ -5,7 +5,7 @@ resolver = "3"
 [workspace.package]
 authors = ["Diego Tavares <[email protected]>"]
 edition = "2024"
-version = "0.1.5"
+version = "0.1.6"
 
 [workspace.dependencies]
 async-trait = "0.1"

rust/config/rqd.yaml

@@ -140,6 +140,24 @@ machine:
   #   # Override the operating system reported to Cuebot
   #   os: "Rocky Linux 9"
 
+  # Memory metric to use for process memory measurement
+  # Options: "rss" (default) or "pss"
+  #
+  # RSS (Resident Set Size): Total physical memory used by a process, including
+  #   shared libraries counted fully. This can overreport memory usage when
+  #   multiple processes share the same libraries.
+  #
+  # PSS (Proportional Set Size): Divides shared memory proportionally among
+  #   processes using it, providing more accurate memory accounting. For example,
+  #   if three processes share a 30MB library, each process reports 10MB instead
+  #   of 30MB.
+  #
+  # Note: PSS requires Linux kernel 4.14+ and reads from /proc/[pid]/smaps_rollup.
+  #       If PSS is unavailable, RQD will automatically fall back to RSS.
+  #
+  # Default: rss
+  # memory_metric: "rss"
+
 # =============================================================================
 # RUNNER CONFIGURATION
 # =============================================================================

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

@@ -14,6 +14,14 @@ lazy_static! {
 }
 //===Config Types===
 
+#[derive(Default, Debug, Serialize, Deserialize, Clone, Copy, PartialEq)]
+#[serde(rename_all = "lowercase")]
+pub enum MemoryMetric {
+    #[default]
+    Rss,
+    Pss,
+}
+
 #[derive(Debug, Deserialize, Clone)]
 #[serde(default)]
 pub struct LoggingConfig {
@@ -90,6 +98,7 @@ pub struct MachineConfig {
     pub nimby_start_retry_interval: Duration,
     pub nimby_display_xauthority_path: String,
     pub memory_oom_margin_percentage: u32,
+    pub memory_metric: MemoryMetric,
 }
 
 impl Default for MachineConfig {
@@ -113,6 +122,7 @@ impl Default for MachineConfig {
             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,
+            memory_metric: MemoryMetric::Rss,
         }
     }
 }

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

@@ -21,10 +21,13 @@ use opencue_proto::{
     report::{ChildrenProcStats, ProcStats, Stat},
 };
 use sysinfo::{DiskRefreshKind, Disks, MemoryRefreshKind, RefreshKind};
-use tracing::debug;
+use tracing::{debug, info};
 use uuid::Uuid;
 
-use crate::{config::MachineConfig, system::reservation::ProcessorStructure};
+use crate::{
+    config::{MachineConfig, MemoryMetric},
+    system::reservation::ProcessorStructure,
+};
 
 use super::manager::{MachineGpuStats, MachineStat, ProcessStats, SystemManager};
 
@@ -144,6 +147,8 @@ impl LinuxSystem {
             .into_diagnostic()
             .wrap_err("SC_CLK_TCK not available")?;
 
+        info!("Memory metric configured: {:?}", config.memory_metric);
+
         Ok(Self {
             config: config.clone(),
             static_info: MachineStaticInfo {
@@ -562,6 +567,35 @@ impl LinuxSystem {
         Ok(())
     }
 
+    /// Reads PSS (Proportional Set Size) from /proc/[pid]/smaps_rollup
+    ///
+    /// PSS divides shared memory proportionally among processes using it,
+    /// providing more accurate memory accounting than RSS.
+    ///
+    /// Requires Linux kernel 4.14+. Returns error if unavailable.
+    fn read_pss(&self, pid: u32) -> Result<u64> {
+        let smaps_rollup_path = format!("/proc/{}/smaps_rollup", pid);
+
+        let content = std::fs::read_to_string(&smaps_rollup_path).into_diagnostic()?;
+
+        for line in content.lines() {
+            if line.starts_with("Pss:") {
+                let parts: Vec<&str> = line.split_whitespace().collect();
+                if parts.len() >= 2 {
+                    if let Ok(pss_kb) = parts[1].parse::<u64>() {
+                        // PSS is in kB, convert to bytes
+                        return Ok(pss_kb * 1024);
+                    }
+                }
+            }
+        }
+
+        Err(miette!(
+            "Could not parse PSS from /proc/{}/smaps_rollup",
+            pid
+        ))
+    }
+
     /// Reads proc data from stat and statm files:
     ///
     /// # Used fields:
@@ -603,14 +637,30 @@ impl LinuxSystem {
                 fields_statm[0].parse::<u64>(), // size
                 fields_statm[1].parse::<u64>(), // rss
             ) {
-                (Ok(vsize), Ok(rss)) => (vsize, rss),
+                (Ok(vsize), Ok(rss)) => (
+                    vsize.saturating_mul(self.static_info.page_size),
+                    rss.saturating_mul(self.static_info.page_size),
+                ),
                 _ => Err(miette!("Invalid /proc/{pid}/statm file"))?,
             };
+            let virtual_memory = vsize.saturating_mul(self.static_info.page_size);
+
+            // Read memory based on configured metric
+            let memory = match self.config.memory_metric {
+                MemoryMetric::Pss => {
+                    // Try PSS, fallback to RSS if unavailable
+                    match self.read_pss(pid) {
+                        Ok(pss) => pss,
+                        Err(_) => rss,
+                    }
+                }
+                MemoryMetric::Rss => {
+                    // Original RSS logic
+                    rss
+                }
+            };
 
             let (start_time, run_time) = self.calculate_process_time(start_time);
-            // Rss is stored in number of pages
-            let memory = rss.saturating_mul(self.static_info.page_size);
-            let virtual_memory = vsize.saturating_mul(self.static_info.page_size);
 
             // Remove ()
             let name = if name.len() > 2 {
@@ -904,7 +954,7 @@ impl SystemManager for LinuxSystem {
 
 #[cfg(test)]
 mod tests {
-    use crate::config::MachineConfig;
+    use crate::config::{MachineConfig, MemoryMetric};
     use std::fs;
     use std::{collections::HashMap, sync::Mutex};