cos_gpu_installer: add flag `-signature-url`
cos-gpu-intaller:v2 uses `-driver-version` to
find driver signature in cos-tools. But for
testing, `-nvidia-installer-url` is used so we
cannot pass `-driver-version` and we cannot
release signature for testing to cos-tools.
So this CL adds a flag `-signature-url` to find
signature at a URL in test mode.
`-signature-url` can only be used together with
`-nvidia-installer-url` and `-test`.
BUG=b/161566626
Change-Id: Ibb55004c6c8f19b1cbb82ec755df10a27d86428f
Reviewed-on: https://cos-review.googlesource.com/c/cos/tools/+/30400
Cloud-Build: GCB Service account <228075978874@cloudbuild.gserviceaccount.com>
Reviewed-by: Robert Kolchmeyer <rkolchmeyer@google.com>
Tested-by: He Gao <hegao@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 07d7541..5a4a50d 100644
--- a/src/cmd/cos_gpu_installer/internal/commands/install.go
+++ b/src/cmd/cos_gpu_installer/internal/commands/install.go
@@ -40,6 +40,7 @@
gcsDownloadBucket string
gcsDownloadPrefix string
nvidiaInstallerURL string
+ signatureURL string
debug bool
test bool
}
@@ -74,11 +75,15 @@
"The GCS path prefix when downloading COS artifacts."+
"If not set then the COS version build number (e.g. 13310.1041.38) will be used.")
f.StringVar(&c.nvidiaInstallerURL, "nvidia-installer-url", "",
- "A URL to an nvidia-installer to use for driver installation. This flag is mutually exclusive with `-version`. This flag must be used with `-allow-unsigned-driver`. This flag is only for debugging.")
+ "A URL to an nvidia-installer to use for driver installation. This flag is mutually exclusive with `-version`. "+
+ "This flag must be used with `-allow-unsigned-driver`. This flag is only for debugging and testing.")
+ f.StringVar(&c.signatureURL, "signature-url", "",
+ "A URL to the driver signature. This flag can only be used together with `-test` and `-nvidia-installer-url` for for debugging and testing.")
f.BoolVar(&c.debug, "debug", false,
"Enable debug mode.")
f.BoolVar(&c.test, "test", false,
- "Enable test mode.")
+ "Enable test mode. "+
+ "In test mode, `-nvidia-installer-url` can be used without `-allow-unsigned-driver`.")
}
func (c *InstallCommand) validateFlags() error {
@@ -88,6 +93,9 @@
if c.nvidiaInstallerURL != "" && c.unsignedDriver == false && c.test == false {
return stderrors.New("-nvidia-installer-url is set, and -allow-unsigned-driver is not; -nvidia-installer-url must be used with -allow-unsigned-driver if not in test mode")
}
+ if c.signatureURL != "" && (c.nvidiaInstallerURL == "" || c.test == false) {
+ return stderrors.New("-signature-url must be used with -nvidia-installer-url and -test")
+ }
return nil
}
@@ -227,7 +235,11 @@
defer func() { callback <- 0 }()
if !c.unsignedDriver {
- if err := signing.DownloadDriverSignatures(downloader, c.driverVersion); err != nil {
+ if c.signatureURL != "" {
+ if err := signing.DownloadDriverSignaturesFromURL(c.signatureURL); err != nil {
+ return errors.Wrap(err, "failed to download driver signature")
+ }
+ } else if err := signing.DownloadDriverSignatures(downloader, c.driverVersion); err != nil {
if strings.Contains(err.Error(), "404 Not Found") {
return fmt.Errorf("The GPU driver is not available for the COS version. Please wait for half a day and retry.")
}
diff --git a/src/cmd/cos_gpu_installer/internal/signing/signature.go b/src/cmd/cos_gpu_installer/internal/signing/signature.go
index 952b1bb..68126eb 100644
--- a/src/cmd/cos_gpu_installer/internal/signing/signature.go
+++ b/src/cmd/cos_gpu_installer/internal/signing/signature.go
@@ -7,6 +7,7 @@
"path/filepath"
"cos.googlesource.com/cos/tools.git/src/pkg/cos"
+ "cos.googlesource.com/cos/tools.git/src/pkg/utils"
log "github.com/golang/glog"
"github.com/pkg/errors"
)
@@ -33,10 +34,38 @@
return errors.Wrapf(err, "failed to download driver signature for version %s", driverVersion)
}
- tarballPath := filepath.Join(gpuDriverSigningDir, driverVersion+".signature.tar.gz")
+ if err := decompressSignature(driverVersion + ".signature.tar.gz"); err != nil {
+ return errors.Wrapf(err, "failed to decompress driver signature for version %s.", driverVersion)
+ }
+
+ return nil
+}
+
+// DownloadDriverSignaturesFromURL downloads GPU driver signatures from a provided URL.
+func DownloadDriverSignaturesFromURL(signatureURL string) error {
+ if err := os.MkdirAll(gpuDriverSigningDir, 0755); err != nil {
+ return errors.Wrapf(err, "failed to create signing dir %s", gpuDriverSigningDir)
+ }
+
+ log.Infof("Downloading driver signature from URL: %s", signatureURL)
+ signatureName := filepath.Base(signatureURL)
+ outputPath := filepath.Join(gpuDriverSigningDir, signatureName)
+ if err := utils.DownloadContentFromURL(signatureURL, outputPath, "driver signature"); err != nil {
+ return errors.Wrapf(err, "failed to download driver signature from URL %s.", signatureURL)
+ }
+
+ if err := decompressSignature(signatureName); err != nil {
+ return errors.Wrapf(err, "failed to decompress driver signature: %s.", signatureName)
+ }
+
+ return nil
+}
+
+func decompressSignature(signatureName string) error {
+ tarballPath := filepath.Join(gpuDriverSigningDir, signatureName)
log.Infof("Decompressing signature %s", tarballPath)
if err := extractSignatures(tarballPath, gpuDriverSigningDir); err != nil {
- return errors.Wrapf(err, "failed to extract driver signatures for version %s", driverVersion)
+ return errors.Wrapf(err, "failed to extract driver signatures %s", signatureName)
}
// Create a dummy private key. We don't need private key to sign the driver
@@ -48,7 +77,6 @@
if err := f.Close(); err != nil {
return errors.Wrapf(err, "failed to close dummy key file in %s", filepath.Join(gpuDriverSigningDir, gpuDriverDummyKey))
}
-
return nil
}
