Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Rust code looks pretty good - minor nits #1

Merged
merged 1 commit into from
Feb 29, 2024
Merged

Rust code looks pretty good - minor nits #1

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions Cargo.lock
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ license = "BSD-3-Clause"
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html

[dependencies]
anyhow = "1.0.80"
base16ct = "0.2.0"
clap = { version = "4.5.0", features = ["derive"] }
fxhash = "0.2.1"
Expand Down
134 changes: 70 additions & 64 deletions src/main.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,10 +1,12 @@
use anyhow::anyhow;
use base16ct;
use clap::Parser;
use core::hash::BuildHasherDefault;
use fxhash::{FxHashMap, FxHasher};
use html_escape::encode_text;
use indexmap::IndexMap;
use md5::{Digest, Md5};
use std::ffi::{OsStr, OsString};

use regex::Regex;
use std::fmt::{self, Display, Formatter};
Expand All @@ -18,7 +20,6 @@ use tinytemplate::TinyTemplate;
use indicatif::{MultiProgress, ProgressBar, ProgressStyle};
use once_cell::sync::Lazy;
use serde::{Deserialize, Serialize};
use std::process::ExitCode;
use std::sync::Mutex;
use std::time::Instant;

Expand Down Expand Up @@ -64,8 +65,8 @@ static TEMPLATE_INDEX: &str = r#"
struct Cli {
path: PathBuf,
/// Output directory, defaults to `tl_out`
#[arg(short)]
out: Option<PathBuf>,
#[arg(short, default_value = "tl_out")]
out: PathBuf,
/// Delete out directory if it already exists
#[arg(long)]
overwrite: bool,
Expand Down Expand Up @@ -217,7 +218,7 @@ struct EmptyMetadata {}

#[derive(Debug, Deserialize)]
struct DynamoOutputGraphMetadata {
sizes: Option<FxHashMap<String, Vec<SymInt>>>,
_sizes: Option<FxHashMap<String, Vec<SymInt>>>,
Copy link
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Does this change serde's behavior? I need the field to match what I am getting in the JSON format

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It will change the field name - so if you need it to match the JSON name then you can prefix it with:

#[allow(dead_code)]
sizes: Option<FxHashMap<String, Vec<SymInt>>>,

so it doesn't complain about the unread field.

Copy link
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ok, i'll need to do this. fortunately these are all defaulted so it doesn't matter too much yet

}

#[derive(Debug, Deserialize)]
Expand All @@ -227,7 +228,7 @@ struct DynamoStartMetadata {

#[derive(Debug, Deserialize)]
struct InductorOutputCodeMetadata {
filename: Option<String>,
filename: Option<PathBuf>,
}

#[derive(Debug, Deserialize)]
Expand All @@ -244,7 +245,7 @@ struct Envelope {
optimize_ddp_split_graph: Option<EmptyMetadata>,
optimize_ddp_split_child: Option<OptimizeDdpSplitChildMetadata>,
compiled_autograd_graph: Option<EmptyMetadata>,
dynamo_guards: Option<EmptyMetadata>,
_dynamo_guards: Option<EmptyMetadata>,
aot_forward_graph: Option<EmptyMetadata>,
aot_backward_graph: Option<EmptyMetadata>,
aot_joint_graph: Option<EmptyMetadata>,
Expand All @@ -259,10 +260,10 @@ struct IndexContext {
stack_trie_html: String,
}

fn main() -> ExitCode {
fn main() -> anyhow::Result<()> {
let cli = Cli::parse();
let path = cli.path;
let out_path = cli.out.unwrap_or(PathBuf::from("tl_out"));
let out_path = cli.out;

if out_path.exists() {
if !cli.overwrite {
Expand All @@ -271,17 +272,17 @@ fn main() -> ExitCode {
out_path.display()
);
}
fs::remove_dir_all(&out_path).unwrap();
fs::remove_dir_all(&out_path)?;
}
fs::create_dir(&out_path).unwrap();
fs::create_dir(&out_path)?;

let file = File::open(path).unwrap();
let metadata = file.metadata().unwrap();
let file = File::open(path)?;
let metadata = file.metadata()?;
let file_size = metadata.len();
let multi = MultiProgress::new();
let pb = multi.add(ProgressBar::new(file_size));
pb.set_style(ProgressStyle::default_bar()
.template("{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} [{bytes_per_sec}] ({eta})").unwrap()
.template("{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} [{bytes_per_sec}] ({eta})")?
.progress_chars("#>-"));
let spinner = multi.add(ProgressBar::new_spinner());
let reader = io::BufReader::new(file);
Expand All @@ -292,8 +293,7 @@ fn main() -> ExitCode {
r"(?<thread>\d+)",
r"(?<pathname>[^:]+):(?<line>\d+)\] ",
r"(?<payload>.)"
))
.unwrap();
))?;

let mut stack_trie = StackTrieNode::default();

Expand Down Expand Up @@ -377,21 +377,26 @@ fn main() -> ExitCode {
}
};

// TODO: borrow only here
let compile_id_dir = e
let compile_id_dir: PathBuf = e
.compile_id
.clone()
.map_or(format!("unknown_{}", lineno), |e: CompileId| {
format!("{}_{}_{}", e.frame_id, e.frame_compile_id, e.attempt)
});
.as_ref()
Copy link
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ohhhh lol

.map_or(
format!("unknown_{lineno}"),
|CompileId {
frame_id,
frame_compile_id,
attempt,
}| { format!("{frame_id}_{frame_compile_id}_{attempt}") },
Copy link
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

huh, f-string style is nice

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just be aware that it only works with identifiers - so you can't do:

format!("{foo.bar}")

but you can do:

format!("{myvar}", myvar = foo.bar)

)
.into();

let subdir = out_path.join(&compile_id_dir);
fs::create_dir_all(&subdir).unwrap();
fs::create_dir_all(&subdir)?;

let mut payload = String::new();
if let Some(expect) = e.has_payload {
let mut first = true;
while let Some((payload_lineno, payload_line)) =
while let Some((_payload_lineno, payload_line)) =
iter.next_if(|(_, l)| l.starts_with('\t'))
{
// Careful! Distinguish between missing EOL and not
Expand Down Expand Up @@ -429,53 +434,55 @@ fn main() -> ExitCode {
};
};

let mut write_dump = |filename: &str, sentinel: Option<EmptyMetadata>| {
if let Some(_r) = sentinel {
let f = subdir.join(filename);
fs::write(&f, &payload).unwrap();
compile_directory.push(Path::new(&compile_id_dir).join(filename));
}
};
let mut write_dump =
|filename: &str, sentinel: Option<EmptyMetadata>| -> anyhow::Result<()> {
if sentinel.is_some() {
let f = subdir.join(filename);
fs::write(&f, &payload)?;
compile_directory.push(compile_id_dir.join(filename));
}
Ok(())
};

write_dump("optimize_ddp_split_graph.txt", e.optimize_ddp_split_graph);
write_dump("compiled_autograd_graph.txt", e.compiled_autograd_graph);
write_dump("aot_forward_graph.txt", e.aot_forward_graph);
write_dump("aot_backward_graph.txt", e.aot_backward_graph);
write_dump("aot_joint_graph.txt", e.aot_joint_graph);
write_dump("inductor_post_grad_graph.txt", e.inductor_post_grad_graph);
write_dump("optimize_ddp_split_graph.txt", e.optimize_ddp_split_graph)?;
write_dump("compiled_autograd_graph.txt", e.compiled_autograd_graph)?;
write_dump("aot_forward_graph.txt", e.aot_forward_graph)?;
write_dump("aot_backward_graph.txt", e.aot_backward_graph)?;
write_dump("aot_joint_graph.txt", e.aot_joint_graph)?;
write_dump("inductor_post_grad_graph.txt", e.inductor_post_grad_graph)?;

if let Some(_metadata) = e.dynamo_output_graph {
if e.dynamo_output_graph.is_some() {
// TODO: dump sizes
let filename = "dynamo_output_graph.txt";
let f = subdir.join(&filename);
fs::write(&f, &payload).unwrap();
compile_directory.push(Path::new(&compile_id_dir).join(filename));
fs::write(&f, &payload)?;
compile_directory.push(compile_id_dir.join(filename));
}

if let Some(metadata) = e.inductor_output_code {
let filename = match metadata.filename {
Some(p) =>
// Bah, where's pattern guards when you need 'em
{
match Path::new(&p).file_stem() {
Some(stem) => {
format!("inductor_output_code_{}.txt", stem.to_str().unwrap())
}
None => "inductor_output_code.txt".to_string(),
}
}
None => "inductor_output_code.txt".to_string(),
};
let filename = metadata
.filename
.as_ref()
.and_then(|p| Path::file_stem(p))
.map_or_else(
|| PathBuf::from("inductor_output_code.txt"),
|stem| {
let mut r = OsString::from("inductor_output_code_");
r.push(stem);
r.push(OsStr::new(".txt"));
r.into()
},
);
let f = subdir.join(&filename);
fs::write(&f, &payload).unwrap();
compile_directory.push(Path::new(&compile_id_dir).join(filename));
fs::write(&f, &payload)?;
compile_directory.push(compile_id_dir.join(filename));
}

if let Some(metadata) = e.optimize_ddp_split_child {
let filename = format!("optimize_ddp_split_child_{}.txt", metadata.name);
let f = subdir.join(&filename);
fs::write(&f, &payload).unwrap();
compile_directory.push(Path::new(&compile_id_dir).join(filename));
fs::write(&f, &payload)?;
compile_directory.push(compile_id_dir.join(filename));
}
}
pb.finish_with_message("done");
Expand All @@ -485,7 +492,7 @@ fn main() -> ExitCode {

let mut tt = TinyTemplate::new();
tt.add_formatter("format_unescaped", tinytemplate::format_unescaped);
tt.add_template("index.html", TEMPLATE_INDEX).unwrap();
tt.add_template("index.html", TEMPLATE_INDEX)?;
let index_context = IndexContext {
css: CSS,
directory: directory
Expand All @@ -496,12 +503,11 @@ fn main() -> ExitCode {
};
fs::write(
out_path.join("index.html"),
tt.render("index.html", &index_context).unwrap(),
)
.unwrap();
tt.render("index.html", &index_context)?,
)?;

if !cli.no_browser {
opener::open(out_path.join("index.html")).unwrap();
opener::open(out_path.join("index.html"))?;
}

// other_rank is included here because you should only have logs from one rank when
Expand All @@ -510,8 +516,8 @@ fn main() -> ExitCode {
&& (stats.fail_glog + stats.fail_json + stats.fail_payload_md5 + stats.other_rank > 0)
{
// Report something went wrong
ExitCode::from(1)
} else {
ExitCode::from(0)
return Err(anyhow!("Something went wrong"));
}

Ok(())
}
Loading