src/cmd/cos_gpu_installer/internal/commands/install_test.go

// Package commands implements subcommands of cos_gpu_installer.
package commands

import (
	"fmt"
	"os"
	"path/filepath"
	"testing"

	"cos.googlesource.com/cos/tools.git/src/cmd/cos_gpu_installer/deviceinfo"
	"cos.googlesource.com/cos/tools.git/src/cmd/cos_gpu_installer/internal/installer"
	"cos.googlesource.com/cos/tools.git/src/pkg/modules"
)

func TestSetCoherentGPUMemoryMode(t *testing.T) {
	testCases := []struct {
		name               string
		gpuType            deviceinfo.GPUType
		driverVersion      string
		initialParams      modules.ModuleParameters
		expectParamSet     bool
		expectedParamValue string
		userParamSet       bool
	}{
		{
			name:               "GB200 with R580 driver should set default",
			gpuType:            deviceinfo.GB200,
			driverVersion:      "580.01.01",
			initialParams:      modules.NewModuleParameters(),
			expectParamSet:     true,
			expectedParamValue: "NVreg_CoherentGPUMemoryMode=driver",
			userParamSet:       false,
		},
		{
			name:               "GB200 with R585 driver should set default",
			gpuType:            deviceinfo.GB200,
			driverVersion:      "585.01.01",
			initialParams:      modules.NewModuleParameters(),
			expectParamSet:     true,
			expectedParamValue: "NVreg_CoherentGPUMemoryMode=driver",
			userParamSet:       false,
		},
		{
			name:           "GB200 with pre-R580 driver should not set default",
			gpuType:        deviceinfo.GB200,
			driverVersion:  "535.01.01",
			initialParams:  modules.NewModuleParameters(),
			expectParamSet: false,
			userParamSet:   false,
		},
		{
			name:          "GB200 with user-provided value should not override",
			gpuType:       deviceinfo.GB200,
			driverVersion: "580.01.01",
			initialParams: func() modules.ModuleParameters {
				p := modules.NewModuleParameters()
				p.Set("nvidia.NVreg_CoherentGPUMemoryMode=user_value")
				return p
			}(),
			expectParamSet: false,
			userParamSet:   false,
		},
		{
			name:           "Non-GB200 GPU with R580 driver should not set default",
			gpuType:        deviceinfo.H200,
			driverVersion:  "580.01.01",
			initialParams:  modules.NewModuleParameters(),
			expectParamSet: false,
			userParamSet:   false,
		},
		{
			name:               "GB300 with R580 driver should set default",
			gpuType:            deviceinfo.GB300,
			driverVersion:      "580.01.01",
			initialParams:      modules.NewModuleParameters(),
			expectParamSet:     true,
			expectedParamValue: "NVreg_CoherentGPUMemoryMode=driver",
			userParamSet:       false,
		},
		{
			name:               "GB300 with R585 driver should set default",
			gpuType:            deviceinfo.GB300,
			driverVersion:      "585.01.01",
			initialParams:      modules.NewModuleParameters(),
			expectParamSet:     true,
			expectedParamValue: "NVreg_CoherentGPUMemoryMode=driver",
			userParamSet:       false,
		},
		{
			name:           "GB300 with pre-R580 driver should not set default",
			gpuType:        deviceinfo.GB300,
			driverVersion:  "535.01.01",
			initialParams:  modules.NewModuleParameters(),
			expectParamSet: false,
			userParamSet:   false,
		},
		{
			name:          "GB300 with user-provided value should not override",
			gpuType:       deviceinfo.GB300,
			driverVersion: "580.01.01",
			initialParams: func() modules.ModuleParameters {
				p := modules.NewModuleParameters()
				p.Set("nvidia.NVreg_CoherentGPUMemoryMode=user_value")
				return p
			}(),
			expectParamSet: false,
			userParamSet:   false,
		},
		{
			name:           "Non-GB300 GPU with R580 driver should not set default",
			gpuType:        deviceinfo.H200,
			driverVersion:  "580.01.01",
			initialParams:  modules.NewModuleParameters(),
			expectParamSet: false,
			userParamSet:   false,
		},
		{
			name:               "VR200 with R580 driver should set default",
			gpuType:            deviceinfo.VR200,
			driverVersion:      "580.01.01",
			initialParams:      modules.NewModuleParameters(),
			expectParamSet:     true,
			expectedParamValue: "NVreg_CoherentGPUMemoryMode=driver",
			userParamSet:       false,
		},
		{
			name:               "VR200 with R585 driver should set default",
			gpuType:            deviceinfo.VR200,
			driverVersion:      "585.01.01",
			initialParams:      modules.NewModuleParameters(),
			expectParamSet:     true,
			expectedParamValue: "NVreg_CoherentGPUMemoryMode=driver",
			userParamSet:       false,
		},
		{
			name:           "VR200 with pre-R580 driver should not set default",
			gpuType:        deviceinfo.VR200,
			driverVersion:  "535.01.01",
			initialParams:  modules.NewModuleParameters(),
			expectParamSet: false,
			userParamSet:   false,
		},
		{
			name:          "VR200 with user-provided value should not override",
			gpuType:       deviceinfo.VR200,
			driverVersion: "580.01.01",
			initialParams: func() modules.ModuleParameters {
				p := modules.NewModuleParameters()
				p.Set("nvidia.NVreg_CoherentGPUMemoryMode=user_value")
				return p
			}(),
			expectParamSet: false,
			userParamSet:   false,
		},
		{
			name:           "Non-VR200 GPU with R580 driver should not set default",
			gpuType:        deviceinfo.H200,
			driverVersion:  "580.01.01",
			initialParams:  modules.NewModuleParameters(),
			expectParamSet: false,
			userParamSet:   false,
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			cmd := &InstallCommand{
				driverVersion:      tc.driverVersion,
				kernelModuleParams: tc.initialParams,
			}

			if err := cmd.enableCDMM(tc.gpuType); err != nil {
				t.Fatalf("enableCDMM failed: %v", err)
			}

			paramWasSet := false
			if params, ok := cmd.kernelModuleParams["nvidia"]; ok {
				for _, param := range params {
					if param == "NVreg_CoherentGPUMemoryMode=driver" {
						paramWasSet = true
						break
					}
				}
			}

			if tc.expectParamSet != paramWasSet {
				t.Errorf("Expected param set state to be %v, but got %v", tc.expectParamSet, paramWasSet)
			}

			// Verify user-set value is preserved
			if tc.userParamSet {
				foundUserValue := false
				if params, ok := cmd.kernelModuleParams["nvidia"]; ok {
					for _, param := range params {
						if param == "NVreg_CoherentGPUMemoryMode=user_value" {
							foundUserValue = true
							break
						}
					}
				}
				if !foundUserValue {
					t.Errorf("User-provided kernel module parameter was not preserved.")
				}
			}
		})
	}
}

