This is an automated email from the ASF dual-hosted git repository.
dongjoon-hyun pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/spark-kubernetes-operator.git
The following commit(s) were added to refs/heads/main by this push:
new cb873488 [SPARK-60042] Support Kueue admission check `podSetUpdates`
cb873488 is described below
commit cb87348882a28506a62f2830dce19de2041900c2
Author: Dongjoon Hyun <[email protected]>
AuthorDate: Wed Oct 7 11:49:39 2026 -0700
[SPARK-60042] Support Kueue admission check `podSetUpdates`
### What changes were proposed in this pull request?
This PR aims to support Kueue admission check `podSetUpdates` like the
Kueue built-in integrations. The `labels`, `annotations`, `nodeSelector` and
`tolerations` which the admission checks report in
`status.admissionChecks[].podSetUpdates` of a `Workload` are merged with the
flavors and applied to the pods of `SparkApplication` and `SparkCluster`. A
conflicting value fails the resource with `SchedulingFailure`, like a node
label conflict of the flavors.
### Why are the changes needed?
Previously, the operator applied only the flavors. So, the pods admitted
through a
[ProvisioningRequest](https://kueue.sigs.k8s.io/docs/concepts/admission_check/provisioning_request/)
admission check lacked its annotations and node selector, and were not bound
to the provisioned capacity.
### Does this PR introduce _any_ user-facing change?
No for the released versions, because the Kueue integration is not released
yet in
[1.0.0](https://github.com/apache/spark-kubernetes-operator/releases/tag/1.0.0)
(2026-07-23).
### How was this patch tested?
Pass the CIs with the newly added test cases.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Opus 5.5
Closes #951 from dongjoon-hyun/SPARK-60042.
Authored-by: Dongjoon Hyun <[email protected]>
Signed-off-by: Dongjoon Hyun <[email protected]>
---
docs/spark_custom_resources.md | 16 ++
.../spark/k8s/operator/context/BaseContext.java | 26 +-
.../k8s/operator/context/SparkAppContext.java | 6 +-
.../k8s/operator/context/SparkClusterContext.java | 26 +-
...KueuePodSetFlavor.java => KueuePodSetInfo.java} | 56 +++-
.../k8s/operator/kueue/KueueWorkloadUtils.java | 218 ++++++++++-----
.../kueue/v1beta2/AdmissionCheckState.java | 45 ++++
.../k8s/operator/kueue/v1beta2/PodSetUpdate.java | 54 ++++
.../k8s/operator/kueue/v1beta2/WorkloadStatus.java | 3 +
.../reconciler/SparkAppResourceSpecFactory.java | 43 ++-
.../reconciler/reconcilesteps/AppInitStep.java | 2 +-
.../reconciler/reconcilesteps/ClusterInitStep.java | 2 +-
.../k8s/operator/context/SparkAppContextTest.java | 8 +-
.../operator/context/SparkClusterContextTest.java | 34 ++-
...SetFlavorTest.java => KueuePodSetInfoTest.java} | 68 ++++-
.../k8s/operator/kueue/KueueWorkloadUtilsTest.java | 299 ++++++++++++++++++---
.../k8s/operator/kueue/v1beta2/WorkloadTest.java | 45 ++++
.../SparkAppResourceSpecFactoryTest.java | 27 +-
.../reconciler/reconcilesteps/AppInitStepTest.java | 20 +-
.../reconcilesteps/ClusterInitStepTest.java | 12 +-
tests/e2e/kueue/chainsaw-test.yaml | 224 +++++++++++++++
.../spark-cluster-example-admission-check.yaml | 37 +++
tests/e2e/kueue/spark-example-admission-check.yaml | 35 +++
23 files changed, 1095 insertions(+), 211 deletions(-)
diff --git a/docs/spark_custom_resources.md b/docs/spark_custom_resources.md
index e5c19516..5ffcb9ac 100644
--- a/docs/spark_custom_resources.md
+++ b/docs/spark_custom_resources.md
@@ -716,6 +716,22 @@ spec:
cluster-scoped, so reading them needs the rules which
`operatorRbac.kueue.enabled` grants
through the ClusterRole, hence `operatorRbac.clusterRole.create` as well.
Until they can be
read, the resource is held and the read is retried.
+* Like Kueue built-in integrations, the `labels`, `annotations`,
`nodeSelector` and `tolerations`
+ which the [admission
checks](https://kueue.sigs.k8s.io/docs/concepts/admission_check/) of the
+ `Workload` report for a pod set in `status.admissionChecks[].podSetUpdates`
are added to its pods
+ as well, once all the checks are `Ready` and Kueue admits the `Workload`.
The executor pods get
+ them through the executor pod template, like the flavors. So a resource can
use a
+
[ProvisioningRequest](https://kueue.sigs.k8s.io/docs/concepts/admission_check/provisioning_request/)
+ admission check, whose `autoscaling.x-k8s.io/consume-provisioning-request`
and
+ `autoscaling.x-k8s.io/provisioning-class-name` annotations, and the node
selector which its
+ `ProvisioningRequestConfig` sets, bind the pods to the capacity provisioned
for them. The
+ operator reads them from the `Workload`, so they need no further RBAC rules.
Like Kueue, an
+ update must not change a value which the flavors or an earlier admission
check set, nor a value
+ of the node selector, labels or annotations of the pods, and such a conflict
fails the resource
+ like a node label conflict above. The labels and annotations which Spark
sets on the driver and
+ executor pods itself, e.g. `spark-role` or the ones of
+ `spark.kubernetes.{driver,executor}.{label,annotation}.*`, are not checked,
and Spark's value
+ wins over an update of the same key.
* `spec.suspend` takes precedence. A suspended resource does not get a
`Workload`, and suspending
a queued resource deletes its `Workload` to release the quota. Suspending a
running
`SparkApplication` deletes its `Workload` only after its driver and executor
pods are gone, or
diff --git
a/spark-operator/src/main/java/org/apache/spark/k8s/operator/context/BaseContext.java
b/spark-operator/src/main/java/org/apache/spark/k8s/operator/context/BaseContext.java
index 312c7424..608a916e 100644
---
a/spark-operator/src/main/java/org/apache/spark/k8s/operator/context/BaseContext.java
+++
b/spark-operator/src/main/java/org/apache/spark/k8s/operator/context/BaseContext.java
@@ -32,7 +32,7 @@ import
io.javaoperatorsdk.operator.api.event.ResourceEventRecorder;
import io.javaoperatorsdk.operator.api.reconciler.Context;
import org.apache.spark.k8s.operator.BaseResource;
-import org.apache.spark.k8s.operator.kueue.KueuePodSetFlavor;
+import org.apache.spark.k8s.operator.kueue.KueuePodSetInfo;
import org.apache.spark.k8s.operator.kueue.v1beta2.Workload;
import org.apache.spark.k8s.operator.utils.PodUtils;
@@ -46,8 +46,8 @@ public abstract class BaseContext<CR extends BaseResource<?,
?, ?, ?, ?>> {
/** The JOSDK context of the reconciliation this context belongs to. */
protected final Context<?> josdkContext;
- /** The flavors which Kueue assigned to the pod sets, applied to the
secondary resources. */
- protected Map<String, KueuePodSetFlavor> kueuePodSetFlavors = Map.of();
+ /** The KueuePodSetInfo by the pod set name, applied to the secondary
resources. */
+ protected Map<String, KueuePodSetInfo> kueuePodSetInfos = Map.of();
/** The ready pods by Spark role, see {@link #countReadyPodsByRole}. */
private Map<String, Long> readyPodsByRole;
@@ -69,24 +69,24 @@ public abstract class BaseContext<CR extends
BaseResource<?, ?, ?, ?, ?>> {
public abstract CR getResource();
/**
- * Sets the flavors which Kueue assigned to the pod sets of the resource, so
that its secondary
- * resources carry their node selectors and tolerations.
+ * Sets the KueuePodSetInfo which Kueue assigned to the pod sets of the
resource, so that its
+ * secondary resources carry their node selectors, tolerations, labels and
annotations.
*
- * @param kueuePodSetFlavors The KueuePodSetFlavor by the pod set name.
+ * @param kueuePodSetInfos The KueuePodSetInfo by the pod set name.
*/
- public void setKueuePodSetFlavors(Map<String, KueuePodSetFlavor>
kueuePodSetFlavors) {
+ public void setKueuePodSetInfos(Map<String, KueuePodSetInfo>
kueuePodSetInfos) {
synchronized (this) {
- this.kueuePodSetFlavors = kueuePodSetFlavors;
- applyKueuePodSetFlavors();
+ this.kueuePodSetInfos = kueuePodSetInfos;
+ applyKueuePodSetInfos();
}
}
/**
- * Applies the flavors to the secondary resource spec which was built before
them, e.g. to find
- * the driver pod or to check whether the master exists. Called while
holding the lock of this
- * context, so that a reader sees either spec with its flavors.
+ * Applies the KueuePodSetInfo to the secondary resource spec which was
built before it, e.g. to
+ * find the driver pod or to check whether the master exists. Called while
holding the lock of
+ * this context, so that a reader sees either spec with it.
*/
- protected abstract void applyKueuePodSetFlavors();
+ protected abstract void applyKueuePodSetInfos();
/**
* Returns the Kubernetes client associated with this context.
diff --git
a/spark-operator/src/main/java/org/apache/spark/k8s/operator/context/SparkAppContext.java
b/spark-operator/src/main/java/org/apache/spark/k8s/operator/context/SparkAppContext.java
index 8030f9c4..cf1491a7 100644
---
a/spark-operator/src/main/java/org/apache/spark/k8s/operator/context/SparkAppContext.java
+++
b/spark-operator/src/main/java/org/apache/spark/k8s/operator/context/SparkAppContext.java
@@ -213,18 +213,18 @@ public class SparkAppContext extends
BaseContext<SparkApplication> {
if (secondaryResourceSpec == null) {
secondaryResourceSpec =
SparkAppResourceSpecFactory.buildResourceSpec(
- sparkApplication, josdkContext.getClient(), submissionWorker,
kueuePodSetFlavors);
+ sparkApplication, josdkContext.getClient(), submissionWorker,
kueuePodSetInfos);
}
return secondaryResourceSpec;
}
}
/**
- * Drops the resource spec built before the flavors, since they reach the
driver and the
+ * Drops the resource spec built before the KueuePodSetInfo, since it
reaches the driver and the
* executors through the pod templates which the submission builds the
resources from.
*/
@Override
- protected void applyKueuePodSetFlavors() {
+ protected void applyKueuePodSetInfos() {
// The caller holds this lock already, which is re-entrant, so that the
field is guarded by
// the same monitor as every other access to it.
synchronized (this) {
diff --git
a/spark-operator/src/main/java/org/apache/spark/k8s/operator/context/SparkClusterContext.java
b/spark-operator/src/main/java/org/apache/spark/k8s/operator/context/SparkClusterContext.java
index 88502da9..0de8c4b7 100644
---
a/spark-operator/src/main/java/org/apache/spark/k8s/operator/context/SparkClusterContext.java
+++
b/spark-operator/src/main/java/org/apache/spark/k8s/operator/context/SparkClusterContext.java
@@ -31,7 +31,7 @@ import io.javaoperatorsdk.operator.api.reconciler.Context;
import org.apache.spark.k8s.operator.SparkCluster;
import org.apache.spark.k8s.operator.SparkClusterResourceSpec;
import org.apache.spark.k8s.operator.SparkClusterSubmissionWorker;
-import org.apache.spark.k8s.operator.kueue.KueuePodSetFlavor;
+import org.apache.spark.k8s.operator.kueue.KueuePodSetInfo;
import org.apache.spark.k8s.operator.kueue.KueueWorkloadFactory;
import
org.apache.spark.k8s.operator.reconciler.SparkClusterResourceSpecFactory;
@@ -67,37 +67,37 @@ public class SparkClusterContext extends
BaseContext<SparkCluster> {
if (secondaryResourceSpec == null) {
secondaryResourceSpec =
SparkClusterResourceSpecFactory.buildResourceSpec(sparkCluster,
submissionWorker);
- applyKueuePodSetFlavors();
+ applyKueuePodSetInfos();
}
return secondaryResourceSpec;
}
}
/**
- * Adds the flavors to the StatefulSets of the resource spec, which is not
built again for them:
- * unlike the driver and the executors, the master and worker pods are
created from the pod
- * templates of these StatefulSets. Applying the same flavors twice is a
no-op, so the spec built
- * with them is left as it is.
+ * Adds the KueuePodSetInfo to the StatefulSets of the resource spec, which
is not built again for
+ * it: unlike the driver and the executors, the master and worker pods are
created from the pod
+ * templates of these StatefulSets. Applying the same KueuePodSetInfo twice
is a no-op, so the
+ * spec built with it is left as it is.
*/
@Override
- protected void applyKueuePodSetFlavors() {
+ protected void applyKueuePodSetInfos() {
// The caller holds this lock already, which is re-entrant, so that the
field is guarded by
// the same monitor as every other access to it.
synchronized (this) {
if (secondaryResourceSpec == null) {
return;
}
- applyKueuePodSetFlavor(
+ applyKueuePodSetInfo(
KueueWorkloadFactory.PODSET_MASTER,
secondaryResourceSpec.getMasterStatefulSet());
- applyKueuePodSetFlavor(
+ applyKueuePodSetInfo(
KueueWorkloadFactory.PODSET_WORKER,
secondaryResourceSpec.getWorkerStatefulSet());
}
}
- private void applyKueuePodSetFlavor(String podSetName, StatefulSet
statefulSet) {
- KueuePodSetFlavor flavor = kueuePodSetFlavors.get(podSetName);
- if (flavor != null) {
- flavor.applyTo(statefulSet.getSpec().getTemplate().getSpec());
+ private void applyKueuePodSetInfo(String podSetName, StatefulSet
statefulSet) {
+ KueuePodSetInfo podSetInfo = kueuePodSetInfos.get(podSetName);
+ if (podSetInfo != null) {
+ podSetInfo.applyTo(statefulSet.getSpec().getTemplate());
}
}
diff --git
a/spark-operator/src/main/java/org/apache/spark/k8s/operator/kueue/KueuePodSetFlavor.java
b/spark-operator/src/main/java/org/apache/spark/k8s/operator/kueue/KueuePodSetInfo.java
similarity index 61%
rename from
spark-operator/src/main/java/org/apache/spark/k8s/operator/kueue/KueuePodSetFlavor.java
rename to
spark-operator/src/main/java/org/apache/spark/k8s/operator/kueue/KueuePodSetInfo.java
index 0019f77c..c2afd2aa 100644
---
a/spark-operator/src/main/java/org/apache/spark/k8s/operator/kueue/KueuePodSetFlavor.java
+++
b/spark-operator/src/main/java/org/apache/spark/k8s/operator/kueue/KueuePodSetInfo.java
@@ -25,27 +25,69 @@ import java.util.List;
import java.util.Map;
import java.util.Objects;
+import io.fabric8.kubernetes.api.model.ObjectMetaBuilder;
import io.fabric8.kubernetes.api.model.PodSpec;
+import io.fabric8.kubernetes.api.model.PodTemplateSpec;
import io.fabric8.kubernetes.api.model.Toleration;
/**
- * The node selector and tolerations of the ResourceFlavors which Kueue
assigned to a pod set.
- * Like Kueue built-in integrations, they are applied to the pods of the pod
set.
+ * The node selector, tolerations, labels and annotations which Kueue assigned
to a pod set, like
+ * Kueue's `podset.PodSetInfo`: the ones of its ResourceFlavors and of the
podSetUpdates of the
+ * admission checks of its Workload. Like Kueue built-in integrations, they
are applied to the pods
+ * of the pod set.
*
- * @param nodeSelector The merged `nodeLabels` of the ResourceFlavors.
- * @param tolerations The merged `tolerations` of the ResourceFlavors.
+ * @param nodeSelector The merged `nodeLabels` of the ResourceFlavors and
`nodeSelector` of the
+ * podSetUpdates.
+ * @param tolerations The merged `tolerations` of the ResourceFlavors and the
podSetUpdates.
+ * @param labels The merged `labels` of the podSetUpdates.
+ * @param annotations The merged `annotations` of the podSetUpdates.
*/
-public record KueuePodSetFlavor(Map<String, String> nodeSelector,
List<Toleration> tolerations) {
+public record KueuePodSetInfo(
+ Map<String, String> nodeSelector,
+ List<Toleration> tolerations,
+ Map<String, String> labels,
+ Map<String, String> annotations) {
private static final String OPERATOR_EQUAL = "Equal";
+ /**
+ * Creates the KueuePodSetInfo of ResourceFlavors only, which have neither
labels nor
+ * annotations.
+ *
+ * @param nodeSelector The merged `nodeLabels` of the ResourceFlavors.
+ * @param tolerations The merged `tolerations` of the ResourceFlavors.
+ */
+ public KueuePodSetInfo(Map<String, String> nodeSelector, List<Toleration>
tolerations) {
+ this(nodeSelector, tolerations, Map.of(), Map.of());
+ }
+
+ /**
+ * Adds the labels and annotations to the metadata of the given pod
template, and the node
+ * selector and tolerations to its pod spec. Missing metadata or a missing
pod spec is created. A
+ * conflict with the pod template must be checked beforehand, as {@link
+ * KueueWorkloadUtils#holdForAdmission} does.
+ *
+ * @param template The pod template to be modified in place.
+ */
+ public void applyTo(final PodTemplateSpec template) {
+ template.setMetadata(
+ new ObjectMetaBuilder(template.getMetadata())
+ .addToLabels(labels)
+ .addToAnnotations(annotations)
+ .build());
+ if (template.getSpec() == null) {
+ template.setSpec(new PodSpec());
+ }
+ applyTo(template.getSpec());
+ }
+
/**
* Adds the node selector and tolerations to the given pod spec. A node
selector conflict must be
- * checked beforehand, see {@link KueueWorkloadUtils#resolvePodSetFlavors}.
+ * checked beforehand, see {@link KueueWorkloadUtils#resolvePodSetInfos}.
*
* @param podSpec The pod spec to be modified in place.
*/
- public void applyTo(final PodSpec podSpec) {
+ void applyTo(final PodSpec podSpec) {
Map<String, String> mergedNodeSelector = new HashMap<>();
if (podSpec.getNodeSelector() != null) {
mergedNodeSelector.putAll(podSpec.getNodeSelector());
diff --git
a/spark-operator/src/main/java/org/apache/spark/k8s/operator/kueue/KueueWorkloadUtils.java
b/spark-operator/src/main/java/org/apache/spark/k8s/operator/kueue/KueueWorkloadUtils.java
index 0e335a93..482fdba1 100644
---
a/spark-operator/src/main/java/org/apache/spark/k8s/operator/kueue/KueueWorkloadUtils.java
+++
b/spark-operator/src/main/java/org/apache/spark/k8s/operator/kueue/KueueWorkloadUtils.java
@@ -48,7 +48,10 @@ import com.fasterxml.jackson.databind.SerializationFeature;
import io.fabric8.kubernetes.api.model.Condition;
import io.fabric8.kubernetes.api.model.ConditionBuilder;
import io.fabric8.kubernetes.api.model.HasMetadata;
+import io.fabric8.kubernetes.api.model.ObjectMeta;
import io.fabric8.kubernetes.api.model.OwnerReference;
+import io.fabric8.kubernetes.api.model.PodSpec;
+import io.fabric8.kubernetes.api.model.PodTemplateSpec;
import io.fabric8.kubernetes.api.model.Toleration;
import io.fabric8.kubernetes.client.KubernetesClient;
import io.fabric8.kubernetes.client.KubernetesClientException;
@@ -57,8 +60,10 @@ import lombok.extern.slf4j.Slf4j;
import org.apache.spark.k8s.operator.Constants;
import org.apache.spark.k8s.operator.client.KubernetesClientFactory;
import org.apache.spark.k8s.operator.context.BaseContext;
+import org.apache.spark.k8s.operator.kueue.v1beta2.AdmissionCheckState;
import org.apache.spark.k8s.operator.kueue.v1beta2.PodSet;
import org.apache.spark.k8s.operator.kueue.v1beta2.PodSetAssignment;
+import org.apache.spark.k8s.operator.kueue.v1beta2.PodSetUpdate;
import org.apache.spark.k8s.operator.kueue.v1beta2.PriorityClassRef;
import org.apache.spark.k8s.operator.kueue.v1beta2.ResourceFlavor;
import org.apache.spark.k8s.operator.kueue.v1beta2.Workload;
@@ -252,7 +257,8 @@ public final class KueueWorkloadUtils {
* queued. An API failure of the request is retried rather than failing the
resource: a transient
* one shortly, a persistent one with the default interval, so that its
event is not rewritten
* every few seconds until a user fixes the cause. Once admitted, the
flavors which Kueue
- * assigned are set to the context, so that the secondary resources are
built with them.
+ * assigned, with the pod set updates of the admission checks, are set to
the context, so that
+ * the secondary resources are built with them.
*
* @param context The context of the resource to be admitted.
* @param desired The Workload built for the resource. The pod sets hash
annotation is added to it
@@ -334,16 +340,16 @@ public final class KueueWorkloadUtils {
requested);
return Optional.of(ReconcileProgress.completeAndDefaultRequeue());
}
- Map<String, KueuePodSetFlavor> flavors;
+ Map<String, KueuePodSetInfo> podSetInfos;
try {
- flavors = resolvePodSetFlavors(context.getClient(),
admission.workload());
- checkNoNodeSelectorConflict(flavors, desired);
+ podSetInfos = resolvePodSetInfos(context.getClient(),
admission.workload());
+ checkNoPodTemplateConflict(podSetInfos, desired);
} catch (KubernetesClientException e) {
return Optional.of(retryAfterFlavorReadFailure(context, e,
workloadName));
} catch (IllegalArgumentException e) {
- // Like Kueue, a node selector conflict is permanent, so the quota is
released before the
- // caller fails the resource, which this step does not reach again. The
release itself is
- // retried, since a Workload left behind holds the quota of a resource
that never starts.
+ // Like Kueue, a conflict is permanent, so the quota is released before
the caller fails the
+ // resource, which this step does not reach again. The release itself is
retried, since a
+ // Workload left behind holds the quota of a resource that never starts.
Optional<ReconcileProgress> retry =
releaseOrRetry(context, "Failed to release the Kueue Workload of a
rejected admission");
if (retry.isPresent()) {
@@ -355,22 +361,23 @@ public final class KueueWorkloadUtils {
context.getEventRecorder(),
EventUtils.REASON_KUEUE_ADMITTED,
"Kueue admitted Workload " + workloadName + ", requesting " +
requested + ".");
- context.setKueuePodSetFlavors(flavors);
+ context.setKueuePodSetInfos(podSetInfos);
return Optional.empty();
}
/**
- * Sets the flavors of the Workload which Kueue admitted before on the
context of a resource
- * whose driver or master exists already. Such a reconcile applies the
secondary resources again
- * without requesting the admission, so rebuilding them without the flavors
would remove the node
- * selector and the tolerations which Kueue assigned to the pods.
+ * Sets the KueuePodSetInfo of the Workload which Kueue admitted before on
the context of a
+ * resource whose driver or master exists already. Such a reconcile applies
the secondary
+ * resources again without requesting the admission, so rebuilding them
without it would remove
+ * the node selector, the tolerations, the labels and the annotations which
Kueue assigned to the
+ * pods.
*
* @param context The context of the resource to be reconciled.
* @return The progress to return while the flavors cannot be read, or empty
to proceed.
- * @throws IllegalArgumentException if a node label of the flavors conflicts
with the node
- * selector of a pod set, like {@link #holdForAdmission}.
+ * @throws IllegalArgumentException if the KueuePodSetInfo conflicts with a
pod set, like {@link
+ * #holdForAdmission}.
*/
- public static Optional<ReconcileProgress> applyAdmittedFlavors(final
BaseContext<?> context) {
+ public static Optional<ReconcileProgress> applyAdmittedPodSetInfos(final
BaseContext<?> context) {
HasMetadata resource = context.getResource();
String workloadName = KueueWorkloadFactory.getWorkloadName(resource);
Workload workload;
@@ -392,12 +399,12 @@ public final class KueueWorkloadUtils {
return Optional.empty();
}
try {
- Map<String, KueuePodSetFlavor> flavors =
resolvePodSetFlavors(context.getClient(), workload);
+ Map<String, KueuePodSetInfo> podSetInfos =
resolvePodSetInfos(context.getClient(), workload);
// The pod set templates of the Workload hold the node selectors as of
the admission, so the
// check reports a ResourceFlavor which was edited since. Unlike the
admission, the quota is
// not released with the failure: the pods it was reserved for are
running already.
- checkNoNodeSelectorConflict(flavors, workload);
- context.setKueuePodSetFlavors(flavors);
+ checkNoPodTemplateConflict(podSetInfos, workload);
+ context.setKueuePodSetInfos(podSetInfos);
} catch (KubernetesClientException e) {
return Optional.of(retryAfterFlavorReadFailure(context, e,
workloadName));
}
@@ -461,7 +468,7 @@ public final class KueueWorkloadUtils {
* into quota which nothing uses. An admitted Workload is kept, since the
resources it was
* admitted for may be running already, and it is released with them like
the Workload of a
* queued resource. If they were requested already, its flavors are applied
again like {@link
- * #applyAdmittedFlavors}, since the secondary resources are applied again
in this reconcile.
+ * #applyAdmittedPodSetInfos}, since the secondary resources are applied
again in this reconcile.
* Otherwise they start without Kueue, so the flavors are not read, which
must not hold them
* back, and an evicted or deactivated Workload is released first rather
than kept for resources
* whose quota Kueue is about to take or no longer counts. Like the
admission request, a failed
@@ -473,9 +480,8 @@ public final class KueueWorkloadUtils {
* an admitted Workload, so that a resource without one costs no lookup.
* @return The progress to return while the release or the read of the
flavors fails, or empty to
* proceed.
- * @throws IllegalArgumentException if a node label of the flavors of an
admitted Workload
- * conflicts with the node selector of a pod set of the requested
resources, like {@link
- * #applyAdmittedFlavors}.
+ * @throws IllegalArgumentException if the flavors of an admitted Workload
conflict with a pod set
+ * of the requested resources, like {@link #applyAdmittedPodSetInfos}.
*/
public static Optional<ReconcileProgress> handleDequeuedWorkload(
final BaseContext<?> context, final BooleanSupplier requested) {
@@ -485,7 +491,7 @@ public final class KueueWorkloadUtils {
}
if (isAdmitted(workload.get())) {
if (requested.getAsBoolean()) {
- return applyAdmittedFlavors(context);
+ return applyAdmittedPodSetInfos(context);
}
if (findEviction(workload.get()).isEmpty() &&
!isDeactivated(workload.get())) {
return Optional.empty();
@@ -541,24 +547,35 @@ public final class KueueWorkloadUtils {
/**
* Resolves the node selector and tolerations of the ResourceFlavors which
Kueue assigned to each
- * pod set of the admitted Workload, like Kueue's `podset.FromAssignment`.
The Topology Aware
- * Scheduling gate and annotation are not handled. The flavors of a pod set
are applied in the
- * order of the resource names, so that a later flavor overwrites a node
label deterministically.
+ * pod set of the admitted Workload, like Kueue's `podset.FromAssignment`,
and merges the
+ * podSetUpdates of its admission checks into them, like Kueue's
`podset.FromUpdate`. The Topology
+ * Aware Scheduling gate and annotation are not handled. The flavors of a
pod set are applied in
+ * the order of the resource names, so that a later flavor overwrites a node
label
+ * deterministically. Like Kueue built-in integrations, each admission check
adds its first
+ * podSetUpdate of the pod set, in the order of `status.admissionChecks`,
which Kueue keeps sorted
+ * by name. Their state is not checked, since Kueue admits a Workload only
once all of its
+ * admission checks are ready.
*
* @param client The KubernetesClient.
* @param admitted The admitted Workload.
- * @return The KueuePodSetFlavor by the pod set name. A pod set without an
assignment, or one
- * whose flavors have neither node labels nor tolerations, is absent, so
that the resources
- * of such a pod set are built as they are without Kueue.
+ * @return The KueuePodSetInfo by the pod set name. A pod set without an
assignment, or one to
+ * which neither its flavors nor the podSetUpdates add anything, is
absent, so that the
+ * resources of such a pod set are built as they are without Kueue.
* @throws KubernetesClientException if a ResourceFlavor cannot be read.
+ * @throws IllegalArgumentException if a podSetUpdate changes a value of the
node selector, the
+ * labels or the annotations which the flavors or an earlier admission
check assigned to the
+ * pod set. Like Kueue built-in integrations, this is permanent. A
conflict with the pod
+ * templates is checked separately, like Kueue's `podset.Merge`.
*/
- public static Map<String, KueuePodSetFlavor> resolvePodSetFlavors(
+ public static Map<String, KueuePodSetInfo> resolvePodSetInfos(
final KubernetesClient client, final Workload admitted) {
- Map<String, KueuePodSetFlavor> result = new HashMap<>();
+ Map<String, KueuePodSetInfo> result = new HashMap<>();
WorkloadStatus status = admitted.getStatus();
if (status == null || status.getAdmission() == null) {
return result;
}
+ List<AdmissionCheckState> admissionChecks =
+ status.getAdmissionChecks() == null ? List.of() :
status.getAdmissionChecks();
Map<String, ResourceFlavor> flavorCache = new HashMap<>();
for (PodSetAssignment assignment :
status.getAdmission().getPodSetAssignments()) {
Map<String, String> nodeSelector = new HashMap<>();
@@ -574,18 +591,91 @@ public final class KueueWorkloadUtils {
nodeSelector.putAll(flavor.getSpec().getNodeLabels());
}
if (flavor.getSpec().getTolerations() != null) {
- KueuePodSetFlavor.addTolerations(tolerations,
flavor.getSpec().getTolerations());
+ KueuePodSetInfo.addTolerations(tolerations,
flavor.getSpec().getTolerations());
}
}
}
- if (nodeSelector.isEmpty() && tolerations.isEmpty()) {
+ Map<String, String> labels = new HashMap<>();
+ Map<String, String> annotations = new HashMap<>();
+ String podSetName = assignment.getName();
+ for (AdmissionCheckState check : admissionChecks) {
+ Optional<PodSetUpdate> found = findPodSetUpdate(check, podSetName);
+ if (found.isEmpty()) {
+ continue;
+ }
+ PodSetUpdate update = found.get();
+ String conflict =
+ "Kueue admission check "
+ + check.getName()
+ + " conflicts with the ResourceFlavors or an earlier admission
check of the "
+ + podSetName
+ + " pods for the ";
+ putAllWithoutConflict(nodeSelector, update.getNodeSelector(), conflict
+ "node selector");
+ putAllWithoutConflict(labels, update.getLabels(), conflict + "labels");
+ putAllWithoutConflict(annotations, update.getAnnotations(), conflict +
"annotations");
+ if (update.getTolerations() != null) {
+ KueuePodSetInfo.addTolerations(tolerations, update.getTolerations());
+ }
+ }
+ if (nodeSelector.isEmpty()
+ && tolerations.isEmpty()
+ && labels.isEmpty()
+ && annotations.isEmpty()) {
continue;
}
- result.put(assignment.getName(), new KueuePodSetFlavor(nodeSelector,
tolerations));
+ result.put(
+ podSetName, new KueuePodSetInfo(nodeSelector, tolerations, labels,
annotations));
}
return result;
}
+ /** Like Kueue, only the first podSetUpdate of a pod set in an admission
check is used. */
+ private static Optional<PodSetUpdate> findPodSetUpdate(
+ final AdmissionCheckState check, final String podSetName) {
+ if (check.getPodSetUpdates() == null) {
+ return Optional.empty();
+ }
+ return check.getPodSetUpdates().stream()
+ .filter(update -> podSetName.equals(update.getName()))
+ .findFirst();
+ }
+
+ /**
+ * Like Kueue's `PodSetInfo.Merge`, adds the values of a podSetUpdate to the
ones which the
+ * flavors and the earlier admission checks assigned to a pod set, which it
must not change.
+ */
+ private static void putAllWithoutConflict(
+ final Map<String, String> assigned,
+ final Map<String, String> updated,
+ final String conflict) {
+ checkNoConflict(conflict, assigned, updated);
+ if (updated != null) {
+ assigned.putAll(updated);
+ }
+ }
+
+ /**
+ * Like Kueue's `maps.HaveConflict`, a key which both maps have must have
the same value.
+ *
+ * @param conflict The start of the message, which names what conflicts.
+ * @param existing The values before the merge, reported as `value1`.
+ * @param added The values to be merged, reported as `value2`.
+ * @throws IllegalArgumentException if a key of both maps has different
values.
+ */
+ private static void checkNoConflict(
+ final String conflict, final Map<String, String> existing, final
Map<String, String> added) {
+ if (existing == null || added == null) {
+ return;
+ }
+ for (Map.Entry<String, String> e : existing.entrySet()) {
+ String value = added.get(e.getKey());
+ if (value != null && !value.equals(e.getValue())) {
+ throw new IllegalArgumentException(
+ conflict + " key=" + e.getKey() + ", value1=" + e.getValue() + ",
value2=" + value);
+ }
+ }
+ }
+
private static ResourceFlavor getResourceFlavor(
final KubernetesClient client, final String name) {
ResourceFlavor flavor =
client.resources(ResourceFlavor.class).withName(name).get();
@@ -596,54 +686,40 @@ public final class KueueWorkloadUtils {
return flavor;
}
- private static Map<String, String> podSetNodeSelector(
- final Workload workload, final String podSetName) {
+ private static PodTemplateSpec podSetTemplate(final Workload workload, final
String podSetName) {
return workload.getSpec().getPodSets().stream()
.filter(podSet -> podSetName.equals(podSet.getName()))
.map(PodSet::getTemplate)
- .filter(template -> template != null && template.getSpec() != null)
- .map(template -> template.getSpec().getNodeSelector())
.filter(Objects::nonNull)
.findFirst()
- .orElse(Map.of());
+ .orElseGet(PodTemplateSpec::new);
}
/**
- * Like Kueue's podset.Merge, a node label must not change the node selector
of the pods.
+ * Like Kueue's podset.Merge, the node selector, labels and annotations
which Kueue assigned must
+ * not change the ones of the pods.
*
- * @param flavors The resolved flavors by the pod set name.
+ * @param podSetInfos The resolved KueuePodSetInfo by the pod set name.
* @param desired The Workload built for the resource, whose pod set
templates have the node
- * selectors of the pods.
- * @throws IllegalArgumentException if a node label of the flavors conflicts
with the node
- * selector of a pod set. Like Kueue built-in integrations, this is
permanent.
+ * selectors, labels and annotations of the pods.
+ * @throws IllegalArgumentException if a value of a KueuePodSetInfo
conflicts with a pod set
+ * template. Like Kueue built-in integrations, this is permanent.
*/
- static void checkNoNodeSelectorConflict(
- final Map<String, KueuePodSetFlavor> flavors, final Workload desired) {
- for (Map.Entry<String, KueuePodSetFlavor> flavor : flavors.entrySet()) {
- checkNoNodeSelectorConflict(
- flavor.getKey(),
- podSetNodeSelector(desired, flavor.getKey()),
- flavor.getValue().nodeSelector());
- }
- }
-
- private static void checkNoNodeSelectorConflict(
- final String podSetName,
- final Map<String, String> podNodeSelector,
- final Map<String, String> flavorNodeLabels) {
- for (Map.Entry<String, String> e : podNodeSelector.entrySet()) {
- String flavorValue = flavorNodeLabels.get(e.getKey());
- if (flavorValue != null && !flavorValue.equals(e.getValue())) {
- throw new IllegalArgumentException(
- "The node labels of the Kueue ResourceFlavors conflict with the
node selector of the "
- + podSetName
- + " pods for key="
- + e.getKey()
- + ", value1="
- + e.getValue()
- + ", value2="
- + flavorValue);
- }
+ static void checkNoPodTemplateConflict(
+ final Map<String, KueuePodSetInfo> podSetInfos, final Workload desired) {
+ for (Map.Entry<String, KueuePodSetInfo> entry : podSetInfos.entrySet()) {
+ PodTemplateSpec template = podSetTemplate(desired, entry.getKey());
+ ObjectMeta metadata =
+ template.getMetadata() == null ? new ObjectMeta() :
template.getMetadata();
+ PodSpec spec = template.getSpec() == null ? new PodSpec() :
template.getSpec();
+ KueuePodSetInfo info = entry.getValue();
+ String conflict =
+ "The Kueue ResourceFlavors and admission checks conflict with the "
+ + entry.getKey()
+ + " pods for the ";
+ checkNoConflict(conflict + "node selector", spec.getNodeSelector(),
info.nodeSelector());
+ checkNoConflict(conflict + "labels", metadata.getLabels(),
info.labels());
+ checkNoConflict(conflict + "annotations", metadata.getAnnotations(),
info.annotations());
}
}
diff --git
a/spark-operator/src/main/java/org/apache/spark/k8s/operator/kueue/v1beta2/AdmissionCheckState.java
b/spark-operator/src/main/java/org/apache/spark/k8s/operator/kueue/v1beta2/AdmissionCheckState.java
new file mode 100644
index 00000000..10a6c9d9
--- /dev/null
+++
b/spark-operator/src/main/java/org/apache/spark/k8s/operator/kueue/v1beta2/AdmissionCheckState.java
@@ -0,0 +1,45 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.spark.k8s.operator.kueue.v1beta2;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
+import com.fasterxml.jackson.annotation.JsonInclude;
+import lombok.AllArgsConstructor;
+import lombok.Builder;
+import lombok.Data;
+import lombok.NoArgsConstructor;
+
+/**
+ * AdmissionCheckState is the state of an admission check of a Workload, which
the controller of
+ * the check reports, e.g. with the pod set updates of a ProvisioningRequest
admission check.
+ */
+@Data
+@NoArgsConstructor
+@AllArgsConstructor
+@Builder
+@JsonInclude(JsonInclude.Include.NON_NULL)
+@JsonIgnoreProperties(ignoreUnknown = true)
+public class AdmissionCheckState {
+ private String name;
+ @Builder.Default
+ private List<PodSetUpdate> podSetUpdates = new ArrayList<>();
+}
diff --git
a/spark-operator/src/main/java/org/apache/spark/k8s/operator/kueue/v1beta2/PodSetUpdate.java
b/spark-operator/src/main/java/org/apache/spark/k8s/operator/kueue/v1beta2/PodSetUpdate.java
new file mode 100644
index 00000000..49b10155
--- /dev/null
+++
b/spark-operator/src/main/java/org/apache/spark/k8s/operator/kueue/v1beta2/PodSetUpdate.java
@@ -0,0 +1,54 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.spark.k8s.operator.kueue.v1beta2;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
+import com.fasterxml.jackson.annotation.JsonInclude;
+import io.fabric8.kubernetes.api.model.Toleration;
+import lombok.AllArgsConstructor;
+import lombok.Builder;
+import lombok.Data;
+import lombok.NoArgsConstructor;
+
+/**
+ * PodSetUpdate is a modification of the pods of a pod set which an admission
check suggests, e.g.
+ * the annotations which let the pods consume the capacity of a
ProvisioningRequest.
+ */
+@Data
+@NoArgsConstructor
+@AllArgsConstructor
+@Builder
+@JsonInclude(JsonInclude.Include.NON_NULL)
+@JsonIgnoreProperties(ignoreUnknown = true)
+public class PodSetUpdate {
+ private String name;
+ @Builder.Default
+ private Map<String, String> labels = new HashMap<>();
+ @Builder.Default
+ private Map<String, String> annotations = new HashMap<>();
+ @Builder.Default
+ private Map<String, String> nodeSelector = new HashMap<>();
+ @Builder.Default
+ private List<Toleration> tolerations = new ArrayList<>();
+}
diff --git
a/spark-operator/src/main/java/org/apache/spark/k8s/operator/kueue/v1beta2/WorkloadStatus.java
b/spark-operator/src/main/java/org/apache/spark/k8s/operator/kueue/v1beta2/WorkloadStatus.java
index 90f6f478..0aea7c1a 100644
---
a/spark-operator/src/main/java/org/apache/spark/k8s/operator/kueue/v1beta2/WorkloadStatus.java
+++
b/spark-operator/src/main/java/org/apache/spark/k8s/operator/kueue/v1beta2/WorkloadStatus.java
@@ -45,6 +45,9 @@ public class WorkloadStatus {
private List<Condition> conditions = new ArrayList<>();
private Admission admission;
private RequeueState requeueState;
+ // Without a default, so that a status which the operator creates to record
a condition does not
+ // write the admission checks, which Kueue owns.
+ private List<AdmissionCheckState> admissionChecks;
/**
* Checks if the workload has been admitted by Kueue.
diff --git
a/spark-operator/src/main/java/org/apache/spark/k8s/operator/reconciler/SparkAppResourceSpecFactory.java
b/spark-operator/src/main/java/org/apache/spark/k8s/operator/reconciler/SparkAppResourceSpecFactory.java
index 53badf6e..03f036fd 100644
---
a/spark-operator/src/main/java/org/apache/spark/k8s/operator/reconciler/SparkAppResourceSpecFactory.java
+++
b/spark-operator/src/main/java/org/apache/spark/k8s/operator/reconciler/SparkAppResourceSpecFactory.java
@@ -35,7 +35,6 @@ import java.util.HashMap;
import java.util.Map;
import java.util.Optional;
-import io.fabric8.kubernetes.api.model.PodSpec;
import io.fabric8.kubernetes.api.model.PodTemplateSpec;
import io.fabric8.kubernetes.client.KubernetesClient;
import lombok.extern.slf4j.Slf4j;
@@ -45,7 +44,7 @@ import org.apache.spark.k8s.operator.SparkAppSubmissionWorker;
import org.apache.spark.k8s.operator.SparkApplication;
import org.apache.spark.k8s.operator.decorators.OwnerResourceDecorator;
import org.apache.spark.k8s.operator.decorators.SecurityContextDecorator;
-import org.apache.spark.k8s.operator.kueue.KueuePodSetFlavor;
+import org.apache.spark.k8s.operator.kueue.KueuePodSetInfo;
import org.apache.spark.k8s.operator.kueue.KueueWorkloadFactory;
import org.apache.spark.k8s.operator.spec.ApplicationSpec;
import org.apache.spark.k8s.operator.spec.BaseApplicationTemplateSpec;
@@ -59,23 +58,23 @@ public final class SparkAppResourceSpecFactory {
private SparkAppResourceSpecFactory() {}
/**
- * Builds a SparkAppResourceSpec for the given SparkApplication with the
flavors which Kueue
- * assigned to the driver and executor pod sets. The flavors are applied to
the pod templates,
- * because the executor pods are created by the driver.
+ * Builds a SparkAppResourceSpec for the given SparkApplication with the
KueuePodSetInfo which
+ * Kueue assigned to the driver and executor pod sets. It is applied to the
pod templates, because
+ * the executor pods are created by the driver.
*
* @param app The SparkApplication.
* @param client The KubernetesClient.
* @param worker The SparkAppSubmissionWorker.
- * @param kueuePodSetFlavors The KueuePodSetFlavor by the pod set name.
+ * @param kueuePodSetInfos The KueuePodSetInfo by the pod set name.
* @return A SparkAppResourceSpec containing the configured resources.
*/
public static SparkAppResourceSpec buildResourceSpec(
final SparkApplication app,
final KubernetesClient client,
final SparkAppSubmissionWorker worker,
- final Map<String, KueuePodSetFlavor> kueuePodSetFlavors) {
+ final Map<String, KueuePodSetInfo> kueuePodSetInfos) {
SparkApplication appToSubmit =
- kueuePodSetFlavors.isEmpty() ? app : applyKueuePodSetFlavors(app,
kueuePodSetFlavors);
+ kueuePodSetInfos.isEmpty() ? app : applyKueuePodSetInfos(app,
kueuePodSetInfos);
Map<String, String> confOverrides = overrideDependencyConf(appToSubmit);
SparkAppResourceSpec resourceSpec = worker.getResourceSpec(appToSubmit,
client, confOverrides);
cleanUpTempResourcesForApp(appToSubmit, confOverrides);
@@ -87,35 +86,31 @@ public final class SparkAppResourceSpecFactory {
/**
* Returns a copy of the SparkApplication whose driver and executor pod
templates have the
- * flavors, so that the SparkApplication and its Kueue Workload are not
changed.
+ * KueuePodSetInfo, so that the SparkApplication and its Kueue Workload are
not changed.
*/
- private static SparkApplication applyKueuePodSetFlavors(
- final SparkApplication app, final Map<String, KueuePodSetFlavor>
kueuePodSetFlavors) {
+ private static SparkApplication applyKueuePodSetInfos(
+ final SparkApplication app, final Map<String, KueuePodSetInfo>
kueuePodSetInfos) {
SparkApplication copy = ReconcilerUtils.clone(app);
ApplicationSpec spec = copy.getSpec();
- KueuePodSetFlavor driverFlavor =
kueuePodSetFlavors.get(KueueWorkloadFactory.PODSET_DRIVER);
- if (driverFlavor != null) {
- spec.setDriverSpec(applyKueuePodSetFlavor(spec.getDriverSpec(),
driverFlavor));
+ KueuePodSetInfo driverInfo =
kueuePodSetInfos.get(KueueWorkloadFactory.PODSET_DRIVER);
+ if (driverInfo != null) {
+ spec.setDriverSpec(applyKueuePodSetInfo(spec.getDriverSpec(),
driverInfo));
}
- KueuePodSetFlavor executorFlavor =
- kueuePodSetFlavors.get(KueueWorkloadFactory.PODSET_EXECUTOR);
- if (executorFlavor != null) {
- spec.setExecutorSpec(applyKueuePodSetFlavor(spec.getExecutorSpec(),
executorFlavor));
+ KueuePodSetInfo executorInfo =
kueuePodSetInfos.get(KueueWorkloadFactory.PODSET_EXECUTOR);
+ if (executorInfo != null) {
+ spec.setExecutorSpec(applyKueuePodSetInfo(spec.getExecutorSpec(),
executorInfo));
}
return copy;
}
- private static BaseApplicationTemplateSpec applyKueuePodSetFlavor(
- final BaseApplicationTemplateSpec templateSpec, final KueuePodSetFlavor
flavor) {
+ private static BaseApplicationTemplateSpec applyKueuePodSetInfo(
+ final BaseApplicationTemplateSpec templateSpec, final KueuePodSetInfo
podSetInfo) {
BaseApplicationTemplateSpec result =
templateSpec == null ? new BaseApplicationTemplateSpec() :
templateSpec;
if (result.getPodTemplateSpec() == null) {
result.setPodTemplateSpec(new PodTemplateSpec());
}
- if (result.getPodTemplateSpec().getSpec() == null) {
- result.getPodTemplateSpec().setSpec(new PodSpec());
- }
- flavor.applyTo(result.getPodTemplateSpec().getSpec());
+ podSetInfo.applyTo(result.getPodTemplateSpec());
return result;
}
diff --git
a/spark-operator/src/main/java/org/apache/spark/k8s/operator/reconciler/reconcilesteps/AppInitStep.java
b/spark-operator/src/main/java/org/apache/spark/k8s/operator/reconciler/reconcilesteps/AppInitStep.java
index e03851ec..c2e9ab1c 100644
---
a/spark-operator/src/main/java/org/apache/spark/k8s/operator/reconciler/reconcilesteps/AppInitStep.java
+++
b/spark-operator/src/main/java/org/apache/spark/k8s/operator/reconciler/reconcilesteps/AppInitStep.java
@@ -217,7 +217,7 @@ public final class AppInitStep extends AppReconcileStep {
if (driverRequested || isDriverRequested(context)) {
// The driver resources are applied again in this reconcile, so the
flavors of the Workload
// which was admitted before are resolved again instead of dropping
them from the resources.
- return KueueWorkloadUtils.applyAdmittedFlavors(context);
+ return KueueWorkloadUtils.applyAdmittedPodSetInfos(context);
}
} catch (KubernetesClientException e) {
// Requesting the admission of a driver which is already running would
be wrong, so the
diff --git
a/spark-operator/src/main/java/org/apache/spark/k8s/operator/reconciler/reconcilesteps/ClusterInitStep.java
b/spark-operator/src/main/java/org/apache/spark/k8s/operator/reconciler/reconcilesteps/ClusterInitStep.java
index 2f5f399b..d5eb1859 100644
---
a/spark-operator/src/main/java/org/apache/spark/k8s/operator/reconciler/reconcilesteps/ClusterInitStep.java
+++
b/spark-operator/src/main/java/org/apache/spark/k8s/operator/reconciler/reconcilesteps/ClusterInitStep.java
@@ -240,7 +240,7 @@ public final class ClusterInitStep extends
ClusterReconcileStep {
// The master and worker StatefulSets are applied again in this
reconcile, so the flavors
// of the Workload which was admitted before are resolved again
instead of dropping them
// from the pod templates.
- return KueueWorkloadUtils.applyAdmittedFlavors(context);
+ return KueueWorkloadUtils.applyAdmittedPodSetInfos(context);
}
} catch (KubernetesClientException e) {
// Requesting the admission of a master which is already running would
be wrong, so the
diff --git
a/spark-operator/src/test/java/org/apache/spark/k8s/operator/context/SparkAppContextTest.java
b/spark-operator/src/test/java/org/apache/spark/k8s/operator/context/SparkAppContextTest.java
index 46688fc5..03343b26 100644
---
a/spark-operator/src/test/java/org/apache/spark/k8s/operator/context/SparkAppContextTest.java
+++
b/spark-operator/src/test/java/org/apache/spark/k8s/operator/context/SparkAppContextTest.java
@@ -59,7 +59,7 @@ import org.apache.spark.k8s.operator.Constants;
import org.apache.spark.k8s.operator.SparkAppResourceSpec;
import org.apache.spark.k8s.operator.SparkAppSubmissionWorker;
import org.apache.spark.k8s.operator.SparkApplication;
-import org.apache.spark.k8s.operator.kueue.KueuePodSetFlavor;
+import org.apache.spark.k8s.operator.kueue.KueuePodSetInfo;
class SparkAppContextTest {
private final SparkApplication application = buildApplication();
@@ -214,7 +214,7 @@ class SparkAppContextTest {
}
@Test
- void kueuePodSetFlavorsRebuildTheCachedResourceSpec() {
+ void kueuePodSetInfosRebuildTheCachedResourceSpec() {
// The driver pod spec can be built before the Kueue admission, e.g. to
find the driver pod of
// the current attempt, so setting the flavors drops the cached one.
KubernetesClient client = mock(KubernetesClient.class);
@@ -231,8 +231,8 @@ class SparkAppContextTest {
verify(worker, times(1)).getResourceSpec(any(), any(), any());
Toleration toleration = new Toleration("NoSchedule", "spot", "Exists",
null, null);
- context.setKueuePodSetFlavors(
- Map.of("driver", new KueuePodSetFlavor(Map.of("pool", "cpu"),
List.of(toleration))));
+ context.setKueuePodSetInfos(
+ Map.of("driver", new KueuePodSetInfo(Map.of("pool", "cpu"),
List.of(toleration))));
context.getDriverPodSpec();
ArgumentCaptor<SparkApplication> captor =
ArgumentCaptor.forClass(SparkApplication.class);
diff --git
a/spark-operator/src/test/java/org/apache/spark/k8s/operator/context/SparkClusterContextTest.java
b/spark-operator/src/test/java/org/apache/spark/k8s/operator/context/SparkClusterContextTest.java
index 57de9004..2806f771 100644
---
a/spark-operator/src/test/java/org/apache/spark/k8s/operator/context/SparkClusterContextTest.java
+++
b/spark-operator/src/test/java/org/apache/spark/k8s/operator/context/SparkClusterContextTest.java
@@ -30,6 +30,7 @@ import java.util.Optional;
import java.util.stream.Stream;
import com.fasterxml.jackson.core.JsonProcessingException;
+import io.fabric8.kubernetes.api.model.ObjectMeta;
import io.fabric8.kubernetes.api.model.ObjectMetaBuilder;
import io.fabric8.kubernetes.api.model.Pod;
import io.fabric8.kubernetes.api.model.PodBuilder;
@@ -44,7 +45,7 @@ import org.apache.spark.k8s.operator.Constants;
import org.apache.spark.k8s.operator.SparkCluster;
import org.apache.spark.k8s.operator.SparkClusterSubmissionWorker;
import org.apache.spark.k8s.operator.config.SparkOperatorConf;
-import org.apache.spark.k8s.operator.kueue.KueuePodSetFlavor;
+import org.apache.spark.k8s.operator.kueue.KueuePodSetInfo;
import org.apache.spark.k8s.operator.kueue.KueueWorkloadFactory;
import org.apache.spark.k8s.operator.kueue.v1beta2.Workload;
import org.apache.spark.k8s.operator.spec.ClusterSpec;
@@ -57,7 +58,7 @@ import org.apache.spark.k8s.operator.utils.TestUtils;
class SparkClusterContextTest {
@Test
- void kueuePodSetFlavorsAreAppliedToStatefulSets() {
+ void kueuePodSetInfosAreAppliedToStatefulSets() {
SparkCluster cluster = buildCluster();
String clusterSpec = asJson(cluster.getSpec());
SparkClusterContext context =
@@ -69,12 +70,20 @@ class SparkClusterContextTest {
Assertions.assertEquals(
List.of(),
podSpec(context.getMasterStatefulSetSpec()).getTolerations());
- context.setKueuePodSetFlavors(
+ context.setKueuePodSetInfos(
Map.of(
KueueWorkloadFactory.PODSET_MASTER,
- new KueuePodSetFlavor(Map.of("pool", "cpu"), List.of(spot)),
+ new KueuePodSetInfo(
+ Map.of("pool", "cpu"),
+ List.of(spot),
+ Map.of("team", "a"),
+ Map.of("provisioning", "pr-master")),
KueueWorkloadFactory.PODSET_WORKER,
- new KueuePodSetFlavor(Map.of("pool", "gpu"), List.of(gpu))));
+ new KueuePodSetInfo(
+ Map.of("pool", "gpu"),
+ List.of(gpu),
+ Map.of("team", "a"),
+ Map.of("provisioning", "pr-worker"))));
PodSpec master = podSpec(context.getMasterStatefulSetSpec());
Assertions.assertEquals(Map.of("pool", "cpu"), master.getNodeSelector());
@@ -82,6 +91,21 @@ class SparkClusterContextTest {
PodSpec worker = podSpec(context.getWorkerStatefulSetSpec());
Assertions.assertEquals(Map.of("pool", "gpu"), worker.getNodeSelector());
Assertions.assertEquals(List.of(gpu), worker.getTolerations());
+ // The labels which the StatefulSets select their pods with are kept
+ ObjectMeta masterMetadata =
+
context.getMasterStatefulSetSpec().getSpec().getTemplate().getMetadata();
+ Assertions.assertEquals("a", masterMetadata.getLabels().get("team"));
+ Assertions.assertEquals("pr-master",
masterMetadata.getAnnotations().get("provisioning"));
+ Assertions.assertEquals(
+ Constants.LABEL_SPARK_ROLE_MASTER_VALUE,
+ masterMetadata.getLabels().get(Constants.LABEL_SPARK_ROLE_NAME));
+ ObjectMeta workerMetadata =
+
context.getWorkerStatefulSetSpec().getSpec().getTemplate().getMetadata();
+ Assertions.assertEquals("a", workerMetadata.getLabels().get("team"));
+ Assertions.assertEquals("pr-worker",
workerMetadata.getAnnotations().get("provisioning"));
+ Assertions.assertEquals(
+ Constants.LABEL_SPARK_ROLE_WORKER_VALUE,
+ workerMetadata.getLabels().get(Constants.LABEL_SPARK_ROLE_NAME));
// The flavors are applied to the built StatefulSets only, so the
SparkCluster keeps the pod
// templates which the pod sets of its Kueue Workload are hashed from
Assertions.assertEquals(clusterSpec, asJson(cluster.getSpec()));
diff --git
a/spark-operator/src/test/java/org/apache/spark/k8s/operator/kueue/KueuePodSetFlavorTest.java
b/spark-operator/src/test/java/org/apache/spark/k8s/operator/kueue/KueuePodSetInfoTest.java
similarity index 60%
rename from
spark-operator/src/test/java/org/apache/spark/k8s/operator/kueue/KueuePodSetFlavorTest.java
rename to
spark-operator/src/test/java/org/apache/spark/k8s/operator/kueue/KueuePodSetInfoTest.java
index 3f29f1b5..2bd996af 100644
---
a/spark-operator/src/test/java/org/apache/spark/k8s/operator/kueue/KueuePodSetFlavorTest.java
+++
b/spark-operator/src/test/java/org/apache/spark/k8s/operator/kueue/KueuePodSetInfoTest.java
@@ -28,17 +28,19 @@ import java.util.Map;
import io.fabric8.kubernetes.api.model.PodSpec;
import io.fabric8.kubernetes.api.model.PodSpecBuilder;
+import io.fabric8.kubernetes.api.model.PodTemplateSpec;
+import io.fabric8.kubernetes.api.model.PodTemplateSpecBuilder;
import io.fabric8.kubernetes.api.model.Toleration;
import org.junit.jupiter.api.Test;
-class KueuePodSetFlavorTest {
+class KueuePodSetInfoTest {
@Test
void applyToAddsNodeSelectorAndTolerationsToEmptyPodSpec() {
PodSpec podSpec = new PodSpec();
Toleration toleration = toleration("spot", "Equal", "true", "NoSchedule",
null);
- new KueuePodSetFlavor(Map.of("instance-type", "spot"),
List.of(toleration)).applyTo(podSpec);
+ new KueuePodSetInfo(Map.of("instance-type", "spot"),
List.of(toleration)).applyTo(podSpec);
assertEquals(Map.of("instance-type", "spot"), podSpec.getNodeSelector());
assertEquals(List.of(toleration), podSpec.getTolerations());
@@ -54,7 +56,7 @@ class KueuePodSetFlavorTest {
.withTolerations(existing)
.build();
- new KueuePodSetFlavor(Map.of("instance-type", "spot", "pool", "p1"),
List.of(added))
+ new KueuePodSetInfo(Map.of("instance-type", "spot", "pool", "p1"),
List.of(added))
.applyTo(podSpec);
assertEquals(
@@ -69,7 +71,7 @@ class KueuePodSetFlavorTest {
Toleration existing = toleration("spot", "Equal", "true", "NoExecute",
30L);
PodSpec podSpec = new PodSpecBuilder().withTolerations(existing).build();
- new KueuePodSetFlavor(
+ new KueuePodSetInfo(
Map.of(),
List.of(
toleration("spot", null, "true", "NoExecute", 60L),
@@ -82,22 +84,72 @@ class KueuePodSetFlavorTest {
podSpec.getTolerations());
}
+ @Test
+ void applyToAddsLabelsAndAnnotationsToEmptyPodTemplate() {
+ // The executor pod template of a SparkApplication may have neither
metadata nor a pod spec
+ PodTemplateSpec template = new PodTemplateSpec();
+ Toleration toleration = toleration("provisioned", "Exists", null,
"NoSchedule", null);
+
+ new KueuePodSetInfo(
+ Map.of("pool", "provisioned"),
+ List.of(toleration),
+ Map.of("team", "a"),
+ Map.of("autoscaling.x-k8s.io/consume-provisioning-request", "pr"))
+ .applyTo(template);
+
+ assertEquals(Map.of("team", "a"), template.getMetadata().getLabels());
+ assertEquals(
+ Map.of("autoscaling.x-k8s.io/consume-provisioning-request", "pr"),
+ template.getMetadata().getAnnotations());
+ assertEquals(Map.of("pool", "provisioned"),
template.getSpec().getNodeSelector());
+ assertEquals(List.of(toleration), template.getSpec().getTolerations());
+ }
+
+ @Test
+ void applyToKeepsExistingLabelsAndAnnotations() {
+ PodTemplateSpec template =
+ new PodTemplateSpecBuilder()
+ .withNewMetadata()
+ .withLabels(Map.of("spark-role", "worker", "team", "a"))
+ .withAnnotations(Map.of("note", "kept"))
+ .endMetadata()
+ .withNewSpec()
+ .withNodeSelector(Map.of("zone", "a"))
+ .endSpec()
+ .build();
+
+ new KueuePodSetInfo(
+ Map.of("pool", "p1"),
+ List.of(),
+ Map.of("team", "a", "tier", "batch"),
+ Map.of("autoscaling.x-k8s.io/provisioning-class-name", "class"))
+ .applyTo(template);
+
+ assertEquals(
+ Map.of("spark-role", "worker", "team", "a", "tier", "batch"),
+ template.getMetadata().getLabels());
+ assertEquals(
+ Map.of("note", "kept", "autoscaling.x-k8s.io/provisioning-class-name",
"class"),
+ template.getMetadata().getAnnotations());
+ assertEquals(Map.of("zone", "a", "pool", "p1"),
template.getSpec().getNodeSelector());
+ }
+
@Test
void isSameToleration() {
assertTrue(
- KueuePodSetFlavor.isSameToleration(
+ KueuePodSetInfo.isSameToleration(
toleration("k", "", "v", "NoSchedule", null),
toleration("k", "Equal", "v", "NoSchedule", 10L)));
assertTrue(
- KueuePodSetFlavor.isSameToleration(
+ KueuePodSetInfo.isSameToleration(
toleration("k", "Exists", null, null, null),
toleration("k", "Exists", null, null, null)));
assertFalse(
- KueuePodSetFlavor.isSameToleration(
+ KueuePodSetInfo.isSameToleration(
toleration("k", "Exists", null, null, null),
toleration("k", "Equal", null, null, null)));
assertFalse(
- KueuePodSetFlavor.isSameToleration(
+ KueuePodSetInfo.isSameToleration(
toleration("k", "Equal", "v", "NoSchedule", null),
toleration("k", "Equal", "v", "NoExecute", null)));
}
diff --git
a/spark-operator/src/test/java/org/apache/spark/k8s/operator/kueue/KueueWorkloadUtilsTest.java
b/spark-operator/src/test/java/org/apache/spark/k8s/operator/kueue/KueueWorkloadUtilsTest.java
index 90efba4a..4cfe4675 100644
---
a/spark-operator/src/test/java/org/apache/spark/k8s/operator/kueue/KueueWorkloadUtilsTest.java
+++
b/spark-operator/src/test/java/org/apache/spark/k8s/operator/kueue/KueueWorkloadUtilsTest.java
@@ -77,8 +77,10 @@ import org.apache.spark.k8s.operator.context.SparkAppContext;
import
org.apache.spark.k8s.operator.kueue.KueueWorkloadUtils.AdmissionResponse;
import org.apache.spark.k8s.operator.kueue.KueueWorkloadUtils.AdmissionResult;
import org.apache.spark.k8s.operator.kueue.v1beta2.Admission;
+import org.apache.spark.k8s.operator.kueue.v1beta2.AdmissionCheckState;
import org.apache.spark.k8s.operator.kueue.v1beta2.PodSet;
import org.apache.spark.k8s.operator.kueue.v1beta2.PodSetAssignment;
+import org.apache.spark.k8s.operator.kueue.v1beta2.PodSetUpdate;
import org.apache.spark.k8s.operator.kueue.v1beta2.PriorityClassRef;
import org.apache.spark.k8s.operator.kueue.v1beta2.RequeueState;
import org.apache.spark.k8s.operator.kueue.v1beta2.ResourceFlavor;
@@ -216,13 +218,13 @@ class KueueWorkloadUtilsTest {
if (requested) {
// They are applied again, so the flavors which Kueue assigned to them
are not dropped
verify(context)
- .setKueuePodSetFlavors(
+ .setKueuePodSetInfos(
Map.of(
"executor",
- new KueuePodSetFlavor(Map.of("pool", "spot"),
List.of(toleration("spot")))));
+ new KueuePodSetInfo(Map.of("pool", "spot"),
List.of(toleration("spot")))));
} else {
// They start without Kueue, so the flavors are not read, which could
hold them back
- verify(context, never()).setKueuePodSetFlavors(any());
+ verify(context, never()).setKueuePodSetInfos(any());
}
}
@@ -445,7 +447,7 @@ class KueueWorkloadUtilsTest {
&& requeueAfter.compareTo(Duration.ofSeconds(60)) <= 0,
requeueAfter.toString());
Assertions.assertNotNull(getWorkload());
- verify(context, never()).setKueuePodSetFlavors(any());
+ verify(context, never()).setKueuePodSetInfos(any());
// The wait is reported with the reason of the eviction
ArgumentCaptor<EventRecord> captor =
ArgumentCaptor.forClass(EventRecord.class);
verify(eventRecorder).record(captor.capture());
@@ -934,7 +936,7 @@ class KueueWorkloadUtilsTest {
}
@Test
- void resolvePodSetFlavorsMergesFlavorsOfEachPodSet() {
+ void resolvePodSetInfosMergesFlavorsOfEachPodSet() {
Toleration spot = toleration("spot");
Toleration gpu = toleration("gpu");
createFlavor("cpu-flavor", Map.of("pool", "cpu", "zone", "a"),
List.of(spot));
@@ -942,47 +944,47 @@ class KueueWorkloadUtilsTest {
// Like Kueue, a flavor assigned to several resources is applied once
Map<String, String> driverFlavors = Map.of("cpu", "cpu-flavor", "memory",
"cpu-flavor");
// Like Kueue, a later flavor overwrites a node label, which is in the
resource name order.
- // The reverse insertion order pins the assertion to the sorting of
`resolvePodSetFlavors`.
+ // The reverse insertion order pins the assertion to the sorting of
`resolvePodSetInfos`.
Map<String, String> executorFlavors = new LinkedHashMap<>();
executorFlavors.put("nvidia.com/gpu", "gpu-flavor");
executorFlavors.put("cpu", "cpu-flavor");
- Map<String, KueuePodSetFlavor> flavors =
- KueueWorkloadUtils.resolvePodSetFlavors(
+ Map<String, KueuePodSetInfo> flavors =
+ KueueWorkloadUtils.resolvePodSetInfos(
kubernetesClient,
admittedWorkload(Map.of("driver", driverFlavors, "executor",
executorFlavors)));
Assertions.assertEquals(
Map.of(
"driver",
- new KueuePodSetFlavor(Map.of("pool", "cpu", "zone", "a"),
List.of(spot)),
+ new KueuePodSetInfo(Map.of("pool", "cpu", "zone", "a"),
List.of(spot)),
"executor",
- new KueuePodSetFlavor(
+ new KueuePodSetInfo(
Map.of("pool", "gpu", "zone", "a", "accelerator", "a100"),
List.of(spot, gpu))),
flavors);
}
@Test
- void resolvePodSetFlavorsWithoutAdmissionIsEmpty() {
+ void resolvePodSetInfosWithoutAdmissionIsEmpty() {
Workload admitted = workload("owner-uid-1", 1);
admitted.setStatus(status("Admitted", "True"));
Assertions.assertEquals(
- Map.of(), KueueWorkloadUtils.resolvePodSetFlavors(kubernetesClient,
admitted));
+ Map.of(), KueueWorkloadUtils.resolvePodSetInfos(kubernetesClient,
admitted));
Assertions.assertEquals(
Map.of(),
- KueueWorkloadUtils.resolvePodSetFlavors(kubernetesClient,
admittedWorkload(Map.of())));
+ KueueWorkloadUtils.resolvePodSetInfos(kubernetesClient,
admittedWorkload(Map.of())));
}
@Test
- void resolvePodSetFlavorsSkipsFlavorsWithoutNodeLabelsAndTolerations() {
+ void resolvePodSetInfosSkipsFlavorsWithoutNodeLabelsAndTolerations() {
// Kueue's stock `default-flavor` has neither, so its pod sets are built
as they are without
// Kueue, instead of being rebuilt around an empty pod template.
createFlavor("default-flavor", Map.of(), List.of());
createFlavor("spot-flavor", Map.of("pool", "spot"), List.of());
- Map<String, KueuePodSetFlavor> flavors =
- KueueWorkloadUtils.resolvePodSetFlavors(
+ Map<String, KueuePodSetInfo> flavors =
+ KueueWorkloadUtils.resolvePodSetInfos(
kubernetesClient,
admittedWorkload(
Map.of(
@@ -990,47 +992,215 @@ class KueueWorkloadUtilsTest {
"executor", Map.of("cpu", "spot-flavor"))));
Assertions.assertEquals(
- Map.of("executor", new KueuePodSetFlavor(Map.of("pool", "spot"),
List.of())), flavors);
+ Map.of("executor", new KueuePodSetInfo(Map.of("pool", "spot"),
List.of())), flavors);
}
@Test
- void checkNoNodeSelectorConflictAllowsTheSameNodeSelector() {
- Map<String, KueuePodSetFlavor> flavors =
- Map.of("executor", new KueuePodSetFlavor(Map.of("pool", "cpu"),
List.of()));
+ void checkNoPodTemplateConflictAllowsTheSameNodeSelector() {
+ Map<String, KueuePodSetInfo> flavors =
+ Map.of("executor", new KueuePodSetInfo(Map.of("pool", "cpu"),
List.of()));
Assertions.assertDoesNotThrow(
() ->
- KueueWorkloadUtils.checkNoNodeSelectorConflict(
+ KueueWorkloadUtils.checkNoPodTemplateConflict(
flavors, workloadWithNodeSelector(Map.of("pool", "cpu",
"zone", "a"))));
}
@Test
- void checkNoNodeSelectorConflictFailsOnDifferentValue() {
- Map<String, KueuePodSetFlavor> flavors =
- Map.of("executor", new KueuePodSetFlavor(Map.of("pool", "cpu"),
List.of()));
+ void checkNoPodTemplateConflictFailsOnDifferentNodeSelector() {
+ Map<String, KueuePodSetInfo> flavors =
+ Map.of("executor", new KueuePodSetInfo(Map.of("pool", "cpu"),
List.of()));
IllegalArgumentException e =
Assertions.assertThrows(
IllegalArgumentException.class,
() ->
- KueueWorkloadUtils.checkNoNodeSelectorConflict(
+ KueueWorkloadUtils.checkNoPodTemplateConflict(
flavors, workloadWithNodeSelector(Map.of("pool", "gpu"))));
Assertions.assertTrue(e.getMessage().contains("executor"), e.getMessage());
Assertions.assertTrue(e.getMessage().contains("pool"), e.getMessage());
}
@Test
- void resolvePodSetFlavorsFailsOnMissingFlavor() {
+ void checkNoPodTemplateConflictFailsOnDifferentLabelOrAnnotation() {
+ // Like Kueue's podset.Merge, the labels and annotations of the admission
checks must not
+ // change the ones of the pods either. The pod set template has no pod
spec here.
+ Workload desired = workloadWithMetadata(Map.of("team", "a"),
Map.of("note", "x"));
+ KueuePodSetInfo same = podSetInfoWithMetadata(Map.of("team", "a"),
Map.of("note", "x"));
+ KueuePodSetInfo otherLabel = podSetInfoWithMetadata(Map.of("team", "b"),
Map.of());
+ KueuePodSetInfo otherAnnotation = podSetInfoWithMetadata(Map.of(),
Map.of("note", "y"));
+
+ Assertions.assertDoesNotThrow(
+ () -> KueueWorkloadUtils.checkNoPodTemplateConflict(Map.of("executor",
same), desired));
+ IllegalArgumentException labels =
+ Assertions.assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ KueueWorkloadUtils.checkNoPodTemplateConflict(
+ Map.of("executor", otherLabel), desired));
+ Assertions.assertTrue(labels.getMessage().contains("labels"),
labels.getMessage());
+ Assertions.assertTrue(labels.getMessage().contains("team"),
labels.getMessage());
+ IllegalArgumentException annotations =
+ Assertions.assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ KueueWorkloadUtils.checkNoPodTemplateConflict(
+ Map.of("executor", otherAnnotation), desired));
+ Assertions.assertTrue(
+ annotations.getMessage().contains("annotations"),
annotations.getMessage());
+ Assertions.assertTrue(annotations.getMessage().contains("note"),
annotations.getMessage());
+ }
+
+ @Test
+ void resolvePodSetInfosFailsOnMissingFlavor() {
KubernetesClientException e =
Assertions.assertThrows(
KubernetesClientException.class,
() ->
- KueueWorkloadUtils.resolvePodSetFlavors(
+ KueueWorkloadUtils.resolvePodSetInfos(
kubernetesClient,
admittedWorkload(Map.of("executor", Map.of("cpu",
"missing-flavor")))));
Assertions.assertEquals(404, e.getCode());
}
+ @Test
+ void resolvePodSetInfosMergesAdmissionCheckUpdates() {
+ Toleration spot = toleration("spot");
+ Toleration provisioned = toleration("provisioned");
+ createFlavor("cpu-flavor", Map.of("pool", "cpu"), List.of(spot));
+ // Like Kueue, the admission checks are merged in the order of the list,
which Kueue keeps
+ // sorted by name, and each of them contributes the first update of a pod
set only.
+ Workload admitted =
+ admittedWorkload(
+ Map.of(
+ "driver", Map.of("cpu", "cpu-flavor"),
+ "executor", Map.of("cpu", "cpu-flavor")),
+ admissionCheck(
+ "a-check",
+ PodSetUpdate.builder()
+ .name("executor")
+ .labels(Map.of("team", "a"))
+ .annotations(Map.of("provisioning", "pr-1"))
+ // The node label of the flavor with the same value is no
conflict
+ .nodeSelector(Map.of("pool", "cpu"))
+ // The toleration of the flavor is not added again
+ .tolerations(List.of(spot, provisioned))
+ .build(),
+ PodSetUpdate.builder().name("executor").labels(Map.of("team",
"b")).build()),
+ admissionCheck(
+ "b-check",
+
PodSetUpdate.builder().name("driver").annotations(Map.of("class", "c")).build(),
+ PodSetUpdate.builder()
+ .name("executor")
+ .labels(Map.of("tier", "batch"))
+ .tolerations(List.of(provisioned))
+ .build()));
+
+ Assertions.assertEquals(
+ Map.of(
+ "driver",
+ new KueuePodSetInfo(
+ Map.of("pool", "cpu"), List.of(spot), Map.of(),
Map.of("class", "c")),
+ "executor",
+ new KueuePodSetInfo(
+ Map.of("pool", "cpu"),
+ List.of(spot, provisioned),
+ Map.of("team", "a", "tier", "batch"),
+ Map.of("provisioning", "pr-1"))),
+ KueueWorkloadUtils.resolvePodSetInfos(kubernetesClient, admitted));
+ }
+
+ @Test
+ void resolvePodSetInfosKeepsAPodSetWithAdmissionCheckUpdatesOnly() {
+ // e.g. a ProvisioningRequest admission check with Kueue's stock
`default-flavor`
+ createFlavor("default-flavor", Map.of(), List.of());
+ Map<String, String> annotations =
+ Map.of(
+ "autoscaling.x-k8s.io/consume-provisioning-request", "pr",
+ "autoscaling.x-k8s.io/provisioning-class-name", "class");
+ Workload admitted =
+ admittedWorkload(
+ Map.of(
+ "driver", Map.of("cpu", "default-flavor"),
+ "executor", Map.of("cpu", "default-flavor")),
+ admissionCheck(
+ "provisioning",
+
PodSetUpdate.builder().name("executor").annotations(annotations).build()));
+
+ Assertions.assertEquals(
+ Map.of("executor", podSetInfoWithMetadata(Map.of(), annotations)),
+ KueueWorkloadUtils.resolvePodSetInfos(kubernetesClient, admitted));
+ }
+
+ @Test
+ void resolvePodSetInfosFailsOnConflictingAdmissionCheckUpdates() {
+ createFlavor("cpu-flavor", Map.of("pool", "cpu"), List.of());
+ Map<String, Map<String, String>> podSetFlavors =
+ Map.of("executor", Map.of("cpu", "cpu-flavor"));
+ // Like Kueue's PodSetInfo.Merge, an update must not change the node
selector of a flavor
+ Workload overFlavor =
+ admittedWorkload(
+ podSetFlavors,
+ admissionCheck(
+ "provisioning",
+ PodSetUpdate.builder()
+ .name("executor")
+ .nodeSelector(Map.of("pool", "provisioned"))
+ .build()));
+ // Nor the labels which an earlier admission check set
+ Workload overCheck =
+ admittedWorkload(
+ podSetFlavors,
+ admissionCheck(
+ "a-check",
+ PodSetUpdate.builder().name("executor").labels(Map.of("team",
"a")).build()),
+ admissionCheck(
+ "b-check",
+ PodSetUpdate.builder().name("executor").labels(Map.of("team",
"b")).build()));
+
+ IllegalArgumentException e =
+ Assertions.assertThrows(
+ IllegalArgumentException.class,
+ () -> KueueWorkloadUtils.resolvePodSetInfos(kubernetesClient,
overFlavor));
+ Assertions.assertTrue(e.getMessage().contains("provisioning"),
e.getMessage());
+ Assertions.assertTrue(e.getMessage().contains("executor"), e.getMessage());
+ Assertions.assertTrue(e.getMessage().contains("pool"), e.getMessage());
+ e =
+ Assertions.assertThrows(
+ IllegalArgumentException.class,
+ () -> KueueWorkloadUtils.resolvePodSetInfos(kubernetesClient,
overCheck));
+ Assertions.assertTrue(e.getMessage().contains("b-check"), e.getMessage());
+ Assertions.assertTrue(e.getMessage().contains("team"), e.getMessage());
+ }
+
+ @Test
+ void admissionCheckUpdateConflictingWithThePodsReleasesTheWorkload() {
+ // Like Kueue's ErrInvalidPodSetUpdate, the conflict is permanent, so the
quota is released
+ // before the caller fails the resource
+ createFlavor("default-flavor", Map.of(), List.of());
+ KueueWorkloadUtils.requestAdmission(
+ kubernetesClient, workloadWithMetadata(Map.of("team", "a"), Map.of()));
+ Workload admitted =
+ admittedWorkload(
+ Map.of("executor", Map.of("cpu", "default-flavor")),
+ admissionCheck(
+ "check",
+ PodSetUpdate.builder().name("executor").labels(Map.of("team",
"b")).build()));
+ admit(admitted.getStatus());
+ SparkAppContext context = context(kubernetesClient);
+
+ IllegalArgumentException e =
+ Assertions.assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ KueueWorkloadUtils.holdForAdmission(
+ context, workloadWithMetadata(Map.of("team", "a"),
Map.of()), "driver"));
+
+ Assertions.assertTrue(e.getMessage().contains("team"), e.getMessage());
+ Assertions.assertNull(getWorkload());
+ verify(context, never()).setKueuePodSetInfos(any());
+ }
+
@Test
@SuppressWarnings("unchecked")
void rejectedAdmissionWithFailedReleaseIsRetried() {
@@ -1070,25 +1240,39 @@ class KueueWorkloadUtilsTest {
Assertions.assertEquals(
Optional.of(ReconcileProgress.completeAndDefaultRequeue()),
KueueWorkloadUtils.holdForAdmission(context, desired, "driver"));
- verify(context, never()).setKueuePodSetFlavors(any());
+ verify(context, never()).setKueuePodSetInfos(any());
}
@Test
void admittedFlavorsAreAppliedWithoutRequestingTheAdmission() {
// A resource whose driver or master exists applies its resources again,
so the flavors of the
- // Workload which was admitted before are resolved again rather than
dropped.
+ // Workload which was admitted before are resolved again rather than
dropped, along with the
+ // updates of its admission checks.
createFlavor("spot-flavor", Map.of("pool", "spot"),
List.of(toleration("spot")));
KueueWorkloadUtils.requestAdmission(kubernetesClient,
workload("owner-uid-1", 1));
- admit(admittedWorkload(Map.of("executor", Map.of("cpu",
"spot-flavor"))).getStatus());
+ Workload admitted =
+ admittedWorkload(
+ Map.of("executor", Map.of("cpu", "spot-flavor")),
+ admissionCheck(
+ "provisioning",
+ PodSetUpdate.builder()
+ .name("executor")
+ .annotations(Map.of("provisioning", "pr-1"))
+ .build()));
+ admit(admitted.getStatus());
SparkAppContext context = context(kubernetesClient);
- Assertions.assertEquals(Optional.empty(),
KueueWorkloadUtils.applyAdmittedFlavors(context));
+ Assertions.assertEquals(Optional.empty(),
KueueWorkloadUtils.applyAdmittedPodSetInfos(context));
verify(context)
- .setKueuePodSetFlavors(
+ .setKueuePodSetInfos(
Map.of(
"executor",
- new KueuePodSetFlavor(Map.of("pool", "spot"),
List.of(toleration("spot")))));
+ new KueuePodSetInfo(
+ Map.of("pool", "spot"),
+ List.of(toleration("spot")),
+ Map.of(),
+ Map.of("provisioning", "pr-1"))));
}
@Test
@@ -1103,11 +1287,12 @@ class KueueWorkloadUtilsTest {
IllegalArgumentException e =
Assertions.assertThrows(
- IllegalArgumentException.class, () ->
KueueWorkloadUtils.applyAdmittedFlavors(context));
+ IllegalArgumentException.class,
+ () -> KueueWorkloadUtils.applyAdmittedPodSetInfos(context));
Assertions.assertTrue(e.getMessage().contains("executor"), e.getMessage());
Assertions.assertTrue(e.getMessage().contains("pool"), e.getMessage());
- verify(context, never()).setKueuePodSetFlavors(any());
+ verify(context, never()).setKueuePodSetInfos(any());
// Unlike the admission, the quota is kept, since the pods it was reserved
for are running
Assertions.assertNotNull(getWorkload());
}
@@ -1117,11 +1302,11 @@ class KueueWorkloadUtilsTest {
SparkAppContext context = context(kubernetesClient);
// The Workload is gone, e.g. evicted and deleted, which must not hold the
running resources
- Assertions.assertEquals(Optional.empty(),
KueueWorkloadUtils.applyAdmittedFlavors(context));
+ Assertions.assertEquals(Optional.empty(),
KueueWorkloadUtils.applyAdmittedPodSetInfos(context));
KueueWorkloadUtils.requestAdmission(kubernetesClient,
workload("owner-uid-1", 1));
- Assertions.assertEquals(Optional.empty(),
KueueWorkloadUtils.applyAdmittedFlavors(context));
- verify(context, never()).setKueuePodSetFlavors(any());
+ Assertions.assertEquals(Optional.empty(),
KueueWorkloadUtils.applyAdmittedPodSetInfos(context));
+ verify(context, never()).setKueuePodSetInfos(any());
}
@Test
@@ -1133,8 +1318,8 @@ class KueueWorkloadUtilsTest {
// Applying no flavors would drop them from the resources which are
applied again
Assertions.assertEquals(
Optional.of(ReconcileProgress.completeAndDefaultRequeue()),
- KueueWorkloadUtils.applyAdmittedFlavors(context));
- verify(context, never()).setKueuePodSetFlavors(any());
+ KueueWorkloadUtils.applyAdmittedPodSetInfos(context));
+ verify(context, never()).setKueuePodSetInfos(any());
}
@Test
@@ -1155,7 +1340,7 @@ class KueueWorkloadUtilsTest {
Assertions.assertEquals(
Optional.of(ReconcileProgress.completeAndDefaultRequeue()),
- KueueWorkloadUtils.applyAdmittedFlavors(context));
+ KueueWorkloadUtils.applyAdmittedPodSetInfos(context));
ArgumentCaptor<EventRecord> captor =
ArgumentCaptor.forClass(EventRecord.class);
verify(eventRecorder).record(captor.capture());
@@ -1193,7 +1378,9 @@ class KueueWorkloadUtilsTest {
return flavor;
}
- private static Workload admittedWorkload(final Map<String, Map<String,
String>> podSetFlavors) {
+ private static Workload admittedWorkload(
+ final Map<String, Map<String, String>> podSetFlavors,
+ final AdmissionCheckState... admissionChecks) {
Workload workload = workload("owner-uid-1", 1);
WorkloadStatus status = status("Admitted", "True");
status.setAdmission(
@@ -1209,6 +1396,7 @@ class KueueWorkloadUtilsTest {
.build())
.toList())
.build());
+ status.setAdmissionChecks(List.of(admissionChecks));
workload.setStatus(status);
return workload;
}
@@ -1217,6 +1405,16 @@ class KueueWorkloadUtilsTest {
return new Toleration("NoSchedule", key, "Exists", null, null);
}
+ private static AdmissionCheckState admissionCheck(
+ final String name, final PodSetUpdate... podSetUpdates) {
+ return
AdmissionCheckState.builder().name(name).podSetUpdates(List.of(podSetUpdates)).build();
+ }
+
+ private static KueuePodSetInfo podSetInfoWithMetadata(
+ final Map<String, String> labels, final Map<String, String> annotations)
{
+ return new KueuePodSetInfo(Map.of(), List.of(), labels, annotations);
+ }
+
@Test
void workloadPriorityClassLabelTakesPrecedence() {
createWorkloadPriorityClass("high", 1000);
@@ -1684,6 +1882,23 @@ class KueueWorkloadUtilsTest {
return workload;
}
+ private static Workload workloadWithMetadata(
+ final Map<String, String> labels, final Map<String, String> annotations)
{
+ Workload workload = workload("owner-uid-1", 1);
+ workload
+ .getSpec()
+ .getPodSets()
+ .get(0)
+ .setTemplate(
+ new PodTemplateSpecBuilder()
+ .withNewMetadata()
+ .withLabels(labels)
+ .withAnnotations(annotations)
+ .endMetadata()
+ .build());
+ return workload;
+ }
+
private static Workload workloadWithDriver() {
Workload workload = workload("owner-uid-1", 2);
workload
diff --git
a/spark-operator/src/test/java/org/apache/spark/k8s/operator/kueue/v1beta2/WorkloadTest.java
b/spark-operator/src/test/java/org/apache/spark/k8s/operator/kueue/v1beta2/WorkloadTest.java
index 1fbaab7f..46e18ecf 100644
---
a/spark-operator/src/test/java/org/apache/spark/k8s/operator/kueue/v1beta2/WorkloadTest.java
+++
b/spark-operator/src/test/java/org/apache/spark/k8s/operator/kueue/v1beta2/WorkloadTest.java
@@ -22,6 +22,7 @@ package org.apache.spark.k8s.operator.kueue.v1beta2;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertNull;
import static org.junit.jupiter.api.Assertions.assertTrue;
import java.util.List;
@@ -34,6 +35,7 @@ import io.fabric8.kubernetes.api.model.ConditionBuilder;
import io.fabric8.kubernetes.api.model.ObjectMetaBuilder;
import io.fabric8.kubernetes.api.model.PodTemplateSpecBuilder;
import io.fabric8.kubernetes.api.model.Quantity;
+import io.fabric8.kubernetes.api.model.Toleration;
import org.junit.jupiter.api.Test;
import org.apache.spark.k8s.operator.Constants;
@@ -188,6 +190,49 @@ class WorkloadTest {
assertFalse(json.has("podsReady"));
}
+ @Test
+ void testWorkloadStatusAdmissionChecksDeserialization() throws Exception {
+ // As Kueue and the controllers of the admission checks write them, with
unmodeled fields
+ String json =
+ "{\"admissionChecks\":[{\"name\":\"provisioning\",\"state\":\"Ready\","
+ +
"\"lastTransitionTime\":\"2026-10-07T00:00:00Z\",\"message\":\"Provisioned\","
+ +
"\"podSetUpdates\":[{\"name\":\"executor\",\"labels\":{\"team\":\"a\"},"
+ +
"\"annotations\":{\"autoscaling.x-k8s.io/consume-provisioning-request\":\"pr\"},"
+ + "\"nodeSelector\":{\"pool\":\"provisioned\"},"
+ +
"\"tolerations\":[{\"key\":\"provisioned\",\"operator\":\"Exists\","
+ + "\"effect\":\"NoSchedule\"}]}]},"
+ +
"{\"name\":\"another\",\"state\":\"Ready\",\"message\":\"\",\"retryCount\":1}]}";
+
+ WorkloadStatus status = OBJECT_MAPPER.readValue(json,
WorkloadStatus.class);
+
+ assertEquals(2, status.getAdmissionChecks().size());
+ AdmissionCheckState check = status.getAdmissionChecks().get(0);
+ assertEquals("provisioning", check.getName());
+ assertEquals(
+ List.of(
+ PodSetUpdate.builder()
+ .name("executor")
+ .labels(Map.of("team", "a"))
+
.annotations(Map.of("autoscaling.x-k8s.io/consume-provisioning-request", "pr"))
+ .nodeSelector(Map.of("pool", "provisioned"))
+ .tolerations(
+ List.of(new Toleration("NoSchedule", "provisioned",
"Exists", null, null)))
+ .build()),
+ check.getPodSetUpdates());
+ assertEquals("another", status.getAdmissionChecks().get(1).getName());
+ assertEquals(List.of(),
status.getAdmissionChecks().get(1).getPodSetUpdates());
+ }
+
+ @Test
+ void testWorkloadStatusWithoutAdmissionChecksOmitsThem() throws Exception {
+ WorkloadStatus status = OBJECT_MAPPER.readValue("{\"conditions\":[]}",
WorkloadStatus.class);
+
+ assertNull(status.getAdmissionChecks());
+ // The operator records a condition on a Workload without status as well,
which must not write
+ // the admission checks which Kueue owns
+ assertFalse(OBJECT_MAPPER.valueToTree(new
WorkloadStatus()).has("admissionChecks"));
+ }
+
@Test
void testInitSpecAndStatus() {
Workload workload = new Workload();
diff --git
a/spark-operator/src/test/java/org/apache/spark/k8s/operator/reconciler/SparkAppResourceSpecFactoryTest.java
b/spark-operator/src/test/java/org/apache/spark/k8s/operator/reconciler/SparkAppResourceSpecFactoryTest.java
index d5dbebeb..c46abdb9 100644
---
a/spark-operator/src/test/java/org/apache/spark/k8s/operator/reconciler/SparkAppResourceSpecFactoryTest.java
+++
b/spark-operator/src/test/java/org/apache/spark/k8s/operator/reconciler/SparkAppResourceSpecFactoryTest.java
@@ -49,7 +49,7 @@ import org.mockito.ArgumentCaptor;
import org.apache.spark.k8s.operator.SparkAppResourceSpec;
import org.apache.spark.k8s.operator.SparkAppSubmissionWorker;
import org.apache.spark.k8s.operator.SparkApplication;
-import org.apache.spark.k8s.operator.kueue.KueuePodSetFlavor;
+import org.apache.spark.k8s.operator.kueue.KueuePodSetInfo;
import org.apache.spark.k8s.operator.spec.BaseApplicationTemplateSpec;
import org.apache.spark.k8s.operator.utils.ModelUtils;
@@ -95,6 +95,9 @@ class SparkAppResourceSpecFactoryTest {
BaseApplicationTemplateSpec.builder()
.podTemplateSpec(
new PodTemplateSpecBuilder()
+ .withNewMetadata()
+ .withLabels(Map.of("team", "a"))
+ .endMetadata()
.withNewSpec()
.withNodeSelector(Map.of("zone", "a"))
.withTolerations(userToleration)
@@ -132,17 +135,35 @@ class SparkAppResourceSpecFactoryTest {
mockClient,
mockWorker,
Map.of(
- "driver", new KueuePodSetFlavor(Map.of("pool", "cpu"),
List.of(spot)),
- "executor", new KueuePodSetFlavor(Map.of("pool", "gpu"),
List.of(gpu))));
+ "driver",
+ new KueuePodSetInfo(
+ Map.of("pool", "cpu"),
+ List.of(spot),
+ Map.of("tier", "batch"),
+ Map.of("provisioning", "pr-driver")),
+ "executor",
+ new KueuePodSetInfo(
+ Map.of("pool", "gpu"),
+ List.of(gpu),
+ Map.of("tier", "batch"),
+ Map.of("provisioning", "pr-executor"))));
// A driver without the pod template gets a new one
Pod driverTemplate = templates.get(DRIVER_SPARK_TEMPLATE_FILE_PROP_KEY);
assertEquals(Map.of("pool", "cpu"),
driverTemplate.getSpec().getNodeSelector());
assertEquals(List.of(spot), driverTemplate.getSpec().getTolerations());
+ assertEquals(Map.of("tier", "batch"),
driverTemplate.getMetadata().getLabels());
+ assertEquals(
+ Map.of("provisioning", "pr-driver"),
driverTemplate.getMetadata().getAnnotations());
// The flavor is merged into the executor pod template of the spec
Pod executorTemplate =
templates.get(EXECUTOR_SPARK_TEMPLATE_FILE_PROP_KEY);
assertEquals(Map.of("zone", "a", "pool", "gpu"),
executorTemplate.getSpec().getNodeSelector());
assertEquals(List.of(userToleration, gpu),
executorTemplate.getSpec().getTolerations());
+ // Spark keeps the labels and annotations of the pod templates on the
driver and executor pods
+ assertEquals(
+ Map.of("team", "a", "tier", "batch"),
executorTemplate.getMetadata().getLabels());
+ assertEquals(
+ Map.of("provisioning", "pr-executor"),
executorTemplate.getMetadata().getAnnotations());
// The temp files are cleaned up, and the SparkApplication is not modified
templateFiles.values().forEach(path ->
assertFalse(Files.exists(Path.of(path)), path));
assertEquals(originalSpec,
ModelUtils.objectMapper.writeValueAsString(app.getSpec()));
diff --git
a/spark-operator/src/test/java/org/apache/spark/k8s/operator/reconciler/reconcilesteps/AppInitStepTest.java
b/spark-operator/src/test/java/org/apache/spark/k8s/operator/reconciler/reconcilesteps/AppInitStepTest.java
index ca185ff3..03b05e44 100644
---
a/spark-operator/src/test/java/org/apache/spark/k8s/operator/reconciler/reconcilesteps/AppInitStepTest.java
+++
b/spark-operator/src/test/java/org/apache/spark/k8s/operator/reconciler/reconcilesteps/AppInitStepTest.java
@@ -81,7 +81,7 @@ import org.apache.spark.k8s.operator.SparkAppSubmissionWorker;
import org.apache.spark.k8s.operator.SparkApplication;
import org.apache.spark.k8s.operator.config.SparkOperatorConf;
import org.apache.spark.k8s.operator.context.SparkAppContext;
-import org.apache.spark.k8s.operator.kueue.KueuePodSetFlavor;
+import org.apache.spark.k8s.operator.kueue.KueuePodSetInfo;
import org.apache.spark.k8s.operator.kueue.KueueWorkloadFactory;
import org.apache.spark.k8s.operator.kueue.KueueWorkloadUtils;
import org.apache.spark.k8s.operator.kueue.v1beta2.Admission;
@@ -1377,7 +1377,7 @@ class AppInitStepTest {
Assertions.assertEquals(ReconcileProgress.completeAndDefaultRequeue(),
progress);
Assertions.assertNull(getWorkload());
- verify(mockContext, never()).setKueuePodSetFlavors(any());
+ verify(mockContext, never()).setKueuePodSetInfos(any());
Assertions.assertEquals(
ApplicationStateSummary.DriverRequested,
application.getStatus().getCurrentState().getCurrentStateSummary());
@@ -1511,7 +1511,7 @@ class AppInitStepTest {
ApplicationStateSummary.DriverRequested,
application.getStatus().getCurrentState().getCurrentStateSummary());
// Like a driver found in the cache, it gets the flavors of the admitted
Workload again
- verify(mockContext).setKueuePodSetFlavors(any());
+ verify(mockContext).setKueuePodSetInfos(any());
// The driver found by the suspend check is not looked up again for the
Kueue admission
verify(mockContext).getCurrentAttemptDriverPod();
}
@@ -1829,9 +1829,9 @@ class AppInitStepTest {
"executor", Map.of("cpu", "spot-flavor", "memory",
"spot-flavor")));
appInitStep.reconcile(mockContext, recorder);
- KueuePodSetFlavor flavor = new KueuePodSetFlavor(Map.of("pool", "spot"),
List.of(spot));
+ KueuePodSetInfo flavor = new KueuePodSetInfo(Map.of("pool", "spot"),
List.of(spot));
InOrder inOrder = inOrder(mockContext);
- inOrder.verify(mockContext).setKueuePodSetFlavors(Map.of("driver", flavor,
"executor", flavor));
+ inOrder.verify(mockContext).setKueuePodSetInfos(Map.of("driver", flavor,
"executor", flavor));
inOrder.verify(mockContext).getDriverPreResourcesSpec();
Assertions.assertNotNull(
kubernetesClient.pods().inNamespace("default").withName("driver-pod").get());
@@ -1858,7 +1858,7 @@ class AppInitStepTest {
// Like Kueue, the conflict is permanent, so the quota is released
Assertions.assertEquals(ReconcileProgress.completeAndImmediateRequeue(),
progress);
Assertions.assertNull(getWorkload());
- verify(mockContext, never()).setKueuePodSetFlavors(any());
+ verify(mockContext, never()).setKueuePodSetInfos(any());
ArgumentCaptor<ApplicationStatus> captor =
ArgumentCaptor.forClass(ApplicationStatus.class);
verify(recorder).persistStatus(any(), captor.capture(), any());
Assertions.assertEquals(
@@ -1888,7 +1888,7 @@ class AppInitStepTest {
// Like a persistent admission failure, it is retried with the default
interval
Assertions.assertEquals(ReconcileProgress.completeAndDefaultRequeue(),
progress);
Assertions.assertNotNull(getWorkload());
- verify(mockContext, never()).setKueuePodSetFlavors(any());
+ verify(mockContext, never()).setKueuePodSetInfos(any());
verifyNoInteractions(recorder);
EventRecord event = captureEvents(2).get(1);
Assertions.assertEquals(EventType.WARNING, event.type());
@@ -1933,8 +1933,8 @@ class AppInitStepTest {
InOrder inOrder = inOrder(mockContext);
inOrder
.verify(mockContext)
- .setKueuePodSetFlavors(
- Map.of("driver", new KueuePodSetFlavor(Map.of("pool", "spot"),
List.of(spot))));
+ .setKueuePodSetInfos(
+ Map.of("driver", new KueuePodSetInfo(Map.of("pool", "spot"),
List.of(spot))));
inOrder.verify(mockContext).getDriverPreResourcesSpec();
Assertions.assertNotNull(getWorkload());
}
@@ -1964,7 +1964,7 @@ class AppInitStepTest {
ReconcileProgress progress = appInitStep.reconcile(mockContext, recorder);
Assertions.assertEquals(ReconcileProgress.completeAndDefaultRequeue(),
progress);
- verify(mockContext, never()).setKueuePodSetFlavors(any());
+ verify(mockContext, never()).setKueuePodSetInfos(any());
Assertions.assertNotNull(
kubernetesClient.pods().inNamespace("default").withName("driver-pod").get());
// The admitted Workload is released with the driver, like the one of a
queued application
diff --git
a/spark-operator/src/test/java/org/apache/spark/k8s/operator/reconciler/reconcilesteps/ClusterInitStepTest.java
b/spark-operator/src/test/java/org/apache/spark/k8s/operator/reconciler/reconcilesteps/ClusterInitStepTest.java
index 3fab2f4f..1f136cb6 100644
---
a/spark-operator/src/test/java/org/apache/spark/k8s/operator/reconciler/reconcilesteps/ClusterInitStepTest.java
+++
b/spark-operator/src/test/java/org/apache/spark/k8s/operator/reconciler/reconcilesteps/ClusterInitStepTest.java
@@ -77,7 +77,7 @@ import org.apache.spark.k8s.operator.SparkCluster;
import org.apache.spark.k8s.operator.SparkClusterSubmissionWorker;
import org.apache.spark.k8s.operator.config.SparkOperatorConf;
import org.apache.spark.k8s.operator.context.SparkClusterContext;
-import org.apache.spark.k8s.operator.kueue.KueuePodSetFlavor;
+import org.apache.spark.k8s.operator.kueue.KueuePodSetInfo;
import org.apache.spark.k8s.operator.kueue.KueueWorkloadFactory;
import org.apache.spark.k8s.operator.kueue.KueueWorkloadUtils;
import org.apache.spark.k8s.operator.kueue.v1beta2.Admission;
@@ -607,7 +607,7 @@ class ClusterInitStepTest {
Assertions.assertEquals(ReconcileProgress.completeAndDefaultRequeue(),
progress);
verify(mockClient, never()).resource(any(Workload.class));
- verify(mockContext, never()).setKueuePodSetFlavors(any());
+ verify(mockContext, never()).setKueuePodSetInfos(any());
ArgumentCaptor<ClusterStatus> captor =
ArgumentCaptor.forClass(ClusterStatus.class);
verify(recorder).persistStatus(any(), captor.capture());
Assertions.assertEquals(
@@ -1047,8 +1047,8 @@ class ClusterInitStepTest {
InOrder inOrder = inOrder(mockContext);
inOrder
.verify(mockContext)
- .setKueuePodSetFlavors(
- Map.of("master", new KueuePodSetFlavor(Map.of("pool", "spot"),
List.of(spot))));
+ .setKueuePodSetInfos(
+ Map.of("master", new KueuePodSetInfo(Map.of("pool", "spot"),
List.of(spot))));
inOrder.verify(mockContext).getMasterStatefulSetSpec();
// The admission is not requested again for a master which exists, nor is
the Workload released
verify(mockClient, never()).resource(any(Workload.class));
@@ -1087,7 +1087,7 @@ class ClusterInitStepTest {
// Like Kueue, the conflict is permanent, so the quota is released
Assertions.assertEquals(ReconcileProgress.completeAndImmediateRequeue(),
progress);
Assertions.assertNull(getWorkload());
- verify(mockContext, never()).setKueuePodSetFlavors(any());
+ verify(mockContext, never()).setKueuePodSetInfos(any());
verify(mockContext, never()).getMasterServiceSpec();
ArgumentCaptor<ClusterStatus> captor =
ArgumentCaptor.forClass(ClusterStatus.class);
verify(recorder).persistStatus(any(), captor.capture(), any());
@@ -1118,7 +1118,7 @@ class ClusterInitStepTest {
// The cluster is not failed permanently, the flavors are read again
Assertions.assertEquals(ReconcileProgress.completeAndDefaultRequeue(),
progress);
Assertions.assertNotNull(getWorkload());
- verify(mockContext, never()).setKueuePodSetFlavors(any());
+ verify(mockContext, never()).setKueuePodSetInfos(any());
verify(mockContext, never()).getMasterServiceSpec();
verifyNoInteractions(recorder);
EventRecord event = captureEvents(2).get(1);
diff --git a/tests/e2e/kueue/chainsaw-test.yaml
b/tests/e2e/kueue/chainsaw-test.yaml
index 17981438..4435a009 100644
--- a/tests/e2e/kueue/chainsaw-test.yaml
+++ b/tests/e2e/kueue/chainsaw-test.yaml
@@ -1149,3 +1149,227 @@ spec:
kubectl delete localqueue "$LOCAL_QUEUE" -n default
--ignore-not-found=true --wait=false
kubectl delete clusterqueue "$CLUSTER_QUEUE" --ignore-not-found=true
--wait=false
kubectl delete resourceflavor "$FLAVOR" --ignore-not-found=true
--wait=false
+ - name: pod-set-updates-of-admission-checks-reach-the-pods
+ # Each step uses its own queue objects, so that a step never waits for
Kueue to finish
+ # deleting the objects of another step.
+ bindings:
+ - name: FLAVOR
+ value: spark-e2e-admission-check-flavor
+ - name: CLUSTER_QUEUE
+ value: spark-e2e-admission-check-cluster-queue
+ - name: LOCAL_QUEUE
+ value: spark-e2e-admission-check-queue
+ - name: ADMISSION_CHECK
+ value: spark-e2e-admission-check
+ try:
+ - apply:
+ file: kueue-queues.yaml
+ # An admission check whose controller is this test
+ - apply:
+ resource:
+ apiVersion: kueue.x-k8s.io/v1beta2
+ kind: AdmissionCheck
+ metadata:
+ name: ($ADMISSION_CHECK)
+ spec:
+ controllerName: spark.apache.org/kueue-e2e
+ # Kueue keeps a ClusterQueue inactive while an admission check of it is
not active, which the
+ # controller of the check reports. The quota fits both resources (1 CPU
for each pod).
+ - script:
+ timeout: 30s
+ env:
+ - name: ADMISSION_CHECK
+ value: ($ADMISSION_CHECK)
+ - name: CLUSTER_QUEUE
+ value: ($CLUSTER_QUEUE)
+ content: |
+ set -e
+ kubectl patch admissioncheck "$ADMISSION_CHECK" --subresource=status
--type merge -p "$(
+ jq -nc --arg now "$(date -u +%Y-%m-%dT%H:%M:%SZ)" '{status:
{conditions: [{
+ type: "Active", status: "True", reason: "Active",
lastTransitionTime: $now,
+ message: "The test acts as the controller of the admission
check"}]}}')"
+ kubectl patch clusterqueue "$CLUSTER_QUEUE" --type json -p "$(
+ jq -nc --arg check "$ADMISSION_CHECK" '[
+ {op: "replace", path:
"/spec/resourceGroups/0/flavors/0/resources/0/nominalQuota",
+ value: "4"},
+ {op: "replace", path:
"/spec/resourceGroups/0/flavors/0/resources/1/nominalQuota",
+ value: "8Gi"},
+ {op: "add", path: "/spec/admissionChecksStrategy",
+ value: {admissionChecks: [{name: $check}]}}]')"
+ - apply:
+ file: spark-example-admission-check.yaml
+ - apply:
+ file: spark-cluster-example-admission-check.yaml
+ # Kueue reserves the quota of both Workloads, while they wait for the
admission check
+ - assert:
+ timeout: 60s
+ resource:
+ apiVersion: kueue.x-k8s.io/v1beta2
+ kind: Workload
+ metadata:
+ name: sparkapplication-spark-job-kueue-admission-check-test
+ namespace: default
+ status:
+ (conditions[?type == 'QuotaReserved' && status == 'True'] |
length(@)): 1
+ admissionChecks:
+ - name: ($ADMISSION_CHECK)
+ state: Pending
+ - assert:
+ timeout: 60s
+ resource:
+ apiVersion: kueue.x-k8s.io/v1beta2
+ kind: Workload
+ metadata:
+ name: sparkcluster-spark-cluster-kueue-admission-check-test
+ namespace: default
+ status:
+ (conditions[?type == 'QuotaReserved' && status == 'True'] |
length(@)): 1
+ admissionChecks:
+ - name: ($ADMISSION_CHECK)
+ state: Pending
+ # Like a ProvisioningRequest admission check, the test sets the check
Ready with an update of
+ # each pod set, whose label and annotation name the pod set. The node
selector is the node
+ # label of the flavor, which is no conflict, and the toleration matches no
taint, so the
+ # scheduling is not changed.
+ - script:
+ timeout: 30s
+ env:
+ - name: ADMISSION_CHECK
+ value: ($ADMISSION_CHECK)
+ content: |
+ set -e
+ now=$(date -u +%Y-%m-%dT%H:%M:%SZ)
+ ready() {
+ workload=$1
+ shift
+ kubectl patch workload "$workload" -n default --subresource=status
--type merge -p "$(
+ jq -nc --arg check "$ADMISSION_CHECK" --arg now "$now" '{status:
{admissionChecks: [{
+ name: $check, state: "Ready", lastTransitionTime: $now,
+ message: "Checked by the test",
+ podSetUpdates: [$ARGS.positional[] | {
+ name: .,
+ labels: {"spark.apache.org/kueue-e2e-admission-check": .},
+ annotations: {"spark.apache.org/kueue-e2e-admission-check":
.},
+ nodeSelector: {"kubernetes.io/os": "linux"},
+ tolerations: [{key:
"spark.apache.org/kueue-e2e-admission-check",
+ operator: "Exists", effect: "NoSchedule"}]}]}]}}' --args
"$@")"
+ }
+ ready sparkapplication-spark-job-kueue-admission-check-test driver
executor
+ ready sparkcluster-spark-cluster-kueue-admission-check-test master
worker
+ # The updates reach the driver and the executor pods along with the
flavor, which the executor
+ # pods get through the executor pod template.
+ - assert:
+ timeout: 5m
+ resource:
+ apiVersion: v1
+ kind: Pod
+ metadata:
+ namespace: default
+ labels:
+ spark.operator/spark-app-name:
spark-job-kueue-admission-check-test
+ spark-role: driver
+ spark.apache.org/kueue-e2e-admission-check: driver
+ annotations:
+ spark.apache.org/kueue-e2e-admission-check: driver
+ spec:
+ (tolerations[?key == 'spark.apache.org/kueue-e2e-admission-check']
| length(@)): 1
+ (tolerations[?key == 'spark.apache.org/kueue-e2e'] | length(@)): 1
+ - assert:
+ timeout: 5m
+ resource:
+ apiVersion: v1
+ kind: Pod
+ metadata:
+ namespace: default
+ labels:
+ spark.operator/spark-app-name:
spark-job-kueue-admission-check-test
+ spark-role: executor
+ spark.apache.org/kueue-e2e-admission-check: executor
+ annotations:
+ spark.apache.org/kueue-e2e-admission-check: executor
+ spec:
+ (tolerations[?key == 'spark.apache.org/kueue-e2e-admission-check']
| length(@)): 1
+ (tolerations[?key == 'spark.apache.org/kueue-e2e'] | length(@)): 1
+ # So do the pod templates of the master and worker StatefulSets
+ - assert:
+ timeout: 5m
+ resource:
+ apiVersion: apps/v1
+ kind: StatefulSet
+ metadata:
+ name: spark-cluster-kueue-admission-check-test-master
+ namespace: default
+ spec:
+ template:
+ metadata:
+ labels:
+ spark.apache.org/kueue-e2e-admission-check: master
+ annotations:
+ spark.apache.org/kueue-e2e-admission-check: master
+ spec:
+ (tolerations[?key ==
'spark.apache.org/kueue-e2e-admission-check'] | length(@)): 1
+ (tolerations[?key == 'spark.apache.org/kueue-e2e'] |
length(@)): 1
+ - assert:
+ timeout: 5m
+ resource:
+ apiVersion: apps/v1
+ kind: StatefulSet
+ metadata:
+ name: spark-cluster-kueue-admission-check-test-worker
+ namespace: default
+ spec:
+ template:
+ metadata:
+ labels:
+ spark.apache.org/kueue-e2e-admission-check: worker
+ annotations:
+ spark.apache.org/kueue-e2e-admission-check: worker
+ spec:
+ (tolerations[?key ==
'spark.apache.org/kueue-e2e-admission-check'] | length(@)): 1
+ (tolerations[?key == 'spark.apache.org/kueue-e2e'] |
length(@)): 1
+ catch:
+ - describe:
+ apiVersion: spark.apache.org/v1
+ kind: SparkApplication
+ namespace: default
+ - describe:
+ apiVersion: spark.apache.org/v1
+ kind: SparkCluster
+ namespace: default
+ - describe:
+ apiVersion: kueue.x-k8s.io/v1beta2
+ kind: Workload
+ namespace: default
+ - describe:
+ apiVersion: kueue.x-k8s.io/v1beta2
+ kind: AdmissionCheck
+ - describe:
+ apiVersion: kueue.x-k8s.io/v1beta2
+ kind: ClusterQueue
+ - describe:
+ apiVersion: v1
+ kind: Pod
+ namespace: default
+ selector:
spark.operator/spark-app-name=spark-job-kueue-admission-check-test
+ - events:
+ namespace: default
+ finally:
+ - script:
+ timeout: 120s
+ env:
+ - name: FLAVOR
+ value: ($FLAVOR)
+ - name: CLUSTER_QUEUE
+ value: ($CLUSTER_QUEUE)
+ - name: LOCAL_QUEUE
+ value: ($LOCAL_QUEUE)
+ - name: ADMISSION_CHECK
+ value: ($ADMISSION_CHECK)
+ content: |
+ # Deleted at once, so that their clean-ups are awaited together
+ kubectl delete sparkapplication/spark-job-kueue-admission-check-test
sparkcluster/spark-cluster-kueue-admission-check-test -n default
--ignore-not-found=true
+ # Kueue finalizers may delay the deletion, which later steps do not
depend on.
+ kubectl delete localqueue "$LOCAL_QUEUE" -n default
--ignore-not-found=true --wait=false
+ kubectl delete clusterqueue "$CLUSTER_QUEUE" --ignore-not-found=true
--wait=false
+ kubectl delete resourceflavor "$FLAVOR" --ignore-not-found=true
--wait=false
+ kubectl delete admissioncheck "$ADMISSION_CHECK"
--ignore-not-found=true --wait=false
diff --git a/tests/e2e/kueue/spark-cluster-example-admission-check.yaml
b/tests/e2e/kueue/spark-cluster-example-admission-check.yaml
new file mode 100644
index 00000000..e3591192
--- /dev/null
+++ b/tests/e2e/kueue/spark-cluster-example-admission-check.yaml
@@ -0,0 +1,37 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+apiVersion: spark.apache.org/v1
+kind: SparkCluster
+metadata:
+ name: spark-cluster-kueue-admission-check-test
+ namespace: default
+ labels:
+ kueue.x-k8s.io/queue-name: ($LOCAL_QUEUE)
+spec:
+ runtimeVersions:
+ sparkVersion: "4.2.0"
+ clusterTolerations:
+ instanceConfig:
+ initWorkers: 1
+ minWorkers: 1
+ maxWorkers: 1
+ sparkConf:
+ spark.kubernetes.container.image: "apache/spark:{{SPARK_VERSION}}-scala"
+ spark.master.ui.title: "Spark Cluster Kueue Admission Check E2E Test"
+ spark.master.rest.host: "0.0.0.0"
+ spark.ui.reverseProxy: "true"
diff --git a/tests/e2e/kueue/spark-example-admission-check.yaml
b/tests/e2e/kueue/spark-example-admission-check.yaml
new file mode 100644
index 00000000..b9241ac2
--- /dev/null
+++ b/tests/e2e/kueue/spark-example-admission-check.yaml
@@ -0,0 +1,35 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+apiVersion: spark.apache.org/v1
+kind: SparkApplication
+metadata:
+ name: spark-job-kueue-admission-check-test
+ namespace: default
+ labels:
+ kueue.x-k8s.io/queue-name: ($LOCAL_QUEUE)
+spec:
+ mainClass: "org.apache.spark.examples.SparkPi"
+ jars: "local:///opt/spark/examples/jars/spark-examples.jar"
+ # Runs long enough for the test to check the executor pod before it is
deleted
+ driverArgs: ["10000"]
+ sparkConf:
+ spark.executor.instances: "1"
+ spark.kubernetes.container.image: "apache/spark:{{SPARK_VERSION}}-scala"
+ spark.kubernetes.authenticate.driver.serviceAccountName: "spark"
+ runtimeVersions:
+ sparkVersion: "4.2.0"
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
