Skip to content

Commit

Permalink
Fix #228 - Abstract the output format and clean up
Browse files Browse the repository at this point in the history
  • Loading branch information
Lars T Hansen committed Jan 13, 2025
1 parent 9c9ced2 commit fcb4bcc
Show file tree
Hide file tree
Showing 5 changed files with 640 additions and 282 deletions.
89 changes: 76 additions & 13 deletions src/main.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ mod interrupt;
mod jobs;
mod log;
mod nvidia;
mod output;
mod procfs;
mod procfsapi;
mod ps;
Expand All @@ -17,6 +18,8 @@ mod time;
mod users;
mod util;

use std::io;

const TIMEOUT_SECONDS: u64 = 5; // For subprocesses
const USAGE_ERROR: i32 = 2; // clap, Python, Go

Expand Down Expand Up @@ -57,9 +60,15 @@ enum Commands {
/// One output record per Sonar invocation will contain a load= field with an encoding of
/// the per-cpu usage since boot.
load: bool,

/// Output JSON, not CSV
json: bool,
},
/// Extract system information
Sysinfo {},
Sysinfo {
/// Output CSV, not JSON
csv: bool,
},
/// Extract slurm job information
Slurmjobs {
/// Set the sacct start time to now-`window` and the end time to now, and dump records that
Expand All @@ -70,6 +79,9 @@ enum Commands {
/// From-to dates on the form yyyy-mm-dd,yyyy-mm-dd (with the comma); from is inclusive,
/// to is exclusive. Precludes -window.
span: Option<String>,

/// Output json, not CSV
json: bool,
},
Version {},
}
Expand All @@ -83,6 +95,7 @@ fn main() {

log::init();

let mut writer = io::stdout();
match &command_line() {
Commands::PS {
rollup,
Expand All @@ -95,6 +108,7 @@ fn main() {
exclude_commands,
lockdir,
load,
json,
} => {
let opts = ps::PsOptions {
rollup: *rollup,
Expand All @@ -115,30 +129,33 @@ fn main() {
vec![]
},
lockdir: lockdir.clone(),
json: *json,
};
if *batchless {
let mut jm = batchless::BatchlessJobManager::new();
ps::create_snapshot(&mut jm, &opts, &timestamp);
ps::create_snapshot(&mut writer, &mut jm, &opts, &timestamp);
} else {
let mut jm = slurm::SlurmJobManager {};
ps::create_snapshot(&mut jm, &opts, &timestamp);
ps::create_snapshot(&mut writer, &mut jm, &opts, &timestamp);
}
}
Commands::Sysinfo {} => {
sysinfo::show_system(&timestamp);
Commands::Sysinfo { csv } => {
sysinfo::show_system(&mut writer, &timestamp, *csv);
}
Commands::Slurmjobs { window, span } => {
slurmjobs::show_slurm_jobs(window, span);
Commands::Slurmjobs { window, span, json } => {
slurmjobs::show_slurm_jobs(&mut writer, window, span, *json);
}
Commands::Version {} => {
show_version(&mut std::io::stdout());
show_version(&mut writer);
}
}
//let _ = writer.flush();
}

// For the sake of simplicity:
// - allow repeated options to overwrite earlier values
// - all error reporting is via a generic "usage" message, without specificity as to what was wrong
// - both --json and --csv are accepted to all commands

fn command_line() -> Commands {
let args = std::env::args().collect::<Vec<String>>();
Expand All @@ -158,6 +175,8 @@ fn command_line() -> Commands {
let mut exclude_commands = None;
let mut lockdir = None;
let mut load = false;
let mut json = false;
let mut csv = false;
while next < args.len() {
let arg = args[next].as_ref();
next += 1;
Expand All @@ -167,6 +186,10 @@ fn command_line() -> Commands {
(next, rollup) = (new_next, true);
} else if let Some(new_next) = bool_arg(arg, &args, next, "--load") {
(next, load) = (new_next, true);
} else if let Some(new_next) = bool_arg(arg, &args, next, "--json") {
(next, json) = (new_next, true);
} else if let Some(new_next) = bool_arg(arg, &args, next, "--csv") {
(next, csv) = (new_next, true);
} else if let Some(new_next) =
bool_arg(arg, &args, next, "--exclude-system-jobs")
{
Expand Down Expand Up @@ -210,6 +233,10 @@ fn command_line() -> Commands {
eprintln!("--rollup and --batchless are incompatible");
std::process::exit(USAGE_ERROR);
}
if json && csv {
eprintln!("--csv and --json are incompatible");
std::process::exit(USAGE_ERROR);
}

Commands::PS {
batchless,
Expand All @@ -222,12 +249,34 @@ fn command_line() -> Commands {
exclude_commands,
lockdir,
load,
json,
}
}
"sysinfo" => {
let mut json = false;
let mut csv = false;
while next < args.len() {
let arg = args[next].as_ref();
next += 1;
if let Some(new_next) = bool_arg(arg, &args, next, "--json") {
(next, json) = (new_next, true);
} else if let Some(new_next) = bool_arg(arg, &args, next, "--csv") {
(next, csv) = (new_next, true);
} else {
usage(true);
}
}
if json && csv {
eprintln!("--csv and --json are incompatible");
std::process::exit(USAGE_ERROR);
}
Commands::Sysinfo { csv }
}
"sysinfo" => Commands::Sysinfo {},
"slurm" => {
let mut window = None;
let mut span = None;
let mut json = false;
let mut csv = false;
while next < args.len() {
let arg = args[next].as_ref();
next += 1;
Expand All @@ -237,14 +286,22 @@ fn command_line() -> Commands {
(next, window) = (new_next, Some(value));
} else if let Some((new_next, value)) = string_arg(arg, &args, next, "--span") {
(next, span) = (new_next, Some(value));
} else if let Some(new_next) = bool_arg(arg, &args, next, "--json") {
(next, json) = (new_next, true);
} else if let Some(new_next) = bool_arg(arg, &args, next, "--csv") {
(next, csv) = (new_next, true);
} else {
usage(true);
}
}
if window.is_some() && span.is_some() {
usage(true);
}
Commands::Slurmjobs { window, span }
if json && csv {
eprintln!("--csv and --json are incompatible");
std::process::exit(USAGE_ERROR);
}
Commands::Slurmjobs { window, span, json }
}
"version" => Commands::Version {},
"help" => {
Expand Down Expand Up @@ -316,9 +373,9 @@ fn usage(is_error: bool) -> ! {
Usage: sonar <COMMAND>
Commands:
ps Take a snapshot of the currently running processes
sysinfo Extract system information
slurm Extract slurm job information for a [start,end) time interval
ps Print process and load information
sysinfo Print system information
slurm Print slurm job information for a [start,end) time interval
help Print this message
Options for `ps`:
Expand All @@ -345,6 +402,10 @@ Options for `ps`:
--lockdir directory
Create a per-host lockfile in this directory and exit early if the file
exists on startup [default: none]
--load
Print per-cpu and per-gpu load data
--json
Format output as JSON, not CSV
Options for `slurm`:
--window minutes
Expand All @@ -353,6 +414,8 @@ Options for `slurm`:
--span start,end
Both `start` and `end` are on the form yyyy-mm-dd. Mostly useful for seeding a
database with older data. Precludes --window
--json
Format output as JSON, not CSV
",
);
let _ = out.flush();
Expand Down
Loading

0 comments on commit fcb4bcc

Please sign in to comment.