fix(validators): make GKE NCCL node selector dynamic#1262
Conversation
|
Note Reviews pausedIt 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 Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughTraining overlay manifests received explanatory comments describing the calibrated nccl-all-reduce-bw absolute floor. The validator added a package constant Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related PRs
Suggested labels
Suggested reviewers
🚥 Pre-merge checks | ✅ 3 | ❌ 1❌ Failed checks (1 inconclusive)
✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
mchmarny
left a comment
mchmarny
left a comment
There was a problem hiding this comment.
Clean fix. The discoverGKENodeConfig helper faithfully mirrors discoverEKSNodeConfig — same single-node-read pattern, same ErrCodeInternal for a missing label, same config.Nodes[0] call site (safe: determineGPUConfig guarantees a non-empty slice before returning). platformWorkerScheduling thread-through is correct; the gkeAccelerator zero-value-when-not-GKE pattern matches instanceType for EKS. Tests are table-driven and cover both valid SKUs plus the missing-label path. CI is fully green.
One nit inline (use the gkeAcceleratorLabel constant in the error message). Two mechanics items before merge:
- Missing type label: add
bug(per PR policy — type labels aren't auto-assigned). - Branch is behind
main: rebase + squash before merging.
| "k8s.io/apimachinery/pkg/runtime/schema" | ||
| "k8s.io/client-go/dynamic" | ||
| ) | ||
|
|
There was a problem hiding this comment.
nit: the error message duplicates the label string that's already in the gkeAcceleratorLabel constant one line above. fmt.Sprintf("GPU node missing %s label", gkeAcceleratorLabel) avoids the drift risk if the label name ever changes.
There was a problem hiding this comment.
Good catch — and self-resolved by the rebase. The new commit removes discoverGKENodeConfig entirely; the GKE lookup now happens inline in platformWorkerScheduling and uses the gkeAcceleratorLabel constant directly (matches the OKE pattern added in #1233). No more duplication to drift on.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
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 `@validators/performance/nccl_all_reduce_bw_constraint.go`:
- Around line 845-855: Ensure GKE accelerator labels are homogeneous before
deriving a nodeSelector: instead of using only nodes[0] in the
recipe.CriteriaServiceGKE branch, iterate over the target node slice
(nodes/config.Nodes and any place WorkerCount is sized) and verify every node
has a non-empty labels[gkeAcceleratorLabel] and that all values are identical;
if any node is missing the label or labels differ, fail fast (return an error)
rather than silently selecting the first node’s label and creating a selector
that may not match the full WorkerCount. Update the logic around nodeSelector,
the homogeneity check, and error handling in the same function where
recipe.CriteriaServiceGKE, nodes, WorkerCount, and gkeAcceleratorLabel are
referenced.
🪄 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: 42ce4ad8-7f3c-4816-9766-271f13073655
📒 Files selected for processing (5)
recipes/overlays/h100-eks-training.yamlrecipes/overlays/h100-gke-cos-training.yamlvalidators/performance/nccl_all_reduce_bw_constraint.govalidators/performance/nccl_gke_utils.govalidators/performance/nccl_test.go
yuanchen8911
left a comment
yuanchen8911
left a comment
There was a problem hiding this comment.
Requesting changes on two items (details in the inline comment and review above):
-
Mixed/missing GKE accelerator labels fail open in
platformWorkerScheduling— whencommonGKEAccelerator(nodes)returns""on a non-empty GKE target set, no nodeSelector is emitted whileWorkerCount/GPUCountPerNodewere sized fromconfig.Nodes, so worker placement can diverge from the counted cohort and the bandwidth result is measured against a mismatched set instead of erroring. Please fail closed (error on a mixed/missing set, mirroringnarrowByAccelerator's zero-match error) or narrow the counted node set to the selected cohort. -
The PR description is stale — it references
discoverGKENodeConfig, a newgkeAcceleratorparameter, andTestDiscoverGKENodeConfigthat don't exist; the code addscommonGKEAcceleratorwith an unchangedplatformWorkerSchedulingsignature. Please update the description before merge.
Open question: do heterogeneous H100 GKE pools (a3-megagpu-8g + a3-highgpu-1g) occur in supported deployments? If pools are always single-SKU, item 1 is latent and can be deferred to #1256; if mixed pools are realistic, the fail-closed guard belongs in this PR.
|
Updated to address review comments: Fixed: Label issue: Branch status: |
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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 `@validators/performance/nccl_all_reduce_bw_constraint.go`:
- Around line 857-860: Replace the fmt.Errorf usage in the error return after
commonGKEAccelerator(nodes) with the structured error constructor aicrErrors.New
using ErrCodeInvalidRequest: change the return that currently uses
fmt.Errorf("GKE nodes have missing or mixed %s labels — cannot derive a
nodeSelector that matches the full WorkerCount cohort", gkeAcceleratorLabel) to
return aicrErrors.New(aicrErrors.ErrCodeInvalidRequest, "<same message with
gkeAcceleratorLabel>") so the error carries the proper code; update the error
return site in the function containing acc := commonGKEAccelerator(nodes) and
ensure callers (notably the caller referenced in the review) will be adjusted to
propagate the structured error directly rather than re-wrapping it.
- Around line 602-605: The call site around platformWorkerScheduling currently
wraps returned errors with aicrErrors.Wrap; once platformWorkerScheduling is
changed to return structured aicrErrors (per the upstream change), remove the
extra wrapping and propagate the error directly (i.e., return the original err
from the function instead of re-wrapping), leaving the
defaultNodeSelector/defaultTolerations assignment and only handling nil-case
locally; update the code that does return aicrErrors.Wrap(...) to simply return
nil, err so structured codes are preserved.
🪄 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: 8b9da964-1d75-4b40-bb61-9a8dafe43735
📒 Files selected for processing (2)
validators/performance/nccl_all_reduce_bw_constraint.govalidators/performance/nccl_test.go
There was a problem hiding this comment.
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 `@validators/performance/nccl_all_reduce_bw_constraint.go`:
- Around line 859-860: Replace the error code used when GKE nodes have
missing/mixed accelerator labels: in the return that currently calls
aicrErrors.New(aicrErrors.ErrCodeInvalidRequest, fmt.Sprintf(...,
gkeAcceleratorLabel)), change ErrCodeInvalidRequest to
aicrErrors.ErrCodeInternal so the error matches other cluster-state validation
failures (see nearby checks like the ones that report "node selector matches
zero nodes" and "no nodes match accelerator").
🪄 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: 898ecd04-9512-4b1b-a3a8-19e9e129592a
📒 Files selected for processing (1)
validators/performance/nccl_all_reduce_bw_constraint.go
| return nil, nil, aicrErrors.New(aicrErrors.ErrCodeInvalidRequest, | ||
| fmt.Sprintf("GKE nodes have missing or mixed %s labels — cannot derive a nodeSelector that matches the full WorkerCount cohort", gkeAcceleratorLabel)) |
There was a problem hiding this comment.
Use ErrCodeInternal instead of ErrCodeInvalidRequest for consistency.
The missing/mixed GKE accelerator label is a cluster-state validation failure (discovered condition) rather than a malformed user request. Other cluster-state errors in this file use ErrCodeInternal: line 405 (node selector matches zero nodes), line 441 (no nodes match accelerator), line 515 (no schedulable GPU nodes), and line 531 (no GPUs found). ErrCodeInvalidRequest is reserved for user-input parsing failures (lines 204, 337, 623). For consistency and correct error classification, use ErrCodeInternal here.
🔧 Proposed fix
acc := commonGKEAccelerator(nodes)
if acc == "" {
- return nil, nil, aicrErrors.New(aicrErrors.ErrCodeInvalidRequest,
+ return nil, nil, aicrErrors.New(aicrErrors.ErrCodeInternal,
fmt.Sprintf("GKE nodes have missing or mixed %s labels — cannot derive a nodeSelector that matches the full WorkerCount cohort", gkeAcceleratorLabel))
}🤖 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 `@validators/performance/nccl_all_reduce_bw_constraint.go` around lines 859 -
860, Replace the error code used when GKE nodes have missing/mixed accelerator
labels: in the return that currently calls
aicrErrors.New(aicrErrors.ErrCodeInvalidRequest, fmt.Sprintf(...,
gkeAcceleratorLabel)), change ErrCodeInvalidRequest to
aicrErrors.ErrCodeInternal so the error matches other cluster-state validation
failures (see nearby checks like the ones that report "node selector matches
zero nodes" and "no nodes match accelerator").
yuanchen8911
left a comment
yuanchen8911
left a comment
There was a problem hiding this comment.
need another rebase
Derive the GKE NCCL all-reduce node selector from observed node accelerator labels instead of a hardcoded value, and fail with a structured error when nodes have missing or mixed accelerator labels.
yuanchen8911
force-pushed
the
fix/nccl-bw-1256
branch
from
c5ed6e8 to
ea136f0
Compare
Recipe evidence checkAffected leaf overlays: 7
How to refresh evidenceRun on a cluster matching the recipe's aicr snapshot -o snapshot.yaml
aicr validate \
-r recipes/overlays/<slug>.yaml \
-s snapshot.yaml \
--emit-attestation ./out \
--push ghcr.io/<your-fork>/aicr-evidence
cp ./out/pointer.yaml recipes/evidence/<slug>.yamlThis gate is warning-only and never blocks merge. See ADR-007 for the trust model. |
3 participants
Summary
Fixes the GKE NCCL all-reduce bandwidth validator to dynamically discover the
cloud.google.com/gke-acceleratorlabel from target nodes instead of hardcodingnvidia-h100-mega-80gb. Now fails explicitly when GKE nodes have missing or mixed accelerator labels to prevent WorkerCount divergence.Motivation / Context
The NCCL validator hardcoded
cloud.google.com/gke-accelerator=nvidia-h100-mega-80gbinplatformWorkerScheduling, so worker pods could only schedule ona3-megagpu-8gnodes. Any GKE H100 SKU with a different accelerator label (e.g.a3-highgpu-1gwithnvidia-h100-80gb) failed at scheduling.Additionally, when GKE nodes had missing or mixed accelerator labels, the code silently fell through with no nodeSelector while
WorkerCountwas sized fromconfig.Nodes. This could cause worker placement to diverge from the counted cohort, leading to invalid bandwidth measurements.Fixes: #1256
Type of Change
Component(s) Affected
pkg/validator)Implementation Notes
commonGKEAccelerator(nodes []v1.Node) stringhelper that returns thecloud.google.com/gke-acceleratorlabel shared by all nodes, or""when nodes disagree or any node lacks the label (mirrorscommonGPUProductfor OKE).platformWorkerScheduling()signature to return(map[string]string, []v1.Toleration, error)instead of just the first two values.commonGKEAccelerator()returns"", preventing silent WorkerCount/scheduling divergence.gkeAcceleratorLabelconstant as requested in review.Testing
Updated
TestPlatformWorkerScheduling:a3-megagpu-8g,a3-highgpu-1g) updated to handle error returnRisk Assessment
(ns, tols)to(ns, tols, err)but this is an internal function. GKE nodes with homogeneous accelerator labels work as before; mixed/missing labels now fail fast instead of silently diverging.Rollout notes: N/A — non-breaking for valid clusters. Clusters with mixed/missing GKE accelerator labels will now fail explicitly with a clear error message instead of silently producing invalid results.
Checklist
go test ./validators/performance/...)commonGPUProduct)