feat(clustered): resume a preempted gang when Kueue readmits it by mitja-kleider · Pull Request #8258 · flyteorg/flyte

mitja-kleider · 2026-10-27T20:24:29Z

Tracking issue

Why are the changes needed?

When Kueue suspends a clustered task's JobSet after it started, the attempt ends: Preempted uses one of the task's retries. On a shared cluster, preemption of low-priority multi-node training is routine, so retries turns into a preemption budget, and an ordinary preemption fails the task once it is spent. Kueue's own model is to requeue the workload and resume the same JobSet.

What changes were proposed in this pull request?

A per-task opt-in, resume_after_preemption in ClusterFailurePolicy:

false (default): unchanged.
true: a JobSet Kueue preempted after it started keeps its attempt and stays in place. The task reports Running with Kueue's eviction reason, e.g. gang evicted by kueue (Preempted, InClusterQueue): <message>; waiting for re-admission, and goes back to running when Kueue readmits it.

Only a preemption waits. Any other eviction (node failures, a PodsReady timeout, a deactivated Workload) still ends the attempt as today, as does a Workload that can't be found or read, so missing RBAC degrades to today's behaviour instead of waiting forever. If the JobSet is suspended before Kueue has written the Evicted condition, the task reports waiting for that poll rather than ending the attempt.

A Workload deactivated while it waits keeps the Evicted reason of the earlier preemption, so a deactivated Workload whose reason isn't a Deactivated* one is classified as a deactivation (system retry) rather than by that stale reason. This also applies with resume_after_preemption off.

The phase stays Running while waiting because phases can't move backwards within an attempt; the reason changes, which bumps the phase version. max_runtime keeps counting while the gang waits. gang.Eviction gains Summary() (the message without the retry clause) for the waiting reason.

How was this patch tested?

Unit tests in kueue_test.go cover which evictions wait and which end the attempt with the setting on, the stale-reason classification with it off, a missing Workload, the phase version bump on eviction, and readmission back to running.

On a cluster (Kueue v0.17.6, JobSet v0.12.0, flyte-binary v2.0.52 with this commit), a 2-node CPU clustered task with retries=0 and resume_after_preemption: true in a ClusterQueue with withinClusterQueue: LowerPriority:

A higher-priority Job preempted it, the JobSet was suspended, and the task stayed on attempt 1 reporting gang evicted by kueue (Preempted, InClusterQueue): ...; waiting for re-admission. When the preemptor finished, Kueue readmitted it and the task ran to success on attempt 1 (phases Queued, Initializing, Executing, Completed).
Deactivating the Workload (spec.active: false) while it waited after a preemption ended the gang with [GangEvicted] gang evicted by kueue (Deactivated); this does not count against the task's retries, and the task restarted after a system failure with a fresh JobSet, which ran to success.

Labels

added

Check all the applicable boxes

I updated the documentation accordingly.
All new and existing tests passed.
All commits are signed-off.

AdilFayyaz · 2026-10-28T20:00:50Z

+package clustered
+
+import (
+	pluginsConfig "github.com/flyteorg/flyte/v2/flyteplugins/go/tasks/config"
+)
+
+//go:generate pflags Config --default-var=defaultConfig
+
+const (
+	// OnEvictionRetry ends the attempt when Kueue takes back a gang that had started.
+	OnEvictionRetry = "retry"
+	// OnEvictionRequeue keeps the attempt and its JobSet while Kueue requeues the gang, so
+	// a preemption costs no retry. Only a deactivated Workload, which Kueue will not
+	// readmit, ends the attempt. The task's max_runtime keeps counting while it waits.
+	OnEvictionRequeue = "requeue"
+)
+
+var (
+	defaultConfig = Config{
+		Kueue: KueueConfig{OnEviction: OnEvictionRetry},
+	}
+
+	configSection = pluginsConfig.MustRegisterSubSection("clustered", &defaultConfig)
+)
+
+// Config is config for the 'clustered-task' plugin.
+type Config struct {
+	Kueue KueueConfig `json:"kueue" pflag:",Kueue integration"`
+}
+
+type KueueConfig struct {
+	OnEviction string `json:"on-eviction" pflag:",What to do when Kueue takes back a gang that had started: retry or requeue"`
+}
+
+func GetConfig() *Config {
+	return configSection.GetConfig().(*Config)
+}
+
+func SetConfig(cfg *Config) error {
+	return configSection.SetConfig(cfg)
+}


