From d539abea62dc604aa78d1d348118316daa4ccf0d Mon Sep 17 00:00:00 2001 From: Andrurachi Date: Tue, 12 May 2026 21:04:16 -0500 Subject: [PATCH] feat(server): add optional GPU metrics collection during prove closes #265 --- README.md | 22 +++++ crates/server/cli/src/commands/server.rs | 77 +++++++++++++-- crates/server/cli/src/gpu_metrics.rs | 114 +++++++++++++++++++++++ crates/server/cli/src/main.rs | 17 +++- 4 files changed, 221 insertions(+), 9 deletions(-) create mode 100644 crates/server/cli/src/gpu_metrics.rs diff --git a/README.md b/README.md index ced84f7c..a12d7c3a 100644 --- a/README.md +++ b/README.md @@ -399,6 +399,9 @@ fn main() -> Result<(), Box> { | `ERE_FORCE_REBUILD_DOCKER_IMAGE` | Force to rebuild docker images locally even they exist, it also prevents pulling image from registry. | `false` | | `ERE_GPU_DEVICES` | Specifies which GPU devices to use when running Docker containers for GPU-enabled zkVMs. The value is passed to Docker's `--gpus` flag. | `all` | | `ERE_DOCKER_NETWORK` | Specifies the Docker network being used (if any) so spawned `ere-server-*` containers will join that network. | `` | +| `ERE_COLLECT_GPU_METRICS` | Enable GPU metrics collection during prove operations (GPU resource only). Outputs CSV files with utilization, memory, and temperature data. | `false` | +| `ERE_GPU_METRICS_DIR` | Directory for GPU metrics CSV files. Created automatically if missing. | `.` | +| `ERE_GPU_METRICS` | Comma-separated `nvidia-smi` query fields (overrides the defaults). See `nvidia-smi --help-query-gpu` for available options. | `timestamp,name,...` | Example usage: @@ -416,6 +419,25 @@ ERE_GPU_DEVICES="device=0,1" ere prove ... ERE_GPU_DEVICES="4" ere prove ... ``` +## GPU Metrics Collection + +When running `ere-server` with `ProverResource::Gpu`, you can optionally collect hardware metrics during the proving process using `nvidia-smi`. + +```bash +# Basic usage (saves to current directory) +ere-server --elf-path app.elf --collect-gpu-metrics gpu + +# Custom output directory and specific metrics +ERE_GPU_METRICS="timestamp,power.draw,temperature.gpu" \ + ere-server --elf-path app.elf --collect-gpu-metrics --gpu-metrics-dir /data/metrics gpu + +``` + +Metrics are dumped into a new CSV file (`metrics_{zkvm}_{timestamp}.csv`) for each prove request, sampled at 1-second intervals. By default, it records the timestamp, GPU name, compute/memory utilization, VRAM usage, and temperature. + +*Note: `nvidia-smi` must be available in your system's PATH. If it is not found, the server will log a warning and continue proving normally.* + + ## Directory Layout ``` diff --git a/crates/server/cli/src/commands/server.rs b/crates/server/cli/src/commands/server.rs index edcaf328..f3182e1c 100644 --- a/crates/server/cli/src/commands/server.rs +++ b/crates/server/cli/src/commands/server.rs @@ -1,5 +1,6 @@ use std::{ net::{Ipv4Addr, SocketAddr}, + path::PathBuf, sync::Arc, time::{Duration, Instant}, }; @@ -24,7 +25,7 @@ use tokio::{ }; use tower::ServiceBuilder; use tower_http::{catch_panic::CatchPanicLayer, trace::TraceLayer}; -use tracing::info; +use tracing::{info, warn}; use twirp::{ Request, Response, Router, TwirpErrorResponse, async_trait::async_trait, @@ -34,24 +35,40 @@ use twirp::{ server::not_found_handler, }; -use crate::{metrics, otel}; +use crate::{gpu_metrics, metrics, otel}; pub async fn run( port: u16, elf: Elf, resource: ProverResource, prove_timeout: Option, + collect_gpu_metrics: bool, + gpu_metrics_dir: PathBuf, ) -> Result<(), Error> { let resource_kind = resource.kind(); - let zkvm = crate::construct_zkvm(elf, resource)?; + let zkvm = crate::construct_zkvm(elf, resource.clone())?; info!("initialized zkVMProver with {resource_kind} prover"); + // GPU metrics config + let gpu_metrics_config = if collect_gpu_metrics && resource == ProverResource::Gpu { + Some(gpu_metrics::GpuMetricsConfig { + zkvm_name: zkvm.name(), + output_dir: gpu_metrics_dir, + }) + } else { + None + }; + let metrics_handle = metrics::init(zkvm.name(), zkvm.sdk_version()) .context("failed to install metrics recorder")?; metrics::spawn_upkeep(metrics_handle.clone()); let prove_state = Arc::new(ProveState::new(prove_timeout)); - let server = Arc::new(zkVMServer::new(zkvm, Arc::clone(&prove_state))); + let server = Arc::new(zkVMServer::new( + zkvm, + Arc::clone(&prove_state), + gpu_metrics_config, + )); let api_middleware = ServiceBuilder::new() .layer( @@ -145,14 +162,20 @@ pub struct zkVMServer { zkvm: Arc, prove_sem: Arc, prove_state: Arc, + gpu_metrics_config: Option, } impl zkVMServer { - pub fn new(zkvm: T, prove_state: Arc) -> Self { + pub fn new( + zkvm: T, + prove_state: Arc, + gpu_metrics_config: Option, + ) -> Self { Self { zkvm: Arc::new(zkvm), prove_sem: Arc::new(Semaphore::new(1)), prove_state, + gpu_metrics_config, } } @@ -175,15 +198,53 @@ impl zkVMServer { .await .context("prove semaphore closed unexpectedly")?; + // Start GPU metrics collection if enabled + let metrics_collector = if let Some(ref config) = self.gpu_metrics_config { + match gpu_metrics::GpuMetricsCollector::start(config) { + Ok(Some(collector)) => { + info!( + "Started GPU metrics collection: {}", + collector.output_file().display() + ); + Some(collector) + } + Ok(None) => { + // nvidia-smi not available + None + } + Err(e) => { + warn!("Failed to start GPU metrics collection: {:#}", e); + None + } + } + } else { + None + }; + + // Execute prove in blocking thread let zkvm = Arc::clone(&self.zkvm); let prove_state = Arc::clone(&self.prove_state); - tokio::task::spawn_blocking(move || { + let result = tokio::task::spawn_blocking(move || { let _permit = permit; let _in_flight = ProveInFlight::new(prove_state); - Ok(zkvm.prove(&input)?) + zkvm.prove(&input) }) .await - .context("prove panicked")? + .context("prove panicked")?; + + // Stop GPU metrics collection + if let Some(collector) = metrics_collector { + match collector.stop() { + Ok(path) => { + info!("GPU metrics saved to: {}", path.display()); + } + Err(e) => { + warn!("Error stopping GPU metrics collection: {:#}", e); + } + } + } + + result.map_err(Into::into) } async fn verify(&self, proof: Proof) -> anyhow::Result { diff --git a/crates/server/cli/src/gpu_metrics.rs b/crates/server/cli/src/gpu_metrics.rs new file mode 100644 index 00000000..2e96c237 --- /dev/null +++ b/crates/server/cli/src/gpu_metrics.rs @@ -0,0 +1,114 @@ +use std::{ + env, + fs::{self, File}, + io, + path::{Path, PathBuf}, + process::{Child, Command, Stdio}, + time::{SystemTime, UNIX_EPOCH}, +}; + +use anyhow::{Context, Result}; +use tracing::{debug, warn}; + +/// Configuration for GPU metrics collection +#[derive(Clone, Debug)] +pub struct GpuMetricsConfig { + /// Name of the zkVM (e.g., "sp1", "risc0") + pub zkvm_name: &'static str, + /// Output directory for CSV files + pub output_dir: PathBuf, +} + +/// Handle for a running GPU metrics collection process +pub struct GpuMetricsCollector { + process: Child, + output_file: PathBuf, +} + +impl GpuMetricsCollector { + /// Start collecting GPU metrics + /// + /// Spawns `nvidia-smi` subprocess that writes CSV to disk every second. + /// Returns `Ok(None)` if nvidia-smi is not available. + /// Returns `Err` only on unexpected errors. + pub fn start(config: &GpuMetricsConfig) -> Result> { + // Create output directory if it doesn't exist + fs::create_dir_all(&config.output_dir).with_context(|| { + format!( + "Failed to create directory: {}", + config.output_dir.display() + ) + })?; + + // Generate filename: metrics_{zkvm}_{timestamp}.csv + let timestamp = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_secs(); + let filename = format!("metrics_{}_{}.csv", config.zkvm_name, timestamp); + let output_file = config.output_dir.join(filename); + + // Get query fields (from env or default) + let fields = env::var("ERE_GPU_METRICS").unwrap_or_else(|_| { + "timestamp,name,utilization.gpu,utilization.memory,memory.used,memory.total,temperature.gpu" + .to_string() + }); + + let file = File::create(&output_file) + .with_context(|| format!("Failed to create metrics file: {}", output_file.display()))?; + + // Build and spawn nvidia-smi command + let process = match Command::new("nvidia-smi") + .args(["--query-gpu", &fields]) + .args(["--format", "csv"]) + .args(["-l", "1"]) // Sample every 1 second + .stdout(Stdio::from(file)) + .stderr(Stdio::null()) + .spawn() + { + Ok(process) => process, + Err(e) if e.kind() == io::ErrorKind::NotFound => { + warn!("nvidia-smi not found; skipping GPU metrics collection"); + return Ok(None); + } + Err(e) => return Err(e).context("Failed to spawn nvidia-smi process"), + }; + + debug!( + "Started GPU metrics collection (PID: {}): {}", + process.id(), + output_file.display() + ); + + Ok(Some(Self { + process, + output_file, + })) + } + + /// Stop collecting GPU metrics and return the file path + pub fn stop(mut self) -> Result { + debug!( + "Stopping GPU metrics collection (PID: {})", + self.process.id() + ); + + let _ = self.process.kill(); + let _ = self.process.wait(); + + Ok(self.output_file.clone()) + } + + /// Get the output file path + pub fn output_file(&self) -> &Path { + &self.output_file + } +} + +impl Drop for GpuMetricsCollector { + fn drop(&mut self) { + // Ensure process is killed and reaped + let _ = self.process.kill(); + let _ = self.process.wait(); + } +} diff --git a/crates/server/cli/src/main.rs b/crates/server/cli/src/main.rs index c07adf0b..489cbd99 100644 --- a/crates/server/cli/src/main.rs +++ b/crates/server/cli/src/main.rs @@ -13,6 +13,7 @@ use tracing::info; use tracing_subscriber::{EnvFilter, Layer, layer::SubscriberExt, util::SubscriberInitExt}; mod commands; +mod gpu_metrics; mod metrics; mod otel; @@ -44,6 +45,12 @@ struct Args { /// milliseconds. Disabled when not set. #[arg(long, env = "ERE_PROVE_TIMEOUT_MS")] prove_timeout_ms: Option, + /// Collect GPU metrics during prove operations using nvidia-smi + #[arg(long, env = "ERE_COLLECT_GPU_METRICS")] + collect_gpu_metrics: bool, + /// Directory where GPU metrics CSV files are written + #[arg(long, env = "ERE_GPU_METRICS_DIR", default_value = ".")] + gpu_metrics_dir: PathBuf, #[command( flatten, next_help_heading = "ELF source (read from stdin if none set)" @@ -99,7 +106,15 @@ async fn main() -> Result<(), Error> { match args.command { Command::Server(resource) => { let prove_timeout = args.prove_timeout_ms.map(Duration::from_millis); - commands::server::run(args.port, elf, resource, prove_timeout).await? + commands::server::run( + args.port, + elf, + resource, + prove_timeout, + args.collect_gpu_metrics, + args.gpu_metrics_dir, + ) + .await? } Command::Keygen { program_vk_path } => commands::keygen::run(elf, &program_vk_path)?, }