const (
	testBuildNumber        = "19999.44.28"
	testDriverVersion      = "580.95.05"
	testOtherDriverVersion = "570.10.01"
	// testMarkerLibrary is a library every installation holds, with the driver
	// version it belongs to as its content, so that a test can tell which
	// installation a path in the install dir leads to.
	testMarkerLibrary = "lib64/libcuda.so"
)

// writeInstallTree creates a driver installation directory holding a cache
// entry for the given build and driver version, plus the marker library.
func writeInstallTree(t *testing.T, dir, buildNumber, driverVersion string, kernelOpen bool) {
	t.Helper()
	if err := os.MkdirAll(filepath.Join(dir, filepath.Dir(testMarkerLibrary)), 0755); err != nil {
		t.Fatalf("Failed to create %s: %v", dir, err)
	}
	cache := fmt.Sprintf("BUILD_ID=%s\nDRIVER_VERSION=%s\n", buildNumber, driverVersion)
	if kernelOpen {
		cache += "KERNEL_OPEN=Y\n"
	}
	if err := os.WriteFile(filepath.Join(dir, installer.CacheFile), []byte(cache), 0644); err != nil {
		t.Fatalf("Failed to write the cache entry in %s: %v", dir, err)
	}
	if err := os.WriteFile(filepath.Join(dir, testMarkerLibrary), []byte(driverVersion), 0644); err != nil {
		t.Fatalf("Failed to write %s in %s: %v", testMarkerLibrary, dir, err)
	}
}

// installedVersion returns the driver version of the installation that dir
// leads to, as recorded in its marker library.
func installedVersion(t *testing.T, dir string) string {
	t.Helper()
	content, err := os.ReadFile(filepath.Join(dir, testMarkerLibrary))
	if err != nil {
		t.Fatalf("Failed to read %s in %s: %v", testMarkerLibrary, dir, err)
	}
	return string(content)
}

