dongjoon-hyun commented on code in PR #823: URL: https://github.com/apache/spark-kubernetes-operator/pull/823#discussion_r4008253849
########## spark-operator/src/main/java/org/apache/spark/k8s/operator/kueue/KueueWorkloadFactory.java: ########## @@ -0,0 +1,438 @@ +/* + * 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; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Locale; +import java.util.Map; + +import io.fabric8.kubernetes.api.model.Container; +import io.fabric8.kubernetes.api.model.ContainerBuilder; +import io.fabric8.kubernetes.api.model.HasMetadata; +import io.fabric8.kubernetes.api.model.ObjectMetaBuilder; +import io.fabric8.kubernetes.api.model.OwnerReference; +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.Quantity; +import io.fabric8.kubernetes.api.model.ResourceRequirements; +import io.fabric8.kubernetes.api.model.ResourceRequirementsBuilder; +import io.fabric8.kubernetes.api.model.apps.StatefulSet; +import lombok.extern.slf4j.Slf4j; + +import org.apache.spark.k8s.operator.Constants; +import org.apache.spark.k8s.operator.SparkApplication; +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.v1beta1.PodSet; +import org.apache.spark.k8s.operator.kueue.v1beta1.Workload; +import org.apache.spark.k8s.operator.kueue.v1beta1.WorkloadSpec; +import org.apache.spark.k8s.operator.reconciler.SparkClusterResourceSpecFactory; +import org.apache.spark.k8s.operator.spec.ApplicationSpec; +import org.apache.spark.k8s.operator.spec.BaseApplicationTemplateSpec; +import org.apache.spark.k8s.operator.utils.ModelUtils; +import org.apache.spark.k8s.operator.utils.ReconcilerUtils; +import org.apache.spark.k8s.operator.utils.StringUtils; +import org.apache.spark.network.util.JavaUtils; + +/** + * Factory for creating Kueue Workload resources from Spark custom resources. + * This factory supports both {@link SparkApplication} (driver and executor pod sets) + * and {@link SparkCluster} (master and worker pod sets). + */ +@Slf4j +@SuppressWarnings("PMD.GodClass") +public final class KueueWorkloadFactory { + + public static final String PODSET_DRIVER = "driver"; + public static final String PODSET_EXECUTOR = "executor"; + public static final String PODSET_MASTER = "master"; + public static final String PODSET_WORKER = "worker"; + + private static final String DEFAULT_CORES = "1"; + private static final String DEFAULT_MEMORY = "1g"; + private static final String DEFAULT_MIN_MEMORY_OVERHEAD = "384m"; + private static final double DEFAULT_MEMORY_OVERHEAD_FACTOR = 0.10; + private static final double NON_JVM_MEMORY_OVERHEAD_FACTOR = 0.40; + private static final int DEFAULT_EXECUTOR_INSTANCES = 2; + + private static final String NODE_SELECTOR_PREFIX = "spark.kubernetes.node.selector."; + + private KueueWorkloadFactory() {} + + /** + * Builds a Kueue Workload from a SparkApplication resource. + * + * @param app The SparkApplication. + * @return The constructed Kueue Workload. + */ + public static Workload buildWorkload(final SparkApplication app) { + ApplicationSpec spec = app.getSpec(); + Map<String, String> sparkConf = spec.getSparkConf(); + if ("true".equalsIgnoreCase(sparkConf.get("spark.dynamicAllocation.enabled"))) { + throw new UnsupportedOperationException( + "Kueue does not support SparkApplication with dynamic allocation " + + "(spark.dynamicAllocation.enabled=true) yet."); + } + // Like Spark's KubernetesClusterManager, `local[*]` runs the driver only without executors. + boolean driverOnly = + sparkConf.getOrDefault("spark.kubernetes.driver.master", "").startsWith("local"); + List<PodSet> podSets = + driverOnly + ? List.of(buildDriverPodSet(app, sparkConf)) + : List.of(buildDriverPodSet(app, sparkConf), buildExecutorPodSet(app, sparkConf)); + return buildWorkload(app, Constants.LABEL_SPARK_APPLICATION_NAME, spec.isSuspend(), podSets); + } + + /** + * Builds a Kueue Workload from a SparkCluster resource. + * + * @param cluster The SparkCluster. + * @return The constructed Kueue Workload. + */ + public static Workload buildWorkload(final SparkCluster cluster) { + // Use the same StatefulSets which the operator creates for the cluster. + SparkClusterResourceSpec resourceSpec = + SparkClusterResourceSpecFactory.buildResourceSpec( + cluster, new SparkClusterSubmissionWorker()); + if (resourceSpec.getHorizontalPodAutoscaler().isPresent()) { + throw new UnsupportedOperationException( + "Kueue does not support SparkCluster with HorizontalPodAutoscaler " + + "(minWorkers < maxWorkers) yet."); + } + List<PodSet> podSets = + List.of( + buildPodSet(PODSET_MASTER, resourceSpec.getMasterStatefulSet()), + buildPodSet(PODSET_WORKER, resourceSpec.getWorkerStatefulSet())); + return buildWorkload( + cluster, Constants.LABEL_SPARK_CLUSTER_NAME, cluster.getSpec().isSuspend(), podSets); + } + + private static Workload buildWorkload( + final HasMetadata owner, + final String nameLabelKey, + final boolean suspend, + final List<PodSet> podSets) { + Map<String, String> labels = new HashMap<>(); + if (owner.getMetadata().getLabels() != null) { + labels.putAll(owner.getMetadata().getLabels()); + } + labels.put(nameLabelKey, owner.getMetadata().getName()); + + OwnerReference ownerReference = ModelUtils.buildOwnerReferenceTo(owner); + ownerReference.setController(true); + + Workload workload = new Workload(); + workload.setMetadata( + new ObjectMetaBuilder() + .withName(getWorkloadName(owner)) + .withNamespace(owner.getMetadata().getNamespace()) + .withLabels(labels) + .withOwnerReferences(ownerReference) + .build()); + workload.setSpec( + WorkloadSpec.builder() + .queueName(getQueueName(owner)) + .active(!suspend) + .podSets(podSets) + .build()); + return workload; + } + + private static PodSet buildPodSet(final String name, final StatefulSet statefulSet) { + return PodSet.builder() + .name(name) + .count(statefulSet.getSpec().getReplicas()) + .template(statefulSet.getSpec().getTemplate()) + .build(); + } + + /** + * Returns the Workload name prefixed with the lower-cased kind of the owner resource, like Kueue + * built-in integrations, to avoid name collisions between SparkApplication and SparkCluster. + */ + private static String getWorkloadName(final HasMetadata resource) { + return resource.getKind().toLowerCase(Locale.ROOT) + "-" + resource.getMetadata().getName(); + } + + /** + * Extracts the Kueue queue name from the resource metadata labels. + * + * @param resource The Kubernetes resource (e.g. SparkApplication or SparkCluster). + * @return The queue name, or null if not found. + */ + public static String getQueueName(final HasMetadata resource) { + if (resource != null + && resource.getMetadata() != null + && resource.getMetadata().getLabels() != null) { + String queue = resource.getMetadata().getLabels().get(Constants.LABEL_QUEUE_NAME); + if (StringUtils.isNotEmpty(queue)) { + return queue; + } + } + return null; + } + + /** + * Checks whether the resource is configured to use Kueue. + * + * @param resource The Kubernetes resource (SparkApplication or SparkCluster). + * @return true if a Kueue queue name is specified. + */ + public static boolean hasQueueName(final HasMetadata resource) { + return StringUtils.isNotEmpty(getQueueName(resource)); + } + + static PodSet buildDriverPodSet( + final SparkApplication app, final Map<String, String> sparkConf) { + ApplicationSpec spec = app.getSpec(); + return buildPodSet( + PODSET_DRIVER, + 1, + spec.getDriverSpec(), + sparkConf, + Constants.DRIVER_SPARK_CONTAINER_PROP_KEY, + calculateDriverMemoryMiB(sparkConf, isNonJvmApp(spec))); + } + + static PodSet buildExecutorPodSet( + final SparkApplication app, final Map<String, String> sparkConf) { + ApplicationSpec spec = app.getSpec(); + return buildPodSet( + PODSET_EXECUTOR, + parseInt(sparkConf.get("spark.executor.instances"), DEFAULT_EXECUTOR_INSTANCES), + spec.getExecutorSpec(), + sparkConf, + Constants.EXECUTOR_SPARK_CONTAINER_PROP_KEY, + calculateExecutorMemoryMiB(sparkConf, isNonJvmApp(spec), isPythonApp(spec))); + } + + /** + * Builds a PodSet for a Spark role (`driver` or `executor`) from the pod template of the + * SparkApplication and the resource configurations in the same way as Spark does. + */ + private static PodSet buildPodSet( + final String role, + final int count, + final BaseApplicationTemplateSpec roleSpec, + final Map<String, String> sparkConf, + final String containerNameKey, + final long memoryMiB) { + PodTemplateSpec templateSpec = + roleSpec != null && roleSpec.getPodTemplateSpec() != null + ? ReconcilerUtils.clone(roleSpec.getPodTemplateSpec()) + : new PodTemplateSpecBuilder().build(); Review Comment: Thank you. Fixed in f83e218. `spark.kubernetes.{driver,executor}.podTemplateFile` without a pod template in the spec is rejected with `UnsupportedOperationException`, next to the dynamic allocation check. The executor key is not checked in driver-only mode. `testBuildWorkloadWithPodTemplateFile` covers the rejection, the driver-only case, and the precedence of the pod template spec. Reading the file in the operator can be a follow-up. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
