feat: injecting the GetConnectionURLEnv environment variable (#13)

Author: 0x5457

go.mod

@@ -5,15 +5,17 @@ go 1.22.0
 require (
 	github.com/gin-contrib/gzip v1.0.1
 	github.com/gin-gonic/gin v1.10.0
+	github.com/lithammer/shortuuid/v4 v4.2.0
 	github.com/onsi/ginkgo/v2 v2.19.0
 	github.com/onsi/gomega v1.33.1
+	github.com/samber/lo v1.47.0
 	gomodules.xyz/jsonpatch/v2 v2.4.0
-	gopkg.in/yaml.v2 v2.4.0
 	k8s.io/api v0.31.0
 	k8s.io/apimachinery v0.31.0
 	k8s.io/client-go v0.31.0
 	k8s.io/utils v0.0.0-20240711033017-18e509b52bc8
 	sigs.k8s.io/controller-runtime v0.19.1
+	sigs.k8s.io/yaml v1.4.0
 )
 
 require (
@@ -106,6 +108,7 @@ require (
 	google.golang.org/protobuf v1.34.2 // indirect
 	gopkg.in/evanphx/json-patch.v4 v4.12.0 // indirect
 	gopkg.in/inf.v0 v0.9.1 // indirect
+	gopkg.in/yaml.v2 v2.4.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 	k8s.io/apiextensions-apiserver v0.31.0 // indirect
 	k8s.io/apiserver v0.31.0 // indirect
@@ -115,5 +118,4 @@ require (
 	sigs.k8s.io/apiserver-network-proxy/konnectivity-client v0.30.3 // indirect
 	sigs.k8s.io/json v0.0.0-20221116044647-bc3834ca7abd // indirect
 	sigs.k8s.io/structured-merge-diff/v4 v4.4.1 // indirect
-	sigs.k8s.io/yaml v1.4.0 // indirect
 )

go.sum

@@ -115,6 +115,8 @@ github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
 github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
 github.com/leodido/go-urn v1.4.0 h1:WT9HwE9SGECu3lg4d/dIA+jxlljEa1/ffXKmRjqdmIQ=
 github.com/leodido/go-urn v1.4.0/go.mod h1:bvxc+MVxLKB4z00jd1z+Dvzr47oO32F/QSNjSBOlFxI=
+github.com/lithammer/shortuuid/v4 v4.2.0 h1:LMFOzVB3996a7b8aBuEXxqOBflbfPQAiVzkIcHO0h8c=
+github.com/lithammer/shortuuid/v4 v4.2.0/go.mod h1:D5noHZ2oFw/YaKCfGy0YxyE7M0wMbezmMjPdhyEFe6Y=
 github.com/mailru/easyjson v0.7.7 h1:UGYAvKxe3sBsEDzO8ZeWOSlIQfWFlxbzLZe7hwFURr0=
 github.com/mailru/easyjson v0.7.7/go.mod h1:xzfreul335JAWq5oZzymOObrkdz5UnU4kGfJJLY9Nlc=
 github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
@@ -148,6 +150,8 @@ github.com/prometheus/procfs v0.15.1/go.mod h1:fB45yRUv8NstnjriLhBQLuOUt+WW4BsoG
 github.com/rogpeppe/go-internal v1.12.0 h1:exVL4IDcn6na9z1rAb56Vxr+CgyK3nn3O+epU5NdKM8=
 github.com/rogpeppe/go-internal v1.12.0/go.mod h1:E+RYuTGaKKdloAfM02xzb0FW3Paa99yedzYV+kq4uf4=
 github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
+github.com/samber/lo v1.47.0 h1:z7RynLwP5nbyRscyvcD043DWYoOcYRv3mV8lBeqOCLc=
+github.com/samber/lo v1.47.0/go.mod h1:RmDH9Ct32Qy3gduHQuKJ3gW1fMHAnE/fAzQuf6He5cU=
 github.com/spf13/cobra v1.8.1 h1:e5/vxKd/rZsfSJMUX1agtjeTDf+qv1/JdBF8gg5k9ZM=
 github.com/spf13/cobra v1.8.1/go.mod h1:wHxEcudfqmLYa8iTfL+OuZPbBZkmvliBWKIezN3kD9Y=
 github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA=

internal/config/config.go

@@ -20,6 +20,7 @@ type Worker struct {
 }
 
 type PodMutation struct {
+	OperatorEndpoint string         `json:"operatorEndpoint"`
 	PatchToPod       map[string]any `json:"patchToPod"`
 	PatchToContainer map[string]any `json:"patchToContainer"`
 }
@@ -58,6 +59,7 @@ func NewDefaultConfig() *Config {
 			},
 		},
 		PodMutation: PodMutation{
+			OperatorEndpoint: "http://localhost:8080",
 			PatchToPod: map[string]any{
 				"spec": map[string]any{
 					"initContainers": []corev1.Container{

internal/constants/constants.go

@@ -16,4 +16,8 @@ const (
 	VRAMContainerAnnotationFormat   = TensorFusionDomain + "/vram-%s"
 
 	PendingRequeueDuration = time.Second * 3
+
+	GetConnectionURLEnv    = "TENSOR_FUSION_OPERATOR_GET_CONNECTION_URL"
+	ConnectionNameEnv      = "TENSOR_FUSION_CONNECTION_NAME"
+	ConnectionNamespaceEnv = "TENSOR_FUSION_CONNECTION_NAMESPACE"
 )

internal/controller/pod_controller.go

@@ -82,8 +82,8 @@ func GenerateTensorFusionConnection(pod *corev1.Pod, tfReq []webhookv1.TFReq) []
 	for _, req := range tfReq {
 		connection := &tfv1.TensorFusionConnection{
 			ObjectMeta: metav1.ObjectMeta{
-				Name:      fmt.Sprintf("%s-tf-%s", pod.Name, req.ContainerName),
-				Namespace: pod.Namespace,
+				Name:      req.ConnectionName,
+				Namespace: req.ConnectionNamespace,
 				OwnerReferences: []metav1.OwnerReference{
 					{
 						APIVersion: "v1",

internal/webhook/v1/pod_webhook.go

@@ -21,6 +21,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"net/http"
+	"path"
 
 	"gomodules.xyz/jsonpatch/v2"
 	corev1 "k8s.io/api/core/v1"
@@ -34,6 +35,8 @@ import (
 
 	"github.com/NexusGPU/tensor-fusion-operator/internal/config"
 	"github.com/NexusGPU/tensor-fusion-operator/internal/constants"
+	"github.com/lithammer/shortuuid/v4"
+	"github.com/samber/lo"
 )
 
 // SetupPodWebhookWithManager registers the webhook for Pod in the manager.
@@ -65,7 +68,7 @@ func (m *TensorFusionPodMutator) Handle(ctx context.Context, req admission.Reque
 	}
 
 	log := log.FromContext(ctx)
-	log.Info("Mutating pod", "name", pod.Name, "namespace", pod.Namespace)
+	log.Info("Mutating pod", "generateName", pod.GenerateName, "namespace", pod.Namespace)
 
 	reqs := ParseTFReq(pod)
 	if len(reqs) == 0 {
@@ -88,9 +91,11 @@ func (m *TensorFusionPodMutator) InjectDecoder(d admission.Decoder) error {
 }
 
 type TFReq struct {
-	ContainerName string
-	Tflops        resource.Quantity
-	Vram          resource.Quantity
+	ContainerName       string
+	ConnectionName      string
+	ConnectionNamespace string
+	Tflops              resource.Quantity
+	Vram                resource.Quantity
 }
 
 func ParseTFReq(pod *corev1.Pod) []TFReq {
@@ -114,8 +119,25 @@ func ParseTFReq(pod *corev1.Pod) []TFReq {
 			continue
 		}
 
+		connectionNameEnv, ok := lo.Find(container.Env, func(e corev1.EnvVar) bool {
+			return e.Name == constants.ConnectionNameEnv
+		})
+
+		if !ok {
+			continue
+		}
+		connectionNamespaceEnv, ok := lo.Find(container.Env, func(e corev1.EnvVar) bool {
+			return e.Name == constants.ConnectionNamespaceEnv
+		})
+
+		if !ok {
+			continue
+		}
+
 		req := TFReq{
-			ContainerName: containerName,
+			ContainerName:       containerName,
+			ConnectionName:      connectionNameEnv.Name,
+			ConnectionNamespace: connectionNamespaceEnv.Name,
 		}
 
 		// Parse TFLOPS requirement
@@ -152,6 +174,7 @@ func (m *TensorFusionPodMutator) patchTFClient(pod *corev1.Pod, tfReq []TFReq) (
 		for i := range pod.Spec.Containers {
 			container := &pod.Spec.Containers[i]
 			if container.Name == req.ContainerName {
+				// patch from config
 				containerJSON, err := json.Marshal(container)
 				if err != nil {
 					return nil, fmt.Errorf("marshal container: %v", err)
@@ -165,10 +188,26 @@ func (m *TensorFusionPodMutator) patchTFClient(pod *corev1.Pod, tfReq []TFReq) (
 				if err != nil {
 					return nil, fmt.Errorf("apply strategic merge patch to container: %v", err)
 				}
-
 				if err := json.Unmarshal(patchedJSON, container); err != nil {
 					return nil, fmt.Errorf("unmarshal patched container: %v", err)
 				}
+
+				// add connection env
+				connectionName := fmt.Sprintf("%s-tf-worker-%s", pod.GenerateName+container.Name, shortuuid.New())
+				connectionNamespace := pod.Namespace
+
+				container.Env = append(container.Env, corev1.EnvVar{
+					Name:  constants.ConnectionNameEnv,
+					Value: connectionName,
+				})
+				container.Env = append(container.Env, corev1.EnvVar{
+					Name:  constants.ConnectionNamespaceEnv,
+					Value: connectionNamespace,
+				})
+				container.Env = append(container.Env, corev1.EnvVar{
+					Name:  constants.GetConnectionURLEnv,
+					Value: path.Join(m.Config.OperatorEndpoint, fmt.Sprintf("/api/connection?name=%s,namespace=%s", connectionName, connectionNamespace)),
+				})
 			}
 		}
 	}

internal/webhook/v1/pod_webhook_test.go

@@ -78,6 +78,16 @@ var _ = Describe("TensorFusionPodMutator", func() {
 						{
 							Name:  "main",
 							Image: "test-image",
+							Env: []corev1.EnvVar{
+								{
+									Name:  constants.ConnectionNameEnv,
+									Value: "conn1",
+								},
+								{
+									Name:  constants.ConnectionNamespaceEnv,
+									Value: "ns",
+								},
+							},
 						},
 					},
 				},
@@ -163,6 +173,16 @@ var _ = Describe("TensorFusionPodMutator", func() {
 					Containers: []corev1.Container{
 						{
 							Name: "test-container",
+							Env: []corev1.EnvVar{
+								{
+									Name:  constants.ConnectionNameEnv,
+									Value: "conn1",
+								},
+								{
+									Name:  constants.ConnectionNamespaceEnv,
+									Value: "ns",
+								},
+							},
 						},
 					},
 				},