Added GPU_PORTION env var for GPU sharing pods (#312)

Added GPU_PORTION env var for GPU sharing pods

Author: romanbaron

CHANGELOG.md

@@ -6,6 +6,9 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
 
 ## [Unreleased]
 
+### Added
+- Added GPU_PORTION env var for GPU sharing pods
+
 ### Changed
 - Changed RUNAI-VISIBLE-DEVICES key in GPU sharing configmap to NVIDIA_VISIBLE_DEVICES
 

deployments/crds/internal/scheduling.run.ai_bindrequests.yaml

@@ -56,7 +56,7 @@ spec:
                     type: integer
                   portion:
                     description: |-
-                      This is the portion size that the pod will receive from each connected gpu device
+                      This is the portion size that the pod will receive from each connected GPU device
                       This is a serialized float that should be written as a decimal point number.
                     type: string
                 type: object

hack/update-client.sh

@@ -26,7 +26,7 @@ kube::codegen::gen_client \
 
 rm -f generate-dep.go && rm -r vendor && go mod tidy
 
-changed_files=$(git diff --name-only | grep v1alpha2)
+changed_files=$(git diff --name-only | grep pkg/apis/client | grep v1alpha2)
 ${SDK_HACK_DIR}/replace_headers.sh \
   ${SDK_HACK_DIR}/boilerplate.go.txt \
   ${changed_files}

pkg/apis/scheduling/v1alpha2/bindrequest_types.go

@@ -37,7 +37,7 @@ type ReceivedGPU struct {
 	// Count is the amount of GPUs devices that were received
 	Count int `json:"count,omitempty"`
 
-	// This is the portion size that the pod will receive from each connected gpu device
+	// This is the portion size that the pod will receive from each connected GPU device
 	// This is a serialized float that should be written as a decimal point number.
 	Portion string `json:"portion,omitempty"`
 }

pkg/binder/binding/default_binder_test.go

@@ -118,10 +118,10 @@ func TestBindApplyResourceReceivedType(t *testing.T) {
 							},
 						},
 						{
-							Name: common.NumOfGpusEnvVar,
+							Name: common.GPUPortion,
 							ValueFrom: &v1.EnvVarSource{
 								ConfigMapKeyRef: &v1.ConfigMapKeySelector{
-									Key: common.NumOfGpusEnvVar,
+									Key: common.GPUPortion,
 									LocalObjectReference: v1.LocalObjectReference{
 										Name: "my-config-0",
 									},
@@ -159,7 +159,10 @@ func TestBindApplyResourceReceivedType(t *testing.T) {
 			SelectedNode:         "my-node",
 			ReceivedResourceType: common.ReceivedTypeFraction,
 			SelectedGPUGroups:    []string{"group1"},
-			ReceivedGPU:          &v1alpha2.ReceivedGPU{Count: 1, Portion: "1"},
+			ReceivedGPU: &v1alpha2.ReceivedGPU{
+				Count:   1,
+				Portion: "1",
+			},
 		},
 	}
 

pkg/binder/binding/fraction_binder_test.go

@@ -217,7 +217,8 @@ var _ = Describe("FractionBinder", func() {
 					Fail(fmt.Sprintf("Failed to read configmap: %v", err))
 				} else {
 					Expect(configMap.Data[common.NvidiaVisibleDevices]).To(Equal(testData.gpuIndexByGroupIndex))
-					Expect(configMap.Data[common.NumOfGpusEnvVar]).To(Equal("0.5"))
+					Expect(configMap.Data[common.NumOfGpusEnvVarBC]).To(Equal("0.5"))
+					Expect(configMap.Data[common.GPUPortion]).To(Equal("0.5"))
 				}
 			})
 		}

pkg/binder/common/constants.go

@@ -5,7 +5,7 @@ package common
 
 const (
 	NvidiaVisibleDevices = "NVIDIA_VISIBLE_DEVICES"
-	NumOfGpusEnvVar      = "RUNAI_NUM_OF_GPUS"
+	GPUPortion           = "GPU_PORTION"
 	ReceivedTypeFraction = "Fraction"
 	ReceivedTypeRegular  = "Regular"
 )

pkg/binder/common/gpu_access.go

@@ -16,10 +16,11 @@ import (
 )
 
 const (
-	visibleDevicesBC = "RUNAI-VISIBLE-DEVICES" // Deprecated, this value was replaced with NVIDIA_VISIBLE_DEVICES
+	visibleDevicesBC  = "RUNAI-VISIBLE-DEVICES" // Deprecated, this value was replaced with NVIDIA_VISIBLE_DEVICES
+	NumOfGpusEnvVarBC = "RUNAI_NUM_OF_GPUS"     // Deprecated, please use GPU_PORTION env var instead
 )
 
