package controller import ( "context" "fmt" "sort" "strings" "time" appsv1 "k8s.io/api/apps/v1" corev1 "k8s.io/api/core/v1" apierrors "k8s.io/apimachinery/pkg/api/errors" "k8s.io/apimachinery/pkg/api/resource" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/runtime" "k8s.io/apimachinery/pkg/types" ctrl "sigs.k8s.io/controller-runtime" "sigs.k8s.io/controller-runtime/pkg/client" "sigs.k8s.io/controller-runtime/pkg/controller/controllerutil" "sigs.k8s.io/controller-runtime/pkg/log" orcv1 "github.com/cyrus/spark-orchestrator/api/v1alpha1" "github.com/cyrus/spark-orchestrator/internal/k8s" "github.com/cyrus/spark-orchestrator/internal/util" ) type SparkVirtualClusterReconciler struct { client.Client Scheme *runtime.Scheme } const ( sparkPrePullStateConfigMap = "spark-prepull-state" sparkNamespaceOwnerKindLabel = "spark" ) func (r *SparkVirtualClusterReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) { logger := log.FromContext(ctx).WithValues("sparkVirtualCluster", req.Name) var cluster orcv1.SparkVirtualCluster if err := r.Get(ctx, req.NamespacedName, &cluster); err != nil { return ctrl.Result{}, client.IgnoreNotFound(err) } if cluster.Spec.DesiredState == "" { cluster.Spec.DesiredState = orcv1.VirtualClusterDesiredActive } if cluster.Spec.ScaleInPolicy == "" { cluster.Spec.ScaleInPolicy = orcv1.ScaleInPolicyBlockIfRunning } if cluster.DeletionTimestamp.IsZero() { if !controllerutil.ContainsFinalizer(&cluster, util.FinalizerVirtualCluster) { controllerutil.AddFinalizer(&cluster, util.FinalizerVirtualCluster) if err := r.Update(ctx, &cluster); err != nil { return ctrl.Result{}, err } } } else { if controllerutil.ContainsFinalizer(&cluster, util.FinalizerVirtualCluster) { if err := r.cleanupCluster(ctx, &cluster); err != nil { return ctrl.Result{}, err } controllerutil.RemoveFinalizer(&cluster, util.FinalizerVirtualCluster) if err := r.Update(ctx, &cluster); err != nil { return ctrl.Result{}, err } } return ctrl.Result{}, nil } if cluster.Spec.DesiredState == orcv1.VirtualClusterDesiredReleased { cluster.Status.Phase = orcv1.SparkClusterPhaseReleasing cluster.Status.ObservedAt = metav1.Now() if err := r.Status().Update(ctx, &cluster); err != nil { return ctrl.Result{}, err } if err := r.cleanupCluster(ctx, &cluster); err != nil { cluster.Status.Phase = orcv1.SparkClusterPhaseFailed cluster.Status.LastError = err.Error() cluster.Status.ObservedAt = metav1.Now() _ = r.Status().Update(ctx, &cluster) return ctrl.Result{}, err } cluster.Status.Phase = orcv1.SparkClusterPhaseReleased cluster.Status.AssignedNodes = nil cluster.Status.LastError = "" cluster.Status.ObservedAt = metav1.Now() setCondition(&cluster.Status.Conditions, orcv1.NewCondition(orcv1.ConditionReady, "Released", "Cluster released and nodes are idle", metav1.ConditionFalse, cluster.Generation)) if err := r.Status().Update(ctx, &cluster); err != nil { return ctrl.Result{}, err } return ctrl.Result{}, nil } ns := cluster.Spec.Namespace if ns == "" { ns = util.SafeName("spark", cluster.Name) if err := k8s.EnsureNamespaceOwnedByCluster(ctx, r.Client, ns, sparkNamespaceOwnerKindLabel, cluster.Name); err != nil { return ctrl.Result{}, err } } else { if err := k8s.EnsureNamespace(ctx, r.Client, ns); err != nil { return ctrl.Result{}, err } } if err := k8s.EnsureSparkOperatorRBAC(ctx, r.Client, ns); err != nil { return ctrl.Result{}, err } targetNodes := normalizeNodeList(cluster.Spec.NodeRefs) previousNodes := normalizeNodeList(cluster.Status.AssignedNodes) toRelease := diffNodes(previousNodes, targetNodes) for _, nodeID := range toRelease { if err := r.releaseNode(ctx, nodeID); err != nil { return ctrl.Result{}, err } } for _, nodeID := range targetNodes { if err := r.allocateNode(ctx, &cluster, nodeID); err != nil { cluster.Status.Phase = orcv1.SparkClusterPhaseFailed cluster.Status.LastError = err.Error() cluster.Status.ObservedAt = metav1.Now() _ = r.Status().Update(ctx, &cluster) logger.Error(err, "allocate node failed", "nodeID", nodeID) return ctrl.Result{RequeueAfter: 10 * time.Second}, nil } } prePullReady, prePullMessage, err := r.ensureSparkImagePrePull(ctx, &cluster, ns, targetNodes) if err != nil { cluster.Status.Phase = orcv1.SparkClusterPhaseFailed cluster.Status.LastError = err.Error() cluster.Status.Namespace = ns cluster.Status.AssignedNodes = targetNodes cluster.Status.ObservedAt = metav1.Now() setCondition(&cluster.Status.Conditions, orcv1.NewCondition(orcv1.ConditionFailed, "PrePullFailed", err.Error(), metav1.ConditionTrue, cluster.Generation)) _ = r.Status().Update(ctx, &cluster) return ctrl.Result{RequeueAfter: 10 * time.Second}, nil } if !prePullReady { cluster.Status.Phase = orcv1.SparkClusterPhaseCreating cluster.Status.Namespace = ns cluster.Status.AssignedNodes = targetNodes cluster.Status.LastError = "" cluster.Status.ObservedAt = metav1.Now() setCondition(&cluster.Status.Conditions, orcv1.NewCondition(orcv1.ConditionRunning, "PrePulling", prePullMessage, metav1.ConditionTrue, cluster.Generation)) if err := r.Status().Update(ctx, &cluster); err != nil { return ctrl.Result{}, err } return ctrl.Result{RequeueAfter: 5 * time.Second}, nil } warmupReady, warmupMessage, err := r.ensureSparkIvyWarmup(ctx, &cluster, ns, targetNodes) if err != nil { cluster.Status.Phase = orcv1.SparkClusterPhaseFailed cluster.Status.LastError = err.Error() cluster.Status.Namespace = ns cluster.Status.AssignedNodes = targetNodes cluster.Status.ObservedAt = metav1.Now() setCondition(&cluster.Status.Conditions, orcv1.NewCondition(orcv1.ConditionFailed, "IvyWarmupFailed", err.Error(), metav1.ConditionTrue, cluster.Generation)) _ = r.Status().Update(ctx, &cluster) return ctrl.Result{RequeueAfter: 10 * time.Second}, nil } if !warmupReady { cluster.Status.Phase = orcv1.SparkClusterPhaseCreating cluster.Status.Namespace = ns cluster.Status.AssignedNodes = targetNodes cluster.Status.LastError = "" cluster.Status.ObservedAt = metav1.Now() setCondition(&cluster.Status.Conditions, orcv1.NewCondition(orcv1.ConditionRunning, "IvyWarmup", warmupMessage, metav1.ConditionTrue, cluster.Generation)) if err := r.Status().Update(ctx, &cluster); err != nil { return ctrl.Result{}, err } return ctrl.Result{RequeueAfter: 5 * time.Second}, nil } cluster.Status.Namespace = ns cluster.Status.AssignedNodes = targetNodes if len(toRelease) > 0 || !equalNodeLists(previousNodes, targetNodes) { cluster.Status.Phase = orcv1.SparkClusterPhaseScaling } else { cluster.Status.Phase = orcv1.SparkClusterPhaseReady } cluster.Status.LastError = "" cluster.Status.ObservedAt = metav1.Now() setCondition(&cluster.Status.Conditions, orcv1.NewCondition(orcv1.ConditionReady, "ClusterReady", "Cluster nodes are isolated and ready", metav1.ConditionTrue, cluster.Generation)) if err := r.Status().Update(ctx, &cluster); err != nil { return ctrl.Result{}, err } if cluster.Status.Phase == orcv1.SparkClusterPhaseScaling { cluster.Status.Phase = orcv1.SparkClusterPhaseReady cluster.Status.ObservedAt = metav1.Now() _ = r.Status().Update(ctx, &cluster) } return ctrl.Result{}, nil } func (r *SparkVirtualClusterReconciler) cleanupCluster(ctx context.Context, cluster *orcv1.SparkVirtualCluster) error { ns := cluster.Status.Namespace if ns == "" { ns = cluster.Spec.Namespace } if ns == "" { ns = util.SafeName("spark", cluster.Name) } if err := r.cleanupSparkPrePullArtifacts(ctx, ns, cluster.Name); err != nil { return err } if err := r.cleanupSparkIvyWarmupArtifacts(ctx, ns, cluster.Name); err != nil { return err } nodes := cluster.Status.AssignedNodes if len(nodes) == 0 { nodes = cluster.Spec.NodeRefs } for _, nodeID := range normalizeNodeList(nodes) { if err := r.releaseNode(ctx, nodeID); err != nil { return err } } return nil } func (r *SparkVirtualClusterReconciler) allocateNode(ctx context.Context, cluster *orcv1.SparkVirtualCluster, nodeID string) error { var nodeInv orcv1.NodeInventory if err := r.Get(ctx, types.NamespacedName{Name: nodeID}, &nodeInv); err != nil { return err } if nodeInv.Status.KubernetesNodeName == "" { return fmt.Errorf("node %s has no kubernetesNodeName yet", nodeID) } clusterRef := util.SparkClusterRef(cluster.Name) if nodeInv.Status.Phase != orcv1.NodePhaseIdle && !(nodeInv.Status.Phase == orcv1.NodePhaseAllocated && (nodeInv.Status.ClusterRef == clusterRef || nodeInv.Status.ClusterRef == cluster.Name)) { return fmt.Errorf("node %s is not idle", nodeID) } if err := k8s.EnsureNodeIsolation(ctx, r.Client, nodeInv.Status.KubernetesNodeName, cluster.Name); err != nil { return err } nodeInv.Status.Phase = orcv1.NodePhaseAllocated nodeInv.Status.ClusterRef = clusterRef nodeInv.Status.LastError = "" nodeInv.Status.ObservedAt = metav1.Now() setCondition(&nodeInv.Status.Conditions, orcv1.NewCondition(orcv1.ConditionReady, "Allocated", fmt.Sprintf("Allocated to cluster %s", cluster.Name), metav1.ConditionTrue, nodeInv.Generation)) return r.Status().Update(ctx, &nodeInv) } func (r *SparkVirtualClusterReconciler) releaseNode(ctx context.Context, nodeID string) error { var nodeInv orcv1.NodeInventory if err := r.Get(ctx, types.NamespacedName{Name: nodeID}, &nodeInv); err != nil { if apierrors.IsNotFound(err) { return nil } return err } if nodeInv.Status.KubernetesNodeName != "" { if err := k8s.ClearNodeIsolation(ctx, r.Client, nodeInv.Status.KubernetesNodeName); err != nil { return err } } nodeInv.Status.Phase = orcv1.NodePhaseIdle nodeInv.Status.ClusterRef = "" nodeInv.Status.ObservedAt = metav1.Now() setCondition(&nodeInv.Status.Conditions, orcv1.NewCondition(orcv1.ConditionReady, "Idle", "Node is idle and available", metav1.ConditionTrue, nodeInv.Generation)) return r.Status().Update(ctx, &nodeInv) } func (r *SparkVirtualClusterReconciler) ensureSparkImagePrePull( ctx context.Context, cluster *orcv1.SparkVirtualCluster, namespace string, nodeIDs []string, ) (bool, string, error) { if len(nodeIDs) == 0 { return true, "No target nodes for pre-pull", nil } enabled, image, pullPolicy, err := k8s.ResolveSparkPrePullConfig(cluster.Spec.SparkConfig) if err != nil { return false, "", err } if !enabled { if err := r.cleanupSparkPrePullArtifacts(ctx, namespace, cluster.Name); err != nil { return false, "", err } return true, "Spark pre-pull is disabled", nil } token := util.RequestHash(image + "|" + strings.Join(nodeIDs, ",")) stateName := util.SafeName(sparkPrePullStateConfigMap, cluster.Name) stateToken, err := r.readSparkPrePullState(ctx, namespace, stateName) if err != nil { return false, "", err } daemonSetName := util.SafeName("spark-prepull", cluster.Name, token) if err := r.cleanupSparkPrePullDaemonSets(ctx, namespace, cluster.Name, daemonSetName); err != nil { return false, "", err } if stateToken == token { return true, fmt.Sprintf("Spark image pre-pull already completed for %s", image), nil } desired := buildSparkPrePullDaemonSet(namespace, cluster.Name, daemonSetName, image, pullPolicy) var existing appsv1.DaemonSet err = r.Get(ctx, types.NamespacedName{Name: daemonSetName, Namespace: namespace}, &existing) if apierrors.IsNotFound(err) { if err := r.Create(ctx, desired); err != nil { return false, "", err } return false, fmt.Sprintf("Spark pre-pull DaemonSet created for image %s", image), nil } if err != nil { return false, "", err } if existing.Status.DesiredNumberScheduled == 0 || existing.Status.NumberReady < int32(len(nodeIDs)) { return false, fmt.Sprintf("Spark image pre-pull in progress: ready=%d/%d", existing.Status.NumberReady, len(nodeIDs)), nil } if err := r.writeSparkPrePullState(ctx, namespace, stateName, token, image); err != nil { return false, "", err } if err := r.Delete(ctx, &existing); err != nil && !apierrors.IsNotFound(err) { return false, "", err } return true, fmt.Sprintf("Spark image pre-pull completed for %s", image), nil } func (r *SparkVirtualClusterReconciler) ensureSparkIvyWarmup( ctx context.Context, cluster *orcv1.SparkVirtualCluster, namespace string, nodeIDs []string, ) (bool, string, error) { if len(nodeIDs) == 0 { return true, "No target nodes for Ivy warmup", nil } token := k8s.ResolveSparkIvyWarmupToken(cluster.Spec.SparkConfig) stateName := util.SafeName(k8s.SparkIvyCacheStateConfigMapName, cluster.Name) stateToken, err := r.readSparkIvyWarmupState(ctx, namespace, stateName) if err != nil { return false, "", err } daemonSetName := util.SafeName("spark-ivy-warmup", cluster.Name, token) if err := r.cleanupSparkIvyWarmupDaemonSets(ctx, namespace, cluster.Name, daemonSetName); err != nil { return false, "", err } if stateToken == token { return true, "Spark Ivy cache warmup already completed", nil } desired, err := k8s.BuildSparkIvyWarmupDaemonSet(namespace, cluster.Name, daemonSetName, cluster.Spec.SparkConfig) if err != nil { return false, "", err } var existing appsv1.DaemonSet err = r.Get(ctx, types.NamespacedName{Name: daemonSetName, Namespace: namespace}, &existing) if apierrors.IsNotFound(err) { if err := r.Create(ctx, desired); err != nil { return false, "", err } return false, fmt.Sprintf("Spark Ivy warmup DaemonSet created for %d nodes", len(nodeIDs)), nil } if err != nil { return false, "", err } if !k8s.SparkIvyWarmupPodSpecMatches(existing.Spec.Template.Spec, desired.Spec.Template.Spec) { patchBase := existing.DeepCopy() existing.Spec = desired.Spec if err := r.Patch(ctx, &existing, client.MergeFrom(patchBase)); err != nil { return false, "", err } return false, "Spark Ivy warmup DaemonSet updated", nil } if existing.Status.DesiredNumberScheduled == 0 || existing.Status.NumberReady < int32(len(nodeIDs)) { return false, fmt.Sprintf("Spark Ivy warmup in progress: ready=%d/%d", existing.Status.NumberReady, len(nodeIDs)), nil } if err := r.writeSparkIvyWarmupState(ctx, namespace, stateName, token); err != nil { return false, "", err } if err := r.Delete(ctx, &existing); err != nil && !apierrors.IsNotFound(err) { return false, "", err } return true, "Spark Ivy cache warmup completed", nil } func (r *SparkVirtualClusterReconciler) cleanupSparkIvyWarmupArtifacts(ctx context.Context, namespace, clusterName string) error { if namespace == "" { return nil } if err := r.cleanupSparkIvyWarmupDaemonSets(ctx, namespace, clusterName, ""); err != nil { return err } state := &corev1.ConfigMap{ObjectMeta: metav1.ObjectMeta{ Name: util.SafeName(k8s.SparkIvyCacheStateConfigMapName, clusterName), Namespace: namespace, }} if err := r.Delete(ctx, state); err != nil && !apierrors.IsNotFound(err) { return err } return nil } func (r *SparkVirtualClusterReconciler) cleanupSparkIvyWarmupDaemonSets(ctx context.Context, namespace, clusterName, exceptName string) error { return k8s.DeleteSparkIvyWarmupDaemonSets(ctx, r.Client, namespace, clusterName, exceptName) } func (r *SparkVirtualClusterReconciler) readSparkIvyWarmupState(ctx context.Context, namespace, stateName string) (string, error) { state := &corev1.ConfigMap{} if err := r.Get(ctx, types.NamespacedName{Name: stateName, Namespace: namespace}, state); err != nil { if apierrors.IsNotFound(err) { return "", nil } return "", err } return strings.TrimSpace(state.Data["token"]), nil } func (r *SparkVirtualClusterReconciler) writeSparkIvyWarmupState(ctx context.Context, namespace, stateName, token string) error { if namespace == "" || stateName == "" { return nil } state := &corev1.ConfigMap{} key := types.NamespacedName{Name: stateName, Namespace: namespace} err := r.Get(ctx, key, state) if apierrors.IsNotFound(err) { state = &corev1.ConfigMap{ ObjectMeta: metav1.ObjectMeta{ Name: stateName, Namespace: namespace, Labels: map[string]string{ util.LabelManagedBy: util.ManagedByValue, }, }, Data: map[string]string{ "token": token, }, } return r.Create(ctx, state) } if err != nil { return err } if state.Data == nil { state.Data = map[string]string{} } state.Data["token"] = token return r.Update(ctx, state) } func (r *SparkVirtualClusterReconciler) cleanupSparkPrePullArtifacts(ctx context.Context, namespace, clusterName string) error { if namespace == "" { return nil } if err := r.cleanupSparkPrePullDaemonSets(ctx, namespace, clusterName, ""); err != nil { return err } state := &corev1.ConfigMap{ObjectMeta: metav1.ObjectMeta{ Name: util.SafeName(sparkPrePullStateConfigMap, clusterName), Namespace: namespace, }} if err := r.Delete(ctx, state); err != nil && !apierrors.IsNotFound(err) { return err } return nil } func (r *SparkVirtualClusterReconciler) cleanupSparkPrePullDaemonSets(ctx context.Context, namespace, clusterName, exceptName string) error { if namespace == "" { return nil } var list appsv1.DaemonSetList if err := r.List( ctx, &list, client.InNamespace(namespace), client.MatchingLabels{ util.LabelManagedBy: util.ManagedByValue, util.LabelClusterID: clusterName, util.LabelComponent: util.ComponentSparkPrePull, }, ); err != nil { return err } for i := range list.Items { ds := &list.Items[i] if exceptName != "" && ds.Name == exceptName { continue } if err := r.Delete(ctx, ds); err != nil && !apierrors.IsNotFound(err) { return err } } return nil } func (r *SparkVirtualClusterReconciler) readSparkPrePullState(ctx context.Context, namespace, stateName string) (string, error) { state := &corev1.ConfigMap{} if err := r.Get(ctx, types.NamespacedName{Name: stateName, Namespace: namespace}, state); err != nil { if apierrors.IsNotFound(err) { return "", nil } return "", err } return strings.TrimSpace(state.Data["token"]), nil } func (r *SparkVirtualClusterReconciler) writeSparkPrePullState(ctx context.Context, namespace, stateName, token, image string) error { if namespace == "" || stateName == "" { return nil } state := &corev1.ConfigMap{} key := types.NamespacedName{Name: stateName, Namespace: namespace} err := r.Get(ctx, key, state) if apierrors.IsNotFound(err) { state = &corev1.ConfigMap{ ObjectMeta: metav1.ObjectMeta{ Name: stateName, Namespace: namespace, Labels: map[string]string{ util.LabelManagedBy: util.ManagedByValue, }, }, Data: map[string]string{ "token": token, "image": image, }, } return r.Create(ctx, state) } if err != nil { return err } if state.Data == nil { state.Data = map[string]string{} } state.Data["token"] = token state.Data["image"] = image return r.Update(ctx, state) } func buildSparkPrePullDaemonSet(namespace, clusterName, daemonSetName, image string, pullPolicy corev1.PullPolicy) *appsv1.DaemonSet { labels := map[string]string{ util.LabelManagedBy: util.ManagedByValue, util.LabelClusterID: clusterName, util.LabelComponent: util.ComponentSparkPrePull, "orchestrator.spark.cyrus.io/prepull-daemon": daemonSetName, } return &appsv1.DaemonSet{ ObjectMeta: metav1.ObjectMeta{ Name: daemonSetName, Namespace: namespace, Labels: labels, }, Spec: appsv1.DaemonSetSpec{ Selector: &metav1.LabelSelector{ MatchLabels: map[string]string{ "orchestrator.spark.cyrus.io/prepull-daemon": daemonSetName, }, }, Template: corev1.PodTemplateSpec{ ObjectMeta: metav1.ObjectMeta{ Labels: labels, }, Spec: corev1.PodSpec{ RestartPolicy: corev1.RestartPolicyAlways, NodeSelector: map[string]string{ util.LabelClusterID: clusterName, }, Tolerations: []corev1.Toleration{ { Key: util.TaintKeySparkCluster, Operator: corev1.TolerationOpEqual, Value: clusterName, Effect: corev1.TaintEffectNoSchedule, }, }, Containers: []corev1.Container{ { Name: "prepull", Image: image, ImagePullPolicy: pullPolicy, Command: []string{"/bin/sh", "-c", "trap : TERM INT; sleep 3600 & wait"}, Resources: corev1.ResourceRequirements{ Requests: corev1.ResourceList{ corev1.ResourceCPU: resourceMustParse("10m"), corev1.ResourceMemory: resourceMustParse("32Mi"), }, Limits: corev1.ResourceList{ corev1.ResourceCPU: resourceMustParse("100m"), corev1.ResourceMemory: resourceMustParse("128Mi"), }, }, }, }, }, }, }, } } func resourceMustParse(v string) resource.Quantity { return resource.MustParse(v) } func (r *SparkVirtualClusterReconciler) SetupWithManager(mgr ctrl.Manager) error { return ctrl.NewControllerManagedBy(mgr). For(&orcv1.SparkVirtualCluster{}). Complete(r) } func normalizeNodeList(nodes []string) []string { unique := map[string]struct{}{} for _, node := range nodes { if node == "" { continue } unique[node] = struct{}{} } out := make([]string, 0, len(unique)) for node := range unique { out = append(out, node) } sort.Strings(out) return out } func diffNodes(oldNodes, newNodes []string) []string { set := map[string]struct{}{} for _, n := range newNodes { set[n] = struct{}{} } out := make([]string, 0) for _, old := range oldNodes { if _, ok := set[old]; !ok { out = append(out, old) } } return out } func equalNodeLists(a, b []string) bool { if len(a) != len(b) { return false } for i := range a { if a[i] != b[i] { return false } } return true }