fix: add shm device for shared limiter among process, communicate with hypervisor (#276)

Author: Code2Life

internal/constants/constants.go

@@ -159,7 +159,7 @@ const (
 	ProvisionerNamePlaceholder = "__GPU_NODE_RESOURCE_NAME__"
 )
 
-const TFDataPath = "/tmp/tensor-fusion/data"
+const TFDataPath = "/run/tensor-fusion"
 const DataVolumeName = "tf-data"
 const TensorFusionPoolManualCompaction = Domain + "/manual-compaction"
 const AlertJobName = "tensor-fusion"

internal/constants/env.go

@@ -93,6 +93,8 @@ const (
 	LdPreloadEnv     = "LD_PRELOAD"
 	LdPreloadLimiter = "/home/app/libcuda_limiter.so"
 
+	SharedMemResName = "tensor-fusion.ai/shm"
+
 	// disable GPU limiter, for emergency use
 	DisableGpuLimiterEnv = "DISABLE_GPU_LIMITER"
 	// directly forward CUDA calls to GPU driver in nGPU mode, for emergency use

internal/utils/compose.go

@@ -243,6 +243,8 @@ func AddTFDefaultClientConfBeforePatch(
 				Value: constants.NGPUPathValue,
 			})
 
+			pod.Spec.Containers[injectContainerIndex].Resources.Limits[constants.SharedMemResName] = resource.MustParse("1")
+
 			// disable GPU limiter killer switch
 			if pod.Annotations[constants.DisableFeaturesAnnotation] != "" {
 				features := strings.Split(pod.Annotations[constants.DisableFeaturesAnnotation], ",")
@@ -334,13 +336,29 @@ func AddTFHypervisorConfAfterTemplate(ctx context.Context, spec *v1.PodSpec, poo
 		},
 	})
 
+	composeHypervisorInitContainer(spec, pool)
 	composeHypervisorContainer(spec, pool)
 
 	if enableVector {
 		composeVectorContainer(spec, pool)
 	}
 }
 
+func composeHypervisorInitContainer(spec *v1.PodSpec, pool *tfv1.GPUPool) {
+	spec.InitContainers = append(spec.InitContainers, v1.Container{
+		Name:    "init-shm",
+		Image:   pool.Spec.ComponentConfig.Hypervisor.Image,
+		Command: []string{"hypervisor", "mount-shm"},
+		VolumeMounts: []v1.VolumeMount{
+			{
+				Name:      constants.DataVolumeName,
+				ReadOnly:  false,
+				MountPath: constants.TFDataPath,
+			},
+		},
+	})
+}
+
 func composeHypervisorContainer(spec *v1.PodSpec, pool *tfv1.GPUPool) {
 	spec.Containers[0].VolumeMounts = append(spec.Containers[0].VolumeMounts, v1.VolumeMount{
 		Name:      constants.DataVolumeName,
@@ -553,6 +571,8 @@ func AddWorkerConfAfterTemplate(ctx context.Context, spec *v1.PodSpec, workerCon
 		},
 	})
 
+	spec.Containers[0].Resources.Limits[constants.SharedMemResName] = resource.MustParse("1")
+
 	// Add volume from host for CUDA hot migration and snapshot
 	spec.Volumes = append(spec.Volumes, v1.Volume{
 		Name: constants.DataVolumeName,