Address reviewers requested changes

Signed-off-by: Carlos Eduardo Arango Gutierrez <eduardoa@nvidia.com>

controllers/object_controls.go

@@ -135,12 +135,6 @@ const (
 	DeviceListStrategyEnvName = "DEVICE_LIST_STRATEGY"
 	// CDIAnnotationPrefixEnvName is the name of the device-plugin envvar for configuring the CDI annotation prefix
 	CDIAnnotationPrefixEnvName = "CDI_ANNOTATION_PREFIX"
-	// NvidiaCtrRuntimeSocket is the EnvVar name for the Container toolkit runtime socket
-	NvidiaCtrRuntimeSocket = "RUNTIME_SOCKET"
-	// NvidiaCtrRuntimeSocket is the EnvVar name for the Container toolkit runtime config
-	NvidiaCtrRuntimeConfig = "RUNTIME_CONFIG"
-	// NvidiaCtrRuntimeSocket is the EnvVar name for the Container toolkit runtime class name
-	NvidiaCtrRuntimeName = "NVIDIA_RUNTIME_NAME"
 )
 
 // ContainerProbe defines container probe types
@@ -1111,7 +1105,7 @@ func TransformToolkit(obj *appsv1.DaemonSet, config *gpuv1.ClusterPolicySpec, n
 
 	if runtime == gpuv1.Containerd.String() {
 		// Set the runtime class name that is to be configured for containerd
-		setContainerEnv(&(obj.Spec.Template.Spec.Containers[0]), NvidiaCtrRuntimeName, getRuntimeClass(config))
+		setContainerEnv(&(obj.Spec.Template.Spec.Containers[0]), "CONTAINERD_RUNTIME_CLASS", getRuntimeClass(config))
 	}
 
 	// setup mounts for runtime config file
@@ -1121,7 +1115,13 @@ func TransformToolkit(obj *appsv1.DaemonSet, config *gpuv1.ClusterPolicySpec, n
 	}
 	sourceConfigFileName := path.Base(runtimeConfigFile)
 	// update runtime args
-	setContainerEnv(&(obj.Spec.Template.Spec.Containers[0]), NvidiaCtrRuntimeConfig, DefaultRuntimeConfigTargetDir+sourceConfigFileName)
+	if runtime == gpuv1.Containerd.String() {
+		setContainerEnv(&(obj.Spec.Template.Spec.Containers[0]), "CONTAINERD_CONFIG", DefaultRuntimeConfigTargetDir+sourceConfigFileName)
+	} else if runtime == gpuv1.Docker.String() {
+		setContainerEnv(&(obj.Spec.Template.Spec.Containers[0]), "DOCKER_CONFIG", DefaultRuntimeConfigTargetDir+sourceConfigFileName)
+	} else if runtime == gpuv1.CRIO.String() {
+		setContainerEnv(&(obj.Spec.Template.Spec.Containers[0]), "CRIO_CONFIG", DefaultRuntimeConfigTargetDir+sourceConfigFileName)
+	}
 
 	volMountConfigName := fmt.Sprintf("%s-config", runtime)
 	volMountConfig := corev1.VolumeMount{Name: volMountConfigName, MountPath: DefaultRuntimeConfigTargetDir}
@@ -1138,8 +1138,15 @@ func TransformToolkit(obj *appsv1.DaemonSet, config *gpuv1.ClusterPolicySpec, n
 		}
 		sourceSocketFileName := path.Base(runtimeSocketFile)
 		// update runtime args
-		setContainerEnv(&(obj.Spec.Template.Spec.Containers[0]), NvidiaCtrRuntimeSocket, DefaultRuntimeSocketTargetDir+sourceSocketFileName)
-
+		// update runtime args
+		if runtime == gpuv1.Containerd.String() {
+			setContainerEnv(&(obj.Spec.Template.Spec.Containers[0]), "CONTAINERD_SOCKET", DefaultRuntimeSocketTargetDir+sourceSocketFileName)
+		} else if runtime == gpuv1.Docker.String() {
+			setContainerEnv(&(obj.Spec.Template.Spec.Containers[0]), "DOCKER_SOCKET", DefaultRuntimeSocketTargetDir+sourceSocketFileName)
+		} else if runtime == gpuv1.CRIO.String() {
+			runtimeArgs := " --socket " + DefaultRuntimeSocketTargetDir + sourceSocketFileName
+			setContainerEnv(&(obj.Spec.Template.Spec.Containers[0]), "RUNTIME_ARGS", runtimeArgs)
+		}
 		volMountSocketName := fmt.Sprintf("%s-socket", runtime)
 		volMountSocket := corev1.VolumeMount{Name: volMountSocketName, MountPath: DefaultRuntimeSocketTargetDir}
 		obj.Spec.Template.Spec.Containers[0].VolumeMounts = append(obj.Spec.Template.Spec.Containers[0].VolumeMounts, volMountSocket)