address review comments

meetrajvala · meetrajvala · commit 75e8f5b212be · 2025-04-16T09:49:49.000Z
diff --git a/launcher/image/preload.sh b/launcher/image/preload.sh
@@ -132,7 +132,23 @@ set_gpu_driver_ref_values() {
 
   mkdir ${GPU_REF_VALUES_PATH}
   cos_gpu_installer_image_ref=$(cos-extensions list -- --gpu-installer)
+  if [ -z "${cos_gpu_installer_image_ref}" ]; then
+    echo "Error: cos-extensions list returned an empty image reference." >&2
+    return 1
+  fi
+
   cos_gpu_installer_image_digest=$(get_cos_gpu_installer_image_digest ${cos_gpu_installer_image_ref})
+  if [ -z "${cos_gpu_installer_image_ref}" ]; then
+    echo "Error: get_cos_gpu_installer_image_digest returned an empty or invalid digest for: ${cos_gpu_installer_image_ref}." >&2
+    return 1
+  fi
+
+  image_digest_hex_part=$(echo "${cos_gpu_installer_image_digest}" | sed 's/^sha256://' | tr -d '\n')
+  # Check for the expected length of the SHA256 digest (64 hex characters)
+  if [ ${#image_digest_hex_part} -ne 64 ]; then
+    echo "Error: cos_gpu_installer image digest has an unexpected length: ${#image_digest_hex_part}, Expected 64." >&2
+    return 1
+  fi
   
   echo ${cos_gpu_installer_image_ref} >> ${COS_GPU_INSTALLER_IMAGE_REF}
   echo ${cos_gpu_installer_image_digest} >> ${COS_GPU_INSTALLER_IMAGE_DIGEST}
diff --git a/launcher/image/test/test_gpu_driver_installation_cloudbuild.yaml b/launcher/image/test/test_gpu_driver_installation_cloudbuild.yaml
@@ -6,6 +6,34 @@ substitutions:
   '_ZONE': 'us-central1-f'
   '_WORKLOAD_IMAGE': 'us-west1-docker.pkg.dev/confidential-space-images-dev/cs-integ-test-images/gpu/cuda-vector-add:latest'
 steps:
+- name: 'gcr.io/cloud-builders/gcloud'
+  id: CreateShieldedVMWithSingleGPU
+  entrypoint: 'bash'
+  env:
+  - 'BUILD_ID=$BUILD_ID'
+  args: ['create_gpu_vm.sh','-i', '${_IMAGE_NAME}',
+          '-p', '${_IMAGE_PROJECT}',
+          '-m', 'tee-image-reference=${_WORKLOAD_IMAGE},tee-container-log-redirect=true,tee-install-gpu-driver=true',
+          '-n', '${_VM_NAME_PREFIX}-${BUILD_ID}',
+          '-z', '${_ZONE}',
+          '-v', 'n1-standard-4',
+          '-g', 'nvidia-tesla-t4',
+          '-c', '1'
+        ]
+- name: 'gcr.io/cloud-builders/gcloud'
+  id: CreateShieldedVMWithMultipleGPU
+  entrypoint: 'bash'
+  env:
+  - 'BUILD_ID=$BUILD_ID'
+  args: ['create_gpu_vm.sh','-i', '${_IMAGE_NAME}',
+          '-p', '${_IMAGE_PROJECT}',
+          '-m', 'tee-image-reference=${_WORKLOAD_IMAGE},tee-container-log-redirect=true,tee-install-gpu-driver=true',
+          '-n', '${_VM_NAME_PREFIX}-${BUILD_ID}-mul',
+          '-z', '${_ZONE}',
+          '-v', 'n1-standard-4',
+          '-g', 'nvidia-tesla-t4',
+          '-c', '2'
+        ]
 - name: 'gcr.io/cloud-builders/gcloud'
   id:  CreateTDXCVMWithUnsupportedGPU
   entrypoint: 'bash'
@@ -46,6 +74,14 @@ steps:
           '-c', '1',
           '-x', 'TDX',
         ]
+- name: 'gcr.io/cloud-builders/gcloud'
+  id: SingleGpuWorkloadTest
+  entrypoint: 'bash'
+  args: ['scripts/gpu/test_gpu_workload.sh', '${_VM_NAME_PREFIX}-${BUILD_ID}', '${_ZONE}']
+- name: 'gcr.io/cloud-builders/gcloud'
+  id: MultipleGpuWorkloadTest
+  entrypoint: 'bash'
+  args: ['scripts/gpu/test_gpu_workload.sh', '${_VM_NAME_PREFIX}-${BUILD_ID}-mul', '${_ZONE}']
 - name: 'gcr.io/cloud-builders/gcloud'
   id: UnsupportedGpuWorkloadTest
   entrypoint: 'bash'
@@ -58,6 +94,18 @@ steps:
   id: ConfidentialGpuWorkloadTest
   entrypoint: 'bash'
   args: ['scripts/gpu/test_gpu_workload.sh', '${_VM_NAME_PREFIX}-${BUILD_ID}-tdxcvm-cgpu', 'us-east5-a']
+- name: 'gcr.io/cloud-builders/gcloud'
+  id: SingleGpuCleanUp
+  entrypoint: 'bash'
+  env:
+  - 'CLEANUP=$_CLEANUP'
+  args: ['cleanup.sh', '${_VM_NAME_PREFIX}-${BUILD_ID}', '${_ZONE}']
+- name: 'gcr.io/cloud-builders/gcloud'
+  id: MultipleGpuCleanUp
+  entrypoint: 'bash'
+  env:
+  - 'CLEANUP=$_CLEANUP'
+  args: ['cleanup.sh', '${_VM_NAME_PREFIX}-${BUILD_ID}-mul', '${_ZONE}']
 - name: 'gcr.io/cloud-builders/gcloud'
   id: UnsupportedGpuVmCleanUp
   entrypoint: 'bash'
diff --git a/launcher/internal/gpu/driverinstaller.go b/launcher/internal/gpu/driverinstaller.go
@@ -23,11 +23,16 @@ import (
 const (
 	CCModeON             CCMode = "ON"
 	CCModeOFF            CCMode = "OFF"
+	CCModeDevTools       CCMode = "DEVTOOLS"
 	installerContainerID        = "tee-gpu-driver-installer-container"
 	installerSnapshotID         = "tee-gpu-driver-installer-snapshot"
 )
 
 var supportedCGPUTypes = []deviceinfo.GPUType{
+	deviceinfo.L4,
+	deviceinfo.T4,
+	deviceinfo.A100_40GB,
+	deviceinfo.A100_80GB,
 	deviceinfo.H100,
 }
 
@@ -36,7 +41,7 @@ type CCMode string
 
 func (ccm CCMode) isValid() error {
 	switch ccm {
-	case CCModeOFF, CCModeON:
+	case CCModeOFF, CCModeON, CCModeDevTools:
 		return nil
 	}
 	return fmt.Errorf("invalid gpu cc mode: %s", ccm)
@@ -74,11 +79,11 @@ func (di *DriverInstaller) InstallGPUDrivers(ctx context.Context) error {
 	}
 
 	if !gpuType.OpenSupported() {
-		return fmt.Errorf("unsupported GPU type %s, please retry with one of the supported confidential GPU types: %v", gpuType.String(), supportedCGPUTypes)
+		return fmt.Errorf("unsupported open sourced kernel modules for GPU type %s, please retry with one of the supported GPU types: %v", gpuType.String(), supportedCGPUTypes)
 	}
 
 	ctx = namespaces.WithNamespace(ctx, namespaces.Default)
-	installerImageRef, err := getInstallerImageReference()
+	installerImageRef, err := getInstallerImageReference(InstallerImageRefFile)
 	if err != nil {
 		di.logger.Error(fmt.Sprintf("failed to get the installer container image reference: %v", err))
 		return err
@@ -90,8 +95,7 @@ func (di *DriverInstaller) InstallGPUDrivers(ctx context.Context) error {
 		return fmt.Errorf("failed to pull installer image: %v", err)
 	}
 
-	installerDigest := image.Target().Digest.String()
-	if err := verifyInstallerImageDigest(installerDigest); err != nil {
+	if err := verifyInstallerImageDigest(image, InstallerImageDigestFile); err != nil {
 		return err
 	}
 
@@ -178,25 +182,27 @@ func (di *DriverInstaller) InstallGPUDrivers(ctx context.Context) error {
 		if err = setGPUStateToReady(); err != nil {
 			return fmt.Errorf("failed to set the GPU state to ready: %v", err)
 		}
-	} else {
-		return fmt.Errorf("confidential compute is not enabled for the gpu type %s", gpuType)
 	}
 
 	di.logger.Info("GPU driver installation completed successfully")
 	return nil
 }
 
-func getInstallerImageReference() (string, error) {
-	imageRefBytes, err := os.ReadFile(InstallerImageRefFile)
+func getInstallerImageReference(installerImageRefFile string) (string, error) {
+	imageRefBytes, err := os.ReadFile(installerImageRefFile)
 	if err != nil {
 		return "", fmt.Errorf("failed to get the cos-gpu-installer version: %v", err)
 	}
 	installerImageRef := strings.TrimSpace(string(imageRefBytes))
+	if len(installerImageRef) < 1 {
+		return "", fmt.Errorf("empty value of cos-gpu-installer image reference")
+	}
 	return installerImageRef, nil
 }
 
-func verifyInstallerImageDigest(installerDigest string) error {
-	imageDigestBytes, err := os.ReadFile(InstallerImageDigestFile)
+func verifyInstallerImageDigest(image containerd.Image, referenceDigestFile string) error {
+	installerDigest := image.Target().Digest.String()
+	imageDigestBytes, err := os.ReadFile(referenceDigestFile)
 	if err != nil {
 		return fmt.Errorf("failed to get the cos-gpu-installer image digest: %v", err)
 	}
@@ -239,16 +245,16 @@ func setGPUStateToReady() error {
 }
 
 func isGPUCCModeEnabled() (bool, error) {
-	ccMode, err := GetGPUCCMode()
+	ccMode, err := QueryCCMode()
 	if err != nil {
 		return false, err
 	}
 	return ccMode == CCModeON, nil
 }
 
-// GetGPUCCMode executes nvidia-smi to determine the current Confidential Computing (CC) mode status of the GPU.
-// It returns the CC mode ("ON" or "OFF") and an error if the command fails or if the output cannot be parsed.
-func GetGPUCCMode() (CCMode, error) {
+// QueryCCMode executes nvidia-smi to determine the current Confidential Computing (CC) mode status of the GPU.
+// If DEVTOOLS mode is enabled, it would override CC mode as DEVTOOLS. DEVTOOLS mode would be enabled only when CC mode is ON.
+func QueryCCMode() (CCMode, error) {
 	// Run nvidia-smi conf-compute command to get the confidential computing mode status.
 	nvidiaSmiCmd := fmt.Sprintf("%s/bin/nvidia-smi", InstallationHostDir)
 	ccModeOutput, err := exec.Command(nvidiaSmiCmd, "conf-compute", "-f").Output()
@@ -259,7 +265,17 @@ func GetGPUCCMode() (CCMode, error) {
 	if err != nil {
 		return "", err
 	}
-	return CCMode(ccMode), nil
+
+	devToolsEnabled, err := isDevToolsModeEnabled()
+	if err != nil {
+		return "", err
+	}
+
+	if devToolsEnabled {
+		ccMode = CCModeDevTools
+	}
+
+	return ccMode, nil
 }
 
 func parseCCStatus(output string) (CCMode, error) {
@@ -276,6 +292,21 @@ func parseCCStatus(output string) (CCMode, error) {
 	return ccMode, nil
 }
 
+func isDevToolsModeEnabled() (bool, error) {
+	nvidiaSmiCmd := fmt.Sprintf("%s/bin/nvidia-smi", InstallationHostDir)
+	output, err := exec.Command(nvidiaSmiCmd, "conf-compute", "-d").Output()
+	if err != nil {
+		return false, err
+	}
+	re := regexp.MustCompile(`DevTools Mode:\s*(ON|OFF)`)
+	match := re.FindStringSubmatch(string(output))
+
+	if len(match) < 2 {
+		return false, fmt.Errorf("DevTools mode not found in output: %s", output)
+	}
+	return match[1] == "ON", nil
+}
+
 func launchNvidiaPersistencedProcess(logger logging.Logger) error {
 	nvidiaPersistencedCmd := fmt.Sprintf("%s/bin/nvidia-persistenced", InstallationHostDir)
 	logger.Info("Starting nvidia-persistenced process")
diff --git a/launcher/internal/gpu/driverinstaller_test.go b/launcher/internal/gpu/driverinstaller_test.go
@@ -1,9 +1,79 @@
 package gpu
 
 import (
+	"os"
+	"strings"
 	"testing"
 )
 
+func TestGetInstallerImageReference(t *testing.T) {
+	tests := []struct {
+		name        string
+		fileContent string
+		filePath    string
+		wantRef     string
+		wantErr     bool
+		errSubstr   string
+	}{
+		{
+			name:        "Successful read",
+			fileContent: "gcr.io/google-containers/cos-gpu-installer:v1.2.3",
+			filePath:    "/tmp/test_installer_ref_success.txt",
+			wantRef:     "gcr.io/google-containers/cos-gpu-installer:v1.2.3",
+			wantErr:     false,
+		},
+		{
+			name:        "Successful read with whitespace",
+			fileContent: "  gcr.io/google-containers/cos-gpu-installer:v1.2.4  \n",
+			filePath:    "/tmp/test_installer_ref_whitespace.txt",
+			wantRef:     "gcr.io/google-containers/cos-gpu-installer:v1.2.4",
+			wantErr:     false,
+		},
+		{
+			name:        "File does not exist",
+			fileContent: "",
+			filePath:    "/tmp/non_existent_file.txt",
+			wantRef:     "",
+			wantErr:     true,
+			errSubstr:   "no such file or directory",
+		},
+		{
+			name:        "Empty file",
+			fileContent: "",
+			filePath:    "/tmp/test_installer_ref_empty.txt",
+			wantRef:     "",
+			wantErr:     true,
+			errSubstr:   "empty value",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if tt.filePath != "/tmp/non_existent_file.txt" {
+				err := os.WriteFile(tt.filePath, []byte(tt.fileContent), 0644)
+				if err != nil {
+					t.Errorf("failed to write to the testfile %s", tt.filePath)
+				}
+				defer os.Remove(tt.filePath)
+			}
+
+			ref, err := getInstallerImageReference(tt.filePath)
+
+			if (err != nil) != tt.wantErr {
+				t.Errorf("getInstallerImageReference() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+			if tt.wantErr && !strings.Contains(err.Error(), tt.errSubstr) {
+				t.Errorf("getInstallerImageReference() err message %s is expected to contain %s", err.Error(), tt.errSubstr)
+				return
+			}
+			if ref != tt.wantRef {
+				t.Errorf("getInstallerImageReference() got = %v, want %v", ref, tt.wantRef)
+			}
+		})
+	}
+}
+
 func TestParseCCStatus(t *testing.T) {
 	tests := []struct {
 		name    string
