fix: inject env NVSHARE_MANAGED_MEMORY (#127)

beclab · Jan 10, 2025 · bcad95f · bcad95f
1 parent 022eb24
commit bcad95f
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 5 deletions.
diff --git a/pkg/apiserver/handler_webhook.go b/pkg/apiserver/handler_webhook.go
@@ -11,6 +11,7 @@ import (
 	"bytetrade.io/web3os/app-service/pkg/appinstaller"
 	"bytetrade.io/web3os/app-service/pkg/constants"
 	"bytetrade.io/web3os/app-service/pkg/provider"
+	"bytetrade.io/web3os/app-service/pkg/upgrade"
 	"bytetrade.io/web3os/app-service/pkg/users/userspace"
 	"bytetrade.io/web3os/app-service/pkg/utils"
 	"bytetrade.io/web3os/app-service/pkg/webhook"
@@ -296,8 +297,21 @@ func (h *Handler) gpuLimitMutate(ctx context.Context, req *admissionv1.Admission
 	if !gpuRequired.IsZero() && GPUType == "" {
 		return h.sidecarWebhook.AdmissionError(req.UID, api.ErrGPUNodeNotFound)
 	}
+	terminus, err := upgrade.GetTerminusVersion(ctx, h.ctrlClient)
+	if err != nil {
+		return h.sidecarWebhook.AdmissionError(req.UID, err)
+	}
+	nvshareManagedMemory := ""
+	if terminus.Spec.Settings != nil {
+		nvshareManagedMemory = terminus.Spec.Settings[constants.EnvNvshareManagedMemory]
+	}
+
+	envs := []webhook.EnvKeyValue{{
+		Key:   constants.EnvNvshareManagedMemory,
+		Value: nvshareManagedMemory,
+	}}
 
-	patchBytes, err := webhook.CreatePatchForDeployment(tpl, req.Namespace, gpuRequired, GPUType)
+	patchBytes, err := webhook.CreatePatchForDeployment(tpl, req.Namespace, gpuRequired, GPUType, envs)
 	if err != nil {
 		klog.Errorf("create patch error %v", err)
 		return h.sidecarWebhook.AdmissionError(req.UID, err)

diff --git a/pkg/webhook/webhook.go b/pkg/webhook/webhook.go
@@ -420,16 +420,21 @@ var resourcePath = "/spec/template/spec/containers/%d/resources/limits"
 var envPath = "/spec/template/spec/containers/%d/env/%s"
 var runtimeClassPath = "/spec/template/spec/runtimeClassName"
 
+type EnvKeyValue struct {
+	Key   string
+	Value string
+}
+
 // CreatePatchForDeployment add gpu env for deployment and returns patch bytes.
-func CreatePatchForDeployment(tpl *corev1.PodTemplateSpec, namespace string, gpuRequired *resource.Quantity, typeKey string) ([]byte, error) {
-	patches, err := addResourceLimits(tpl, namespace, gpuRequired, typeKey)
+func CreatePatchForDeployment(tpl *corev1.PodTemplateSpec, namespace string, gpuRequired *resource.Quantity, typeKey string, envKeyValues []EnvKeyValue) ([]byte, error) {
+	patches, err := addResourceLimits(tpl, namespace, gpuRequired, typeKey, envKeyValues)
 	if err != nil {
 		return []byte{}, err
 	}
 	return json.Marshal(patches)
 }
 
-func addResourceLimits(tpl *corev1.PodTemplateSpec, namespace string, gpuRequired *resource.Quantity, typeKey string) (patch []patchOp, err error) {
+func addResourceLimits(tpl *corev1.PodTemplateSpec, namespace string, gpuRequired *resource.Quantity, typeKey string, envKeyValues []EnvKeyValue) (patch []patchOp, err error) {
 	if typeKey == constants.NvidiaGPU || typeKey == constants.NvshareGPU {
 		if tpl.Spec.RuntimeClassName != nil {
 			patch = append(patch, patchOp{
@@ -523,8 +528,25 @@ func addResourceLimits(tpl *corev1.PodTemplateSpec, namespace string, gpuRequire
 					}
 				}
 			}
-
 		}
+		envNames := make([]string, 0)
+		for envIdx, env := range container.Env {
+			for _, e := range envKeyValues {
+				if e.Value == "" {
+					continue
+				}
+				if env.Name == e.Key {
+					envNames = append(envNames, env.Name)
+					patch = append(patch, genPatchesForEnv(constants.PatchOpReplace, i, envIdx, e.Key, e.Value)...)
+				}
+			}
+		}
+		for _, env := range envKeyValues {
+			if !funk.Contains(envNames, env.Key) {
+				patch = append(patch, genPatchesForEnv(constants.PatchOpAdd, i, -1, env.Key, env.Value)...)
+			}
+		}
+
 	}
 
 	return patch, nil