Aic Auto Collect

AIC Auto Collect

Start with the project playbook

Read docs/perf_database/collector-upgrade-playbook.md before changing collector code: it is the canonical upgrade/bring-up workflow. This skill is the run pipeline; the repository-owned policy in .claude/rules/collector/ is authoritative over anything restated here.

Goal

You are already on a GPU node with one or more GPUs. Your job is to collect AIC perf data for one target (GPU system, framework/backend, framework version), fix collector failures, recollect until the perf files are good enough, verify AIC can consume the new data, and maintain a draft PR to ai-dynamo/aiconfigurator as the checkpoint.

This is a long-running workflow. Be persistent. Iterate carefully. Do not hide failures by shrinking coverage or deleting hard cases unless the configuration is genuinely unsupported and documented.

Operating Mode

Assume the user wants the complete bring-up unless they explicitly ask for a smoke test only.

• Run collectors inside the target backend runtime containers. Do not claim data was collected for a backend/version unless the collector code ran in that container on the target GPU.
• Full collection means every supported registry op for the requested backend/version has either finished with real rows or is explicitly unsupported by that runtime/GPU with evidence.
• Smoke and --limit runs are only gates before the full run; never present them as final data.
• Keep working through multi-hour or multi-day runs. Use stable resume checkpoints, inspect errors, patch, and retry.
• The unit of work is one draft PR per (system, backend, backend_version). Use that PR as the checkpoint, not a private local branch.
• Make small sign-off commits as progress checkpoints after each collect_xx.py/op-family collection finishes. If a single collection runs longer than about one hour, checkpoint safe collector fixes, current perf outputs, and the failure/status summary before continuing.
• If multiple frameworks are requested, finish and checkpoint one framework/version PR first, then create a new draft PR for the next framework/version.
• Do not synthesize missing rows. Unsupported runtime/kernel paths should be filtered or documented, not filled with fake numbers.
• If the PR body says customer support is ready, verify the default AIC workflow can actually choose a valid configuration for at least the expected representative model/backend/system path.

Required Inputs

Establish these before collecting:

• Target AIC repo and branch.
• Target backend and version: sglang, vllm, or trtllm. Treat "all three" as an ordered queue of separate draft PRs.
• Target runtime container image for each backend, such as:
  • nvcr.io/nvidia/ai-dynamo/tensorrtllm-runtime:<tag>
  • nvcr.io/nvidia/ai-dynamo/sglang-runtime:<tag>
  • nvcr.io/nvidia/ai-dynamo/vllm-runtime:<tag>
• Target system name, such as rtx_pro_6000_server.
• Target ops: start narrow, then expand to all relevant registry ops.
• Whether to collect power columns.
• Whether the GPU type is already supported by AIC.

Discover the node:

nvidia-smi
nvidia-smi --query-gpu=name,memory.total,power.limit --format=csv
python3 - <<'PY'
import torch
print("cuda available:", torch.cuda.is_available())
print("device count:", torch.cuda.device_count())
for i in range(torch.cuda.device_count()):
    print(i, torch.cuda.get_device_name(i), torch.cuda.get_device_capability(i))
PY

Protect the host and other workloads

Run framework inspection and collection inside the pinned container. Treat the host driver, CUDA toolkit, kernel, container runtime, and system packages as immutable.

Before every GPU run, inspect:

docker ps
nvidia-smi
nvidia-smi --query-compute-apps=pid,process_name,used_gpu_memory --format=csv
df -h
df -h /dev/shm

Use task-unique container names and persistent output/checkpoint/log roots. Clean only task-owned workers and containers. Never kill or restart another container, restart Docker, or run a global prune. If another workload blocks the run, report it.

When the runtime root is read-only, mount every framework JIT/cache location that the selected path may create, not only the generic XDG cache. This can include DeepGEMM (~/.deep_gemm or DG_JIT_CACHE_DIR), TensorRT-LLM, TileLang, Triton/TorchInductor, CUDA, and FlashInfer caches. A family-wide EROFS failure while creating one of these directories is runner plumbing, not evidence that its shapes or kernel are unsupported. Preserve the rejected attempt, fix only the task runner, and rerun in a fresh namespace.

Treat an OOM as unclassified until the same case fails on a clean GPU. Check for stale workers, retained weights, descriptor/JIT caches, and oversized dummy allocations before adding a capacity rule.

Phase 1: Prepare Repo

1. Clone or update ai-dynamo/aiconfigurator.
2. Create a branch named like data/<system>-<backend>-<version> or codex/<system>-<backend>-<version>.
3. Immediately create a draft PR for this (system, backend, version) once the branch exists and the scope is known. Use it as the remote checkpoint even before all data is ready.
4. Install only lightweight local dev dependencies needed for tests. Do not install the target backend locally unless the collector is not running in a backend container.
5. Set:

export PYTHONPATH="$PWD"
export COLLECTOR_LOG_DIR="$PWD/collector_logs"
export COLLECTOR_CHECKPOINT_DIR="$PWD/collector_checkpoints"
mkdir -p "$COLLECTOR_LOG_DIR"
mkdir -p "$COLLECTOR_CHECKPOINT_DIR"

