WIP: feat(collector): project AKS GPU pool driver mode into the snapshot

[yuanchen8911]

Summary

Add the K8s.aks-gpu-pools.gpu-driver snapshot reading: aicr snapshot --aks-gpu-pools <file> projects an operator-supplied az aks nodepool list -o json dump into the K8s measurement, recording each GPU agent pool's gpuProfile.driver ownership mode per ADR-015 DD3.

Motivation / Context

The projection half of rollout PR 2 under the ADR-015 implementation umbrella #1761 (ADR-015: recipe configuration profiles, Deferred Decision 3), split out per the sequencing recorded on #1761. The AKS gpuStack profile (#1967) requires a distinguishing snapshot signal for its two values; DD3 resolves that signal as the AgentPool gpuProfile.driver property, and PR 2 must supply "a bounded, explicit way to add that projection to a snapshot without adding Azure SDK/auth coupling to the generic profile core."

This PR is that mechanism: explicit operator input, no Azure SDK, no auth. The reading is inert until a declaration references it — no in-tree recipe does yet — so this can merge independently of #1933.

Fixes: N/A
Related: #1761, #1967 (consumer; merges after this and #1933)

Type of Change

• New feature (non-breaking change that adds functionality)

Component(s) Affected

• CLI (cmd/aicr, pkg/cli)
• Collectors / snapshotter (pkg/collector, pkg/snapshotter)
• Docs/examples (docs/, examples/)

Implementation Notes

Projection semantics (pkg/collector/k8s/aksgpupools.go):

• NVIDIA GPU pools are identified by VM-size family (Standard_NC*/ND*/NV*), minus AMD accelerators: the NG family, MI300X/MI325X (ND sizes AKS requires creating with --gpu-driver none), and Radeon Pro V620/V710 (NV sizes). Without the marker exclusion, a supported NVIDIA-Install + AMD-MI300X cluster would falsely project Mixed. Known limitation (documented in code): detection is prefix+marker based, so a hypothetical future NVIDIA family outside NC/ND/NV would be skipped; the AgentPool object carries no vendor field to do better.
• Per pool: gpuProfile.driver when set; absent or null gpuProfile follows the provider's documented Install default; a pool with gpuProfile.nvidia present projects as Managed (fully AKS-managed, out of AICR scope).
• Aggregate: all pools agree → that value; disagreement → Mixed. Managed, Mixed, and unknown provider strings deliberately match no declared profile constraint, so profile-qualified resolution fails closed naming the observed state.
• No GPU pools → the gpu-driver key is omitted; constraint evaluation reports the reading unavailable and fails closed.
• gpu-pool-count and a sorted gpu-pools roster (name=mode,...) accompany the reading for diagnostics.

The projection runs at the orchestration layer, not in a collector. It is pure file processing, so it never enters the cluster: pkg/snapshotter projects the file up front (a bad file fails the run in milliseconds, before any deploy or collection) and attaches the subtype to the K8s measurement. In agent Job mode the controller-side CLI merges the projection into the snapshot the Job returns — the file never needs to be visible in-pod, so Job mode works with no ConfigMap forwarding. The Job's result ConfigMap is then rewritten with the merged bytes (Cleanup deletes Job+RBAC but never that ConfigMap, so without the rewrite a stale projection-less snapshot would persist there — including when the ConfigMap is the user-requested output). This replaces the earlier collector-field + env-forwarding design and its Job-mode rejection: the reading is now producible by the ordinary aicr snapshot workflow in both modes, which #1967's mandatory constraint requires. aicr validate accepts the same flag for its live-capture path.

Fail-loud contract. The file arrives via an explicit flag, so open/decode/size failures (including a top-level JSON null, which json.Unmarshal would otherwise silently accept into a slice) are errors that fail the snapshot run — never a degraded measurement, and never subject to the snapshotter's degrade-to-warning collector policy. The read is size-bounded (os.Open + io.LimitReader, 1 MiB cap in pkg/defaults) and gated to regular files (Lstat: no FIFOs, no symlinks).

Design choice for review — explicit input vs. automatic acquisition. This lands the ADR-friendly floor: the operator runs az aks nodepool list -g <rg> --cluster-name <name> -o json and passes the file. An az shell-out or IMDS-based acquisition would be a convenience follow-up, not a replacement — the explicit path stays as the airgap-friendly, auth-free baseline either way.

Generalization path (GKE and beyond). The concept extends to other clouds, split deliberately across two layers:

• Per-provider, never shared: the projection. Each cloud expresses driver ownership in its own control-plane object (GKE: node-pool gpuDriverInstallationConfig). A new provider adds its own projector emitting its own namespaced subtype (gke-gpu-pools beside aks-gpu-pools) — no cross-provider pools schema, so each subtype's fail-closed constraint semantics stay exact.
• Shared: everything around the projection. readBoundedPoolsFile (pkg/collector/k8s/providerpools.go) carries the bounded fail-loud read (Lstat regular-file gate + os.Open + io.LimitReader); the flag plumbing pattern and the orchestration pre-flight are reused as-is.
• Additive flags: new providers get sibling flags (--gke-gpu-pools, …), never a generic flag with a provider discriminator — the flag name tells the operator which cloud CLI command produces the input, and the parser knows the schema without sniffing.

