From 6af4a9a799760a431bbb9edb4dde6be319ef25d1 Mon Sep 17 00:00:00 2001 From: Hoyt Koepke Date: Thu, 16 Jan 2025 12:00:45 -0800 Subject: [PATCH] Address PR feedback. --- hf_xet/Cargo.lock | 1 - hf_xet/Cargo.toml | 1 - hf_xet/src/profiling.rs | 74 ++++++++++++++++++++++++----------------- 3 files changed, 44 insertions(+), 32 deletions(-) diff --git a/hf_xet/Cargo.lock b/hf_xet/Cargo.lock index 935da637..18754afe 100644 --- a/hf_xet/Cargo.lock +++ b/hf_xet/Cargo.lock @@ -1023,7 +1023,6 @@ dependencies = [ "data", "error_printer", "lazy_static", - "once_cell", "parutils", "pprof", "pyo3", diff --git a/hf_xet/Cargo.toml b/hf_xet/Cargo.toml index 252e5513..5b06e927 100644 --- a/hf_xet/Cargo.toml +++ b/hf_xet/Cargo.toml @@ -32,7 +32,6 @@ reqwest = "0.11.27" serde = { version = "1.0.215", features = ["derive"] } serde_json = "1.0.133" lazy_static = "1.5" -once_cell = "1.17" chrono = "0.4" pprof = { version = "0.14", features = [ diff --git a/hf_xet/src/profiling.rs b/hf_xet/src/profiling.rs index f030a880..9801eb5b 100644 --- a/hf_xet/src/profiling.rs +++ b/hf_xet/src/profiling.rs @@ -1,3 +1,27 @@ +// Enables an easy way to profile xet-core operations. Internally, the code uses pprof-rs, which links in the +// GPerfTools suite. +// +// To use, compile with --features=profiling; with maturin, install with `maturin develop --release +// --features=profiling` +// +// When running, it saves the profile to profilers/, writing both a gperf protobuf file and a flamegraph svg +// file. The protobuf file can be read using the pprof tool: +// +// ```bash +// go install github.com/google/pprof@latest # Installs to ~/go/bin/ +// # Replace this with the correct path to the hf_xet library file. +// ~/go/bin/pprof ./venv/lib/python3.10/site-packages/hf_xet/hf_xet.abi3.so profiles//pprof.pb +// ``` +// +// On OSX, I like using qcachegrind as the interactive visualizer of the results, which requires the profile data to be +// exported to the callgrind format: +// +// ```bash +// brew install qcachegrind # If needed +// ~/go/bin/pprof -callgrind ./venv/lib/python3.10/site-packages/hf_xet/hf_xet.abi3.so profiles//pprof.pb +// qcachegrind +// ``` + use std::fs; use std::path::PathBuf; @@ -5,28 +29,29 @@ use chrono::Local; use pprof::protos::Message; use pprof::{ProfilerGuard, ProfilerGuardBuilder}; -/// A simple pprof-rs integration assuming only one function type is active at once. -/// We store a single global session (with a single call_name). All callers -/// share that session if it exists, or create a new one if not. Once the -/// last handle is dropped, the profiler stops and results are saved. - -/// A global reference to the **current** session (if any). -/// Since we only allow one active session name, this is just an `Option`. const SAMPLING_FREQUENCY: i32 = 100; // 100 Hz +// A global reference to the current profiling session. The python at_exit function dumps this out. lazy_static::lazy_static! { static ref CURRENT_SESSION: ProfilingSession<'static> = ProfilingSession::new(); } -/// Holds the actual ProfilerGuard plus the associated call_name. -/// When this struct is dropped, we build the report and dump it to disk. struct ProfilingSession<'a> { guard: Option>, } -/// In `Drop`, we stop the profiler by building a report, then save -/// the flamegraph & a protobuf (optional) to `profiles///`. impl ProfilingSession<'_> { + /// Create a new SessionState by starting a `ProfilerGuard`. + fn new() -> Self { + let guard = ProfilerGuardBuilder::default() + .frequency(SAMPLING_FREQUENCY) + .build() + .inspect_err(|e| eprintln!("Profiler: Warning: Failed to start profiler: {e:?}")) + .ok(); + + ProfilingSession { guard } + } + fn save_report(&self) { let Some(guard) = &self.guard else { return; @@ -36,7 +61,7 @@ impl ProfilingSession<'_> { let report = match guard.report().build() { Ok(r) => r, Err(e) => { - eprintln!("Error building profiler report: {e:?}"); + eprintln!("Profiler: Error building profiler report: {e:?}"); return; }, }; @@ -45,7 +70,7 @@ impl ProfilingSession<'_> { let output_dir = PathBuf::from("profiles").join(date_str); let Ok(_) = fs::create_dir_all(&output_dir) - .inspect_err(|e| eprintln!("Error creating profile output directory: {e:?}")) + .inspect_err(|e| eprintln!("Profiler: Error creating profile output directory: {e:?}")) else { return; }; @@ -53,39 +78,28 @@ impl ProfilingSession<'_> { // Write flamegraph.svg let flame_path = output_dir.join("flamegraph.svg"); if let Ok(mut fg_file) = - fs::File::create(&flame_path).inspect_err(|e| eprintln!("Error writing flamegraph file: {e:?}")) + fs::File::create(&flame_path).inspect_err(|e| eprintln!("Profiler: Error writing flamegraph file: {e:?}")) { let _ = report.flamegraph(&mut fg_file).inspect_err(|e| { - eprintln!("Failed writing flamegraph: {e:?}"); + eprintln!("Profiler: Failed writing flamegraph: {e:?}"); }); } let pb_path = output_dir.join("pprof.pb"); if let Ok(mut pb_file) = - fs::File::create(pb_path).inspect_err(|e| eprintln!("Failed opening pperf out file: {e:?}")) + fs::File::create(pb_path).inspect_err(|e| eprintln!("Profiler: Failed opening pperf out file: {e:?}")) { if let Ok(profile) = report .pprof() - .inspect_err(|e| eprintln!("Error creating pprof profile report: {e:?}")) + .inspect_err(|e| eprintln!("Profiler: Error creating pprof profile report: {e:?}")) { let _ = profile.write_to_writer(&mut pb_file).inspect_err(|e| { - eprintln!("Failed writing protobuf perf out file: {e:?}"); + eprintln!("Profiler: Failed writing protobuf gperf out file: {e:?}"); }); } } - eprintln!("Saved run profile to {output_dir:?}"); - } - - /// Create a new SessionState by starting a `ProfilerGuard`. - fn new() -> Self { - let guard = ProfilerGuardBuilder::default() - .frequency(SAMPLING_FREQUENCY) - .build() - .inspect_err(|e| eprintln!("Warning: Failed to start profiler: {e:?}")) - .ok(); - - ProfilingSession { guard } + eprintln!("Profiler: Saved run profile to {output_dir:?}"); } }