Local validation often works with the repo's pinned environment:

uv run --frozen ruff check <files>
uv run --frozen ruff format --check <files>
uv run --frozen pytest <tests> -q

If the local uv environment lacks framework packages such as torch, run framework-dependent collector smoke tests in the runtime container instead of trying to mutate the host environment.

Container pattern:

docker run --rm -it --gpus all --ipc=host --network host \
  --ulimit memlock=-1 --ulimit stack=67108864 \
  -v "$PWD:/workspace" -w /workspace \
  <runtime-image> bash

export PYTHONPATH=/workspace
export COLLECTOR_LOG_DIR=/workspace/collector_logs/<backend>/<version>/<op>
mkdir -p "$COLLECTOR_LOG_DIR"
python collector/collect.py --backend <backend> --ops <op> --smoke
python collector/collect.py --backend <backend> --ops <op> \
  --checkpoint-dir "$COLLECTOR_CHECKPOINT_DIR/<backend>/<version>/<op>" --resume

Inside the container, record the real framework version:

python - <<'PY'
import importlib.metadata as m
for pkg in ("tensorrt_llm", "sglang", "vllm"):
    try:
        print(pkg, m.version(pkg))
    except Exception:
        pass
PY

Update the draft PR body whenever scope changes or a meaningful checkpoint lands. Include current status, what is still partial, and links or paths to failure summaries.

Phase 2: Add Or Validate System Metadata

If the system is new or incomplete:

1. Add aic-core/src/aiconfigurator_core/systems/<system>.yaml.
2. Add <system> to SupportedSystems in aic-core/src/aiconfigurator_core/sdk/common.py.
3. Create aic-core/src/aiconfigurator_core/systems/data/<system>/.
4. Populate YAML with conservative, documented values:
   • gpu.mem_bw
   • gpu.mem_capacity
   • tensor core FLOPS fields relevant to the architecture
   • gpu.power
   • gpu.sm_version
   • node.num_gpus_per_node
   • inter-node, intra-node, PCIe bandwidth
   • misc.nccl_version
   • misc.other_mem

Prefer verified values from nvidia-smi, node docs, or nearby existing YAML files. If a value is an estimate, leave a comment.

Run:

pytest tests/unit/sdk/test_common.py -q

Phase 3: Framework Version And Kernel Preflight

Before collecting a new framework version or a new GPU SM target, check whether collector code matches the runtime. This avoids spending hours on known-bad grids.

1. Record the exact runtime package version from inside the container. Do not trust only the image tag.
2. Inspect collector/<backend>/registry.py for version routes and the concrete collect_xx.py files registered for this version.
3. Read the top of each relevant collect_xx.py file: docstring, __compat__, version notes, SM gates, and known unsupported filters.
4. Compare against the framework source for the exact version/tag when there is any sign of API or kernel churn:
   • SGLang: attention backends, DeepGEMM, FlashInfer, DeepSeek/DSV module code, MoE runner paths.
   • vLLM: attention/MLA modules, quantized GEMM paths, GDN/Mamba modules, custom all-reduce.
   • TRT-LLM: attention plugins, MLA/DSA modules, MoE, compute-scale, GDN/Mamba modules.
5. For a new SM such as SM120, check kernel architecture support before full collection: tiny smoke cases, import probes, and direct framework source guards.
6. Patch collector routing or test-case generation before full collection when the runtime cannot support a case. Use version routing and explicit SM filters rather than discovering the same failure across thousands of tasks.
7. If a special framework image exists for a model family or kernel family, try that image before declaring the op unsupported. Record its exact package version and image digest.

Commit and push preflight collector fixes before starting long full runs.

Phase 4: Choose Collection Plan

Read the backend registry:

python3 - <<'PY'
from collector.sglang.registry import REGISTRY as SGLANG
from collector.vllm.registry import REGISTRY as VLLM
from collector.trtllm.registry import REGISTRY as TRTLLM
for name, reg in [("sglang", SGLANG), ("vllm", VLLM), ("trtllm", TRTLLM)]:
    print(name, [e.op for e in reg])
PY

Recommended order:

1. gemm
2. moe
3. attention / MLA ops
4. module-level ops such as DSA, GDN, MHC, WideEP
5. communication ops separately, because they may require multi-GPU or full-node ownership

For new GPUs, prioritize ops used by the target models and backend first, but do not claim full support until all expected ops for that backend/version are collected or explicitly marked unsupported.

For an all-backend bring-up, run this as separate draft PRs:

1. Finish one GPU/framework/framework_version PR.
2. Push final data and collector fixes for that framework.
3. Create the next draft PR for the next framework/version.

For each backend/version, enumerate registry ops and create a tracking table with: op, collector file, perf filename, smoke status, full status, row count, duplicate-key status, AIC sanity status, and unsupported notes. Keep this table in a local scratch file and summarize it in the PR body.

Phase 5: Progressive Collection Loop

For each op:

