Eric Coissac
c0e1968654
This change introduces a comprehensive profiling system, allowing for periodic resource usage monitoring (CPU, I/O, memory) via a background thread. It also exposes new public APIs for budget management, locking, NUMA partitioning, and progress tracking.
405 lines
15 KiB
Rust
405 lines
15 KiB
Rust
use std::sync::OnceLock;
|
|
use std::time::Instant;
|
|
|
|
use libc::{RUSAGE_SELF, getrusage, rusage, timeval};
|
|
use sysinfo::System;
|
|
use tracing::debug;
|
|
|
|
// ── Memory query ──────────────────────────────────────────────────────────────
|
|
|
|
/// Returns the number of bytes available for allocation in the current process context.
|
|
///
|
|
/// On Linux, cgroup memory limits (SLURM, containers) are checked first: the
|
|
/// process may be constrained to far less than the host's available RAM.
|
|
/// Returns `min(cgroup_available, host_available)` when a finite limit is found.
|
|
///
|
|
/// On macOS, `available_memory()` can return 0 when the memory compressor
|
|
/// inflates the page count; in that case we fall back to half of total memory.
|
|
/// Returns the process peak RSS (high-water mark since process start).
|
|
/// Monotonically increasing — use delta before/after a phase to measure its RAM cost.
|
|
pub fn peak_rss_bytes() -> u64 {
|
|
rss_to_bytes(&get_rusage())
|
|
}
|
|
|
|
pub fn available_memory_bytes() -> u64 {
|
|
let sys = System::new_all();
|
|
let host_avail = match sys.available_memory() {
|
|
0 => sys.total_memory() / 2,
|
|
n => n,
|
|
};
|
|
#[cfg(target_os = "linux")]
|
|
if let Some(cg) = cgroup_v2_available().or_else(cgroup_v1_available) {
|
|
return cg.min(host_avail);
|
|
}
|
|
host_avail
|
|
}
|
|
|
|
/// cgroup v2 (unified hierarchy): reads memory.max and memory.current for the
|
|
/// current process's cgroup. Returns None if unlimited or on any parse error.
|
|
#[cfg(target_os = "linux")]
|
|
fn cgroup_v2_available() -> Option<u64> {
|
|
let cgroup = std::fs::read_to_string("/proc/self/cgroup").ok()?;
|
|
let rel = cgroup
|
|
.lines()
|
|
.find(|l| l.starts_with("0::"))?
|
|
.strip_prefix("0::")?
|
|
.trim();
|
|
let base = format!("/sys/fs/cgroup{rel}");
|
|
// "max" means no limit → parse::<u64>() fails → None
|
|
let limit: u64 = std::fs::read_to_string(format!("{base}/memory.max"))
|
|
.ok()?
|
|
.trim()
|
|
.parse()
|
|
.ok()?;
|
|
let used: u64 = std::fs::read_to_string(format!("{base}/memory.current"))
|
|
.ok()?
|
|
.trim()
|
|
.parse()
|
|
.ok()?;
|
|
Some(limit.saturating_sub(used))
|
|
}
|
|
|
|
/// cgroup v1 (memory subsystem): reads memory.limit_in_bytes and
|
|
/// memory.usage_in_bytes. Returns None if unlimited or on any parse error.
|
|
#[cfg(target_os = "linux")]
|
|
fn cgroup_v1_available() -> Option<u64> {
|
|
let cgroup = std::fs::read_to_string("/proc/self/cgroup").ok()?;
|
|
let path = cgroup
|
|
.lines()
|
|
.find(|l| l.contains(":memory:"))?
|
|
.split(':')
|
|
.nth(2)?;
|
|
let base = format!("/sys/fs/cgroup/memory{path}");
|
|
let limit: u64 = std::fs::read_to_string(format!("{base}/memory.limit_in_bytes"))
|
|
.ok()?
|
|
.trim()
|
|
.parse()
|
|
.ok()?;
|
|
// Kernel uses 2^63 (rounded to page) as "no limit" sentinel
|
|
if limit > (1u64 << 62) {
|
|
return None;
|
|
}
|
|
let used: u64 = std::fs::read_to_string(format!("{base}/memory.usage_in_bytes"))
|
|
.ok()?
|
|
.trim()
|
|
.parse()
|
|
.ok()?;
|
|
Some(limit.saturating_sub(used))
|
|
}
|
|
|
|
// ── CPU parallelism query ────────────────────────────────────────────────────
|
|
|
|
/// Returns the number of cores this process can actually use concurrently.
|
|
///
|
|
/// `std::thread::available_parallelism()` reads CPU affinity
|
|
/// (`sched_getaffinity`), not the container's CPU quota — a Docker/cgroup
|
|
/// container commonly reports the *host's* full core count this way while
|
|
/// actually being throttled (via `cpu.max`/`cpu.cfs_quota_us`) to a fraction
|
|
/// of a core. Sizing a thread/worker pool off the unthrottled count causes
|
|
/// severe oversubscription: dozens of threads contending for a sliver of
|
|
/// real CPU time, which can look indistinguishable from a hang for minutes
|
|
/// or hours (observed in CI). On Linux, this reads the cgroup CPU quota
|
|
/// first and returns `min(cgroup_quota, host_parallelism)` when a finite
|
|
/// quota is found; falls back to `available_parallelism()` otherwise (same
|
|
/// convention as [`available_memory_bytes`]).
|
|
pub fn effective_parallelism() -> usize {
|
|
let host = std::thread::available_parallelism().map(|n| n.get()).unwrap_or(1);
|
|
#[cfg(target_os = "linux")]
|
|
{
|
|
if let Some(quota) = cgroup_v2_cpu_quota() {
|
|
let effective = quota.clamp(1, host);
|
|
tracing::debug!(host, quota, effective, source = "cgroup v2", "effective_parallelism");
|
|
return effective;
|
|
}
|
|
if let Some(quota) = cgroup_v1_cpu_quota() {
|
|
let effective = quota.clamp(1, host);
|
|
tracing::debug!(host, quota, effective, source = "cgroup v1", "effective_parallelism");
|
|
return effective;
|
|
}
|
|
}
|
|
tracing::debug!(host, effective = host, source = "available_parallelism (no cgroup quota found)", "effective_parallelism");
|
|
host
|
|
}
|
|
|
|
// ── Process-wide CPU budget (hard ceiling) ───────────────────────────────────
|
|
|
|
static CPU_CAP: OnceLock<usize> = OnceLock::new();
|
|
|
|
/// Install a hard ceiling on the number of cores the process may use.
|
|
///
|
|
/// The value is clamped to `[1, effective_parallelism()]` — a cap can only
|
|
/// lower the budget, never raise it above what the process is actually
|
|
/// allowed to run on. First call wins; later calls are ignored.
|
|
///
|
|
/// Call this once at startup (before sizing any worker pool or the global
|
|
/// rayon pool) when the user passes an explicit `--cpu-max`.
|
|
pub fn set_cpu_cap(n: usize) {
|
|
let capped = n.clamp(1, effective_parallelism());
|
|
let _ = CPU_CAP.set(capped);
|
|
}
|
|
|
|
/// The process-wide CPU budget: the ceiling set via [`set_cpu_cap`] if any,
|
|
/// otherwise [`effective_parallelism`].
|
|
///
|
|
/// Every worker pool and the global rayon pool must size themselves against
|
|
/// this value rather than calling [`effective_parallelism`] directly.
|
|
pub fn cpu_budget() -> usize {
|
|
CPU_CAP.get().copied().unwrap_or_else(effective_parallelism)
|
|
}
|
|
|
|
/// cgroup v2 (unified hierarchy): reads `cpu.max` ("<quota> <period>", or
|
|
/// "max <period>" when unlimited) for the current process's cgroup, rounded
|
|
/// up to whole cores. Returns `None` if unlimited or on any parse error.
|
|
#[cfg(target_os = "linux")]
|
|
fn cgroup_v2_cpu_quota() -> Option<usize> {
|
|
let cgroup = std::fs::read_to_string("/proc/self/cgroup").ok()?;
|
|
let rel = cgroup
|
|
.lines()
|
|
.find(|l| l.starts_with("0::"))?
|
|
.strip_prefix("0::")?
|
|
.trim();
|
|
let base = format!("/sys/fs/cgroup{rel}");
|
|
let raw = std::fs::read_to_string(format!("{base}/cpu.max")).ok()?;
|
|
let mut parts = raw.split_whitespace();
|
|
let quota_str = parts.next()?;
|
|
let period: f64 = parts.next()?.parse().ok()?;
|
|
if quota_str == "max" {
|
|
return None; // unlimited
|
|
}
|
|
let quota: f64 = quota_str.parse().ok()?;
|
|
Some((quota / period).ceil().max(1.0) as usize)
|
|
}
|
|
|
|
/// cgroup v1 (cpu subsystem): reads `cpu.cfs_quota_us`/`cpu.cfs_period_us`,
|
|
/// rounded up to whole cores. Returns `None` if unlimited (quota <= 0) or on
|
|
/// any parse error.
|
|
#[cfg(target_os = "linux")]
|
|
fn cgroup_v1_cpu_quota() -> Option<usize> {
|
|
let cgroup = std::fs::read_to_string("/proc/self/cgroup").ok()?;
|
|
let path = cgroup
|
|
.lines()
|
|
.find(|l| l.contains(":cpu:") || l.contains(":cpu,cpuacct:"))?
|
|
.split(':')
|
|
.nth(2)?;
|
|
let base = format!("/sys/fs/cgroup/cpu{path}");
|
|
let quota: i64 = std::fs::read_to_string(format!("{base}/cpu.cfs_quota_us"))
|
|
.ok()?
|
|
.trim()
|
|
.parse()
|
|
.ok()?;
|
|
if quota <= 0 {
|
|
return None; // unlimited
|
|
}
|
|
let period: i64 = std::fs::read_to_string(format!("{base}/cpu.cfs_period_us"))
|
|
.ok()?
|
|
.trim()
|
|
.parse()
|
|
.ok()?;
|
|
if period <= 0 {
|
|
return None;
|
|
}
|
|
Some(((quota as f64) / (period as f64)).ceil().max(1.0) as usize)
|
|
}
|
|
|
|
// ── raw helpers ───────────────────────────────────────────────────────────────
|
|
|
|
pub(crate) fn get_rusage() -> rusage {
|
|
let mut ru = unsafe { std::mem::zeroed::<rusage>() };
|
|
unsafe { getrusage(RUSAGE_SELF, &mut ru) };
|
|
ru
|
|
}
|
|
|
|
pub(crate) fn tv_to_secs(tv: timeval) -> f64 {
|
|
tv.tv_sec as f64 + tv.tv_usec as f64 * 1e-6
|
|
}
|
|
|
|
#[cfg(target_os = "macos")]
|
|
pub(crate) fn rss_to_bytes(ru: &rusage) -> u64 {
|
|
ru.ru_maxrss as u64
|
|
}
|
|
|
|
#[cfg(not(target_os = "macos"))]
|
|
pub(crate) fn rss_to_bytes(ru: &rusage) -> u64 {
|
|
ru.ru_maxrss as u64 * 1024
|
|
}
|
|
|
|
// Monotonically increasing counters — negative delta would be a kernel bug.
|
|
pub(crate) fn delta(end: i64, start: i64) -> u64 {
|
|
(end - start).max(0) as u64
|
|
}
|
|
|
|
// ── CpuSample ─────────────────────────────────────────────────────────────────
|
|
|
|
/// Snapshot of process-wide CPU time + wall clock at a point in time.
|
|
/// Use [`cpu_efficiency`](Self::cpu_efficiency) to measure the fraction of
|
|
/// available cores used since the snapshot was taken.
|
|
pub struct CpuSample {
|
|
wall: Instant,
|
|
user_secs: f64,
|
|
sys_secs: f64,
|
|
previous: f64,
|
|
}
|
|
|
|
impl CpuSample {
|
|
pub fn now() -> Self {
|
|
let ru = get_rusage();
|
|
Self {
|
|
wall: Instant::now(),
|
|
user_secs: tv_to_secs(ru.ru_utime),
|
|
sys_secs: tv_to_secs(ru.ru_stime),
|
|
previous: 0.0,
|
|
}
|
|
}
|
|
|
|
/// (user_delta + sys_delta) / (wall_delta × n_cores) since this snapshot.
|
|
/// Returns 0.0 if less than 100 ms have elapsed (too noisy).
|
|
pub fn cpu_efficiency(&self, n_cores: usize) -> f64 {
|
|
let ru = get_rusage();
|
|
let wall = self.wall.elapsed().as_secs_f64();
|
|
if wall < 0.1 {
|
|
return 0.0;
|
|
}
|
|
let cpu =
|
|
(tv_to_secs(ru.ru_utime) - self.user_secs) + (tv_to_secs(ru.ru_stime) - self.sys_secs);
|
|
cpu / (wall * n_cores as f64)
|
|
}
|
|
|
|
/// Cores of CPU time consumed per wall-clock second since the last call
|
|
/// (or since [`now`](Self::now) for the first call) — an absolute core
|
|
/// count, unlike [`cpu_efficiency`](Self::cpu_efficiency) which
|
|
/// normalizes to `n_cores` and never resets its baseline. Resets the
|
|
/// window on every call that returns `Some`; returns `None` (state left
|
|
/// untouched, so the window keeps accumulating) when less than 100 ms
|
|
/// have elapsed — too short to be meaningful.
|
|
pub fn sample_cores(&mut self) -> Option<f64> {
|
|
let delta_wall = self.wall.elapsed().as_secs_f64();
|
|
if delta_wall < 0.1 {
|
|
return None;
|
|
}
|
|
let n = CpuSample::now();
|
|
let delta_ru = (n.user_secs - self.user_secs) + (n.sys_secs - self.sys_secs);
|
|
self.user_secs = n.user_secs;
|
|
self.sys_secs = n.sys_secs;
|
|
self.wall = n.wall;
|
|
Some(delta_ru / delta_wall)
|
|
}
|
|
|
|
pub fn do_i_activate(&mut self, threshold: f64) -> bool {
|
|
let Some(efficiency) = self.sample_cores() else {
|
|
return false;
|
|
};
|
|
|
|
let activate = 0f64.max(efficiency - self.previous) >= threshold;
|
|
|
|
debug!(
|
|
"Do I activate : {} -> {} = {} Activate: {}",
|
|
self.previous,
|
|
efficiency,
|
|
0f64.max(efficiency - self.previous),
|
|
activate
|
|
);
|
|
self.previous = efficiency;
|
|
|
|
activate
|
|
}
|
|
}
|
|
|
|
// ── IoSample ──────────────────────────────────────────────────────────────────
|
|
|
|
/// Snapshot of process-wide block I/O (bytes read + written) + wall clock.
|
|
///
|
|
/// Same activation protocol as [`CpuSample`], but the growth check in
|
|
/// [`do_i_activate`](Self::do_i_activate) is *relative* rather than absolute:
|
|
/// raw I/O throughput has no portable scale across storage devices, unlike a
|
|
/// core count.
|
|
pub struct IoSample {
|
|
wall: Instant,
|
|
bytes: u64,
|
|
previous_rate: f64,
|
|
}
|
|
|
|
impl IoSample {
|
|
pub fn now() -> Self {
|
|
Self {
|
|
wall: Instant::now(),
|
|
bytes: Self::read_bytes(),
|
|
previous_rate: 0.0,
|
|
}
|
|
}
|
|
|
|
/// Bytes actually submitted to the block layer (read + write), summed
|
|
/// process-wide. Returns 0 if unavailable — degrades gracefully to a
|
|
/// signal that never triggers activation (CPU-only heuristic).
|
|
#[cfg(target_os = "linux")]
|
|
fn read_bytes() -> u64 {
|
|
let Ok(io) = std::fs::read_to_string("/proc/self/io") else {
|
|
return 0;
|
|
};
|
|
io.lines()
|
|
.filter_map(|l| {
|
|
l.strip_prefix("read_bytes: ")
|
|
.or_else(|| l.strip_prefix("write_bytes: "))
|
|
})
|
|
.filter_map(|v| v.trim().parse::<u64>().ok())
|
|
.sum()
|
|
}
|
|
|
|
#[cfg(target_os = "macos")]
|
|
fn read_bytes() -> u64 {
|
|
use libc::{RUSAGE_INFO_V4, getpid, proc_pid_rusage, rusage_info_v4};
|
|
let mut info: rusage_info_v4 = unsafe { std::mem::zeroed() };
|
|
let ret =
|
|
unsafe { proc_pid_rusage(getpid(), RUSAGE_INFO_V4, &mut info as *mut _ as *mut _) };
|
|
if ret != 0 {
|
|
return 0;
|
|
}
|
|
info.ri_diskio_bytesread + info.ri_diskio_byteswritten
|
|
}
|
|
|
|
#[cfg(not(any(target_os = "linux", target_os = "macos")))]
|
|
fn read_bytes() -> u64 {
|
|
0
|
|
}
|
|
|
|
/// Bytes/second submitted to the block layer since the last call (or
|
|
/// since [`now`](Self::now) for the first call). Resets the window on
|
|
/// every call that returns `Some`; returns `None` (state left
|
|
/// untouched) when less than 100 ms have elapsed — too short to be
|
|
/// meaningful.
|
|
pub fn sample_rate(&mut self) -> Option<f64> {
|
|
let elapsed = self.wall.elapsed().as_secs_f64();
|
|
if elapsed < 0.1 {
|
|
return None;
|
|
}
|
|
let n = Self::read_bytes();
|
|
let rate = n.saturating_sub(self.bytes) as f64 / elapsed;
|
|
self.bytes = n;
|
|
self.wall = Instant::now();
|
|
Some(rate)
|
|
}
|
|
|
|
/// Same protocol as [`CpuSample::do_i_activate`] (0.1 s minimum window,
|
|
/// state untouched on early return), but growth is measured relative to
|
|
/// the previous rate. `threshold` is a fraction, e.g. `0.2` for a 20 %
|
|
/// increase in throughput since the last real sample.
|
|
pub fn do_i_activate(&mut self, threshold: f64) -> bool {
|
|
let Some(rate) = self.sample_rate() else {
|
|
return false;
|
|
};
|
|
|
|
let activate = if self.previous_rate == 0.0 {
|
|
rate > 0.0 // bootstrap: any measured throughput is signal enough
|
|
} else {
|
|
(rate - self.previous_rate) / self.previous_rate >= threshold
|
|
};
|
|
|
|
debug!(
|
|
"Do I activate (I/O) : {} -> {} Activate: {}",
|
|
self.previous_rate, rate, activate
|
|
);
|
|
self.previous_rate = rate;
|
|
|
|
activate
|
|
}
|
|
}
|