resource_tracker/
config.rs

1use clap::{ArgAction, Parser, ValueEnum};
2use serde::Deserialize;
3
4const DEFAULT_INTERVAL_SECS: u64 = 1;
5const RENICE_MIN: i64 = -20;
6const RENICE_MAX: i64 = 19;
7const DEFAULT_CONFIG_FILE: &str = "resource-tracker.toml";
8
9// ---------------------------------------------------------------------------
10// Output format
11// ---------------------------------------------------------------------------
12//
13/// Output format emitted to stdout on each polling interval.
14#[derive(Debug, Clone, Copy, PartialEq, ValueEnum)]
15pub enum OutputFormat {
16    /// JSON Lines - one JSON object per line (default).
17    Json,
18    /// CSV - header on first line, one row per interval.
19    /// Columns mirror Python resource-tracker's SystemTracker output.
20    Csv,
21}
22
23// ---------------------------------------------------------------------------
24// TOML file structure
25// ---------------------------------------------------------------------------
26//
27#[derive(Debug, Default, Deserialize)]
28struct TomlConfig {
29    job: Option<TomlJob>,
30    tracker: Option<TomlTracker>,
31}
32
33#[derive(Debug, Deserialize)]
34struct TomlJob {
35    /// Human-readable label attached to every sample (e.g. "benchmark-run-42").
36    name: Option<String>,
37    /// Root PID of the process tree whose CPU usage should be attributed.
38    pid: Option<i32>,
39}
40
41#[derive(Debug, Deserialize)]
42struct TomlTracker {
43    /// How often to emit a sample, in seconds. Default: 1.
44    interval_secs: Option<u64>,
45    /// Resource tracker nice value. Default: no change.
46    renice: Option<i32>,
47    /// Aggregate CPU steal value. Default: yes.
48    aggregate_cpu_steal: Option<bool>,
49}
50
51// ---------------------------------------------------------------------------
52// Job metadata (Section 9.3) - sent to Sentinel API at run registration
53// ---------------------------------------------------------------------------
54//
55/// All optional metadata fields from Section 9.3 of the spec.
56/// Accepted via CLI flags and TRACKER_* environment variables.
57/// Used when registering a run with the Sentinel API (Priority 4).
58#[derive(Debug, Clone, Default)]
59pub struct JobMetadata {
60    pub project_name: Option<String>,
61    pub job_name: Option<String>,
62    pub stage_name: Option<String>,
63    pub task_name: Option<String>,
64    pub team: Option<String>,
65    pub env: Option<String>,
66    pub language: Option<String>,
67    pub orchestrator: Option<String>,
68    pub executor: Option<String>,
69    pub external_run_id: Option<String>,
70    pub container_image: Option<String>,
71    /// Arbitrary key=value tags supplied via repeated --tag flags.
72    pub tags: Vec<String>,
73    /// Shell-wrapper command as a token list, e.g. ["stress", "--cpu", "4"].
74    /// Empty when not running in shell-wrapper mode.
75    pub command: Vec<String>,
76}
77
78// ---------------------------------------------------------------------------
79// CLI arguments (clap derive)
80// ---------------------------------------------------------------------------
81//
82#[derive(Debug, Parser)]
83#[command(
84    name = "resource-tracker",
85    about = "Lightweight Linux resource & GPU tracker.\n\n\
86             Shell-wrapper mode: resource-tracker [FLAGS] -- <command> [args...]\n\
87             The tracker will spawn <command>, monitor it, and exit when it exits.",
88    version
89)]
90struct Cli {
91    // -- Core flags ----------------------------------------------------------
92    /// Root PID of the process tree to track CPU usage for.
93    /// Overridden automatically in shell-wrapper mode.
94    #[arg(short = 'p', long, value_name = "PID")]
95    pid: Option<i32>,
96
97    /// Polling interval in seconds (must be >= 1).
98    #[arg(short = 'i', long, value_name = "SECS")]
99    interval: Option<u64>,
100
101    /// Nice value for the tracker process: -20 .. 19.
102    /// Bare --renice uses the default 19.
103    /// Increasing priority requires root privileges (see: man nice).
104    #[arg(
105        short = 'r',
106        long = "renice",
107        value_name = "VALUE",
108        env = "TRACKER_RENICE",
109        num_args = 0..=1,
110        default_missing_value = "19",
111        value_parser = clap::value_parser!(i32).range(RENICE_MIN..=RENICE_MAX),
112        verbatim_doc_comment,
113    )]
114    renice: Option<i32>,
115
116    /// Aggregate CPU steal value (default: true)
117    #[arg(
118        long = "aggregate-cpu-steal",
119        value_name = "AGGREGATE_CPU_STEAL",
120        env = "TRACKER_AGGREGATE_CPU_STEAL",
121        default_missing_value = "true"
122    )]
123    aggregate_cpu_steal: Option<bool>,
124
125    /// Path to TOML config file.
126    #[arg(short = 'c', long, value_name = "FILE", default_value = DEFAULT_CONFIG_FILE)]
127    config: String,
128
129    /// Output format: json (default) or csv.
130    #[arg(short = 'f', long, value_name = "FORMAT", default_value = "json")]
131    format: OutputFormat,
132
133    /// Write metric output to FILE instead of stdout.
134    /// Useful in shell-wrapper mode to keep the tracked app's stdout clean.
135    #[arg(short = 'o', long, value_name = "FILE", env = "TRACKER_OUTPUT")]
136    output: Option<String>,
137
138    /// Suppress metric output entirely (no stdout, no file).
139    /// Useful when streaming to Sentinel and local output is not needed.
140    #[arg(long, env = "TRACKER_QUIET")]
141    quiet: bool,
142
143    // -- Section 9.3 metadata flags ------------------------------------------
144    /// Project name for Sentinel run registration.
145    #[arg(long, value_name = "NAME", env = "TRACKER_PROJECT_NAME")]
146    project_name: Option<String>,
147
148    /// Job name attached to every sample and to the Sentinel run record.
149    #[arg(short = 'n', long, value_name = "NAME", env = "TRACKER_JOB_NAME")]
150    job_name: Option<String>,
151
152    /// Stage name (e.g. "train", "eval") for Sentinel run registration.
153    #[arg(long, value_name = "NAME", env = "TRACKER_STAGE_NAME")]
154    stage_name: Option<String>,
155
156    /// Task name for Sentinel run registration.
157    #[arg(long, value_name = "NAME", env = "TRACKER_TASK_NAME")]
158    task_name: Option<String>,
159
160    /// Team name for Sentinel run registration.
161    #[arg(long, value_name = "NAME", env = "TRACKER_TEAM")]
162    team: Option<String>,
163
164    /// Environment label (e.g. "prod", "staging") for Sentinel run registration.
165    #[arg(long, value_name = "ENV", env = "TRACKER_ENV")]
166    env: Option<String>,
167
168    /// Programming language label for Sentinel run registration.
169    #[arg(long, value_name = "LANG", env = "TRACKER_LANGUAGE")]
170    language: Option<String>,
171
172    /// Orchestrator label (e.g. "airflow", "prefect") for Sentinel run registration.
173    #[arg(long, value_name = "NAME", env = "TRACKER_ORCHESTRATOR")]
174    orchestrator: Option<String>,
175
176    /// Executor label (e.g. "kubernetes", "slurm") for Sentinel run registration.
177    #[arg(long, value_name = "NAME", env = "TRACKER_EXECUTOR")]
178    executor: Option<String>,
179
180    /// External run ID from the calling system for Sentinel run registration.
181    #[arg(long, value_name = "ID", env = "TRACKER_EXTERNAL_RUN_ID")]
182    external_run_id: Option<String>,
183
184    /// Container image name/tag for Sentinel run registration.
185    #[arg(long, value_name = "IMAGE", env = "TRACKER_CONTAINER_IMAGE")]
186    container_image: Option<String>,
187
188    /// Arbitrary key=value tag. May be repeated: --tag key1=val1 --tag key2=val2
189    #[arg(long = "tag", value_name = "KEY=VALUE", action = ArgAction::Append)]
190    tags: Vec<String>,
191
192    // -- Shell-wrapper mode --------------------------------------------------
193    /// Command to spawn and monitor. All tokens after -- are the command + args.
194    /// Example: resource-tracker -- Rscript model.R --epochs 10
195    #[arg(
196        trailing_var_arg = true,
197        allow_hyphen_values = true,
198        value_name = "COMMAND"
199    )]
200    command: Vec<String>,
201}
202
203// ---------------------------------------------------------------------------
204// Merged config
205// ---------------------------------------------------------------------------
206//
207/// Resolved configuration after merging CLI args > TOML file > defaults.
208#[derive(Debug, Clone)]
209pub struct Config {
210    /// Root PID for per-process CPU attribution. None = system-wide only.
211    /// Set automatically from the spawned child PID in shell-wrapper mode.
212    pub pid: Option<i32>,
213    /// Polling interval in seconds.
214    pub interval_secs: u64,
215    /// Nice value applied to the tracker process itself (0..=19).
216    pub renice: Option<i32>,
217    // CPU steal tracking mode
218    pub aggregate_cpu_steal: bool,
219    /// Output format (JSON or CSV).
220    pub format: OutputFormat,
221    /// Write metric output to this file path instead of stdout.
222    /// None = write to stdout.
223    pub output_file: Option<String>,
224    /// Suppress all metric output (no stdout, no file).
225    pub quiet: bool,
226    /// Section 9.3 job metadata (used for Sentinel API registration).
227    pub metadata: JobMetadata,
228    /// Shell-wrapper command. Empty = standalone mode.
229    pub command: Vec<String>,
230}
231
232impl Config {
233    /// Parse CLI args, optionally load the TOML config file, and merge with
234    /// defaults.  CLI flags always win; config file wins over defaults.
235    pub fn load() -> Self {
236        let cli = Cli::parse();
237
238        // Silently skip missing or unparseable config files.
239        let toml: TomlConfig = std::fs::read_to_string(&cli.config)
240            .ok()
241            .and_then(|s| toml::from_str(&s).ok())
242            .unwrap_or_default();
243
244        let interval_secs = cli
245            .interval
246            .or_else(|| toml.tracker.as_ref().and_then(|t| t.interval_secs))
247            .unwrap_or(DEFAULT_INTERVAL_SECS);
248
249        if interval_secs == 0 {
250            eprintln!("error: --interval must be >= 1 (got 0)");
251            std::process::exit(1);
252        }
253
254        let renice = cli
255            .renice
256            .or_else(|| toml.tracker.as_ref().and_then(|t| t.renice));
257
258        let aggregate_cpu_steal = cli
259            .aggregate_cpu_steal
260            .or_else(|| toml.tracker.as_ref().and_then(|t| t.aggregate_cpu_steal))
261            .unwrap_or(true);
262
263        let pid = cli.pid.or_else(|| toml.job.as_ref().and_then(|j| j.pid));
264
265        let metadata = JobMetadata {
266            project_name: cli.project_name,
267            job_name: cli
268                .job_name
269                .or_else(|| toml.job.as_ref().and_then(|j| j.name.clone())),
270            stage_name: cli.stage_name,
271            task_name: cli.task_name,
272            team: cli.team,
273            env: cli.env,
274            language: cli.language,
275            orchestrator: cli.orchestrator,
276            executor: cli.executor,
277            external_run_id: cli.external_run_id,
278            container_image: cli.container_image,
279            tags: cli.tags,
280            command: cli.command.clone(),
281        };
282
283        Config {
284            pid,
285            interval_secs,
286            renice,
287            aggregate_cpu_steal,
288            format: cli.format,
289            output_file: cli.output,
290            quiet: cli.quiet,
291            metadata,
292            command: cli.command,
293        }
294    }
295}
296
297// ---------------------------------------------------------------------------
298// Unit tests
299// ---------------------------------------------------------------------------
300
301#[cfg(test)]
302mod tests {
303    use super::*;
304
305    // T-CFG-01: TomlConfig deserializes from a valid TOML string.
306    #[test]
307    fn test_toml_config_deserializes() {
308        let toml_str = r#"
309[job]
310name = "benchmark"
311pid = 12345
312
313[tracker]
314interval_secs = 5
315"#;
316        let cfg: TomlConfig = toml::from_str(toml_str).expect("TOML parse failed");
317        let job = cfg.job.as_ref().expect("job section missing");
318        assert_eq!(job.name.as_deref(), Some("benchmark"));
319        assert_eq!(job.pid, Some(12345));
320        let tracker = cfg.tracker.as_ref().expect("tracker section missing");
321        assert_eq!(tracker.interval_secs, Some(5));
322    }
323
324    // T-CFG-02: TomlConfig defaults to None fields when the file is empty.
325    #[test]
326    fn test_toml_config_default_is_all_none() {
327        let cfg = TomlConfig::default();
328        assert!(cfg.job.is_none(), "job must be None in default TomlConfig");
329        assert!(
330            cfg.tracker.is_none(),
331            "tracker must be None in default TomlConfig"
332        );
333    }
334
335    // T-CFG-03: JobMetadata default produces all-None/empty fields.
336    #[test]
337    fn test_job_metadata_default_all_none() {
338        let m = JobMetadata::default();
339        assert!(m.project_name.is_none());
340        assert!(m.job_name.is_none());
341        assert!(m.stage_name.is_none());
342        assert!(m.task_name.is_none());
343        assert!(m.team.is_none());
344        assert!(m.env.is_none());
345        assert!(m.language.is_none());
346        assert!(m.orchestrator.is_none());
347        assert!(m.executor.is_none());
348        assert!(m.external_run_id.is_none());
349        assert!(m.container_image.is_none());
350        assert!(
351            m.tags.is_empty(),
352            "tags must be empty in default JobMetadata"
353        );
354    }
355
356    // T-CFG-04: OutputFormat variants compare correctly.
357    #[test]
358    fn test_output_format_equality() {
359        assert_eq!(OutputFormat::Json, OutputFormat::Json);
360        assert_eq!(OutputFormat::Csv, OutputFormat::Csv);
361        assert_ne!(OutputFormat::Json, OutputFormat::Csv);
362    }
363
364    // T-CFG-05: TomlConfig gracefully ignores unknown keys.
365    #[test]
366    fn test_toml_config_ignores_unknown_keys() {
367        let toml_str = r#"
368[job]
369name = "run1"
370unknown_field = "ignored"
371"#;
372        // Should not panic; unknown fields are silently dropped by serde.
373        let result: Result<TomlConfig, _> = toml::from_str(toml_str);
374        // toml crate returns error for unknown fields by default unless
375        // serde is configured to ignore them. If this fails, the test still
376        // documents the expected behavior.
377        let _ = result; // accept either Ok or Err
378    }
379}
resource_tracker/config.rs

resource_tracker/
config.rs
