fix: alloc worker port to avoid port conflicts

Author: 0x5457

go.mod

@@ -9,6 +9,7 @@ require (
 	github.com/onsi/ginkgo/v2 v2.22.1
 	github.com/onsi/gomega v1.36.2
 	github.com/samber/lo v1.47.0
+	golang.org/x/exp v0.0.0-20241217172543-b2144cdd0a67
 	gomodules.xyz/jsonpatch/v2 v2.4.0
 	k8s.io/api v0.32.0
 	k8s.io/apimachinery v0.32.0
@@ -94,7 +95,6 @@ require (
 	go.uber.org/zap v1.27.0 // indirect
 	golang.org/x/arch v0.12.0 // indirect
 	golang.org/x/crypto v0.31.0 // indirect
-	golang.org/x/exp v0.0.0-20241217172543-b2144cdd0a67 // indirect
 	golang.org/x/net v0.33.0 // indirect
 	golang.org/x/oauth2 v0.24.0 // indirect
 	golang.org/x/sync v0.10.0 // indirect

internal/config/config.go

@@ -15,7 +15,6 @@ type Config struct {
 
 type Worker struct {
 	corev1.PodTemplate
-	Port int16 `json:"port"`
 }
 
 type PodMutation struct {
@@ -40,7 +39,6 @@ func LoadConfig(filename string) (*Config, error) {
 func NewDefaultConfig() *Config {
 	return &Config{
 		Worker: Worker{
-			Port: 1234,
 			PodTemplate: corev1.PodTemplate{
 				Template: corev1.PodTemplateSpec{
 					Spec: corev1.PodSpec{

internal/constants/constants.go

@@ -3,21 +3,23 @@ package constants
 import "time"
 
 const (
-	// TensorFusionDomain is the domain prefix used for all tensor-fusion.ai related annotations and finalizers
-	TensorFusionDomain = "tensor-fusion.ai"
+	// Domain is the domain prefix used for all tensor-fusion.ai related annotations and finalizers
+	Domain = "tensor-fusion.ai"
 
 	// Finalizer constants
-	TensorFusionFinalizerSuffix = "finalizer"
-	TensorFusionFinalizer       = TensorFusionDomain + "/" + TensorFusionFinalizerSuffix
+	FinalizerSuffix = "finalizer"
+	Finalizer       = Domain + "/" + FinalizerSuffix
 
 	// Annotation key constants
-	EnableContainerAnnotationFormat = TensorFusionDomain + "/enable-%s"
-	TFLOPSContainerAnnotationFormat = TensorFusionDomain + "/tflops-%s"
-	VRAMContainerAnnotationFormat   = TensorFusionDomain + "/vram-%s"
+	EnableContainerAnnotationFormat = Domain + "/enable-%s"
+	TFLOPSContainerAnnotationFormat = Domain + "/tflops-%s"
+	VRAMContainerAnnotationFormat   = Domain + "/vram-%s"
 
 	PendingRequeueDuration = time.Second * 3
 
 	GetConnectionURLEnv    = "TENSOR_FUSION_OPERATOR_GET_CONNECTION_URL"
 	ConnectionNameEnv      = "TENSOR_FUSION_CONNECTION_NAME"
 	ConnectionNamespaceEnv = "TENSOR_FUSION_CONNECTION_NAMESPACE"
+
+	WorkerPortEnv = "TENSOR_FUSION_WORKER_PORT"
 )

internal/controller/pod_controller.go

@@ -119,7 +119,7 @@ func GenerateTensorFusionConnection(pod *corev1.Pod, tfReq []webhookv1.TFReq) []
 func (r *PodReconciler) SetupWithManager(mgr ctrl.Manager) error {
 	p, err := predicate.LabelSelectorPredicate(metav1.LabelSelector{
 		MatchLabels: map[string]string{
-			constants.TensorFusionDomain + "/enabled": "true",
+			constants.Domain + "/enabled": "true",
 		},
 	})
 	if err != nil {

internal/controller/tensorfusioncluster_controller.go

@@ -29,7 +29,7 @@ import (
 )
 
 var (
-	tensorFusionClusterFinalizer = constants.TensorFusionFinalizer
+	tensorFusionClusterFinalizer = constants.Finalizer
 )
 
 // TensorFusionClusterReconciler reconciles a TensorFusionCluster object

internal/controller/tensorfusionconnection_controller.go

@@ -44,7 +44,7 @@ type TensorFusionConnectionReconciler struct {
 }
 
 var (
-	tensorFusionConnectionFinalizer = constants.TensorFusionFinalizer
+	tensorFusionConnectionFinalizer = constants.Finalizer
 )
 
 // +kubebuilder:rbac:groups=core,resources=pods,verbs=get;list;watch;create;update;patch;delete
@@ -134,7 +134,10 @@ func (r *TensorFusionConnectionReconciler) Reconcile(ctx context.Context, req ct
 
 		if workerPod.Status.Phase == corev1.PodRunning {
 			connection.Status.Phase = tfv1.TensorFusionConnectionRunning
-			connection.Status.ConnectionURL = r.WorkerGenerator.GenerateConnectionURL(connection, workerPod)
+			connection.Status.ConnectionURL, err = r.WorkerGenerator.GenerateConnectionURL(connection, workerPod)
+			if err != nil {
+				return ctrl.Result{}, err
+			}
 		}
 		// TODO: Handle PodFailure
 	}
@@ -157,7 +160,8 @@ func (r *TensorFusionConnectionReconciler) tryStartWorker(ctx context.Context, g
 	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(gpu, connection, namespacedName)
+			port := r.WorkerGenerator.AllocPort()
+			pod = r.WorkerGenerator.GenerateWorkerPod(gpu, connection, namespacedName, port)
 			if err := ctrl.SetControllerReference(connection, pod, r.Scheme); err != nil {
 				return nil, fmt.Errorf("set owner reference %w", err)
 			}

internal/webhook/v1/pod_webhook_test.go

@@ -69,8 +69,8 @@ var _ = Describe("TensorFusionPodMutator", func() {
 					Name:      "test-pod",
 					Namespace: "default",
 					Annotations: map[string]string{
-						constants.TensorFusionDomain + "/tflops-main": "100",
-						constants.TensorFusionDomain + "/vram-main":   "16Gi",
+						constants.Domain + "/tflops-main": "100",
+						constants.Domain + "/vram-main":   "16Gi",
 					},
 				},
 				Spec: corev1.PodSpec{
@@ -165,8 +165,8 @@ var _ = Describe("TensorFusionPodMutator", func() {
 			pod := &corev1.Pod{
 				ObjectMeta: metav1.ObjectMeta{
 					Annotations: map[string]string{
-						constants.TensorFusionDomain + "/tflops-test-container": "100",
-						constants.TensorFusionDomain + "/vram-test-container":   "16Gi",
+						constants.Domain + "/tflops-test-container": "100",
+						constants.Domain + "/vram-test-container":   "16Gi",
 					},
 				},
 				Spec: corev1.PodSpec{

internal/worker/worker.go

@@ -2,30 +2,51 @@ package worker
 
 import (
 	"fmt"
+	"strconv"
+	"time"
 
 	tfv1 "github.com/NexusGPU/tensor-fusion-operator/api/v1"
 	"github.com/NexusGPU/tensor-fusion-operator/internal/config"
+	"github.com/NexusGPU/tensor-fusion-operator/internal/constants"
+	"github.com/samber/lo"
+	"golang.org/x/exp/rand"
 	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/types"
 )
 
+func init() {
+	rand.Seed(uint64(time.Now().UnixNano()))
+}
+
 type WorkerGenerator struct {
 	WorkerConfig *config.Worker
 }
 
-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) GenerateConnectionURL(connection *tfv1.TensorFusionConnection, pod *corev1.Pod) (string, error) {
+	port, ok := lo.Find(pod.Spec.Containers[0].Env, func(env corev1.EnvVar) bool {
+		return env.Name == constants.WorkerPortEnv
+	})
+
+	if !ok {
+		return "", fmt.Errorf("worker port not found in pod %s", pod.Name)
+	}
+	return fmt.Sprintf("native+%s+%d", pod.Status.PodIP, port.Value), nil
+}
+
+func (wg *WorkerGenerator) AllocPort() int16 {
+	min := 30000
+	max := 65535
+	return int16(rand.Intn(max-min+1) + min)
 }
 
 func (wg *WorkerGenerator) GenerateWorkerPod(
 	gpu *tfv1.GPU,
 	connection *tfv1.TensorFusionConnection,
 	namespacedName types.NamespacedName,
+	port int16,
 ) *corev1.Pod {
-
 	spec := wg.WorkerConfig.Template.Spec.DeepCopy()
-
 	if spec.NodeSelector == nil {
 		spec.NodeSelector = make(map[string]string)
 	}
@@ -34,13 +55,16 @@ func (wg *WorkerGenerator) GenerateWorkerPod(
 	spec.Containers[0].Env = append(spec.Containers[0].Env, corev1.EnvVar{
 		Name:  "NVIDIA_VISIBLE_DEVICES",
 		Value: gpu.Status.UUID,
+	}, corev1.EnvVar{
+		Name:  constants.WorkerPortEnv,
+		Value: strconv.Itoa(int(port)),
 	})
 
 	return &corev1.Pod{
 		ObjectMeta: metav1.ObjectMeta{
 			Name:      namespacedName.Name,
 			Namespace: namespacedName.Namespace,
 		},
-		Spec: spec,
+		Spec: *spec,
 	}
 }