func TestActivatePreloadedInstall(t *testing.T) {
	testCases := []struct {
		name string
		// preloadedDirs are subdirectories of the install dir to populate with
		// preloaded installations of testDriverVersion.
		preloadedDirs []string
		// preloadedBuildNumber is the COS build recorded in those installations.
		preloadedBuildNumber string
		// preloadedOpen is the kernel module flavor recorded in them.
		preloadedOpen bool
		kernelOpen    bool
		wantActivated bool
	}{
		{
			name:                 "activates a matching open preloaded installation",
			preloadedDirs:        []string{testDriverVersion + "-open"},
			preloadedBuildNumber: testBuildNumber,
			preloadedOpen:        true,
			kernelOpen:           true,
			wantActivated:        true,
		},
		{
			name:                 "activates a matching proprietary preloaded installation",
			preloadedDirs:        []string{testDriverVersion + "-closed"},
			preloadedBuildNumber: testBuildNumber,
			kernelOpen:           false,
			wantActivated:        true,
		},
		{
			name:                 "activates a preloaded installation under an arbitrary directory name",
			preloadedDirs:        []string{"preloaded-driver-0"},
			preloadedBuildNumber: testBuildNumber,
			preloadedOpen:        true,
			kernelOpen:           true,
			wantActivated:        true,
		},
		{
			name:                 "ignores a preloaded installation of the other kernel module flavor",
			preloadedDirs:        []string{testDriverVersion + "-closed"},
			preloadedBuildNumber: testBuildNumber,
			kernelOpen:           true,
			wantActivated:        false,
		},
		{
			name:                 "ignores a preloaded installation from another COS build",
			preloadedDirs:        []string{testDriverVersion + "-open"},
			preloadedBuildNumber: "12345.6.7",
			preloadedOpen:        true,
			kernelOpen:           true,
			wantActivated:        false,
		},
		{
			name:          "does nothing without preloaded installations",
			kernelOpen:    true,
			wantActivated: false,
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			installDir := t.TempDir()
			// The versions proto lives next to the activated installation and must
			// survive activation.
			protoPath := filepath.Join(installDir, "gpu_driver_versions.bin")
			if err := os.WriteFile(protoPath, []byte("proto"), 0644); err != nil {
				t.Fatalf("Failed to write %s: %v", protoPath, err)
			}
			for _, dir := range tc.preloadedDirs {
				writeInstallTree(t, filepath.Join(installDir, dir), tc.preloadedBuildNumber, testDriverVersion, tc.preloadedOpen)
			}

			if err := activatePreloadedInstall(installDir, testBuildNumber, testDriverVersion, tc.kernelOpen); err != nil {
				t.Fatalf("activatePreloadedInstall(%q, %q, %q, %t) returned error: %v", installDir, testBuildNumber, testDriverVersion, tc.kernelOpen, err)
			}

			isCached, isOpen, err := installer.NewCacher(installDir, testBuildNumber, testDriverVersion).IsCached()
			gotActivated := isCached && err == nil
			if gotActivated != tc.wantActivated {
				t.Errorf("After activatePreloadedInstall, install dir cached = %t, want %t", gotActivated, tc.wantActivated)
			}
			if tc.wantActivated {
				if isOpen != tc.kernelOpen {
					t.Errorf("After activatePreloadedInstall, install dir kernel open = %t, want %t", isOpen, tc.kernelOpen)
				}
				if got := installedVersion(t, installDir); got != testDriverVersion {
					t.Errorf("Install dir holds driver version %q, want %q", got, testDriverVersion)
				}
				// The payload is linked into the install dir rather than copied,
				// with relative targets, as the install dir is bind mounted under
				// other paths.
				for _, name := range []string{filepath.Dir(testMarkerLibrary), installer.CacheFile} {
					path := filepath.Join(installDir, name)
					target, err := os.Readlink(path)
					if err != nil {
						t.Errorf("Failed to read the link %s: %v", path, err)
						continue
					}
					if filepath.IsAbs(target) {
						t.Errorf("Link %s points to %q, want a target relative to the install dir", path, target)
					}
				}
			}
			if _, err := os.Stat(protoPath); err != nil {
				t.Errorf("Expected %s to be preserved: %v", protoPath, err)
			}
			// Preloaded installations must remain available for later activations.
			for _, dir := range tc.preloadedDirs {
				if _, err := os.Stat(filepath.Join(installDir, dir, installer.CacheFile)); err != nil {
					t.Errorf("Expected preloaded installation %s to be preserved: %v", dir, err)
				}
			}
		})
	}
}

