|
| 1 | +/* |
| 2 | + * Licensed to the Apache Software Foundation (ASF) under one |
| 3 | + * or more contributor license agreements. See the NOTICE file |
| 4 | + * distributed with this work for additional information |
| 5 | + * regarding copyright ownership. The ASF licenses this file |
| 6 | + * to you under the Apache License, Version 2.0 (the |
| 7 | + * "License"); you may not use this file except in compliance |
| 8 | + * with the License. You may obtain a copy of the License at |
| 9 | + * |
| 10 | + * http://www.apache.org/licenses/LICENSE-2.0 |
| 11 | + * |
| 12 | + * Unless required by applicable law or agreed to in writing, |
| 13 | + * software distributed under the License is distributed on an |
| 14 | + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| 15 | + * KIND, either express or implied. See the License for the |
| 16 | + * specific language governing permissions and limitations |
| 17 | + * under the License. |
| 18 | + */ |
| 19 | + |
| 20 | +package org.apache.spark.k8s.operator.kueue; |
| 21 | + |
| 22 | +import static java.net.HttpURLConnection.HTTP_FORBIDDEN; |
| 23 | + |
| 24 | +import java.util.Comparator; |
| 25 | +import java.util.Map; |
| 26 | + |
| 27 | +import io.fabric8.kubernetes.api.model.PodTemplateSpec; |
| 28 | +import io.fabric8.kubernetes.api.model.scheduling.v1.PriorityClass; |
| 29 | +import io.fabric8.kubernetes.client.KubernetesClient; |
| 30 | +import io.fabric8.kubernetes.client.KubernetesClientException; |
| 31 | +import lombok.extern.slf4j.Slf4j; |
| 32 | + |
| 33 | +import org.apache.spark.k8s.operator.Constants; |
| 34 | +import org.apache.spark.k8s.operator.kueue.v1beta2.PodSet; |
| 35 | +import org.apache.spark.k8s.operator.kueue.v1beta2.PriorityClassRef; |
| 36 | +import org.apache.spark.k8s.operator.kueue.v1beta2.Workload; |
| 37 | +import org.apache.spark.k8s.operator.kueue.v1beta2.WorkloadPriorityClass; |
| 38 | +import org.apache.spark.k8s.operator.kueue.v1beta2.WorkloadSpec; |
| 39 | +import org.apache.spark.k8s.operator.utils.StringUtils; |
| 40 | + |
| 41 | +/** Utilities to resolve the priority of a Kueue Workload. */ |
| 42 | +@Slf4j |
| 43 | +public final class KueueWorkloadPriority { |
| 44 | + |
| 45 | + /** Like Kueue, the priority of a Workload without any priority class. */ |
| 46 | + private static final int DEFAULT_PRIORITY = 0; |
| 47 | + |
| 48 | + /** The API group of the Kubernetes PriorityClass, which a pod template refers to. */ |
| 49 | + private static final String SCHEDULING_API_GROUP = "scheduling.k8s.io"; |
| 50 | + |
| 51 | + private KueueWorkloadPriority() {} |
| 52 | + |
| 53 | + /** |
| 54 | + * Checks whether the priority class of the given Workload is followed and whether Kueue accepts |
| 55 | + * the change. Like Kueue's `classifyWorkloadsForPriorityUpdate`, only a Workload without a |
| 56 | + * priority class or with a WorkloadPriorityClass follows the owner, so the priority of a Workload |
| 57 | + * backed by a Kubernetes PriorityClass is never rewritten, e.g. when the global default changes. |
| 58 | + * Once the quota is reserved, the Workload CEL rules freeze the presence, the group and the kind |
| 59 | + * of the priorityClassRef. The name of a WorkloadPriorityClass stays mutable, which is what Kueue |
| 60 | + * relies on to raise the priority of a Workload waiting for its admission checks. |
| 61 | + */ |
| 62 | + static boolean isPriorityClassChangeAllowed( |
| 63 | + final Workload workload, final PriorityClassRef desired) { |
| 64 | + PriorityClassRef current = workload.getSpec().getPriorityClassRef(); |
| 65 | + if (current != null && !Constants.KUEUE_API_GROUP.equals(current.getGroup())) { |
| 66 | + return false; |
| 67 | + } |
| 68 | + if (workload.getStatus() == null || !workload.getStatus().isQuotaReserved()) { |
| 69 | + return true; |
| 70 | + } |
| 71 | + return current != null |
| 72 | + && desired != null |
| 73 | + && Constants.KUEUE_API_GROUP.equals(desired.getGroup()); |
| 74 | + } |
| 75 | + |
| 76 | + /** |
| 77 | + * Sets the priority of the desired Workload in the same way as Kueue built-in integrations. The |
| 78 | + * WorkloadPriorityClass of the `kueue.x-k8s.io/priority-class` label takes precedence over the |
| 79 | + * PriorityClass of the first pod set which has one, and the global default PriorityClass is used |
| 80 | + * without both. Without any of them, the priority is 0 without a priority class. The priority is |
| 81 | + * left unset if the operator is not allowed to read the cluster-scoped priority classes. |
| 82 | + * |
| 83 | + * @param client The KubernetesClient. |
| 84 | + * @param workload The Workload whose priority is set in place. |
| 85 | + * @throws IllegalStateException if the priority class does not exist. |
| 86 | + */ |
| 87 | + static void setPriority(final KubernetesClient client, final Workload workload) { |
| 88 | + WorkloadSpec spec = workload.getSpec(); |
| 89 | + try { |
| 90 | + Map<String, String> labels = workload.getMetadata().getLabels(); |
| 91 | + String workloadPriorityClassName = |
| 92 | + labels == null ? null : labels.get(Constants.LABEL_WORKLOAD_PRIORITY_CLASS); |
| 93 | + if (StringUtils.isNotEmpty(workloadPriorityClassName)) { |
| 94 | + WorkloadPriorityClass workloadPriorityClass = |
| 95 | + client.resources(WorkloadPriorityClass.class).withName(workloadPriorityClassName).get(); |
| 96 | + if (workloadPriorityClass == null) { |
| 97 | + throw new IllegalStateException( |
| 98 | + "Kueue WorkloadPriorityClass " + workloadPriorityClassName + " is not found."); |
| 99 | + } |
| 100 | + spec.setPriorityClassRef( |
| 101 | + new PriorityClassRef( |
| 102 | + Constants.KUEUE_API_GROUP, "WorkloadPriorityClass", workloadPriorityClassName)); |
| 103 | + spec.setPriority(workloadPriorityClass.getValue()); |
| 104 | + return; |
| 105 | + } |
| 106 | + PriorityClass priorityClass = getPodPriorityClass(client, spec); |
| 107 | + if (priorityClass == null) { |
| 108 | + spec.setPriorityClassRef(null); |
| 109 | + spec.setPriority(DEFAULT_PRIORITY); |
| 110 | + } else { |
| 111 | + spec.setPriorityClassRef( |
| 112 | + new PriorityClassRef( |
| 113 | + SCHEDULING_API_GROUP, "PriorityClass", priorityClass.getMetadata().getName())); |
| 114 | + spec.setPriority(priorityClass.getValue()); |
| 115 | + } |
| 116 | + } catch (KubernetesClientException e) { |
| 117 | + if (e.getCode() != HTTP_FORBIDDEN) { |
| 118 | + throw e; |
| 119 | + } |
| 120 | + log.warn( |
| 121 | + "Requesting the Kueue Workload {} without priority because the operator is not allowed " |
| 122 | + + "to read the priority classes.", |
| 123 | + workload.getMetadata().getName()); |
| 124 | + } |
| 125 | + } |
| 126 | + |
| 127 | + /** |
| 128 | + * Returns the PriorityClass of the first pod set which has one, or the global default |
| 129 | + * PriorityClass. Like Kueue, the lowest one wins if there are more than one global default. |
| 130 | + */ |
| 131 | + private static PriorityClass getPodPriorityClass( |
| 132 | + final KubernetesClient client, final WorkloadSpec spec) { |
| 133 | + for (PodSet podSet : spec.getPodSets()) { |
| 134 | + PodTemplateSpec template = podSet.getTemplate(); |
| 135 | + String name = |
| 136 | + template == null || template.getSpec() == null |
| 137 | + ? null |
| 138 | + : template.getSpec().getPriorityClassName(); |
| 139 | + if (StringUtils.isNotEmpty(name)) { |
| 140 | + PriorityClass priorityClass = |
| 141 | + client.scheduling().v1().priorityClasses().withName(name).get(); |
| 142 | + if (priorityClass == null) { |
| 143 | + throw new IllegalStateException("PriorityClass " + name + " is not found."); |
| 144 | + } |
| 145 | + return priorityClass; |
| 146 | + } |
| 147 | + } |
| 148 | + return client.scheduling().v1().priorityClasses().list().getItems().stream() |
| 149 | + .filter(priorityClass -> Boolean.TRUE.equals(priorityClass.getGlobalDefault())) |
| 150 | + .min(Comparator.comparing(PriorityClass::getValue)) |
| 151 | + .orElse(null); |
| 152 | + } |
| 153 | +} |
0 commit comments