dongjoon-hyun commented on code in PR #823: URL: https://github.com/apache/spark-kubernetes-operator/pull/823#discussion_r4007507900
########## spark-operator/src/main/java/org/apache/spark/k8s/operator/kueue/KueueWorkloadFactory.java: ########## @@ -0,0 +1,475 @@ +/* + * 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 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.ClusterSpec; +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). + */ +@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"; + + public static final String DEFAULT_CORES = "1"; + public static final String DEFAULT_MEMORY = "1g"; + public static final String DEFAULT_MIN_MEMORY_OVERHEAD = "384m"; + public static final double DEFAULT_MEMORY_OVERHEAD_FACTOR = 0.10; + public static final double NON_JVM_MEMORY_OVERHEAD_FACTOR = 0.40; + public static final int DEFAULT_EXECUTOR_INSTANCES = 2; + + 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) { + String queueName = getQueueName(app); + ApplicationSpec appSpec = app.getSpec(); + Map<String, String> sparkConf = + appSpec != null && appSpec.getSparkConf() != null + ? appSpec.getSparkConf() + : Map.of(); + + PodSet driverPodSet = buildDriverPodSet(app, sparkConf); + PodSet executorPodSet = buildExecutorPodSet(app, sparkConf); + + List<PodSet> podSets = new ArrayList<>(); + podSets.add(driverPodSet); + podSets.add(executorPodSet); + + boolean active = appSpec == null || !appSpec.isSuspend(); + + Map<String, String> labels = new HashMap<>(); + if (app.getMetadata().getLabels() != null) { + labels.putAll(app.getMetadata().getLabels()); + } + labels.put(Constants.LABEL_SPARK_APPLICATION_NAME, app.getMetadata().getName()); + if (StringUtils.isNotEmpty(queueName)) { + labels.put(Constants.LABEL_QUEUE_NAME, queueName); + } + + OwnerReference ownerReference = ModelUtils.buildOwnerReferenceTo(app); + ownerReference.setController(true); + + Workload workload = new Workload(); + workload.setMetadata( + new ObjectMetaBuilder() + .withName(getWorkloadName(app)) + .withNamespace(app.getMetadata().getNamespace()) + .withLabels(labels) + .withOwnerReferences(ownerReference) + .build()); + + workload.setSpec( + WorkloadSpec.builder() + .queueName(queueName) + .active(active) + .podSets(podSets) + .build()); + + return workload; + } + + /** + * Builds a Kueue Workload from a SparkCluster resource. + * + * @param cluster The SparkCluster. + * @return The constructed Kueue Workload. + */ + public static Workload buildWorkload(final SparkCluster cluster) { + String queueName = getQueueName(cluster); + ClusterSpec clusterSpec = cluster.getSpec(); + + // 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 = new ArrayList<>(); + podSets.add(buildPodSet(PODSET_MASTER, resourceSpec.getMasterStatefulSet())); + podSets.add(buildPodSet(PODSET_WORKER, resourceSpec.getWorkerStatefulSet())); + + boolean active = clusterSpec == null || !clusterSpec.isSuspend(); + + Map<String, String> labels = new HashMap<>(); + if (cluster.getMetadata().getLabels() != null) { + labels.putAll(cluster.getMetadata().getLabels()); + } + labels.put(Constants.LABEL_SPARK_CLUSTER_NAME, cluster.getMetadata().getName()); + if (StringUtils.isNotEmpty(queueName)) { + labels.put(Constants.LABEL_QUEUE_NAME, queueName); + } + + OwnerReference ownerReference = ModelUtils.buildOwnerReferenceTo(cluster); + ownerReference.setController(true); + + Workload workload = new Workload(); + workload.setMetadata( + new ObjectMetaBuilder() + .withName(getWorkloadName(cluster)) + .withNamespace(cluster.getMetadata().getNamespace()) + .withLabels(labels) + .withOwnerReferences(ownerReference) + .build()); + + workload.setSpec( + WorkloadSpec.builder() + .queueName(queueName) + .active(active) + .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)); + } + + /** + * Builds the driver PodSet. + */ + public static PodSet buildDriverPodSet( + final SparkApplication app, final Map<String, String> sparkConf) { + PodTemplateSpec templateSpec = null; + if (app.getSpec() != null + && app.getSpec().getDriverSpec() != null + && app.getSpec().getDriverSpec().getPodTemplateSpec() != null) { + templateSpec = ReconcilerUtils.clone(app.getSpec().getDriverSpec().getPodTemplateSpec()); + } + if (templateSpec == null) { + templateSpec = new PodTemplateSpecBuilder().build(); + } + ensurePodSpec(templateSpec); + + String cpu = + sparkConf.getOrDefault( + "spark.kubernetes.driver.request.cores", + sparkConf.getOrDefault("spark.driver.cores", DEFAULT_CORES)); + long memoryMiB = calculateDriverMemoryMiB(sparkConf, isNonJvmApp(app.getSpec())); + String gpuAmount = sparkConf.get("spark.driver.resource.gpu.amount"); + String gpuVendor = sparkConf.get("spark.driver.resource.gpu.vendor"); + + decorateTemplateResources( Review Comment: Thank you. Fixed in db33ec0. `decorateNodeSelector` applies `spark.kubernetes.node.selector.*` first and `spark.kubernetes.{driver,executor}.node.selector.*` last, like Spark, while keeping the selectors of the pod template. `testBuildWorkloadWithNodeSelector` covers the precedence. ########## spark-operator/src/main/java/org/apache/spark/k8s/operator/kueue/KueueWorkloadFactory.java: ########## @@ -0,0 +1,475 @@ +/* + * 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 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.ClusterSpec; +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). + */ +@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"; + + public static final String DEFAULT_CORES = "1"; + public static final String DEFAULT_MEMORY = "1g"; + public static final String DEFAULT_MIN_MEMORY_OVERHEAD = "384m"; + public static final double DEFAULT_MEMORY_OVERHEAD_FACTOR = 0.10; + public static final double NON_JVM_MEMORY_OVERHEAD_FACTOR = 0.40; + public static final int DEFAULT_EXECUTOR_INSTANCES = 2; + + 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) { + String queueName = getQueueName(app); + ApplicationSpec appSpec = app.getSpec(); + Map<String, String> sparkConf = + appSpec != null && appSpec.getSparkConf() != null + ? appSpec.getSparkConf() + : Map.of(); + + PodSet driverPodSet = buildDriverPodSet(app, sparkConf); + PodSet executorPodSet = buildExecutorPodSet(app, sparkConf); + + List<PodSet> podSets = new ArrayList<>(); + podSets.add(driverPodSet); + podSets.add(executorPodSet); + + boolean active = appSpec == null || !appSpec.isSuspend(); + + Map<String, String> labels = new HashMap<>(); + if (app.getMetadata().getLabels() != null) { + labels.putAll(app.getMetadata().getLabels()); + } + labels.put(Constants.LABEL_SPARK_APPLICATION_NAME, app.getMetadata().getName()); + if (StringUtils.isNotEmpty(queueName)) { + labels.put(Constants.LABEL_QUEUE_NAME, queueName); + } + + OwnerReference ownerReference = ModelUtils.buildOwnerReferenceTo(app); + ownerReference.setController(true); + + Workload workload = new Workload(); + workload.setMetadata( + new ObjectMetaBuilder() + .withName(getWorkloadName(app)) + .withNamespace(app.getMetadata().getNamespace()) + .withLabels(labels) + .withOwnerReferences(ownerReference) + .build()); + + workload.setSpec( + WorkloadSpec.builder() + .queueName(queueName) + .active(active) + .podSets(podSets) + .build()); + + return workload; + } + + /** + * Builds a Kueue Workload from a SparkCluster resource. + * + * @param cluster The SparkCluster. + * @return The constructed Kueue Workload. + */ + public static Workload buildWorkload(final SparkCluster cluster) { + String queueName = getQueueName(cluster); + ClusterSpec clusterSpec = cluster.getSpec(); + + // 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 = new ArrayList<>(); + podSets.add(buildPodSet(PODSET_MASTER, resourceSpec.getMasterStatefulSet())); + podSets.add(buildPodSet(PODSET_WORKER, resourceSpec.getWorkerStatefulSet())); + + boolean active = clusterSpec == null || !clusterSpec.isSuspend(); + + Map<String, String> labels = new HashMap<>(); + if (cluster.getMetadata().getLabels() != null) { + labels.putAll(cluster.getMetadata().getLabels()); + } + labels.put(Constants.LABEL_SPARK_CLUSTER_NAME, cluster.getMetadata().getName()); + if (StringUtils.isNotEmpty(queueName)) { + labels.put(Constants.LABEL_QUEUE_NAME, queueName); + } + + OwnerReference ownerReference = ModelUtils.buildOwnerReferenceTo(cluster); + ownerReference.setController(true); + + Workload workload = new Workload(); + workload.setMetadata( + new ObjectMetaBuilder() + .withName(getWorkloadName(cluster)) + .withNamespace(cluster.getMetadata().getNamespace()) + .withLabels(labels) + .withOwnerReferences(ownerReference) + .build()); + + workload.setSpec( + WorkloadSpec.builder() + .queueName(queueName) + .active(active) + .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)); + } + + /** + * Builds the driver PodSet. + */ + public static PodSet buildDriverPodSet( + final SparkApplication app, final Map<String, String> sparkConf) { + PodTemplateSpec templateSpec = null; + if (app.getSpec() != null + && app.getSpec().getDriverSpec() != null + && app.getSpec().getDriverSpec().getPodTemplateSpec() != null) { + templateSpec = ReconcilerUtils.clone(app.getSpec().getDriverSpec().getPodTemplateSpec()); + } + if (templateSpec == null) { + templateSpec = new PodTemplateSpecBuilder().build(); + } + ensurePodSpec(templateSpec); + + String cpu = + sparkConf.getOrDefault( + "spark.kubernetes.driver.request.cores", + sparkConf.getOrDefault("spark.driver.cores", DEFAULT_CORES)); + long memoryMiB = calculateDriverMemoryMiB(sparkConf, isNonJvmApp(app.getSpec())); + String gpuAmount = sparkConf.get("spark.driver.resource.gpu.amount"); + String gpuVendor = sparkConf.get("spark.driver.resource.gpu.vendor"); + + decorateTemplateResources( + templateSpec, + sparkConf.get(Constants.DRIVER_SPARK_CONTAINER_PROP_KEY), + "spark-driver", + cpu, + memoryMiB, + gpuAmount, + gpuVendor); + + return PodSet.builder() + .name(PODSET_DRIVER) + .count(1) + .template(templateSpec) + .build(); + } + + /** + * Builds the executor PodSet. + */ + public static PodSet buildExecutorPodSet( + final SparkApplication app, final Map<String, String> sparkConf) { + if ("true".equalsIgnoreCase(sparkConf.get("spark.dynamicAllocation.enabled"))) { + throw new UnsupportedOperationException( + "Kueue does not support SparkApplication with dynamic allocation " + + "(spark.dynamicAllocation.enabled=true) yet."); + } + PodTemplateSpec templateSpec = null; + if (app.getSpec() != null + && app.getSpec().getExecutorSpec() != null + && app.getSpec().getExecutorSpec().getPodTemplateSpec() != null) { + templateSpec = ReconcilerUtils.clone(app.getSpec().getExecutorSpec().getPodTemplateSpec()); + } + if (templateSpec == null) { + templateSpec = new PodTemplateSpecBuilder().build(); + } + ensurePodSpec(templateSpec); + + String cpu = + sparkConf.getOrDefault( + "spark.kubernetes.executor.request.cores", + sparkConf.getOrDefault("spark.executor.cores", DEFAULT_CORES)); + long memoryMiB = + calculateExecutorMemoryMiB( + sparkConf, isNonJvmApp(app.getSpec()), isPythonApp(app.getSpec())); + String gpuAmount = sparkConf.get("spark.executor.resource.gpu.amount"); + String gpuVendor = sparkConf.get("spark.executor.resource.gpu.vendor"); + + decorateTemplateResources( + templateSpec, + sparkConf.get("spark.kubernetes.executor.podTemplateContainerName"), + "spark-executor", + cpu, + memoryMiB, + gpuAmount, + gpuVendor); + + return PodSet.builder() + .name(PODSET_EXECUTOR) + .count(parseInt(sparkConf.get("spark.executor.instances"), DEFAULT_EXECUTOR_INSTANCES)) + .template(templateSpec) + .build(); + } + + /** + * Calculates total driver memory in MiB including overhead. + */ + public static long calculateDriverMemoryMiB( + final Map<String, String> sparkConf, final boolean isNonJvm) { + long memMiB = + JavaUtils.byteStringAsMb(sparkConf.getOrDefault("spark.driver.memory", DEFAULT_MEMORY)); + return memMiB + calculateMemoryOverheadMiB(sparkConf, "spark.driver", memMiB, isNonJvm); + } + + /** + * Calculates total executor memory in MiB including overhead, off-heap memory and PySpark + * memory. + */ + public static long calculateExecutorMemoryMiB( + final Map<String, String> sparkConf, final boolean isNonJvm, final boolean isPython) { + long memMiB = + JavaUtils.byteStringAsMb(sparkConf.getOrDefault("spark.executor.memory", DEFAULT_MEMORY)); + long total = + memMiB + calculateMemoryOverheadMiB(sparkConf, "spark.executor", memMiB, isNonJvm); + if ("true".equalsIgnoreCase(sparkConf.get("spark.memory.offHeap.enabled"))) { + // `spark.memory.offHeap.size` is in bytes unless otherwise specified. + total += + JavaUtils.byteStringAsBytes(sparkConf.getOrDefault("spark.memory.offHeap.size", "0")) + / 1024 + / 1024; + } + if (isPython && sparkConf.containsKey("spark.executor.pyspark.memory")) { + total += JavaUtils.byteStringAsMb(sparkConf.get("spark.executor.pyspark.memory")); + } + return total; + } + + private static long calculateMemoryOverheadMiB( + final Map<String, String> sparkConf, + final String prefix, + final long memMiB, + final boolean isNonJvm) { + String overhead = sparkConf.get(prefix + ".memoryOverhead"); + if (StringUtils.isNotEmpty(overhead)) { + return JavaUtils.byteStringAsMb(overhead); + } + // Like Spark's BasicDriverFeatureStep, the deprecated `spark.kubernetes.memoryOverheadFactor` + // or the default factor (0.4 for non-JVM applications) is used if not set explicitly. + double defaultFactor = + parseDouble( + sparkConf.get("spark.kubernetes.memoryOverheadFactor"), + isNonJvm ? NON_JVM_MEMORY_OVERHEAD_FACTOR : DEFAULT_MEMORY_OVERHEAD_FACTOR); + double factor = parseDouble(sparkConf.get(prefix + ".memoryOverheadFactor"), defaultFactor); + long minOverheadMiB = + JavaUtils.byteStringAsMb( + sparkConf.getOrDefault(prefix + ".minMemoryOverhead", DEFAULT_MIN_MEMORY_OVERHEAD)); + return Math.max((int) (factor * memMiB), minOverheadMiB); + } + + /** Follows the main application resource selection of SparkAppSubmissionWorker. */ + private static boolean isPythonApp(final ApplicationSpec spec) { + return StringUtils.isEmpty(spec.getJars()) + && ("org.apache.spark.deploy.PythonRunner".equals(spec.getMainClass()) + || StringUtils.isNotEmpty(spec.getPyFiles())); + } + + private static boolean isNonJvmApp(final ApplicationSpec spec) { + return isPythonApp(spec) + || (StringUtils.isEmpty(spec.getJars()) && StringUtils.isNotEmpty(spec.getSparkRFiles())); + } + + private static void ensurePodSpec(final PodTemplateSpec templateSpec) { + if (templateSpec.getSpec() == null) { + templateSpec.setSpec(new PodSpecBuilder().build()); + } + if (templateSpec.getSpec().getContainers() == null) { + templateSpec.getSpec().setContainers(new ArrayList<>()); + } + } + + private static void decorateTemplateResources( + final PodTemplateSpec templateSpec, + final String containerName, + final String defaultContainerName, + final String cpu, + final long memoryMiB, + final String gpuAmount, + final String gpuVendor) { + PodSpec podSpec = templateSpec.getSpec(); + Container container; + if (podSpec.getContainers().isEmpty()) { + container = new ContainerBuilder().withName(defaultContainerName).build(); + podSpec.getContainers().add(container); + } else { + // Like Spark's KubernetesUtils.selectSparkContainer, select the container by name and + // fall back to the first container. + container = + podSpec.getContainers().stream() + .filter(c -> containerName != null && containerName.equals(c.getName())) + .findFirst() + .orElse(podSpec.getContainers().get(0)); + } + + ResourceRequirements resources = container.getResources(); + if (resources == null) { + resources = new ResourceRequirementsBuilder().build(); + container.setResources(resources); + } + + Map<String, Quantity> requests = resources.getRequests(); + if (requests == null) { + requests = new HashMap<>(); + resources.setRequests(requests); + } + + // Like Spark, overwrite the requests of the pod template. + requests.put("cpu", new Quantity(cpu)); + requests.put("memory", new Quantity(memoryMiB + "Mi")); + + if (StringUtils.isNotEmpty(gpuAmount)) { + if (StringUtils.isEmpty(gpuVendor)) { + throw new IllegalArgumentException( + "Resource: gpu was requested, but vendor was not specified."); + } + // Like Spark's KubernetesConf.buildKubernetesResourceName, e.g., `nvidia.com/gpu`. + String gpuResourceName = gpuVendor + "/gpu"; Review Comment: Fixed. Every `spark.{driver,executor}.resource.<name>.amount` is converted into `<vendor>/<name>` requests and limits, and the `gpu`-specific reads in the callers are gone. The test now covers `xilinx.com/fpga` together with `nvidia.com/gpu`. ########## spark-operator/src/main/java/org/apache/spark/k8s/operator/kueue/KueueWorkloadFactory.java: ########## @@ -0,0 +1,475 @@ +/* + * 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 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.ClusterSpec; +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). + */ +@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"; + + public static final String DEFAULT_CORES = "1"; + public static final String DEFAULT_MEMORY = "1g"; + public static final String DEFAULT_MIN_MEMORY_OVERHEAD = "384m"; + public static final double DEFAULT_MEMORY_OVERHEAD_FACTOR = 0.10; + public static final double NON_JVM_MEMORY_OVERHEAD_FACTOR = 0.40; + public static final int DEFAULT_EXECUTOR_INSTANCES = 2; + + 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) { + String queueName = getQueueName(app); + ApplicationSpec appSpec = app.getSpec(); + Map<String, String> sparkConf = + appSpec != null && appSpec.getSparkConf() != null + ? appSpec.getSparkConf() + : Map.of(); + + PodSet driverPodSet = buildDriverPodSet(app, sparkConf); + PodSet executorPodSet = buildExecutorPodSet(app, sparkConf); + + List<PodSet> podSets = new ArrayList<>(); + podSets.add(driverPodSet); + podSets.add(executorPodSet); + + boolean active = appSpec == null || !appSpec.isSuspend(); + + Map<String, String> labels = new HashMap<>(); + if (app.getMetadata().getLabels() != null) { + labels.putAll(app.getMetadata().getLabels()); + } + labels.put(Constants.LABEL_SPARK_APPLICATION_NAME, app.getMetadata().getName()); + if (StringUtils.isNotEmpty(queueName)) { + labels.put(Constants.LABEL_QUEUE_NAME, queueName); + } + + OwnerReference ownerReference = ModelUtils.buildOwnerReferenceTo(app); + ownerReference.setController(true); + + Workload workload = new Workload(); + workload.setMetadata( + new ObjectMetaBuilder() + .withName(getWorkloadName(app)) + .withNamespace(app.getMetadata().getNamespace()) + .withLabels(labels) + .withOwnerReferences(ownerReference) + .build()); + + workload.setSpec( + WorkloadSpec.builder() + .queueName(queueName) + .active(active) + .podSets(podSets) + .build()); + + return workload; + } + + /** + * Builds a Kueue Workload from a SparkCluster resource. + * + * @param cluster The SparkCluster. + * @return The constructed Kueue Workload. + */ + public static Workload buildWorkload(final SparkCluster cluster) { + String queueName = getQueueName(cluster); + ClusterSpec clusterSpec = cluster.getSpec(); + + // 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 = new ArrayList<>(); + podSets.add(buildPodSet(PODSET_MASTER, resourceSpec.getMasterStatefulSet())); + podSets.add(buildPodSet(PODSET_WORKER, resourceSpec.getWorkerStatefulSet())); + + boolean active = clusterSpec == null || !clusterSpec.isSuspend(); + + Map<String, String> labels = new HashMap<>(); + if (cluster.getMetadata().getLabels() != null) { + labels.putAll(cluster.getMetadata().getLabels()); + } + labels.put(Constants.LABEL_SPARK_CLUSTER_NAME, cluster.getMetadata().getName()); + if (StringUtils.isNotEmpty(queueName)) { + labels.put(Constants.LABEL_QUEUE_NAME, queueName); + } + + OwnerReference ownerReference = ModelUtils.buildOwnerReferenceTo(cluster); + ownerReference.setController(true); + + Workload workload = new Workload(); + workload.setMetadata( + new ObjectMetaBuilder() + .withName(getWorkloadName(cluster)) + .withNamespace(cluster.getMetadata().getNamespace()) + .withLabels(labels) + .withOwnerReferences(ownerReference) + .build()); + + workload.setSpec( + WorkloadSpec.builder() + .queueName(queueName) + .active(active) + .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)); + } + + /** + * Builds the driver PodSet. + */ + public static PodSet buildDriverPodSet( + final SparkApplication app, final Map<String, String> sparkConf) { + PodTemplateSpec templateSpec = null; + if (app.getSpec() != null + && app.getSpec().getDriverSpec() != null + && app.getSpec().getDriverSpec().getPodTemplateSpec() != null) { + templateSpec = ReconcilerUtils.clone(app.getSpec().getDriverSpec().getPodTemplateSpec()); + } + if (templateSpec == null) { + templateSpec = new PodTemplateSpecBuilder().build(); + } + ensurePodSpec(templateSpec); + + String cpu = + sparkConf.getOrDefault( + "spark.kubernetes.driver.request.cores", + sparkConf.getOrDefault("spark.driver.cores", DEFAULT_CORES)); + long memoryMiB = calculateDriverMemoryMiB(sparkConf, isNonJvmApp(app.getSpec())); + String gpuAmount = sparkConf.get("spark.driver.resource.gpu.amount"); + String gpuVendor = sparkConf.get("spark.driver.resource.gpu.vendor"); + + decorateTemplateResources( + templateSpec, + sparkConf.get(Constants.DRIVER_SPARK_CONTAINER_PROP_KEY), + "spark-driver", + cpu, + memoryMiB, + gpuAmount, + gpuVendor); + + return PodSet.builder() + .name(PODSET_DRIVER) + .count(1) + .template(templateSpec) + .build(); + } + + /** + * Builds the executor PodSet. + */ + public static PodSet buildExecutorPodSet( + final SparkApplication app, final Map<String, String> sparkConf) { + if ("true".equalsIgnoreCase(sparkConf.get("spark.dynamicAllocation.enabled"))) { + throw new UnsupportedOperationException( + "Kueue does not support SparkApplication with dynamic allocation " + + "(spark.dynamicAllocation.enabled=true) yet."); + } + PodTemplateSpec templateSpec = null; + if (app.getSpec() != null + && app.getSpec().getExecutorSpec() != null + && app.getSpec().getExecutorSpec().getPodTemplateSpec() != null) { + templateSpec = ReconcilerUtils.clone(app.getSpec().getExecutorSpec().getPodTemplateSpec()); + } + if (templateSpec == null) { + templateSpec = new PodTemplateSpecBuilder().build(); + } + ensurePodSpec(templateSpec); + + String cpu = + sparkConf.getOrDefault( + "spark.kubernetes.executor.request.cores", + sparkConf.getOrDefault("spark.executor.cores", DEFAULT_CORES)); + long memoryMiB = + calculateExecutorMemoryMiB( + sparkConf, isNonJvmApp(app.getSpec()), isPythonApp(app.getSpec())); + String gpuAmount = sparkConf.get("spark.executor.resource.gpu.amount"); + String gpuVendor = sparkConf.get("spark.executor.resource.gpu.vendor"); + + decorateTemplateResources( + templateSpec, + sparkConf.get("spark.kubernetes.executor.podTemplateContainerName"), + "spark-executor", + cpu, + memoryMiB, + gpuAmount, + gpuVendor); + + return PodSet.builder() + .name(PODSET_EXECUTOR) + .count(parseInt(sparkConf.get("spark.executor.instances"), DEFAULT_EXECUTOR_INSTANCES)) + .template(templateSpec) + .build(); + } + + /** + * Calculates total driver memory in MiB including overhead. + */ + public static long calculateDriverMemoryMiB( + final Map<String, String> sparkConf, final boolean isNonJvm) { + long memMiB = + JavaUtils.byteStringAsMb(sparkConf.getOrDefault("spark.driver.memory", DEFAULT_MEMORY)); + return memMiB + calculateMemoryOverheadMiB(sparkConf, "spark.driver", memMiB, isNonJvm); + } + + /** + * Calculates total executor memory in MiB including overhead, off-heap memory and PySpark + * memory. + */ + public static long calculateExecutorMemoryMiB( + final Map<String, String> sparkConf, final boolean isNonJvm, final boolean isPython) { + long memMiB = + JavaUtils.byteStringAsMb(sparkConf.getOrDefault("spark.executor.memory", DEFAULT_MEMORY)); + long total = + memMiB + calculateMemoryOverheadMiB(sparkConf, "spark.executor", memMiB, isNonJvm); + if ("true".equalsIgnoreCase(sparkConf.get("spark.memory.offHeap.enabled"))) { + // `spark.memory.offHeap.size` is in bytes unless otherwise specified. + total += + JavaUtils.byteStringAsBytes(sparkConf.getOrDefault("spark.memory.offHeap.size", "0")) + / 1024 + / 1024; + } + if (isPython && sparkConf.containsKey("spark.executor.pyspark.memory")) { + total += JavaUtils.byteStringAsMb(sparkConf.get("spark.executor.pyspark.memory")); + } + return total; + } + + private static long calculateMemoryOverheadMiB( + final Map<String, String> sparkConf, + final String prefix, + final long memMiB, + final boolean isNonJvm) { + String overhead = sparkConf.get(prefix + ".memoryOverhead"); + if (StringUtils.isNotEmpty(overhead)) { + return JavaUtils.byteStringAsMb(overhead); + } + // Like Spark's BasicDriverFeatureStep, the deprecated `spark.kubernetes.memoryOverheadFactor` + // or the default factor (0.4 for non-JVM applications) is used if not set explicitly. + double defaultFactor = + parseDouble( + sparkConf.get("spark.kubernetes.memoryOverheadFactor"), + isNonJvm ? NON_JVM_MEMORY_OVERHEAD_FACTOR : DEFAULT_MEMORY_OVERHEAD_FACTOR); + double factor = parseDouble(sparkConf.get(prefix + ".memoryOverheadFactor"), defaultFactor); + long minOverheadMiB = + JavaUtils.byteStringAsMb( + sparkConf.getOrDefault(prefix + ".minMemoryOverhead", DEFAULT_MIN_MEMORY_OVERHEAD)); + return Math.max((int) (factor * memMiB), minOverheadMiB); + } + + /** Follows the main application resource selection of SparkAppSubmissionWorker. */ + private static boolean isPythonApp(final ApplicationSpec spec) { + return StringUtils.isEmpty(spec.getJars()) + && ("org.apache.spark.deploy.PythonRunner".equals(spec.getMainClass()) + || StringUtils.isNotEmpty(spec.getPyFiles())); + } + + private static boolean isNonJvmApp(final ApplicationSpec spec) { + return isPythonApp(spec) + || (StringUtils.isEmpty(spec.getJars()) && StringUtils.isNotEmpty(spec.getSparkRFiles())); + } + + private static void ensurePodSpec(final PodTemplateSpec templateSpec) { + if (templateSpec.getSpec() == null) { + templateSpec.setSpec(new PodSpecBuilder().build()); + } + if (templateSpec.getSpec().getContainers() == null) { + templateSpec.getSpec().setContainers(new ArrayList<>()); + } + } + + private static void decorateTemplateResources( + final PodTemplateSpec templateSpec, + final String containerName, + final String defaultContainerName, + final String cpu, + final long memoryMiB, + final String gpuAmount, + final String gpuVendor) { + PodSpec podSpec = templateSpec.getSpec(); + Container container; + if (podSpec.getContainers().isEmpty()) { + container = new ContainerBuilder().withName(defaultContainerName).build(); + podSpec.getContainers().add(container); + } else { + // Like Spark's KubernetesUtils.selectSparkContainer, select the container by name and + // fall back to the first container. + container = + podSpec.getContainers().stream() + .filter(c -> containerName != null && containerName.equals(c.getName())) + .findFirst() + .orElse(podSpec.getContainers().get(0)); + } + + ResourceRequirements resources = container.getResources(); + if (resources == null) { + resources = new ResourceRequirementsBuilder().build(); + container.setResources(resources); + } + + Map<String, Quantity> requests = resources.getRequests(); + if (requests == null) { + requests = new HashMap<>(); + resources.setRequests(requests); + } + + // Like Spark, overwrite the requests of the pod template. + requests.put("cpu", new Quantity(cpu)); + requests.put("memory", new Quantity(memoryMiB + "Mi")); + + if (StringUtils.isNotEmpty(gpuAmount)) { + if (StringUtils.isEmpty(gpuVendor)) { + throw new IllegalArgumentException( + "Resource: gpu was requested, but vendor was not specified."); + } + // Like Spark's KubernetesConf.buildKubernetesResourceName, e.g., `nvidia.com/gpu`. + String gpuResourceName = gpuVendor + "/gpu"; + requests.put(gpuResourceName, new Quantity(gpuAmount)); + Map<String, Quantity> limits = resources.getLimits(); + if (limits == null) { + limits = new HashMap<>(); + resources.setLimits(limits); + } + limits.put(gpuResourceName, new Quantity(gpuAmount)); + } + } + + private static int parseInt(final String str, final int defaultValue) { Review Comment: Fixed. `parseInt` / `parseDouble` let `NumberFormatException` propagate now, and only an empty value falls back to the default. `testBuildWorkloadWithMalformedNumbers` covers both `spark.executor.instances` and `spark.driver.memoryOverheadFactor`. In the same spirit, `spark.memory.offHeap.enabled=true` without a positive `spark.memory.offHeap.size` is rejected like Spark's `Utils.checkOffHeapEnabled`. ########## spark-operator/src/test/java/org/apache/spark/k8s/operator/kueue/KueueWorkloadFactoryTest.java: ########## @@ -0,0 +1,497 @@ +/* + * 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 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.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.util.HashMap; +import java.util.Map; + +import io.fabric8.kubernetes.api.model.Container; +import io.fabric8.kubernetes.api.model.ContainerBuilder; +import io.fabric8.kubernetes.api.model.ObjectMetaBuilder; +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.ResourceRequirementsBuilder; +import io.fabric8.kubernetes.api.model.apps.StatefulSetSpecBuilder; +import org.junit.jupiter.api.Test; + +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.kueue.v1beta1.PodSet; +import org.apache.spark.k8s.operator.kueue.v1beta1.Workload; +import org.apache.spark.k8s.operator.spec.ApplicationSpec; +import org.apache.spark.k8s.operator.spec.BaseApplicationTemplateSpec; +import org.apache.spark.k8s.operator.spec.ClusterSpec; +import org.apache.spark.k8s.operator.spec.ClusterTolerations; +import org.apache.spark.k8s.operator.spec.MasterSpec; +import org.apache.spark.k8s.operator.spec.RuntimeVersions; +import org.apache.spark.k8s.operator.spec.WorkerInstanceConfig; +import org.apache.spark.k8s.operator.spec.WorkerSpec; + +class KueueWorkloadFactoryTest { + + @Test + void testCalculateDriverMemoryMiB() { + // 1g memory -> 1024MiB. Overhead factor 0.10 -> 102MiB -> min 384MiB. + // Total = 1024 + 384 = 1408 + Map<String, String> conf = new HashMap<>(); + conf.put("spark.driver.memory", "1g"); + assertEquals(1408L, KueueWorkloadFactory.calculateDriverMemoryMiB(conf, false)); + + // Explicit overhead 512m -> 1024 + 512 = 1536 + conf.put("spark.driver.memoryOverhead", "512m"); + assertEquals(1536L, KueueWorkloadFactory.calculateDriverMemoryMiB(conf, false)); + + // Custom overhead factor 0.50 -> 1024 * 0.50 = 512 -> 1024 + 512 = 1536 + conf.remove("spark.driver.memoryOverhead"); + conf.put("spark.driver.memoryOverheadFactor", "0.50"); + assertEquals(1536L, KueueWorkloadFactory.calculateDriverMemoryMiB(conf, false)); + + // Custom minimum overhead 1g -> 1024 + max(512, 1024) = 2048 + conf.put("spark.driver.minMemoryOverhead", "1g"); + assertEquals(2048L, KueueWorkloadFactory.calculateDriverMemoryMiB(conf, false)); + } + + @Test + void testCalculateDriverMemoryMiBForNonJvmApp() { + Map<String, String> conf = new HashMap<>(); + conf.put("spark.driver.memory", "8g"); + // Non-JVM: 8192 + 0.40 * 8192 = 8192 + 3276 = 11468 + assertEquals(11468L, KueueWorkloadFactory.calculateDriverMemoryMiB(conf, true)); + // JVM: 8192 + 0.10 * 8192 = 8192 + 819 = 9011 + assertEquals(9011L, KueueWorkloadFactory.calculateDriverMemoryMiB(conf, false)); + + // The deprecated `spark.kubernetes.memoryOverheadFactor` 0.20 -> 8192 + 1638 = 9830 + conf.put("spark.kubernetes.memoryOverheadFactor", "0.20"); + assertEquals(9830L, KueueWorkloadFactory.calculateDriverMemoryMiB(conf, true)); + assertEquals(9830L, KueueWorkloadFactory.calculateDriverMemoryMiB(conf, false)); + } + + @Test + void testCalculateExecutorMemoryMiB() { + Map<String, String> conf = new HashMap<>(); + // 2048MiB. Overhead: max(204, 384) = 384. Total = 2432 + conf.put("spark.executor.memory", "2g"); + assertEquals(2432L, KueueWorkloadFactory.calculateExecutorMemoryMiB(conf, false, false)); + + // PySpark memory is added only for Python applications + conf.put("spark.executor.pyspark.memory", "1g"); + assertEquals(2432L, KueueWorkloadFactory.calculateExecutorMemoryMiB(conf, false, false)); + // Non-JVM overhead: max(819, 384) = 819. Total = 2048 + 819 + 1024 = 3891 + assertEquals(3891L, KueueWorkloadFactory.calculateExecutorMemoryMiB(conf, true, true)); + } + + @Test + void testCalculateExecutorMemoryMiBWithOffHeap() { + Map<String, String> conf = new HashMap<>(); + conf.put("spark.executor.memory", "4g"); + conf.put("spark.memory.offHeap.size", "4g"); + // Off-heap memory is ignored if disabled. 4096 + 409 = 4505 + assertEquals(4505L, KueueWorkloadFactory.calculateExecutorMemoryMiB(conf, false, false)); + + // 4096 + 409 + 4096 = 8601 + conf.put("spark.memory.offHeap.enabled", "true"); + assertEquals(8601L, KueueWorkloadFactory.calculateExecutorMemoryMiB(conf, false, false)); + + // `spark.memory.offHeap.size` is in bytes unless otherwise specified + conf.put("spark.memory.offHeap.size", "4294967296"); + assertEquals(8601L, KueueWorkloadFactory.calculateExecutorMemoryMiB(conf, false, false)); + } + + @Test + void testCalculateMemoryWithUnits() { + // 2048 + 384 = 2432 + assertEquals( + 2432L, + KueueWorkloadFactory.calculateExecutorMemoryMiB( + Map.of("spark.executor.memory", "2Gi"), false, false)); + assertThrows( + NumberFormatException.class, + () -> + KueueWorkloadFactory.calculateDriverMemoryMiB( + Map.of("spark.driver.memory", "10zz"), false)); + assertThrows( + NumberFormatException.class, + () -> + KueueWorkloadFactory.calculateExecutorMemoryMiB( + Map.of("spark.executor.memory", "1.5g"), false, false)); + } + + @Test + void testHasAndGetQueueName() { + SparkApplication app = new SparkApplication(); + assertFalse(KueueWorkloadFactory.hasQueueName(app)); + assertNull(KueueWorkloadFactory.getQueueName(app)); + + // From label + app.setMetadata( + new ObjectMetaBuilder() + .withLabels(Map.of(Constants.LABEL_QUEUE_NAME, "test-queue")) + .build()); + assertTrue(KueueWorkloadFactory.hasQueueName(app)); + assertEquals("test-queue", KueueWorkloadFactory.getQueueName(app)); + + // Empty label + SparkApplication app2 = new SparkApplication(); + app2.setMetadata( + new ObjectMetaBuilder() + .withLabels(Map.of(Constants.LABEL_QUEUE_NAME, "")) + .build()); + assertFalse(KueueWorkloadFactory.hasQueueName(app2)); + assertNull(KueueWorkloadFactory.getQueueName(app2)); + } + + @Test + void testBuildWorkloadStaticAllocation() { + SparkApplication app = new SparkApplication(); + app.setMetadata( + new ObjectMetaBuilder() + .withName("spark-pi") + .withNamespace("spark-jobs") + .withUid("app-uid-123") + .withLabels(Map.of(Constants.LABEL_QUEUE_NAME, "team-a-queue")) + .build()); + + ApplicationSpec spec = new ApplicationSpec(); + spec.setSparkConf( + Map.of( + "spark.driver.cores", "2", + "spark.driver.memory", "2g", + "spark.executor.cores", "4", + "spark.executor.memory", "4g", + "spark.executor.instances", "3")); + app.setSpec(spec); + + Workload workload = KueueWorkloadFactory.buildWorkload(app); + assertNotNull(workload); + assertEquals("sparkapplication-spark-pi", workload.getMetadata().getName()); + assertEquals("spark-jobs", workload.getMetadata().getNamespace()); + assertEquals("team-a-queue", workload.getSpec().getQueueName()); + assertTrue(workload.getSpec().getActive()); + assertEquals(1, workload.getMetadata().getOwnerReferences().size()); + assertEquals("spark-pi", workload.getMetadata().getOwnerReferences().get(0).getName()); + assertTrue(workload.getMetadata().getOwnerReferences().get(0).getController()); + + assertEquals(2, workload.getSpec().getPodSets().size()); + + PodSet driverPodSet = workload.getSpec().getPodSets().get(0); + assertEquals("driver", driverPodSet.getName()); + assertEquals(1, driverPodSet.getCount()); + assertNull(driverPodSet.getMinCount()); + Container driverContainer = + driverPodSet.getTemplate().getSpec().getContainers().get(0); + assertEquals(new Quantity("2"), driverContainer.getResources().getRequests().get("cpu")); + assertEquals( + new Quantity("2432Mi"), + driverContainer.getResources().getRequests().get("memory")); // 2048 + 384 = 2432 + + PodSet executorPodSet = workload.getSpec().getPodSets().get(1); + assertEquals("executor", executorPodSet.getName()); + assertEquals(3, executorPodSet.getCount()); + assertNull(executorPodSet.getMinCount()); + Container executorContainer = + executorPodSet.getTemplate().getSpec().getContainers().get(0); + assertEquals(new Quantity("4"), executorContainer.getResources().getRequests().get("cpu")); + assertEquals( + new Quantity("4505Mi"), + executorContainer.getResources().getRequests().get("memory")); // 4096 + 409 = 4505 + } + + @Test + void testBuildWorkloadDefaultExecutorInstances() { + SparkApplication app = new SparkApplication(); + app.setMetadata( + new ObjectMetaBuilder().withName("spark-default").withNamespace("default").build()); + + Workload workload = KueueWorkloadFactory.buildWorkload(app); + PodSet executorPodSet = workload.getSpec().getPodSets().get(1); + assertEquals(2, executorPodSet.getCount()); + assertNull(executorPodSet.getMinCount()); + } + + @Test + void testBuildWorkloadDynamicAllocation() { + SparkApplication app = new SparkApplication(); + app.setMetadata( + new ObjectMetaBuilder() + .withName("spark-elastic") + .withNamespace("default") + .withLabels(Map.of(Constants.LABEL_QUEUE_NAME, "elastic-queue")) + .build()); + + ApplicationSpec spec = new ApplicationSpec(); + spec.setSparkConf( + Map.of( + "spark.dynamicAllocation.enabled", "true", + "spark.dynamicAllocation.minExecutors", "2", + "spark.dynamicAllocation.maxExecutors", "10")); + app.setSpec(spec); + + assertThrows( + UnsupportedOperationException.class, () -> KueueWorkloadFactory.buildWorkload(app)); + } + + @Test + void testBuildWorkloadForPythonApp() { + SparkApplication app = new SparkApplication(); + app.setMetadata( + new ObjectMetaBuilder().withName("pi-python").withNamespace("default").build()); + + ApplicationSpec spec = new ApplicationSpec(); + spec.setPyFiles("local:///opt/spark/examples/src/main/python/pi.py"); + spec.setSparkConf( + Map.of( + "spark.driver.memory", "8g", + "spark.executor.memory", "8g", + "spark.executor.pyspark.memory", "1g")); + app.setSpec(spec); + + Workload workload = KueueWorkloadFactory.buildWorkload(app); + Container driverContainer = + workload.getSpec().getPodSets().get(0).getTemplate().getSpec().getContainers().get(0); + assertEquals( + new Quantity("11468Mi"), + driverContainer.getResources().getRequests().get("memory")); // 8192 + 3276 = 11468 + Container executorContainer = + workload.getSpec().getPodSets().get(1).getTemplate().getSpec().getContainers().get(0); + assertEquals( + new Quantity("12492Mi"), + executorContainer.getResources().getRequests().get("memory")); // 8192 + 3276 + 1024 + } + + @Test + void testBuildWorkloadWithGpuAndTemplateSpec() { + SparkApplication app = new SparkApplication(); + app.setMetadata( + new ObjectMetaBuilder() + .withName("spark-gpu") + .withNamespace("default") + .withLabels(Map.of(Constants.LABEL_QUEUE_NAME, "gpu-queue")) + .build()); + + PodTemplateSpec driverTemplate = + new PodTemplateSpecBuilder() + .withNewSpec() + .withContainers( + new ContainerBuilder() + .withName("sidecar") + .withResources( + new ResourceRequirementsBuilder() + .withRequests(Map.of("cpu", new Quantity("100m"))) + .build()) + .build(), + new ContainerBuilder() + .withName("custom-driver") + .withResources( + new ResourceRequirementsBuilder() + .withRequests(Map.of("cpu", new Quantity("3"))) + .build()) + .build()) + .endSpec() + .build(); + + ApplicationSpec spec = new ApplicationSpec(); + spec.setDriverSpec(new BaseApplicationTemplateSpec(driverTemplate)); + spec.setSparkConf( + Map.of( + Constants.DRIVER_SPARK_CONTAINER_PROP_KEY, "custom-driver", Review Comment: Added `testBuildWorkloadWithExecutorTemplateSpecAndRequestCores`, which selects `custom-executor` behind a leading sidecar via `spark.kubernetes.executor.podTemplateContainerName`, checks that the sidecar's requests are untouched, and covers `spark.kubernetes.{driver,executor}.request.cores`. In addition, a warning is logged like Spark's `selectSparkContainer` when the specified container is not found. ########## spark-operator/src/main/java/org/apache/spark/k8s/operator/kueue/KueueWorkloadFactory.java: ########## @@ -0,0 +1,475 @@ +/* + * 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 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.ClusterSpec; +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). + */ +@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"; + + public static final String DEFAULT_CORES = "1"; + public static final String DEFAULT_MEMORY = "1g"; + public static final String DEFAULT_MIN_MEMORY_OVERHEAD = "384m"; + public static final double DEFAULT_MEMORY_OVERHEAD_FACTOR = 0.10; + public static final double NON_JVM_MEMORY_OVERHEAD_FACTOR = 0.40; + public static final int DEFAULT_EXECUTOR_INSTANCES = 2; + + 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) { + String queueName = getQueueName(app); + ApplicationSpec appSpec = app.getSpec(); + Map<String, String> sparkConf = + appSpec != null && appSpec.getSparkConf() != null + ? appSpec.getSparkConf() + : Map.of(); + + PodSet driverPodSet = buildDriverPodSet(app, sparkConf); + PodSet executorPodSet = buildExecutorPodSet(app, sparkConf); + + List<PodSet> podSets = new ArrayList<>(); + podSets.add(driverPodSet); + podSets.add(executorPodSet); + + boolean active = appSpec == null || !appSpec.isSuspend(); + + Map<String, String> labels = new HashMap<>(); + if (app.getMetadata().getLabels() != null) { + labels.putAll(app.getMetadata().getLabels()); + } + labels.put(Constants.LABEL_SPARK_APPLICATION_NAME, app.getMetadata().getName()); + if (StringUtils.isNotEmpty(queueName)) { + labels.put(Constants.LABEL_QUEUE_NAME, queueName); + } + + OwnerReference ownerReference = ModelUtils.buildOwnerReferenceTo(app); + ownerReference.setController(true); + + Workload workload = new Workload(); + workload.setMetadata( + new ObjectMetaBuilder() + .withName(getWorkloadName(app)) + .withNamespace(app.getMetadata().getNamespace()) + .withLabels(labels) + .withOwnerReferences(ownerReference) + .build()); + + workload.setSpec( + WorkloadSpec.builder() + .queueName(queueName) + .active(active) + .podSets(podSets) + .build()); + + return workload; + } + + /** + * Builds a Kueue Workload from a SparkCluster resource. + * + * @param cluster The SparkCluster. + * @return The constructed Kueue Workload. + */ + public static Workload buildWorkload(final SparkCluster cluster) { + String queueName = getQueueName(cluster); + ClusterSpec clusterSpec = cluster.getSpec(); + + // 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 = new ArrayList<>(); + podSets.add(buildPodSet(PODSET_MASTER, resourceSpec.getMasterStatefulSet())); + podSets.add(buildPodSet(PODSET_WORKER, resourceSpec.getWorkerStatefulSet())); + + boolean active = clusterSpec == null || !clusterSpec.isSuspend(); + + Map<String, String> labels = new HashMap<>(); + if (cluster.getMetadata().getLabels() != null) { + labels.putAll(cluster.getMetadata().getLabels()); + } + labels.put(Constants.LABEL_SPARK_CLUSTER_NAME, cluster.getMetadata().getName()); + if (StringUtils.isNotEmpty(queueName)) { + labels.put(Constants.LABEL_QUEUE_NAME, queueName); + } + + OwnerReference ownerReference = ModelUtils.buildOwnerReferenceTo(cluster); + ownerReference.setController(true); + + Workload workload = new Workload(); + workload.setMetadata( + new ObjectMetaBuilder() + .withName(getWorkloadName(cluster)) + .withNamespace(cluster.getMetadata().getNamespace()) + .withLabels(labels) + .withOwnerReferences(ownerReference) + .build()); + + workload.setSpec( + WorkloadSpec.builder() + .queueName(queueName) + .active(active) + .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)); + } + + /** + * Builds the driver PodSet. + */ + public static PodSet buildDriverPodSet( + final SparkApplication app, final Map<String, String> sparkConf) { + PodTemplateSpec templateSpec = null; + if (app.getSpec() != null + && app.getSpec().getDriverSpec() != null + && app.getSpec().getDriverSpec().getPodTemplateSpec() != null) { + templateSpec = ReconcilerUtils.clone(app.getSpec().getDriverSpec().getPodTemplateSpec()); + } + if (templateSpec == null) { + templateSpec = new PodTemplateSpecBuilder().build(); + } + ensurePodSpec(templateSpec); + + String cpu = + sparkConf.getOrDefault( + "spark.kubernetes.driver.request.cores", + sparkConf.getOrDefault("spark.driver.cores", DEFAULT_CORES)); + long memoryMiB = calculateDriverMemoryMiB(sparkConf, isNonJvmApp(app.getSpec())); + String gpuAmount = sparkConf.get("spark.driver.resource.gpu.amount"); + String gpuVendor = sparkConf.get("spark.driver.resource.gpu.vendor"); + + decorateTemplateResources( + templateSpec, + sparkConf.get(Constants.DRIVER_SPARK_CONTAINER_PROP_KEY), + "spark-driver", + cpu, + memoryMiB, + gpuAmount, + gpuVendor); + + return PodSet.builder() + .name(PODSET_DRIVER) + .count(1) + .template(templateSpec) + .build(); + } + + /** + * Builds the executor PodSet. + */ + public static PodSet buildExecutorPodSet( + final SparkApplication app, final Map<String, String> sparkConf) { + if ("true".equalsIgnoreCase(sparkConf.get("spark.dynamicAllocation.enabled"))) { + throw new UnsupportedOperationException( + "Kueue does not support SparkApplication with dynamic allocation " + + "(spark.dynamicAllocation.enabled=true) yet."); + } + PodTemplateSpec templateSpec = null; + if (app.getSpec() != null + && app.getSpec().getExecutorSpec() != null + && app.getSpec().getExecutorSpec().getPodTemplateSpec() != null) { + templateSpec = ReconcilerUtils.clone(app.getSpec().getExecutorSpec().getPodTemplateSpec()); + } + if (templateSpec == null) { + templateSpec = new PodTemplateSpecBuilder().build(); + } + ensurePodSpec(templateSpec); + + String cpu = + sparkConf.getOrDefault( + "spark.kubernetes.executor.request.cores", + sparkConf.getOrDefault("spark.executor.cores", DEFAULT_CORES)); + long memoryMiB = + calculateExecutorMemoryMiB( + sparkConf, isNonJvmApp(app.getSpec()), isPythonApp(app.getSpec())); + String gpuAmount = sparkConf.get("spark.executor.resource.gpu.amount"); + String gpuVendor = sparkConf.get("spark.executor.resource.gpu.vendor"); + + decorateTemplateResources( + templateSpec, + sparkConf.get("spark.kubernetes.executor.podTemplateContainerName"), Review Comment: Added `Constants.EXECUTOR_SPARK_CONTAINER_PROP_KEY` and used it here. ########## spark-operator/src/main/java/org/apache/spark/k8s/operator/kueue/KueueWorkloadFactory.java: ########## @@ -0,0 +1,475 @@ +/* + * 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 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.ClusterSpec; +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). + */ +@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"; + + public static final String DEFAULT_CORES = "1"; + public static final String DEFAULT_MEMORY = "1g"; + public static final String DEFAULT_MIN_MEMORY_OVERHEAD = "384m"; + public static final double DEFAULT_MEMORY_OVERHEAD_FACTOR = 0.10; + public static final double NON_JVM_MEMORY_OVERHEAD_FACTOR = 0.40; + public static final int DEFAULT_EXECUTOR_INSTANCES = 2; + + 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) { + String queueName = getQueueName(app); + ApplicationSpec appSpec = app.getSpec(); + Map<String, String> sparkConf = + appSpec != null && appSpec.getSparkConf() != null + ? appSpec.getSparkConf() + : Map.of(); + + PodSet driverPodSet = buildDriverPodSet(app, sparkConf); + PodSet executorPodSet = buildExecutorPodSet(app, sparkConf); + + List<PodSet> podSets = new ArrayList<>(); + podSets.add(driverPodSet); + podSets.add(executorPodSet); + + boolean active = appSpec == null || !appSpec.isSuspend(); + + Map<String, String> labels = new HashMap<>(); + if (app.getMetadata().getLabels() != null) { + labels.putAll(app.getMetadata().getLabels()); + } + labels.put(Constants.LABEL_SPARK_APPLICATION_NAME, app.getMetadata().getName()); + if (StringUtils.isNotEmpty(queueName)) { + labels.put(Constants.LABEL_QUEUE_NAME, queueName); + } + + OwnerReference ownerReference = ModelUtils.buildOwnerReferenceTo(app); + ownerReference.setController(true); + + Workload workload = new Workload(); + workload.setMetadata( + new ObjectMetaBuilder() + .withName(getWorkloadName(app)) + .withNamespace(app.getMetadata().getNamespace()) + .withLabels(labels) + .withOwnerReferences(ownerReference) + .build()); + + workload.setSpec( + WorkloadSpec.builder() + .queueName(queueName) + .active(active) + .podSets(podSets) + .build()); + + return workload; + } + + /** + * Builds a Kueue Workload from a SparkCluster resource. + * + * @param cluster The SparkCluster. + * @return The constructed Kueue Workload. + */ + public static Workload buildWorkload(final SparkCluster cluster) { + String queueName = getQueueName(cluster); + ClusterSpec clusterSpec = cluster.getSpec(); + + // 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 = new ArrayList<>(); + podSets.add(buildPodSet(PODSET_MASTER, resourceSpec.getMasterStatefulSet())); + podSets.add(buildPodSet(PODSET_WORKER, resourceSpec.getWorkerStatefulSet())); + + boolean active = clusterSpec == null || !clusterSpec.isSuspend(); + + Map<String, String> labels = new HashMap<>(); + if (cluster.getMetadata().getLabels() != null) { + labels.putAll(cluster.getMetadata().getLabels()); + } + labels.put(Constants.LABEL_SPARK_CLUSTER_NAME, cluster.getMetadata().getName()); + if (StringUtils.isNotEmpty(queueName)) { + labels.put(Constants.LABEL_QUEUE_NAME, queueName); + } + + OwnerReference ownerReference = ModelUtils.buildOwnerReferenceTo(cluster); + ownerReference.setController(true); + + Workload workload = new Workload(); + workload.setMetadata( + new ObjectMetaBuilder() + .withName(getWorkloadName(cluster)) + .withNamespace(cluster.getMetadata().getNamespace()) + .withLabels(labels) + .withOwnerReferences(ownerReference) + .build()); + + workload.setSpec( + WorkloadSpec.builder() + .queueName(queueName) + .active(active) + .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)); + } + + /** + * Builds the driver PodSet. + */ + public static PodSet buildDriverPodSet( + final SparkApplication app, final Map<String, String> sparkConf) { + PodTemplateSpec templateSpec = null; + if (app.getSpec() != null + && app.getSpec().getDriverSpec() != null + && app.getSpec().getDriverSpec().getPodTemplateSpec() != null) { + templateSpec = ReconcilerUtils.clone(app.getSpec().getDriverSpec().getPodTemplateSpec()); + } + if (templateSpec == null) { + templateSpec = new PodTemplateSpecBuilder().build(); + } + ensurePodSpec(templateSpec); + + String cpu = + sparkConf.getOrDefault( + "spark.kubernetes.driver.request.cores", + sparkConf.getOrDefault("spark.driver.cores", DEFAULT_CORES)); + long memoryMiB = calculateDriverMemoryMiB(sparkConf, isNonJvmApp(app.getSpec())); + String gpuAmount = sparkConf.get("spark.driver.resource.gpu.amount"); + String gpuVendor = sparkConf.get("spark.driver.resource.gpu.vendor"); + + decorateTemplateResources( + templateSpec, + sparkConf.get(Constants.DRIVER_SPARK_CONTAINER_PROP_KEY), + "spark-driver", Review Comment: Fixed. The default names are `spark-kubernetes-driver` and `spark-kubernetes-executor` now. ########## spark-operator/src/main/java/org/apache/spark/k8s/operator/kueue/KueueWorkloadFactory.java: ########## @@ -0,0 +1,475 @@ +/* + * 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 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.ClusterSpec; +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). + */ +@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"; + + public static final String DEFAULT_CORES = "1"; + public static final String DEFAULT_MEMORY = "1g"; + public static final String DEFAULT_MIN_MEMORY_OVERHEAD = "384m"; + public static final double DEFAULT_MEMORY_OVERHEAD_FACTOR = 0.10; + public static final double NON_JVM_MEMORY_OVERHEAD_FACTOR = 0.40; + public static final int DEFAULT_EXECUTOR_INSTANCES = 2; + + 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) { + String queueName = getQueueName(app); + ApplicationSpec appSpec = app.getSpec(); + Map<String, String> sparkConf = + appSpec != null && appSpec.getSparkConf() != null + ? appSpec.getSparkConf() + : Map.of(); + + PodSet driverPodSet = buildDriverPodSet(app, sparkConf); + PodSet executorPodSet = buildExecutorPodSet(app, sparkConf); + + List<PodSet> podSets = new ArrayList<>(); + podSets.add(driverPodSet); + podSets.add(executorPodSet); + + boolean active = appSpec == null || !appSpec.isSuspend(); + + Map<String, String> labels = new HashMap<>(); + if (app.getMetadata().getLabels() != null) { + labels.putAll(app.getMetadata().getLabels()); + } + labels.put(Constants.LABEL_SPARK_APPLICATION_NAME, app.getMetadata().getName()); + if (StringUtils.isNotEmpty(queueName)) { + labels.put(Constants.LABEL_QUEUE_NAME, queueName); + } + + OwnerReference ownerReference = ModelUtils.buildOwnerReferenceTo(app); + ownerReference.setController(true); + + Workload workload = new Workload(); + workload.setMetadata( + new ObjectMetaBuilder() + .withName(getWorkloadName(app)) + .withNamespace(app.getMetadata().getNamespace()) + .withLabels(labels) + .withOwnerReferences(ownerReference) + .build()); + + workload.setSpec( + WorkloadSpec.builder() + .queueName(queueName) + .active(active) + .podSets(podSets) + .build()); + + return workload; + } + + /** + * Builds a Kueue Workload from a SparkCluster resource. + * + * @param cluster The SparkCluster. + * @return The constructed Kueue Workload. + */ + public static Workload buildWorkload(final SparkCluster cluster) { + String queueName = getQueueName(cluster); + ClusterSpec clusterSpec = cluster.getSpec(); + + // 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 = new ArrayList<>(); + podSets.add(buildPodSet(PODSET_MASTER, resourceSpec.getMasterStatefulSet())); + podSets.add(buildPodSet(PODSET_WORKER, resourceSpec.getWorkerStatefulSet())); + + boolean active = clusterSpec == null || !clusterSpec.isSuspend(); + + Map<String, String> labels = new HashMap<>(); + if (cluster.getMetadata().getLabels() != null) { + labels.putAll(cluster.getMetadata().getLabels()); + } + labels.put(Constants.LABEL_SPARK_CLUSTER_NAME, cluster.getMetadata().getName()); + if (StringUtils.isNotEmpty(queueName)) { + labels.put(Constants.LABEL_QUEUE_NAME, queueName); + } + + OwnerReference ownerReference = ModelUtils.buildOwnerReferenceTo(cluster); + ownerReference.setController(true); + + Workload workload = new Workload(); + workload.setMetadata( + new ObjectMetaBuilder() + .withName(getWorkloadName(cluster)) + .withNamespace(cluster.getMetadata().getNamespace()) + .withLabels(labels) + .withOwnerReferences(ownerReference) + .build()); + + workload.setSpec( + WorkloadSpec.builder() + .queueName(queueName) + .active(active) + .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)); + } + + /** + * Builds the driver PodSet. + */ + public static PodSet buildDriverPodSet( + final SparkApplication app, final Map<String, String> sparkConf) { + PodTemplateSpec templateSpec = null; + if (app.getSpec() != null + && app.getSpec().getDriverSpec() != null + && app.getSpec().getDriverSpec().getPodTemplateSpec() != null) { + templateSpec = ReconcilerUtils.clone(app.getSpec().getDriverSpec().getPodTemplateSpec()); + } + if (templateSpec == null) { + templateSpec = new PodTemplateSpecBuilder().build(); + } + ensurePodSpec(templateSpec); + + String cpu = + sparkConf.getOrDefault( + "spark.kubernetes.driver.request.cores", + sparkConf.getOrDefault("spark.driver.cores", DEFAULT_CORES)); + long memoryMiB = calculateDriverMemoryMiB(sparkConf, isNonJvmApp(app.getSpec())); + String gpuAmount = sparkConf.get("spark.driver.resource.gpu.amount"); + String gpuVendor = sparkConf.get("spark.driver.resource.gpu.vendor"); + + decorateTemplateResources( + templateSpec, + sparkConf.get(Constants.DRIVER_SPARK_CONTAINER_PROP_KEY), + "spark-driver", + cpu, + memoryMiB, + gpuAmount, + gpuVendor); + + return PodSet.builder() + .name(PODSET_DRIVER) + .count(1) + .template(templateSpec) + .build(); + } + + /** + * Builds the executor PodSet. + */ + public static PodSet buildExecutorPodSet( + final SparkApplication app, final Map<String, String> sparkConf) { + if ("true".equalsIgnoreCase(sparkConf.get("spark.dynamicAllocation.enabled"))) { Review Comment: Fixed. The four methods are package-private now, and the dynamic allocation check moved to the top of `buildWorkload(SparkApplication)` next to the `SparkCluster` HPA rejection. In d7cb986, I also merged the duplicated driver / executor logic into a single role-based `buildPodSet` and the duplicated Workload assembly into one helper. -- 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]
