use std::collections::BTreeSet;
use std::ffi::OsString;
use std::io;
#[cfg(feature = "client")]
use std::path::Path;
use std::path::PathBuf;
use clap::{Args, Parser, Subcommand};
use rayon::prelude::*;
#[cfg(feature = "client")]
use url::Url;
#[cfg(feature = "client")]
use legion_prof_viewer::{
app, deferred_data::DeferredDataSource, file_data::FileDataSource,
http::client::HTTPClientDataSource, parallel_data::ParallelDeferredDataSource,
};
#[cfg(feature = "archiver")]
use legion_prof::backend::archiver;
#[cfg(feature = "duckdb")]
use legion_prof::backend::duckdb;
#[cfg(feature = "nvtxw")]
use legion_prof::backend::nvtxw;
#[cfg(feature = "server")]
use legion_prof::backend::server;
#[cfg(feature = "viewer")]
use legion_prof::backend::viewer;
use legion_prof::backend::{analyze, dump, trace_viewer, visualize};
use legion_prof::serialize::deserialize;
use legion_prof::state::{Config, NodeID, State, Timestamp};
#[derive(Debug, Clone, Args)]
struct ParserArgs {
#[arg(required = true, help = "input Legion Prof log filenames")]
filenames: Vec<OsString>,
#[arg(long, help = "start time in microseconds to trim the profile")]
start_trim: Option<u64>,
#[arg(long, help = "stop time in microseconds to trim the profile")]
stop_trim: Option<u64>,
#[arg(
long,
default_value_t = 1000.0,
help = "threshold for warning about message latencies in microseconds"
)]
message_threshold: f64,
#[arg(
long,
default_value_t = 5.0,
help = "perentage of messages that must be over the threshold to trigger a warning"
)]
message_percentage: f64,
#[arg(
long,
default_value_t = 0,
help = "minimum threshold (in microseconds) for visualizing function calls"
)]
call_threshold: u64,
#[arg(long, help = "a list of nodes that will be visualized")]
nodes: Option<String>,
#[arg(
long,
hide = true,
help = "parse all log files, even when a subset of nodes are being shown (uses more memory)"
)]
no_filter_input: bool,
#[arg(short, long, help = "print verbose profiling information")]
verbose: bool,
}
#[derive(Debug, Clone, Args)]
struct OutputArgs {
#[arg(
short,
long,
default_value = "legion_prof",
help = "output directory pathname"
)]
output: OsString,
#[arg(short, long, help = "overwrite output directory if it exists")]
force: bool,
}
#[derive(Debug, Clone, Subcommand)]
enum Commands {
#[command(about = "save an archive of the profile for sharing")]
Archive {
#[command(flatten)]
args: ParserArgs,
#[command(flatten)]
out: OutputArgs,
#[arg(long, default_value_t = 4, help = "number of zoom levels to archive")]
levels: u32,
#[arg(long, default_value_t = 4, help = "branch factor for archive")]
branch_factor: u64,
#[arg(
long,
default_value_t = 1000,
help = "stop subdividing tiles at or below this size"
)]
min_tile_size: u64,
#[arg(
long,
default_value_t = 10,
help = "zstd compression factor for archive"
)]
zstd_compression: i32,
},
#[command(about = "connect viewer to an HTTP server or (local/remote) archive")]
Attach {
#[arg(required = true, help = "URL(s) or path(s) to attach to")]
args: Vec<OsString>,
},
#[command(name = "duckdb", about = "save profile to DuckDB database")]
DuckDB {
#[command(flatten)]
args: ParserArgs,
#[command(flatten)]
out: OutputArgs,
},
#[command(about = "dump parsed log files in a JSON format")]
Dump {
#[command(flatten)]
args: ParserArgs,
},
#[command(about = "dump a legacy format profile for sharing")]
Legacy {
#[command(flatten)]
args: ParserArgs,
#[command(flatten)]
out: OutputArgs,
},
#[command(about = "process data through NVTXW for NVIDIA Nsight Systems")]
NVTXW {
#[command(flatten)]
args: ParserArgs,
#[arg(long, help = "path to NVTXW backend implementation")]
backend: Option<OsString>,
#[arg(long, help = "output nsys-rep filename")]
output: OsString,
#[arg(short, long, help = "overwrite output file if it exists")]
force: bool,
#[arg(long, help = "input nsys-rep filename to merge with Legion Prof data")]
merge: Option<OsString>,
},
#[command(about = "start profile HTTP server")]
Serve {
#[command(flatten)]
args: ParserArgs,
#[arg(
long,
default_value = "127.0.0.1",
help = "host to bind for HTTP server"
)]
host: String,
#[arg(long, default_value_t = 8080, help = "port to bind for HTTP server")]
port: u16,
},
#[command(about = "start interactive profile viewer")]
View {
#[command(flatten)]
args: ParserArgs,
},
#[command(about = "print statistics")]
Statistics {
#[command(flatten)]
args: ParserArgs,
},
#[command(about = "emit JSON for Google Trace Viewer")]
Trace {
#[command(flatten)]
args: ParserArgs,
#[command(flatten)]
out: OutputArgs,
},
}
#[derive(Debug, Clone, Parser)]
#[command(author, version, about, long_about = None)]
struct Cli {
#[command(subcommand)]
command: Commands,
}
fn main() -> io::Result<()> {
let env = env_logger::Env::default().filter_or("RUST_LOG", "info");
env_logger::init_from_env(env);
let cli = Cli::parse();
match cli.command {
Commands::Archive { .. } => {
#[cfg(not(feature = "archiver"))]
panic!(
"Legion Prof was not built with the \"archiver\" feature. \
Rebuild with --features=archiver to enable."
);
}
Commands::Attach { .. } => {
#[cfg(not(feature = "client"))]
panic!(
"Legion Prof was not built with the \"client\" feature. \
Rebuild with --features=client to enable."
);
}
Commands::DuckDB { .. } => {
#[cfg(not(feature = "duckdb"))]
panic!(
"Legion Prof was not built with the \"duckdb\" feature. \
Rebuild with --features=duckdb to enable."
);
}
Commands::NVTXW { .. } => {
#[cfg(not(feature = "nvtxw"))]
panic!(
"Legion Prof was not built with the \"nvtxw\" feature. \
Rebuild with --features=nvtxw to enable."
);
}
Commands::Serve { .. } => {
#[cfg(not(feature = "server"))]
panic!(
"Legion Prof was not built with the \"server\" feature. \
Rebuild with --features=server to enable."
);
}
Commands::View { .. } => {
#[cfg(not(feature = "viewer"))]
panic!(
"Legion Prof was not built with the \"viewer\" feature. \
Rebuild with --features=viewer to enable."
);
}
_ => {}
}
match cli.command {
Commands::Attach { args } => {
#[cfg(feature = "client")]
{
fn http_ds(url: Url) -> Box<dyn DeferredDataSource> {
Box::new(HTTPClientDataSource::new(url))
}
fn file_ds(path: impl AsRef<Path>) -> Box<dyn DeferredDataSource> {
Box::new(ParallelDeferredDataSource::new(FileDataSource::new(path)))
}
let data_sources: Vec<_> = args
.into_iter()
.map(|arg| {
arg.into_string()
.map(|s| Url::parse(&s).map(http_ds).unwrap_or_else(|_| {
println!("The argument '{}' does not appear to be a valid URL. Attempting to open it as a local file...", &s);
file_ds(&s)
}))
.unwrap_or_else(file_ds)
})
.collect();
app::start(data_sources);
}
return Ok(());
}
_ => {}
}
let args = match cli.command {
Commands::Archive { ref args, .. }
| Commands::Dump { ref args, .. }
| Commands::Legacy { ref args, .. }
| Commands::DuckDB { ref args, .. }
| Commands::NVTXW { ref args, .. }
| Commands::View { ref args, .. }
| Commands::Serve { ref args, .. }
| Commands::Statistics { ref args, .. }
| Commands::Trace { ref args, .. } => args,
Commands::Attach { .. } => unreachable!(),
};
let start_trim = args.start_trim.map(Timestamp::from_us);
let stop_trim = args.stop_trim.map(Timestamp::from_us);
let message_threshold = args.message_threshold;
let message_percentage = args.message_percentage;
let mut node_list: Vec<NodeID> = Vec::new();
let mut filter_input = false;
if let Some(nodes_str) = &args.nodes {
node_list = nodes_str
.split(",")
.flat_map(|x| {
let splits: Vec<_> = x
.splitn(2, "-")
.map(|x| x.parse::<u64>().unwrap())
.collect();
if splits.len() == 2 {
(splits[0]..=splits[1]).into_iter().map(NodeID)
} else {
(splits[0]..=splits[0]).into_iter().map(NodeID)
}
})
.collect();
filter_input = !args.no_filter_input;
}
let records: Result<Vec<_>, _> = args
.filenames
.par_iter()
.map(|filename| {
println!("Reading log file {:?}...", filename);
deserialize(filename, &node_list, filter_input)
})
.collect();
match cli.command {
Commands::Dump { .. } => {
for record in records? {
dump::dump_record(&record)?;
}
return Ok(());
}
_ => {}
}
let mut state = State::default();
let paths: Vec<_> = args.filenames.iter().map(PathBuf::from).collect();
let mut unique_paths = BTreeSet::<String>::new();
for p in paths {
if let Some(base) = p.parent() {
unique_paths.insert(base.to_string_lossy().to_string());
}
}
state.source_locator.extend(unique_paths.into_iter());
state.visible_nodes = node_list;
if filter_input {
println!("Filtering profiles to nodes: {:?}", state.visible_nodes);
}
for record in records? {
println!("Matched {} objects", record.len());
state.process_records(&record, Timestamp::from_us(args.call_threshold));
}
if !state.complete_parse() {
println!("Nothing to do");
return Ok(());
}
let mut have_alllogs = true;
let num_nodes: usize = state.num_nodes.try_into().unwrap();
if num_nodes > args.filenames.len() {
println!(
"Warning: This run involved {:?} nodes, but only {:?} log files were provided. If --verbose is enabled, subsequent warnings may not indicate a true error.",
num_nodes,
args.filenames.len()
);
have_alllogs = false;
}
if state.visible_nodes.len() < num_nodes && filter_input {
println!(
"Warning: This run involved {:?} nodes, but only {:?} log files were used. If --verbose ie enabled, subsequent warnings may not indicate a true error.",
num_nodes,
state.visible_nodes.len()
);
have_alllogs = false;
}
Config::set_config(filter_input, args.verbose, have_alllogs);
state.trim_time_range(start_trim, stop_trim);
println!("Sorting time ranges");
state.sort_time_range();
state.check_message_latencies(message_threshold, message_percentage);
state.filter_output();
println!("Calculating critical paths");
state.compute_critical_paths();
match cli.command {
Commands::Archive {
out,
levels,
branch_factor,
min_tile_size,
zstd_compression,
..
} => {
#[cfg(feature = "archiver")]
{
state.stack_time_points();
state.assign_colors();
archiver::write(
state,
levels,
branch_factor,
min_tile_size,
out.output,
out.force,
zstd_compression,
)?;
}
}
Commands::DuckDB { out, .. } => {
#[cfg(feature = "duckdb")]
{
state.stack_time_points();
state.assign_colors();
duckdb::write(state, out.output, out.force)?;
}
}
Commands::Legacy { out, .. } => {
state.assign_colors();
visualize::emit_interactive_visualization(&state, out.output, out.force)?;
}
Commands::NVTXW {
backend,
output,
force,
merge,
..
} => {
#[cfg(feature = "nvtxw")]
{
state.stack_time_points();
state.assign_colors();
let zero_time = state.zero_time;
nvtxw::write(state, backend, output, force, merge, zero_time)?;
}
}
Commands::View { .. } => {
#[cfg(feature = "viewer")]
{
state.stack_time_points();
state.assign_colors();
viewer::start(state);
}
}
Commands::Serve { host, port, .. } => {
#[cfg(feature = "server")]
{
state.stack_time_points();
state.assign_colors();
server::start(state, &host, port);
}
}
Commands::Statistics { .. } => {
analyze::analyze_statistics(&state);
}
Commands::Trace { out, .. } => {
trace_viewer::emit_trace(&state, out.output, out.force)?;
}
Commands::Attach { .. } | Commands::Dump { .. } => unreachable!(),
}
Ok(())
}