Skip to content

fix: propagate PodCliqueSet annotations to PodGang #500

Closed
w1t wants to merge 2 commits into
ai-dynamo:mainfrom
w1t:main
Closed

fix: propagate PodCliqueSet annotations to PodGang #500
w1t wants to merge 2 commits into
ai-dynamo:mainfrom
w1t:main

Conversation

@w1t

@w1t w1t commented Mar 28, 2026

Copy link
Copy Markdown
Contributor

What type of PR is this?

/kind bug

What this PR does / why we need it:

Problem

When creating a disaggregated instance via Dynamo with a custom
nvidia.com/kai-scheduler-queue annotation, the resulting PodGang
is created without that annotation.

The kai-scheduler reads queue assignment from the PodGang, so it
raises QueueDoesNotExist even though the queue exists and was
validated successfully by the Dynamo operator.

Root cause

buildResource in podgang.go constructs PodGang metadata but does
not copy pcs.Annotations into pg.Annotations.

This differs from the existing behavior in pod.go, where annotations
from the parent resource are propagated to each Pod.

Fix

Merge pcs.Annotations into pg.Annotations using lo.Assign, with
computed annotations taking precedence over user-provided values when
the same key is present.

topologyAnnotations := map[string]string{}
if r.tasConfig.Enabled {
    topologyAnnotations[apicommonconstants.AnnotationTopologyName] =
        grovecorev1alpha1.DefaultClusterTopologyName
}
pg.Annotations = lo.Assign(pcs.Annotations, topologyAnnotations)

Test

Added TestPodGangAnnotationsFromPodCliqueSet in podgang_test.go.

The test calls buildResource with a PodCliqueSet carrying scheduler
annotations and asserts that the resulting PodGang contains the same
annotations. Covers three cases:

  • single annotation (nvidia.com/kai-scheduler-queue)

  • multiple scheduler annotations

  • mixed custom and scheduler annotations

    Which issue(s) this PR fixes:

Fixes #490

Special notes for your reviewer:

The merge order in lo.Assign(pcs.Annotations, topologyAnnotations)
is intentional — topology annotations derived during reconciliation
overwrite any user-provided annotation with the same key.

Does this PR introduce an API change?

NONE

buildResource in podgang.go was constructing PodGang metadata without
copying pcs.Annotations. As a result, scheduler annotations such as
nvidia.com/kai-scheduler-queue set on a PodCliqueSet were not propagated
to the PodGang.

The kai-scheduler reads queue assignment from the PodGang, so the missing
annotation caused QueueDoesNotExist even when the queue itself was valid.

Fix this by merging pcs.Annotations into pg.Annotations with lo.Assign,
while keeping computed annotations such as topology authoritative when
keys overlap. This matches the existing annotation propagation pattern
used for Pods.
@copy-pr-bot

copy-pr-bot Bot commented Mar 28, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@Ronkahn21 Ronkahn21 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for raising the issue and the fix! The root cause analysis is spot on. One concern: the current lo.Assign call replaces all existing PodGang annotations (and labels) instead of merging — see inline comments.

Comment thread operator/internal/controller/podcliqueset/components/podgang/podgang.go Outdated

@Ronkahn21 Ronkahn21 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for raising the issue and the fix! The root cause analysis is spot on. One concern: the current lo.Assign call replaces all existing PodGang annotations instead of merging — see inline comments.

  Preserve existing PodGang annotations when buildResource merges
  PodCliqueSet annotations, and expand buildResource test coverage to
  cover both create and update paths, including TAS-disabled
  preservation and TAS-enabled topology override behavior.

  Co-authored-by: Ron Kahn <122778260+Ronkahn21@users.noreply.github.com>
  Signed-off-by: w1t <1955451+w1t@users.noreply.github.com>
@w1t

w1t commented Apr 4, 2026

Copy link
Copy Markdown
Contributor Author

Thanks for raising the issue and the fix! The root cause analysis is spot on. One concern: the current lo.Assign call replaces all existing PodGang annotations instead of merging — see inline comments.

Good catch, thanks. I’ve updated the fix and addressed the inline comments.

@Ronkahn21 Ronkahn21 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@Ronkahn21

Copy link
Copy Markdown
Contributor

@w1t Please rebase your branch (fixing the conflicts),
thanks

@gflarity gflarity left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good. Thanks for the contribution. Please see the minor comments, and you'll need to rebase.

Afterwards we can run the full test suite, re-approve and merge :)

}
pg.Annotations[apicommonconstants.AnnotationTopologyName] = grovecorev1alpha1.DefaultClusterTopologyName
topologyAnnotations[apicommonconstants.AnnotationTopologyName] = grovecorev1alpha1.DefaultClusterTopologyName
}

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There's missing space after the first comma. goimports (part of CI) will catch this.

Suggested change
}
pg.Annotations = lo.Assign(pg.Annotations, pcs.Annotations, topologyAnnotations)

require.NoError(t, err)

for expectedKey, expectedValue := range test.expectedLabels {
actualValue, exists := pg.Labels[expectedKey]

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These assertions verify that every expected key is present with the right value, but they don't verify there are no unexpected extra entries. A stray annotation slipping through would go undetected.

Consider using a direct map comparison instead:

assert.Equal(t, test.expectedLabels, pg.Labels)
assert.Equal(t, test.expectedAnnotations, pg.Annotations)

Note: there will be some rebase issues you'll also need to account for.

@danbar2

danbar2 commented May 3, 2026

Copy link
Copy Markdown
Contributor

stale PR, will be resolved in #573

@danbar2 danbar2 closed this May 3, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[BUG] Can not find queue from kai-scheduler when creating an instance via dynamo

4 participants