python3 collector/collect.py --backend <backend> --ops <op> --smoke
python3 collector/collect.py --backend <backend> --ops <op> --shuffle --limit 20
python3 collector/collect.py --backend <backend> --ops <op> --shuffle --limit 100
python3 collector/collect.py --backend <backend> --ops <op> \
  --checkpoint-dir "$COLLECTOR_CHECKPOINT_DIR/<backend>/<version>/<op>" --resume

Use a separate log directory per op when iterating:

export COLLECTOR_LOG_DIR="$PWD/collector_logs/<backend>/<version>/<op>"
export COLLECTOR_CHECKPOINT_DIR="$PWD/collector_checkpoints"
mkdir -p "$COLLECTOR_LOG_DIR"
mkdir -p "$COLLECTOR_CHECKPOINT_DIR/<backend>/<version>/<op>"

If the process times out or crashes after partial progress, rerun with the same --checkpoint-dir and --resume. Keep checkpoints until the op is accepted. If the process has been running for about an hour, checkpoint safe progress:

• Commit collector fixes and tests.
• Copy currently produced perf files only if they are clearly labeled or known to be append-safe for replacement by a later full run.
• Push the draft PR branch.
• Update the PR body or a tracked status note with current op, completed cases, row counts, and active failures.

Full-collection acceptance for an op:

• The final collector run has zero unclassified errors, or every remaining skipped case is intentionally filtered before execution.
• The output perf file was copied from the full output directory, not from an earlier smoke/limited run.
• Row counts are plausible compared with nearby systems or the generated case count.
• Duplicate shape keys are removed only when they are true repeated measurements for the same key; keep the latest or best-defined row consistently and document the policy.
• The op is not required by AIC's default query path, or if it is required, representative AIC CLI/chart generation can consume it.

For very long runs, commit and push safe progress periodically:

git status --short
git add collector tests src/aiconfigurator/systems
git commit --signoff -m "<backend>: <op family> progress"
git push

Use sign-off commits. If DCO fails because a commit lacks Signed-off-by, amend only the latest commit if appropriate:

git commit --amend --no-edit --signoff
git push --force-with-lease

Phase 6: Fix Collector Errors

When a run fails, inspect:

• collection_summary_<backend>.json
• errors_*.json
• worker logs in COLLECTOR_LOG_DIR
• traceback module and task parameters

Maintain a failure ledger for the active PR. For each failure group record: op, collector file, runtime image/version, SM, task params, exact error, classification, action taken, and whether an upstream issue was filed.

Classify each error group:

• collector_bug: import/API/mock object/signature/test generation issue.
• unsupported_config: generated case is invalid for this SM/backend/kernel.
• framework_bug: backend kernel crashes or rejects a valid production-like case.
• resource_issue: OOM, shared memory, graph capture, worker restart, timeout.
• transient: rare cache/race/infra issue that passes on retry.

Fix policy:

• For API/import/signature errors:
  1. Read collector/<backend>/registry.py to map the failing op to its collector module.
  2. Read the failing collector/<backend>/collect_*.py.
  3. Search the framework source checkout for the missing class, function, import path, or attribute.
  4. Verify the real API in source. Do not trust Python's "Did you mean" suggestions blindly.
  5. Patch the collector with a minimal compatibility-preserving change.
• For API changes, use version routing and __compat__ where needed.
• For GPU capability gaps, filter test cases before execution using get_sm_version().
• For Blackwell/SM100+ features, gate FP4/NVFP4/FP8 paths carefully.
• For dimension constraints, filter on per-rank dimensions such as inter_size // tp.
• For CUDA-fatal errors, prefer worker restart and skip only the exact invalid region.
• Preserve data quality. Do not reduce benchmark repetitions or broad test dimensions just to pass.
• Keep older backend versions working.
• Keep shared test-case generators deterministic when unit tests depend on them. Put runtime availability probes in backend collector wrappers when the skip is caused by a partially installed framework package.
• Guard importlib.util.find_spec("a.b.c") with try/except ModuleNotFoundError; partial packages can make find_spec raise instead of return None.
• If a collector route only exists for future framework versions, use VersionRoute in collector/<backend>/registry.py and add a wrapper module with an explicit __compat__.
• If a collector helper silently returns None after a dry-run or subprocess failure, make it raise or record the skip explicitly so failed coverage is visible.
• If a valid production-like case fails in the framework kernel, file or prepare an upstream bug for SGLang, vLLM, or TRT-LLM with exact image/version, SM, minimal repro, and traceback. Link it in the PR or failure ledger.
• If the case is not production-like or the runtime explicitly does not support that SM/path, skip it in collect_xx.py before execution with a clear comment and a narrow predicate.

After a fix:

pytest tests/unit/collector -q
python3 collector/collect.py --backend <backend> --ops <op> --smoke
python3 collector/collect.py --backend <backend> --ops <op> --shuffle --limit 100
python3 collector/collect.py --backend <backend> --ops <op> \
  --checkpoint-dir "$COLLECTOR_CHECKPOINT_DIR/<backend>/<version>/<op>" --resume

Commit small fixes:

git add collector tests
git commit --signoff -m "fix <backend> <op> collector for <system>"