From baa287e0c5d43123f3da2344b95a9d84c3ec910b Mon Sep 17 00:00:00 2001 From: itdevwu Date: Sat, 14 Nov 2026 04:10:15 +0800 Subject: [PATCH 1/4] =?UTF-8?q?=F0=9F=93=8C=20build:=20pin=20the=20Rust=20?= =?UTF-8?q?toolchain?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CI and release jobs installed whatever `stable` was current, so Rust 1.99 silently introduced new Clippy lints and changed the compiler that builds shipped binaries. Pin 1.99.1 in rust-toolchain.toml and install that toolchain in every workflow so local, CI, and release builds match. --- .github/workflows/ci.yml | 8 ++++---- .github/workflows/release.yml | 3 ++- docs/development.md | 4 +++- rust-toolchain.toml | 4 ++++ 4 files changed, 13 insertions(+), 6 deletions(-) create mode 100644 rust-toolchain.toml 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/docs/development.md b/docs/development.md index 6698344..7a26608 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 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"] From 1a99bc158bdf42282d445d01096733ab8564c7a0 Mon Sep 17 00:00:00 2001 From: itdevwu Date: Sat, 14 Nov 2026 04:15:11 +0800 Subject: [PATCH 2/4] =?UTF-8?q?=F0=9F=A7=AA=20test:=20pin=20the=20PyTorch?= =?UTF-8?q?=20live=20fixture=20environment?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The release live suite depended on an unrecorded local PyTorch env, and the compiled-Triton case failed late on Python 3.12, missing NumPy, or a runtime that registers only libcuda.so.1. Add environment-pytorch.yml (Python 3.11, torch 2.3.0, Triton, NumPy 1.x), check Dynamo, Triton, and NumPy before starting the workload, and point Triton at the CUDA link stub only when no unversioned libcuda.so is registered. --- docs/development.md | 14 +++++--- environment-pytorch.yml | 17 ++++++++++ tests/integration/test_pytorch_cuda.py | 47 ++++++++++++++++++++++++-- 3 files changed, 72 insertions(+), 6 deletions(-) create mode 100644 environment-pytorch.yml diff --git a/docs/development.md b/docs/development.md index 7a26608..8c58140 100644 --- a/docs/development.md +++ b/docs/development.md @@ -57,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 @@ -129,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/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( [ From 354a144fc17f4b0f9f77aa2eb9fe54dcfd55b36b Mon Sep 17 00:00:00 2001 From: itdevwu Date: Sat, 14 Nov 2026 04:15:11 +0800 Subject: [PATCH 3/4] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20refactor:=20split=20CL?= =?UTF-8?q?I=20orchestration=20into=20modules?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Move clap arguments, command failures, CPU and syscall inventory, live collection, hidden development commands, and human rendering out of main.rs into dedicated modules with explicit imports. main.rs keeps the entry point and the public measure, validate, doctor, and discover flow. No behavior changes. --- xprobe/cli/src/args.rs | 448 +++++ xprobe/cli/src/dev.rs | 525 ++++++ xprobe/cli/src/failure.rs | 166 ++ xprobe/cli/src/inventory.rs | 460 +++++ xprobe/cli/src/live.rs | 1432 +++++++++++++++ xprobe/cli/src/main.rs | 3431 ++--------------------------------- xprobe/cli/src/render.rs | 227 +++ 7 files changed, 3425 insertions(+), 3264 deletions(-) create mode 100644 xprobe/cli/src/args.rs create mode 100644 xprobe/cli/src/dev.rs create mode 100644 xprobe/cli/src/failure.rs create mode 100644 xprobe/cli/src/inventory.rs create mode 100644 xprobe/cli/src/live.rs create mode 100644 xprobe/cli/src/render.rs 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", + } +} From 6c997b2afe4e020a263e27cd57dc82d958534b53 Mon Sep 17 00:00:00 2001 From: itdevwu Date: Sat, 14 Nov 2026 04:30:07 +0800 Subject: [PATCH 4/4] =?UTF-8?q?=F0=9F=93=9D=20docs:=20record=20toolchain?= =?UTF-8?q?=20and=20fixture=20rules=20for=20agents?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- AGENTS.md | 7 +++++++ 1 file changed, 7 insertions(+) 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.