feat: impl GenerateConnectionURL func (#12)

* fix lint

* feat: impl GenerateConnectionURL func

Author: 0x5457

cmd/main.go

@@ -167,7 +167,7 @@ func main() {
 		Scheme:    mgr.GetScheme(),
 		Scheduler: scheduler,
 		WorkerGenerator: &worker.WorkerGenerator{
-			PodTemplate: &config.Worker,
+			WorkerConfig: &config.Worker,
 		},
 	}).SetupWithManager(mgr); err != nil {
 		setupLog.Error(err, "unable to create controller", "controller", "TensorFusionConnection")
@@ -185,7 +185,7 @@ func main() {
 
 	// nolint:goconst
 	if os.Getenv("ENABLE_WEBHOOKS") != "false" {
-		if err = webhookcorev1.SetupPodWebhookWithManager(mgr, &config.PodMutator); err != nil {
+		if err = webhookcorev1.SetupPodWebhookWithManager(mgr, &config.PodMutation); err != nil {
 			setupLog.Error(err, "unable to create webhook", "webhook", "Pod")
 			os.Exit(1)
 		}

internal/config/config.go

@@ -9,11 +9,17 @@ import (
 )
 
 type Config struct {
-	Worker     corev1.PodTemplate `json:"worker"`
-	PodMutator PodMutator         `json:"podMutator"`
+	Worker      Worker      `json:"worker"`
+	PodMutation PodMutation `json:"podMutation"`
 }
 
-type PodMutator struct {
+type Worker struct {
+	corev1.PodTemplate
+	SendPort    int16 `json:"sendPort"`
+	ReceivePort int16 `json:"receivePort"`
+}
+
+type PodMutation struct {
 	PatchToPod       map[string]any `json:"patchToPod"`
 	PatchToContainer map[string]any `json:"patchToContainer"`
 }
@@ -33,21 +39,25 @@ func LoadConfig(filename string) (*Config, error) {
 
 func NewDefaultConfig() *Config {
 	return &Config{
-		Worker: corev1.PodTemplate{
-			Template: corev1.PodTemplateSpec{
-				Spec: corev1.PodSpec{
-					TerminationGracePeriodSeconds: ptr.To[int64](0),
-					Containers: []corev1.Container{
-						{
-							Name:    "tensorfusion-worker",
-							Image:   "busybox:stable-glibc",
-							Command: []string{"sleep", "infinity"},
+		Worker: Worker{
+			SendPort:    1234,
+			ReceivePort: 4321,
+			PodTemplate: corev1.PodTemplate{
+				Template: corev1.PodTemplateSpec{
+					Spec: corev1.PodSpec{
+						TerminationGracePeriodSeconds: ptr.To[int64](0),
+						Containers: []corev1.Container{
+							{
+								Name:    "tensorfusion-worker",
+								Image:   "busybox:stable-glibc",
+								Command: []string{"sleep", "infinity"},
+							},
 						},
 					},
 				},
 			},
 		},
-		PodMutator: PodMutator{
+		PodMutation: PodMutation{
 			PatchToPod: map[string]any{
 				"spec": map[string]any{
 					"initContainers": []corev1.Container{

internal/controller/tensorfusionconnection_controller.go

@@ -117,15 +117,15 @@ func (r *TensorFusionConnectionReconciler) Reconcile(ctx context.Context, req ct
 	}
 
 	// Start worker job
-	phase, err := r.tryStartWorker(ctx, connection, types.NamespacedName{Name: connection.Name, Namespace: connection.Namespace})
+	workerPod, err := r.tryStartWorker(ctx, 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 phase == corev1.PodRunning {
+	if workerPod.Status.Phase == corev1.PodRunning {
 		connection.Status.Phase = tfv1.TensorFusionConnectionRunning
-		connection.Status.ConnectionURL = r.WorkerGenerator.GenerateConnectionURL(gpu, connection)
+		connection.Status.ConnectionURL = r.WorkerGenerator.GenerateConnectionURL(gpu, connection, workerPod)
 	}
 	// TODO: Handle PodFailure
 
@@ -141,23 +141,23 @@ 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.PodPhase, error) {
+func (r *TensorFusionConnectionReconciler) tryStartWorker(ctx context.Context, 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)
 			if err := ctrl.SetControllerReference(connection, pod, r.Scheme); err != nil {
-				return "", fmt.Errorf("set owner reference %w", err)
+				return nil, fmt.Errorf("set owner reference %w", err)
 			}
 			if err := r.Create(ctx, pod); err != nil {
-				return "", fmt.Errorf("create pod %w", err)
+				return nil, fmt.Errorf("create pod %w", err)
 			}
-			return corev1.PodPending, nil
+			return pod, nil
 		}
 	}
-	return pod.Status.Phase, nil
+	return pod, nil
 }
 
 // handleDeletion handles cleanup of external dependencies

internal/controller/tensorfusionconnection_controller_test.go

@@ -75,7 +75,7 @@ var _ = Describe("TensorFusionConnection Controller", func() {
 				Client: k8sClient,
 				Scheme: k8sClient.Scheme(),
 				WorkerGenerator: &worker.WorkerGenerator{
-					PodTemplate: &config.Worker,
+					WorkerConfig: &config.Worker,
 				},
 			}
 			_, err := controllerReconciler.Reconcile(ctx, reconcile.Request{

internal/webhook/v1/pod_webhook.go

@@ -37,7 +37,7 @@ import (
 )
 
 // SetupPodWebhookWithManager registers the webhook for Pod in the manager.
-func SetupPodWebhookWithManager(mgr ctrl.Manager, config *config.PodMutator) error {
+func SetupPodWebhookWithManager(mgr ctrl.Manager, config *config.PodMutation) error {
 	webhookServer := mgr.GetWebhookServer()
 
 	webhookServer.Register("/mutate-v1-pod",
@@ -53,7 +53,7 @@ func SetupPodWebhookWithManager(mgr ctrl.Manager, config *config.PodMutator) err
 
 type TensorFusionPodMutator struct {
 	Client  client.Client
-	Config  *config.PodMutator
+	Config  *config.PodMutation
 	decoder admission.Decoder
 }
 

internal/webhook/v1/pod_webhook_test.go

@@ -57,7 +57,7 @@ var _ = Describe("TensorFusionPodMutator", func() {
 		config := config.NewDefaultConfig()
 		mutator = &TensorFusionPodMutator{
 			Client: client,
-			Config: &config.PodMutator,
+			Config: &config.PodMutation,
 		}
 		Expect(mutator.InjectDecoder(decoder)).To(Succeed())
 	})

internal/webhook/v1/webhook_suite_test.go

@@ -119,7 +119,7 @@ var _ = BeforeSuite(func() {
 	Expect(err).NotTo(HaveOccurred())
 
 	conf := config.NewDefaultConfig()
-	err = SetupPodWebhookWithManager(mgr, &conf.PodMutator)
+	err = SetupPodWebhookWithManager(mgr, &conf.PodMutation)
 	Expect(err).NotTo(HaveOccurred())
 
 	// +kubebuilder:scaffold:webhook

internal/worker/worker.go

@@ -1,18 +1,21 @@
 package worker
 
 import (
+	"fmt"
+
 	tfv1 "github.com/NexusGPU/tensor-fusion-operator/api/v1"
+	"github.com/NexusGPU/tensor-fusion-operator/internal/config"
 	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/types"
 )
 
 type WorkerGenerator struct {
-	PodTemplate *corev1.PodTemplate
+	WorkerConfig *config.Worker
 }
 
-func (wg *WorkerGenerator) GenerateConnectionURL(_gpu *tfv1.GPU, _connection *tfv1.TensorFusionConnection) string {
-	return "TODO://"
+func (wg *WorkerGenerator) GenerateConnectionURL(_gpu *tfv1.GPU, connection *tfv1.TensorFusionConnection, pod *corev1.Pod) string {
+	return fmt.Sprintf("native+%s+%d+%d", pod.Status.PodIP, wg.WorkerConfig.SendPort, wg.WorkerConfig.ReceivePort)
 }
 
 func (wg *WorkerGenerator) GenerateWorkerPod(
@@ -24,6 +27,6 @@ func (wg *WorkerGenerator) GenerateWorkerPod(
 			Name:      namespacedName.Name,
 			Namespace: namespacedName.Namespace,
 		},
-		Spec: wg.PodTemplate.Template.Spec,
+		Spec: wg.WorkerConfig.Template.Spec,
 	}
 }