feat: create worker pod using GenerateName (#65)

Author: 0x5457

internal/constants/constants.go

@@ -44,11 +44,12 @@ const (
 	ConnectionNameEnv      = "TENSOR_FUSION_CONNECTION_NAME"
 	ConnectionNamespaceEnv = "TENSOR_FUSION_CONNECTION_NAMESPACE"
 
-	WorkerPortEnv        = "TENSOR_FUSION_WORKER_PORT"
-	WokerCudaUpLimitEnv  = "TENSOR_FUSION_CUDA_UP_LIMIT"
-	WokerCudaMemLimitEnv = "TENSOR_FUSION_CUDA_MEM_LIMIT"
-	NamespaceEnv         = "OPERATOR_NAMESPACE"
-	NamespaceDefaultVal  = "tensor-fusion"
+	WorkerPortEnv         = "TENSOR_FUSION_WORKER_PORT"
+	WorkerCudaUpLimitEnv  = "TENSOR_FUSION_CUDA_UP_LIMIT"
+	WorkerCudaMemLimitEnv = "TENSOR_FUSION_CUDA_MEM_LIMIT"
+	WorkerPodNameEnv      = "POD_NAME"
+	NamespaceEnv          = "OPERATOR_NAMESPACE"
+	NamespaceDefaultVal   = "tensor-fusion"
 )
 
 const (

internal/controller/tensorfusionworkload_controller.go

@@ -24,7 +24,6 @@ import (
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/errors"
 	"k8s.io/apimachinery/pkg/runtime"
-	"k8s.io/apimachinery/pkg/types"
 	"k8s.io/client-go/tools/record"
 	ctrl "sigs.k8s.io/controller-runtime"
 	"sigs.k8s.io/controller-runtime/pkg/client"
@@ -129,7 +128,7 @@ func (r *TensorFusionWorkloadReconciler) Reconcile(ctx context.Context, req ctrl
 
 		// Calculate how many pods need to be added
 		podsToAdd := int(desiredReplicas - currentReplicas)
-		if err := r.scaleUpWorkers(ctx, workerGenerator, workload, podsToAdd, req.Namespace); err != nil {
+		if err := r.scaleUpWorkers(ctx, workerGenerator, workload, podsToAdd); err != nil {
 			return ctrl.Result{}, err
 		}
 	} else if currentReplicas > desiredReplicas {
@@ -159,38 +158,28 @@ func (r *TensorFusionWorkloadReconciler) tryStartWorker(
 	workerGenerator *worker.WorkerGenerator,
 	gpu *tfv1.GPU,
 	workload *tfv1.TensorFusionWorkload,
-	namespacedName types.NamespacedName,
 ) (*corev1.Pod, error) {
-	// Try to get the Pod
-	pod := &corev1.Pod{}
-	if err := r.Get(ctx, namespacedName, pod); err != nil {
-		if errors.IsNotFound(err) {
-			// Pod doesn't exist, create a new one
-			port := workerGenerator.AllocPort()
-			pod, err = workerGenerator.GenerateWorkerPod(gpu, namespacedName, port, workload.Spec.Resources.Limits)
-			if err != nil {
-				return nil, fmt.Errorf("generate worker pod %w", err)
-			}
+	port := workerGenerator.AllocPort()
+	pod, err := workerGenerator.GenerateWorkerPod(gpu, workload.Name, workload.Namespace, port, workload.Spec.Resources.Limits)
+	if err != nil {
+		return nil, fmt.Errorf("generate worker pod %w", err)
+	}
 
-			// Add labels to identify this pod as part of the workload
-			if pod.Labels == nil {
-				pod.Labels = make(map[string]string)
-			}
-			pod.Labels[constants.WorkloadKey] = workload.Name
-			pod.Labels[constants.GpuKey] = gpu.Name
+	// Add labels to identify this pod as part of the workload
+	if pod.Labels == nil {
+		pod.Labels = make(map[string]string)
+	}
+	pod.Labels[constants.WorkloadKey] = workload.Name
+	pod.Labels[constants.GpuKey] = gpu.Name
 
-			// Add finalizer for GPU resource cleanup
-			pod.Finalizers = append(pod.Finalizers, constants.Finalizer)
+	// Add finalizer for GPU resource cleanup
+	pod.Finalizers = append(pod.Finalizers, constants.Finalizer)
 
-			if err := ctrl.SetControllerReference(workload, pod, r.Scheme); err != nil {
-				return nil, fmt.Errorf("set owner reference %w", err)
-			}
-			if err := r.Create(ctx, pod); err != nil {
-				return nil, fmt.Errorf("create pod %w", err)
-			}
-			return pod, nil
-		}
-		return nil, err
+	if err := ctrl.SetControllerReference(workload, pod, r.Scheme); err != nil {
+		return nil, fmt.Errorf("set owner reference %w", err)
+	}
+	if err := r.Create(ctx, pod); err != nil {
+		return nil, fmt.Errorf("create pod %w", err)
 	}
 	return pod, nil
 }
@@ -270,27 +259,19 @@ func (r *TensorFusionWorkloadReconciler) deletePod(ctx context.Context, pod *cor
 }
 
 // scaleUpWorkers handles the scaling up of worker pods
-func (r *TensorFusionWorkloadReconciler) scaleUpWorkers(ctx context.Context, workerGenerator *worker.WorkerGenerator, workload *tfv1.TensorFusionWorkload, count int, namespace string) error {
+func (r *TensorFusionWorkloadReconciler) scaleUpWorkers(ctx context.Context, workerGenerator *worker.WorkerGenerator, workload *tfv1.TensorFusionWorkload, count int) error {
 	log := log.FromContext(ctx)
 
 	// Create worker pods
-	currentCount := int(workload.Status.Replicas)
-	for i := range count {
+	for range count {
 		// Schedule GPU for the worker
 		gpu, err := r.Scheduler.Schedule(ctx, workload.Spec.PoolName, workload.Spec.Resources.Requests)
 		if err != nil {
 			r.Recorder.Eventf(workload, corev1.EventTypeWarning, "ScheduleGPUFailed", "Failed to schedule GPU: %v", err)
 			return fmt.Errorf("schedule GPU: %w", err)
 		}
 
-		// Create worker pod
-		workerName := fmt.Sprintf("%s-worker-%d", workload.Name, currentCount+i)
-		namespacedName := types.NamespacedName{
-			Namespace: namespace,
-			Name:      workerName,
-		}
-
-		_, err = r.tryStartWorker(ctx, workerGenerator, gpu, workload, namespacedName)
+		_, err = r.tryStartWorker(ctx, workerGenerator, gpu, workload)
 		if err != nil {
 			// Try to release the GPU resource if pod creation fails
 			releaseErr := r.Scheduler.Release(ctx, workload.Spec.Resources.Requests, gpu)
@@ -299,8 +280,6 @@ func (r *TensorFusionWorkloadReconciler) scaleUpWorkers(ctx context.Context, wor
 			}
 			return fmt.Errorf("create worker pod: %w", err)
 		}
-
-		log.Info("Created worker pod", "name", workerName)
 	}
 
 	return nil

internal/worker/worker.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
-	"path"
 	"strconv"
 	"time"
 
@@ -14,7 +13,6 @@ import (
 	"golang.org/x/exp/rand"
 	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
-	"k8s.io/apimachinery/pkg/types"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 )
 
@@ -46,7 +44,8 @@ func (wg *WorkerGenerator) AllocPort() int {
 
 func (wg *WorkerGenerator) GenerateWorkerPod(
 	gpu *tfv1.GPU,
-	namespacedName types.NamespacedName,
+	generateName string,
+	namespace string,
 	port int,
 	limits tfv1.Resource,
 ) (*corev1.Pod, error) {
@@ -64,13 +63,14 @@ func (wg *WorkerGenerator) GenerateWorkerPod(
 		Name: constants.DataVolumeName,
 		VolumeSource: corev1.VolumeSource{
 			HostPath: &corev1.HostPathVolumeSource{
-				Path: path.Join(constants.TFDataPath, namespacedName.Name),
+				Path: constants.TFDataPath,
 			},
 		},
 	})
 	spec.Containers[0].VolumeMounts = append(spec.Containers[0].VolumeMounts, corev1.VolumeMount{
-		Name:      constants.DataVolumeName,
-		MountPath: constants.TFDataPath,
+		Name:        constants.DataVolumeName,
+		MountPath:   constants.TFDataPath,
+		SubPathExpr: fmt.Sprintf("${%s}", constants.WorkerPodNameEnv),
 	})
 
 	spec.Containers[0].Env = append(spec.Containers[0].Env, corev1.EnvVar{
@@ -80,19 +80,26 @@ func (wg *WorkerGenerator) GenerateWorkerPod(
 		Name:  constants.WorkerPortEnv,
 		Value: strconv.Itoa(port),
 	}, corev1.EnvVar{
-		Name: constants.WokerCudaUpLimitEnv,
+		Name: constants.WorkerCudaUpLimitEnv,
 		// TODO: convert tflops to percent
 		Value: "100",
 	}, corev1.EnvVar{
-		Name: constants.WokerCudaMemLimitEnv,
+		Name: constants.WorkerCudaMemLimitEnv,
 		// bytesize
 		Value: strconv.FormatInt(limits.Vram.Value(), 10),
+	}, corev1.EnvVar{
+		Name: constants.WorkerPodNameEnv,
+		ValueFrom: &corev1.EnvVarSource{
+			FieldRef: &corev1.ObjectFieldSelector{
+				FieldPath: "metadata.name",
+			},
+		},
 	})
 
 	return &corev1.Pod{
 		ObjectMeta: metav1.ObjectMeta{
-			Name:      namespacedName.Name,
-			Namespace: namespacedName.Namespace,
+			GenerateName: generateName,
+			Namespace:    namespace,
 		},
 		Spec: spec,
 	}, nil