-func AddVisibleDevicesEnvVars(container *v1.Container, sharedGpuConfigMapName string) {
+func AddGPUSharingEnvVars(container *v1.Container, sharedGpuConfigMapName string) {
 	AddEnvVarToContainer(container, v1.EnvVar{
 		Name: NvidiaVisibleDevices,
 		ValueFrom: &v1.EnvVarSource{
@@ -33,10 +34,22 @@ func AddVisibleDevicesEnvVars(container *v1.Container, sharedGpuConfigMapName st
 	})
 
 	AddEnvVarToContainer(container, v1.EnvVar{
-		Name: NumOfGpusEnvVar,
+		Name: NumOfGpusEnvVarBC,
 		ValueFrom: &v1.EnvVarSource{
 			ConfigMapKeyRef: &v1.ConfigMapKeySelector{
-				Key: NumOfGpusEnvVar,
+				Key: NumOfGpusEnvVarBC,
+				LocalObjectReference: v1.LocalObjectReference{
+					Name: sharedGpuConfigMapName,
+				},
+			},
+		},
+	})
+
+	AddEnvVarToContainer(container, v1.EnvVar{
+		Name: GPUPortion,
+		ValueFrom: &v1.EnvVarSource{
+			ConfigMapKeyRef: &v1.ConfigMapKeySelector{
+				Key: GPUPortion,
 				LocalObjectReference: v1.LocalObjectReference{
 					Name: sharedGpuConfigMapName,
 				},
@@ -69,6 +82,8 @@ func SetNvidiaVisibleDevices(
 		return err
 	}
 	updateFunc = func(data map[string]string) error {
+		// BC for pods that were created with NVIDIA_VISIBLE_DEVICES env var
+		// with value from RUNAI-VISIBLE-DEVICES entry in GPU sharing configmap
 		if _, found := data[visibleDevicesBC]; found {
 			data[visibleDevicesBC] = visibleDevicesValue
 		}
@@ -83,12 +98,13 @@ func SetNvidiaVisibleDevices(
 	return nil
 }
 
-func SetNumOfGPUDevices(
+func SetGPUPortion(
 	ctx context.Context, kubeClient client.Client, pod *v1.Pod, containerRef *gpusharingconfigmap.PodContainerRef,
-	numOfGPUs string,
+	gpuPortionStr string,
 ) error {
 	updateFunc := func(data map[string]string) error {
-		data[NumOfGpusEnvVar] = numOfGPUs
+		data[NumOfGpusEnvVarBC] = gpuPortionStr
+		data[GPUPortion] = gpuPortionStr
 		return nil
 	}
 	capabilitiesMapName, err := gpusharingconfigmap.ExtractCapabilitiesConfigMapName(pod, containerRef)
@@ -98,7 +114,7 @@ func SetNumOfGPUDevices(
 
 	err = UpdateConfigMapEnvironmentVariable(ctx, kubeClient, pod, capabilitiesMapName, updateFunc)
 	if err != nil {
-		return fmt.Errorf("failed to update gpu sharing configmap for pod <%s/%s>: %v",
+		return fmt.Errorf("failed to update GPU_PORTION value in gpu sharing configmap for pod <%s/%s>: %v",
 			pod.Namespace, pod.Name, err)
 	}
 	return nil

pkg/binder/plugins/gpusharing/gpu_sharing.go

@@ -74,7 +74,7 @@ func (p *GPUSharing) Mutate(pod *v1.Pod) error {
 		return err
 	}
 
-	common.AddVisibleDevicesEnvVars(containerRef.Container, capabilitiesConfigMapName)
+	common.AddGPUSharingEnvVars(containerRef.Container, capabilitiesConfigMapName)
 	common.SetConfigMapVolume(pod, capabilitiesConfigMapName)
 	common.AddDirectEnvVarsConfigMapSource(containerRef.Container, directEnvVarsMapName)
 
@@ -116,8 +116,7 @@ func (p *GPUSharing) PreBind(
 		return err
 	}
 
-	numOfGPUDevices := fmt.Sprintf("%v", bindRequest.Spec.ReceivedGPU.Portion)
-	return common.SetNumOfGPUDevices(ctx, p.kubeClient, pod, containerRef, numOfGPUDevices)
+	return common.SetGPUPortion(ctx, p.kubeClient, pod, containerRef, bindRequest.Spec.ReceivedGPU.Portion)
 }
 
 func (p *GPUSharing) createCapabilitiesConfigMapIfMissing(ctx context.Context, pod *v1.Pod,