Repository navigation
fix(tas): preserve leader assignment during elastic workload scale-up - #13574
Conversation
kubernetes-prow
Bot
added
release-note
✅ Deploy Preview for kubernetes-sigs-kueue canceled.
|
kubernetes-prow
Bot
added
cncf-cla: yes
|
Hi @RooobinYe. Thanks for your PR. I'm waiting for a kubernetes-sigs member to verify that this patch is reasonable to test. If it is, they should reply with Tip We noticed you've done this a few times! Consider joining the org to skip this step and gain Once the patch is verified, the new status will be reflected by the I understand the commands that are listed here. DetailsInstructions for interacting with me using PR comments are available here. If you have questions or suggestions related to my behavior, please file an issue against the kubernetes-sigs/prow repository. |
kubernetes-prow
Bot
added
the
size/L
📝 WalkthroughWalkthroughElastic TAS scale-up now preserves valid leader assignments, detects stale leader assignments, and falls back to fresh placement when needed. Unit and integration tests cover leader stability, worker domain expansion, and current-capacity placement. ChangesElastic TAS scale-up
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant WorkloadSlice
participant TAScheduler
participant PodSetAssignments
WorkloadSlice->>TAScheduler: replace slice with increased worker count
TAScheduler->>PodSetAssignments: inspect leader previousAssignment
TAScheduler->>PodSetAssignments: place worker delta
PodSetAssignments-->>TAScheduler: combined leader and worker assignments
TAScheduler-->>WorkloadSlice: admit replacement slice
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Pull request overview
Fixes a TAS elastic scale-up bug where a leader PodSet’s TopologyAssignment could be recomputed (and overwritten) during worker scale-up, potentially forcing an already-running leader/head pod to restart on a different node and lose state. The change keeps the leader’s prior placement stable while still placing only the delta workers, and adds regression coverage.
Changes:
- Preserve the leader PodSet’s previous
TopologyAssignmentduring elastic scale-up and account for the leader’s resource usage before placing delta workers. - Detect stale leader
PreviousAssignmentdomains early and fall back to full “fresh placement” for the whole group when the leader’s prior domain no longer exists. - Add unit + integration tests covering preserved-leader scale-up and stale-leader fallback behavior.
Note: If you want me to consult AGENTS.md, please acknowledge that its usage is experimental and comes with no guarantees of backwards/future compatibility.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated no comments.
| File | Description |
|---|---|
pkg/cache/scheduler/tas_elastic_workloads.go |
Preserves leader assignment during elastic scale-up; adds stale-leader detection to force fresh placement when needed. |
pkg/cache/scheduler/tas_cache_test.go |
Adds unit tests validating leader preservation on scale-up and fallback behavior when the leader’s previous domain is stale. |
test/integration/singlecluster/tas/tas_test.go |
Adds an integration scenario ensuring leader assignment is unchanged across worker scale-up via workload slice replacement. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
test/integration/singlecluster/tas/tas_test.go (1)
6939-7070: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsider adding an integration-level negative scenario (stale leader domain).
This new spec only covers the positive path (leader assignment preserved during scale-up). The stale-leader fallback (previous leader domain referencing a deleted node) is covered at the unit level (
tas_cache_test.go) but has no integration-level counterpart here, even though this describe block is otherwise dedicated to elastic scale-up leader/worker behavior. Consider adding a companion spec that deletes the leader's backing node before creating the replacement slice and asserts fresh placement for the whole group, reusing the existing node-deletion pattern already used elsewhere in this file (e.g. "should evict workload when multiple assigned nodes are deleted").As per path instructions,
test/integration/**/*.go: "Verify tests use envtest and follow the existing table-driven test patterns with Ginkgo/Gomega. Ensure new test cases cover both positive and negative scenarios."🤖 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 `@test/integration/singlecluster/tas/tas_test.go` around lines 6939 - 7070, The elastic scale-up integration coverage lacks a stale-leader-domain negative scenario. Add a companion Ginkgo spec near “should preserve the leader assignment when scaling up workers in a pod set group” that deletes the leader’s backing node using the existing node-deletion pattern, creates the replacement slice, and asserts the entire group receives fresh placement rather than reusing the stale leader domain; follow the file’s envtest Ginkgo/Gomega and table-driven conventions where applicable.Source: Path instructions
🤖 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.
Nitpick comments:
In `@test/integration/singlecluster/tas/tas_test.go`:
- Around line 6939-7070: The elastic scale-up integration coverage lacks a
stale-leader-domain negative scenario. Add a companion Ginkgo spec near “should
preserve the leader assignment when scaling up workers in a pod set group” that
deletes the leader’s backing node using the existing node-deletion pattern,
creates the replacement slice, and asserts the entire group receives fresh
placement rather than reusing the stale leader domain; follow the file’s envtest
Ginkgo/Gomega and table-driven conventions where applicable.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: ab963d54-c775-4d2a-8103-ab031123e1cc
📒 Files selected for processing (3)
pkg/cache/scheduler/tas_cache_test.gopkg/cache/scheduler/tas_elastic_workloads.gotest/integration/singlecluster/tas/tas_test.go
|
/ok-to-test |
kubernetes-prow
Bot
added
ok-to-test
| for _, psa := range wl1.Status.Admission.PodSetAssignments { | ||
| switch psa.Name { | ||
| case "leader": | ||
| originalLeaderAssignment = utiltas.InternalFrom(psa.TopologyAssignment) | ||
| case "workers": | ||
| originalWorkersAssignment = utiltas.InternalFrom(psa.TopologyAssignment) | ||
| } | ||
| } |
There was a problem hiding this comment.
Let's berak it down using some helpers like podSetAssignmentByName. I suppose we already should have such helpers
There was a problem hiding this comment.
I added a test-local podSetAssignmentByName and used it through topologyAssignmentByName to simplify the assertions
| for _, psa := range wl2.Status.Admission.PodSetAssignments { | ||
| g.Expect(psa.TopologyAssignment).ShouldNot(gomega.BeNil()) | ||
| newAssignment := utiltas.InternalFrom(psa.TopologyAssignment) | ||
| switch psa.Name { | ||
| case "leader": | ||
| g.Expect(newAssignment).Should(gomega.BeComparableTo(originalLeaderAssignment)) | ||
| case "workers": | ||
| g.Expect(utiltas.CountPodsInAssignment(newAssignment)).Should(gomega.Equal(int32(4))) | ||
| newCounts := make(map[utiltas.TopologyDomainID]int32, len(newAssignment.Domains)) | ||
| for _, domain := range newAssignment.Domains { | ||
| newCounts[utiltas.DomainID(domain.Values)] += domain.Count | ||
| } | ||
| for _, domain := range originalWorkersAssignment.Domains { | ||
| g.Expect(newCounts[utiltas.DomainID(domain.Values)]).Should( | ||
| gomega.BeNumerically(">=", domain.Count), | ||
| "previous workers domain %v must be preserved", domain.Values) | ||
| } | ||
| } |
There was a problem hiding this comment.
Break this down using helper functions.
There was a problem hiding this comment.
Done, I used podSetAssignmentByName.
| features.ElasticJobsViaWorkloadSlicesWithTAS: true, | ||
| }, | ||
| }, | ||
| "elastic workload scale up with leader: stale leader assignment falls back to fresh placement": { |
There was a problem hiding this comment.
I think this test has one node with 8 CPU and pods requesting 1 CPU each, so everything lands on x1 trivially. I think it proves the fallback path works but it doesn't really exercise the group replacement logic since there's only one place to go!?
There was a problem hiding this comment.
Now the test has x1 node with 2 CPU, x2 node with 3 CPU and x3 node has 3 CPU. Workers are expected has x1 *1 + x2 * 3
During elastic workload scale-up, handleScaleUp recomputed the leader's placement and unconditionally overwrote its TopologyAssignment, which could relocate (and restart) the running leader pod. Preserve the leader's previous assignment instead and account for its usage before placing the delta workers. If the leader's previous assignment points to a domain that no longer exists, fall back to fresh placement for the whole group. Covered by a unit test at the snapshot layer and an integration test that exercises the full flavorassigner path with a leader+workers pod set group.
1caef59 to
90b0235
Compare
There was a problem hiding this comment.
🧹 Nitpick comments (1)
test/integration/singlecluster/tas/tas_test.go (1)
6971-7081: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winWell-constructed positive-path test; consider adding a negative (stale-leader) integration scenario for full path-instruction coverage.
The scale-up test itself correctly forces and verifies the exact regression scenario (leader's node fully packed, forcing delta workers elsewhere) and asserts leader-assignment stability plus worker-domain preservation. However, only the positive path is covered here — the stale-leader-assignment fallback (already unit-tested in
pkg/cache/scheduler/tas_cache_test.go) has no integration-level counterpart in this file.As per path instructions, "Ensure new test cases cover both positive and negative scenarios." Consider adding a companion
Itblock (e.g., leader previousAssignment pointing at a deleted node) to confirm the fallback wiring end-to-end, though the unit test already validates the core logic precisely.🤖 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 `@test/integration/singlecluster/tas/tas_test.go` around lines 6971 - 7081, The integration coverage only exercises valid leader-assignment preservation and lacks the stale-leader fallback path. Add a companion `ginkgo.It` scenario near the existing scale-up test that makes the leader’s previous assignment reference a deleted or unavailable node, then verifies replacement admission succeeds and the leader receives a valid fallback assignment while worker placement remains correct. Reuse the existing workload, topology-assignment, and domain-count helpers and preserve the current positive-path test.Source: Path instructions
🤖 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.
Nitpick comments:
In `@test/integration/singlecluster/tas/tas_test.go`:
- Around line 6971-7081: The integration coverage only exercises valid
leader-assignment preservation and lacks the stale-leader fallback path. Add a
companion `ginkgo.It` scenario near the existing scale-up test that makes the
leader’s previous assignment reference a deleted or unavailable node, then
verifies replacement admission succeeds and the leader receives a valid fallback
assignment while worker placement remains correct. Reuse the existing workload,
topology-assignment, and domain-count helpers and preserve the current
positive-path test.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 40a95891-25f7-4db5-b0ae-8d74fec7da84
📒 Files selected for processing (3)
pkg/cache/scheduler/tas_cache_test.gopkg/cache/scheduler/tas_elastic_workloads.gotest/integration/singlecluster/tas/tas_test.go
|
/area elastic-jobs |
kubernetes-prow
Bot
added
the
area/elastic-jobs
|
/approve |
|
@mimowo: once the present PR merges, I will cherry-pick it on top of DetailsIn response to this:
Instructions for interacting with me using PR comments are available here. If you have questions or suggestions related to my behavior, please file an issue against the kubernetes-sigs/prow repository. |
kubernetes-prow
Bot
added
the
approved
kubernetes-prow
Bot
added
the
lgtm
|
LGTM label has been added. DetailsGit tree hash: 557a91f9ac932f7d79d432b65702d755e862daf7 |
|
[APPROVALNOTIFIER] This PR is APPROVED This pull-request has been approved by: mimowo, RooobinYe, sohankunkerkar The full list of commands accepted by this bot can be found here. The pull request process is described here DetailsNeeds approval from an approver in each of these files:
Approvers can indicate their approval by writing |
|
@mimowo: new pull request created: #13813 DetailsIn response to this:
Instructions for interacting with me using PR comments are available here. If you have questions or suggestions related to my behavior, please file an issue against the kubernetes-sigs/prow repository. |
|
@mimowo: #13574 failed to apply on top of branch "release-0.18": DetailsIn response to this:
Instructions for interacting with me using PR comments are available here. If you have questions or suggestions related to my behavior, please file an issue against the kubernetes-sigs/prow repository. |
|
@RooobinYe could you please prepare the CP for 0.18? |
What type of PR is this?
/kind bug
/area tas
What this PR does / why we need it:
When an elastic workload (
ElasticJobsViaWorkloadSlices+ElasticJobsViaWorkloadSlicesWithTAS) with a leader/workers pod set group (e.g. a RayCluster head + workers) scales up,handleScaleUppassed the leader tofindTopologyAssignment, which ignores the leader'sPreviousAssignmentand computes a brand-new placement for it. The result unconditionally overwrote the leader'sTopologyAssignmentin the new workload slice. Since the assignment translates to NodeAffinity, the running leader pod could be terminated and restarted on another node, destroying cluster state and aborting running jobs. In practice the recomputed leader placement could even come out empty, tripping the internalunexpected remainingCountassertion in the snapshot code.This PR fixes
handleScaleUpto:TopologyAssignmentinstead of recomputing it, and account for the leader's usage inassumedUsagebefore placing the delta workers, so the delta placement does not oversubscribe the leader's domain;findTopologyAssignment(its only effect for the leader was computing the new placement, i.e. the bug itself; elastic slices with TAS only supportUnconstrained, so no co-location constraint is lost);handleElasticWorkloadso it covers the scale-up, scale-down, and same-count paths alike.Test coverage:
TestFindTopologyAssignments: scale-up with a leader preserving the leader assignment, and the stale-leader fallback. Both fail without the fix.getPreviousTopologyAssignment→handleScaleUp) with a leader+workers pod set group. The scenario packs the group onto a fully-used node so that a recomputed leader placement could never coincide with the original one; the test fails onmainwith a leader-assignment diff and passes with the fix.Which issue(s) this PR fixes:
Fixes #12528
Special notes for your reviewer:
FindTopologyAssignmentsForWorkloadis invoked withoutSimulateWorkloadRemovalfor the replaced slice whilehandleScaleUpre-adds the previous pods' usage toassumedUsage). That is a pre-existing problem independent of the leader reassignment fixed here: the preemption path accounts correctly, so a legitimate scale-up is still admitted through slice replacement. This PR didn't fix this problem.kueue.x-k8s.io/elastic-job: "true"annotation on the workloads, so it exercises fresh placement rather than the elastic delta path (and its assertions only check that the assignment is non-empty). The new test sets the annotation, matching what the jobframework reconciler propagates from the job. Fixing the existing test is left out of this PR as well.Does this PR introduce a user-facing change?
Summary by CodeRabbit
Summary by CodeRabbit
Bug Fixes
Tests