diff --git a/skills/data-loading-bottleneck/SKILL.md b/skills/data-loading-bottleneck/SKILL.md new file mode 100644 index 00000000000..9528afa2ba6 --- /dev/null +++ b/skills/data-loading-bottleneck/SKILL.md @@ -0,0 +1,245 @@ +--- +name: data-loading-bottleneck +description: "Diagnose input-bound PyTorch training. Use for low or bursty GPU utilization, slow batches, num_workers tuning, preprocessing regressions, or input stalls. Not for model/kernel optimization." +license: Apache-2.0 +compatibility: Requires Linux, Python, Git, an NVIDIA GPU and driver, and CUDA-enabled PyTorch. +permissions: + - file_read + - file_write + - shell + - network + - env +metadata: + author: "DALI Team " + tags: + - pytorch + - training + - performance + - data-loading + - profiling + - dali + languages: + - python + team: dali + domain: deep-learning + version: "1.0.0" +--- + +# Data Loading Bottleneck + +## Purpose + +Determine whether PyTorch training is input-bound and, if so, localize one cause without +changing the production path, lifecycle, data flow, or topology. + +## Prerequisites + +Requires Linux, Git, CUDA, and PyTorch. DALI replay and Nsight/NVTX profiling are optional. +Preflight attempts their allowed setup and routes around anything unavailable. + +## Instructions + +### 1. Preflight and route + +Create an artifact directory for commands and raw output, then run preflight with the +production Python: + +```bash + /scripts/collect_preflight.py \ + --source-dir --artifact-dir \ + --data-path \ + [--expected-visible-gpus N] +``` + +For remote or custom input, use `--data-source ` instead of `--data-path`. +Preflight records only the description, while the production run validates access. Stop on a +hard blocker. If a sandbox hides the accelerator, rerun preflight and GPU work in production. +CPU execution is not a substitute. Record any environment change made in response to a +warning. + +Start with the replay-support result from preflight. If replay is unavailable, use +`torch.version.cuda` to choose one package for a single isolated installation attempt: + +- CUDA 12.x: `nvidia-dali-cuda120` +- CUDA 13.x: `nvidia-dali-cuda130` +- Other or unknown: record the unsupported runtime and skip replay. + +```bash + -m pip install --target /dali-deps +``` + +After a successful installation, append the target with `site.addsitedir()` and retry +`from nvidia.dali.plugin.pytorch.loader_evaluator import LoaderEvaluator`. Keep +target-installed dependencies behind production packages. If the import succeeds, use the +same setup for Real and Replay. If it fails, preserve the failure, leave production unchanged, +and profile after Real. Do not try another installation strategy. Use the production Python +and worktree `PYTHONPATH` for every run. + +After preflight, record the original checkout status and diff, then create a disposable Git +worktree. Reproduce the canonical code and configuration, including staged, unstaged, and +relevant untracked changes. Put its package root or `src` directory first on `PYTHONPATH`, +point out-of-tree builds to it, and record representative module `__file__` paths. Stop if it +cannot reproduce the workload. Leave the original checkout untouched. Remove the worktree +after diagnosis and keep the artifacts. + +### 2. Instrument one bounded production run + +Instrument the production training path in the disposable worktree. Use a standalone harness +only when that path cannot be bounded or instrumented, as described in Troubleshooting. +Locate the last blocking loader retrieval at the intended replay boundary and the complete +optimizer update that consumes its batch. Record device transfer relative to the boundary, +batch/sample accounting, distributed topology, and implicit defaults. When +`LoaderEvaluator` is available, read `references/pytorch-dali.md` and build its paired Real +and Replay loaders. Without it, use the bounded production loader directly. + +Set `prefetch_depth` to the batches that can be ready at the replay boundary, including +production wrapper buffers. Per rank, it must be at least `num_workers * prefetch_factor`. +Use 2 when `prefetch_factor` is unset and one when `num_workers == 0`. Choose +`measured_batches > prefetch_depth`. Set warmup and drain to at least that depth, then set +`total_batches = warmup_batches + measured_batches + drain_batches`. Bound the source to +that total, warm the first part, time the measured window, and drain outside it. + +Keep the same workload, integration boundary, and timed window through localization. Add +only semantic ranges. Preserve work-affecting batch +structure, routing, topology, synchronization, and update behavior. Time complete steps and +every blocking loader retrieval. Synchronize the device and count samples at both window +boundaries. With multiple ranks, add boundary barriers and record each rank. Include normal +wait variability. In Real and Profile, do not manipulate the page cache or replace production +input with synthetic, repeated, modified, or deliberately pre-cached data. + +If instrumentation fails, follow Troubleshooting. If no valid Real window remains after +those attempts, report `INCONCLUSIVE` and continue at §6. + +Runs with material changes to the data source, sampling rules, batching, preprocessing, or +work-affecting input distribution are substitutes and cannot support canonical +`NOT DETECTED`. + +### 3. Run Real + +In a fresh process, run the bounded window through `LoaderEvaluator(mode="log")` when +available. Otherwise, use the bounded production loader. Per rank, record measured +samples, complete-step time, exposed loader wait, and the timestamp immediately before each +window barrier. Classify the source cache state as known warm before Real, warmed only by +this run's normal access, or unknown. + +```text +aggregate throughput = sum(samples across ranks) / max(rank window duration) +``` + +Do not infer balanced ranks from final arrival skew alone. Collectives can repeatedly +reconverge imbalanced ranks. Continue at §4 when `LoaderEvaluator` is available. Otherwise, +continue at §5. + +### 4. Run Replay and classify + +Start another fresh process with the same `LoaderEvaluator` wrapper in `replay` mode. Change +only the mode. Any other change invalidates the comparison. Apply the post-boundary work +equivalence and lifecycle checks in `references/pytorch-dali.md` before classification. + +If either run emits fewer than `total_batches`, report `INCONCLUSIVE` and continue at §6. +For any other unavailable or invalid Replay, continue at §5. + +For a valid comparison: + +```text +speedup = replay aggregate throughput / real aggregate throughput +``` + +| speedup | Verdict | +|---|---| +| `>1.50x` | `DETECTED` | +| `>1.10x` and `<=1.50x` | `POTENTIAL` | +| `<=1.10x` | `NOT DETECTED` | + +These fixed heuristics follow DALI's [Data Loading Bottleneck Detection tutorial](https://docs.nvidia.com/deeplearning/dali/user-guide/docs/examples/frameworks/pytorch/loader_evaluator/pytorch_data_loader_evaluator.html). +They are not estimates of run-to-run noise. + +Do not repeat Real or Replay. Continue at §6 after `NOT DETECTED`, and at §5 after +`DETECTED` or `POTENTIAL`. + +### 5. Profile and localize + +When §3 or §4 routes here, read `references/profiling.md`. If Nsight Systems, NVTX, or profile +summarization is still unavailable after its allowed setup, skip capture. Keep a valid Replay +verdict and mark localization unavailable. Without valid Replay, report `INCONCLUSIVE`. +Continue at §6. + +Otherwise, reuse a valid Real run as the unprofiled baseline. Follow the reference's +single-capture workflow and shared limit of one recapture for any reason. + +Apply the reference's structural and Real-baseline checks to each conclusion. A structurally +invalid capture is unusable. A work mismatch restricts only the conclusions it could affect, +while a valid Replay verdict remains authoritative. Without valid Replay, report `DETECTED` +only when validated profile evidence shows an input-path stage materially delaying full-step +progress in the measured window. Otherwise, report `INCONCLUSIVE`. + +Use that delay as detection evidence. Attribute a cause only to a measured concrete stage or +supported capacity limit. Leave an unsplit limiting stage as `unresolved composite`. When +localization supports an optimization, give one ranked, evidence-backed recommendation, +mark it untested, and do not implement or benchmark it. Otherwise, state the missing evidence. + +### 6. Report + +Complete `assets/report-template.md` and return it in the final response, *not* as a path to a +Markdown file. Remove unused sections and placeholders. Always keep the decision table, +Workload, Detection, and Confidence and scope. Keep Workload to one paragraph when the +canonical command ran as-is. + +- For `DETECTED` or `POTENTIAL`, give Cause and Next action. Include Localization when + profiling ran and Recommendation when supported. Without localization, set Cause to + `unresolved composite`. +- For `NOT DETECTED`, set Cause and Next action to `not applicable`, then omit Localization + and Recommendation. +- For `INCONCLUSIVE`, set Cause to `unresolved composite`, name the missing evidence under + Next action, and add Localization only when profiling ran. + +Add Substitutions only when the measured run differed from canonical. Add Distributed +behavior and Trace navigation only for `world_size > 1`. Use Real timing for every rank, +profiled GPU or NCCL values only for captured ranks, and map profiled PIDs to ranks. Mark +missing measurements `invalid` or `unavailable`. Report excluded input costs in Detection +with their value and frequency, never as Cause. A recommendation must cite its measurement, +mechanism, feasibility constraint, and `untested` status. + +Under Confidence and scope, keep only facts that could change the interpretation. Include +invalid and superseded artifacts with their rejection reasons. Use absolute artifact paths, +including the full `.nsys-rep` path. Cite the evidence for the verdict and recommendation, +and limit the result to the measured workload and recorded source cache state. + +## Available Scripts + +| Script | Purpose | Arguments | +|---|---|---| +| `scripts/collect_preflight.py` | Record environment, source, data, and optional-tool readiness in `preflight.json` | `--source-dir`, `--artifact-dir`, one of `--data-path` or `--data-source`; optional `--expected-visible-gpus` | +| `scripts/summarize_nsys.py` | Summarize the measured NVTX window, CUDA activity, and loader-wait/GPU-idle overlap | input `.nsys-rep` and required `--output` JSON path | + +Both scripts write to the requested path, print status to stdout, and send diagnostics to +stderr. Preflight returns `0` when ready, `1` when blocked, and `2` on an operational or +usage error. The summarizer returns `0` on success, `1` for an invalid report or output, and +`2` for invalid arguments. Where a `run_script` helper exists, use it with the same +repo-relative script and arguments. Otherwise, use the documented production-Python commands. + +## Examples + +- “Find out whether data loading is causing bursty GPU utilization in this PyTorch training + job, and identify the limiting stage.” +- “Optimize this model's attention kernels” is outside this skill's scope. + +## Limitations + +- Covers steady-state CUDA-enabled PyTorch training. Inference, startup, epoch transitions, + checkpointing, and offline data preparation are outside the workflow. +- Determines whether the input path limits full-step training throughput. It does not + optimize model, kernel, optimizer, or communication performance. +- Produces workload-specific conclusions. It does not predict behavior under another + topology, data path, or cache state. + +## Troubleshooting + +- **Instrumented run failure:** rerun the canonical command without instrumentation and save + both commands and outputs to separate workload failure from instrumentation failure. +- **Resource failure:** make at most one nearest runnable attempt that changes one resource + setting. Label it a substitute and restrict the verdict accordingly. +- **Production path cannot be bounded or instrumented:** save the failed production attempt + before using a standalone harness. Treat the harness as a substitute. State what differs or + is missing and limit every conclusion that depends on it. Without the complete training + step, it cannot establish that production is input-bound. diff --git a/skills/data-loading-bottleneck/agents/openai.yaml b/skills/data-loading-bottleneck/agents/openai.yaml new file mode 100644 index 00000000000..6126aa2f4e9 --- /dev/null +++ b/skills/data-loading-bottleneck/agents/openai.yaml @@ -0,0 +1,4 @@ +interface: + display_name: "Data Loading Bottleneck" + short_description: "Measure, profile, and diagnose training input stalls" + default_prompt: "Use $data-loading-bottleneck to check whether this training workload is input-bound." diff --git a/skills/data-loading-bottleneck/assets/report-template.md b/skills/data-loading-bottleneck/assets/report-template.md new file mode 100644 index 00000000000..49186547ad4 --- /dev/null +++ b/skills/data-loading-bottleneck/assets/report-template.md @@ -0,0 +1,93 @@ +# Data-loading result + +| Decision | Result | +|---|---| +| **Verdict** | **[DETECTED / POTENTIAL / NOT DETECTED / INCONCLUSIVE]** | +| **Cause** | [Supported concrete stage or capacity limit, or `unresolved composite`] | +| **Next action** | [One evidence-backed recommendation, or the evidence needed to resolve the result] | + +## Workload + +[Command, model, settings, environment, dataset source, format, cardinality, and cache state] + +**Substitutions** + +| Change | Why it was required | Expected bias and verdict scope | +|---|---|---| +| [canonical -> measured] | [reason] | [effect and where the conclusion applies] | + +## Detection + +| Run | Full-step window | Throughput | Exposed loader wait | +|---|---:|---:|---:| +| Real | [value, invalid, or unavailable] | [value, invalid, or unavailable] | [value and percent, invalid, or unavailable] | +| Replay | [value, invalid, or unavailable] | [value, invalid, or unavailable] | [value and percent, invalid, or unavailable] | + +**Result:** [Measured speedup, wait-only prediction from Real, their difference, and threshold +interpretation. If Replay is invalid or unavailable, report the primary effect] + +**Replay boundary:** [What replay bypassed and what remained, or why it was unavailable] + +**Distributed behavior** + +| Rank | Samples | Full-step window | Exposed loader wait | Pre-barrier timestamp | Useful GPU / NCCL | +|---:|---:|---:|---:|---:|---:| +| [rank] | [value] | [value] | [value and percent] | [value] | [value / value or unavailable] | + +**Rank result:** [Slowest-window aggregation, wait and arrival skew, and whether collectives +mask starvation] + +## Localization + +**Primary Nsight Systems report** + +```text +/absolute/path/to/trace.nsys-rep +``` + +**Profile summary** + +```text +/absolute/path/to/profile-summary.json +``` + +**Trace navigation** + +| Domain | PID(s) | Range names | +|---|---|---| +| [domain] | [PID list] | [names] | + +**GPU timeline** + +- [Active and idle result] +- [Transfer or other critical-path result] + +**End-to-end attribution** + +| Stage (main / worker / GPU) | Status | P50 | Critical-path evidence | Artifact | +|---|---|---:|---|---| +| [stage] | [measured / absent / composite / unavailable] | [value, N, and unit] | [dependency, observed overlap and denominator, or why unavailable] | [summary or timeline] | + +**Cause result:** [Supported actionable cause, unresolved composite and limit, or missing +evidence.] + +## Recommendation + +[Recommendation] + +## Confidence and scope + +- **Equivalence:** [Real <=> Replay: `pass` (`exact` or `work-equivalent`), `fail`, or + `unavailable`. Real <=> Profile: give one of those statuses for each conclusion. Cite the + relevant identities or batch signatures, steps, lifecycle, operating-regime evidence, and + any restrictions] +- **Missing coverage:** [Main-process ranges or worker PIDs without ranges] +- **Profiler perturbation:** [Profiled versus unprofiled difference, or unavailable] +- **Limits:** [For INCONCLUSIVE, name the evidence needed] + +**Artifacts** + +| Absolute path | Status | Purpose / reason | +|---|---|---| +| [/path/to/primary-artifact] | primary | [detection or diagnosis] | +| [/path/to/artifact] | [supporting / restricted / topology only / superseded / invalid] | [purpose, restriction, or rejection reason] | diff --git a/skills/data-loading-bottleneck/evals/config.yml b/skills/data-loading-bottleneck/evals/config.yml new file mode 100644 index 00000000000..38c7f347b3f --- /dev/null +++ b/skills/data-loading-bottleneck/evals/config.yml @@ -0,0 +1,16 @@ +schema_version: 1 + +harbor: + task_source: evals_json + custom_dockerfile_mode: preserve + base_image_mode: disabled + n_attempts: 1 + n_concurrent: 1 + max_agents: 1 + stop_on_pass: false + timeout_multiplier: 6 + auto_scale_timeout: true + sandbox: + template: harbor-eval-claude-code-gpu + pre_agent_setup: + - cp -a /opt/eval/workloads /workspace/workloads diff --git a/skills/data-loading-bottleneck/evals/environment/Dockerfile b/skills/data-loading-bottleneck/evals/environment/Dockerfile new file mode 100644 index 00000000000..3ab4f392e9a --- /dev/null +++ b/skills/data-loading-bottleneck/evals/environment/Dockerfile @@ -0,0 +1,41 @@ +FROM nvidia/cuda:13.3.1-base-ubuntu24.04 + +ENV DEBIAN_FRONTEND=noninteractive \ + VIRTUAL_ENV=/opt/venv \ + PATH=/opt/venv/bin:${PATH} + +RUN apt-get update && \ + apt-get install -y --no-install-recommends ca-certificates curl git gnupg python3 python3-venv && \ + curl -fsSL https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/7fa2af80.pub \ + -o /tmp/nvidia-devtools.pub && \ + gpg --dearmor --output /usr/share/keyrings/nvidia-devtools-keyring.gpg /tmp/nvidia-devtools.pub && \ + echo "deb [signed-by=/usr/share/keyrings/nvidia-devtools-keyring.gpg] https://developer.download.nvidia.com/devtools/repos/ubuntu2404/amd64/ /" \ + > /etc/apt/sources.list.d/nvidia-devtools.list && \ + apt-get update && \ + apt-get install -y --no-install-recommends nsight-systems-cli && \ + rm -rf /tmp/nvidia-devtools.pub /var/lib/apt/lists/* + +RUN python3 -m venv /opt/venv && \ + python -m pip install --no-cache-dir 'torch~=2.13' 'torchvision~=0.28' --index-url https://download.pytorch.org/whl/cu130 + +RUN python -m pip install --no-cache-dir 'nvidia-dali-cuda130~=2.3' 'nvtx~=0.2' timm==1.0.28 + +RUN mkdir -p /opt/eval/workloads && \ + git clone --branch v1.0.28 --depth 1 \ + https://github.com/huggingface/pytorch-image-models.git \ + /opt/eval/workloads/pytorch-image-models && \ + git clone --branch v0.28.0 --depth 1 \ + https://github.com/pytorch/vision.git \ + /opt/eval/workloads/vision && \ + chown -R 1000:1000 /opt/eval/workloads + +RUN mkdir -p /opt/eval/data && \ + curl -fsSL https://s3.amazonaws.com/fast-ai-imageclas/imagenette2.tgz \ + -o /tmp/imagenette2.tgz && \ + tar -xzf /tmp/imagenette2.tgz -C /opt/eval/data && \ + rm -f /tmp/imagenette2.tgz + +RUN command -v nsys && \ + python -c "import nvtx, timm, torch, torchvision; from nvidia.dali.plugin.pytorch.loader_evaluator import LoaderEvaluator; assert torch.version.cuda" + +WORKDIR /workspace diff --git a/skills/data-loading-bottleneck/evals/evals.json b/skills/data-loading-bottleneck/evals/evals.json new file mode 100644 index 00000000000..b366471efe0 --- /dev/null +++ b/skills/data-loading-bottleneck/evals/evals.json @@ -0,0 +1,48 @@ +{ + "skill_name": "data-loading-bottleneck", + "evals": [ + { + "id": "mobilenet-input-bound", + "prompt": "Training throughput dropped after our vision team moved its single-GPU jobs to a new worker image. The GPU now runs in short bursts and sits idle between them, especially with smaller models. Determine whether the input path materially limits full-step throughput for the reproduction below.\n\nRepository: /workspace/workloads/pytorch-image-models\nDataset: /opt/eval/data/imagenette2\nCommand (run from the repository root):\npython train.py --data-dir /opt/eval/data/imagenette2 --train-split train --val-split val --model mobilenetv3_small_100 --num-classes 10 --batch-size 32 --workers 4 --epochs 1 --amp --pin-mem --output /workspace/output", + "expected_output": "An inline data-loading report based on the stated timm training path. It gives a valid DETECTED verdict and follows the required profiling route before naming a cause or leaving the limiter unresolved.", + "files": [], + "assertions": [ + "The agent ran collect_preflight.py and diagnosed the supplied timm train.py path rather than replacing it with an independent harness.", + "The agent worked in a disposable Git worktree, preserved the original checkout, and bounded enough production batches for warmup, measurement, and drain.", + "The DETECTED verdict is supported by valid fresh-process Real and Replay windows whose only run-mode change was log versus replay and whose timed post-boundary work was shown to be equivalent.", + "The agent captured and inspected the required Nsight Systems profile, attributed only what the evidence supports, and returned the completed report inline rather than only linking to a Markdown file." + ], + "expected_skill": "data-loading-bottleneck", + "expected_script": "collect_preflight.py", + "should_trigger": true + }, + { + "id": "vit-not-input-bound", + "prompt": "GPU utilization is below what the training team expected and throughput is disappointing. One group blames JPEG decoding while another thinks the large ViT dominates the step. Determine whether the input path materially limits full-step throughput before we spend time rewriting the loader.\n\nRepository: /workspace/workloads/vision\nDataset: /opt/eval/data/imagenette2\nCommand (run from /workspace/workloads/vision/references/classification):\npython train.py --data-path /opt/eval/data/imagenette2 --model vit_l_16 --epochs 1 --batch-size 16 --workers 4 --amp --print-freq 20 --output-dir ''", + "expected_output": "An inline data-loading report based on the torchvision reference training path. It gives a valid NOT DETECTED verdict and stops without unsupported localization or optimization advice.", + "files": [], + "assertions": [ + "The agent ran collect_preflight.py and diagnosed the supplied torchvision reference train.py path rather than replacing it with an independent harness.", + "The agent used a disposable Git worktree and valid bounded Real and Replay windows with equivalent timed post-boundary work.", + "A Replay speedup at or below 1.10x leads to NOT DETECTED, and the agent does not profile, localize, or recommend a data-loading optimization after that verdict.", + "The final response contains the completed report inline with Cause and Next action marked not applicable." + ], + "expected_skill": "data-loading-bottleneck", + "expected_script": "collect_preflight.py", + "should_trigger": true + }, + { + "id": "inference-latency", + "prompt": "An online image-classification service built from /workspace/workloads/vision has developed a p99 inference-latency regression after deployment. Diagnose the inference problem.", + "expected_output": "The agent treats the request as an inference problem and does not run the training-focused data-loading diagnosis.", + "files": [], + "assertions": [ + "The agent recognizes that this incident concerns deployed inference rather than training input throughput.", + "The agent does not run collect_preflight.py, Real/Replay training comparisons, or a training Nsight workflow." + ], + "expected_skill": null, + "expected_script": null, + "should_trigger": false + } + ] +} diff --git a/skills/data-loading-bottleneck/references/profiling.md b/skills/data-loading-bottleneck/references/profiling.md new file mode 100644 index 00000000000..ea75af89eff --- /dev/null +++ b/skills/data-loading-bottleneck/references/profiling.md @@ -0,0 +1,356 @@ +# Profiling and Localization + +## Select the production target + +Profile the production command before isolating stages. Use the Nsight Systems result from +preflight and test the `nvtx` import. If needed, make one project-appropriate installation +attempt. If either tool remains unavailable, return to `SKILL.md` before instrumenting. + +For distributed training, choose the target from per-rank Real timing. Use rank 0 when +exposed waits are balanced. Otherwise, profile the rank with the largest exposed wait and its +workers. Per-rank timing carries skew, while `summarize_nsys.py` reads one report. An +NCCL-blocked rank may only show a starving peer's symptom. Keep ranges rank-local and +aggregate samples with the slowest-rank window. + +Select the rank with a launcher wrapper: + +```bash +#!/bin/sh +# torchrun --no-python --nproc_per_node= /bin/sh ./profile_rank.sh train.py ... +production_python=${1:?production Python required} +shift +if [ "${RANK:?torchrun did not set RANK}" = "${PROFILE_RANK:-0}" ]; then + exec nsys profile --output="${PROFILE_OUTPUT:-data-loading-rank-${RANK}}" "$production_python" "$@" +fi +exec "$production_python" "$@" +``` + +Retain the command, revision or diagnostic diff, profiler version, summary, and artifact +paths. Allow one capture and at most one recapture. Spend the recapture on an unusable +capture, excessive overhead, or one dominant composite. Then report any remaining limit. + +## Add semantic ranges + +Use manual annotations in the domain `data-loading`. Bare gaps in an unlabeled GPU trace +conflate loader waits, synchronization, and scheduler noise. Add these main-process ranges: + +- `capture_session`: outer capture activation before loader iteration and warmup +- `profile_window`: synchronized measurement boundary after warmup +- `batch_wait`: each blocking loader retrieval +- `train_step`: full consumer step exposed by the integration boundary +- `forward` and `backward`: separate compute phases. Use `compute` only when separation is + impossible and report that limitation + +Add `device_transfer`, `optimizer`, or synchronization when useful. Name boundaries +precisely. A callback after device transfer, for example, must not imply that `train_step` +includes transfer. + +```python +import nvtx + +DOMAIN = "data-loading" + +with nvtx.annotate("batch_wait", domain=DOMAIN): + batch = next(iterator) + +with nvtx.annotate("train_step", domain=DOMAIN): + with nvtx.annotate("forward", domain=DOMAIN): + loss = model(batch) + with nvtx.annotate("backward", domain=DOMAIN): + loss.backward() + optimizer.step() +``` + +`forward` and `backward` measure host-side launch of asynchronous CUDA work, not GPU +execution. The device usually synchronizes later, at loss readback or the step boundary. +Take GPU compute time from the CUDA activity in the profile summary, not from these range +widths. + +For managed loops, place equivalent ranges at the loader wrapper, batch callbacks, +production forward, and backward hooks. Use `nvtx.start_range`/`end_range` across +callbacks, and close every handle from a `finally` path. + +### Semantic colors + +Apply this Solarized mapping consistently to main and worker NVTX ranges. Annotate every +material workload-specific operation needed for attribution, even if it is absent from the +table, and use the nearest semantic color. Colors encode broad categories. Range names +identify exact operations. + +| Meaning | Color | +|---|---| +| Structural parents | `base00` `#657b83` | +| Source open/read | `cyan` `#2aa198` | +| Decode/conversion | `violet` `#6c71c4` | +| Preprocessing | `orange` `#cb4b16` | +| Batch assembly | `green` `#859900` | +| Loader wait | `yellow` `#b58900` | +| Compute | `blue` `#268bd2` | +| Device transfer | `red` `#dc322f` | +| Synchronization | `magenta` `#d33682` | + +Pass the colors as integers, not strings, to NVTX. For instance: `color=0x6c71c4`. + +## Mark workers + +With active workers, annotate the material stages present, such as source access, +materialization or parsing, per-sample processing, batch assembly, and handoff. Choose +boundaries from the workload code before capture. Split safely separable composites that +prevent attribution. Generic labels such as fetch, transform, and preprocess cannot support +a cause when their material children can be annotated safely. Refine a dominant composite +only if the recommendation could change. Otherwise, leave it unsplit and record the limit. +Image decode, augmentation, tokenization, and patchify are examples rather than requirements. + +Place ranges where work executes. Lazy APIs can move cost. Pillow `Image.open` parses +metadata but normally defers pixel decoding to `load()`, `convert()`, or pixel access. End +work ranges before `yield`, `await`, or blocking queue operations, and attribute asynchronous +work from its execution events. If instrumentation forces earlier materialization, confirm +that it preserves consumer-visible output, work, ordering, memory behavior, and relevant +timing. + +Annotate concrete operations rather than one range around the whole item: + +```python +import io +import nvtx +from PIL import Image +from torch.utils.data import Dataset + +WORKER_DOMAIN = "data-loading-worker" + +class ExampleDataset(Dataset): + def __getitem__(self, index): + with nvtx.annotate("read_file", domain=WORKER_DOMAIN, color=0x2AA198): + encoded = self.paths[index].read_bytes() + with nvtx.annotate("decode_jpeg", domain=WORKER_DOMAIN, color=0x6C71C4): + image = Image.open(io.BytesIO(encoded)) + image.load() + with nvtx.annotate("resize", domain=WORKER_DOMAIN, color=0xCB4B16): + image = self.resize(image) + with nvtx.annotate("random_crop", domain=WORKER_DOMAIN, color=0xCB4B16): + image = self.crop(image) + with nvtx.annotate("to_tensor", domain=WORKER_DOMAIN, color=0xCB4B16): + tensor = self.to_tensor(image) + with nvtx.annotate("normalize", domain=WORKER_DOMAIN, color=0xCB4B16): + return self.normalize(tensor) +``` + +Each range names one operation a recommendation can act on. A range around the entire +transform pipeline may dominate while identifying nothing removable. For an opaque +`transforms.Compose`, annotate each entry in its `transforms` list instead of the outer call. + +`image.load()` forces the decode inside its own range. Without it, Pillow defers that cost +into the next one. + +When the dataset cannot be edited, wrap it and split the wrapped range into concrete stages +before naming a cause. Preserve all loader behavior: + +- Define wrapper classes and functions at module scope for spawn and forkserver. This applies + only to definitions, not constructed `nvtx.annotate`, domain, or registered-message objects. + See the fork rule below. +- Preserve batched `Dataset.__getitems__`. Wrapping only `__getitem__` may disable a fast + path, and defining `__getitems__` on a dataset that lacks one forces the batched path + production never runs, because PyTorch selects it with `hasattr`. +- Wrap, do not replace, custom collation. +- Forward required attributes explicitly. Generic `__getattr__` can recurse during worker + reconstruction. + +Prefer direct stage annotations if wrappers would affect type checks, sharding, sampling, +batched fetch, or lifecycle. If worker instrumentation is unsafe or impossible, state why. +Main-process waits cannot identify worker sub-stages. + +### Register annotations in workers + +`nvtx.get_domain` and `Domain.get_registered_string` cache process-owned handles with +`lru_cache`. A forked worker inherits its ancestor's cached handles, so constructing the +annotation after the fork can still reuse an invalid handle. This also applies to a forkserver +that imported the module. Spawn starts a fresh interpreter and is unaffected. + +The failure is silent. The message appears in `StringIds` without child ranges, errors, or a +dropped-event signal, and an unmatched pop can corrupt neighboring counts. Keep only wrapper +definitions, domains, and message strings at module scope. Construct annotations inside the +worker without registering their pair in an ancestor. If that is impossible, use a +worker-only pair first registered in the child or report worker-stage coverage as +unavailable. Leave the caches intact. + +## Capture with Nsight Systems + +```bash +nsys profile \ + --trace=cuda,nvtx,osrt \ + --sample=none --cpuctxsw=process-tree \ + --capture-range=nvtx \ + --nvtx-capture=capture_session@data-loading \ + --capture-range-end=stop \ + --force-overwrite=true \ + --output=/data-loading \ + train.py ... +``` + +Add `--trace-fork-before-exec=true` only when worker-stage attribution requires following +active fork or forkserver workers. Omit it for `num_workers == 0` and spawn. This option can +substantially perturb, crash, or deadlock the target, especially alongside other process-tree +tracing. Omitting it for a fork-based loader makes worker-stage coverage unavailable. + +Start `capture_session` immediately before creating the profiled loader iterator so +activation cannot pause the consumer at the measurement boundary and fill the prefetch queue. +Warm inside the session, enclose the synchronized measurement in `profile_window`, and end +`capture_session` afterward. If a managed loop prevents activation before iterator startup, +activate early, discard the transient, and warm through at least one full prefetch cycle +before `profile_window`. An empty capture means the configured range name or domain never +fired. If `--cpuctxsw=process-tree` is unsupported, use `none` and do not infer worker CPU +pressure from scheduling. + +Keep `--capture-range-end=stop`: it ends collection after `capture_session` while allowing +training to continue. Nsight otherwise defaults to `stop-shutdown`, which terminates the +target. Close both structural ranges normally. The summarizer rejects an open +`profile_window`. + +Use manual ranges in this process-tree capture. Do not add `--python-functions-trace`. +Python tracing plus pre-exec worker following can materially perturb fork/forkserver loaders. + +If NVTX capture triggering fails, spend the recapture on one fallback. Prefer retaining the +outer range and triggering with the CUDA Profiler API: + +```python +torch.cuda.profiler.start() +try: + with nvtx.annotate("capture_session", domain="data-loading"): + # create the iterator and warm up + with nvtx.annotate("profile_window", domain="data-loading"): + # measured work + pass +finally: + torch.cuda.profiler.stop() +``` + +Replace `--capture-range=nvtx` with `--capture-range=cudaProfilerApi`. Retain +`--capture-range-end=stop`. If that trigger is unavailable, use a time-bounded capture +instead. Do not try both. + +## Summarize the measured window + +After a detailed capture, extract clipped range and CUDA statistics without loading raw +profiler dumps into context: + +```bash + /scripts/summarize_nsys.py \ + /data-loading.nsys-rep \ + --output /profile-summary.json +``` + +Use the JSON for range counts and distributions, worker PIDs, stored colors, GPU activity, +CUDA copies, and loader-wait/GPU-idle overlap. The summary establishes neither hierarchy nor +causality, so inspect the timeline before assigning a critical path, cause, or recommendation. +`batch_wait_overlapping_gpu_idle_percent` uses loader wait as its denominator. +`gpu_idle_overlapping_batch_wait_percent` uses GPU idle. + +## Validate the capture and overhead + +Use the profile summary and timeline to confirm: + +- one `capture_session` encloses warmup and measurement +- one `profile_window` encloses the measurement +- main-process range counts match the bounded window on every captured rank +- every worker expected to be captured has worker ranges +- capture extends beyond the prefetched queue. Otherwise, the trace shows queue drain, not + production behavior. + +An empty worker PID list with a limitation about ranges that "could not be named" means +coverage is indeterminate. Those ranges remain in the trace but the summary could not label +them. Missing worker PIDs without that limitation are a missing-marker failure unless worker +following was deliberately omitted. + +Missing required markers or invalid structural ranges make a capture unusable. Spend the +recapture if available, or report the gap. Worker ranges overlap, so never sum their +durations. Correlate them with queue consumption and main-process `batch_wait`. + +Restrict statistics to `profile_window` and state every overlap or utilization denominator. +Range `sum_ns` is an invocation sum rather than elapsed time. Never add nested parent and +child ranges, parallel workers, or separate CUDA category unions. Compute GPU active time +from the union of kernel, copy, and memset intervals. Use CUDA copy events, not host-range +width, to attribute transfers. + +Compare the profiled window directly with the retained unprofiled Real baseline. After +warmup, profiling must preserve the amount and path of measured full-step work and its +operating regime. Match batch count and work-affecting signatures, completed steps, producer +lifecycle, routing and topology, synchronization, and update behavior. Different records, +random values, resulting numerical state, or exact stall positions are acceptable unless +they change the work or critical-path relationship being reported. When content, order, or +locality can change a claim, compare the relevant workload or record identities. Use +identifiers when event correlation requires them. + +Accept profiler perturbation only when throughput stays within ~10% and no systematic stall +or lifecycle pattern appears. For a larger difference, use the recapture, if available. Keep +`cuda`, `nvtx`, manual ranges, and `--sample=none`. Drop `osrt` and set `--cpuctxsw=none`. +Retain `--trace-fork-before-exec=true` only for worker-stage attribution. Without it, limit +conclusions to the main process and GPU. After recapture, apply a work or operating-regime +mismatch only to the conclusions it could affect. Preserve conclusions supported by +independent equivalent work. + +## Interpret + +Build one end-to-end hierarchy from the workload code and trace, with one row per profile +summary range group. Include each applicable source/read, parsing or materialization, +per-record operation, batching or packing, handoff or IPC, exposed wait, device copy, +forward/loss, backward, optimizer, synchronization, and residual stage. Indent children and +mark every row measured, absent, composite, or unavailable. + +### Attribute and stop + +Begin with the observed loss of full-step progress. Correlate `batch_wait` with GPU-idle time, +device transfers with delayed model work, and producer stages with delivered batches. Trace +backward with batch or record identifiers when timing alone is insufficient. Report +wall-clock overlap with its denominator and reserve range sums for per-invocation service +demand. Name a limiting stage only when the timeline or measured capacity shows that it +materially gated batch delivery or model work. + +Use starvation as detection evidence only. Establish the cause separately. A dominant unsplit +stage remains `unresolved composite`, even when its aggregate demand predicts throughput. + +A call stack, blocked-thread state, scheduler event, or overlap identifies the wait site. It +does not quantify removable causal cost. Report `blocked at ` until the mechanism is +supported. + +Name an inferred mechanism such as a lock, GIL, allocator, IPC, or storage as the cause only +when measured service demand and available capacity predict observed throughput while +excluding competing paths. A controlled one-factor intervention that consistently improves +equivalent full-step windows provides the same support. Until then, label it +`candidate: ` and state the cheapest falsification test. Stop once a validated +trace identifies a concrete critical stage with sufficient evidence. A dominant stack frame +alone is insufficient. + +- Use OS/runtime events as corroboration for semantic ranges, not as a substitute. +- Continuously scheduled workers suggest CPU or memory/IPC pressure. Low scheduled time + plus long reads suggests storage. Queue waits can mean idleness or backpressure. +- Attribute transfer from GPU copy events and whether they delay model work, not + host-range width. +- In DDP, compare per-rank exposed waits and useful non-NCCL GPU work before blaming + collectives. Boundary arrival skew may stay small because ranks reconverge inside the + window. +- Compare equivalent kernel work before claiming input contention slowed compute. + +Map the measured critical path to one recommendation. Do not privilege worker scaling: + +- **Source access dominates:** recommend the indicated storage investigation or change, + such as latency, bandwidth, cache, format, or local staging. +- **Decode or preprocessing dominates:** recommend optimizing or eliminating the hot CPU + operation. When several consecutive compatible stages dominate and GPU headroom exists, + consider a contiguous accelerator pipeline from encoded input through model-ready GPU + tensors. Note operation and variable-shape support, batching, layout, dtype, randomness, + semantic equivalence, and GPU contention. Avoid partial offload that adds host-device + round trips. DALI replay proves bypass headroom, not production-pipeline feasibility. +- **Batch assembly or IPC dominates:** recommend examining payload bytes, early dtype + inflation, copies, shared memory, or packing. +- **Workers stay busy and the work is safely parallelizable:** recommend a bounded + `num_workers` experiment only when CPU, memory, and IPC headroom make it plausible. +- **Device transfer is critical:** recommend pinning or overlap based on GPU copy events. + +## Conditional cases + +- **Iterable/custom producer:** mark producer and queue boundaries. +- **Asynchronous/device preprocessing:** synchronize measurement boundaries and record + whether replay contains host or device data. +- **Remote storage:** mark the client/read path. Local file events may miss + network-library work. +- **Variable batches:** compare samples or bytes and verify shape distributions. diff --git a/skills/data-loading-bottleneck/references/pytorch-dali.md b/skills/data-loading-bottleneck/references/pytorch-dali.md new file mode 100644 index 00000000000..4eb835072ee --- /dev/null +++ b/skills/data-loading-bottleneck/references/pytorch-dali.md @@ -0,0 +1,123 @@ +# DALI Loader Evaluation + +## Build paired loaders + +Use identically bounded production sources and the same `LoaderEvaluator` configuration. +Apply the bound at loader construction or the production entrypoint: + +```python +import itertools +import torch + +class BoundedDataLoader(torch.utils.data.DataLoader): + def __init__(self, max_batches, **kwargs): + super().__init__(**kwargs) + self.max_batches = max_batches + + def __iter__(self): + return itertools.islice(super().__iter__(), self.max_batches) + + def __len__(self): + try: + return min(self.max_batches, super().__len__()) + except TypeError: + return self.max_batches + +total_batches = warmup_batches + measured_batches + drain_batches +bounded = BoundedDataLoader(total_batches, dataset=dataset, **production_loader_kwargs) +cache_batches = len(bounded) +source = LoaderEvaluator( + bounded, + mode="replay" if replay_mode else "log", + num_cached_batches=cache_batches, +) +``` + +Use `mode="log"` for Real and `mode="replay"` for Replay. Construction exhausts the bounded +source because `num_cached_batches` limits retention rather than consumption. Exclude +construction from timing. An OOM or unbounded construction is invalid. + +`LoaderEvaluator` requires a finite `len()`, though the dataset may remain unsized. The +adapter returns `max_batches`, while Real's emitted count still detects early exhaustion. +Replay emits only the declared length, so validate retained cache count and post-boundary +work separately. + +In both modes, warm `warmup_batches`, time `measured_batches`, then consume `drain_batches` +outside the window. The drain keeps source exhaustion and worker shutdown out of timing. + +Keep production loader kwargs, dataset, sampler, batch sampler, sharding, and `set_epoch` +unchanged. Before Replay, estimate retained memory from the largest Real boundary batch and +configured cache size. Account for every local rank when batches use shared host memory and +every device when they remain on GPU. If the cache clearly cannot coexist with training +state, record the estimate, skip Replay, and profile. + +Cache the whole bound when practical. A smaller cache cycles batches and is valid only when +the retained batches represent Real's timed post-boundary work and their objects are safe to +reuse. Mutations or aliasing may fail only on wraparound. Require +`0 < cache_batches <= len(bounded)`. Record configured and retained counts and whether Replay +cycles them. + +When a framework type-checks or re-instantiates loaders, use a framework-specific +`DataLoader` shim. Delegate iteration to the evaluator and preserve required production +attributes such as `dataset`, `sampler`, and the bound. Demonstrate equivalent behavior. For +Lightning DDP, disable automatic sampler replacement and validate each rank's lifecycle. + +## Prove equivalent timed work + +After warmup, Real and Replay are equivalent when cached substitution preserves the amount +and path of timed work after the replay boundary. Preserve batch structure, routing, topology, +synchronization, update behavior, and safe object use. Different records, random values, or +resulting numerical state are acceptable unless they change that work. + +Record sample counts and the work-affecting boundary signature: container and tensor +structure, shapes, dtypes, and bytes. For variable batches, record actual samples or bytes +and compare the shape sequence or distribution required by the consumer. Stable record keys +help when values or order affect work and when correlating stalls. Identity is supporting +evidence rather than a universal requirement. + +When exact order matters, use a dedicated sampler generator for map-style sources. For +iterable and framework-owned sources, use their seed or epoch controls and stable keys. Keep +the sampler generator separate from model randomness. An ordered shard list may substitute +for keys only when it determines the consumed content. + +Inspect the consumer from the replay boundary through the complete update for value- or +order-dependent branches, routing, sparsity, compilation, skipped updates, and stateful +caches. Match their controlling inputs or show unchanged execution. Keep the model, +optimizer, precision, initial state, steps, and synchronization fixed. Ordinary numerical +divergence is acceptable when timed work stays the same. + +For multiple ranks, cache the bounded source independently per rank while preserving +rank-local sharding and `set_epoch`. Confirm producers are quiescent during replay. Inspect +persistent-worker PIDs rather than assuming. + +Cache construction may advance RNG, epoch, or library state. Let the required warmup absorb +ordinary transients. Reject only persistent changes that survive warmup and alter timed work +or consumer resources beyond Replay's intended removal of producer work and contention. + +## Reject invalid Replay + +If either run emits fewer than `total_batches`, report `INCONCLUSIVE`. Reject Replay for: + +- unequal bounds, samples, or steps, or a cache that does not represent the timed + post-boundary work; +- a material difference in batch structure, routing, topology, synchronization, update + behavior, or other work that remains after warmup; +- mutated or aliased cached objects, unsafe asynchronous reuse, or cycling that changes the + timed work; or +- an unexpected producer or framework lifecycle change beyond Replay's intentional producer + quiescence. + +Do not place the emitted-count check after the last `yield` inside `__iter__`: a consumer +that stops at the bound never resumes the generator. Ordinary wait variability does not +invalidate Replay. + +## Interpret replay + +Compare measured speedup with the wait-only prediction `1 / (1 - w)`, where `w` is exposed +wait fraction. Use it as a reference rather than a limit. Replay may also free worker CPU +capacity, while a measured loader call may include retained work such as device copies in a +blocking prefetch wrapper. Report measured and expected speedup and their difference. +Agreement checks consistency but does not prove equivalent work. + +State what the cache bypasses and what remains. Report speedup as diagnostic headroom, never +as a promised gain. Prefetch can make queue drain look faster than worker supply. diff --git a/skills/data-loading-bottleneck/scripts/collect_preflight.py b/skills/data-loading-bottleneck/scripts/collect_preflight.py new file mode 100644 index 00000000000..2a1480f1f65 --- /dev/null +++ b/skills/data-loading-bottleneck/scripts/collect_preflight.py @@ -0,0 +1,321 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import argparse +import importlib +import importlib.metadata +import json +import os +import platform +import resource +import shutil +import subprocess +import sys +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + + +def _error(exc: BaseException) -> str: + return f"{type(exc).__name__}: {exc}" + + +def _limit(value: int) -> int | str: + return "unlimited" if value == resource.RLIM_INFINITY else value + + +def _space(path: Path) -> dict[str, int] | None: + try: + usage = shutil.disk_usage(path) + except OSError: + return None + return {"total_bytes": usage.total, "available_bytes": usage.free} + + +def _unescape_mount(value: str) -> str: + for escaped, plain in (("\\040", " "), ("\\011", "\t"), ("\\012", "\n"), ("\\134", "\\")): + value = value.replace(escaped, plain) + return value + + +def _mount_info(path: Path) -> dict[str, str] | None: + try: + resolved = path.resolve(strict=False) + best: tuple[int, dict[str, str]] | None = None + for line in Path("/proc/self/mountinfo").read_text().splitlines(): + before, after = line.split(" - ", 1) + left, right = before.split(), after.split() + mount_point = Path(_unescape_mount(left[4])) + try: + resolved.relative_to(mount_point) + except ValueError: + continue + info = { + "mount_point": str(mount_point), + "filesystem": right[0], + "source": _unescape_mount(right[1]), + } + candidate = len(str(mount_point)), info + if best is None or candidate[0] > best[0]: + best = candidate + return best[1] if best else None + except (OSError, ValueError, IndexError): + return None + + +def _path_info(path: Path) -> dict[str, Any]: + exists = path.exists() + probe = path if exists else path.parent + is_directory = path.is_dir() + # Directories need +x to be usable, and a missing path is judged by its parent directory. + execute = 0 if exists and not is_directory else os.X_OK + return { + "path": str(path.absolute()), + "exists": exists, + "is_directory": is_directory, + "readable": exists and os.access(path, os.R_OK | execute), + "writable": os.access(probe, os.W_OK | execute), + "space": _space(probe), + "mount": _mount_info(probe), + } + + +def _command_probe(command: str, args: list[str]) -> dict[str, Any]: + executable = shutil.which(command) + if not executable: + return {"available": False, "executable": None} + try: + proc = subprocess.run( + [executable, *args], capture_output=True, text=True, timeout=15, check=False + ) + text = (proc.stdout or proc.stderr).strip() + return { + "available": proc.returncode == 0, + "executable": executable, + "returncode": proc.returncode, + "version": text, + } + except (OSError, subprocess.SubprocessError) as exc: + return {"available": False, "executable": executable, "error": _error(exc)} + + +def _module_probe( + module: str, + distribution_names: tuple[str, ...] = (), + required_attribute: str | None = None, +) -> dict[str, Any]: + try: + loaded = importlib.import_module(module) + if required_attribute is not None and not hasattr(loaded, required_attribute): + return { + "available": False, + "module": module, + "error": f"Missing required attribute: {required_attribute}", + } + version = getattr(loaded, "__version__", None) + if version is None: + for distribution in distribution_names: + try: + version = importlib.metadata.version(distribution) + break + except importlib.metadata.PackageNotFoundError: + pass + return { + "available": True, + "module": module, + "required_attribute": required_attribute, + "version": version, + "path": getattr(loaded, "__file__", None), + } + except Exception as exc: + return {"available": False, "module": module, "error": _error(exc)} + + +def _torch_probe() -> dict[str, Any]: + try: + import torch + + cuda_available = torch.cuda.is_available() + devices = [ + { + "index": index, + "name": torch.cuda.get_device_name(index), + "memory_mib": torch.cuda.get_device_properties(index).total_memory >> 20, + } + for index in (range(torch.cuda.device_count()) if cuda_available else ()) + ] + return { + "available": True, + "version": torch.__version__, + "cuda_runtime": torch.version.cuda, + "cuda_available": cuda_available, + "visible_devices": devices, + } + except Exception as exc: + return { + "available": False, + "error": _error(exc), + "cuda_available": False, + "visible_devices": [], + } + + +def _arguments() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + data = parser.add_mutually_exclusive_group(required=True) + data.add_argument("--data-path", type=Path, help="local dataset path") + data.add_argument("--data-source", help="remote or custom dataset source") + parser.add_argument("--source-dir", required=True, type=Path, help="production Git checkout") + parser.add_argument("--artifact-dir", required=True, type=Path) + parser.add_argument("--expected-visible-gpus", type=int, help="minimum visible GPU count") + args = parser.parse_args() + if args.data_source is not None and not args.data_source.strip(): + parser.error("--data-source cannot be empty") + if args.expected_visible_gpus is not None and args.expected_visible_gpus < 1: + parser.error("--expected-visible-gpus must be at least 1") + return args + + +def main() -> int: + args = _arguments() + blockers: list[dict[str, str]] = [] + warnings: list[dict[str, str]] = [] + + def block(code: str, message: str) -> None: + blockers.append({"code": code, "message": message}) + + try: + args.artifact_dir.mkdir(parents=True, exist_ok=True) + except OSError as exc: + print( + f"preflight blocked: cannot create artifact directory: {_error(exc)}", file=sys.stderr + ) + return 2 + + data = ( + _path_info(args.data_path) + if args.data_path is not None + else {"kind": "remote_or_custom", "description": args.data_source} + ) + source = _path_info(args.source_dir) + artifact = _path_info(args.artifact_dir) + shared_memory = _path_info(Path("/dev/shm")) + torch_info = _torch_probe() + git = _command_probe("git", ["--version"]) + + if args.data_path is not None and (not data["exists"] or not data["readable"]): + block("data_unavailable", "Data path is missing or unreadable.") + source_ready = source["exists"] and source["is_directory"] and source["readable"] + if not source_ready: + block("source_unavailable", "Source directory is missing, unreadable, or not a directory.") + if not git["available"]: + block("git_unavailable", "Git is unavailable or its probe failed.") + if git["available"] and source_ready: + repo = _command_probe("git", ["-C", str(args.source_dir), "rev-parse", "--show-toplevel"]) + if repo["available"]: + source["git_root"] = repo["version"] + else: + block("source_not_git", "Source directory is not inside a Git worktree.") + if not torch_info["available"]: + block("torch_unavailable", "PyTorch cannot be imported by the production Python.") + elif not torch_info["cuda_available"]: + block("cuda_unavailable", "CUDA is not visible to the production Python.") + if ( + torch_info["cuda_available"] + and args.expected_visible_gpus is not None + and len(torch_info["visible_devices"]) < args.expected_visible_gpus + ): + block("insufficient_visible_gpus", "Visible GPU count is smaller than expected.") + + dali = _module_probe( + "nvidia.dali.plugin.pytorch.loader_evaluator", + ("nvidia-dali-cuda130", "nvidia-dali-cuda120", "nvidia-dali"), + "LoaderEvaluator", + ) + nvtx = _module_probe("nvtx", ("nvtx",)) + nsys = _command_probe("nsys", ["--version"]) + for code, label, probe in ( + ("dali_unavailable", "DALI replay", dali), + ("nvtx_unavailable", "NVTX", nvtx), + ("nsys_unavailable", "Nsight Systems", nsys), + ): + if not probe["available"]: + warnings.append( + {"code": code, "message": f"{label} is unavailable or its probe failed."} + ) + nofile = resource.getrlimit(resource.RLIMIT_NOFILE) + memlock = resource.getrlimit(resource.RLIMIT_MEMLOCK) + result = { + "schema_version": 1, + "timestamp_utc": datetime.now(timezone.utc).isoformat(), + "status": "blocked" if blockers else "ready", + "hard_blockers": blockers, + "warnings": warnings, + "invocation": { + "data_path": str(args.data_path) if args.data_path is not None else None, + "data_source": args.data_source, + "source_dir": str(args.source_dir), + "artifact_dir": str(args.artifact_dir), + "expected_visible_gpus": args.expected_visible_gpus, + }, + "python": { + "executable": sys.executable, + "version": platform.python_version(), + "implementation": platform.python_implementation(), + }, + "host": { + "platform": platform.platform(), + "cpu_count": os.cpu_count(), + "allowed_cpus": len(os.sched_getaffinity(0)), + "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"), + "limits": { + "open_files": [_limit(value) for value in nofile], + "locked_memory_bytes": [_limit(value) for value in memlock], + }, + }, + "paths": { + "data": data, + "source": source, + "artifacts": artifact, + "shared_memory": shared_memory, + }, + "torch": torch_info, + "dependencies": { + "dali_loader_evaluator": dali, + "nvtx": nvtx, + "nsys": nsys, + "git": git, + }, + } + + output = args.artifact_dir / "preflight.json" + try: + output.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n") + except OSError as exc: + print(f"preflight blocked: cannot write {output}: {_error(exc)}", file=sys.stderr) + return 2 + + print(f"preflight status={result['status']} json={output}") + # argparse rejects anything below 1, so the count is either a truthy int or None. + expected = args.expected_visible_gpus or "unspecified" + print(f"gpu visible={len(torch_info['visible_devices'])} expected={expected}") + print(f"dali={dali['available']} nvtx={nvtx['available']} nsys={nsys['available']}") + for item in blockers + warnings: + print(f"{item['code']}: {item['message']}", file=sys.stderr) + return 1 if blockers else 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/skills/data-loading-bottleneck/scripts/summarize_nsys.py b/skills/data-loading-bottleneck/scripts/summarize_nsys.py new file mode 100644 index 00000000000..664bfd2ae99 --- /dev/null +++ b/skills/data-loading-bottleneck/scripts/summarize_nsys.py @@ -0,0 +1,563 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import argparse +import json +import math +import sqlite3 +import subprocess +import sys +import tempfile +from collections import defaultdict +from contextlib import closing +from itertools import groupby +from operator import itemgetter +from pathlib import Path +from typing import Any, Iterable +from urllib.parse import quote + +# Nsight packs a process ID into the high bits of every global thread ID. +_PID_BITS = 24 +_PID_MASK = (1 << _PID_BITS) - 1 +_PROCESS_MASK = ~_PID_MASK +_RANGE_TYPES = {"NvtxPushPopRange", "NvtxStartEndRange"} +_INCOMPLETE_RANGE_TYPES = {"NvtxPushRange", "NvtxStartRange"} +_WINDOW_NAME = "profile_window" +_WINDOW_DOMAIN = "data-loading" +_WORKER_DOMAIN = "data-loading-worker" +_REQUIRED_COLUMNS = { + "StringIds": {"id", "value"}, + "ENUM_NSYS_EVENT_TYPE": {"id", "name"}, + "PROCESSES": {"globalPid", "pid", "name"}, + "NVTX_EVENTS": { + "start", + "end", + "eventType", + "color", + "text", + "globalTid", + "textId", + "domainId", + }, +} + + +class SummaryError(Exception): + pass + + +def _sql_list(names: Iterable[str]) -> str: + """Quote event-type names for an SQL IN clause. They are ASCII literals, never user input.""" + return ", ".join(repr(name) for name in sorted(names)) + + +def _row_pid(row: sqlite3.Row) -> int: + """The real PID when Nsight recorded the process, else the packed global ID.""" + return row["pid"] if row["pid"] is not None else row["global_pid"] + + +def _arguments() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("trace", type=Path, help="Nsight Systems .nsys-rep") + parser.add_argument("--output", required=True, type=Path, help="JSON summary path") + return parser.parse_args() + + +def _resolve_paths(args: argparse.Namespace) -> tuple[Path, Path]: + try: + trace = args.trace.expanduser().resolve(strict=True) + except OSError as exc: + raise SummaryError(f"input is unavailable: {exc}") from exc + output = args.output.expanduser().resolve(strict=False) + if output == trace: + raise SummaryError("output must differ from the input trace") + return trace, output + + +def _export_sqlite(trace: Path, directory: Path) -> Path: + output = directory / "trace.sqlite" + command = [ + "nsys", + "export", + "--type=sqlite", + "--quiet=true", + "--force-overwrite=true", + f"--output={output}", + str(trace), + ] + try: + result = subprocess.run(command, capture_output=True, text=True, check=False) + except OSError as exc: + raise SummaryError(f"failed to run nsys export: {exc}") from exc + if result.returncode != 0 or not output.is_file(): + detail = (result.stderr or result.stdout).strip() + raise SummaryError(f"nsys export failed: {detail or f'exit {result.returncode}'}") + return output + + +def _schema(connection: sqlite3.Connection) -> dict[str, set[str]]: + tables = { + row[0] for row in connection.execute("SELECT name FROM sqlite_master WHERE type = 'table'") + } + required = {"NVTX_EVENTS", "StringIds", "ENUM_NSYS_EVENT_TYPE", "PROCESSES"} + missing = sorted(required - tables) + if missing: + raise SummaryError(f"unsupported Nsight schema; missing tables: {', '.join(missing)}") + return { + table: {row[1] for row in connection.execute(f'PRAGMA table_info("{table}")')} + for table in tables + } + + +def _require_columns(schema: dict[str, set[str]], table: str, columns: set[str]) -> None: + missing = sorted(columns - schema.get(table, set())) + if missing: + raise SummaryError(f"unsupported Nsight schema; {table} lacks: {', '.join(missing)}") + + +def _metadata(connection: sqlite3.Connection, schema: dict[str, set[str]]) -> dict[str, str]: + if "META_DATA_EXPORT" not in schema or not {"name", "value"} <= schema["META_DATA_EXPORT"]: + return {} + return dict( + connection.execute("SELECT name, value FROM META_DATA_EXPORT WHERE value IS NOT NULL") + ) + + +def _merge_intervals(intervals: Iterable[tuple[int, int]]) -> list[tuple[int, int]]: + merged: list[list[int]] = [] + for start, end in sorted((start, end) for start, end in intervals if end > start): + if not merged or start > merged[-1][1]: + merged.append([start, end]) + else: + merged[-1][1] = max(merged[-1][1], end) + return [(start, end) for start, end in merged] + + +def _duration(intervals: Iterable[tuple[int, int]]) -> int: + return sum(end - start for start, end in _merge_intervals(intervals)) + + +def _intersection_duration( + left: Iterable[tuple[int, int]], right: Iterable[tuple[int, int]] +) -> int: + left, right = _merge_intervals(left), _merge_intervals(right) + total = i = j = 0 + while i < len(left) and j < len(right): + total += max(0, min(left[i][1], right[j][1]) - max(left[i][0], right[j][0])) + if left[i][1] < right[j][1]: + i += 1 + else: + j += 1 + return total + + +def _prepare_nvtx(connection: sqlite3.Connection) -> None: + names = _sql_list(_RANGE_TYPES | _INCOMPLETE_RANGE_TYPES | {"NvtxDomainCreate"}) + connection.executescript(f""" + CREATE TEMP VIEW normalized_nvtx AS + WITH raw AS ( + SELECT e.rowid AS event_id, e.start, e.end, e.color, + COALESCE(e.text, s.value) AS name, e.domainId, + (e.globalTid & {_PROCESS_MASK}) AS global_pid, t.name AS event_type + FROM NVTX_EVENTS e + JOIN ENUM_NSYS_EVENT_TYPE t ON t.id = e.eventType + LEFT JOIN StringIds s ON s.id = e.textId + WHERE t.name IN ({names}) + ), + latest_domain AS ( + SELECT global_pid, domainId, MAX(event_id) AS event_id + FROM raw WHERE event_type = 'NvtxDomainCreate' AND name IS NOT NULL + GROUP BY global_pid, domainId + ), + own_domain AS ( + SELECT r.global_pid, r.domainId, r.name + FROM raw r JOIN latest_domain d USING (global_pid, domainId, event_id) + ), + unambiguous_domain AS ( + SELECT domainId, MIN(name) AS name + FROM raw WHERE event_type = 'NvtxDomainCreate' AND name IS NOT NULL + GROUP BY domainId HAVING COUNT(DISTINCT name) = 1 + ) + SELECT r.*, COALESCE(p.pid, (r.global_pid >> {_PID_BITS}) & {_PID_MASK}) AS pid, + p.name AS process, COALESCE(own.name, any.name) AS domain + FROM raw r + LEFT JOIN own_domain own USING (global_pid, domainId) + LEFT JOIN unambiguous_domain any USING (domainId) + LEFT JOIN PROCESSES p ON p.globalPid = r.global_pid; + """) + + +def _select_window( + connection: sqlite3.Connection, trace_duration_ns: int +) -> tuple[int, dict[str, Any]]: + """Find the single measurement window, rejecting absent, unresolved, or duplicate ones.""" + candidates = list( + connection.execute( + f"""SELECT * FROM normalized_nvtx + WHERE event_type IN ({_sql_list(_RANGE_TYPES | _INCOMPLETE_RANGE_TYPES)}) + AND name = ? AND domain = ?""", + (_WINDOW_NAME, _WINDOW_DOMAIN), + ) + ) + incomplete = [ + str(_row_pid(row)) + for row in candidates + if row["end"] is None or int(row["end"]) > trace_duration_ns + ] + if incomplete: + raise SummaryError( + f"incomplete {_WINDOW_NAME!r} range for PID(s): {', '.join(sorted(incomplete))}" + ) + windows = [ + row + for row in candidates + if row["event_type"] in _RANGE_TYPES + and row["end"] is not None + and int(row["end"]) > int(row["start"]) + ] + if not windows: + raise SummaryError(f"no completed {_WINDOW_NAME!r}@{_WINDOW_DOMAIN} range") + if any(row["pid"] is None for row in windows): + pids = sorted(str(row["global_pid"]) for row in windows if row["pid"] is None) + raise SummaryError(f"cannot resolve profile-window process ID(s): {', '.join(pids)}") + if len(windows) != 1: + pids = ", ".join(str(row["pid"]) for row in sorted(windows, key=lambda row: row["pid"])) + raise SummaryError( + f"expected exactly one {_WINDOW_NAME!r}@{_WINDOW_DOMAIN} range; " + f"found {len(windows)} for PID(s): {pids}" + ) + selected = windows[0] + start, end = int(selected["start"]), int(selected["end"]) + window = { + "pid": int(selected["pid"]), + "process": selected["process"], + "start_ns": start, + "end_ns": end, + "duration_ns": end - start, + } + return int(selected["global_pid"]), window + + +def _unnamed_limitations(connection: sqlite3.Connection) -> list[str]: + """Ranges whose domain ID resolved to no name leave worker coverage indeterminate.""" + unnamed = list(connection.execute(f"""SELECT pid, global_pid FROM normalized_nvtx + WHERE event_type IN ({_sql_list(_RANGE_TYPES)}) AND end IS NOT NULL AND end > start + AND domain IS NULL AND domainId IS NOT NULL AND domainId != 0""")) + if not unnamed: + return [] + pids = sorted({str(_row_pid(row)) for row in unnamed}) + return [ + f"{len(unnamed)} completed NVTX range(s) in non-default domains could not be named " + f"for PID(s): {', '.join(pids)}; domain IDs were absent or ambiguous " + "across processes." + ] + + +def _clip_to_window( + connection: sqlite3.Connection, global_pid: int, window: dict[str, Any] +) -> None: + """Expose in-window ranges, clipped to the window, as the `profile_pieces` view.""" + start, end = window["start_ns"], window["end_ns"] + connection.execute(f"""CREATE TEMP VIEW profile_pieces AS + SELECT *, MAX(start, {start}) AS clipped_start, + MIN(end, {end}) AS clipped_end + FROM normalized_nvtx + WHERE event_type IN ({_sql_list(_RANGE_TYPES)}) AND end IS NOT NULL + AND end > {start} AND start < {end} + AND domain IN ('{_WINDOW_DOMAIN}', '{_WORKER_DOMAIN}') + AND (global_pid = {global_pid} OR domain = '{_WORKER_DOMAIN}')""") + + +def _range_summaries(connection: sqlite3.Connection) -> list[dict[str, Any]]: + """One entry per (pid, process, domain, name) group, with clipped duration statistics.""" + summaries = [] + rows = connection.execute("""SELECT pid, process, domain, COALESCE(name, '') AS name, + color, start, end, clipped_start, clipped_end + FROM profile_pieces WHERE clipped_end > clipped_start + ORDER BY domain, pid, name, process""") + # Grouping needs only contiguity, and ORDER BY covers the same four columns. + range_key = itemgetter("pid", "process", "domain", "name") + for (resolved_pid, process, domain, name), entries in groupby(rows, range_key): + entries = list(entries) + durations = [int(row["clipped_end"] - row["clipped_start"]) for row in entries] + partial_count = sum( + duration != row["end"] - row["start"] for duration, row in zip(durations, entries) + ) + durations.sort() + total = sum(durations) + summaries.append( + { + "pid": resolved_pid, + "process": process, + "domain": domain, + "name": name, + "uncolored_count": sum(row["color"] is None for row in entries), + "partial_count": partial_count, + "colors_argb": sorted( + { + f"#{int(row['color']) & 0xFFFFFFFF:08x}" + for row in entries + if row["color"] is not None + } + ), + "count": len(durations), + "sum_ns": total, + "mean_ns": total / len(durations), + # Nearest-rank percentiles: p95 is the maximum below 20 samples. + "p50_ns": durations[math.ceil(0.50 * len(durations)) - 1], + "p95_ns": durations[math.ceil(0.95 * len(durations)) - 1], + } + ) + return summaries + + +def _wait_intervals(connection: sqlite3.Connection, global_pid: int) -> list[tuple[int, int]]: + """Merged main-process loader waits inside the window.""" + return _merge_intervals( + (int(row[0]), int(row[1])) + for row in connection.execute( + """SELECT clipped_start, clipped_end FROM profile_pieces + WHERE global_pid = ? AND domain = ? AND name = 'batch_wait' + AND clipped_end > clipped_start""", + (global_pid, _WINDOW_DOMAIN), + ) + ) + + +def _coverage(connection: sqlite3.Connection, summaries: list[dict[str, Any]]) -> dict[str, Any]: + return { + "pids_by_domain": { + domain: sorted( + { + item["pid"] + for item in summaries + if item["domain"] == domain and item["pid"] is not None + } + ) + for domain in (_WINDOW_DOMAIN, _WORKER_DOMAIN) + }, + "unresolved_pid_events": connection.execute( + "SELECT COUNT(*) FROM profile_pieces WHERE clipped_end > clipped_start AND pid IS NULL" + ).fetchone()[0], + } + + +def _enum_labels( + connection: sqlite3.Connection, + schema: dict[str, set[str]], + table: str, +) -> dict[int, str]: + if table not in schema or not {"id", "label"} <= schema[table]: + return {} + return dict(connection.execute(f'SELECT id, label FROM "{table}"')) + + +def _enum_label(labels: dict[int, str], value: Any) -> str: + return labels.get(value, f"Unknown ({value})") + + +def _load_gpu( + connection: sqlite3.Connection, + schema: dict[str, set[str]], + global_pid: int, + window: dict[str, Any], + wait_intervals: list[tuple[int, int]], +) -> tuple[list[dict[str, Any]], list[str]]: + specifications = { + "kernel": ("CUPTI_ACTIVITY_KIND_KERNEL", "start end deviceId globalPid"), + "memcpy": ( + "CUPTI_ACTIVITY_KIND_MEMCPY", + "start end deviceId globalPid bytes copyKind srcKind dstKind", + ), + "memset": ("CUPTI_ACTIVITY_KIND_MEMSET", "start end deviceId globalPid"), + } + available: dict[str, tuple[str, list[str]]] = {} + limitations: list[str] = [] + for kind, (table, column_names) in specifications.items(): + columns = column_names.split() + if table not in schema or not set(columns) <= schema[table]: + limitations.append(f"{table} is unavailable; {kind} activity is omitted.") + else: + available[kind] = table, columns + if not available: + return [], limitations + + window_start, window_end = window["start_ns"], window["end_ns"] + copy_kinds = _enum_labels(connection, schema, "ENUM_CUDA_MEMCPY_OPER") + memory_kinds = _enum_labels(connection, schema, "ENUM_CUDA_MEM_KIND") + intervals: dict[int, dict[str, list[tuple[int, int]]]] = defaultdict( + lambda: {kind: [] for kind in specifications} + ) + copy_groups: dict[int, dict[tuple[str, str, str], dict[str, Any]]] = defaultdict(dict) + + for kind, (table, columns) in available.items(): + rows = connection.execute( + f"""SELECT {", ".join(columns)} FROM "{table}" + WHERE globalPid = ? AND end > ? AND start < ?""", + (global_pid, window_start, window_end), + ) + for row in rows: + start, end = max(int(row["start"]), window_start), min(int(row["end"]), window_end) + if end <= start: + continue + device_id = int(row["deviceId"]) + intervals[device_id][kind].append((start, end)) + if kind != "memcpy": + continue + key = ( + _enum_label(copy_kinds, row["copyKind"]), + _enum_label(memory_kinds, row["srcKind"]), + _enum_label(memory_kinds, row["dstKind"]), + ) + group = copy_groups[device_id].setdefault( + key, + {"count": 0, "bytes": 0, "partial_count": 0, "intervals": []}, + ) + group["count"] += 1 + group["bytes"] += int(row["bytes"]) + group["partial_count"] += (end - start) != int(row["end"]) - int(row["start"]) + group["intervals"].append((start, end)) + + wait = sum(end - start for start, end in wait_intervals) + summaries = [] + for device_id in sorted(intervals): + by_kind = intervals[device_id] + all_intervals = [interval for values in by_kind.values() for interval in values] + active = _duration(all_intervals) + wait_idle = wait - _intersection_duration(wait_intervals, all_intervals) + copies = [] + for (copy_kind, source, destination), group in sorted(copy_groups[device_id].items()): + copies.append( + { + "kind": copy_kind, + "source_memory": source, + "destination_memory": destination, + "count": group["count"], + "bytes": group["bytes"], + "partial_count": group["partial_count"], + "sum_ns": sum(end - start for start, end in group["intervals"]), + "union_ns": _duration(group["intervals"]), + } + ) + duration = window["duration_ns"] + idle = duration - active + summaries.append( + { + "pid": window["pid"], + "device_id": device_id, + "window_duration_ns": duration, + "active_union_ns": active, + "idle_ns": idle, + "active_percent": 100.0 * active / duration, + "kernel_union_ns": _duration(by_kind["kernel"]), + "memcpy_union_ns": _duration(by_kind["memcpy"]), + "memset_union_ns": _duration(by_kind["memset"]), + "batch_wait_union_ns": wait, + "batch_wait_gpu_idle_overlap_ns": wait_idle, + "batch_wait_overlapping_gpu_idle_percent": ( + 100.0 * wait_idle / wait if wait else None + ), + "gpu_idle_overlapping_batch_wait_percent": ( + 100.0 * wait_idle / idle if idle else None + ), + "copies": copies, + } + ) + if not summaries: + limitations.append( + "No CUDA kernel, memcpy, or memset events overlapped the selected window(s)." + ) + return summaries, limitations + + +def _summarize(connection: sqlite3.Connection, trace: Path) -> dict[str, Any]: + schema = _schema(connection) + metadata = _metadata(connection, schema) + if "ANALYSIS_DETAILS" not in schema or "duration" not in schema["ANALYSIS_DETAILS"]: + raise SummaryError("unsupported Nsight schema; ANALYSIS_DETAILS.duration is unavailable") + trace_duration_ns = connection.execute("SELECT MAX(duration) FROM ANALYSIS_DETAILS").fetchone()[ + 0 + ] + if not trace_duration_ns: + raise SummaryError("Nsight trace duration is unavailable") + for table, columns in _REQUIRED_COLUMNS.items(): + _require_columns(schema, table, columns) + _prepare_nvtx(connection) + global_pid, window = _select_window(connection, int(trace_duration_ns)) + limitations = _unnamed_limitations(connection) + _clip_to_window(connection, global_pid, window) + ranges = _range_summaries(connection) + coverage = _coverage(connection, ranges) + gpu, gpu_limitations = _load_gpu( + connection, schema, global_pid, window, _wait_intervals(connection, global_pid) + ) + return { + "source": { + "trace": str(trace), + "nsys_version": metadata.get("EXPORT_PRODUCT_VERSION"), + "export_schema_version": metadata.get("EXPORT_SCHEMA_VERSION"), + }, + "selection": { + "window_name": _WINDOW_NAME, + "window_domain": _WINDOW_DOMAIN, + "worker_domain": _WORKER_DOMAIN, + "windows": [window], + "window_union_duration_ns": float(window["duration_ns"]), + }, + "nvtx_ranges": ranges, + "gpu": gpu, + "coverage": coverage, + "limitations": limitations + gpu_limitations, + } + + +def main() -> int: + args = _arguments() + try: + trace, output = _resolve_paths(args) + with tempfile.TemporaryDirectory(prefix="nsys-summary-") as directory: + sqlite_path = _export_sqlite(trace, Path(directory)) + uri = f"file:{quote(str(sqlite_path))}?mode=ro" + with closing(sqlite3.connect(uri, uri=True)) as connection: + connection.row_factory = sqlite3.Row + summary = _summarize(connection, trace) + try: + with output.open("w", encoding="utf-8") as stream: + json.dump(summary, stream, indent=2) + stream.write("\n") + except OSError as exc: + raise SummaryError(f"cannot write summary: {exc}") from exc + except SummaryError as exc: + print(f"error: {exc}", file=sys.stderr) + return 1 + except sqlite3.Error as exc: + print(f"error: cannot read Nsight SQLite export: {exc}", file=sys.stderr) + return 1 + print( + json.dumps( + { + "output": str(output), + "windows": len(summary["selection"]["windows"]), + "nvtx_groups": len(summary["nvtx_ranges"]), + "gpu_groups": len(summary["gpu"]), + }, + separators=(",", ":"), + ) + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main())