The pattern is documented for contributors in docs/contributor/collector.md ("Provider Node-Pool Projections").

Public SDK surface. The pkg/client/v1 facade AgentConfig carries AKSGPUPoolsPath (translated in toInternalAgentConfig, pinned by an SDK-level test), so Client.CollectSnapshot can produce the reading through the documented collect-then-resolve workflow. This is deliberately different from ClusterConfigPath/DiscoverNetwork, which stay off the facade — those are optional in-pod network enrichments, whereas this is controller-side input that AKS profile-qualified resolution requires.

Shared contract with #1967. TestReadingShapeMatchesProfileContract pins the reading to the literal subtype/key/values the AKS declaration's constraints name, so drift between the two PRs surfaces as a test failure.

Remaining work before undraft

• Review sign-off on the explicit-input design (vs. az shell-out / node-label alternatives).
• docs/integrator/aks-gpu-setup.md snippet showing the az command + flag in the snapshot flow ("Recording the pool mode in snapshots" under GPU Driver Setup).
• make qualify full gate.

Testing

go test ./pkg/... ./cmd/...   # 75 packages ok
golangci-lint run -c .golangci.yaml ./...   # 0 issues

Table-driven projection tests cover: all-Install, all-None, absent and null gpuProfile defaulting to Install, Mixed, Managed, unknown-value preservation, VM-family filtering (NV in; NG, MI300X-class ND, and Radeon NV AMD sizes out; AMD-only cluster omits the reading), and the no-GPU-pools omission. Fail-loud tests cover missing file, wrong JSON shape (the {"value":[]} trap), top-level JSON null, malformed JSON, the size cap, and non-regular files. Orchestration tests pin both modes: local-mode Measure attaches the subtype and fails before any collector runs on a bad file; the Job-mode mergeAKSGPUPools round-trip lands the subtype on the returned snapshot. Coverage: pkg/collector/k8s 89.5%.

Risk Assessment

• Low — Isolated change, well-tested, easy to revert