// TestActivatePreloadedInstallSwitchesVersion checks that activating another
// version repoints the install dir at it, and that what is written to the
// install dir afterwards reaches that version alone.
func TestActivatePreloadedInstallSwitchesVersion(t *testing.T) {
	installDir := t.TempDir()
	otherDir := filepath.Join(installDir, testOtherDriverVersion+"-open")
	writeInstallTree(t, otherDir, testBuildNumber, testOtherDriverVersion, true)
	writeInstallTree(t, filepath.Join(installDir, testDriverVersion+"-open"), testBuildNumber, testDriverVersion, true)

	if err := activatePreloadedInstall(installDir, testBuildNumber, testOtherDriverVersion, true); err != nil {
		t.Fatalf("activatePreloadedInstall(%q, %q, %q, true) returned error: %v", installDir, testBuildNumber, testOtherDriverVersion, err)
	}
	if err := activatePreloadedInstall(installDir, testBuildNumber, testDriverVersion, true); err != nil {
		t.Fatalf("activatePreloadedInstall(%q, %q, %q, true) returned error: %v", installDir, testBuildNumber, testDriverVersion, err)
	}
	if err := installer.NewCacher(installDir, testBuildNumber, testDriverVersion).Cache(true); err != nil {
		t.Fatalf("Failed to write the cache entry of %s: %v", installDir, err)
	}

	if got := installedVersion(t, installDir); got != testDriverVersion {
		t.Errorf("Install dir holds driver version %q, want %q", got, testDriverVersion)
	}
	if got := installedVersion(t, otherDir); got != testOtherDriverVersion {
		t.Errorf("Preloaded installation %s holds driver version %q, want %q", otherDir, got, testOtherDriverVersion)
	}
	if isCached, _, err := installer.NewCacher(otherDir, testBuildNumber, testOtherDriverVersion).IsCached(); err != nil || !isCached {
		t.Errorf("Preloaded installation %s cached for %s = %t (err %v), want true", otherDir, testOtherDriverVersion, isCached, err)
	}
}

// TestActivatePreloadedInstallDeactivates checks that a driver version without
// a preloaded installation takes down the activated one, so that it is not
// installed over a preloaded installation of another version.
func TestActivatePreloadedInstallDeactivates(t *testing.T) {
	installDir := t.TempDir()
	preloadedDir := filepath.Join(installDir, testDriverVersion+"-open")
	writeInstallTree(t, preloadedDir, testBuildNumber, testDriverVersion, true)

	if err := activatePreloadedInstall(installDir, testBuildNumber, testDriverVersion, true); err != nil {
		t.Fatalf("activatePreloadedInstall(%q, %q, %q, true) returned error: %v", installDir, testBuildNumber, testDriverVersion, err)
	}
	if err := activatePreloadedInstall(installDir, testBuildNumber, testOtherDriverVersion, true); err != nil {
		t.Fatalf("activatePreloadedInstall(%q, %q, %q, true) returned error: %v", installDir, testBuildNumber, testOtherDriverVersion, err)
	}

	entries, err := os.ReadDir(installDir)
	if err != nil {
		t.Fatalf("Failed to read %s: %v", installDir, err)
	}
	for _, entry := range entries {
		if entry.Type()&os.ModeSymlink != 0 {
			t.Errorf("Install dir holds %s, want the links of the deactivated installation to be removed", entry.Name())
		}
	}
	// Installing the driver the regular way must now write to the install dir
	// itself rather than into the preloaded installation.
	writeInstallTree(t, installDir, testBuildNumber, testOtherDriverVersion, true)
	if got := installedVersion(t, preloadedDir); got != testDriverVersion {
		t.Errorf("Preloaded installation %s holds driver version %q, want %q", preloadedDir, got, testDriverVersion)
	}
}

// TestActivatePreloadedInstallKeepsRegularInstall checks that a driver the
// install dir already holds is not replaced, since it was not put there by an
// earlier activation.
func TestActivatePreloadedInstallKeepsRegularInstall(t *testing.T) {
	installDir := t.TempDir()
	writeInstallTree(t, installDir, testBuildNumber, testOtherDriverVersion, true)
	writeInstallTree(t, filepath.Join(installDir, testDriverVersion+"-open"), testBuildNumber, testDriverVersion, true)

	if err := activatePreloadedInstall(installDir, testBuildNumber, testDriverVersion, true); err != nil {
		t.Fatalf("activatePreloadedInstall(%q, %q, %q, true) returned error: %v", installDir, testBuildNumber, testDriverVersion, err)
	}

	if got := installedVersion(t, installDir); got != testOtherDriverVersion {
		t.Errorf("Install dir holds driver version %q, want the installation %q it already held", got, testOtherDriverVersion)
	}
}