feat: add GPU node selector and UUID support (#23)

- Add UUID and NodeSelector fields to GPUStatus
- Update worker pod generation to use GPU node selector
- Add NVIDIA_VISIBLE_DEVICES env var with GPU UUID
- Simplify connection URL generation by removing unused GPU parameter

Author: 0x5457

api/v1/gpu_types.go

@@ -22,8 +22,10 @@ import (
 
 // GPUStatus defines the observed state of GPU.
 type GPUStatus struct {
-	Capacity  Resource `json:"capacity"`
-	Available Resource `json:"available"`
+	UUID         string            `json:"uuid"`
+	NodeSelector map[string]string `json:"nodeSelector"`
+	Capacity     Resource          `json:"capacity"`
+	Available    Resource          `json:"available"`
 }
 
 // +kubebuilder:object:root=true

api/v1/zz_generated.deepcopy.go

@@ -75,9 +75,17 @@ spec:
                 - tflops
                 - vram
                 type: object
+              nodeSelector:
+                additionalProperties:
+                  type: string
+                type: object
+              uuid:
+                type: string
             required:
             - available
             - capacity
+            - nodeSelector
+            - uuid
             type: object
         type: object
     served: true

config/crd/bases/tensor-fusion.ai_gpus.yaml

@@ -118,15 +118,15 @@ func (r *TensorFusionConnectionReconciler) Reconcile(ctx context.Context, req ct
 
 	if connection.Status.Phase != tfv1.TensorFusionConnectionPending {
 		// Start worker job
-		workerPod, err := r.tryStartWorker(ctx, connection, types.NamespacedName{Name: connection.Name, Namespace: connection.Namespace})
+		workerPod, err := r.tryStartWorker(ctx, gpu, connection, types.NamespacedName{Name: connection.Name, Namespace: connection.Namespace})
 		if err != nil {
 			log.Error(err, "Failed to start worker pod")
 			return ctrl.Result{}, err
 		}
 
 		if workerPod.Status.Phase == corev1.PodRunning {
 			connection.Status.Phase = tfv1.TensorFusionConnectionRunning
-			connection.Status.ConnectionURL = r.WorkerGenerator.GenerateConnectionURL(gpu, connection, workerPod)
+			connection.Status.ConnectionURL = r.WorkerGenerator.GenerateConnectionURL(connection, workerPod)
 		}
 		// TODO: Handle PodFailure
 	}
@@ -143,13 +143,13 @@ func (r *TensorFusionConnectionReconciler) Reconcile(ctx context.Context, req ct
 	return ctrl.Result{}, nil
 }
 
-func (r *TensorFusionConnectionReconciler) tryStartWorker(ctx context.Context, connection *tfv1.TensorFusionConnection, namespacedName types.NamespacedName) (*corev1.Pod, error) {
+func (r *TensorFusionConnectionReconciler) tryStartWorker(ctx context.Context, gpu *tfv1.GPU, connection *tfv1.TensorFusionConnection, 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
-			pod = r.WorkerGenerator.GenerateWorkerPod(connection, namespacedName)
+			pod = r.WorkerGenerator.GenerateWorkerPod(gpu, connection, namespacedName)
 			if err := ctrl.SetControllerReference(connection, pod, r.Scheme); err != nil {
 				return nil, fmt.Errorf("set owner reference %w", err)
 			}

internal/controller/tensorfusionconnection_controller.go

@@ -14,19 +14,32 @@ type WorkerGenerator struct {
 	WorkerConfig *config.Worker
 }
 
-func (wg *WorkerGenerator) GenerateConnectionURL(_gpu *tfv1.GPU, connection *tfv1.TensorFusionConnection, pod *corev1.Pod) string {
+func (wg *WorkerGenerator) GenerateConnectionURL(connection *tfv1.TensorFusionConnection, pod *corev1.Pod) string {
 	return fmt.Sprintf("native+%s+%d", pod.Status.PodIP, wg.WorkerConfig.Port)
 }
 
 func (wg *WorkerGenerator) GenerateWorkerPod(
+	gpu *tfv1.GPU,
 	connection *tfv1.TensorFusionConnection,
 	namespacedName types.NamespacedName,
 ) *corev1.Pod {
+
+	spec := wg.WorkerConfig.Template.Spec
+	if spec.NodeSelector == nil {
+		spec.NodeSelector = make(map[string]string)
+	}
+	spec.NodeSelector = gpu.Status.NodeSelector
+
+	spec.Containers[0].Env = append(spec.Containers[0].Env, corev1.EnvVar{
+		Name:  "NVIDIA_VISIBLE_DEVICES",
+		Value: gpu.Status.UUID,
+	})
+
 	return &corev1.Pod{
 		ObjectMeta: metav1.ObjectMeta{
 			Name:      namespacedName.Name,
 			Namespace: namespacedName.Namespace,
 		},
-		Spec: wg.WorkerConfig.Template.Spec,
+		Spec: spec,
 	}
 }