Rollout notes: No behavior change unless --aks-gpu-pools is passed; the reading is inert until a recipe declaration references it (#1967).

Checklist

• Tests pass locally (make test with -race)
• Linter passes (make lint)
• I did not skip/disable tests to make CI green
• I added/updated tests for new functionality
• I updated docs if user-facing behavior changed
• Changes follow existing patterns in the codebase
• Commits are cryptographically signed (git commit -S) — GPG signing info

[github-actions]

Recipe evidence check

No leaf overlays affected by this PR.

This gate is warning-only and never blocks merge.

[coderabbitai]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

📝 Walkthrough

Walkthrough

Adds an --aks-gpu-pools option for snapshot and validate commands. The option reads a bounded local AKS nodepool JSON dump, projects NVIDIA GPU pool driver data into a new Kubernetes subtype, and supports fail-closed input validation. Local snapshots attach the projection during collection; agent-backed snapshots merge it into returned YAML and ConfigMap output. Tests cover projection, filtering, errors, collection, merging, and rewriting. Documentation describes the CLI and contributor integration.

Estimated code review effort: 4 (Complex) | ~60 minutes

Suggested reviewers: njhensley

🚥 Pre-merge checks | ✅ 4

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly describes the main change: projecting AKS GPU pool driver mode into snapshots.
Description check	✅ Passed	The description is directly related to the changeset and accurately explains the new AKS GPU pool projection flow.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 4

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@docs/user/cli-reference.md`:
- Line 103: Update the --aks-gpu-pools documentation entry to include a runnable
example that exports node pools with az aks nodepool list -o json into
pools.json and invokes aicr snapshot --aks-gpu-pools pools.json. Also document
the 1 MiB input-size limit while preserving the existing behavior and mode
details.

In `@pkg/collector/k8s/aksgpupools.go`:
- Around line 151-162: Update readAKSAgentPools and its projection callers to
accept and propagate ctx, reject non-regular inputs before opening the path, and
preserve the existing error behavior for invalid files. Replace the
uninterruptible read flow with cancellation-aware processing that checks
ctx.Done() during long-running input handling, while retaining the
MaxAKSGPUPoolsBytes limit.

In `@pkg/collector/k8s/k8s.go`:
- Around line 80-91: Update Collect to pass the successfully parsed aksGPUPools
subtype into emptyK8sMeasurement() on both getClient() and discovery failure
paths. Preserve the existing parsing and error propagation while ensuring
Kubernetes fallback measurements retain the explicit pool projection.

In `@pkg/snapshotter/snapshot.go`:
- Around line 154-156: Update the snapshot collection flow around
parseAKSGPUPoolsPathEnv and collector.WithAKSGPUPoolsPath so missing, malformed,
or oversized explicitly configured AKS pool files propagate as fatal errors
instead of being swallowed by collectSafe. Pre-validate the configured pool
input or distinguish its INVALID_REQUEST result during collection, while
preserving collectSafe’s non-fatal behavior for ordinary collector degradation.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Enterprise

Run ID: b272054f-7bbc-4e4a-9f23-eea141e8c1be

📥 Commits

Reviewing files that changed from the base of the PR and between 963aa9f and 9c65244.

📒 Files selected for processing (11)

• docs/user/cli-reference.md
• pkg/cli/snapshot.go
• pkg/collector/factory.go
• pkg/collector/k8s/aksgpupools.go
• pkg/collector/k8s/aksgpupools_test.go
• pkg/collector/k8s/k8s.go
• pkg/defaults/timeouts.go
• pkg/k8s/agent/job.go
• pkg/k8s/agent/types.go
• pkg/snapshotter/agent.go
• pkg/snapshotter/snapshot.go

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@pkg/collector/k8s/aksgpupools.go`:
- Around line 162-196: Harden readAKSAgentPools against pathname swaps by
replacing the separate Lstat-then-Open validation with an atomic open using
no-follow/nonblocking flags, then validate the opened file descriptor is regular
before reading. Preserve the existing size limit and error behavior, and ensure
FIFO or symlink substitutions cannot cause a blocking read.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Enterprise

Run ID: cdcab19f-8940-4ad1-aa9c-c86019d8cf6b

📥 Commits

Reviewing files that changed from the base of the PR and between dcfe7e0 and 50b5fbb.

📒 Files selected for processing (11)

• docs/user/cli-reference.md
• pkg/cli/snapshot.go
• pkg/collector/factory.go
• pkg/collector/k8s/aksgpupools.go
• pkg/collector/k8s/aksgpupools_test.go
• pkg/collector/k8s/k8s.go
• pkg/defaults/timeouts.go
• pkg/k8s/agent/job.go
• pkg/k8s/agent/types.go
• pkg/snapshotter/agent.go
• pkg/snapshotter/snapshot.go

[coderabbitai]

♻️ Duplicate comments (1)

pkg/snapshotter/snapshot.go (1)

154-163: 🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift

Preserve fail-loud behavior for the read that actually feeds the snapshot.

Pre-validating the path does not protect the subsequent collector read. A file replacement or deletion between reads is converted by collectSafe into a successful snapshot with no AKS GPU-pool measurement, contradicting the documented explicit-input contract.

• pkg/snapshotter/snapshot.go#L154-L163: pass the parsed input into collection or propagate failures for explicitly configured AKS input.
• pkg/cli/snapshot.go#L517-L523: perform host-file validation only in local mode so Job mode reaches its intended host-path INVALID_REQUEST rejection.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@pkg/snapshotter/snapshot.go` around lines 154 - 163, The explicit AKS
GPU-pool input must fail if the collector’s actual read fails, rather than being
softened by collectSafe. In pkg/snapshotter/snapshot.go lines 154-163, pass the
parsed path into collection or propagate read errors for explicitly configured
input. In pkg/cli/snapshot.go lines 517-523, restrict host-file validation to
local mode so Job mode preserves its host-path INVALID_REQUEST behavior.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Duplicate comments:
In `@pkg/snapshotter/snapshot.go`:
- Around line 154-163: The explicit AKS GPU-pool input must fail if the
collector’s actual read fails, rather than being softened by collectSafe. In
pkg/snapshotter/snapshot.go lines 154-163, pass the parsed path into collection
or propagate read errors for explicitly configured input. In pkg/cli/snapshot.go
lines 517-523, restrict host-file validation to local mode so Job mode preserves
its host-path INVALID_REQUEST behavior.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Enterprise

Run ID: c7d64cd1-5b8a-4e02-852e-41f1dc8d69dc

📥 Commits

Reviewing files that changed from the base of the PR and between 50b5fbb and 4dd7e51.

📒 Files selected for processing (13)

• docs/contributor/collector.md
• docs/user/cli-reference.md
• pkg/cli/snapshot.go
• pkg/collector/factory.go
• pkg/collector/k8s/aksgpupools.go
• pkg/collector/k8s/aksgpupools_test.go
• pkg/collector/k8s/k8s.go
• pkg/collector/k8s/providerpools.go
• pkg/defaults/timeouts.go
• pkg/k8s/agent/job.go
• pkg/k8s/agent/types.go
• pkg/snapshotter/agent.go
• pkg/snapshotter/snapshot.go

[yuanchen8911]

Recombined into #1967 per the sequencing discussion: the projection and the adopter proved inseparable in practice (the reading is inert without the declaration, the declaration's snapshot-qualified path is broken without the reading, and review findings repeatedly crossed the PR boundary). The final content of this branch (head d7b94ce, all review rounds addressed) is contained verbatim in #1967's combined commit. Review history stays here for reference; #1761's rollout plan is updated accordingly.