diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 37f7626..3ba18e0 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -13,9 +13,8 @@ jobs: timeout-minutes: 30 steps: - uses: actions/checkout@v6 - - uses: dtolnay/rust-toolchain@stable - with: - components: clippy,rustfmt + - name: Install pinned Rust + run: rustup toolchain install - uses: taiki-e/install-action@just - name: Install native build tools run: sudo apt-get update && sudo apt-get install -y autopoint clang cmake ninja-build @@ -84,7 +83,8 @@ jobs: clang curl flex gawk libtool pkg-config python3 - name: Install Rust run: | - curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --profile minimal --default-toolchain stable + curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --profile minimal --default-toolchain none + "$HOME/.cargo/bin/rustup" toolchain install echo "$HOME/.cargo/bin" >> "$GITHUB_PATH" - uses: actions/download-artifact@v8 with: diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 10aeab9..b90fe3c 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -69,7 +69,8 @@ jobs: clang curl flex gawk libtool pkg-config python3 - name: Install Rust run: | - curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --profile minimal --default-toolchain stable + curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --profile minimal --default-toolchain none + "$HOME/.cargo/bin/rustup" toolchain install echo "$HOME/.cargo/bin" >> "$GITHUB_PATH" - uses: actions/download-artifact@v8 with: diff --git a/AGENTS.md b/AGENTS.md index e42cc4e..f6d0244 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -75,9 +75,16 @@ digest, testing installation, and inspecting all shipped ELF compatibility. - Retry CI only when logs identify a transient infrastructure failure. Preserve real compiler, test, packaging, and compatibility failures for diagnosis. +- Pin Rust in `rust-toolchain.toml`. Change the pin only in a dedicated change + and rerun lint and tests under the new compiler. ## Verification +- Run `just` recipes inside the `xprobe-dev` Mamba environment from + `environment.yml`, for example `mamba run -n xprobe-dev just test`. A bare + host shell may lack Clang, Ninja, and autotools. +- Set `PYTORCH_ENV` to an environment created from `environment-pytorch.yml` + for PyTorch live gates and `just test-release-live`. - Run `just fmt-check`, `just lint`, and `just test` for Rust changes. - Run `just test-bpf-live` for BPF attachment changes. - Run `just test-cupti-live` for CUPTI ABI or callback changes. diff --git a/docs/development.md b/docs/development.md index 6698344..8c58140 100644 --- a/docs/development.md +++ b/docs/development.md @@ -2,7 +2,9 @@ ## Native toolchain -Rust is managed with `rustup`. Native build tools are isolated in the checked-in +Rust is managed with `rustup` and pinned by `rust-toolchain.toml`; local +builds, CI, and release packaging use the same compiler and Clippy. Update the +pin in a dedicated change. Native build tools are isolated in the checked-in Mamba environment: ```bash @@ -55,9 +57,14 @@ creating a tag, run the complete live suite on the local NVIDIA development machine: ```bash -PYTORCH_ENV=/path/to/pytorch-env just test-release-live +mamba env create --file environment-pytorch.yml +PYTORCH_ENV="$(mamba run -n xprobe-pytorch python -c 'import sys; print(sys.prefix)')" \ + just test-release-live ``` +`environment-pytorch.yml` pins the PyTorch fixture: Python 3.11, because +torch 2.3 Dynamo does not support 3.12, plus Triton and NumPy 1.x. + This local gate covers BPF and perf attachment, CUDA 12 and 13 behavior, injection, NVTX, mixed host/GPU collection, PyTorch, concurrent workers, and the release benchmarks. The tag workflow independently rebuilds both CUPTI @@ -127,9 +134,10 @@ just test-pytorch-live By default the live recipes use a pinned NVIDIA PyTorch image with Ubuntu 24.04 and CUDA 12.9, so live checks do not depend on the host Mamba environment. -During local development, set `PYTORCH_ENV=/path/to/env` to mount -an existing environment into the already-pinned CUDA fixtures instead of -pulling the PyTorch image. Run eager matrix multiplication, convolution, +During local development, set `PYTORCH_ENV` to the `xprobe-pytorch` +environment to mount it into the already-pinned CUDA fixtures instead of +pulling the PyTorch image. The compiled-Triton case checks Dynamo, Triton, and +NumPy support before starting and fails with an explicit message otherwise. Run eager matrix multiplication, convolution, compiled Triton, bidirectional transfer, selected-kernel, and stream-synchronization profiling on the local GPU. The test also wraps an eager operation in an NVTX range and verifies exact range-ID matching: diff --git a/environment-pytorch.yml b/environment-pytorch.yml new file mode 100644 index 0000000..9225f74 --- /dev/null +++ b/environment-pytorch.yml @@ -0,0 +1,17 @@ +name: xprobe-pytorch + +# Local PyTorch fixture for PYTORCH_ENV live tests and benchmarks. Dynamo +# (torch.compile) in torch 2.3 supports Python 3.11 but not 3.12. +channels: + - https://conda.anaconda.org/conda-forge + - nodefaults + +dependencies: + - python=3.11 + - pip + - pip: + - --index-url https://download.pytorch.org/whl/cu121 + - --extra-index-url https://pypi.org/simple + - torch==2.3.0 + - triton==2.3.0 + - numpy<2 diff --git a/rust-toolchain.toml b/rust-toolchain.toml new file mode 100644 index 0000000..57ef39f --- /dev/null +++ b/rust-toolchain.toml @@ -0,0 +1,4 @@ +[toolchain] +channel = "1.99.1" +profile = "minimal" +components = ["clippy", "rustfmt"] diff --git a/tests/integration/test_pytorch_cuda.py b/tests/integration/test_pytorch_cuda.py index f98d7b0..652f7b4 100644 --- a/tests/integration/test_pytorch_cuda.py +++ b/tests/integration/test_pytorch_cuda.py @@ -73,10 +73,10 @@ def run_inner(xprobe: pathlib.Path) -> None: temp = pathlib.Path(temp_dir) agent = temp / "libxprobe-cupti.so" mode = temp / "mode" + require_compile_support() build_agent(workspace, agent) set_mode(mode, "mm") - environment = os.environ.copy() - environment["NVTX_INJECTION64_PATH"] = str(agent) + environment = workload_environment(agent) workload = subprocess.Popen( [ sys.executable, @@ -188,6 +188,49 @@ def run_inner(xprobe: pathlib.Path) -> None: ) +def require_compile_support() -> None: + check = subprocess.run( + [ + sys.executable, + "-c", + "import numpy, triton, torch._dynamo as dynamo; " + "raise SystemExit(0 if dynamo.is_dynamo_supported() else 3)", + ], + check=False, + capture_output=True, + text=True, + ) + if check.returncode != 0: + raise SystemExit( + "PyTorch environment cannot run torch.compile; it needs numpy, triton, " + "and a Python version supported by Dynamo (see environment-pytorch.yml)\n" + + check.stderr + ) + + +def workload_environment(agent: pathlib.Path) -> dict[str, str]: + environment = os.environ.copy() + environment["NVTX_INJECTION64_PATH"] = str(agent) + # Triton links its launcher against an unversioned libcuda.so. Some container + # runtimes (for example WSL) register only libcuda.so.1, so fall back to the + # toolkit link stub; the real driver is still loaded through its SONAME. + stub = pathlib.Path("/usr/local/cuda/lib64/stubs/libcuda.so") + if ( + "TRITON_LIBCUDA_PATH" not in environment + and not unversioned_libcuda_registered() + and stub.is_file() + ): + environment["TRITON_LIBCUDA_PATH"] = str(stub.parent) + return environment + + +def unversioned_libcuda_registered() -> bool: + cache = subprocess.run( + ["/sbin/ldconfig", "-p"], check=True, capture_output=True, text=True + ).stdout + return any(line.split()[:1] == ["libcuda.so"] for line in cache.splitlines()) + + def build_agent(workspace: pathlib.Path, output: pathlib.Path) -> None: subprocess.run( [ diff --git a/xprobe/cli/src/args.rs b/xprobe/cli/src/args.rs new file mode 100644 index 0000000..e71ce84 --- /dev/null +++ b/xprobe/cli/src/args.rs @@ -0,0 +1,448 @@ +use clap::Args; +use clap::Parser; +use clap::Subcommand; +use clap::ValueEnum; +use std::path::PathBuf; + +#[derive(Debug, Parser)] +#[command(name = "xprobe", version, about = "Runtime host-to-GPU latency probe")] +pub(crate) struct Cli { + #[command(subcommand)] + pub(crate) command: Command, +} + +#[derive(Debug, Subcommand)] +pub(crate) enum Command { + /// Inspect local tracing and GPU capabilities. + Doctor(DoctorArgs), + /// List NVML-confirmed CUDA context holders under a process-tree root. + Discover(DiscoverArgs), + /// Inspect a target process without attaching probes. + #[command(hide = true)] + Inspect(InspectArgs), + /// Resolve a userspace event selector against a target process. + #[command(hide = true)] + Resolve(ResolveArgs), + /// Validate two event selectors and their correlation policy. + Validate(ValidateArgs), + /// Measure latency between two events in files or a running process. + Measure(Box), + /// Run a live measurement from a versioned specification. + #[command(hide = true)] + Trace(TraceArgs), + /// Export completed captures to a stable trace format. + #[command(hide = true)] + Export(ExportArgs), + /// Run bounded event collectors. + #[command(hide = true)] + Capture(DevArgs), + /// Run low-level development probes. + #[command(hide = true)] + Dev(DevArgs), +} + +#[derive(Debug, Args)] +pub(crate) struct DevArgs { + #[command(subcommand)] + pub(crate) command: DevCommand, +} + +#[derive(Debug, Args)] +pub(crate) struct TraceArgs { + /// Versioned `MeasurementSpec` JSON file. + #[arg(long)] + pub(crate) spec: PathBuf, + + /// Unix socket exposed by the target's xprobe CUPTI agent. + #[arg(long, hide = true)] + pub(crate) cupti_socket: Option, + + /// Emit only the versioned JSON result on stdout. + #[arg(long)] + pub(crate) json: bool, + + /// Disable colored output. + #[arg(long)] + pub(crate) no_color: bool, + + /// Never wait for user input. + #[arg(long)] + pub(crate) non_interactive: bool, +} + +#[derive(Debug, Clone, Copy, ValueEnum)] +pub(crate) enum ExportFormatArg { + Jsonl, + Chrome, +} + +#[derive(Debug, Args)] +pub(crate) struct ExportArgs { + /// Completed CUPTI binary, host capture JSON, or Event JSONL; repeat to merge. + #[arg(long, required = true)] + pub(crate) input: Vec, + + /// Artifact format. + #[arg(long, value_enum)] + pub(crate) format: ExportFormatArg, + + /// Destination artifact path. + #[arg(long)] + pub(crate) output: PathBuf, + + /// Emit only the versioned export result on stdout. + #[arg(long)] + pub(crate) json: bool, + + /// Disable colored output. + #[arg(long)] + pub(crate) no_color: bool, + + /// Never wait for user input. + #[arg(long)] + pub(crate) non_interactive: bool, +} + +#[derive(Debug, Subcommand)] +pub(crate) enum DevCommand { + /// Capture entries to one userspace function. + Uprobe(UprobeArgs), + /// Decode an xprobe CUPTI capture as Event JSONL. + Cupti(CuptiArgs), +} + +#[derive(Debug, Clone, Copy, Args)] +pub(crate) struct DoctorArgs { + /// Emit only the versioned JSON result on stdout. + #[arg(long)] + pub(crate) json: bool, + + /// Disable colored output. + #[arg(long)] + pub(crate) no_color: bool, + + /// Never wait for user input. + #[arg(long)] + pub(crate) non_interactive: bool, +} + +#[derive(Debug, Clone, Copy, Args)] +pub(crate) struct DiscoverArgs { + /// Target process ID. + #[arg(long)] + pub(crate) pid: u32, + + /// Maximum number of CUDA worker candidates to return. + #[arg(long, default_value_t = 200)] + pub(crate) limit: usize, + + /// Emit only the versioned JSON result on stdout. + #[arg(long)] + pub(crate) json: bool, + + /// Disable colored output. + #[arg(long)] + pub(crate) no_color: bool, + + /// Never wait for user input. + #[arg(long)] + pub(crate) non_interactive: bool, +} + +#[derive(Debug, Clone, Copy, Args)] +pub(crate) struct InspectArgs { + /// Target process ID. + #[arg(long)] + pub(crate) pid: u32, + + /// Emit only the versioned JSON result on stdout. + #[arg(long)] + pub(crate) json: bool, + + /// Disable colored output. + #[arg(long)] + pub(crate) no_color: bool, + + /// Never wait for user input. + #[arg(long)] + pub(crate) non_interactive: bool, +} + +#[derive(Debug, Args)] +pub(crate) struct ResolveArgs { + /// Target process ID. + #[arg(long)] + pub(crate) pid: u32, + + /// Event selector for a mapped ELF symbol, C++ signature, or file offset. + #[arg(long, alias = "event")] + pub(crate) selector: String, + + /// Emit only the versioned JSON result on stdout. + #[arg(long)] + pub(crate) json: bool, + + /// Disable colored output. + #[arg(long)] + pub(crate) no_color: bool, + + /// Never wait for user input. + #[arg(long)] + pub(crate) non_interactive: bool, +} + +#[derive(Debug, Args)] +pub(crate) struct ValidateArgs { + /// Target process ID. + #[arg(long)] + pub(crate) pid: u32, + + /// Start event selector. + #[arg(long)] + pub(crate) from: Option, + + /// End event selector. + #[arg(long)] + pub(crate) to: Option, + + /// Correlation policy. + #[arg(long = "match")] + pub(crate) match_policy: Option, + + /// Validate bounded PID-scoped CPU sampling instead of an event pair. + #[arg(long, conflicts_with_all = ["from", "to", "match_policy", "syscall_aggregate"])] + pub(crate) cpu_sample: bool, + + /// Validate bounded PID-scoped syscall aggregation instead of an event pair. + #[arg(long, conflicts_with_all = ["from", "to", "match_policy", "cpu_sample"])] + pub(crate) syscall_aggregate: bool, + + #[command(flatten)] + pub(crate) output: CommonOutputArgs, +} + +#[derive(Debug, Args)] +pub(crate) struct MeasureArgs { + /// Versioned measurement JSON for a live target. + #[arg(long, conflicts_with_all = ["input", "pid", "from", "to", "match_policy", "cpu_sample", "syscall_aggregate", "samples", "duration_ms", "timeout_ms", "max_events", "max_samples", "frequency_hz", "stack_depth", "max_threads", "max_inflight", "aggregate", "max_groups", "name"])] + pub(crate) spec: Option, + + /// Completed CUPTI binary, host capture JSON, or Event JSONL; repeat to merge. + #[arg(long)] + pub(crate) input: Vec, + + /// Collect from a running target process instead of completed capture files. + #[arg(long)] + pub(crate) pid: Option, + + /// Unix socket exposed by the target's xprobe CUPTI agent. + #[arg(long, hide = true)] + pub(crate) cupti_socket: Option, + + /// CUPTI agent shared object used for automatic online injection. + #[arg(long)] + pub(crate) agent: Option, + + /// Write the complete bounded event capture to this file. + #[arg(long)] + pub(crate) events_out: Option, + + /// Evidence event format; defaults to jsonl when --events-out is set. + #[arg(long, value_enum)] + pub(crate) format: Option, + + /// Bound foreground collection and cleanup to this many milliseconds. + #[arg(long, default_value_t = 30_000)] + pub(crate) timeout_ms: u64, + + /// Start event selector. + #[arg(long)] + pub(crate) from: Option, + + /// End event selector. + #[arg(long)] + pub(crate) to: Option, + + /// Correlation policy. + #[arg(long = "match")] + pub(crate) match_policy: Option, + + /// Collect a bounded PID-scoped CPU hotspot inventory. + #[arg(long, conflicts_with_all = ["input", "from", "to", "match_policy", "samples", "aggregate", "events_out", "format", "cupti_socket", "agent"])] + pub(crate) cpu_sample: bool, + + /// Target CPU sampling frequency in hertz. + #[arg( + long, + requires = "cpu_sample", + default_value_if("cpu_sample", "true", "99") + )] + pub(crate) frequency_hz: Option, + + /// Bound all observed CPU sample records. + #[arg( + long, + requires = "cpu_sample", + default_value_if("cpu_sample", "true", "10000") + )] + pub(crate) max_samples: Option, + + /// Bound user-space callchain depth. + #[arg( + long, + requires = "cpu_sample", + default_value_if("cpu_sample", "true", "64") + )] + pub(crate) stack_depth: Option, + + /// Bound target threads attached by CPU sampling. + #[arg( + long, + requires = "cpu_sample", + default_value_if("cpu_sample", "true", "1024") + )] + pub(crate) max_threads: Option, + + /// Aggregate all syscall lifecycles in bounded BPF maps. + #[arg(long, conflicts_with_all = ["input", "from", "to", "match_policy", "samples", "cpu_sample", "aggregate", "events_out", "format", "cupti_socket", "agent"])] + pub(crate) syscall_aggregate: bool, + + /// Bound concurrent syscall lifecycle starts retained in BPF. + #[arg( + long, + requires = "syscall_aggregate", + default_value_if("syscall_aggregate", "true", "1024") + )] + pub(crate) max_inflight: Option, + + /// Stop after this many matched samples. + #[arg(long)] + pub(crate) samples: Option, + + /// Restrict events to this many milliseconds from the first selected event. + #[arg(long)] + pub(crate) duration_ms: Option, + + /// Reject captures containing more events than this limit. + #[arg(long, default_value_t = 100_000)] + pub(crate) max_events: usize, + + /// Return bounded activity aggregates instead of exact event evidence. + #[arg(long, conflicts_with_all = ["input", "samples", "events_out", "format"])] + pub(crate) aggregate: bool, + + /// Bound distinct groups retained by aggregate or CPU inventory mode. + #[arg( + long, + default_value_if("aggregate", "true", "4096"), + default_value_if("cpu_sample", "true", "256"), + default_value_if("syscall_aggregate", "true", "256") + )] + pub(crate) max_groups: Option, + + /// Optional measurement name. + #[arg(long)] + pub(crate) name: Option, + + #[command(flatten)] + pub(crate) output: CommonOutputArgs, +} + +#[derive(Debug, Clone, Copy, Args)] +pub(crate) struct CommonOutputArgs { + /// Emit only the versioned JSON result on stdout. + #[arg(long)] + pub(crate) json: bool, + + /// Disable colored output. + #[arg(long)] + pub(crate) no_color: bool, + + /// Never wait for user input. + #[arg(long)] + pub(crate) non_interactive: bool, +} + +#[derive(Debug, Args)] +pub(crate) struct UprobeArgs { + /// Target process ID. + #[arg(long)] + pub(crate) pid: u32, + + /// Executable or mapped shared library containing the symbol. + #[arg(long)] + pub(crate) binary: PathBuf, + + /// ELF function symbol to attach to. + #[arg(long)] + pub(crate) symbol: String, + + /// Attach at function return instead of function entry. + #[arg(long = "return")] + pub(crate) return_probe: bool, + + /// Probe identifier stored in each event. + #[arg(long, default_value_t = 1)] + pub(crate) probe_id: u32, + + /// Stop after this many events. + #[arg(long, default_value_t = 1)] + pub(crate) samples: usize, + + /// Stop waiting after this many milliseconds. + #[arg(long, default_value_t = 5_000)] + pub(crate) timeout_ms: u64, + + #[command(flatten)] + pub(crate) output: EventOutputArgs, + + /// Disable colored output. + #[arg(long)] + pub(crate) no_color: bool, + + /// Never wait for user input. + #[arg(long)] + pub(crate) non_interactive: bool, +} + +#[derive(Debug, Clone, Copy, Args)] +pub(crate) struct EventOutputArgs { + /// Emit only the versioned JSON result on stdout. + #[arg(long)] + pub(crate) json: bool, + + /// Emit one versioned Event JSON object per line. + #[arg(long, conflicts_with = "json")] + pub(crate) jsonl: bool, +} + +#[derive(Debug, Args)] +pub(crate) struct CuptiArgs { + /// Raw binary capture written by the xprobe CUPTI agent. + #[arg(long, required_unless_present = "socket", conflicts_with = "socket")] + pub(crate) input: Option, + + /// Unix socket exposed by a running xprobe CUPTI agent. + #[arg(long, required_unless_present = "input", conflicts_with = "input")] + pub(crate) socket: Option, + + /// Stop waiting for an online snapshot after this many milliseconds. + #[arg(long, default_value_t = 10_000)] + pub(crate) timeout_ms: u64, + + /// Session identifier written into every Event. + #[arg(long)] + pub(crate) session_id: String, + + /// Emit one versioned Event JSON object per line. + #[arg(long)] + pub(crate) json: bool, + + /// Disable colored output. + #[arg(long)] + pub(crate) no_color: bool, + + /// Never wait for user input. + #[arg(long)] + pub(crate) non_interactive: bool, +} diff --git a/xprobe/cli/src/dev.rs b/xprobe/cli/src/dev.rs new file mode 100644 index 0000000..505ed6b --- /dev/null +++ b/xprobe/cli/src/dev.rs @@ -0,0 +1,525 @@ +use crate::EXPORT_SEQUENCE; +use crate::args::CuptiArgs; +use crate::args::DevArgs; +use crate::args::DevCommand; +use crate::args::EventOutputArgs; +use crate::args::ExportArgs; +use crate::args::ExportFormatArg; +use crate::args::InspectArgs; +use crate::args::ResolveArgs; +use crate::args::TraceArgs; +use crate::args::UprobeArgs; +use crate::emit_measurement; +use crate::failure::CommandFailure; +use crate::failure::emit_command_failure; +use crate::failure::emit_error; +use crate::live::LiveMeasureRequest; +use crate::live::collect_live_measurement; +use crate::live::match_policy_name; +use crate::render::print_host_capture; +use crate::render::print_process_report; +use crate::render::print_resolved_probe; +use std::fs; +use std::fs::OpenOptions; +use std::io::Write; +use std::os::unix::fs::OpenOptionsExt; +use std::path::Path; +use std::path::PathBuf; +use std::process::ExitCode; +use std::sync::Arc; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; +use std::time::Duration; +use xprobe_collector::completed; +use xprobe_collector::cupti; +use xprobe_collector::uprobe; +use xprobe_collector::uprobe::UprobeRequest; +use xprobe_core::inspect; +use xprobe_core::resolve; +use xprobe_correlator::MeasureOptions; +use xprobe_exporter::events_to_chrome_trace; +use xprobe_exporter::events_to_jsonl; +use xprobe_protocol::ErrorCode; +use xprobe_protocol::ExportFormat; +use xprobe_protocol::MeasurementMode; +use xprobe_protocol::MeasurementSpec; +use xprobe_protocol::ProcessReport; +use xprobe_protocol::SchemaVersion; +use xprobe_protocol::TraceExportResult; + +pub(crate) fn run_capture(args: DevArgs) -> ExitCode { + match args.command { + DevCommand::Uprobe(args) => run_uprobe(args), + DevCommand::Cupti(args) => run_cupti(args), + } +} + +pub(crate) fn run_trace(args: TraceArgs) -> ExitCode { + let TraceArgs { + spec, + cupti_socket, + json, + no_color: _, + non_interactive: _, + } = args; + let bytes = match fs::read(&spec) { + Ok(bytes) => bytes, + Err(error) => { + return emit_error( + ErrorCode::TraceExportFailed, + format!("failed to read {}: {error}", spec.display()), + false, + json, + ); + } + }; + let spec: MeasurementSpec = match serde_json::from_slice(&bytes) { + Ok(spec) => spec, + Err(error) => { + return emit_error( + ErrorCode::InvalidEventSelector, + format!("invalid MeasurementSpec {}: {error}", spec.display()), + true, + json, + ); + } + }; + if spec.measurement_mode != MeasurementMode::Exact || spec.max_groups.is_some() { + return emit_error( + ErrorCode::InvalidEventSelector, + "legacy trace accepts only exact MeasurementSpec requests".to_owned(), + true, + json, + ); + } + let samples = match spec.samples.map(usize::try_from).transpose() { + Ok(samples) => samples, + Err(error) => { + return emit_error( + ErrorCode::SessionLimitExceeded, + format!("MeasurementSpec samples exceed this platform: {error}"), + true, + json, + ); + } + }; + let Some(max_events) = spec.max_events else { + return emit_error( + ErrorCode::SessionLimitExceeded, + "exact MeasurementSpec requires max_events".to_owned(), + true, + json, + ); + }; + let max_events = match usize::try_from(max_events) { + Ok(max_events) => max_events, + Err(error) => { + return emit_error( + ErrorCode::SessionLimitExceeded, + format!("MeasurementSpec max_events exceed this platform: {error}"), + true, + json, + ); + } + }; + let request = LiveMeasureRequest { + pid: spec.target.pid, + expected_target: Some(spec.target), + cupti_socket, + agent_path: None, + timeout: Duration::from_millis(spec.timeout_ms), + match_policy_text: match_policy_name(spec.match_policy), + mode: MeasurementMode::Exact, + options: MeasureOptions { + session_id: format!("xp_trace_{}", std::process::id()), + name: spec.name, + start_selector: spec.start_selector, + end_selector: spec.end_selector, + match_policy: spec.match_policy, + samples, + duration: spec.duration_ms.map(Duration::from_millis), + max_events, + dropped_events: 0, + }, + }; + match collect_live_measurement(&request) { + Ok(execution) => emit_measurement(&execution.result, json), + Err(error) => emit_command_failure(error, json), + } +} + +pub(crate) fn run_export(args: ExportArgs) -> ExitCode { + let ExportArgs { + input, + format, + output, + json, + no_color: _, + non_interactive: _, + } = args; + let session_id = format!("xp_export_{}", std::process::id()); + let capture = match load_completed_inputs(&input, &session_id) { + Ok(capture) => capture, + Err(error) => return emit_command_failure(error, json), + }; + let (format, artifact) = match format { + ExportFormatArg::Jsonl => ( + ExportFormat::Jsonl, + events_to_jsonl(&capture.events).map_err(|error| error.to_string()), + ), + ExportFormatArg::Chrome => ( + ExportFormat::Chrome, + events_to_chrome_trace(&capture.events).map_err(|error| error.to_string()), + ), + }; + let artifact = match artifact { + Ok(artifact) => artifact, + Err(error) => { + return emit_error(ErrorCode::TraceExportFailed, error, false, json); + } + }; + if let Err(error) = write_export_file(&output, artifact.as_bytes()) { + return emit_error( + ErrorCode::TraceExportFailed, + format!("failed to write {}: {error}", output.display()), + false, + json, + ); + } + let event_count = u64::try_from(capture.events.len()).expect("event count fits u64"); + let result = TraceExportResult { + schema_version: SchemaVersion::current(), + ok: true, + format, + output: output.to_string_lossy().into_owned(), + event_count, + }; + if json { + println!( + "{}", + serde_json::to_string_pretty(&result).expect("export result must serialize") + ); + } else { + println!( + "Exported {} events to {} ({:?})", + result.event_count, result.output, result.format + ); + } + ExitCode::SUCCESS +} + +pub(crate) fn load_completed_inputs( + input: &[PathBuf], + session_id: &str, +) -> Result { + let mut captures = Vec::with_capacity(input.len()); + for path in input { + let bytes = fs::read(path).map_err(|error| { + CommandFailure::new( + ErrorCode::TraceExportFailed, + format!("failed to read {}: {error}", path.display()), + false, + ) + })?; + captures.push(completed::decode(&bytes, session_id).map_err(|error| { + CommandFailure::new( + ErrorCode::TraceExportFailed, + format!("failed to decode {}: {error}", path.display()), + false, + ) + })?); + } + let capture = completed::merge(captures, session_id).map_err(|error| { + CommandFailure::new(ErrorCode::TraceExportFailed, error.to_string(), false) + })?; + if capture.unknown_records != 0 { + return Err(CommandFailure::new( + ErrorCode::TraceExportFailed, + format!( + "capture contains {} unknown CUPTI activity records", + capture.unknown_records + ), + false, + ) + .with_detail("unknown_records", capture.unknown_records) + .with_hint("rebuild xprobe and the CUPTI Agent from the same release")); + } + Ok(capture) +} + +pub(crate) fn write_export_file(path: &Path, bytes: &[u8]) -> Result<(), String> { + let parent = path + .parent() + .filter(|parent| !parent.as_os_str().is_empty()) + .unwrap_or_else(|| Path::new(".")); + let file_name = path + .file_name() + .ok_or_else(|| format!("artifact path {} has no file name", path.display()))?; + let temporary = parent.join(format!( + ".{}.xprobe-{}-{}.tmp", + file_name.to_string_lossy(), + std::process::id(), + EXPORT_SEQUENCE.fetch_add(1, Ordering::Relaxed) + )); + let mut file = OpenOptions::new() + .write(true) + .create_new(true) + .mode(0o600) + .open(&temporary) + .map_err(|error| format!("failed to create {}: {error}", temporary.display()))?; + if let Err(error) = file.write_all(bytes).and_then(|()| file.sync_all()) { + let original = format!("failed to write {}: {error}", temporary.display()); + return Err(remove_temporary_artifact(&temporary, original)); + } + drop(file); + if let Err(error) = fs::rename(&temporary, path) { + let original = format!( + "failed to atomically replace {} with {}: {error}", + path.display(), + temporary.display() + ); + return Err(remove_temporary_artifact(&temporary, original)); + } + fs::File::open(parent) + .and_then(|directory| directory.sync_all()) + .map_err(|error| { + format!( + "failed to sync artifact directory {}: {error}", + parent.display() + ) + }) +} + +pub(crate) fn remove_temporary_artifact(path: &Path, original: String) -> String { + match fs::remove_file(path) { + Ok(()) => original, + Err(cleanup) => format!( + "{original}; failed to remove temporary artifact {}: {cleanup}", + path.display() + ), + } +} + +pub(crate) fn run_resolve(args: ResolveArgs) -> ExitCode { + let ResolveArgs { + pid, + selector, + json, + no_color: _, + non_interactive: _, + } = args; + let report = match inspect::run(pid) { + Ok(report) => report, + Err(error) => { + return emit_error(error.code(), error.to_string(), error.recoverable(), json); + } + }; + + match resolve::run(&report, &selector) { + Ok(resolved) => { + if json { + println!( + "{}", + serde_json::to_string_pretty(&resolved).expect("resolved probe must serialize") + ); + } else { + print_resolved_probe(&resolved); + } + ExitCode::SUCCESS + } + Err(error) => emit_error(error.code(), error.to_string(), error.recoverable(), json), + } +} + +pub(crate) fn run_inspect(args: InspectArgs) -> ExitCode { + let InspectArgs { + pid, + json, + no_color: _, + non_interactive: _, + } = args; + + match inspect::run(pid) { + Ok(report) => { + if json { + println!( + "{}", + serde_json::to_string_pretty(&report).expect("process report must serialize") + ); + } else { + print_process_report(&report); + } + ExitCode::SUCCESS + } + Err(error) => emit_error(error.code(), error.to_string(), error.recoverable(), json), + } +} + +pub(crate) fn run_uprobe(args: UprobeArgs) -> ExitCode { + let UprobeArgs { + pid, + binary, + symbol, + return_probe, + probe_id, + samples, + timeout_ms, + output: EventOutputArgs { json, jsonl }, + no_color: _, + non_interactive: _, + } = args; + let machine_output = json || jsonl; + + let report = match inspect::run(pid) { + Ok(report) => report, + Err(error) => { + return emit_error( + error.code(), + error.to_string(), + error.recoverable(), + machine_output, + ); + } + }; + let binary = match mapped_binary(&report, &binary) { + Ok(binary) => binary, + Err(message) => { + return emit_error(ErrorCode::BinaryNotMapped, message, true, machine_output); + } + }; + let request = UprobeRequest { + target: report.target.clone(), + binary, + symbol: Some(symbol), + symbol_demangled: None, + offset: 0, + probe_kind: if return_probe { + xprobe_protocol::HostProbeKind::Uretprobe + } else { + xprobe_protocol::HostProbeKind::Uprobe + }, + probe_id, + samples, + timeout: Duration::from_millis(timeout_ms), + cancelled: Arc::new(AtomicBool::new(false)), + ready: None, + }; + let result = match uprobe::capture(&request) { + Ok(result) => result, + Err(error) => { + return emit_error( + error.code(), + error.to_string(), + error.recoverable(), + machine_output, + ); + } + }; + if let Err(error) = inspect::verify_target(&report.target) { + return emit_error( + error.code(), + error.to_string(), + error.recoverable(), + machine_output, + ); + } + + if jsonl { + match events_to_jsonl(&result.events) { + Ok(output) => print!("{output}"), + Err(error) => { + return emit_error( + ErrorCode::TraceExportFailed, + format!("failed to serialize host events: {error}"), + false, + true, + ); + } + } + } else if json { + println!( + "{}", + serde_json::to_string_pretty(&result).expect("host capture result must serialize") + ); + } else { + print_host_capture(&result); + } + ExitCode::SUCCESS +} + +pub(crate) fn run_cupti(args: CuptiArgs) -> ExitCode { + let CuptiArgs { + input, + socket, + timeout_ms, + session_id, + json, + no_color: _, + non_interactive: _, + } = args; + let capture = if let Some(input) = input { + let bytes = match fs::read(&input) { + Ok(bytes) => bytes, + Err(error) => { + return emit_error( + ErrorCode::TraceExportFailed, + format!("failed to read {}: {error}", input.display()), + false, + json, + ); + } + }; + match cupti::decode_capture(&bytes, &session_id) { + Ok(capture) => capture, + Err(error) => { + return emit_error(ErrorCode::TraceExportFailed, error.to_string(), false, json); + } + } + } else { + let socket = socket.expect("clap requires one CUPTI capture source"); + match cupti::snapshot(&socket, Duration::from_millis(timeout_ms), &session_id, 0) { + Ok(capture) => capture, + Err(error) => { + return emit_error(ErrorCode::TraceExportFailed, error.to_string(), true, json); + } + } + }; + let output = match events_to_jsonl(&capture.events) { + Ok(output) => output, + Err(error) => { + return emit_error( + ErrorCode::TraceExportFailed, + format!("failed to serialize CUPTI events: {error}"), + false, + json, + ); + } + }; + + if capture.dropped_records != 0 || capture.unknown_records != 0 { + eprintln!( + "CUPTI capture diagnostics: dropped={}, unknown={}", + capture.dropped_records, capture.unknown_records + ); + } + print!("{output}"); + ExitCode::SUCCESS +} + +pub(crate) fn mapped_binary(report: &ProcessReport, binary: &Path) -> Result { + let canonical = fs::canonicalize(binary) + .map_err(|error| format!("failed to resolve {}: {error}", binary.display()))?; + let is_executable = Path::new(&report.executable) == canonical; + let is_library = report + .loaded_libraries + .iter() + .any(|mapped| Path::new(mapped) == canonical); + if !is_executable && !is_library { + return Err(format!( + "{} is not mapped by target PID {}", + canonical.display(), + report.target.pid + )); + } + Ok(canonical) +} diff --git a/xprobe/cli/src/failure.rs b/xprobe/cli/src/failure.rs new file mode 100644 index 0000000..94bb806 --- /dev/null +++ b/xprobe/cli/src/failure.rs @@ -0,0 +1,166 @@ +use crate::clock_domain_name; +use crate::live::match_policy_name; +use std::collections::BTreeMap; +use std::process::ExitCode; +use xprobe_correlator::MeasureError; +use xprobe_correlator::MeasureOptions; +use xprobe_protocol::ErrorCode; +use xprobe_protocol::ErrorResponse; +use xprobe_protocol::Event; +use xprobe_protocol::EventSource; +use xprobe_protocol::XprobeError; + +#[derive(Debug)] +pub(crate) struct CommandFailure { + pub(crate) code: ErrorCode, + pub(crate) message: String, + pub(crate) recoverable: bool, + pub(crate) details: BTreeMap, + pub(crate) hints: Vec, + pub(crate) artifact_events: Option>, +} + +impl CommandFailure { + pub(crate) fn new(code: ErrorCode, message: impl Into, recoverable: bool) -> Self { + Self { + code, + message: message.into(), + recoverable, + details: BTreeMap::new(), + hints: Vec::new(), + artifact_events: None, + } + } + + pub(crate) fn with_detail( + mut self, + key: impl Into, + value: impl Into, + ) -> Self { + self.details.insert(key.into(), value.into()); + self + } + + pub(crate) fn with_hint(mut self, hint: impl Into) -> Self { + self.hints.push(hint.into()); + self + } + + pub(crate) fn with_artifact_events(mut self, events: Vec) -> Self { + self.artifact_events = Some(events); + self + } +} + +pub(crate) fn measurement_failure( + error: &MeasureError, + options: &MeasureOptions, + events: &[Event], + preserve_artifact: bool, +) -> CommandFailure { + let mut failure = CommandFailure::new(error.code(), error.to_string(), error.recoverable()) + .with_detail("start_selector", options.start_selector.clone()) + .with_detail("end_selector", options.end_selector.clone()) + .with_detail("match_policy", match_policy_name(options.match_policy)); + if preserve_artifact { + failure = failure.with_artifact_events(events.to_vec()); + } + match error { + MeasureError::InvalidSelector(_) => { + failure = failure.with_hint("run xprobe validate with the exact selector text"); + } + MeasureError::InvalidPolicy(_) => { + failure = failure + .with_hint("use policy_recommendation.compatible_policies from xprobe validate"); + } + MeasureError::InvalidLimit(_) => { + failure = failure + .with_hint("set positive samples or duration, timeout, and max-events bounds"); + } + MeasureError::EventLimitExceeded { actual, maximum } => { + failure = failure + .with_detail("observed_events", *actual as u64) + .with_detail("max_events", *maximum as u64) + .with_hint("narrow the selectors or explicitly increase max-events"); + } + MeasureError::EventsDropped { count } => { + failure = failure + .with_detail("dropped_events", *count) + .with_hint("reduce the event rate or explicitly increase the capture capacity"); + } + MeasureError::ClockDomainsDiffer { start, end } => { + failure = failure + .with_detail("start_clock", clock_domain_name(start)) + .with_detail("end_clock", clock_domain_name(end)) + .with_hint( + "use endpoints in one clock domain or a CUPTI capture normalized to host monotonic time", + ); + } + MeasureError::NoMatchedSamples => { + let host_events = events + .iter() + .filter(|event| event.source == EventSource::Ebpf) + .count() as u64; + failure = failure + .with_detail("captured_events", events.len() as u64) + .with_detail("host_events", host_events) + .with_detail("cuda_events", events.len() as u64 - host_events) + .with_hint( + "inspect the events-out artifact, then adjust selectors or choose an explicitly compatible policy", + ); + } + } + failure +} + +pub(crate) fn emit_error( + code: ErrorCode, + message: String, + recoverable: bool, + json: bool, +) -> ExitCode { + emit_xprobe_error( + XprobeError { + code, + message, + recoverable, + details: BTreeMap::new(), + hints: Vec::new(), + }, + json, + ) +} + +pub(crate) fn emit_command_failure(error: CommandFailure, json: bool) -> ExitCode { + emit_xprobe_error( + XprobeError { + code: error.code, + message: error.message, + recoverable: error.recoverable, + details: error.details, + hints: error.hints, + }, + json, + ) +} + +pub(crate) fn emit_xprobe_error(error: XprobeError, json: bool) -> ExitCode { + let code = error.code; + let response = ErrorResponse::new(error); + if json { + println!( + "{}", + serde_json::to_string_pretty(&response).expect("error response must serialize") + ); + } else { + eprintln!("{code}: {}", response.error.message); + } + + match code { + ErrorCode::TargetNotFound | ErrorCode::TargetExited | ErrorCode::TargetReused => { + ExitCode::from(3) + } + ErrorCode::PermissionDenied => ExitCode::from(4), + _ => ExitCode::from(1), + } +} diff --git a/xprobe/cli/src/inventory.rs b/xprobe/cli/src/inventory.rs new file mode 100644 index 0000000..696f21b --- /dev/null +++ b/xprobe/cli/src/inventory.rs @@ -0,0 +1,460 @@ +use crate::failure::CommandFailure; +use crate::failure::emit_command_failure; +use crate::failure::emit_error; +use std::path::Path; +use std::process::ExitCode; +use std::time::Duration; +use xprobe_collector::cpu_sampling as cpu_collector; +use xprobe_collector::cpu_sampling::CpuSamplingRequest; +use xprobe_collector::syscall_aggregate as syscall_collector; +use xprobe_collector::syscall_aggregate::SyscallAggregateRequest; +use xprobe_core::cpu_sampling; +use xprobe_core::cpu_sampling_analysis; +use xprobe_core::cpu_sampling_analysis::CpuCaptureData; +use xprobe_core::cpu_sampling_analysis::RawCpuSample; +use xprobe_core::inspect; +use xprobe_core::syscall_aggregate; +use xprobe_protocol::CpuSampleEvent; +use xprobe_protocol::CpuSampleInventoryResult; +use xprobe_protocol::CpuSamplingSpec; +use xprobe_protocol::ErrorCode; +use xprobe_protocol::ProcessReport; +use xprobe_protocol::SchemaVersion; +use xprobe_protocol::SyscallAggregateResult; +use xprobe_protocol::SyscallAggregateSpec; + +#[allow(clippy::too_many_arguments)] +pub(crate) fn run_direct_cpu_sampling( + pid: u32, + name: Option, + frequency_hz: u64, + duration_ms: u64, + timeout_ms: u64, + max_samples: u64, + max_groups: u64, + stack_depth: u32, + max_threads: u32, + json: bool, +) -> ExitCode { + let report = match inspect::run(pid) { + Ok(report) => report, + Err(error) => { + return emit_error(error.code(), error.to_string(), error.recoverable(), json); + } + }; + let spec = CpuSamplingSpec { + schema_version: SchemaVersion::current(), + name, + target: report.target.clone(), + sample_event: CpuSampleEvent::CpuClock, + frequency_hz, + duration_ms, + timeout_ms, + max_samples, + max_groups, + stack_depth, + max_threads, + }; + run_cpu_sampling_spec(&spec, Some(report), json) +} + +pub(crate) fn run_cpu_sampling_spec( + spec: &CpuSamplingSpec, + inspected: Option, + json: bool, +) -> ExitCode { + let report = match inspected.map_or_else(|| inspect::run(spec.target.pid), Ok) { + Ok(report) => report, + Err(error) => { + return emit_error(error.code(), error.to_string(), error.recoverable(), json); + } + }; + if let Err(error) = inspect::verify_target(&spec.target) { + return emit_error(error.code(), error.to_string(), error.recoverable(), json); + } + let validation = match cpu_sampling::validate(&report) { + Ok(validation) => validation, + Err(error) => { + return emit_error(error.code(), error.to_string(), error.recoverable(), json); + } + }; + if !validation.valid { + let message = validation.issues.first().map_or_else( + || "CPU sampling requirements are not satisfied".to_owned(), + |issue| issue.message.clone(), + ); + return emit_error(ErrorCode::PermissionDenied, message, true, json); + } + let request = match cpu_collector_request(spec) { + Ok(request) => request, + Err(error) => return emit_command_failure(error, json), + }; + let capture = match cpu_collector::collect(&request) { + Ok(capture) => capture, + Err(error) => return emit_command_failure(cpu_collector_failure(&error), json), + }; + if capture.observed_samples == 0 { + return emit_error( + ErrorCode::NoMatchedSamples, + "CPU sampling completed without observing an on-CPU sample".to_owned(), + true, + json, + ); + } + let capture = CpuCaptureData { + samples: capture + .stacks + .into_iter() + .map(|stack| RawCpuSample { + thread_id: stack.thread_id, + addresses: stack.addresses, + }) + .collect(), + observed_samples: capture.observed_samples, + lost_samples: capture.lost_samples, + truncated_stacks: capture.truncated_stacks, + throttled_records: capture.throttled_records, + threads_observed: capture.threads_observed, + threads_attached: capture.threads_attached, + skipped_threads: capture.skipped_threads, + capacity_reached: capture.capacity_reached, + python_status: validation.python_status, + }; + let result = match cpu_sampling_analysis::analyze( + &report, + spec, + format!("xp_cpu_{}", std::process::id()), + &capture, + ) { + Ok(result) => result, + Err(error) => { + return emit_error(error.code(), error.to_string(), error.recoverable(), json); + } + }; + emit_cpu_inventory(&result, json) +} + +pub(crate) fn cpu_collector_request( + spec: &CpuSamplingSpec, +) -> Result { + if spec.frequency_hz == 0 + || spec.duration_ms == 0 + || spec.max_samples == 0 + || spec.max_groups == 0 + || spec.stack_depth == 0 + || spec.max_threads == 0 + { + return Err(CommandFailure::new( + ErrorCode::SessionLimitExceeded, + "CPU sampling frequency, duration, samples, groups, stack depth, and threads must be positive", + true, + )); + } + let max_samples = usize::try_from(spec.max_samples).map_err(|error| { + CommandFailure::new( + ErrorCode::SessionLimitExceeded, + format!("CPU max_samples exceed this platform: {error}"), + true, + ) + })?; + let max_threads = usize::try_from(spec.max_threads).map_err(|error| { + CommandFailure::new( + ErrorCode::SessionLimitExceeded, + format!("CPU max_threads exceed this platform: {error}"), + true, + ) + })?; + let stack_depth = u16::try_from(spec.stack_depth).map_err(|error| { + CommandFailure::new( + ErrorCode::SessionLimitExceeded, + format!("CPU stack_depth exceeds the perf ABI: {error}"), + true, + ) + })?; + let thread_ids = cpu_sampling::target_thread_ids(&spec.target).map_err(|error| { + CommandFailure::new(error.code(), error.to_string(), error.recoverable()) + })?; + Ok(CpuSamplingRequest { + thread_ids, + frequency_hz: spec.frequency_hz, + duration: Duration::from_millis(spec.duration_ms), + timeout: Duration::from_millis(spec.timeout_ms), + max_samples, + max_threads, + stack_depth, + }) +} + +pub(crate) fn cpu_collector_failure(error: &cpu_collector::CpuSamplingError) -> CommandFailure { + let code = match error { + cpu_collector::CpuSamplingError::InvalidRequest(_) + | cpu_collector::CpuSamplingError::Timeout => ErrorCode::SessionLimitExceeded, + cpu_collector::CpuSamplingError::Open { source, .. } + | cpu_collector::CpuSamplingError::Arm { source, .. } + if source.kind() == std::io::ErrorKind::PermissionDenied => + { + ErrorCode::PermissionDenied + } + cpu_collector::CpuSamplingError::Cleanup { .. } => ErrorCode::CleanupFailed, + cpu_collector::CpuSamplingError::NoThreadsAttached { .. } => ErrorCode::TargetExited, + cpu_collector::CpuSamplingError::Open { .. } + | cpu_collector::CpuSamplingError::Ring { .. } + | cpu_collector::CpuSamplingError::Arm { .. } => ErrorCode::Internal, + }; + let recoverable = !matches!(code, ErrorCode::Internal | ErrorCode::CleanupFailed); + CommandFailure::new(code, error.to_string(), recoverable) +} + +pub(crate) fn emit_cpu_inventory(result: &CpuSampleInventoryResult, json: bool) -> ExitCode { + if json { + println!( + "{}", + serde_json::to_string_pretty(result).expect("CPU inventory must serialize") + ); + } else { + println!( + "CPU sample inventory: {} samples", + result.collection.observed_samples + ); + println!( + "Stacks: {} groups, {} lost, {} unresolved frames", + result.collection.groups, + result.collection.lost_samples, + result.symbolization.unresolved_frames + ); + for hotspot in result.inventory.hotspots.iter().take(10) { + let label = hotspot + .frame + .symbol + .as_deref() + .or(hotspot.frame.module_path.as_deref()) + .map_or_else(|| format!("{:#x}", hotspot.frame.address), str::to_owned); + println!( + " {:.1}% inclusive, {:.1}% self {label}", + hotspot.inclusive_proportion * 100.0, + hotspot.exclusive_proportion * 100.0 + ); + } + for warning in &result.warnings { + println!("Warning: {}: {}", warning.code, warning.message); + } + } + ExitCode::SUCCESS +} + +#[allow(clippy::too_many_arguments)] +pub(crate) fn run_direct_syscall_aggregate( + pid: u32, + name: Option, + duration_ms: u64, + timeout_ms: u64, + max_groups: u64, + max_inflight: u64, + json: bool, +) -> ExitCode { + let report = match inspect::run(pid) { + Ok(report) => report, + Err(error) => { + return emit_error(error.code(), error.to_string(), error.recoverable(), json); + } + }; + let spec = SyscallAggregateSpec { + schema_version: SchemaVersion::current(), + name, + target: report.target.clone(), + duration_ms, + timeout_ms, + max_groups, + max_inflight, + }; + run_syscall_aggregate_spec(&spec, Some(report), json) +} + +pub(crate) fn run_syscall_aggregate_spec( + spec: &SyscallAggregateSpec, + inspected: Option, + json: bool, +) -> ExitCode { + let report = match inspected.map_or_else(|| inspect::run(spec.target.pid), Ok) { + Ok(report) => report, + Err(error) => { + return emit_error(error.code(), error.to_string(), error.recoverable(), json); + } + }; + if let Err(error) = inspect::verify_target(&spec.target) { + return emit_error(error.code(), error.to_string(), error.recoverable(), json); + } + let validation = match syscall_aggregate::validate(&report) { + Ok(validation) => validation, + Err(error) => { + return emit_error(error.code(), error.to_string(), error.recoverable(), json); + } + }; + if !validation.valid { + let message = validation.issues.first().map_or_else( + || "syscall aggregate requirements are not satisfied".to_owned(), + |issue| issue.message.clone(), + ); + return emit_error(ErrorCode::PermissionDenied, message, true, json); + } + let request = match syscall_collector_request(spec) { + Ok(request) => request, + Err(error) => return emit_command_failure(error, json), + }; + let capture = match syscall_collector::collect(&request) { + Ok(capture) => capture, + Err(error) => { + return emit_error(error.code(), error.to_string(), error.recoverable(), json); + } + }; + let capture = syscall_aggregate::SyscallCaptureData { + groups: capture + .groups + .into_iter() + .map(|group| syscall_aggregate::RawSyscallGroup { + syscall_number: group.syscall_number, + count: group.count, + errors: group.errors, + total_duration_ns: group.total_duration_ns, + min_duration_ns: group.min_duration_ns, + max_duration_ns: group.max_duration_ns, + }) + .collect(), + observed_entries: capture.observed_entries, + matched_exits: capture.matched_exits, + unmatched_exits: capture.unmatched_exits, + dropped_aggregates: capture.dropped_aggregates, + inflight_at_end: capture.inflight_at_end, + }; + let result = match syscall_aggregate::analyze( + &report, + spec, + format!("xp_syscalls_{}", std::process::id()), + &capture, + ) { + Ok(result) => result, + Err(error) => { + return emit_error(error.code(), error.to_string(), error.recoverable(), json); + } + }; + emit_syscall_inventory(&result, json) +} + +pub(crate) fn syscall_collector_request( + spec: &SyscallAggregateSpec, +) -> Result { + if spec.duration_ms == 0 || spec.max_groups == 0 || spec.max_inflight == 0 { + return Err(CommandFailure::new( + ErrorCode::SessionLimitExceeded, + "syscall aggregate duration, groups, and inflight capacity must be positive", + true, + )); + } + let max_groups = u32::try_from(spec.max_groups).map_err(|error| { + CommandFailure::new( + ErrorCode::SessionLimitExceeded, + format!("syscall max_groups exceed the BPF map ABI: {error}"), + true, + ) + })?; + let max_inflight = u32::try_from(spec.max_inflight).map_err(|error| { + CommandFailure::new( + ErrorCode::SessionLimitExceeded, + format!("syscall max_inflight exceeds the BPF map ABI: {error}"), + true, + ) + })?; + Ok(SyscallAggregateRequest { + target: spec.target.clone(), + duration: Duration::from_millis(spec.duration_ms), + timeout: Duration::from_millis(spec.timeout_ms), + max_groups, + max_inflight, + }) +} + +pub(crate) fn emit_syscall_inventory(result: &SyscallAggregateResult, json: bool) -> ExitCode { + if json { + println!( + "{}", + serde_json::to_string_pretty(result).expect("syscall inventory must serialize") + ); + } else { + println!( + "Syscall inventory: {} entries, {} matched exits", + result.collection.observed_entries, result.collection.matched_exits + ); + for group in result.inventory.groups.iter().take(10) { + let label = group + .syscall_name + .clone() + .unwrap_or_else(|| format!("syscall#{}", group.syscall_number)); + println!( + " {label}: count={} total={}ns mean={:.0}ns errors={}", + group.count, group.duration_ns.total, group.duration_ns.mean, group.errors + ); + } + for warning in &result.warnings { + println!("Warning: {}: {}", warning.code, warning.message); + } + } + ExitCode::SUCCESS +} + +pub(crate) fn run_cpu_sampling_file_spec( + value: serde_json::Value, + path: &Path, + has_device_options: bool, + has_artifact_options: bool, + json: bool, +) -> ExitCode { + if has_device_options || has_artifact_options { + return emit_error( + ErrorCode::InvalidEventSelector, + "CpuSamplingSpec does not accept CUPTI or event artifact options".to_owned(), + true, + json, + ); + } + let spec: CpuSamplingSpec = match serde_json::from_value(value) { + Ok(spec) => spec, + Err(error) => { + return emit_error( + ErrorCode::InvalidEventSelector, + format!("invalid CpuSamplingSpec {}: {error}", path.display()), + true, + json, + ); + } + }; + run_cpu_sampling_spec(&spec, None, json) +} + +pub(crate) fn run_syscall_aggregate_file_spec( + value: serde_json::Value, + path: &Path, + has_device_options: bool, + has_artifact_options: bool, + json: bool, +) -> ExitCode { + if has_device_options || has_artifact_options { + return emit_error( + ErrorCode::InvalidEventSelector, + "SyscallAggregateSpec does not accept CUPTI or event artifact options".to_owned(), + true, + json, + ); + } + let spec: SyscallAggregateSpec = match serde_json::from_value(value) { + Ok(spec) => spec, + Err(error) => { + return emit_error( + ErrorCode::InvalidEventSelector, + format!("invalid SyscallAggregateSpec {}: {error}", path.display()), + true, + json, + ); + } + }; + run_syscall_aggregate_spec(&spec, None, json) +} diff --git a/xprobe/cli/src/live.rs b/xprobe/cli/src/live.rs new file mode 100644 index 0000000..4e1fabb --- /dev/null +++ b/xprobe/cli/src/live.rs @@ -0,0 +1,1432 @@ +use crate::MeasurementExecution; +use crate::failure::CommandFailure; +use crate::failure::measurement_failure; +use std::collections::BTreeMap; +use std::path::Path; +use std::path::PathBuf; +use std::sync::Arc; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; +use std::sync::mpsc; +use std::thread; +use std::thread::JoinHandle; +use std::time::Duration; +use std::time::Instant; +use xprobe_collector::completed; +use xprobe_collector::cupti; +use xprobe_collector::linux::LinuxCaptureRequest; +use xprobe_collector::uprobe::UprobeRequest; +use xprobe_collector::{linux, uprobe}; +use xprobe_core::cupti_compat; +use xprobe_core::gpu_aggregate; +use xprobe_core::inject; +use xprobe_core::inspect; +use xprobe_core::validate; +use xprobe_correlator::MeasureOptions; +use xprobe_correlator::measure; +use xprobe_protocol::AggregateActivity; +use xprobe_protocol::AggregateInventoryResult; +use xprobe_protocol::CuptiCollectionSummary; +use xprobe_protocol::ErrorCode; +use xprobe_protocol::EventType; +use xprobe_protocol::HostCaptureResult; +use xprobe_protocol::MatchPolicy; +use xprobe_protocol::MeasurementMode; +use xprobe_protocol::MeasurementResult; +use xprobe_protocol::ProcessReport; +use xprobe_protocol::ResolvedLinuxSelector; +use xprobe_protocol::ResolvedProbe; +use xprobe_protocol::SessionStatus; +use xprobe_protocol::TargetIdentity; +use xprobe_protocol::ValidationResult; +use xprobe_protocol::Warning; + +pub(crate) struct LiveMeasureRequest { + pub(crate) pid: u32, + pub(crate) expected_target: Option, + pub(crate) cupti_socket: Option, + pub(crate) agent_path: Option, + pub(crate) timeout: Duration, + pub(crate) match_policy_text: &'static str, + pub(crate) mode: MeasurementMode, + pub(crate) options: MeasureOptions, +} + +type HostCaptureHandle = JoinHandle>; +pub(crate) const SNAPSHOT_QUIET_PERIOD: Duration = Duration::from_millis(100); + +pub(crate) struct HostCollectors { + pub(crate) cancelled: Arc, + pub(crate) handles: Vec, +} + +impl HostCollectors { + pub(crate) fn finished(&self) -> bool { + self.handles.iter().all(JoinHandle::is_finished) + } + + pub(crate) fn cancel(&self) { + self.cancelled.store(true, Ordering::Release); + } + + pub(crate) fn join(&mut self) -> Result, CommandFailure> { + let mut captures = Vec::with_capacity(self.handles.len()); + for handle in self.handles.drain(..) { + let result = handle.join().map_err(|_| { + CommandFailure::new(ErrorCode::Internal, "host collector thread panicked", false) + })?; + captures.push(result?); + } + Ok(captures) + } +} + +impl Drop for HostCollectors { + fn drop(&mut self) { + self.cancel(); + for handle in self.handles.drain(..) { + if handle.join().is_err() { + eprintln!("xprobe: host collector thread panicked during cleanup"); + } + } + } +} + +pub(crate) struct LiveCollection { + pub(crate) report: ProcessReport, + pub(crate) socket: Option, + pub(crate) deadline: Instant, + pub(crate) collection_end: Option, + pub(crate) baseline_timestamp: Option, + pub(crate) baseline_dropped: u64, + pub(crate) baseline_unknown: u64, + pub(crate) baseline_observed: u64, + pub(crate) collectors: HostCollectors, + pub(crate) host_captures: Option>, + pub(crate) latest_cuda: Option, + pub(crate) cuda_record_offset: u64, + pub(crate) managed_agent: bool, + pub(crate) cupti_armed: bool, + pub(crate) agent_injected: bool, + pub(crate) agent_major: Option, + pub(crate) agent_path: Option, + pub(crate) cupti_path: Option, +} + +pub(crate) fn collect_live_measurement( + request: &LiveMeasureRequest, +) -> Result { + assert_eq!(request.mode, MeasurementMode::Exact); + validate_live_limits(request)?; + let mut collection = prepare_live_collection(request)?; + + let outcome = (|| { + loop { + refresh_live_sources(&mut collection, request)?; + let now = Instant::now(); + if let Some(result) = evaluate_live_result(&collection, request, now)? { + break Ok(result); + } + if now >= collection.deadline { + break finish_live_timeout(&mut collection, request); + } + if collection.socket.is_none() { + thread::sleep(Duration::from_millis(10)); + } + } + })(); + let cuda_capture = collection.cupti_armed; + let stop_result = stop_cupti_agent(&mut collection, request); + let mut execution = match (outcome, stop_result) { + (Ok(mut execution), Ok(())) => { + if cuda_capture && collection.host_captures.is_some() { + let status = execution.result.status; + execution = measure_live_capture( + collection.host_captures.as_deref().unwrap_or(&[]), + &collection, + &request.options, + true, + )?; + execution.result.status = status; + } + execution + } + (Err(error), Ok(())) => { + if cuda_capture && collection.host_captures.is_some() { + return match measure_live_capture( + collection.host_captures.as_deref().unwrap_or(&[]), + &collection, + &request.options, + true, + ) { + Ok(_) => Err(error), + Err(final_error) => Err(final_error), + }; + } + return Err(error); + } + (Ok(execution), Err(cleanup)) => { + return Err(cleanup.with_artifact_events(execution.events)); + } + (Err(mut error), Err(cleanup)) => { + let artifact_events = error.artifact_events.take(); + let mut failure = CommandFailure::new( + ErrorCode::CleanupFailed, + format!( + "{}; original measurement error: {}", + cleanup.message, error.message + ), + cleanup.recoverable, + ) + .with_detail("cleanup_phase", "stop_cupti") + .with_detail("original_error_code", error.code.to_string()) + .with_detail("cleanup_error_code", cleanup.code.to_string()) + .with_hint("verify the target state before starting another measurement"); + if let Some(events) = artifact_events { + failure = failure.with_artifact_events(events); + } + return Err(failure); + } + }; + if collection.agent_injected { + let mut details = BTreeMap::new(); + if let Some(major) = collection.agent_major { + details.insert("cuda_major".to_owned(), serde_json::Value::from(major)); + } + if let Some(path) = collection.agent_path.as_deref() { + details.insert( + "agent_path".to_owned(), + serde_json::Value::from(path.display().to_string()), + ); + } + if let Some(path) = collection.cupti_path.as_deref() { + details.insert( + "cupti_path".to_owned(), + serde_json::Value::from(path.display().to_string()), + ); + } + execution.result.warnings.push(Warning { + code: "CUPTI_AGENT_INJECTED".to_owned(), + message: "xprobe injected the CUPTI agent and left the shared object mapped".to_owned(), + details, + }); + } + Ok(execution) +} + +pub(crate) fn collect_live_aggregate( + request: &LiveMeasureRequest, +) -> Result { + assert_eq!(request.mode, MeasurementMode::Aggregate); + validate_live_limits(request)?; + let mut collection = prepare_live_collection(request)?; + let collection_end = collection + .collection_end + .expect("aggregate validation requires duration"); + if collection_end > collection.deadline { + return Err(CommandFailure::new( + ErrorCode::SessionLimitExceeded, + "aggregate duration must not exceed timeout", + true, + )); + } + + while Instant::now() < collection_end { + let remaining = collection_end.saturating_duration_since(Instant::now()); + thread::sleep(remaining.min(Duration::from_millis(10))); + } + let target_result = inspect::verify_target(&collection.report.target) + .map_err(|error| CommandFailure::new(error.code(), error.to_string(), error.recoverable())); + let stop_result = stop_cupti_agent(&mut collection, request); + match (target_result, stop_result) { + (Err(error), Ok(())) | (Ok(()), Err(error)) => return Err(error), + (Err(target), Err(cleanup)) => { + return Err(CommandFailure::new( + ErrorCode::CleanupFailed, + format!( + "{}; target verification also failed: {}", + cleanup.message, target.message + ), + cleanup.recoverable, + ) + .with_detail("cleanup_phase", "stop_cupti") + .with_detail("target_error_code", target.code.to_string()) + .with_hint("verify the target state before starting another measurement")); + } + (Ok(()), Ok(())) => {} + } + + let capture = collection.latest_cuda.as_ref().ok_or_else(|| { + CommandFailure::new( + ErrorCode::CuptiNotAvailable, + "CUPTI Agent returned no aggregate capture", + false, + ) + })?; + let mut result = aggregate_inventory_result(capture, request)?; + if collection.agent_injected { + result.warnings.push(injection_warning(&collection)); + } + Ok(result) +} + +pub(crate) fn validate_live_limits(request: &LiveMeasureRequest) -> Result<(), CommandFailure> { + if request.timeout.is_zero() { + return Err(CommandFailure::new( + ErrorCode::SessionLimitExceeded, + "timeout must be greater than zero", + true, + )); + } + if request.options.samples == Some(0) + || request + .options + .duration + .is_some_and(|duration| duration.is_zero()) + || request.options.max_events == 0 + || (request.options.samples.is_none() && request.options.duration.is_none()) + { + return Err(CommandFailure::new( + ErrorCode::SessionLimitExceeded, + "live measurement requires positive samples or duration and max-events limits", + true, + )); + } + if request.mode == MeasurementMode::Aggregate { + if request.options.samples.is_some() + || request.options.duration.is_none() + || request.options.match_policy != MatchPolicy::Exact + { + return Err(CommandFailure::new( + ErrorCode::InvalidCorrelationPolicy, + "aggregate measurement requires duration, exact activity endpoints, and no samples limit", + true, + )); + } + if request.options.duration > Some(request.timeout) { + return Err(CommandFailure::new( + ErrorCode::SessionLimitExceeded, + "aggregate duration must not exceed timeout", + true, + )); + } + } + Ok(()) +} + +pub(crate) fn prepare_live_collection( + request: &LiveMeasureRequest, +) -> Result { + let report = inspect::run(request.pid).map_err(|error| { + CommandFailure::new(error.code(), error.to_string(), error.recoverable()) + })?; + if request + .expected_target + .as_ref() + .is_some_and(|expected| expected != &report.target) + { + return Err(CommandFailure::new( + ErrorCode::TargetReused, + "trace specification target identity no longer matches the process", + true, + )); + } + let validation = validate::run( + &report, + &request.options.start_selector, + &request.options.end_selector, + request.match_policy_text, + ) + .map_err(|error| CommandFailure::new(error.code(), error.to_string(), error.recoverable()))?; + if let Some(issue) = validation.issues.first() { + return Err(CommandFailure::new(issue.code, issue.message.clone(), true) + .with_detail( + "recommended_policy", + match_policy_name(validation.policy_recommendation.policy), + ) + .with_detail( + "compatible_policies", + serde_json::to_value(&validation.policy_recommendation.compatible_policies) + .expect("match policies must serialize"), + ) + .with_hint("rerun validate with an explicitly compatible policy")); + } + if request.mode == MeasurementMode::Aggregate { + validate_aggregate_selection(&validation)?; + } + + let activation = prepare_cupti_activation(&report, &validation, request)?; + let arm_config = cupti_arm_config(&validation, request.options.max_events, request.mode)?; + let baseline = arm_cupti_capture( + &activation, + arm_config.as_ref(), + validation.requirements.needs_nvtx, + request, + )?; + let socket = activation.socket; + let baseline_timestamp = baseline + .as_ref() + .and_then(|capture| capture.events.last()) + .map(|event| event.timestamp_ns); + let baseline_dropped = baseline + .as_ref() + .map_or(0, |capture| capture.dropped_records); + let baseline_unknown = baseline + .as_ref() + .map_or(0, |capture| capture.unknown_records); + let baseline_observed = baseline + .as_ref() + .map_or(0, |capture| capture.observed_records); + let started = Instant::now(); + let deadline = started.checked_add(request.timeout).ok_or_else(|| { + CommandFailure::new( + ErrorCode::SessionLimitExceeded, + "timeout exceeds Instant range", + true, + ) + })?; + let collection_end = request + .options + .duration + .and_then(|duration| started.checked_add(duration)); + let collectors = start_host_collectors(&report, &validation, request)?; + let host_captures = collectors.handles.is_empty().then(Vec::new); + + Ok(LiveCollection { + report, + socket, + deadline, + collection_end, + baseline_timestamp, + baseline_dropped, + baseline_unknown, + baseline_observed, + collectors, + host_captures, + latest_cuda: None, + cuda_record_offset: 0, + managed_agent: activation.managed, + cupti_armed: baseline.is_some(), + agent_injected: activation.injected, + agent_major: activation.major, + agent_path: activation.agent_path, + cupti_path: activation.cupti_path, + }) +} + +pub(crate) fn validate_aggregate_selection( + validation: &ValidationResult, +) -> Result<(), CommandFailure> { + let (Some(start), Some(end)) = (validation.start.cuda.as_ref(), validation.end.cuda.as_ref()) + else { + return Err(CommandFailure::new( + ErrorCode::InvalidEventSelector, + "aggregate measurement supports CUPTI activity endpoints only", + true, + )); + }; + let paired = matches!( + (&start.event_type, &end.event_type), + (&EventType::GpuKernelStart, &EventType::GpuKernelEnd) + | (&EventType::GpuMemcpyStart, &EventType::GpuMemcpyEnd) + | (&EventType::GpuMemsetStart, &EventType::GpuMemsetEnd) + ); + if !paired + || start.api_domain != end.api_domain + || start.api_name != end.api_name + || start.kernel_name_regex != end.kernel_name_regex + || start.memcpy_kind != end.memcpy_kind + { + return Err(CommandFailure::new( + ErrorCode::InvalidEventSelector, + "aggregate measurement requires matching kernel, memcpy, or memset start/end selectors", + true, + )); + } + if let Some(pattern) = start.kernel_name_regex.as_deref() + && pattern != ".*" + && matches!( + cupti::bounded_name_filter(pattern), + cupti::CuptiNameFilter::Any + ) + { + return Err(CommandFailure::new( + ErrorCode::InvalidEventSelector, + "aggregate kernel regex must be an exact, prefix, suffix, or contains pattern that the CUPTI Agent can apply", + true, + ) + .with_hint("use ^literal$, ^literal.*, .*literal$, or .*literal.*")); + } + Ok(()) +} + +pub(crate) fn arm_cupti_capture( + activation: &CuptiActivation, + config: Option<&cupti::CuptiArmConfig>, + needs_nvtx: bool, + request: &LiveMeasureRequest, +) -> Result, CommandFailure> { + let (Some(path), Some(config)) = (activation.socket.as_deref(), config) else { + assert!(activation.socket.is_none() && config.is_none()); + return Ok(None); + }; + let capture = match cupti::arm(path, request.timeout, &request.options.session_id, config) { + Ok(capture) => capture, + Err(error) => { + let failure = + CommandFailure::new(ErrorCode::CuptiNotAvailable, error.to_string(), true); + return cleanup_failed_arm(activation, request, failure); + } + }; + if needs_nvtx && !capture.nvtx_callbacks_active { + let failure = CommandFailure::new( + ErrorCode::CuptiNotAvailable, + "the mapped CUPTI agent was not initialized through NVTX; restart the target with NVTX_INJECTION64_PATH set to the same agent before the first NVTX call", + true, + ) + .with_hint("online CUDA injection cannot retrofit an initialized NVTX dispatch table"); + return cleanup_failed_arm(activation, request, failure); + } + if capture.state != cupti::CuptiCaptureState::Active { + let failure = CommandFailure::new( + ErrorCode::CuptiNotAvailable, + format!( + "CUPTI Agent did not enter active state after ARM: {:?}", + capture.state + ), + true, + ); + return cleanup_failed_arm(activation, request, failure); + } + Ok(Some(capture)) +} + +pub(crate) fn cleanup_failed_arm( + activation: &CuptiActivation, + request: &LiveMeasureRequest, + failure: CommandFailure, +) -> Result, CommandFailure> { + if !activation.managed { + return Err(failure); + } + let path = activation + .socket + .as_deref() + .expect("managed activation has a socket"); + match cupti::close(path, request.timeout, &request.options.session_id, 0) { + Ok(_) => Err(failure), + Err(cleanup) => Err(CommandFailure::new( + ErrorCode::CleanupFailed, + format!( + "failed to close CUPTI agent after ARM error: {cleanup}; original error: {}", + failure.message + ), + true, + ) + .with_detail("cleanup_phase", "close_after_arm_failure") + .with_detail("original_error_code", failure.code.to_string()) + .with_hint("verify the target state before starting another measurement")), + } +} + +pub(crate) fn cupti_arm_config( + validation: &ValidationResult, + max_events: usize, + mode: MeasurementMode, +) -> Result, CommandFailure> { + if !validation.requirements.needs_cupti { + return Ok(None); + } + let filters = [&validation.start, &validation.end] + .into_iter() + .filter_map(|endpoint| endpoint.cuda.as_ref()) + .map(|selector| { + cupti::event_filter(selector).map_err(|error| { + CommandFailure::new(ErrorCode::InvalidEventSelector, error.to_string(), false) + }) + }) + .collect::, _>>()?; + let record_capacity = u64::try_from(max_events).map_err(|error| { + CommandFailure::new( + ErrorCode::SessionLimitExceeded, + format!("max-events exceeds the CUPTI Agent range: {error}"), + true, + ) + })?; + Ok(Some(cupti::CuptiArmConfig { + record_capacity, + capture_mode: match mode { + MeasurementMode::Exact => cupti::CuptiCaptureMode::Exact, + MeasurementMode::Aggregate => cupti::CuptiCaptureMode::Aggregate, + }, + filters, + })) +} + +pub(crate) struct CuptiActivation { + pub(crate) socket: Option, + pub(crate) managed: bool, + pub(crate) injected: bool, + pub(crate) major: Option, + pub(crate) agent_path: Option, + pub(crate) cupti_path: Option, +} + +pub(crate) struct AgentSelection { + pub(crate) path: PathBuf, + pub(crate) major: Option, + pub(crate) cupti_path: Option, +} + +pub(crate) fn prepare_cupti_activation( + report: &ProcessReport, + validation: &ValidationResult, + request: &LiveMeasureRequest, +) -> Result { + if !validation.requirements.needs_cupti { + if request.cupti_socket.is_some() { + return Err(CommandFailure::new( + ErrorCode::InvalidEventSelector, + "--cupti-socket is only valid when a CUDA endpoint is selected", + true, + )); + } + return Ok(CuptiActivation { + socket: None, + managed: false, + injected: false, + major: None, + agent_path: None, + cupti_path: None, + }); + } + if let Some(socket) = request.cupti_socket.clone() { + return Ok(CuptiActivation { + socket: Some(socket), + managed: false, + injected: false, + major: None, + agent_path: None, + cupti_path: None, + }); + } + + let socket = std::env::temp_dir().join(format!( + "xprobe-{}-{}.sock", + request.pid, + std::process::id() + )); + let selection = if report.cuda.xprobe_cupti_loaded { + AgentSelection { + path: PathBuf::new(), + major: cupti_compat::target_major(report) + .map_err(|error| CommandFailure::new(error.code(), error.to_string(), true))?, + cupti_path: None, + } + } else { + resolve_agent_path(report, request.agent_path.as_deref())? + }; + eprintln!( + "xprobe: warning: activating the CUPTI agent modifies target PID {}{}{}", + request.pid, + selection + .major + .map_or_else(String::new, |major| format!(" with CUDA {major} Agent")), + selection + .cupti_path + .as_deref() + .map_or_else(String::new, |path| format!(" using {}", path.display())) + ); + let record_capacity = u64::try_from(request.options.max_events).map_err(|error| { + CommandFailure::new( + ErrorCode::SessionLimitExceeded, + format!("max-events exceeds the CUPTI Agent range: {error}"), + true, + ) + })?; + let activation = inject::activate( + report, + &selection.path, + &socket, + record_capacity, + request.timeout, + ) + .map_err(|error| CommandFailure::new(error.code(), error.to_string(), error.recoverable()))?; + Ok(CuptiActivation { + socket: Some(activation.socket_path), + managed: true, + injected: activation.injected, + major: selection.major, + agent_path: activation.injected.then_some(selection.path), + cupti_path: selection.cupti_path, + }) +} + +pub(crate) fn resolve_agent_path( + report: &ProcessReport, + configured: Option<&Path>, +) -> Result { + if let Some(path) = configured { + return Ok(AgentSelection { + path: path.to_owned(), + major: cupti_compat::target_major(report) + .map_err(|error| CommandFailure::new(error.code(), error.to_string(), true))?, + cupti_path: None, + }); + } + if let Some(path) = std::env::var_os("XPROBE_CUPTI_AGENT_PATH") { + return Ok(AgentSelection { + path: PathBuf::from(path), + major: cupti_compat::target_major(report) + .map_err(|error| CommandFailure::new(error.code(), error.to_string(), true))?, + cupti_path: None, + }); + } + let cupti = cupti_compat::resolve_library(report) + .map_err(|error| CommandFailure::new(error.code(), error.to_string(), true))?; + let executable = std::env::current_exe().map_err(|error| { + CommandFailure::new( + ErrorCode::CuptiNotAvailable, + format!("failed to locate xprobe executable: {error}"), + true, + ) + })?; + let parent = executable + .parent() + .expect("executable has a parent directory"); + let candidates = [ + parent.join(format!( + "../lib/xprobe/cuda{}/libxprobe-cupti.so", + cupti.major + )), + parent.join(format!("cuda{}/libxprobe-cupti.so", cupti.major)), + PathBuf::from(format!( + "build/cupti/cuda{}/libxprobe-cupti.so", + cupti.major + )), + ]; + let path = candidates + .into_iter() + .find(|path| path.is_file()) + .ok_or_else(|| { + CommandFailure::new( + ErrorCode::CuptiNotAvailable, + format!( + "CUDA {} CUPTI agent was not found; pass --agent or set XPROBE_CUPTI_AGENT_PATH", + cupti.major + ), + true, + ) + })?; + Ok(AgentSelection { + path, + major: Some(cupti.major), + cupti_path: Some(cupti.path), + }) +} + +pub(crate) fn stop_cupti_agent( + collection: &mut LiveCollection, + request: &LiveMeasureRequest, +) -> Result<(), CommandFailure> { + if !collection.cupti_armed { + return Ok(()); + } + let socket = collection + .socket + .as_deref() + .expect("armed CUPTI agent has a socket"); + collection.cupti_armed = false; + let capture = if collection.managed_agent { + cupti::close( + socket, + request.timeout, + &request.options.session_id, + collection.cuda_record_offset, + ) + } else { + cupti::stop( + socket, + request.timeout, + &request.options.session_id, + collection.cuda_record_offset, + ) + } + .map_err(|error| { + CommandFailure::new(ErrorCode::CleanupFailed, error.to_string(), true) + .with_detail("cleanup_phase", "stop_cupti") + .with_hint("verify the target state before starting another measurement") + })?; + append_cuda_capture(collection, capture)?; + collection.managed_agent = false; + Ok(()) +} + +impl Drop for LiveCollection { + fn drop(&mut self) { + if self.cupti_armed { + if let Some(socket) = self.socket.as_deref() { + let result = if self.managed_agent { + cupti::close( + socket, + Duration::from_secs(2), + "xp_cleanup", + self.cuda_record_offset, + ) + } else { + cupti::stop( + socket, + Duration::from_secs(2), + "xp_cleanup", + self.cuda_record_offset, + ) + }; + if let Err(error) = result { + eprintln!("xprobe: failed to stop CUPTI agent during cleanup: {error}"); + } + } + } + } +} + +pub(crate) fn start_host_collectors( + report: &ProcessReport, + validation: &ValidationResult, + request: &LiveMeasureRequest, +) -> Result { + inspect::verify_target(&report.target).map_err(|error| { + CommandFailure::new(error.code(), error.to_string(), error.recoverable()) + })?; + let cancelled = Arc::new(AtomicBool::new(false)); + let mut probes = Vec::new(); + let mut linux_probes = Vec::new(); + for endpoint in [&validation.start, &validation.end] { + if let Some(probe) = endpoint.host.as_ref() + && !probes + .iter() + .any(|existing: &ResolvedProbe| existing.selector == probe.selector) + { + probes.push(probe.clone()); + } + if let Some(probe) = endpoint.linux.as_ref() + && !linux_probes + .iter() + .any(|existing: &ResolvedLinuxSelector| existing == probe) + { + linux_probes.push(probe.clone()); + } + } + let host_timeout = request + .options + .duration + .map_or(request.timeout, |duration| duration.min(request.timeout)); + let host_samples = request + .options + .samples + .unwrap_or(request.options.max_events); + let mut receivers = Vec::with_capacity(probes.len() + usize::from(!linux_probes.is_empty())); + let mut handles = probes + .into_iter() + .enumerate() + .map(|(index, probe)| { + let (ready, receiver) = mpsc::sync_channel(1); + receivers.push(receiver); + let offset = if probe.symbol.is_some() { + 0 + } else { + probe.file_offset + }; + let capture_request = UprobeRequest { + target: report.target.clone(), + binary: PathBuf::from(&probe.binary_path), + symbol: probe.symbol, + symbol_demangled: probe.symbol_demangled, + offset, + probe_kind: probe.probe_kind, + probe_id: u32::try_from(index + 1).expect("two endpoints fit u32"), + samples: host_samples, + timeout: host_timeout, + cancelled: Arc::clone(&cancelled), + ready: Some(ready), + }; + thread::spawn(move || { + uprobe::capture(&capture_request).map_err(|error| { + CommandFailure::new(error.code(), error.to_string(), error.recoverable()) + }) + }) + }) + .collect::>(); + if !linux_probes.is_empty() { + let (ready, receiver) = mpsc::sync_channel(1); + receivers.push(receiver); + let endpoint_count = linux_probes.len(); + let event_limit = linux_capture_event_limit(&request.options, endpoint_count); + let capture_request = LinuxCaptureRequest { + target: report.target.clone(), + probes: linux_probes, + event_limit, + capacity_limit: request.options.samples.is_none(), + timeout: host_timeout, + cancelled: Arc::clone(&cancelled), + ready: Some(ready), + }; + handles.push(thread::spawn(move || { + linux::capture(&capture_request).map_err(|error| { + CommandFailure::new(error.code(), error.to_string(), error.recoverable()) + }) + })); + } + let mut collectors = HostCollectors { cancelled, handles }; + wait_for_host_collectors(&mut collectors, receivers, request.timeout)?; + if let Err(error) = inspect::verify_target(&report.target) { + collectors.cancel(); + collectors.join()?; + return Err(CommandFailure::new( + error.code(), + error.to_string(), + error.recoverable(), + )); + } + Ok(collectors) +} + +pub(crate) fn linux_capture_event_limit(options: &MeasureOptions, endpoint_count: usize) -> usize { + options + .samples + .and_then(|samples| { + samples + .checked_mul(endpoint_count) + .and_then(|events| events.checked_add(endpoint_count)) + }) + .unwrap_or(options.max_events) + .min(options.max_events) +} + +pub(crate) fn wait_for_host_collectors( + collectors: &mut HostCollectors, + receivers: Vec>, + timeout: Duration, +) -> Result<(), CommandFailure> { + for receiver in receivers { + match receiver.recv_timeout(timeout) { + Ok(()) => {} + Err(mpsc::RecvTimeoutError::Disconnected) => { + collectors.cancel(); + return match collectors.join() { + Err(error) => Err(error), + Ok(_) => Err(CommandFailure::new( + ErrorCode::Internal, + "host collector exited before reporting readiness", + false, + )), + }; + } + Err(mpsc::RecvTimeoutError::Timeout) => { + collectors.cancel(); + collectors.join()?; + return Err(CommandFailure::new( + ErrorCode::SessionLimitExceeded, + "timed out waiting for host collector readiness", + true, + )); + } + } + } + Ok(()) +} + +pub(crate) fn refresh_live_sources( + collection: &mut LiveCollection, + request: &LiveMeasureRequest, +) -> Result<(), CommandFailure> { + if collection.host_captures.is_none() && collection.collectors.finished() { + collection.host_captures = Some(collection.collectors.join()?); + } + let Some(path) = collection.socket.as_deref() else { + return Ok(()); + }; + let remaining = collection + .deadline + .saturating_duration_since(Instant::now()); + if remaining <= SNAPSHOT_QUIET_PERIOD { + collection.collectors.cancel(); + return Ok(()); + } + let capture = cupti::snapshot( + path, + remaining, + &request.options.session_id, + collection.cuda_record_offset, + ) + .map_err(|error| CommandFailure::new(ErrorCode::CuptiNotAvailable, error.to_string(), true))?; + append_cuda_capture(collection, capture)?; + Ok(()) +} + +pub(crate) fn append_cuda_capture( + collection: &mut LiveCollection, + mut capture: cupti::CuptiCapture, +) -> Result<(), CommandFailure> { + if capture.record_offset != collection.cuda_record_offset { + return Err(CommandFailure::new( + ErrorCode::TraceExportFailed, + format!( + "CUPTI incremental capture started at record {}, expected {}", + capture.record_offset, collection.cuda_record_offset + ), + false, + )); + } + let returned = u64::try_from(capture.events.len()).map_err(|error| { + CommandFailure::new( + ErrorCode::TraceExportFailed, + format!("CUPTI incremental event count exceeds u64: {error}"), + false, + ) + })?; + let next_offset = capture.record_offset.checked_add(returned).ok_or_else(|| { + CommandFailure::new( + ErrorCode::TraceExportFailed, + "CUPTI incremental record offset overflowed", + false, + ) + })?; + if let Some(accumulated) = collection.latest_cuda.as_mut() { + if accumulated.record_capacity != capture.record_capacity { + return Err(CommandFailure::new( + ErrorCode::TraceExportFailed, + format!( + "CUPTI record capacity changed from {} to {} during capture", + accumulated.record_capacity, capture.record_capacity + ), + false, + )); + } + if capture.observed_records < accumulated.observed_records + || capture.agent_dropped_records < accumulated.agent_dropped_records + || capture.cupti_dropped_records < accumulated.cupti_dropped_records + || capture.unknown_records < accumulated.unknown_records + { + return Err(CommandFailure::new( + ErrorCode::TraceExportFailed, + "CUPTI incremental counters moved backwards", + false, + )); + } + accumulated.state = capture.state; + accumulated.stop_reason = capture.stop_reason; + accumulated.observed_records = capture.observed_records; + accumulated.agent_dropped_records = capture.agent_dropped_records; + accumulated.cupti_dropped_records = capture.cupti_dropped_records; + accumulated.dropped_records = capture.dropped_records; + accumulated.unknown_records = capture.unknown_records; + accumulated.events.append(&mut capture.events); + } else { + capture.record_offset = 0; + collection.latest_cuda = Some(capture); + } + collection.cuda_record_offset = next_offset; + Ok(()) +} + +pub(crate) fn aggregate_inventory_result( + capture: &cupti::CuptiCapture, + request: &LiveMeasureRequest, +) -> Result { + validate_aggregate_capture(capture)?; + let capture = gpu_aggregate::GpuAggregateCapture { + groups: capture + .aggregate_groups + .iter() + .map(|group| gpu_aggregate::RawGpuAggregateGroup { + activity: match group.activity { + cupti::CuptiAggregateActivity::Kernel => AggregateActivity::Kernel, + cupti::CuptiAggregateActivity::Memcpy => AggregateActivity::Memcpy, + cupti::CuptiAggregateActivity::Memset => AggregateActivity::Memset, + }, + name: group.name.clone(), + name_complete: group.name_complete, + device_id: group.device_id, + memcpy_kind: group.memcpy_kind.clone(), + count: group.count, + total_duration_ns: group.total_duration_ns, + min_duration_ns: group.min_duration_ns, + max_duration_ns: group.max_duration_ns, + total_bytes: group.total_bytes, + }) + .collect(), + observed_activities: capture.observed_records, + group_capacity: capture.record_capacity, + }; + let options = &request.options; + gpu_aggregate::analyze( + gpu_aggregate::GpuAggregateRequest { + session_id: options.session_id.clone(), + name: options.name.clone(), + start_selector: options.start_selector.clone(), + end_selector: options.end_selector.clone(), + }, + capture, + ) + .map_err(|error| CommandFailure::new(error.code(), error.to_string(), error.recoverable())) +} + +pub(crate) fn validate_aggregate_capture( + capture: &cupti::CuptiCapture, +) -> Result<(), CommandFailure> { + if capture.state == cupti::CuptiCaptureState::LimitReached { + return Err(CommandFailure::new( + ErrorCode::EventRateTooHigh, + "aggregate group table reached its configured capacity", + true, + ) + .with_detail("group_capacity", capture.record_capacity) + .with_detail("observed_activities", capture.observed_records) + .with_hint("narrow the selector or explicitly increase max-groups")); + } + if capture.state != cupti::CuptiCaptureState::Stopped + || capture.stop_reason != cupti::CuptiStopReason::Requested + { + return Err(CommandFailure::new( + ErrorCode::CuptiNotAvailable, + format!( + "aggregate CUPTI capture ended in state {:?} with reason {:?}", + capture.state, capture.stop_reason + ), + true, + )); + } + if capture.dropped_records != 0 { + return Err(CommandFailure::new( + ErrorCode::EventsDropped, + "aggregate CUPTI capture dropped activities", + true, + ) + .with_detail("dropped_activities", capture.dropped_records) + .with_hint("reduce capture duration or narrow the activity selector")); + } + if capture.unknown_records != 0 { + return Err(CommandFailure::new( + ErrorCode::TraceExportFailed, + "aggregate CUPTI capture contains unknown activities", + false, + ) + .with_detail("unknown_records", capture.unknown_records) + .with_hint("rebuild xprobe and the CUPTI Agent from the same release")); + } + if !capture.events.is_empty() { + return Err(CommandFailure::new( + ErrorCode::TraceExportFailed, + "aggregate CUPTI capture unexpectedly contains exact events", + false, + )); + } + Ok(()) +} + +pub(crate) fn injection_warning(collection: &LiveCollection) -> Warning { + let mut details = BTreeMap::new(); + if let Some(major) = collection.agent_major { + details.insert("cuda_major".to_owned(), serde_json::Value::from(major)); + } + if let Some(path) = collection.agent_path.as_deref() { + details.insert( + "agent_path".to_owned(), + serde_json::Value::from(path.display().to_string()), + ); + } + if let Some(path) = collection.cupti_path.as_deref() { + details.insert( + "cupti_path".to_owned(), + serde_json::Value::from(path.display().to_string()), + ); + } + Warning { + code: "CUPTI_AGENT_INJECTED".to_owned(), + message: "xprobe injected the CUPTI agent and left the shared object mapped".to_owned(), + details, + } +} + +pub(crate) fn evaluate_live_result( + collection: &LiveCollection, + request: &LiveMeasureRequest, + now: Instant, +) -> Result, CommandFailure> { + let Some(host) = collection.host_captures.as_ref() else { + return Ok(None); + }; + if collection.socket.is_some() && collection.latest_cuda.is_none() { + return Ok(None); + } + let mut execution = match measure_live_capture(host, collection, &request.options, false) { + Ok(execution) => execution, + Err(error) if error.code == ErrorCode::NoMatchedSamples => return Ok(None), + Err(error) => return Err(error), + }; + let reached_samples = request.options.samples.is_some_and(|samples| { + execution.result.measurement.samples.matched + >= u64::try_from(samples).expect("sample limit fits u64") + }); + let reached_duration = collection.collection_end.is_some_and(|end| now >= end); + if reached_samples || reached_duration { + inspect::verify_target(&collection.report.target).map_err(|error| { + CommandFailure::new(error.code(), error.to_string(), error.recoverable()) + })?; + return Ok(Some(execution)); + } + if now >= collection.deadline { + execution.result.status = SessionStatus::TimedOut; + return Ok(Some(execution)); + } + Ok(None) +} + +pub(crate) fn finish_live_timeout( + collection: &mut LiveCollection, + request: &LiveMeasureRequest, +) -> Result { + collection.collectors.cancel(); + if collection.host_captures.is_none() { + collection.host_captures = Some(collection.collectors.join()?); + } + let mut execution = measure_live_capture( + collection + .host_captures + .as_ref() + .expect("host captures were assigned"), + collection, + &request.options, + true, + ) + .map_err(|mut error| { + if error.code == ErrorCode::NoMatchedSamples { + let host_events = collection.host_captures.as_ref().map_or(0, |captures| { + captures.iter().map(|capture| capture.events.len()).sum() + }); + let cuda_events = collection.latest_cuda.as_ref().map_or(0, |capture| { + capture + .events + .iter() + .filter(|event| { + collection + .baseline_timestamp + .is_none_or(|start| event.timestamp_ns > start) + }) + .count() + }); + error.message = format!( + "no event pairs matched before the live measurement timeout \ + (host events: {host_events}, CUDA events: {cuda_events})" + ); + error + .with_detail("timed_out", true) + .with_detail("host_events", host_events as u64) + .with_detail("cuda_events", cuda_events as u64) + } else { + error + } + })?; + execution.result.status = SessionStatus::TimedOut; + Ok(execution) +} + +pub(crate) fn measure_live_capture( + host_captures: &[HostCaptureResult], + collection: &LiveCollection, + options: &MeasureOptions, + require_complete: bool, +) -> Result { + let capture = completed_live_capture(host_captures, collection, &options.session_id)?; + reject_incomplete_capture(&capture, true, require_complete) + .map_err(|error| error.with_artifact_events(capture.events.clone()))?; + if capture.unknown_records != 0 { + return Err(CommandFailure::new( + ErrorCode::TraceExportFailed, + format!( + "live capture contains {} unknown CUPTI activity records", + capture.unknown_records + ), + false, + ) + .with_detail("unknown_records", capture.unknown_records) + .with_hint("rebuild xprobe and the CUPTI Agent from the same release") + .with_artifact_events(capture.events)); + } + let options = MeasureOptions { + dropped_events: capture.dropped_records, + ..options.clone() + }; + let mut result = measure(&capture.events, &options).map_err(|error| { + measurement_failure(&error, &options, &capture.events, require_complete) + })?; + apply_collection_summary(&mut result, &capture); + Ok(MeasurementExecution { + result, + events: capture.events, + }) +} + +pub(crate) fn completed_live_capture( + host_captures: &[HostCaptureResult], + collection: &LiveCollection, + session_id: &str, +) -> Result { + let mut captures = host_captures + .iter() + .map(|capture| completed::CompletedCapture { + dropped_records: capture.dropped, + unknown_records: 0, + record_limit_reached: capture.record_limit_reached.then_some(capture.captured), + capture_failed: false, + cupti: None, + events: capture.events.clone(), + }) + .collect::>(); + if let Some(cuda) = collection.latest_cuda.as_ref() { + let dropped_records = cuda + .dropped_records + .checked_sub(collection.baseline_dropped) + .ok_or_else(|| { + CommandFailure::new( + ErrorCode::TraceExportFailed, + "CUPTI dropped-record counter moved backwards", + false, + ) + })?; + let unknown_records = cuda + .unknown_records + .checked_sub(collection.baseline_unknown) + .ok_or_else(|| { + CommandFailure::new( + ErrorCode::TraceExportFailed, + "CUPTI unknown-record counter moved backwards", + false, + ) + })?; + let observed_records = cuda + .observed_records + .checked_sub(collection.baseline_observed) + .ok_or_else(|| { + CommandFailure::new( + ErrorCode::TraceExportFailed, + "CUPTI observed-record counter moved backwards", + false, + ) + })?; + captures.push(completed::CompletedCapture { + dropped_records, + unknown_records, + record_limit_reached: (cuda.state == cupti::CuptiCaptureState::LimitReached) + .then_some(cuda.record_capacity), + capture_failed: matches!( + cuda.state, + cupti::CuptiCaptureState::Idle | cupti::CuptiCaptureState::Failed + ), + cupti: Some(completed::CompletedCuptiStatistics { + complete: cuda.state == cupti::CuptiCaptureState::Stopped, + record_capacity: cuda.record_capacity, + observed_records, + dropped_records, + }), + events: cuda + .events + .iter() + .filter(|event| { + collection + .baseline_timestamp + .is_none_or(|start| event.timestamp_ns > start) + }) + .cloned() + .collect(), + }); + } + completed::merge(captures, session_id).map_err(|error| { + CommandFailure::new(ErrorCode::TraceExportFailed, error.to_string(), false) + }) +} + +pub(crate) fn reject_incomplete_capture( + capture: &completed::CompletedCapture, + live: bool, + require_complete: bool, +) -> Result<(), CommandFailure> { + if capture.capture_failed { + return Err(CommandFailure::new( + ErrorCode::CuptiNotAvailable, + format!( + "{}CUPTI capture entered a failed state", + if live { "live " } else { "" } + ), + true, + ) + .with_detail("capture_state", "failed") + .with_hint("inspect target CUDA/CUPTI compatibility and retry a fresh bounded capture")); + } + if let Some(capacity) = capture.record_limit_reached { + let mut failure = CommandFailure::new( + ErrorCode::EventRateTooHigh, + format!("capture reached its configured limit of {capacity} records"), + true, + ) + .with_detail("record_capacity", capacity) + .with_hint("narrow the selectors or explicitly increase max-events"); + if let Some(cupti) = capture.cupti.as_ref() { + failure = failure + .with_detail("observed_records", cupti.observed_records) + .with_detail("retained_records", capture.events.len() as u64) + .with_detail("dropped_records", cupti.dropped_records); + } + return Err(failure); + } + if require_complete && capture.cupti.as_ref().is_some_and(|cupti| !cupti.complete) { + return Err(CommandFailure::new( + ErrorCode::CuptiNotAvailable, + format!( + "{}CUPTI capture was not stopped before measurement", + if live { "live " } else { "" } + ), + true, + ) + .with_detail("capture_state", "active") + .with_hint("stop the CUPTI capture before using it as completed input")); + } + Ok(()) +} + +pub(crate) fn apply_collection_summary( + result: &mut MeasurementResult, + capture: &completed::CompletedCapture, +) { + let Some(cupti) = capture.cupti.as_ref() else { + return; + }; + let retained_records = result.collection.cuda_events; + let buffer_utilization = utilization(retained_records, cupti.record_capacity); + result.collection.cupti = Some(CuptiCollectionSummary { + record_capacity: cupti.record_capacity, + observed_records: cupti.observed_records, + retained_records, + dropped_records: cupti.dropped_records, + buffer_utilization, + }); +} + +#[allow(clippy::cast_precision_loss)] +pub(crate) fn utilization(retained: u64, capacity: u64) -> f64 { + if capacity == 0 { + 0.0 + } else { + retained as f64 / capacity as f64 + } +} + +pub(crate) const fn match_policy_name(policy: MatchPolicy) -> &'static str { + match policy { + MatchPolicy::Exact => "exact", + MatchPolicy::FirstAfter => "first-after", + MatchPolicy::Nearest => "nearest", + MatchPolicy::StackNested => "stack-nested", + MatchPolicy::StreamOrder => "stream-order", + } +} diff --git a/xprobe/cli/src/main.rs b/xprobe/cli/src/main.rs index 1adc0e3..5573467 100644 --- a/xprobe/cli/src/main.rs +++ b/xprobe/cli/src/main.rs @@ -1,490 +1,49 @@ use std::{ - collections::BTreeMap, - fs::{self, OpenOptions}, - io::Write, - os::unix::fs::OpenOptionsExt, + fs::{self}, path::{Path, PathBuf}, process::ExitCode, - sync::{ - Arc, - atomic::{AtomicBool, AtomicU64, Ordering}, - mpsc, - }, - thread::{self, JoinHandle}, - time::{Duration, Instant}, + sync::atomic::AtomicU64, + time::Duration, }; static EXPORT_SEQUENCE: AtomicU64 = AtomicU64::new(0); -use clap::{Args, Parser, Subcommand, ValueEnum}; -use xprobe_collector::{ - completed, - cpu_sampling::{self as cpu_collector, CpuSamplingRequest}, - cupti, - linux::{self, LinuxCaptureRequest}, - syscall_aggregate::{self as syscall_collector, SyscallAggregateRequest}, - uprobe::{self, UprobeRequest}, -}; -use xprobe_core::{ - cpu_sampling, - cpu_sampling_analysis::{self, CpuCaptureData, RawCpuSample}, - cupti_compat, discover, doctor, gpu_aggregate, inject, inspect, resolve, syscall_aggregate, - validate, -}; -use xprobe_correlator::{MeasureError, MeasureOptions, measure}; +use clap::Parser; +use xprobe_collector::completed; +use xprobe_core::{cpu_sampling, discover, doctor, inspect, syscall_aggregate, validate}; +use xprobe_correlator::{MeasureOptions, measure}; use xprobe_exporter::{events_to_chrome_trace, events_to_jsonl}; use xprobe_protocol::{ - AggregateActivity, AggregateInventoryResult, CapabilityReport, CheckResult, ClockDomain, - CpuSampleEvent, CpuSampleInventoryResult, CpuSamplingSpec, CpuSamplingValidationResult, - CuptiCollectionSummary, DiscoveryResult, ErrorCode, ErrorResponse, Event, EventSource, - EventType, ExportFormat, HostCaptureResult, MatchPolicy, MeasurementMode, MeasurementResult, - MeasurementSpec, ProcessReport, ResolvedLinuxSelector, ResolvedProbe, SchemaVersion, - SessionStatus, SyscallAggregateResult, SyscallAggregateSpec, SyscallAggregateValidationResult, - TargetIdentity, TraceExportResult, ValidationResult, Warning, XprobeError, + AggregateInventoryResult, ClockDomain, CpuSamplingValidationResult, ErrorCode, Event, + MatchPolicy, MeasurementMode, MeasurementResult, MeasurementSpec, + SyscallAggregateValidationResult, }; -#[derive(Debug, Parser)] -#[command(name = "xprobe", version, about = "Runtime host-to-GPU latency probe")] -struct Cli { - #[command(subcommand)] - command: Command, -} - -#[derive(Debug, Subcommand)] -enum Command { - /// Inspect local tracing and GPU capabilities. - Doctor(DoctorArgs), - /// List NVML-confirmed CUDA context holders under a process-tree root. - Discover(DiscoverArgs), - /// Inspect a target process without attaching probes. - #[command(hide = true)] - Inspect(InspectArgs), - /// Resolve a userspace event selector against a target process. - #[command(hide = true)] - Resolve(ResolveArgs), - /// Validate two event selectors and their correlation policy. - Validate(ValidateArgs), - /// Measure latency between two events in files or a running process. - Measure(Box), - /// Run a live measurement from a versioned specification. - #[command(hide = true)] - Trace(TraceArgs), - /// Export completed captures to a stable trace format. - #[command(hide = true)] - Export(ExportArgs), - /// Run bounded event collectors. - #[command(hide = true)] - Capture(DevArgs), - /// Run low-level development probes. - #[command(hide = true)] - Dev(DevArgs), -} - -#[derive(Debug, Args)] -struct DevArgs { - #[command(subcommand)] - command: DevCommand, -} - -#[derive(Debug, Args)] -struct TraceArgs { - /// Versioned `MeasurementSpec` JSON file. - #[arg(long)] - spec: PathBuf, - - /// Unix socket exposed by the target's xprobe CUPTI agent. - #[arg(long, hide = true)] - cupti_socket: Option, - - /// Emit only the versioned JSON result on stdout. - #[arg(long)] - json: bool, - - /// Disable colored output. - #[arg(long)] - no_color: bool, - - /// Never wait for user input. - #[arg(long)] - non_interactive: bool, -} - -#[derive(Debug, Clone, Copy, ValueEnum)] -enum ExportFormatArg { - Jsonl, - Chrome, -} - -#[derive(Debug, Args)] -struct ExportArgs { - /// Completed CUPTI binary, host capture JSON, or Event JSONL; repeat to merge. - #[arg(long, required = true)] - input: Vec, - - /// Artifact format. - #[arg(long, value_enum)] - format: ExportFormatArg, - - /// Destination artifact path. - #[arg(long)] - output: PathBuf, - - /// Emit only the versioned export result on stdout. - #[arg(long)] - json: bool, - - /// Disable colored output. - #[arg(long)] - no_color: bool, - - /// Never wait for user input. - #[arg(long)] - non_interactive: bool, -} - -#[derive(Debug, Subcommand)] -enum DevCommand { - /// Capture entries to one userspace function. - Uprobe(UprobeArgs), - /// Decode an xprobe CUPTI capture as Event JSONL. - Cupti(CuptiArgs), -} - -#[derive(Debug, Clone, Copy, Args)] -struct DoctorArgs { - /// Emit only the versioned JSON result on stdout. - #[arg(long)] - json: bool, - - /// Disable colored output. - #[arg(long)] - no_color: bool, - - /// Never wait for user input. - #[arg(long)] - non_interactive: bool, -} - -#[derive(Debug, Clone, Copy, Args)] -struct DiscoverArgs { - /// Target process ID. - #[arg(long)] - pid: u32, - - /// Maximum number of CUDA worker candidates to return. - #[arg(long, default_value_t = 200)] - limit: usize, - - /// Emit only the versioned JSON result on stdout. - #[arg(long)] - json: bool, - - /// Disable colored output. - #[arg(long)] - no_color: bool, - - /// Never wait for user input. - #[arg(long)] - non_interactive: bool, -} - -#[derive(Debug, Clone, Copy, Args)] -struct InspectArgs { - /// Target process ID. - #[arg(long)] - pid: u32, - - /// Emit only the versioned JSON result on stdout. - #[arg(long)] - json: bool, - - /// Disable colored output. - #[arg(long)] - no_color: bool, - - /// Never wait for user input. - #[arg(long)] - non_interactive: bool, -} - -#[derive(Debug, Args)] -struct ResolveArgs { - /// Target process ID. - #[arg(long)] - pid: u32, - - /// Event selector for a mapped ELF symbol, C++ signature, or file offset. - #[arg(long, alias = "event")] - selector: String, - - /// Emit only the versioned JSON result on stdout. - #[arg(long)] - json: bool, - - /// Disable colored output. - #[arg(long)] - no_color: bool, - - /// Never wait for user input. - #[arg(long)] - non_interactive: bool, -} - -#[derive(Debug, Args)] -struct ValidateArgs { - /// Target process ID. - #[arg(long)] - pid: u32, - - /// Start event selector. - #[arg(long)] - from: Option, - - /// End event selector. - #[arg(long)] - to: Option, - - /// Correlation policy. - #[arg(long = "match")] - match_policy: Option, - - /// Validate bounded PID-scoped CPU sampling instead of an event pair. - #[arg(long, conflicts_with_all = ["from", "to", "match_policy", "syscall_aggregate"])] - cpu_sample: bool, - - /// Validate bounded PID-scoped syscall aggregation instead of an event pair. - #[arg(long, conflicts_with_all = ["from", "to", "match_policy", "cpu_sample"])] - syscall_aggregate: bool, - - #[command(flatten)] - output: CommonOutputArgs, -} - -#[derive(Debug, Args)] -struct MeasureArgs { - /// Versioned measurement JSON for a live target. - #[arg(long, conflicts_with_all = ["input", "pid", "from", "to", "match_policy", "cpu_sample", "syscall_aggregate", "samples", "duration_ms", "timeout_ms", "max_events", "max_samples", "frequency_hz", "stack_depth", "max_threads", "max_inflight", "aggregate", "max_groups", "name"])] - spec: Option, +mod args; +mod dev; +mod failure; +mod inventory; +mod live; +mod render; - /// Completed CUPTI binary, host capture JSON, or Event JSONL; repeat to merge. - #[arg(long)] - input: Vec, - - /// Collect from a running target process instead of completed capture files. - #[arg(long)] - pid: Option, - - /// Unix socket exposed by the target's xprobe CUPTI agent. - #[arg(long, hide = true)] - cupti_socket: Option, - - /// CUPTI agent shared object used for automatic online injection. - #[arg(long)] - agent: Option, - - /// Write the complete bounded event capture to this file. - #[arg(long)] - events_out: Option, - - /// Evidence event format; defaults to jsonl when --events-out is set. - #[arg(long, value_enum)] - format: Option, - - /// Bound foreground collection and cleanup to this many milliseconds. - #[arg(long, default_value_t = 30_000)] - timeout_ms: u64, - - /// Start event selector. - #[arg(long)] - from: Option, - - /// End event selector. - #[arg(long)] - to: Option, - - /// Correlation policy. - #[arg(long = "match")] - match_policy: Option, - - /// Collect a bounded PID-scoped CPU hotspot inventory. - #[arg(long, conflicts_with_all = ["input", "from", "to", "match_policy", "samples", "aggregate", "events_out", "format", "cupti_socket", "agent"])] - cpu_sample: bool, - - /// Target CPU sampling frequency in hertz. - #[arg( - long, - requires = "cpu_sample", - default_value_if("cpu_sample", "true", "99") - )] - frequency_hz: Option, - - /// Bound all observed CPU sample records. - #[arg( - long, - requires = "cpu_sample", - default_value_if("cpu_sample", "true", "10000") - )] - max_samples: Option, - - /// Bound user-space callchain depth. - #[arg( - long, - requires = "cpu_sample", - default_value_if("cpu_sample", "true", "64") - )] - stack_depth: Option, - - /// Bound target threads attached by CPU sampling. - #[arg( - long, - requires = "cpu_sample", - default_value_if("cpu_sample", "true", "1024") - )] - max_threads: Option, - - /// Aggregate all syscall lifecycles in bounded BPF maps. - #[arg(long, conflicts_with_all = ["input", "from", "to", "match_policy", "samples", "cpu_sample", "aggregate", "events_out", "format", "cupti_socket", "agent"])] - syscall_aggregate: bool, - - /// Bound concurrent syscall lifecycle starts retained in BPF. - #[arg( - long, - requires = "syscall_aggregate", - default_value_if("syscall_aggregate", "true", "1024") - )] - max_inflight: Option, - - /// Stop after this many matched samples. - #[arg(long)] - samples: Option, - - /// Restrict events to this many milliseconds from the first selected event. - #[arg(long)] - duration_ms: Option, - - /// Reject captures containing more events than this limit. - #[arg(long, default_value_t = 100_000)] - max_events: usize, - - /// Return bounded activity aggregates instead of exact event evidence. - #[arg(long, conflicts_with_all = ["input", "samples", "events_out", "format"])] - aggregate: bool, - - /// Bound distinct groups retained by aggregate or CPU inventory mode. - #[arg( - long, - default_value_if("aggregate", "true", "4096"), - default_value_if("cpu_sample", "true", "256"), - default_value_if("syscall_aggregate", "true", "256") - )] - max_groups: Option, - - /// Optional measurement name. - #[arg(long)] - name: Option, - - #[command(flatten)] - output: CommonOutputArgs, -} - -#[derive(Debug, Clone, Copy, Args)] -struct CommonOutputArgs { - /// Emit only the versioned JSON result on stdout. - #[arg(long)] - json: bool, - - /// Disable colored output. - #[arg(long)] - no_color: bool, - - /// Never wait for user input. - #[arg(long)] - non_interactive: bool, -} - -#[derive(Debug, Args)] -struct UprobeArgs { - /// Target process ID. - #[arg(long)] - pid: u32, - - /// Executable or mapped shared library containing the symbol. - #[arg(long)] - binary: PathBuf, - - /// ELF function symbol to attach to. - #[arg(long)] - symbol: String, - - /// Attach at function return instead of function entry. - #[arg(long = "return")] - return_probe: bool, - - /// Probe identifier stored in each event. - #[arg(long, default_value_t = 1)] - probe_id: u32, - - /// Stop after this many events. - #[arg(long, default_value_t = 1)] - samples: usize, - - /// Stop waiting after this many milliseconds. - #[arg(long, default_value_t = 5_000)] - timeout_ms: u64, - - #[command(flatten)] - output: EventOutputArgs, - - /// Disable colored output. - #[arg(long)] - no_color: bool, - - /// Never wait for user input. - #[arg(long)] - non_interactive: bool, -} - -#[derive(Debug, Clone, Copy, Args)] -struct EventOutputArgs { - /// Emit only the versioned JSON result on stdout. - #[arg(long)] - json: bool, - - /// Emit one versioned Event JSON object per line. - #[arg(long, conflicts_with = "json")] - jsonl: bool, -} - -#[derive(Debug, Args)] -struct CuptiArgs { - /// Raw binary capture written by the xprobe CUPTI agent. - #[arg(long, required_unless_present = "socket", conflicts_with = "socket")] - input: Option, - - /// Unix socket exposed by a running xprobe CUPTI agent. - #[arg(long, required_unless_present = "input", conflicts_with = "input")] - socket: Option, - - /// Stop waiting for an online snapshot after this many milliseconds. - #[arg(long, default_value_t = 10_000)] - timeout_ms: u64, - - /// Session identifier written into every Event. - #[arg(long)] - session_id: String, - - /// Emit one versioned Event JSON object per line. - #[arg(long)] - json: bool, - - /// Disable colored output. - #[arg(long)] - no_color: bool, - - /// Never wait for user input. - #[arg(long)] - non_interactive: bool, -} +use args::{ + Cli, Command, CommonOutputArgs, DiscoverArgs, DoctorArgs, ExportFormatArg, MeasureArgs, + ValidateArgs, +}; +use dev::{run_capture, run_export, run_inspect, run_resolve, run_trace, write_export_file}; +use failure::{CommandFailure, emit_command_failure, emit_error, measurement_failure}; +use inventory::{ + run_cpu_sampling_file_spec, run_direct_cpu_sampling, run_direct_syscall_aggregate, + run_syscall_aggregate_file_spec, +}; +use live::{ + LiveMeasureRequest, apply_collection_summary, collect_live_aggregate, collect_live_measurement, + match_policy_name, reject_incomplete_capture, +}; +use render::{ + print_check, print_discovery_result, print_human_report, print_measurement_result, + print_validation_result, +}; fn main() -> ExitCode { let cli = Cli::parse(); @@ -501,372 +60,119 @@ fn main() -> ExitCode { } } -fn run_capture(args: DevArgs) -> ExitCode { - match args.command { - DevCommand::Uprobe(args) => run_uprobe(args), - DevCommand::Cupti(args) => run_cupti(args), - } -} - -fn run_trace(args: TraceArgs) -> ExitCode { - let TraceArgs { +#[allow(clippy::too_many_lines)] +fn run_measure(args: MeasureArgs) -> ExitCode { + let MeasureArgs { spec, + input, + pid, cupti_socket, - json, - no_color: _, - non_interactive: _, + agent, + events_out, + format, + timeout_ms, + from, + to, + match_policy, + cpu_sample, + syscall_aggregate, + frequency_hz, + max_samples, + stack_depth, + max_threads, + max_inflight, + samples, + duration_ms, + max_events, + aggregate, + max_groups, + name, + output: + CommonOutputArgs { + json, + no_color: _, + non_interactive: _, + }, } = args; - let bytes = match fs::read(&spec) { - Ok(bytes) => bytes, - Err(error) => { + if let Some(spec) = spec { + return run_measure_spec( + &spec, + cupti_socket, + agent, + events_out.as_deref(), + format, + json, + ); + } + if cpu_sample { + let Some(pid) = pid else { return emit_error( - ErrorCode::TraceExportFailed, - format!("failed to read {}: {error}", spec.display()), - false, + ErrorCode::InvalidEventSelector, + "--cpu-sample requires --pid".to_owned(), + true, json, ); - } - }; - let spec: MeasurementSpec = match serde_json::from_slice(&bytes) { - Ok(spec) => spec, - Err(error) => { + }; + let Some(duration_ms) = duration_ms else { return emit_error( - ErrorCode::InvalidEventSelector, - format!("invalid MeasurementSpec {}: {error}", spec.display()), + ErrorCode::SessionLimitExceeded, + "--cpu-sample requires --duration-ms".to_owned(), true, json, ); - } - }; - if spec.measurement_mode != MeasurementMode::Exact || spec.max_groups.is_some() { - return emit_error( - ErrorCode::InvalidEventSelector, - "legacy trace accepts only exact MeasurementSpec requests".to_owned(), - true, + }; + return run_direct_cpu_sampling( + pid, + name, + frequency_hz.expect("clap supplies CPU sampling frequency"), + duration_ms, + timeout_ms, + max_samples.expect("clap supplies CPU sample capacity"), + u64::try_from(max_groups.expect("clap supplies CPU group capacity")) + .unwrap_or(u64::MAX), + stack_depth.expect("clap supplies CPU stack depth"), + max_threads.expect("clap supplies CPU thread capacity"), json, ); } - let samples = match spec.samples.map(usize::try_from).transpose() { - Ok(samples) => samples, - Err(error) => { + if syscall_aggregate { + let Some(pid) = pid else { + return emit_error( + ErrorCode::InvalidEventSelector, + "--syscall-aggregate requires --pid".to_owned(), + true, + json, + ); + }; + let Some(duration_ms) = duration_ms else { return emit_error( ErrorCode::SessionLimitExceeded, - format!("MeasurementSpec samples exceed this platform: {error}"), + "--syscall-aggregate requires --duration-ms".to_owned(), true, json, ); - } - }; - let Some(max_events) = spec.max_events else { + }; + return run_direct_syscall_aggregate( + pid, + name, + duration_ms, + timeout_ms, + u64::try_from(max_groups.expect("clap supplies syscall group capacity")) + .unwrap_or(u64::MAX), + max_inflight.expect("clap supplies syscall inflight capacity"), + json, + ); + } + if max_groups.is_some() && !aggregate { return emit_error( - ErrorCode::SessionLimitExceeded, - "exact MeasurementSpec requires max_events".to_owned(), + ErrorCode::InvalidEventSelector, + "--max-groups requires --aggregate, --cpu-sample, or --syscall-aggregate".to_owned(), true, json, ); - }; - let max_events = match usize::try_from(max_events) { - Ok(max_events) => max_events, - Err(error) => { - return emit_error( - ErrorCode::SessionLimitExceeded, - format!("MeasurementSpec max_events exceed this platform: {error}"), - true, - json, - ); - } - }; - let request = LiveMeasureRequest { - pid: spec.target.pid, - expected_target: Some(spec.target), - cupti_socket, - agent_path: None, - timeout: Duration::from_millis(spec.timeout_ms), - match_policy_text: match_policy_name(spec.match_policy), - mode: MeasurementMode::Exact, - options: MeasureOptions { - session_id: format!("xp_trace_{}", std::process::id()), - name: spec.name, - start_selector: spec.start_selector, - end_selector: spec.end_selector, - match_policy: spec.match_policy, - samples, - duration: spec.duration_ms.map(Duration::from_millis), - max_events, - dropped_events: 0, - }, - }; - match collect_live_measurement(&request) { - Ok(execution) => emit_measurement(&execution.result, json), - Err(error) => emit_command_failure(error, json), - } -} - -fn run_export(args: ExportArgs) -> ExitCode { - let ExportArgs { - input, - format, - output, - json, - no_color: _, - non_interactive: _, - } = args; - let session_id = format!("xp_export_{}", std::process::id()); - let capture = match load_completed_inputs(&input, &session_id) { - Ok(capture) => capture, - Err(error) => return emit_command_failure(error, json), - }; - let (format, artifact) = match format { - ExportFormatArg::Jsonl => ( - ExportFormat::Jsonl, - events_to_jsonl(&capture.events).map_err(|error| error.to_string()), - ), - ExportFormatArg::Chrome => ( - ExportFormat::Chrome, - events_to_chrome_trace(&capture.events).map_err(|error| error.to_string()), - ), - }; - let artifact = match artifact { - Ok(artifact) => artifact, - Err(error) => { - return emit_error(ErrorCode::TraceExportFailed, error, false, json); - } - }; - if let Err(error) = write_export_file(&output, artifact.as_bytes()) { - return emit_error( - ErrorCode::TraceExportFailed, - format!("failed to write {}: {error}", output.display()), - false, - json, - ); - } - let event_count = u64::try_from(capture.events.len()).expect("event count fits u64"); - let result = TraceExportResult { - schema_version: SchemaVersion::current(), - ok: true, - format, - output: output.to_string_lossy().into_owned(), - event_count, - }; - if json { - println!( - "{}", - serde_json::to_string_pretty(&result).expect("export result must serialize") - ); - } else { - println!( - "Exported {} events to {} ({:?})", - result.event_count, result.output, result.format - ); - } - ExitCode::SUCCESS -} - -fn load_completed_inputs( - input: &[PathBuf], - session_id: &str, -) -> Result { - let mut captures = Vec::with_capacity(input.len()); - for path in input { - let bytes = fs::read(path).map_err(|error| { - CommandFailure::new( - ErrorCode::TraceExportFailed, - format!("failed to read {}: {error}", path.display()), - false, - ) - })?; - captures.push(completed::decode(&bytes, session_id).map_err(|error| { - CommandFailure::new( - ErrorCode::TraceExportFailed, - format!("failed to decode {}: {error}", path.display()), - false, - ) - })?); - } - let capture = completed::merge(captures, session_id).map_err(|error| { - CommandFailure::new(ErrorCode::TraceExportFailed, error.to_string(), false) - })?; - if capture.unknown_records != 0 { - return Err(CommandFailure::new( - ErrorCode::TraceExportFailed, - format!( - "capture contains {} unknown CUPTI activity records", - capture.unknown_records - ), - false, - ) - .with_detail("unknown_records", capture.unknown_records) - .with_hint("rebuild xprobe and the CUPTI Agent from the same release")); - } - Ok(capture) -} - -fn write_export_file(path: &Path, bytes: &[u8]) -> Result<(), String> { - let parent = path - .parent() - .filter(|parent| !parent.as_os_str().is_empty()) - .unwrap_or_else(|| Path::new(".")); - let file_name = path - .file_name() - .ok_or_else(|| format!("artifact path {} has no file name", path.display()))?; - let temporary = parent.join(format!( - ".{}.xprobe-{}-{}.tmp", - file_name.to_string_lossy(), - std::process::id(), - EXPORT_SEQUENCE.fetch_add(1, Ordering::Relaxed) - )); - let mut file = OpenOptions::new() - .write(true) - .create_new(true) - .mode(0o600) - .open(&temporary) - .map_err(|error| format!("failed to create {}: {error}", temporary.display()))?; - if let Err(error) = file.write_all(bytes).and_then(|()| file.sync_all()) { - let original = format!("failed to write {}: {error}", temporary.display()); - return Err(remove_temporary_artifact(&temporary, original)); - } - drop(file); - if let Err(error) = fs::rename(&temporary, path) { - let original = format!( - "failed to atomically replace {} with {}: {error}", - path.display(), - temporary.display() - ); - return Err(remove_temporary_artifact(&temporary, original)); - } - fs::File::open(parent) - .and_then(|directory| directory.sync_all()) - .map_err(|error| { - format!( - "failed to sync artifact directory {}: {error}", - parent.display() - ) - }) -} - -fn remove_temporary_artifact(path: &Path, original: String) -> String { - match fs::remove_file(path) { - Ok(()) => original, - Err(cleanup) => format!( - "{original}; failed to remove temporary artifact {}: {cleanup}", - path.display() - ), - } -} - -#[allow(clippy::too_many_lines)] -fn run_measure(args: MeasureArgs) -> ExitCode { - let MeasureArgs { - spec, - input, - pid, - cupti_socket, - agent, - events_out, - format, - timeout_ms, - from, - to, - match_policy, - cpu_sample, - syscall_aggregate, - frequency_hz, - max_samples, - stack_depth, - max_threads, - max_inflight, - samples, - duration_ms, - max_events, - aggregate, - max_groups, - name, - output: - CommonOutputArgs { - json, - no_color: _, - non_interactive: _, - }, - } = args; - if let Some(spec) = spec { - return run_measure_spec( - &spec, - cupti_socket, - agent, - events_out.as_deref(), - format, - json, - ); - } - if cpu_sample { - let Some(pid) = pid else { - return emit_error( - ErrorCode::InvalidEventSelector, - "--cpu-sample requires --pid".to_owned(), - true, - json, - ); - }; - let Some(duration_ms) = duration_ms else { - return emit_error( - ErrorCode::SessionLimitExceeded, - "--cpu-sample requires --duration-ms".to_owned(), - true, - json, - ); - }; - return run_direct_cpu_sampling( - pid, - name, - frequency_hz.expect("clap supplies CPU sampling frequency"), - duration_ms, - timeout_ms, - max_samples.expect("clap supplies CPU sample capacity"), - u64::try_from(max_groups.expect("clap supplies CPU group capacity")) - .unwrap_or(u64::MAX), - stack_depth.expect("clap supplies CPU stack depth"), - max_threads.expect("clap supplies CPU thread capacity"), - json, - ); - } - if syscall_aggregate { - let Some(pid) = pid else { - return emit_error( - ErrorCode::InvalidEventSelector, - "--syscall-aggregate requires --pid".to_owned(), - true, - json, - ); - }; - let Some(duration_ms) = duration_ms else { - return emit_error( - ErrorCode::SessionLimitExceeded, - "--syscall-aggregate requires --duration-ms".to_owned(), - true, - json, - ); - }; - return run_direct_syscall_aggregate( - pid, - name, - duration_ms, - timeout_ms, - u64::try_from(max_groups.expect("clap supplies syscall group capacity")) - .unwrap_or(u64::MAX), - max_inflight.expect("clap supplies syscall inflight capacity"), - json, - ); - } - if max_groups.is_some() && !aggregate { - return emit_error( - ErrorCode::InvalidEventSelector, - "--max-groups requires --aggregate, --cpu-sample, or --syscall-aggregate".to_owned(), - true, - json, - ); - } - let (from, to, match_policy) = match parse_measure_selection(from, to, match_policy) { - Ok(selection) => selection, + } + let (from, to, match_policy) = match parse_measure_selection(from, to, match_policy) { + Ok(selection) => selection, Err(error) => { return emit_command_failure(error, json); } @@ -932,389 +238,13 @@ fn run_measure(args: MeasureArgs) -> ExitCode { ) } -#[allow(clippy::too_many_arguments)] -fn run_direct_cpu_sampling( - pid: u32, - name: Option, - frequency_hz: u64, - duration_ms: u64, - timeout_ms: u64, - max_samples: u64, - max_groups: u64, - stack_depth: u32, - max_threads: u32, - json: bool, -) -> ExitCode { - let report = match inspect::run(pid) { - Ok(report) => report, - Err(error) => { - return emit_error(error.code(), error.to_string(), error.recoverable(), json); - } - }; - let spec = CpuSamplingSpec { - schema_version: SchemaVersion::current(), - name, - target: report.target.clone(), - sample_event: CpuSampleEvent::CpuClock, - frequency_hz, - duration_ms, - timeout_ms, - max_samples, - max_groups, - stack_depth, - max_threads, - }; - run_cpu_sampling_spec(&spec, Some(report), json) -} - -fn run_cpu_sampling_spec( - spec: &CpuSamplingSpec, - inspected: Option, - json: bool, -) -> ExitCode { - let report = match inspected.map_or_else(|| inspect::run(spec.target.pid), Ok) { - Ok(report) => report, - Err(error) => { - return emit_error(error.code(), error.to_string(), error.recoverable(), json); - } - }; - if let Err(error) = inspect::verify_target(&spec.target) { - return emit_error(error.code(), error.to_string(), error.recoverable(), json); - } - let validation = match cpu_sampling::validate(&report) { - Ok(validation) => validation, - Err(error) => { - return emit_error(error.code(), error.to_string(), error.recoverable(), json); - } - }; - if !validation.valid { - let message = validation.issues.first().map_or_else( - || "CPU sampling requirements are not satisfied".to_owned(), - |issue| issue.message.clone(), - ); - return emit_error(ErrorCode::PermissionDenied, message, true, json); - } - let request = match cpu_collector_request(spec) { - Ok(request) => request, - Err(error) => return emit_command_failure(error, json), - }; - let capture = match cpu_collector::collect(&request) { - Ok(capture) => capture, - Err(error) => return emit_command_failure(cpu_collector_failure(&error), json), - }; - if capture.observed_samples == 0 { - return emit_error( - ErrorCode::NoMatchedSamples, - "CPU sampling completed without observing an on-CPU sample".to_owned(), +fn validate_measure_source(pid: Option, input: &[PathBuf]) -> Result<(), CommandFailure> { + if pid.is_some() && !input.is_empty() { + return Err(CommandFailure::new( + ErrorCode::InvalidEventSelector, + "--pid and --input select different collection modes", true, - json, - ); - } - let capture = CpuCaptureData { - samples: capture - .stacks - .into_iter() - .map(|stack| RawCpuSample { - thread_id: stack.thread_id, - addresses: stack.addresses, - }) - .collect(), - observed_samples: capture.observed_samples, - lost_samples: capture.lost_samples, - truncated_stacks: capture.truncated_stacks, - throttled_records: capture.throttled_records, - threads_observed: capture.threads_observed, - threads_attached: capture.threads_attached, - skipped_threads: capture.skipped_threads, - capacity_reached: capture.capacity_reached, - python_status: validation.python_status, - }; - let result = match cpu_sampling_analysis::analyze( - &report, - spec, - format!("xp_cpu_{}", std::process::id()), - &capture, - ) { - Ok(result) => result, - Err(error) => { - return emit_error(error.code(), error.to_string(), error.recoverable(), json); - } - }; - emit_cpu_inventory(&result, json) -} - -fn cpu_collector_request(spec: &CpuSamplingSpec) -> Result { - if spec.frequency_hz == 0 - || spec.duration_ms == 0 - || spec.max_samples == 0 - || spec.max_groups == 0 - || spec.stack_depth == 0 - || spec.max_threads == 0 - { - return Err(CommandFailure::new( - ErrorCode::SessionLimitExceeded, - "CPU sampling frequency, duration, samples, groups, stack depth, and threads must be positive", - true, - )); - } - let max_samples = usize::try_from(spec.max_samples).map_err(|error| { - CommandFailure::new( - ErrorCode::SessionLimitExceeded, - format!("CPU max_samples exceed this platform: {error}"), - true, - ) - })?; - let max_threads = usize::try_from(spec.max_threads).map_err(|error| { - CommandFailure::new( - ErrorCode::SessionLimitExceeded, - format!("CPU max_threads exceed this platform: {error}"), - true, - ) - })?; - let stack_depth = u16::try_from(spec.stack_depth).map_err(|error| { - CommandFailure::new( - ErrorCode::SessionLimitExceeded, - format!("CPU stack_depth exceeds the perf ABI: {error}"), - true, - ) - })?; - let thread_ids = cpu_sampling::target_thread_ids(&spec.target).map_err(|error| { - CommandFailure::new(error.code(), error.to_string(), error.recoverable()) - })?; - Ok(CpuSamplingRequest { - thread_ids, - frequency_hz: spec.frequency_hz, - duration: Duration::from_millis(spec.duration_ms), - timeout: Duration::from_millis(spec.timeout_ms), - max_samples, - max_threads, - stack_depth, - }) -} - -fn cpu_collector_failure(error: &cpu_collector::CpuSamplingError) -> CommandFailure { - let code = match error { - cpu_collector::CpuSamplingError::InvalidRequest(_) - | cpu_collector::CpuSamplingError::Timeout => ErrorCode::SessionLimitExceeded, - cpu_collector::CpuSamplingError::Open { source, .. } - | cpu_collector::CpuSamplingError::Arm { source, .. } - if source.kind() == std::io::ErrorKind::PermissionDenied => - { - ErrorCode::PermissionDenied - } - cpu_collector::CpuSamplingError::Cleanup { .. } => ErrorCode::CleanupFailed, - cpu_collector::CpuSamplingError::NoThreadsAttached { .. } => ErrorCode::TargetExited, - cpu_collector::CpuSamplingError::Open { .. } - | cpu_collector::CpuSamplingError::Ring { .. } - | cpu_collector::CpuSamplingError::Arm { .. } => ErrorCode::Internal, - }; - let recoverable = !matches!(code, ErrorCode::Internal | ErrorCode::CleanupFailed); - CommandFailure::new(code, error.to_string(), recoverable) -} - -fn emit_cpu_inventory(result: &CpuSampleInventoryResult, json: bool) -> ExitCode { - if json { - println!( - "{}", - serde_json::to_string_pretty(result).expect("CPU inventory must serialize") - ); - } else { - println!( - "CPU sample inventory: {} samples", - result.collection.observed_samples - ); - println!( - "Stacks: {} groups, {} lost, {} unresolved frames", - result.collection.groups, - result.collection.lost_samples, - result.symbolization.unresolved_frames - ); - for hotspot in result.inventory.hotspots.iter().take(10) { - let label = hotspot - .frame - .symbol - .as_deref() - .or(hotspot.frame.module_path.as_deref()) - .map_or_else(|| format!("{:#x}", hotspot.frame.address), str::to_owned); - println!( - " {:.1}% inclusive, {:.1}% self {label}", - hotspot.inclusive_proportion * 100.0, - hotspot.exclusive_proportion * 100.0 - ); - } - for warning in &result.warnings { - println!("Warning: {}: {}", warning.code, warning.message); - } - } - ExitCode::SUCCESS -} - -#[allow(clippy::too_many_arguments)] -fn run_direct_syscall_aggregate( - pid: u32, - name: Option, - duration_ms: u64, - timeout_ms: u64, - max_groups: u64, - max_inflight: u64, - json: bool, -) -> ExitCode { - let report = match inspect::run(pid) { - Ok(report) => report, - Err(error) => { - return emit_error(error.code(), error.to_string(), error.recoverable(), json); - } - }; - let spec = SyscallAggregateSpec { - schema_version: SchemaVersion::current(), - name, - target: report.target.clone(), - duration_ms, - timeout_ms, - max_groups, - max_inflight, - }; - run_syscall_aggregate_spec(&spec, Some(report), json) -} - -fn run_syscall_aggregate_spec( - spec: &SyscallAggregateSpec, - inspected: Option, - json: bool, -) -> ExitCode { - let report = match inspected.map_or_else(|| inspect::run(spec.target.pid), Ok) { - Ok(report) => report, - Err(error) => { - return emit_error(error.code(), error.to_string(), error.recoverable(), json); - } - }; - if let Err(error) = inspect::verify_target(&spec.target) { - return emit_error(error.code(), error.to_string(), error.recoverable(), json); - } - let validation = match syscall_aggregate::validate(&report) { - Ok(validation) => validation, - Err(error) => { - return emit_error(error.code(), error.to_string(), error.recoverable(), json); - } - }; - if !validation.valid { - let message = validation.issues.first().map_or_else( - || "syscall aggregate requirements are not satisfied".to_owned(), - |issue| issue.message.clone(), - ); - return emit_error(ErrorCode::PermissionDenied, message, true, json); - } - let request = match syscall_collector_request(spec) { - Ok(request) => request, - Err(error) => return emit_command_failure(error, json), - }; - let capture = match syscall_collector::collect(&request) { - Ok(capture) => capture, - Err(error) => { - return emit_error(error.code(), error.to_string(), error.recoverable(), json); - } - }; - let capture = syscall_aggregate::SyscallCaptureData { - groups: capture - .groups - .into_iter() - .map(|group| syscall_aggregate::RawSyscallGroup { - syscall_number: group.syscall_number, - count: group.count, - errors: group.errors, - total_duration_ns: group.total_duration_ns, - min_duration_ns: group.min_duration_ns, - max_duration_ns: group.max_duration_ns, - }) - .collect(), - observed_entries: capture.observed_entries, - matched_exits: capture.matched_exits, - unmatched_exits: capture.unmatched_exits, - dropped_aggregates: capture.dropped_aggregates, - inflight_at_end: capture.inflight_at_end, - }; - let result = match syscall_aggregate::analyze( - &report, - spec, - format!("xp_syscalls_{}", std::process::id()), - &capture, - ) { - Ok(result) => result, - Err(error) => { - return emit_error(error.code(), error.to_string(), error.recoverable(), json); - } - }; - emit_syscall_inventory(&result, json) -} - -fn syscall_collector_request( - spec: &SyscallAggregateSpec, -) -> Result { - if spec.duration_ms == 0 || spec.max_groups == 0 || spec.max_inflight == 0 { - return Err(CommandFailure::new( - ErrorCode::SessionLimitExceeded, - "syscall aggregate duration, groups, and inflight capacity must be positive", - true, - )); - } - let max_groups = u32::try_from(spec.max_groups).map_err(|error| { - CommandFailure::new( - ErrorCode::SessionLimitExceeded, - format!("syscall max_groups exceed the BPF map ABI: {error}"), - true, - ) - })?; - let max_inflight = u32::try_from(spec.max_inflight).map_err(|error| { - CommandFailure::new( - ErrorCode::SessionLimitExceeded, - format!("syscall max_inflight exceeds the BPF map ABI: {error}"), - true, - ) - })?; - Ok(SyscallAggregateRequest { - target: spec.target.clone(), - duration: Duration::from_millis(spec.duration_ms), - timeout: Duration::from_millis(spec.timeout_ms), - max_groups, - max_inflight, - }) -} - -fn emit_syscall_inventory(result: &SyscallAggregateResult, json: bool) -> ExitCode { - if json { - println!( - "{}", - serde_json::to_string_pretty(result).expect("syscall inventory must serialize") - ); - } else { - println!( - "Syscall inventory: {} entries, {} matched exits", - result.collection.observed_entries, result.collection.matched_exits - ); - for group in result.inventory.groups.iter().take(10) { - let label = group - .syscall_name - .clone() - .unwrap_or_else(|| format!("syscall#{}", group.syscall_number)); - println!( - " {label}: count={} total={}ns mean={:.0}ns errors={}", - group.count, group.duration_ns.total, group.duration_ns.mean, group.errors - ); - } - for warning in &result.warnings { - println!("Warning: {}: {}", warning.code, warning.message); - } - } - ExitCode::SUCCESS -} - -fn validate_measure_source(pid: Option, input: &[PathBuf]) -> Result<(), CommandFailure> { - if pid.is_some() && !input.is_empty() { - return Err(CommandFailure::new( - ErrorCode::InvalidEventSelector, - "--pid and --input select different collection modes", - true, - )); + )); } Ok(()) } @@ -1551,64 +481,6 @@ fn run_measure_spec( } } -fn run_cpu_sampling_file_spec( - value: serde_json::Value, - path: &Path, - has_device_options: bool, - has_artifact_options: bool, - json: bool, -) -> ExitCode { - if has_device_options || has_artifact_options { - return emit_error( - ErrorCode::InvalidEventSelector, - "CpuSamplingSpec does not accept CUPTI or event artifact options".to_owned(), - true, - json, - ); - } - let spec: CpuSamplingSpec = match serde_json::from_value(value) { - Ok(spec) => spec, - Err(error) => { - return emit_error( - ErrorCode::InvalidEventSelector, - format!("invalid CpuSamplingSpec {}: {error}", path.display()), - true, - json, - ); - } - }; - run_cpu_sampling_spec(&spec, None, json) -} - -fn run_syscall_aggregate_file_spec( - value: serde_json::Value, - path: &Path, - has_device_options: bool, - has_artifact_options: bool, - json: bool, -) -> ExitCode { - if has_device_options || has_artifact_options { - return emit_error( - ErrorCode::InvalidEventSelector, - "SyscallAggregateSpec does not accept CUPTI or event artifact options".to_owned(), - true, - json, - ); - } - let spec: SyscallAggregateSpec = match serde_json::from_value(value) { - Ok(spec) => spec, - Err(error) => { - return emit_error( - ErrorCode::InvalidEventSelector, - format!("invalid SyscallAggregateSpec {}: {error}", path.display()), - true, - json, - ); - } - }; - run_syscall_aggregate_spec(&spec, None, json) -} - fn measurement_spec_capacity(spec: &MeasurementSpec) -> Result { let capacity = match spec.measurement_mode { MeasurementMode::Exact if spec.max_groups.is_some() => { @@ -1719,1374 +591,42 @@ fn export_measurement_capture( const fn export_format_name(format: ExportFormatArg) -> &'static str { match format { ExportFormatArg::Jsonl => "jsonl", - ExportFormatArg::Chrome => "chrome", - } -} - -fn measure_completed_inputs( - input: &[PathBuf], - options: &MeasureOptions, -) -> Result { - let mut captures = Vec::with_capacity(input.len()); - for path in input { - let bytes = fs::read(path).map_err(|error| { - CommandFailure::new( - ErrorCode::TraceExportFailed, - format!("failed to read {}: {error}", path.display()), - false, - ) - })?; - let capture = completed::decode(&bytes, &options.session_id).map_err(|error| { - CommandFailure::new( - ErrorCode::TraceExportFailed, - format!("failed to decode {}: {error}", path.display()), - false, - ) - })?; - captures.push(capture); - } - let capture = completed::merge(captures, &options.session_id).map_err(|error| { - CommandFailure::new(ErrorCode::TraceExportFailed, error.to_string(), false) - })?; - reject_incomplete_capture(&capture, false, true) - .map_err(|error| error.with_artifact_events(capture.events.clone()))?; - if capture.unknown_records != 0 { - return Err(CommandFailure::new( - ErrorCode::TraceExportFailed, - format!( - "capture contains {} unknown CUPTI activity records", - capture.unknown_records - ), - false, - ) - .with_detail("unknown_records", capture.unknown_records) - .with_hint("rebuild xprobe and the CUPTI Agent from the same release") - .with_artifact_events(capture.events)); - } - let options = MeasureOptions { - dropped_events: capture.dropped_records, - ..options.clone() - }; - let mut result = measure(&capture.events, &options) - .map_err(|error| measurement_failure(&error, &options, &capture.events, true))?; - apply_collection_summary(&mut result, &capture); - Ok(MeasurementExecution { - result, - events: capture.events, - }) -} - -#[derive(Debug)] -struct MeasurementExecution { - result: MeasurementResult, - events: Vec, -} - -#[derive(Debug)] -struct CommandFailure { - code: ErrorCode, - message: String, - recoverable: bool, - details: BTreeMap, - hints: Vec, - artifact_events: Option>, -} - -impl CommandFailure { - fn new(code: ErrorCode, message: impl Into, recoverable: bool) -> Self { - Self { - code, - message: message.into(), - recoverable, - details: BTreeMap::new(), - hints: Vec::new(), - artifact_events: None, - } - } - - fn with_detail(mut self, key: impl Into, value: impl Into) -> Self { - self.details.insert(key.into(), value.into()); - self - } - - fn with_hint(mut self, hint: impl Into) -> Self { - self.hints.push(hint.into()); - self - } - - fn with_artifact_events(mut self, events: Vec) -> Self { - self.artifact_events = Some(events); - self - } -} - -fn measurement_failure( - error: &MeasureError, - options: &MeasureOptions, - events: &[Event], - preserve_artifact: bool, -) -> CommandFailure { - let mut failure = CommandFailure::new(error.code(), error.to_string(), error.recoverable()) - .with_detail("start_selector", options.start_selector.clone()) - .with_detail("end_selector", options.end_selector.clone()) - .with_detail("match_policy", match_policy_name(options.match_policy)); - if preserve_artifact { - failure = failure.with_artifact_events(events.to_vec()); - } - match error { - MeasureError::InvalidSelector(_) => { - failure = failure.with_hint("run xprobe validate with the exact selector text"); - } - MeasureError::InvalidPolicy(_) => { - failure = failure - .with_hint("use policy_recommendation.compatible_policies from xprobe validate"); - } - MeasureError::InvalidLimit(_) => { - failure = failure - .with_hint("set positive samples or duration, timeout, and max-events bounds"); - } - MeasureError::EventLimitExceeded { actual, maximum } => { - failure = failure - .with_detail("observed_events", *actual as u64) - .with_detail("max_events", *maximum as u64) - .with_hint("narrow the selectors or explicitly increase max-events"); - } - MeasureError::EventsDropped { count } => { - failure = failure - .with_detail("dropped_events", *count) - .with_hint("reduce the event rate or explicitly increase the capture capacity"); - } - MeasureError::ClockDomainsDiffer { start, end } => { - failure = failure - .with_detail("start_clock", clock_domain_name(start)) - .with_detail("end_clock", clock_domain_name(end)) - .with_hint( - "use endpoints in one clock domain or a CUPTI capture normalized to host monotonic time", - ); - } - MeasureError::NoMatchedSamples => { - let host_events = events - .iter() - .filter(|event| event.source == EventSource::Ebpf) - .count() as u64; - failure = failure - .with_detail("captured_events", events.len() as u64) - .with_detail("host_events", host_events) - .with_detail("cuda_events", events.len() as u64 - host_events) - .with_hint( - "inspect the events-out artifact, then adjust selectors or choose an explicitly compatible policy", - ); - } - } - failure -} - -const fn clock_domain_name(domain: &ClockDomain) -> &'static str { - match domain { - ClockDomain::HostMonotonic => "host_monotonic", - ClockDomain::Cupti => "cupti", - ClockDomain::CuptiNormalizedToHostMonotonic => "cupti_normalized_to_host_monotonic", - } -} - -struct LiveMeasureRequest { - pid: u32, - expected_target: Option, - cupti_socket: Option, - agent_path: Option, - timeout: Duration, - match_policy_text: &'static str, - mode: MeasurementMode, - options: MeasureOptions, -} - -type HostCaptureHandle = JoinHandle>; -const SNAPSHOT_QUIET_PERIOD: Duration = Duration::from_millis(100); - -struct HostCollectors { - cancelled: Arc, - handles: Vec, -} - -impl HostCollectors { - fn finished(&self) -> bool { - self.handles.iter().all(JoinHandle::is_finished) - } - - fn cancel(&self) { - self.cancelled.store(true, Ordering::Release); - } - - fn join(&mut self) -> Result, CommandFailure> { - let mut captures = Vec::with_capacity(self.handles.len()); - for handle in self.handles.drain(..) { - let result = handle.join().map_err(|_| { - CommandFailure::new(ErrorCode::Internal, "host collector thread panicked", false) - })?; - captures.push(result?); - } - Ok(captures) - } -} - -impl Drop for HostCollectors { - fn drop(&mut self) { - self.cancel(); - for handle in self.handles.drain(..) { - if handle.join().is_err() { - eprintln!("xprobe: host collector thread panicked during cleanup"); - } - } - } -} - -struct LiveCollection { - report: ProcessReport, - socket: Option, - deadline: Instant, - collection_end: Option, - baseline_timestamp: Option, - baseline_dropped: u64, - baseline_unknown: u64, - baseline_observed: u64, - collectors: HostCollectors, - host_captures: Option>, - latest_cuda: Option, - cuda_record_offset: u64, - managed_agent: bool, - cupti_armed: bool, - agent_injected: bool, - agent_major: Option, - agent_path: Option, - cupti_path: Option, -} - -fn collect_live_measurement( - request: &LiveMeasureRequest, -) -> Result { - assert_eq!(request.mode, MeasurementMode::Exact); - validate_live_limits(request)?; - let mut collection = prepare_live_collection(request)?; - - let outcome = (|| { - loop { - refresh_live_sources(&mut collection, request)?; - let now = Instant::now(); - if let Some(result) = evaluate_live_result(&collection, request, now)? { - break Ok(result); - } - if now >= collection.deadline { - break finish_live_timeout(&mut collection, request); - } - if collection.socket.is_none() { - thread::sleep(Duration::from_millis(10)); - } - } - })(); - let cuda_capture = collection.cupti_armed; - let stop_result = stop_cupti_agent(&mut collection, request); - let mut execution = match (outcome, stop_result) { - (Ok(mut execution), Ok(())) => { - if cuda_capture && collection.host_captures.is_some() { - let status = execution.result.status; - execution = measure_live_capture( - collection.host_captures.as_deref().unwrap_or(&[]), - &collection, - &request.options, - true, - )?; - execution.result.status = status; - } - execution - } - (Err(error), Ok(())) => { - if cuda_capture && collection.host_captures.is_some() { - return match measure_live_capture( - collection.host_captures.as_deref().unwrap_or(&[]), - &collection, - &request.options, - true, - ) { - Ok(_) => Err(error), - Err(final_error) => Err(final_error), - }; - } - return Err(error); - } - (Ok(execution), Err(cleanup)) => { - return Err(cleanup.with_artifact_events(execution.events)); - } - (Err(mut error), Err(cleanup)) => { - let artifact_events = error.artifact_events.take(); - let mut failure = CommandFailure::new( - ErrorCode::CleanupFailed, - format!( - "{}; original measurement error: {}", - cleanup.message, error.message - ), - cleanup.recoverable, - ) - .with_detail("cleanup_phase", "stop_cupti") - .with_detail("original_error_code", error.code.to_string()) - .with_detail("cleanup_error_code", cleanup.code.to_string()) - .with_hint("verify the target state before starting another measurement"); - if let Some(events) = artifact_events { - failure = failure.with_artifact_events(events); - } - return Err(failure); - } - }; - if collection.agent_injected { - let mut details = BTreeMap::new(); - if let Some(major) = collection.agent_major { - details.insert("cuda_major".to_owned(), serde_json::Value::from(major)); - } - if let Some(path) = collection.agent_path.as_deref() { - details.insert( - "agent_path".to_owned(), - serde_json::Value::from(path.display().to_string()), - ); - } - if let Some(path) = collection.cupti_path.as_deref() { - details.insert( - "cupti_path".to_owned(), - serde_json::Value::from(path.display().to_string()), - ); - } - execution.result.warnings.push(Warning { - code: "CUPTI_AGENT_INJECTED".to_owned(), - message: "xprobe injected the CUPTI agent and left the shared object mapped".to_owned(), - details, - }); - } - Ok(execution) -} - -fn collect_live_aggregate( - request: &LiveMeasureRequest, -) -> Result { - assert_eq!(request.mode, MeasurementMode::Aggregate); - validate_live_limits(request)?; - let mut collection = prepare_live_collection(request)?; - let collection_end = collection - .collection_end - .expect("aggregate validation requires duration"); - if collection_end > collection.deadline { - return Err(CommandFailure::new( - ErrorCode::SessionLimitExceeded, - "aggregate duration must not exceed timeout", - true, - )); - } - - while Instant::now() < collection_end { - let remaining = collection_end.saturating_duration_since(Instant::now()); - thread::sleep(remaining.min(Duration::from_millis(10))); - } - let target_result = inspect::verify_target(&collection.report.target) - .map_err(|error| CommandFailure::new(error.code(), error.to_string(), error.recoverable())); - let stop_result = stop_cupti_agent(&mut collection, request); - match (target_result, stop_result) { - (Err(error), Ok(())) | (Ok(()), Err(error)) => return Err(error), - (Err(target), Err(cleanup)) => { - return Err(CommandFailure::new( - ErrorCode::CleanupFailed, - format!( - "{}; target verification also failed: {}", - cleanup.message, target.message - ), - cleanup.recoverable, - ) - .with_detail("cleanup_phase", "stop_cupti") - .with_detail("target_error_code", target.code.to_string()) - .with_hint("verify the target state before starting another measurement")); - } - (Ok(()), Ok(())) => {} - } - - let capture = collection.latest_cuda.as_ref().ok_or_else(|| { - CommandFailure::new( - ErrorCode::CuptiNotAvailable, - "CUPTI Agent returned no aggregate capture", - false, - ) - })?; - let mut result = aggregate_inventory_result(capture, request)?; - if collection.agent_injected { - result.warnings.push(injection_warning(&collection)); - } - Ok(result) -} - -fn validate_live_limits(request: &LiveMeasureRequest) -> Result<(), CommandFailure> { - if request.timeout.is_zero() { - return Err(CommandFailure::new( - ErrorCode::SessionLimitExceeded, - "timeout must be greater than zero", - true, - )); - } - if request.options.samples == Some(0) - || request - .options - .duration - .is_some_and(|duration| duration.is_zero()) - || request.options.max_events == 0 - || (request.options.samples.is_none() && request.options.duration.is_none()) - { - return Err(CommandFailure::new( - ErrorCode::SessionLimitExceeded, - "live measurement requires positive samples or duration and max-events limits", - true, - )); - } - if request.mode == MeasurementMode::Aggregate { - if request.options.samples.is_some() - || request.options.duration.is_none() - || request.options.match_policy != MatchPolicy::Exact - { - return Err(CommandFailure::new( - ErrorCode::InvalidCorrelationPolicy, - "aggregate measurement requires duration, exact activity endpoints, and no samples limit", - true, - )); - } - if request.options.duration > Some(request.timeout) { - return Err(CommandFailure::new( - ErrorCode::SessionLimitExceeded, - "aggregate duration must not exceed timeout", - true, - )); - } - } - Ok(()) -} - -fn prepare_live_collection(request: &LiveMeasureRequest) -> Result { - let report = inspect::run(request.pid).map_err(|error| { - CommandFailure::new(error.code(), error.to_string(), error.recoverable()) - })?; - if request - .expected_target - .as_ref() - .is_some_and(|expected| expected != &report.target) - { - return Err(CommandFailure::new( - ErrorCode::TargetReused, - "trace specification target identity no longer matches the process", - true, - )); - } - let validation = validate::run( - &report, - &request.options.start_selector, - &request.options.end_selector, - request.match_policy_text, - ) - .map_err(|error| CommandFailure::new(error.code(), error.to_string(), error.recoverable()))?; - if let Some(issue) = validation.issues.first() { - return Err(CommandFailure::new(issue.code, issue.message.clone(), true) - .with_detail( - "recommended_policy", - match_policy_name(validation.policy_recommendation.policy), - ) - .with_detail( - "compatible_policies", - serde_json::to_value(&validation.policy_recommendation.compatible_policies) - .expect("match policies must serialize"), - ) - .with_hint("rerun validate with an explicitly compatible policy")); - } - if request.mode == MeasurementMode::Aggregate { - validate_aggregate_selection(&validation)?; - } - - let activation = prepare_cupti_activation(&report, &validation, request)?; - let arm_config = cupti_arm_config(&validation, request.options.max_events, request.mode)?; - let baseline = arm_cupti_capture( - &activation, - arm_config.as_ref(), - validation.requirements.needs_nvtx, - request, - )?; - let socket = activation.socket; - let baseline_timestamp = baseline - .as_ref() - .and_then(|capture| capture.events.last()) - .map(|event| event.timestamp_ns); - let baseline_dropped = baseline - .as_ref() - .map_or(0, |capture| capture.dropped_records); - let baseline_unknown = baseline - .as_ref() - .map_or(0, |capture| capture.unknown_records); - let baseline_observed = baseline - .as_ref() - .map_or(0, |capture| capture.observed_records); - let started = Instant::now(); - let deadline = started.checked_add(request.timeout).ok_or_else(|| { - CommandFailure::new( - ErrorCode::SessionLimitExceeded, - "timeout exceeds Instant range", - true, - ) - })?; - let collection_end = request - .options - .duration - .and_then(|duration| started.checked_add(duration)); - let collectors = start_host_collectors(&report, &validation, request)?; - let host_captures = collectors.handles.is_empty().then(Vec::new); - - Ok(LiveCollection { - report, - socket, - deadline, - collection_end, - baseline_timestamp, - baseline_dropped, - baseline_unknown, - baseline_observed, - collectors, - host_captures, - latest_cuda: None, - cuda_record_offset: 0, - managed_agent: activation.managed, - cupti_armed: baseline.is_some(), - agent_injected: activation.injected, - agent_major: activation.major, - agent_path: activation.agent_path, - cupti_path: activation.cupti_path, - }) -} - -fn validate_aggregate_selection(validation: &ValidationResult) -> Result<(), CommandFailure> { - let (Some(start), Some(end)) = (validation.start.cuda.as_ref(), validation.end.cuda.as_ref()) - else { - return Err(CommandFailure::new( - ErrorCode::InvalidEventSelector, - "aggregate measurement supports CUPTI activity endpoints only", - true, - )); - }; - let paired = matches!( - (&start.event_type, &end.event_type), - (&EventType::GpuKernelStart, &EventType::GpuKernelEnd) - | (&EventType::GpuMemcpyStart, &EventType::GpuMemcpyEnd) - | (&EventType::GpuMemsetStart, &EventType::GpuMemsetEnd) - ); - if !paired - || start.api_domain != end.api_domain - || start.api_name != end.api_name - || start.kernel_name_regex != end.kernel_name_regex - || start.memcpy_kind != end.memcpy_kind - { - return Err(CommandFailure::new( - ErrorCode::InvalidEventSelector, - "aggregate measurement requires matching kernel, memcpy, or memset start/end selectors", - true, - )); - } - if let Some(pattern) = start.kernel_name_regex.as_deref() - && pattern != ".*" - && matches!( - cupti::bounded_name_filter(pattern), - cupti::CuptiNameFilter::Any - ) - { - return Err(CommandFailure::new( - ErrorCode::InvalidEventSelector, - "aggregate kernel regex must be an exact, prefix, suffix, or contains pattern that the CUPTI Agent can apply", - true, - ) - .with_hint("use ^literal$, ^literal.*, .*literal$, or .*literal.*")); - } - Ok(()) -} - -fn arm_cupti_capture( - activation: &CuptiActivation, - config: Option<&cupti::CuptiArmConfig>, - needs_nvtx: bool, - request: &LiveMeasureRequest, -) -> Result, CommandFailure> { - let (Some(path), Some(config)) = (activation.socket.as_deref(), config) else { - assert!(activation.socket.is_none() && config.is_none()); - return Ok(None); - }; - let capture = match cupti::arm(path, request.timeout, &request.options.session_id, config) { - Ok(capture) => capture, - Err(error) => { - let failure = - CommandFailure::new(ErrorCode::CuptiNotAvailable, error.to_string(), true); - return cleanup_failed_arm(activation, request, failure); - } - }; - if needs_nvtx && !capture.nvtx_callbacks_active { - let failure = CommandFailure::new( - ErrorCode::CuptiNotAvailable, - "the mapped CUPTI agent was not initialized through NVTX; restart the target with NVTX_INJECTION64_PATH set to the same agent before the first NVTX call", - true, - ) - .with_hint("online CUDA injection cannot retrofit an initialized NVTX dispatch table"); - return cleanup_failed_arm(activation, request, failure); - } - if capture.state != cupti::CuptiCaptureState::Active { - let failure = CommandFailure::new( - ErrorCode::CuptiNotAvailable, - format!( - "CUPTI Agent did not enter active state after ARM: {:?}", - capture.state - ), - true, - ); - return cleanup_failed_arm(activation, request, failure); - } - Ok(Some(capture)) -} - -fn cleanup_failed_arm( - activation: &CuptiActivation, - request: &LiveMeasureRequest, - failure: CommandFailure, -) -> Result, CommandFailure> { - if !activation.managed { - return Err(failure); - } - let path = activation - .socket - .as_deref() - .expect("managed activation has a socket"); - match cupti::close(path, request.timeout, &request.options.session_id, 0) { - Ok(_) => Err(failure), - Err(cleanup) => Err(CommandFailure::new( - ErrorCode::CleanupFailed, - format!( - "failed to close CUPTI agent after ARM error: {cleanup}; original error: {}", - failure.message - ), - true, - ) - .with_detail("cleanup_phase", "close_after_arm_failure") - .with_detail("original_error_code", failure.code.to_string()) - .with_hint("verify the target state before starting another measurement")), - } -} - -fn cupti_arm_config( - validation: &ValidationResult, - max_events: usize, - mode: MeasurementMode, -) -> Result, CommandFailure> { - if !validation.requirements.needs_cupti { - return Ok(None); - } - let filters = [&validation.start, &validation.end] - .into_iter() - .filter_map(|endpoint| endpoint.cuda.as_ref()) - .map(|selector| { - cupti::event_filter(selector).map_err(|error| { - CommandFailure::new(ErrorCode::InvalidEventSelector, error.to_string(), false) - }) - }) - .collect::, _>>()?; - let record_capacity = u64::try_from(max_events).map_err(|error| { - CommandFailure::new( - ErrorCode::SessionLimitExceeded, - format!("max-events exceeds the CUPTI Agent range: {error}"), - true, - ) - })?; - Ok(Some(cupti::CuptiArmConfig { - record_capacity, - capture_mode: match mode { - MeasurementMode::Exact => cupti::CuptiCaptureMode::Exact, - MeasurementMode::Aggregate => cupti::CuptiCaptureMode::Aggregate, - }, - filters, - })) -} - -struct CuptiActivation { - socket: Option, - managed: bool, - injected: bool, - major: Option, - agent_path: Option, - cupti_path: Option, -} - -struct AgentSelection { - path: PathBuf, - major: Option, - cupti_path: Option, -} - -fn prepare_cupti_activation( - report: &ProcessReport, - validation: &ValidationResult, - request: &LiveMeasureRequest, -) -> Result { - if !validation.requirements.needs_cupti { - if request.cupti_socket.is_some() { - return Err(CommandFailure::new( - ErrorCode::InvalidEventSelector, - "--cupti-socket is only valid when a CUDA endpoint is selected", - true, - )); - } - return Ok(CuptiActivation { - socket: None, - managed: false, - injected: false, - major: None, - agent_path: None, - cupti_path: None, - }); - } - if let Some(socket) = request.cupti_socket.clone() { - return Ok(CuptiActivation { - socket: Some(socket), - managed: false, - injected: false, - major: None, - agent_path: None, - cupti_path: None, - }); - } - - let socket = std::env::temp_dir().join(format!( - "xprobe-{}-{}.sock", - request.pid, - std::process::id() - )); - let selection = if report.cuda.xprobe_cupti_loaded { - AgentSelection { - path: PathBuf::new(), - major: cupti_compat::target_major(report) - .map_err(|error| CommandFailure::new(error.code(), error.to_string(), true))?, - cupti_path: None, - } - } else { - resolve_agent_path(report, request.agent_path.as_deref())? - }; - eprintln!( - "xprobe: warning: activating the CUPTI agent modifies target PID {}{}{}", - request.pid, - selection - .major - .map_or_else(String::new, |major| format!(" with CUDA {major} Agent")), - selection - .cupti_path - .as_deref() - .map_or_else(String::new, |path| format!(" using {}", path.display())) - ); - let record_capacity = u64::try_from(request.options.max_events).map_err(|error| { - CommandFailure::new( - ErrorCode::SessionLimitExceeded, - format!("max-events exceeds the CUPTI Agent range: {error}"), - true, - ) - })?; - let activation = inject::activate( - report, - &selection.path, - &socket, - record_capacity, - request.timeout, - ) - .map_err(|error| CommandFailure::new(error.code(), error.to_string(), error.recoverable()))?; - Ok(CuptiActivation { - socket: Some(activation.socket_path), - managed: true, - injected: activation.injected, - major: selection.major, - agent_path: activation.injected.then_some(selection.path), - cupti_path: selection.cupti_path, - }) -} - -fn resolve_agent_path( - report: &ProcessReport, - configured: Option<&Path>, -) -> Result { - if let Some(path) = configured { - return Ok(AgentSelection { - path: path.to_owned(), - major: cupti_compat::target_major(report) - .map_err(|error| CommandFailure::new(error.code(), error.to_string(), true))?, - cupti_path: None, - }); - } - if let Some(path) = std::env::var_os("XPROBE_CUPTI_AGENT_PATH") { - return Ok(AgentSelection { - path: PathBuf::from(path), - major: cupti_compat::target_major(report) - .map_err(|error| CommandFailure::new(error.code(), error.to_string(), true))?, - cupti_path: None, - }); - } - let cupti = cupti_compat::resolve_library(report) - .map_err(|error| CommandFailure::new(error.code(), error.to_string(), true))?; - let executable = std::env::current_exe().map_err(|error| { - CommandFailure::new( - ErrorCode::CuptiNotAvailable, - format!("failed to locate xprobe executable: {error}"), - true, - ) - })?; - let parent = executable - .parent() - .expect("executable has a parent directory"); - let candidates = [ - parent.join(format!( - "../lib/xprobe/cuda{}/libxprobe-cupti.so", - cupti.major - )), - parent.join(format!("cuda{}/libxprobe-cupti.so", cupti.major)), - PathBuf::from(format!( - "build/cupti/cuda{}/libxprobe-cupti.so", - cupti.major - )), - ]; - let path = candidates - .into_iter() - .find(|path| path.is_file()) - .ok_or_else(|| { - CommandFailure::new( - ErrorCode::CuptiNotAvailable, - format!( - "CUDA {} CUPTI agent was not found; pass --agent or set XPROBE_CUPTI_AGENT_PATH", - cupti.major - ), - true, - ) - })?; - Ok(AgentSelection { - path, - major: Some(cupti.major), - cupti_path: Some(cupti.path), - }) -} - -fn stop_cupti_agent( - collection: &mut LiveCollection, - request: &LiveMeasureRequest, -) -> Result<(), CommandFailure> { - if !collection.cupti_armed { - return Ok(()); - } - let socket = collection - .socket - .as_deref() - .expect("armed CUPTI agent has a socket"); - collection.cupti_armed = false; - let capture = if collection.managed_agent { - cupti::close( - socket, - request.timeout, - &request.options.session_id, - collection.cuda_record_offset, - ) - } else { - cupti::stop( - socket, - request.timeout, - &request.options.session_id, - collection.cuda_record_offset, - ) - } - .map_err(|error| { - CommandFailure::new(ErrorCode::CleanupFailed, error.to_string(), true) - .with_detail("cleanup_phase", "stop_cupti") - .with_hint("verify the target state before starting another measurement") - })?; - append_cuda_capture(collection, capture)?; - collection.managed_agent = false; - Ok(()) -} - -impl Drop for LiveCollection { - fn drop(&mut self) { - if self.cupti_armed { - if let Some(socket) = self.socket.as_deref() { - let result = if self.managed_agent { - cupti::close( - socket, - Duration::from_secs(2), - "xp_cleanup", - self.cuda_record_offset, - ) - } else { - cupti::stop( - socket, - Duration::from_secs(2), - "xp_cleanup", - self.cuda_record_offset, - ) - }; - if let Err(error) = result { - eprintln!("xprobe: failed to stop CUPTI agent during cleanup: {error}"); - } - } - } - } -} - -fn start_host_collectors( - report: &ProcessReport, - validation: &ValidationResult, - request: &LiveMeasureRequest, -) -> Result { - inspect::verify_target(&report.target).map_err(|error| { - CommandFailure::new(error.code(), error.to_string(), error.recoverable()) - })?; - let cancelled = Arc::new(AtomicBool::new(false)); - let mut probes = Vec::new(); - let mut linux_probes = Vec::new(); - for endpoint in [&validation.start, &validation.end] { - if let Some(probe) = endpoint.host.as_ref() - && !probes - .iter() - .any(|existing: &ResolvedProbe| existing.selector == probe.selector) - { - probes.push(probe.clone()); - } - if let Some(probe) = endpoint.linux.as_ref() - && !linux_probes - .iter() - .any(|existing: &ResolvedLinuxSelector| existing == probe) - { - linux_probes.push(probe.clone()); - } - } - let host_timeout = request - .options - .duration - .map_or(request.timeout, |duration| duration.min(request.timeout)); - let host_samples = request - .options - .samples - .unwrap_or(request.options.max_events); - let mut receivers = Vec::with_capacity(probes.len() + usize::from(!linux_probes.is_empty())); - let mut handles = probes - .into_iter() - .enumerate() - .map(|(index, probe)| { - let (ready, receiver) = mpsc::sync_channel(1); - receivers.push(receiver); - let offset = if probe.symbol.is_some() { - 0 - } else { - probe.file_offset - }; - let capture_request = UprobeRequest { - target: report.target.clone(), - binary: PathBuf::from(&probe.binary_path), - symbol: probe.symbol, - symbol_demangled: probe.symbol_demangled, - offset, - probe_kind: probe.probe_kind, - probe_id: u32::try_from(index + 1).expect("two endpoints fit u32"), - samples: host_samples, - timeout: host_timeout, - cancelled: Arc::clone(&cancelled), - ready: Some(ready), - }; - thread::spawn(move || { - uprobe::capture(&capture_request).map_err(|error| { - CommandFailure::new(error.code(), error.to_string(), error.recoverable()) - }) - }) - }) - .collect::>(); - if !linux_probes.is_empty() { - let (ready, receiver) = mpsc::sync_channel(1); - receivers.push(receiver); - let endpoint_count = linux_probes.len(); - let event_limit = linux_capture_event_limit(&request.options, endpoint_count); - let capture_request = LinuxCaptureRequest { - target: report.target.clone(), - probes: linux_probes, - event_limit, - capacity_limit: request.options.samples.is_none(), - timeout: host_timeout, - cancelled: Arc::clone(&cancelled), - ready: Some(ready), - }; - handles.push(thread::spawn(move || { - linux::capture(&capture_request).map_err(|error| { - CommandFailure::new(error.code(), error.to_string(), error.recoverable()) - }) - })); - } - let mut collectors = HostCollectors { cancelled, handles }; - wait_for_host_collectors(&mut collectors, receivers, request.timeout)?; - if let Err(error) = inspect::verify_target(&report.target) { - collectors.cancel(); - collectors.join()?; - return Err(CommandFailure::new( - error.code(), - error.to_string(), - error.recoverable(), - )); - } - Ok(collectors) -} - -fn linux_capture_event_limit(options: &MeasureOptions, endpoint_count: usize) -> usize { - options - .samples - .and_then(|samples| { - samples - .checked_mul(endpoint_count) - .and_then(|events| events.checked_add(endpoint_count)) - }) - .unwrap_or(options.max_events) - .min(options.max_events) -} - -fn wait_for_host_collectors( - collectors: &mut HostCollectors, - receivers: Vec>, - timeout: Duration, -) -> Result<(), CommandFailure> { - for receiver in receivers { - match receiver.recv_timeout(timeout) { - Ok(()) => {} - Err(mpsc::RecvTimeoutError::Disconnected) => { - collectors.cancel(); - return match collectors.join() { - Err(error) => Err(error), - Ok(_) => Err(CommandFailure::new( - ErrorCode::Internal, - "host collector exited before reporting readiness", - false, - )), - }; - } - Err(mpsc::RecvTimeoutError::Timeout) => { - collectors.cancel(); - collectors.join()?; - return Err(CommandFailure::new( - ErrorCode::SessionLimitExceeded, - "timed out waiting for host collector readiness", - true, - )); - } - } - } - Ok(()) -} - -fn refresh_live_sources( - collection: &mut LiveCollection, - request: &LiveMeasureRequest, -) -> Result<(), CommandFailure> { - if collection.host_captures.is_none() && collection.collectors.finished() { - collection.host_captures = Some(collection.collectors.join()?); - } - let Some(path) = collection.socket.as_deref() else { - return Ok(()); - }; - let remaining = collection - .deadline - .saturating_duration_since(Instant::now()); - if remaining <= SNAPSHOT_QUIET_PERIOD { - collection.collectors.cancel(); - return Ok(()); - } - let capture = cupti::snapshot( - path, - remaining, - &request.options.session_id, - collection.cuda_record_offset, - ) - .map_err(|error| CommandFailure::new(ErrorCode::CuptiNotAvailable, error.to_string(), true))?; - append_cuda_capture(collection, capture)?; - Ok(()) -} - -fn append_cuda_capture( - collection: &mut LiveCollection, - mut capture: cupti::CuptiCapture, -) -> Result<(), CommandFailure> { - if capture.record_offset != collection.cuda_record_offset { - return Err(CommandFailure::new( - ErrorCode::TraceExportFailed, - format!( - "CUPTI incremental capture started at record {}, expected {}", - capture.record_offset, collection.cuda_record_offset - ), - false, - )); - } - let returned = u64::try_from(capture.events.len()).map_err(|error| { - CommandFailure::new( - ErrorCode::TraceExportFailed, - format!("CUPTI incremental event count exceeds u64: {error}"), - false, - ) - })?; - let next_offset = capture.record_offset.checked_add(returned).ok_or_else(|| { - CommandFailure::new( - ErrorCode::TraceExportFailed, - "CUPTI incremental record offset overflowed", - false, - ) - })?; - if let Some(accumulated) = collection.latest_cuda.as_mut() { - if accumulated.record_capacity != capture.record_capacity { - return Err(CommandFailure::new( - ErrorCode::TraceExportFailed, - format!( - "CUPTI record capacity changed from {} to {} during capture", - accumulated.record_capacity, capture.record_capacity - ), - false, - )); - } - if capture.observed_records < accumulated.observed_records - || capture.agent_dropped_records < accumulated.agent_dropped_records - || capture.cupti_dropped_records < accumulated.cupti_dropped_records - || capture.unknown_records < accumulated.unknown_records - { - return Err(CommandFailure::new( - ErrorCode::TraceExportFailed, - "CUPTI incremental counters moved backwards", - false, - )); - } - accumulated.state = capture.state; - accumulated.stop_reason = capture.stop_reason; - accumulated.observed_records = capture.observed_records; - accumulated.agent_dropped_records = capture.agent_dropped_records; - accumulated.cupti_dropped_records = capture.cupti_dropped_records; - accumulated.dropped_records = capture.dropped_records; - accumulated.unknown_records = capture.unknown_records; - accumulated.events.append(&mut capture.events); - } else { - capture.record_offset = 0; - collection.latest_cuda = Some(capture); - } - collection.cuda_record_offset = next_offset; - Ok(()) -} - -fn aggregate_inventory_result( - capture: &cupti::CuptiCapture, - request: &LiveMeasureRequest, -) -> Result { - validate_aggregate_capture(capture)?; - let capture = gpu_aggregate::GpuAggregateCapture { - groups: capture - .aggregate_groups - .iter() - .map(|group| gpu_aggregate::RawGpuAggregateGroup { - activity: match group.activity { - cupti::CuptiAggregateActivity::Kernel => AggregateActivity::Kernel, - cupti::CuptiAggregateActivity::Memcpy => AggregateActivity::Memcpy, - cupti::CuptiAggregateActivity::Memset => AggregateActivity::Memset, - }, - name: group.name.clone(), - name_complete: group.name_complete, - device_id: group.device_id, - memcpy_kind: group.memcpy_kind.clone(), - count: group.count, - total_duration_ns: group.total_duration_ns, - min_duration_ns: group.min_duration_ns, - max_duration_ns: group.max_duration_ns, - total_bytes: group.total_bytes, - }) - .collect(), - observed_activities: capture.observed_records, - group_capacity: capture.record_capacity, - }; - let options = &request.options; - gpu_aggregate::analyze( - gpu_aggregate::GpuAggregateRequest { - session_id: options.session_id.clone(), - name: options.name.clone(), - start_selector: options.start_selector.clone(), - end_selector: options.end_selector.clone(), - }, - capture, - ) - .map_err(|error| CommandFailure::new(error.code(), error.to_string(), error.recoverable())) -} - -fn validate_aggregate_capture(capture: &cupti::CuptiCapture) -> Result<(), CommandFailure> { - if capture.state == cupti::CuptiCaptureState::LimitReached { - return Err(CommandFailure::new( - ErrorCode::EventRateTooHigh, - "aggregate group table reached its configured capacity", - true, - ) - .with_detail("group_capacity", capture.record_capacity) - .with_detail("observed_activities", capture.observed_records) - .with_hint("narrow the selector or explicitly increase max-groups")); - } - if capture.state != cupti::CuptiCaptureState::Stopped - || capture.stop_reason != cupti::CuptiStopReason::Requested - { - return Err(CommandFailure::new( - ErrorCode::CuptiNotAvailable, - format!( - "aggregate CUPTI capture ended in state {:?} with reason {:?}", - capture.state, capture.stop_reason - ), - true, - )); - } - if capture.dropped_records != 0 { - return Err(CommandFailure::new( - ErrorCode::EventsDropped, - "aggregate CUPTI capture dropped activities", - true, - ) - .with_detail("dropped_activities", capture.dropped_records) - .with_hint("reduce capture duration or narrow the activity selector")); - } - if capture.unknown_records != 0 { - return Err(CommandFailure::new( - ErrorCode::TraceExportFailed, - "aggregate CUPTI capture contains unknown activities", - false, - ) - .with_detail("unknown_records", capture.unknown_records) - .with_hint("rebuild xprobe and the CUPTI Agent from the same release")); - } - if !capture.events.is_empty() { - return Err(CommandFailure::new( - ErrorCode::TraceExportFailed, - "aggregate CUPTI capture unexpectedly contains exact events", - false, - )); - } - Ok(()) -} - -fn injection_warning(collection: &LiveCollection) -> Warning { - let mut details = BTreeMap::new(); - if let Some(major) = collection.agent_major { - details.insert("cuda_major".to_owned(), serde_json::Value::from(major)); - } - if let Some(path) = collection.agent_path.as_deref() { - details.insert( - "agent_path".to_owned(), - serde_json::Value::from(path.display().to_string()), - ); - } - if let Some(path) = collection.cupti_path.as_deref() { - details.insert( - "cupti_path".to_owned(), - serde_json::Value::from(path.display().to_string()), - ); - } - Warning { - code: "CUPTI_AGENT_INJECTED".to_owned(), - message: "xprobe injected the CUPTI agent and left the shared object mapped".to_owned(), - details, - } -} - -fn evaluate_live_result( - collection: &LiveCollection, - request: &LiveMeasureRequest, - now: Instant, -) -> Result, CommandFailure> { - let Some(host) = collection.host_captures.as_ref() else { - return Ok(None); - }; - if collection.socket.is_some() && collection.latest_cuda.is_none() { - return Ok(None); - } - let mut execution = match measure_live_capture(host, collection, &request.options, false) { - Ok(execution) => execution, - Err(error) if error.code == ErrorCode::NoMatchedSamples => return Ok(None), - Err(error) => return Err(error), - }; - let reached_samples = request.options.samples.is_some_and(|samples| { - execution.result.measurement.samples.matched - >= u64::try_from(samples).expect("sample limit fits u64") - }); - let reached_duration = collection.collection_end.is_some_and(|end| now >= end); - if reached_samples || reached_duration { - inspect::verify_target(&collection.report.target).map_err(|error| { - CommandFailure::new(error.code(), error.to_string(), error.recoverable()) - })?; - return Ok(Some(execution)); - } - if now >= collection.deadline { - execution.result.status = SessionStatus::TimedOut; - return Ok(Some(execution)); + ExportFormatArg::Chrome => "chrome", } - Ok(None) } -fn finish_live_timeout( - collection: &mut LiveCollection, - request: &LiveMeasureRequest, +fn measure_completed_inputs( + input: &[PathBuf], + options: &MeasureOptions, ) -> Result { - collection.collectors.cancel(); - if collection.host_captures.is_none() { - collection.host_captures = Some(collection.collectors.join()?); + let mut captures = Vec::with_capacity(input.len()); + for path in input { + let bytes = fs::read(path).map_err(|error| { + CommandFailure::new( + ErrorCode::TraceExportFailed, + format!("failed to read {}: {error}", path.display()), + false, + ) + })?; + let capture = completed::decode(&bytes, &options.session_id).map_err(|error| { + CommandFailure::new( + ErrorCode::TraceExportFailed, + format!("failed to decode {}: {error}", path.display()), + false, + ) + })?; + captures.push(capture); } - let mut execution = measure_live_capture( - collection - .host_captures - .as_ref() - .expect("host captures were assigned"), - collection, - &request.options, - true, - ) - .map_err(|mut error| { - if error.code == ErrorCode::NoMatchedSamples { - let host_events = collection.host_captures.as_ref().map_or(0, |captures| { - captures.iter().map(|capture| capture.events.len()).sum() - }); - let cuda_events = collection.latest_cuda.as_ref().map_or(0, |capture| { - capture - .events - .iter() - .filter(|event| { - collection - .baseline_timestamp - .is_none_or(|start| event.timestamp_ns > start) - }) - .count() - }); - error.message = format!( - "no event pairs matched before the live measurement timeout \ - (host events: {host_events}, CUDA events: {cuda_events})" - ); - error - .with_detail("timed_out", true) - .with_detail("host_events", host_events as u64) - .with_detail("cuda_events", cuda_events as u64) - } else { - error - } + let capture = completed::merge(captures, &options.session_id).map_err(|error| { + CommandFailure::new(ErrorCode::TraceExportFailed, error.to_string(), false) })?; - execution.result.status = SessionStatus::TimedOut; - Ok(execution) -} - -fn measure_live_capture( - host_captures: &[HostCaptureResult], - collection: &LiveCollection, - options: &MeasureOptions, - require_complete: bool, -) -> Result { - let capture = completed_live_capture(host_captures, collection, &options.session_id)?; - reject_incomplete_capture(&capture, true, require_complete) + reject_incomplete_capture(&capture, false, true) .map_err(|error| error.with_artifact_events(capture.events.clone()))?; if capture.unknown_records != 0 { return Err(CommandFailure::new( ErrorCode::TraceExportFailed, format!( - "live capture contains {} unknown CUPTI activity records", + "capture contains {} unknown CUPTI activity records", capture.unknown_records ), false, @@ -3099,9 +639,8 @@ fn measure_live_capture( dropped_events: capture.dropped_records, ..options.clone() }; - let mut result = measure(&capture.events, &options).map_err(|error| { - measurement_failure(&error, &options, &capture.events, require_complete) - })?; + let mut result = measure(&capture.events, &options) + .map_err(|error| measurement_failure(&error, &options, &capture.events, true))?; apply_collection_summary(&mut result, &capture); Ok(MeasurementExecution { result, @@ -3109,164 +648,17 @@ fn measure_live_capture( }) } -fn completed_live_capture( - host_captures: &[HostCaptureResult], - collection: &LiveCollection, - session_id: &str, -) -> Result { - let mut captures = host_captures - .iter() - .map(|capture| completed::CompletedCapture { - dropped_records: capture.dropped, - unknown_records: 0, - record_limit_reached: capture.record_limit_reached.then_some(capture.captured), - capture_failed: false, - cupti: None, - events: capture.events.clone(), - }) - .collect::>(); - if let Some(cuda) = collection.latest_cuda.as_ref() { - let dropped_records = cuda - .dropped_records - .checked_sub(collection.baseline_dropped) - .ok_or_else(|| { - CommandFailure::new( - ErrorCode::TraceExportFailed, - "CUPTI dropped-record counter moved backwards", - false, - ) - })?; - let unknown_records = cuda - .unknown_records - .checked_sub(collection.baseline_unknown) - .ok_or_else(|| { - CommandFailure::new( - ErrorCode::TraceExportFailed, - "CUPTI unknown-record counter moved backwards", - false, - ) - })?; - let observed_records = cuda - .observed_records - .checked_sub(collection.baseline_observed) - .ok_or_else(|| { - CommandFailure::new( - ErrorCode::TraceExportFailed, - "CUPTI observed-record counter moved backwards", - false, - ) - })?; - captures.push(completed::CompletedCapture { - dropped_records, - unknown_records, - record_limit_reached: (cuda.state == cupti::CuptiCaptureState::LimitReached) - .then_some(cuda.record_capacity), - capture_failed: matches!( - cuda.state, - cupti::CuptiCaptureState::Idle | cupti::CuptiCaptureState::Failed - ), - cupti: Some(completed::CompletedCuptiStatistics { - complete: cuda.state == cupti::CuptiCaptureState::Stopped, - record_capacity: cuda.record_capacity, - observed_records, - dropped_records, - }), - events: cuda - .events - .iter() - .filter(|event| { - collection - .baseline_timestamp - .is_none_or(|start| event.timestamp_ns > start) - }) - .cloned() - .collect(), - }); - } - completed::merge(captures, session_id).map_err(|error| { - CommandFailure::new(ErrorCode::TraceExportFailed, error.to_string(), false) - }) -} - -fn reject_incomplete_capture( - capture: &completed::CompletedCapture, - live: bool, - require_complete: bool, -) -> Result<(), CommandFailure> { - if capture.capture_failed { - return Err(CommandFailure::new( - ErrorCode::CuptiNotAvailable, - format!( - "{}CUPTI capture entered a failed state", - if live { "live " } else { "" } - ), - true, - ) - .with_detail("capture_state", "failed") - .with_hint("inspect target CUDA/CUPTI compatibility and retry a fresh bounded capture")); - } - if let Some(capacity) = capture.record_limit_reached { - let mut failure = CommandFailure::new( - ErrorCode::EventRateTooHigh, - format!("capture reached its configured limit of {capacity} records"), - true, - ) - .with_detail("record_capacity", capacity) - .with_hint("narrow the selectors or explicitly increase max-events"); - if let Some(cupti) = capture.cupti.as_ref() { - failure = failure - .with_detail("observed_records", cupti.observed_records) - .with_detail("retained_records", capture.events.len() as u64) - .with_detail("dropped_records", cupti.dropped_records); - } - return Err(failure); - } - if require_complete && capture.cupti.as_ref().is_some_and(|cupti| !cupti.complete) { - return Err(CommandFailure::new( - ErrorCode::CuptiNotAvailable, - format!( - "{}CUPTI capture was not stopped before measurement", - if live { "live " } else { "" } - ), - true, - ) - .with_detail("capture_state", "active") - .with_hint("stop the CUPTI capture before using it as completed input")); - } - Ok(()) -} - -fn apply_collection_summary(result: &mut MeasurementResult, capture: &completed::CompletedCapture) { - let Some(cupti) = capture.cupti.as_ref() else { - return; - }; - let retained_records = result.collection.cuda_events; - let buffer_utilization = utilization(retained_records, cupti.record_capacity); - result.collection.cupti = Some(CuptiCollectionSummary { - record_capacity: cupti.record_capacity, - observed_records: cupti.observed_records, - retained_records, - dropped_records: cupti.dropped_records, - buffer_utilization, - }); -} - -#[allow(clippy::cast_precision_loss)] -fn utilization(retained: u64, capacity: u64) -> f64 { - if capacity == 0 { - 0.0 - } else { - retained as f64 / capacity as f64 - } +#[derive(Debug)] +struct MeasurementExecution { + result: MeasurementResult, + events: Vec, } -const fn match_policy_name(policy: MatchPolicy) -> &'static str { - match policy { - MatchPolicy::Exact => "exact", - MatchPolicy::FirstAfter => "first-after", - MatchPolicy::Nearest => "nearest", - MatchPolicy::StackNested => "stack-nested", - MatchPolicy::StreamOrder => "stream-order", +const fn clock_domain_name(domain: &ClockDomain) -> &'static str { + match domain { + ClockDomain::HostMonotonic => "host_monotonic", + ClockDomain::Cupti => "cupti", + ClockDomain::CuptiNormalizedToHostMonotonic => "cupti_normalized_to_host_monotonic", } } @@ -3418,37 +810,6 @@ fn emit_syscall_aggregate_validation( ExitCode::SUCCESS } -fn run_resolve(args: ResolveArgs) -> ExitCode { - let ResolveArgs { - pid, - selector, - json, - no_color: _, - non_interactive: _, - } = args; - let report = match inspect::run(pid) { - Ok(report) => report, - Err(error) => { - return emit_error(error.code(), error.to_string(), error.recoverable(), json); - } - }; - - match resolve::run(&report, &selector) { - Ok(resolved) => { - if json { - println!( - "{}", - serde_json::to_string_pretty(&resolved).expect("resolved probe must serialize") - ); - } else { - print_resolved_probe(&resolved); - } - ExitCode::SUCCESS - } - Err(error) => emit_error(error.code(), error.to_string(), error.recoverable(), json), - } -} - fn run_doctor(args: DoctorArgs) -> ExitCode { let DoctorArgs { json, @@ -3502,461 +863,3 @@ fn run_discover(args: DiscoverArgs) -> ExitCode { Err(error) => emit_error(error.code(), error.to_string(), error.recoverable(), json), } } - -fn run_inspect(args: InspectArgs) -> ExitCode { - let InspectArgs { - pid, - json, - no_color: _, - non_interactive: _, - } = args; - - match inspect::run(pid) { - Ok(report) => { - if json { - println!( - "{}", - serde_json::to_string_pretty(&report).expect("process report must serialize") - ); - } else { - print_process_report(&report); - } - ExitCode::SUCCESS - } - Err(error) => emit_error(error.code(), error.to_string(), error.recoverable(), json), - } -} - -fn run_uprobe(args: UprobeArgs) -> ExitCode { - let UprobeArgs { - pid, - binary, - symbol, - return_probe, - probe_id, - samples, - timeout_ms, - output: EventOutputArgs { json, jsonl }, - no_color: _, - non_interactive: _, - } = args; - let machine_output = json || jsonl; - - let report = match inspect::run(pid) { - Ok(report) => report, - Err(error) => { - return emit_error( - error.code(), - error.to_string(), - error.recoverable(), - machine_output, - ); - } - }; - let binary = match mapped_binary(&report, &binary) { - Ok(binary) => binary, - Err(message) => { - return emit_error(ErrorCode::BinaryNotMapped, message, true, machine_output); - } - }; - let request = UprobeRequest { - target: report.target.clone(), - binary, - symbol: Some(symbol), - symbol_demangled: None, - offset: 0, - probe_kind: if return_probe { - xprobe_protocol::HostProbeKind::Uretprobe - } else { - xprobe_protocol::HostProbeKind::Uprobe - }, - probe_id, - samples, - timeout: Duration::from_millis(timeout_ms), - cancelled: Arc::new(AtomicBool::new(false)), - ready: None, - }; - let result = match uprobe::capture(&request) { - Ok(result) => result, - Err(error) => { - return emit_error( - error.code(), - error.to_string(), - error.recoverable(), - machine_output, - ); - } - }; - if let Err(error) = inspect::verify_target(&report.target) { - return emit_error( - error.code(), - error.to_string(), - error.recoverable(), - machine_output, - ); - } - - if jsonl { - match events_to_jsonl(&result.events) { - Ok(output) => print!("{output}"), - Err(error) => { - return emit_error( - ErrorCode::TraceExportFailed, - format!("failed to serialize host events: {error}"), - false, - true, - ); - } - } - } else if json { - println!( - "{}", - serde_json::to_string_pretty(&result).expect("host capture result must serialize") - ); - } else { - print_host_capture(&result); - } - ExitCode::SUCCESS -} - -fn run_cupti(args: CuptiArgs) -> ExitCode { - let CuptiArgs { - input, - socket, - timeout_ms, - session_id, - json, - no_color: _, - non_interactive: _, - } = args; - let capture = if let Some(input) = input { - let bytes = match fs::read(&input) { - Ok(bytes) => bytes, - Err(error) => { - return emit_error( - ErrorCode::TraceExportFailed, - format!("failed to read {}: {error}", input.display()), - false, - json, - ); - } - }; - match cupti::decode_capture(&bytes, &session_id) { - Ok(capture) => capture, - Err(error) => { - return emit_error(ErrorCode::TraceExportFailed, error.to_string(), false, json); - } - } - } else { - let socket = socket.expect("clap requires one CUPTI capture source"); - match cupti::snapshot(&socket, Duration::from_millis(timeout_ms), &session_id, 0) { - Ok(capture) => capture, - Err(error) => { - return emit_error(ErrorCode::TraceExportFailed, error.to_string(), true, json); - } - } - }; - let output = match events_to_jsonl(&capture.events) { - Ok(output) => output, - Err(error) => { - return emit_error( - ErrorCode::TraceExportFailed, - format!("failed to serialize CUPTI events: {error}"), - false, - json, - ); - } - }; - - if capture.dropped_records != 0 || capture.unknown_records != 0 { - eprintln!( - "CUPTI capture diagnostics: dropped={}, unknown={}", - capture.dropped_records, capture.unknown_records - ); - } - print!("{output}"); - ExitCode::SUCCESS -} - -fn mapped_binary(report: &ProcessReport, binary: &Path) -> Result { - let canonical = fs::canonicalize(binary) - .map_err(|error| format!("failed to resolve {}: {error}", binary.display()))?; - let is_executable = Path::new(&report.executable) == canonical; - let is_library = report - .loaded_libraries - .iter() - .any(|mapped| Path::new(mapped) == canonical); - if !is_executable && !is_library { - return Err(format!( - "{} is not mapped by target PID {}", - canonical.display(), - report.target.pid - )); - } - Ok(canonical) -} - -fn emit_error(code: ErrorCode, message: String, recoverable: bool, json: bool) -> ExitCode { - emit_xprobe_error( - XprobeError { - code, - message, - recoverable, - details: BTreeMap::new(), - hints: Vec::new(), - }, - json, - ) -} - -fn emit_command_failure(error: CommandFailure, json: bool) -> ExitCode { - emit_xprobe_error( - XprobeError { - code: error.code, - message: error.message, - recoverable: error.recoverable, - details: error.details, - hints: error.hints, - }, - json, - ) -} - -fn emit_xprobe_error(error: XprobeError, json: bool) -> ExitCode { - let code = error.code; - let response = ErrorResponse::new(error); - if json { - println!( - "{}", - serde_json::to_string_pretty(&response).expect("error response must serialize") - ); - } else { - eprintln!("{code}: {}", response.error.message); - } - - match code { - ErrorCode::TargetNotFound | ErrorCode::TargetExited | ErrorCode::TargetReused => { - ExitCode::from(3) - } - ErrorCode::PermissionDenied => ExitCode::from(4), - _ => ExitCode::from(1), - } -} - -fn print_human_report(report: &CapabilityReport) { - println!( - "xprobe doctor (schema {})", - schema_version(report.schema_version) - ); - println!(); - println!("Environment:"); - println!(" OS {}", report.environment.operating_system); - println!(" architecture {}", report.environment.architecture); - println!(" kernel {}", report.environment.kernel_release); - println!(" effective UID {}", report.environment.effective_uid); - println!( - " container {}", - report - .environment - .container - .as_deref() - .unwrap_or("none detected") - ); - println!(" PID namespace {}", report.environment.pid_namespace); - println!(); - println!("Capabilities:"); - println!(" uprobe {}", yes_no(report.capabilities.uprobe)); - println!( - " uretprobe {}", - yes_no(report.capabilities.uretprobe) - ); - println!( - " tracepoint {}", - yes_no(report.capabilities.tracepoint) - ); - println!( - " CUDA callback {}", - yes_no(report.capabilities.cuda_callback) - ); - println!( - " CUDA activity {}", - yes_no(report.capabilities.cuda_activity) - ); - println!( - " runtime inject {}", - yes_no(report.capabilities.runtime_injection) - ); - println!(); - println!("Checks:"); - print_check("BTF", &report.checks.btf); - print_check("eBPF permissions", &report.checks.ebpf_permissions); - print_check("kernel lockdown", &report.checks.kernel_lockdown); - print_check("perf paranoid", &report.checks.perf_event_paranoid); - print_check("ptrace scope", &report.checks.ptrace_scope); - print_check("NVIDIA driver", &report.checks.nvidia_driver); - print_check("CUDA driver", &report.checks.cuda_driver); - print_check("CUDA toolkit", &report.checks.cuda_toolkit); - print_check("CUPTI", &report.checks.cupti); - - if !report.warnings.is_empty() { - println!(); - println!("Warnings:"); - for warning in &report.warnings { - println!(" {}: {}", warning.code, warning.message); - } - } -} - -fn print_process_report(report: &ProcessReport) { - println!("Process: PID {}", report.target.pid); - println!("Executable: {}", report.executable); - println!("Start time: {} ticks", report.target.process_start_time); - println!("Namespace PIDs: {:?}", report.namespace_pids); - println!("Mount namespace: {}", report.mount_namespace); - println!( - "Credentials: uid={} gid={}", - report.credentials.effective_uid, report.credentials.effective_gid - ); - println!(); - println!("CUDA:"); - println!(" libcuda {}", yes_no(report.cuda.libcuda_loaded)); - println!( - " libcudart {}", - yes_no(report.cuda.libcudart_loaded) - ); - println!( - " xprobe CUPTI {}", - yes_no(report.cuda.xprobe_cupti_loaded) - ); - println!(" context {:?}", report.cuda.context.status); - println!(); - println!("Loaded shared libraries: {}", report.loaded_libraries.len()); -} - -fn print_discovery_result(result: &DiscoveryResult) { - println!("Root: PID {}", result.root.pid); - println!( - "CUDA workers: {} of {}{}", - result.candidates.len(), - result.total_candidates, - if result.truncated { " (truncated)" } else { "" } - ); - for candidate in &result.candidates { - println!( - " PID {} parent={} GPUs={} {}", - candidate.target.pid, - candidate.parent_pid, - candidate.gpu_uuids.join(","), - candidate.executable - ); - } - for warning in &result.warnings { - eprintln!("{}: {}", warning.code, warning.message); - } -} - -fn print_host_capture(result: &HostCaptureResult) { - println!( - "Captured {} event(s), dropped {}, timed out: {}", - result.captured, result.dropped, result.timed_out - ); - for event in &result.events { - println!( - " {} pid={} tid={} cpu={} probe={}", - event.timestamp_ns, - event.pid, - event.tid, - event - .cpu - .map_or_else(|| "-".to_owned(), |cpu| cpu.to_string()), - result.probe_id - ); - } -} - -fn print_resolved_probe(resolved: &ResolvedProbe) { - println!("Target: PID {}", resolved.target.pid); - println!("Binary: {}", resolved.binary_path); - println!( - "Build ID: {}", - resolved.build_id.as_deref().unwrap_or("unavailable") - ); - println!("Object: {:?}", resolved.object_kind); - println!("Probe: {:?}", resolved.probe_kind); - if let Some(symbol) = &resolved.symbol { - println!("Symbol: {symbol}"); - } - if let Some(symbol) = &resolved.symbol_demangled { - println!("Demangled: {symbol}"); - } - println!("File offset: {:#x}", resolved.file_offset); - println!("Runtime address: {:#x}", resolved.runtime_address); - println!( - "Mapping: {:#x}-{:#x} (file offset {:#x})", - resolved.mapping.start_address, resolved.mapping.end_address, resolved.mapping.file_offset - ); -} - -fn print_validation_result(result: &ValidationResult) { - println!( - "Validation: {}", - if result.valid { "valid" } else { "invalid" } - ); - println!("Target: PID {}", result.target.pid); - println!("Start: {}", result.start.selector); - println!("End: {}", result.end.selector); - println!("Match: {:?}", result.match_policy); - println!( - "Recommended match: {:?} ({:?})", - result.policy_recommendation.policy, result.policy_recommendation.reason - ); - println!( - "Agent activation: {:?}", - result.requirements.agent_activation - ); - println!("Target mutation: {}", result.requirements.target_mutation); - for issue in &result.issues { - println!("Issue: {}: {}", issue.code, issue.message); - } - for warning in &result.warnings { - println!("Warning: {}: {}", warning.code, warning.message); - } -} - -fn print_measurement_result(result: &MeasurementResult) { - let latency = &result.measurement.latency_ns; - println!( - "Measurement: {}", - result.measurement.name.as_deref().unwrap_or("unnamed") - ); - println!("Matched samples: {}", result.measurement.samples.matched); - println!("Evidence pairs: {}", result.evidence.len()); - println!("Correlation: {}", result.correlation.method); - println!("Clock: {}", result.clock.alignment); - println!("Latency (ns):"); - println!(" min {}", latency.min); - println!(" mean {:.2}", latency.mean); - println!(" p50 {}", latency.p50); - println!(" p90 {}", latency.p90); - println!(" p95 {}", latency.p95); - println!(" p99 {}", latency.p99); - println!(" max {}", latency.max); -} - -fn print_check(name: &str, result: &CheckResult) { - println!( - " {name:<17} {:<11} {}", - format!("{:?}", result.status).to_lowercase(), - result.detail.as_deref().unwrap_or_default() - ); -} - -const fn yes_no(value: bool) -> &'static str { - if value { "yes" } else { "no" } -} - -const fn schema_version(version: SchemaVersion) -> &'static str { - match version { - SchemaVersion::V2 => "2.0", - } -} diff --git a/xprobe/cli/src/render.rs b/xprobe/cli/src/render.rs new file mode 100644 index 0000000..217af2c --- /dev/null +++ b/xprobe/cli/src/render.rs @@ -0,0 +1,227 @@ +use xprobe_protocol::CapabilityReport; +use xprobe_protocol::CheckResult; +use xprobe_protocol::DiscoveryResult; +use xprobe_protocol::HostCaptureResult; +use xprobe_protocol::MeasurementResult; +use xprobe_protocol::ProcessReport; +use xprobe_protocol::ResolvedProbe; +use xprobe_protocol::SchemaVersion; +use xprobe_protocol::ValidationResult; + +pub(crate) fn print_human_report(report: &CapabilityReport) { + println!( + "xprobe doctor (schema {})", + schema_version(report.schema_version) + ); + println!(); + println!("Environment:"); + println!(" OS {}", report.environment.operating_system); + println!(" architecture {}", report.environment.architecture); + println!(" kernel {}", report.environment.kernel_release); + println!(" effective UID {}", report.environment.effective_uid); + println!( + " container {}", + report + .environment + .container + .as_deref() + .unwrap_or("none detected") + ); + println!(" PID namespace {}", report.environment.pid_namespace); + println!(); + println!("Capabilities:"); + println!(" uprobe {}", yes_no(report.capabilities.uprobe)); + println!( + " uretprobe {}", + yes_no(report.capabilities.uretprobe) + ); + println!( + " tracepoint {}", + yes_no(report.capabilities.tracepoint) + ); + println!( + " CUDA callback {}", + yes_no(report.capabilities.cuda_callback) + ); + println!( + " CUDA activity {}", + yes_no(report.capabilities.cuda_activity) + ); + println!( + " runtime inject {}", + yes_no(report.capabilities.runtime_injection) + ); + println!(); + println!("Checks:"); + print_check("BTF", &report.checks.btf); + print_check("eBPF permissions", &report.checks.ebpf_permissions); + print_check("kernel lockdown", &report.checks.kernel_lockdown); + print_check("perf paranoid", &report.checks.perf_event_paranoid); + print_check("ptrace scope", &report.checks.ptrace_scope); + print_check("NVIDIA driver", &report.checks.nvidia_driver); + print_check("CUDA driver", &report.checks.cuda_driver); + print_check("CUDA toolkit", &report.checks.cuda_toolkit); + print_check("CUPTI", &report.checks.cupti); + + if !report.warnings.is_empty() { + println!(); + println!("Warnings:"); + for warning in &report.warnings { + println!(" {}: {}", warning.code, warning.message); + } + } +} + +pub(crate) fn print_process_report(report: &ProcessReport) { + println!("Process: PID {}", report.target.pid); + println!("Executable: {}", report.executable); + println!("Start time: {} ticks", report.target.process_start_time); + println!("Namespace PIDs: {:?}", report.namespace_pids); + println!("Mount namespace: {}", report.mount_namespace); + println!( + "Credentials: uid={} gid={}", + report.credentials.effective_uid, report.credentials.effective_gid + ); + println!(); + println!("CUDA:"); + println!(" libcuda {}", yes_no(report.cuda.libcuda_loaded)); + println!( + " libcudart {}", + yes_no(report.cuda.libcudart_loaded) + ); + println!( + " xprobe CUPTI {}", + yes_no(report.cuda.xprobe_cupti_loaded) + ); + println!(" context {:?}", report.cuda.context.status); + println!(); + println!("Loaded shared libraries: {}", report.loaded_libraries.len()); +} + +pub(crate) fn print_discovery_result(result: &DiscoveryResult) { + println!("Root: PID {}", result.root.pid); + println!( + "CUDA workers: {} of {}{}", + result.candidates.len(), + result.total_candidates, + if result.truncated { " (truncated)" } else { "" } + ); + for candidate in &result.candidates { + println!( + " PID {} parent={} GPUs={} {}", + candidate.target.pid, + candidate.parent_pid, + candidate.gpu_uuids.join(","), + candidate.executable + ); + } + for warning in &result.warnings { + eprintln!("{}: {}", warning.code, warning.message); + } +} + +pub(crate) fn print_host_capture(result: &HostCaptureResult) { + println!( + "Captured {} event(s), dropped {}, timed out: {}", + result.captured, result.dropped, result.timed_out + ); + for event in &result.events { + println!( + " {} pid={} tid={} cpu={} probe={}", + event.timestamp_ns, + event.pid, + event.tid, + event + .cpu + .map_or_else(|| "-".to_owned(), |cpu| cpu.to_string()), + result.probe_id + ); + } +} + +pub(crate) fn print_resolved_probe(resolved: &ResolvedProbe) { + println!("Target: PID {}", resolved.target.pid); + println!("Binary: {}", resolved.binary_path); + println!( + "Build ID: {}", + resolved.build_id.as_deref().unwrap_or("unavailable") + ); + println!("Object: {:?}", resolved.object_kind); + println!("Probe: {:?}", resolved.probe_kind); + if let Some(symbol) = &resolved.symbol { + println!("Symbol: {symbol}"); + } + if let Some(symbol) = &resolved.symbol_demangled { + println!("Demangled: {symbol}"); + } + println!("File offset: {:#x}", resolved.file_offset); + println!("Runtime address: {:#x}", resolved.runtime_address); + println!( + "Mapping: {:#x}-{:#x} (file offset {:#x})", + resolved.mapping.start_address, resolved.mapping.end_address, resolved.mapping.file_offset + ); +} + +pub(crate) fn print_validation_result(result: &ValidationResult) { + println!( + "Validation: {}", + if result.valid { "valid" } else { "invalid" } + ); + println!("Target: PID {}", result.target.pid); + println!("Start: {}", result.start.selector); + println!("End: {}", result.end.selector); + println!("Match: {:?}", result.match_policy); + println!( + "Recommended match: {:?} ({:?})", + result.policy_recommendation.policy, result.policy_recommendation.reason + ); + println!( + "Agent activation: {:?}", + result.requirements.agent_activation + ); + println!("Target mutation: {}", result.requirements.target_mutation); + for issue in &result.issues { + println!("Issue: {}: {}", issue.code, issue.message); + } + for warning in &result.warnings { + println!("Warning: {}: {}", warning.code, warning.message); + } +} + +pub(crate) fn print_measurement_result(result: &MeasurementResult) { + let latency = &result.measurement.latency_ns; + println!( + "Measurement: {}", + result.measurement.name.as_deref().unwrap_or("unnamed") + ); + println!("Matched samples: {}", result.measurement.samples.matched); + println!("Evidence pairs: {}", result.evidence.len()); + println!("Correlation: {}", result.correlation.method); + println!("Clock: {}", result.clock.alignment); + println!("Latency (ns):"); + println!(" min {}", latency.min); + println!(" mean {:.2}", latency.mean); + println!(" p50 {}", latency.p50); + println!(" p90 {}", latency.p90); + println!(" p95 {}", latency.p95); + println!(" p99 {}", latency.p99); + println!(" max {}", latency.max); +} + +pub(crate) fn print_check(name: &str, result: &CheckResult) { + println!( + " {name:<17} {:<11} {}", + format!("{:?}", result.status).to_lowercase(), + result.detail.as_deref().unwrap_or_default() + ); +} + +pub(crate) const fn yes_no(value: bool) -> &'static str { + if value { "yes" } else { "no" } +} + +pub(crate) const fn schema_version(version: SchemaVersion) -> &'static str { + match version { + SchemaVersion::V2 => "2.0", + } +}