cos-gpu-installer: Switch generic NVIDIA driver installer download url
For the OSS GPU driver installation the generic NVIDIA installer is
required for installing the userspace components. It is being pulled
from us.download.nvidia... These drivers are mirrored into NVIDIAs GCS
buckets.
BUG=b/292251267
TEST=installed GPU drivers on older and newer COS versions (w and w/o
OSS drivers) on multiple K80 and T4 GPU type
RELEASE_NOTE=Switch generic NVIDIA driver installer download url in
cos-gpu-installer
Change-Id: Ib6eb580d6f52740c5f35a52d8688940fd87198cb
Reviewed-on: https://cos-review.googlesource.com/c/cos/tools/+/52529
Cloud-Build: GCB Service account <228075978874@cloudbuild.gserviceaccount.com>
Tested-by: Arnav Kansal <rnv@google.com>
Reviewed-by: Oleksandr Tymoshenko <ovt@google.com>
diff --git a/src/cmd/cos_gpu_installer/internal/commands/install.go b/src/cmd/cos_gpu_installer/internal/commands/install.go
index 4c8d9d5..a8f639f 100644
--- a/src/cmd/cos_gpu_installer/internal/commands/install.go
+++ b/src/cmd/cos_gpu_installer/internal/commands/install.go
@@ -26,11 +26,10 @@
)
const (
- grepFound = 0
- hostRootPath = "/root"
- kernelSrcDir = "/build/usr/src/linux"
- toolchainPkgDir = "/build/cos-tools"
- installerURLTemplate = "https://us.download.nvidia.com/tesla/%[1]s/NVIDIA-Linux-x86_64-%[1]s.run"
+ grepFound = 0
+ hostRootPath = "/root"
+ kernelSrcDir = "/build/usr/src/linux"
+ toolchainPkgDir = "/build/cos-tools"
)
type GPUType int
@@ -421,8 +420,7 @@
}
defer func() { callback <- 0 }()
- installerURL := fmt.Sprintf(installerURLTemplate, c.driverVersion)
- installerFile, err := installer.DownloadToInstallDir(installerURL, "Downloading driver installer")
+ installerFile, err := installer.DownloadGenericDriverInstaller(c.driverVersion)
if err != nil {
return err
}
diff --git a/src/cmd/cos_gpu_installer/internal/installer/installer.go b/src/cmd/cos_gpu_installer/internal/installer/installer.go
index ef82b4e..ef2b8b8 100644
--- a/src/cmd/cos_gpu_installer/internal/installer/installer.go
+++ b/src/cmd/cos_gpu_installer/internal/installer/installer.go
@@ -34,6 +34,7 @@
prebuiltModuleTemplate = "nvidia-drivers-%s.tgz"
DefaultVersion = "default"
LatestVersion = "latest"
+ installerURLTemplate = "https://storage.googleapis.com/nvidia-drivers-%[1]s-public/tesla/%[2]s/NVIDIA-Linux-x86_64-%[2]s.run"
)
var (
@@ -671,3 +672,23 @@
prebuiltModulesArtifactPath := fmt.Sprintf(prebuiltModuleTemplate, driverVersion)
return downloader.ArtifactExists(prebuiltModulesArtifactPath)
}
+
+func getGenericDriverInstallerURL(driverVersion string) (string, error) {
+ metadataZone, err := utils.GetGCEMetadata("zone")
+ if err != nil {
+ return "", errors.Wrap(err, "failed to get GCE metadata zone")
+ }
+ downloadLocation := getInstallerDownloadLocation(metadataZone)
+
+ return fmt.Sprintf(installerURLTemplate, downloadLocation, driverVersion), nil
+}
+
+// DownloadGenericDriverInstaller downloads the generic GPU driver installer given driver version.
+func DownloadGenericDriverInstaller(driverVersion string) (string, error) {
+ log.Infof("Downloading GPU driver installer version %s", driverVersion)
+ downloadURL, err := getGenericDriverInstallerURL(driverVersion)
+ if err != nil {
+ return "", errors.Wrap(err, "failed to get driver installer URL")
+ }
+ return DownloadToInstallDir(downloadURL, "GPU driver installer")
+}
diff --git a/src/cmd/cos_gpu_installer/internal/installer/installer_test.go b/src/cmd/cos_gpu_installer/internal/installer/installer_test.go
index d283761..221ca08 100644
--- a/src/cmd/cos_gpu_installer/internal/installer/installer_test.go
+++ b/src/cmd/cos_gpu_installer/internal/installer/installer_test.go
@@ -45,3 +45,14 @@
t.Errorf("Unexpected return, want: %s, got: %s", expectedRet, ret)
}
}
+
+func TestGetGenericDriverInstallerURL(t *testing.T) {
+ ret, err := getGenericDriverInstallerURL("525.125.06")
+ if err != nil {
+ t.Errorf("Unexpected err, want: nil, got: %v", err)
+ }
+ expectedRet := "https://storage.googleapis.com/nvidia-drivers-us-public/tesla/525.125.06/NVIDIA-Linux-x86_64-525.125.06.run"
+ if ret != expectedRet {
+ t.Errorf("Unexpected return, want: %s, got: %s", expectedRet, ret)
+ }
+}
