address review comments

Author: meetrajvala

launcher/internal/gpu/driverinstaller.go

@@ -23,6 +23,7 @@ import (
 const (
 	CCModeON             CCMode = "ON"
 	CCModeOFF            CCMode = "OFF"
+	CCModeDevTools       CCMode = "DEVTOOLS"
 	installerContainerID        = "tee-gpu-driver-installer-container"
 	installerSnapshotID         = "tee-gpu-driver-installer-snapshot"
 )
@@ -36,7 +37,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)
@@ -90,8 +91,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); err != nil {
 		return err
 	}
 
@@ -195,7 +195,8 @@ func getInstallerImageReference() (string, error) {
 	return installerImageRef, nil
 }
 
-func verifyInstallerImageDigest(installerDigest string) error {
+func verifyInstallerImageDigest(image containerd.Image) error {
+	installerDigest := image.Target().Digest.String()
 	imageDigestBytes, err := os.ReadFile(InstallerImageDigestFile)
 	if err != nil {
 		return fmt.Errorf("failed to get the cos-gpu-installer image digest: %v", err)
@@ -239,16 +240,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 +260,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 +287,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")