we'd like this per task instead of a plugin config. Could you drop the config.go and the generated config_flags.go file. We can read this from the task spec.
ClusterFailurePolicy in flyteidl2/plugins/clustered.proto already holds max_restarts so a bool there fits: resume_after_preemption. False (default) keeps today's behavior, true waits in place and resumes when Kueue re-admits the gang.

Done: resume_after_preemption is now a field on ClusterFailurePolicy, and config.go / config_flags.go are gone.

AdilFayyaz · 2026-10-28T20:11:53Z

+			logger.Warnf(ctx, "failed to read the Kueue Workload of JobSet %s/%s; ending the attempt: %v",
+				jobSet.Namespace, jobSet.Name, err)
+		}
+		if GetConfig().Kueue.OnEviction == OnEvictionRequeue && awaitsReadmission(wl) {


GetTaskPhase already reads the ClusteredTaskSpec and derives the maxRestarts from it so the new field can be read there and passed into afterStart the same way.

Done, it's read next to maxRestarts and passed into afterStart.

AdilFayyaz · 2026-10-28T20:14:28Z

+// awaitsReadmission reports whether Kueue will readmit a gang it took back. A deactivated
+// Workload stays suspended until someone reactivates it, and without a Workload there is
+// no telling, so neither is worth waiting for.
+func awaitsReadmission(wl *unstructured.Unstructured) bool {


Can we wait only on Preempted? Node failures and PodsReady timeouts aren't preemptions and a gang that keeps failing to get ready shouldnt wait forever. One thing to keep. if a jobset is suspended but the Evicted condition isn't true yet, report waiting for this poll rather than ending the attempt, so we dont race Kueue writing the condition.

Done, only Preempted waits now; node failures, PodsReady timeouts and deactivations end the attempt. Kept the guard: a suspended JobSet without Evicted=True yet reports waiting for that poll.

The same race exists with the setting off, where it predates this PR: if the JobSet informer sees the suspend before the Workload cache has Evicted=True/Preempted, the attempt ends unclassified, so the preemption becomes a system retry instead of a user retry. Applying the one-poll wait regardless of the setting would fix it. Do you prefer that added here, or as a follow-up?

lets keep that a follow up PR

AdilFayyaz · 2026-10-28T20:18:27Z

+	// Deactivating a Workload that is already evicted, e.g. one waiting for re-admission
+	// after a preemption, can leave the earlier reason in place: the deactivation is the
+	// cause now.
+	if isInactive(wl) && !(evictedNow && strings.HasPrefix(evicted.reason, evictedByDeactivationPrefix)) {


nit: this also applies with resume_after_preemption off: a deactivated Workload with a stale Preempted reason becomes a system retry instead of a user retry. I agree with the change. Could you add a test for it with the setting off and mention it in the description?

Added a test with the setting off and mentioned it in the description.

AdilFayyaz · 2026-10-28T20:18:58Z

-// had already fully started. See phase.go. Build-time wiring (queue label and
-// spec.suspend on creation) is gated by plugin config.
+// gate releases it. If the gate revokes a gang that had already fully started, the
+// attempt ends with GangEvicted, or, with kueue.on-eviction set to requeue, waits for


update comment, we're no longer using the config.

Under Kueue, preemption is routine for low-priority multi-node work, and ending the attempt on each one makes a task's retries a preemption budget. With the task's failure_policy.resume_after_preemption set, a gang Kueue preempts after it started keeps its attempt and JobSet until Kueue readmits it. Any other eviction still ends the attempt. Signed-off-by: Mitja Kleider <mitja.kleider@aleph-alpha.com>

…#1728) ## Motivation flyteorg/flyte#8258 lets a clustered task wait out a Kueue preemption in its current attempt. When Kueue readmits the gang, the pods are recreated with the same `JOBSET_RESTART_ATTEMPT`. `TaskContext.checkpoint` only reads the attempt's own checkpoint when that counter is above 0, so a readmitted gang loads the previous attempt's checkpoint (or none) and loses the progress it saved before the preemption. ## Summary - `Checkpoint.load` / `load_sync` try the attempt's own `checkpoint_path` first and fall back to `prev_checkpoint_path`. They no longer return early when there is no previous checkpoint. - `TaskContext.checkpoint` drops the `JOBSET_RESTART_ATTEMPT > 0` special case. Each attempt has its own checkpoint prefix, so the own path is empty on a fresh attempt. It only holds data when the same attempt runs again (JobSet restart, Kueue readmission), and that data is the newest state. - A JobSet restart that crashed before its first save now falls back to the previous attempt's checkpoint instead of starting from scratch. - Logs which checkpoint was loaded (this attempt's or the previous attempt's). ## Test Plan - `pytest tests/flyte/test_checkpoint.py tests/flyte/clustered`: a table of cases run sync and async (readmission on first and later attempts, JobSet restart, fresh attempt, restart before first save, nothing saved), plus a check that `remote_source` is still the previous attempt's path. `make unit_test` passes. - Devbox end-to-end with a flyte-binary built from flyteorg/flyte#8258, JobSet v0.13.0 and Kueue 0.19.8: - Setup: a ClusterQueue with `withinClusterQueue: LowerPriority`, and a 2-worker low-priority clustered task with `retries=0` and `resume_after_preemption: true`. Rank 0 saves `{"step": n}` every step. - Preempt it with a high-priority Job and let Kueue readmit it. - With this branch: same attempt, `restart_attempt=0`, resumed at step 6 (the last saved step), completed on attempt 1. - Control on `main`: same attempt, resumed at step 0. Related: flyteorg/flyte#8258, flyteorg/flyte#8223 Signed-off-by: M. Adil Fayyaz <62440954+AdilFayyaz@users.noreply.github.com>

github-actions Bot added the flyte2 label Oct 27, 2026

AdilFayyaz reviewed Oct 28, 2026

View reviewed changes

mitja-kleider force-pushed the clustered-kueue-requeue branch from 9ce13b7 to 5740589 Compare October 29, 2026 13:49

mitja-kleider changed the title ~~feat(clustered): option to wait out a Kueue eviction~~ feat(clustered): resume a preempted gang when Kueue readmits it Oct 29, 2026

AdilFayyaz mentioned this pull request Oct 29, 2026

fix: load this attempt's own checkpoint before the previous attempt's flyteorg/flyte-sdk#1728

Merged

AdilFayyaz approved these changes Oct 30, 2026

View reviewed changes

AdilFayyaz merged commit ccf67eb into flyteorg:main Oct 30, 2026
24 checks passed

Navigation Menu

Repository navigation

feat(clustered): resume a preempted gang when Kueue readmits it - #8258

feat(clustered): resume a preempted gang when Kueue readmits it#8258
AdilFayyaz merged 1 commit into
flyteorg:mainfrom
mitja-kleider:clustered-kueue-requeue

mitja-kleider commented Oct 27, 2026 •

edited

Loading

Uh oh!

AdilFayyaz Oct 28, 2026

Uh oh!

mitja-kleider Oct 29, 2026

Uh oh!

AdilFayyaz Oct 28, 2026

Uh oh!

mitja-kleider Oct 29, 2026

Uh oh!

AdilFayyaz Oct 28, 2026

Uh oh!

mitja-kleider Oct 29, 2026

Uh oh!

AdilFayyaz Oct 29, 2026

Uh oh!

AdilFayyaz Oct 28, 2026

Uh oh!

mitja-kleider Oct 29, 2026

Uh oh!

AdilFayyaz Oct 28, 2026

Uh oh!

mitja-kleider Oct 29, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

mitja-kleider commented Oct 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Tracking issue

Why are the changes needed?

What changes were proposed in this pull request?

How was this patch tested?

Labels

Check all the applicable boxes

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

mitja-kleider commented Oct 27, 2026 •

edited

Loading