diff --git a/services/node/docs/PROFILING.md b/services/node/docs/PROFILING.md new file mode 100644 index 0000000..644e332 --- /dev/null +++ b/services/node/docs/PROFILING.md @@ -0,0 +1,78 @@ +# MPC Node Session Profiling (#244) + +On-demand CPU/memory profiling per proof-generation session phase. + +## Why not `pprof`? + +The original issue asks for output "as pprof or flamegraph-compatible +format." A sampling profiler like the `pprof` crate walks *this node +process's own* call stack — but the actual MPC compute work for a session +runs in `co-noir` child processes (`session::run_proof_generation` spawns +one per phase: `merge_shares`, `witness_generation`, `proof_generation`). +An in-process profiler attached to the node would show almost nothing — +just the time spent spawning the child and waiting on it — because the +expensive work never executes on this process's stack at all. + +Instead, profiling here samples each `co-noir` child process's OS-reported +CPU% and memory (RSS) on a fixed 200ms interval for the duration of its +phase, and aggregates that into a `PhaseProfile` per phase. This is +exported as JSON, not the pprof wire format: pprof's call-graph model +doesn't apply to an opaque external process whose internals this node has +no visibility into (a flamegraph needs sampled *stack traces*, which +`co-noir` doesn't expose to callers). + +## API + +Profiling is strictly **opt-in per session** — a session nobody asks to +profile pays zero sampling overhead beyond one registry lookup in +`post_generate`. + +``` +POST /session/:id/profile — enable profiling for this session. + Must be called before POST /session/:id/generate; + enabling it after generation has started (or + finished) has nothing left to sample. +GET /session/:id/profile — returns the SessionProfile collected so far, + as JSON. 404 if profiling was never enabled + for this session_id. +``` + +### Example response + +```json +{ + "session_id": "abc123", + "phases": [ + { "phase": "merge_shares", "duration_ms": 812, "peak_memory_bytes": 41943040, "sample_count": 4, "avg_cpu_percent": 12.5, "peak_cpu_percent": 30.0 }, + { "phase": "witness_generation", "duration_ms": 15420, "peak_memory_bytes": 536870912, "sample_count": 77, "avg_cpu_percent": 88.0, "peak_cpu_percent": 100.0 }, + { "phase": "proof_generation", "duration_ms": 42110, "peak_memory_bytes": 2147483648, "sample_count": 210, "avg_cpu_percent": 95.0, "peak_cpu_percent": 100.0 } + ] +} +``` + +A retried `proof_generation` attempt (see the retry loop in +`run_proof_generation` for transient resource errors) appends another +`"proof_generation"` entry rather than overwriting the previous attempt's, +so all attempts remain visible. + +## Precision note + +`duration_ms` is measured at the sampler's 200ms sampling granularity, not +true child-process wall-clock time — a phase that finishes faster than one +sampling interval is reported as taking roughly one interval with zero +samples. In practice, MPC witness/proof generation phases run for seconds +to minutes, well above that granularity, so this is a deliberate +simplicity/precision tradeoff rather than a correctness gap for the phases +this actually profiles. + +## Implementation + +- `src/profiling.rs` — `ProfileRegistry` (which sessions are enabled + what's + been collected), `sample_process_until_exit` (the sampling loop, spawned + as its own task per phase so it runs concurrently with awaiting the + child). +- `src/session.rs`'s `run_profiled` helper wraps each `co-noir` subprocess + call: when profiling isn't enabled for the session it's exactly + `cmd.output().await` (zero extra cost); when it is, it spawns the child + with piped stdio, starts a sampler task against the child's pid, and + awaits both. diff --git a/services/node/src/api.rs b/services/node/src/api.rs index de0c303..c85ce35 100644 --- a/services/node/src/api.rs +++ b/services/node/src/api.rs @@ -392,6 +392,13 @@ pub async fn post_generate( let circuit_label = circuit_name.clone(); let finalized_sessions = state.finalized_sessions.clone(); + // Profiling is strictly opt-in per session (issue #244): only pass the + // registry through when this session_id was explicitly enabled via + // POST /session/:id/profile before generation started. A session + // nobody asked to profile pays no sampling overhead. + let profiling_enabled = state.profiling.is_enabled(&sid).await; + let profile = profiling_enabled.then(|| state.profiling.clone()); + tokio::spawn(async move { let phase_timeouts = session::PhaseTimeouts::from_env(); let proof_future = session::run_proof_generation( @@ -406,6 +413,7 @@ pub async fn post_generate( crs_path, limits, phase_timeouts, + profile, ); // Enforce a per-session wall-clock budget so a hung proof generation can't diff --git a/services/node/src/main.rs b/services/node/src/main.rs index 8bdc97c..6fee5ac 100644 --- a/services/node/src/main.rs +++ b/services/node/src/main.rs @@ -40,6 +40,7 @@ mod limits; mod metrics; mod pool; mod private_table; +mod profiling; mod session; mod tls; mod heartbeat; @@ -48,6 +49,7 @@ mod gossip; use limits::ResourceLimits; use metrics::NodeMetrics; use private_table::PrivateTableState; +use profiling::ProfileRegistry; use session::MpcSessionState; #[derive(Clone)] @@ -256,6 +258,10 @@ async fn main() { .route("/session/:id/generate", post(api::post_generate)) .route("/session/:id/status", get(api::get_status)) .route("/session/:id/proof", get(api::get_proof)) + .route( + "/session/:id/profile", + post(api::post_enable_profiling).get(api::get_profile), + ) .with_state(state); let addr = format!("0.0.0.0:{}", port); diff --git a/services/node/src/profiling.rs b/services/node/src/profiling.rs new file mode 100644 index 0000000..bb20e92 --- /dev/null +++ b/services/node/src/profiling.rs @@ -0,0 +1,301 @@ +//! On-demand CPU/memory profiling per MPC session phase (issue #244). +//! +//! The actual MPC compute work for a session happens in `co-noir` child +//! processes spawned per phase (`merge_shares`, `witness_generation`, +//! `proof_generation` — see `session::run_proof_generation`), not in this +//! Rust process's own call stack. An in-process sampling profiler (e.g. the +//! `pprof` crate) walks *this* process's call stack, so it would never see +//! the expensive work at all — it would only ever show time spent +//! orchestrating (spawning the child, waiting on it, parsing its output), +//! which is not what "profile a session" means here. +//! +//! Instead, profiling samples each child process's OS-reported CPU% and +//! RSS on a fixed interval for the duration of its phase, aggregated into a +//! peak-memory / average-and-peak-CPU% summary per phase, keyed by +//! session_id. This is exported as JSON rather than the pprof wire format: +//! pprof's call-graph model doesn't apply to an opaque external process +//! whose own internals this node has no visibility into — a flamegraph +//! needs sampled stack traces, which this project's `co-noir` dependency +//! doesn't expose to callers. +//! +//! Profiling is strictly opt-in per session (issue #244's "trigger +//! profiling on demand via API") — sampling only happens for a session +//! whose id has been explicitly enabled via `POST /session/:id/profile`, +//! so a session nobody asked to profile pays zero sampling overhead beyond +//! one registry lookup. + +use serde::{Deserialize, Serialize}; +use std::collections::{HashMap, HashSet}; +use std::sync::Arc; +use std::time::{Duration, Instant}; +use sysinfo::{Pid, System}; +use tokio::sync::RwLock; + +/// How often a profiled child process's resource usage is sampled. +pub const SAMPLE_INTERVAL: Duration = Duration::from_millis(200); + +/// Aggregated resource usage for one phase of one session's proof +/// generation run. +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct PhaseProfile { + /// "merge_shares" | "witness_generation" | "proof_generation". A + /// retried proof_generation attempt appends another entry with the + /// same phase name rather than overwriting the previous attempt's. + pub phase: String, + pub duration_ms: u64, + pub peak_memory_bytes: u64, + pub sample_count: u32, + pub avg_cpu_percent: f32, + pub peak_cpu_percent: f32, +} + +/// The full profile collected so far for one session. +#[derive(Clone, Debug, Default, Serialize, Deserialize)] +pub struct SessionProfile { + pub session_id: String, + pub phases: Vec, +} + +/// Registry of which sessions have profiling enabled and what's been +/// collected for them so far. Cheap to clone (Arc-backed) so it can be +/// threaded into the background proof-generation task alongside the other +/// state `post_generate` already captures. +#[derive(Clone)] +pub struct ProfileRegistry { + enabled: Arc>>, + profiles: Arc>>, +} + +impl ProfileRegistry { + pub fn new() -> Self { + Self { + enabled: Arc::new(RwLock::new(HashSet::new())), + profiles: Arc::new(RwLock::new(HashMap::new())), + } + } + + /// Enable profiling for `session_id`. Idempotent — calling it again on + /// an already-enabled session does not clear previously collected + /// phases. + pub async fn enable(&self, session_id: &str) { + self.enabled.write().await.insert(session_id.to_string()); + self.profiles + .write() + .await + .entry(session_id.to_string()) + .or_insert_with(|| SessionProfile { + session_id: session_id.to_string(), + phases: Vec::new(), + }); + } + + pub async fn is_enabled(&self, session_id: &str) -> bool { + self.enabled.read().await.contains(session_id) + } + + pub async fn get(&self, session_id: &str) -> Option { + self.profiles.read().await.get(session_id).cloned() + } + + async fn record_phase(&self, session_id: &str, phase: PhaseProfile) { + let mut profiles = self.profiles.write().await; + profiles + .entry(session_id.to_string()) + .or_insert_with(|| SessionProfile { + session_id: session_id.to_string(), + phases: Vec::new(), + }) + .phases + .push(phase); + } +} + +impl Default for ProfileRegistry { + fn default() -> Self { + Self::new() + } +} + +/// Sample `pid`'s CPU% and memory every [`SAMPLE_INTERVAL`] until the +/// process can no longer be found (i.e. it exited), then record the +/// aggregated [`PhaseProfile`] for `phase` into `registry`. +/// +/// Spawned as its own task by `session::run_profiled` so sampling runs +/// concurrently with (not blocking) awaiting the child process itself. +/// +/// `duration_ms` is measured from when this task starts sampling, at +/// [`SAMPLE_INTERVAL`] granularity — a phase that finishes faster than one +/// sample interval is reported as taking roughly one interval with zero +/// samples, rather than its true (shorter) wall-clock time. In practice +/// MPC witness/proof generation phases run for seconds to minutes, well +/// above that granularity, so this is a deliberate simplicity/precision +/// tradeoff rather than a correctness gap for the phases this actually +/// profiles. +pub async fn sample_process_until_exit( + registry: ProfileRegistry, + session_id: String, + phase: String, + pid: u32, +) { + let start = Instant::now(); + let mut sys = System::new(); + let sysinfo_pid = Pid::from_u32(pid); + let mut peak_memory_bytes: u64 = 0; + let mut cpu_samples: Vec = Vec::new(); + + loop { + tokio::time::sleep(SAMPLE_INTERVAL).await; + sys.refresh_processes(sysinfo::ProcessesToUpdate::Some(&[sysinfo_pid]), true); + let Some(process) = sys.process(sysinfo_pid) else { + break; + }; + // sysinfo's Process::memory() reports KiB in this version (see + // main.rs's identical *1024 conversion for the node's own process), + // not bytes — convert here so PhaseProfile::peak_memory_bytes means + // what its name says. + let memory_bytes = process.memory() * 1024; + peak_memory_bytes = peak_memory_bytes.max(memory_bytes); + cpu_samples.push(process.cpu_usage()); + } + + let avg_cpu_percent = if cpu_samples.is_empty() { + 0.0 + } else { + cpu_samples.iter().sum::() / cpu_samples.len() as f32 + }; + let peak_cpu_percent = cpu_samples.iter().cloned().fold(0.0f32, f32::max); + + registry + .record_phase( + &session_id, + PhaseProfile { + phase, + duration_ms: start.elapsed().as_millis() as u64, + peak_memory_bytes, + sample_count: cpu_samples.len() as u32, + avg_cpu_percent, + peak_cpu_percent, + }, + ) + .await; +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn enable_then_is_enabled_reports_true_only_for_that_session() { + let registry = ProfileRegistry::new(); + assert!(!registry.is_enabled("s1").await); + + registry.enable("s1").await; + + assert!(registry.is_enabled("s1").await); + assert!(!registry.is_enabled("s2").await, "enabling s1 must not affect s2"); + } + + #[tokio::test] + async fn enable_creates_an_empty_profile_immediately() { + let registry = ProfileRegistry::new(); + registry.enable("s1").await; + + let profile = registry.get("s1").await.expect("enable must create a profile entry"); + assert_eq!(profile.session_id, "s1"); + assert!(profile.phases.is_empty()); + } + + #[tokio::test] + async fn get_returns_none_for_a_session_never_enabled() { + let registry = ProfileRegistry::new(); + assert!(registry.get("never-enabled").await.is_none()); + } + + #[tokio::test] + async fn enabling_twice_does_not_clear_previously_recorded_phases() { + let registry = ProfileRegistry::new(); + registry.enable("s1").await; + registry + .record_phase( + "s1", + PhaseProfile { + phase: "merge_shares".to_string(), + duration_ms: 10, + peak_memory_bytes: 1024, + sample_count: 1, + avg_cpu_percent: 5.0, + peak_cpu_percent: 5.0, + }, + ) + .await; + + // Re-enabling (e.g. a second POST /session/:id/profile) must not + // wipe out the phase already recorded. + registry.enable("s1").await; + + let profile = registry.get("s1").await.unwrap(); + assert_eq!(profile.phases.len(), 1); + assert_eq!(profile.phases[0].phase, "merge_shares"); + } + + #[tokio::test] + async fn record_phase_appends_rather_than_overwrites_across_phases() { + let registry = ProfileRegistry::new(); + registry.enable("s1").await; + + for phase in ["merge_shares", "witness_generation", "proof_generation"] { + registry + .record_phase( + "s1", + PhaseProfile { + phase: phase.to_string(), + duration_ms: 1, + peak_memory_bytes: 0, + sample_count: 0, + avg_cpu_percent: 0.0, + peak_cpu_percent: 0.0, + }, + ) + .await; + } + + let profile = registry.get("s1").await.unwrap(); + let phases: Vec<&str> = profile.phases.iter().map(|p| p.phase.as_str()).collect(); + assert_eq!(phases, vec!["merge_shares", "witness_generation", "proof_generation"]); + } + + // Spawns a real short-lived child process and confirms the sampler + // records a phase for it with a non-zero duration once it exits. + // Unix-only: this whole service stack (co-noir, Docker deployment) + // targets Linux, so `sh` is a safe assumption for CI here the same way + // it would be for the production co-noir subprocess calls. + #[cfg(unix)] + #[tokio::test] + async fn sample_process_until_exit_records_a_phase_for_a_real_short_lived_process() { + let registry = ProfileRegistry::new(); + registry.enable("s1").await; + + let mut child = tokio::process::Command::new("sh") + .args(["-c", "sleep 0.3"]) + .spawn() + .expect("failed to spawn sh"); + let pid = child.id().expect("child must have a pid"); + + let sampler = tokio::spawn(sample_process_until_exit( + registry.clone(), + "s1".to_string(), + "test_phase".to_string(), + pid, + )); + + child.wait().await.expect("child process failed"); + sampler.await.expect("sampler task panicked"); + + let profile = registry.get("s1").await.unwrap(); + assert_eq!(profile.phases.len(), 1); + assert_eq!(profile.phases[0].phase, "test_phase"); + // duration is measured in whole SAMPLE_INTERVAL ticks; a ~300ms + // sleep should take at least one full interval. + assert!(profile.phases[0].duration_ms >= SAMPLE_INTERVAL.as_millis() as u64); + } +} diff --git a/services/node/src/session.rs b/services/node/src/session.rs index b882f08..261558c 100644 --- a/services/node/src/session.rs +++ b/services/node/src/session.rs @@ -10,9 +10,54 @@ use serde::{Deserialize, Serialize}; use std::collections::HashMap; use std::path::PathBuf; +use std::process::Stdio; use tokio::process::Command; use tokio::time::{sleep, Duration, Instant}; +/// Runs `cmd` to completion. When `profile` is `Some`, also samples the +/// child process's CPU/memory on a fixed interval for the duration of the +/// run and records it under `phase` in the registry (issue #244). When +/// `profile` is `None` (the default — profiling is opt-in per session), +/// this is exactly `cmd.output().await` with no extra overhead. +async fn run_profiled( + cmd: &mut Command, + session_id: &str, + phase: &str, + profile: Option<&crate::profiling::ProfileRegistry>, +) -> std::io::Result { + let Some(registry) = profile else { + return cmd.output().await; + }; + + // .output() configures piped stdio internally; .spawn() does not, so + // it must be set explicitly here to still capture stdout/stderr for + // the error-reporting paths above. + cmd.stdout(Stdio::piped()); + cmd.stderr(Stdio::piped()); + + let mut child = cmd.spawn()?; + let sampler = child.id().map(|pid| { + tokio::spawn(crate::profiling::sample_process_until_exit( + registry.clone(), + session_id.to_string(), + phase.to_string(), + pid, + )) + }); + + let output = child.wait_with_output().await?; + + // The sampler notices the process exited (its next refresh finds no + // such pid) and finishes on its own; just make sure it has recorded + // the phase before returning so a caller reading the profile + // immediately afterward sees it. + if let Some(sampler) = sampler { + let _ = sampler.await; + } + + Ok(output) +} + #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)] pub enum SessionStatus { /// Shares received, waiting for generate trigger @@ -209,6 +254,7 @@ pub async fn run_proof_generation( crs_path: String, limits: crate::limits::ResourceLimits, phase_timeouts: PhaseTimeouts, + profile: Option, ) -> Result<(Vec, Vec), String> { let circuit_path = format!( "{}/{}/target/{}.json", @@ -261,8 +307,7 @@ pub async fn run_proof_generation( merge_cmd.arg("--out").arg(&share_path); limits.apply_to_command(&mut merge_cmd); - let merge_output = merge_cmd - .output() + let merge_output = run_profiled(&mut merge_cmd, &session_id, "merge_shares", profile.as_ref()) .await .map_err(|e| format!("failed to spawn co-noir merge-input-shares: {}", e))?; @@ -320,8 +365,7 @@ pub async fn run_proof_generation( .arg("--out") .arg(&witness_path); limits.apply_to_command(&mut witness_cmd); - let witness_output = witness_cmd - .output() + let witness_output = run_profiled(&mut witness_cmd, &session_id, "witness_generation", profile.as_ref()) .await .map_err(|e| format!("failed to spawn co-noir generate-witness: {}", e))?; @@ -391,8 +435,7 @@ pub async fn run_proof_generation( .arg(&public_inputs_path) .arg("--fields-as-json"); limits.apply_to_command(&mut proof_cmd); - let proof_output = proof_cmd - .output() + let proof_output = run_profiled(&mut proof_cmd, &session_id, "proof_generation", profile.as_ref()) .await .map_err(|e| format!("failed to spawn co-noir build-and-generate-proof: {}", e))?;