use crate::job::JobStatus;
use crate::paths;
use crate::spec::JobSpec;
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;
const SAMPLES: usize = 5;
const MIN_MEMORY: u64 = 64 << 20;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
#[serde(rename_all = "kebab-case")]
pub enum Measurement {
#[default]
Peak,
LowerBound,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Sample {
#[serde(default)]
pub kind: Measurement,
pub max_rss: u64,
pub cpu_secs: f64,
pub elapsed_secs: u64,
pub at: u64,
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct Entry {
pub name: String,
pub samples: Vec<Sample>,
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct Store {
#[serde(default)]
pub commands: BTreeMap<String, Entry>,
}
#[derive(Debug, Clone, PartialEq)]
pub struct Suggestion {
pub cpu: u64,
pub mem: u64,
pub samples: usize,
}
pub fn key(cwd: &std::path::Path, command: &[String]) -> String {
let mut hash: u64 = 0xcbf29ce484222325;
for byte in cwd.to_string_lossy().as_bytes() {
hash ^= *byte as u64;
hash = hash.wrapping_mul(0x100000001b3);
}
hash ^= 0xfe;
hash = hash.wrapping_mul(0x100000001b3);
for part in command {
for byte in part.as_bytes() {
hash ^= *byte as u64;
hash = hash.wrapping_mul(0x100000001b3);
}
hash ^= 0xff;
hash = hash.wrapping_mul(0x100000001b3);
}
format!("{hash:016x}")
}
fn store_path() -> anyhow::Result<std::path::PathBuf> {
Ok(paths::state_dir()?.join("usage.json"))
}
pub fn load() -> Store {
let Ok(path) = store_path() else {
return Store::default();
};
let Ok(text) = std::fs::read_to_string(path) else {
return Store::default();
};
serde_json::from_str(&text).unwrap_or_default()
}
pub fn record(spec: &JobSpec, status: &JobStatus) {
if status.state != crate::job::JobState::Completed || status.usage.max_rss == 0 {
return;
}
add(spec, status, Measurement::Peak, status.usage.max_rss);
}
pub fn record_lower_bound(spec: &JobSpec, status: &JobStatus) {
if status.state != crate::job::JobState::Oom {
return;
}
add(
spec,
status,
Measurement::LowerBound,
status.usage.max_rss.max(status.mem),
);
}
fn add(spec: &JobSpec, status: &JobStatus, kind: Measurement, bytes: u64) {
if bytes == 0 {
return;
}
let Ok(path) = store_path() else { return };
let Ok(dir) = paths::state_dir() else { return };
if paths::ensure_dir(&dir, 0o700).is_err() {
return;
}
let lock_path = dir.join("usage.lock");
let Ok(lock) = std::fs::OpenOptions::new()
.create(true)
.write(true)
.truncate(false)
.open(&lock_path)
else {
return;
};
use std::os::unix::io::AsRawFd;
unsafe {
libc::flock(lock.as_raw_fd(), libc::LOCK_EX);
}
let against = spec.learn_key.as_deref().unwrap_or(&spec.command);
let mut store = load();
let entry = store.commands.entry(key(&spec.cwd, against)).or_default();
entry.name = spec.name.clone();
if kind == Measurement::LowerBound {
if let Some(pos) = entry
.samples
.iter()
.position(|s| s.kind == Measurement::LowerBound)
{
if entry.samples[pos].max_rss >= bytes {
let earlier = entry.samples.remove(pos);
entry.samples.push(earlier);
write_store(&path, &store, &lock);
return;
}
entry.samples.remove(pos);
}
}
entry.samples.push(Sample {
kind,
max_rss: bytes,
cpu_secs: status.usage.cpu_secs,
elapsed_secs: status.elapsed().map(|d| d.as_secs()).unwrap_or(0),
at: crate::sys::now_secs(),
});
let extra = entry.samples.len().saturating_sub(SAMPLES);
entry.samples.drain(..extra);
write_store(&path, &store, &lock);
}
fn write_store(path: &std::path::Path, store: &Store, lock: &std::fs::File) {
use std::os::unix::io::AsRawFd;
if let Ok(bytes) = serde_json::to_vec_pretty(store) {
crate::job::write_atomic(path, &bytes, 0o600).ok();
}
unsafe {
libc::flock(lock.as_raw_fd(), libc::LOCK_UN);
}
}
pub fn suggest(
store: &Store,
cwd: &std::path::Path,
command: &[String],
margin: f64,
) -> Option<Suggestion> {
let entry = store.commands.get(&key(cwd, command))?;
if entry.samples.is_empty() {
return None;
}
let peak_mem = entry
.samples
.iter()
.filter(|s| s.kind == Measurement::Peak)
.map(|s| s.max_rss)
.max()
.unwrap_or(0);
let mut mem = (peak_mem as f64 * margin) as u64;
let bound = entry
.samples
.iter()
.filter(|s| s.kind == Measurement::LowerBound)
.map(|s| s.max_rss)
.max()
.unwrap_or(0);
if bound > 0 {
let above = ((bound as f64 * margin) as u64).max(bound + bound / 10 + 1);
mem = mem.max(above);
}
let mem = mem.max(MIN_MEMORY);
let cores = entry
.samples
.iter()
.map(|s| {
s.cpu_secs / s.elapsed_secs.max(1) as f64
})
.fold(0.0f64, f64::max);
let cpu = (cores * margin).ceil().max(1.0) as u64;
Some(Suggestion {
cpu,
mem,
samples: entry.samples.len(),
})
}
#[cfg(test)]
mod tests {
use super::*;
fn sample(max_rss: u64, cpu_secs: f64, elapsed_secs: u64) -> Sample {
Sample {
kind: Measurement::Peak,
max_rss,
cpu_secs,
elapsed_secs,
at: 0,
}
}
fn lower_bound(max_rss: u64) -> Sample {
Sample {
kind: Measurement::LowerBound,
max_rss,
cpu_secs: 1.0,
elapsed_secs: 1,
at: 0,
}
}
fn dir() -> std::path::PathBuf {
std::path::PathBuf::from("/project")
}
fn store_with(command: &[&str], samples: Vec<Sample>) -> Store {
let cmd: Vec<String> = command.iter().map(|s| s.to_string()).collect();
let mut store = Store::default();
store.commands.insert(
key(&dir(), &cmd),
Entry {
name: "test".into(),
samples,
},
);
store
}
#[test]
fn with_no_measurement_there_is_no_claim() {
let store = Store::default();
assert_eq!(suggest(&store, &dir(), &["cargo".into()], 1.5), None);
}
#[test]
fn the_claim_uses_the_largest_measurement() {
let store = store_with(
&["cargo", "test"],
vec![
sample(100 << 20, 1.0, 10),
sample(400 << 20, 1.0, 10),
sample(200 << 20, 1.0, 10),
],
);
let cmd: Vec<String> = vec!["cargo".into(), "test".into()];
let s = suggest(&store, &dir(), &cmd, 1.5).unwrap();
assert_eq!(s.mem, (400 << 20) * 3 / 2, "400MB and one half");
assert_eq!(s.samples, 3);
}
#[test]
fn the_cores_come_from_the_cpu_time_and_the_elapsed_time() {
let cmd: Vec<String> = vec!["make".into()];
let store = store_with(&["make"], vec![sample(1 << 20, 20.0, 10)]);
assert_eq!(suggest(&store, &dir(), &cmd, 1.0).unwrap().cpu, 2);
let store = store_with(&["make"], vec![sample(165 << 20, 1.9, 19)]);
let s = suggest(&store, &dir(), &cmd, 1.5).unwrap();
assert_eq!(s.cpu, 1, "a job that waits needs one core");
assert!(
s.mem < (300 << 20),
"the claim must be near the measurement, and it was {}",
crate::units::format_size(s.mem)
);
}
#[test]
fn a_small_measurement_gives_the_smallest_useful_claim() {
let cmd: Vec<String> = vec!["true".into()];
let store = store_with(&["true"], vec![sample(1 << 20, 0.0, 0)]);
let s = suggest(&store, &dir(), &cmd, 1.5).unwrap();
assert_eq!(s.mem, MIN_MEMORY);
assert_eq!(s.cpu, 1);
}
#[test]
fn two_commands_have_two_records() {
let build: Vec<String> = vec!["cargo".into(), "build".into()];
let test: Vec<String> = vec!["cargo".into(), "test".into()];
assert_ne!(key(&dir(), &build), key(&dir(), &test));
let store = store_with(&["cargo", "build"], vec![sample(1 << 30, 1.0, 1)]);
assert!(suggest(&store, &dir(), &build, 1.5).is_some());
assert!(
suggest(&store, &dir(), &test, 1.5).is_none(),
"`cargo test` must not use the record of `cargo build`"
);
}
#[test]
fn the_key_separates_the_arguments() {
let joined: Vec<String> = vec!["a b".into()];
let split: Vec<String> = vec!["a".into(), "b".into()];
assert_ne!(key(&dir(), &joined), key(&dir(), &split));
}
#[test]
fn one_command_in_two_directories_has_two_records() {
let cmd: Vec<String> = vec!["cargo".into(), "test".into()];
let small = std::path::PathBuf::from("/home/me/small-library");
let large = std::path::PathBuf::from("/home/me/large-program");
assert_ne!(key(&small, &cmd), key(&large, &cmd));
let mut store = Store::default();
store.commands.insert(
key(&small, &cmd),
Entry {
name: "test".into(),
samples: vec![sample(100 << 20, 1.0, 10)],
},
);
assert!(suggest(&store, &small, &cmd, 1.5).is_some());
assert!(
suggest(&store, &large, &cmd, 1.5).is_none(),
"a different directory must not use this record"
);
}
#[test]
fn every_line_of_a_fan_out_shares_one_record() {
let template: Vec<String> = vec!["./process".into(), "{}".into()];
let line_a: Vec<String> = vec!["./process".into(), "a.csv".into()];
let line_b: Vec<String> = vec!["./process".into(), "b.csv".into()];
assert_ne!(key(&dir(), &line_a), key(&dir(), &line_b));
assert_eq!(key(&dir(), &template), key(&dir(), &template));
let store = store_with(&["./process", "{}"], vec![sample(400 << 20, 1.0, 10)]);
assert!(
suggest(&store, &dir(), &line_b, 1.5).is_none(),
"the command of the line must not reach the record"
);
let s = suggest(&store, &dir(), &template, 1.5).unwrap();
assert_eq!(s.mem, (400 << 20) * 3 / 2);
}
#[test]
fn the_next_claim_is_above_a_lower_bound() {
let cmd: Vec<String> = vec!["train".into()];
let store = store_with(&["train"], vec![lower_bound(8 << 30)]);
let s = suggest(&store, &dir(), &cmd, 1.5).unwrap();
assert!(
s.mem > (8 << 30),
"the claim must be above a claim that failed, and it was {}",
crate::units::format_size(s.mem)
);
assert_eq!(s.mem, 12 << 30);
let s = suggest(&store, &dir(), &cmd, 1.0).unwrap();
assert!(
s.mem > (8 << 30),
"a margin of 1.0 gave the claim that already failed"
);
}
#[test]
fn a_lower_bound_is_not_averaged_away_by_the_smaller_runs() {
let cmd: Vec<String> = vec!["train".into()];
let store = store_with(
&["train"],
vec![
sample(1 << 30, 1.0, 10),
lower_bound(8 << 30),
sample(1 << 30, 1.0, 10),
sample(1 << 30, 1.0, 10),
],
);
let s = suggest(&store, &dir(), &cmd, 1.5).unwrap();
assert!(
s.mem > (8 << 30),
"the lower bound went away, and the claim is {}",
crate::units::format_size(s.mem)
);
}
#[test]
fn the_largest_evidence_wins_whatever_its_kind() {
let cmd: Vec<String> = vec!["train".into()];
let store = store_with(
&["train"],
vec![lower_bound(2 << 30), sample(6 << 30, 1.0, 10)],
);
let s = suggest(&store, &dir(), &cmd, 1.5).unwrap();
assert_eq!(s.mem, 9 << 30, "6GB and one half");
}
#[test]
fn an_earlier_file_keeps_its_meaning() {
let text = r#"{"commands":{"x":{"name":"t","samples":[
{"max_rss":1073741824,"cpu_secs":1.0,"elapsed_secs":10,"at":0}]}}}"#;
let store: Store = serde_json::from_str(text).unwrap();
let sample = &store.commands["x"].samples[0];
assert_eq!(sample.kind, Measurement::Peak);
assert_eq!(sample.max_rss, 1 << 30);
}
#[test]
fn a_kill_for_memory_records_the_claim_when_it_is_larger() {
use crate::testutil::{env_lock, EnvVar};
let _guard = env_lock();
let dir = std::env::temp_dir().join(format!("qex-usage-oom-{}", std::process::id()));
std::fs::remove_dir_all(&dir).ok();
std::fs::create_dir_all(&dir).unwrap();
let _env = EnvVar::set("XDG_STATE_HOME", dir.to_str().unwrap());
let mut spec = crate::spec::JobSpec {
id: uuid::Uuid::new_v4(),
name: "train".into(),
cwd: "/project".into(),
command: vec!["train".into()],
env: Default::default(),
cpu: 1,
mem: 4 << 30,
timeout: None,
max_queue_time: None,
tags: vec![],
priority: 0,
env_capture: crate::config::EnvCapture::None,
claim_source: "explicit".into(),
learn_key: None,
group: None,
group_name: None,
claims: Default::default(),
locks: vec![],
retries: 0,
nice: None,
needs: vec![],
after: vec![],
submitted_at: 0,
dedupe_key: None,
dedupe_window: 0,
};
let mut status = crate::job::JobStatus::new(&spec);
status.state = crate::job::JobState::Oom;
status.usage.max_rss = 1 << 30;
record_lower_bound(&spec, &status);
record(&spec, &status);
let store = load();
let entry = &store.commands[&key(&spec.cwd, &spec.command)];
assert_eq!(entry.samples.len(), 1);
assert_eq!(entry.samples[0].kind, Measurement::LowerBound);
assert_eq!(
entry.samples[0].max_rss,
4 << 30,
"the claim is the larger evidence"
);
spec.command = vec!["ladder".into()];
let ladder = key(&spec.cwd, &spec.command);
let mut peak = crate::job::JobStatus::new(&spec);
peak.state = crate::job::JobState::Completed;
peak.usage.max_rss = 2 << 30;
record(&spec, &peak);
for claim in [1u64 << 30, 2 << 30, 4 << 30] {
let mut oom = crate::job::JobStatus::new(&spec);
oom.state = crate::job::JobState::Oom;
oom.mem = claim;
oom.usage.max_rss = claim;
record_lower_bound(&spec, &oom);
}
let entry = &load().commands[&ladder];
let bounds: Vec<&Sample> = entry
.samples
.iter()
.filter(|s| s.kind == Measurement::LowerBound)
.collect();
assert_eq!(
bounds.len(),
1,
"a ladder of three attempts must leave one bound: {:?}",
entry.samples
);
assert_eq!(bounds[0].max_rss, 4 << 30, "the largest bound must win");
assert!(
entry
.samples
.iter()
.any(|s| s.kind == Measurement::Peak && s.max_rss == (2 << 30)),
"the measurement of the job that completed must stay: {:?}",
entry.samples
);
spec.command = vec!["other".into()];
let mut killed = crate::job::JobStatus::new(&spec);
killed.state = crate::job::JobState::Killed;
killed.usage.max_rss = 3 << 30;
record(&spec, &killed);
assert!(
!load().commands.contains_key(&key(&spec.cwd, &spec.command)),
"a job that a command stopped must teach the learner nothing"
);
std::fs::remove_dir_all(&dir).ok();
}
#[test]
fn the_key_does_not_hold_the_command() {
let secret: Vec<String> = vec!["deploy".into(), "--token=SECRET123".into()];
let k = key(&dir(), &secret);
assert!(!k.contains("SECRET"), "the key holds the command: {k}");
assert!(!k.contains("project"), "the key holds the directory: {k}");
assert_eq!(k.len(), 16, "the key is a hash of a fixed length");
}
}