Add ProcessConfigs method for GPU precompilation
This method include the implementation of finding all the configs that
have not been completely processed. For each such config it invokes the
builder methods which create the -custom.run self extracting archive and
also uploads this precompiled driver to the output directory as specified in the config.
BUG=b/238215395
Change-Id: I7f787b6520fd76a0d0b3424fbe1b7f4b4c969f11
Reviewed-on: https://cos-review.googlesource.com/c/cos/tools/+/37747
Cloud-Build: GCB Service account <228075978874@cloudbuild.gserviceaccount.com>
Reviewed-by: He Gao <hegao@google.com>
Tested-by: Arnav Kansal <rnv@google.com>
diff --git a/src/cmd/cos_gpu_driver_builder/builder/builder.go b/src/cmd/cos_gpu_driver_builder/builder/builder.go
deleted file mode 100644
index 5f524d0..0000000
--- a/src/cmd/cos_gpu_driver_builder/builder/builder.go
+++ /dev/null
@@ -1,16 +0,0 @@
-package builder
-
-import (
- "context"
- "log"
-
- "cloud.google.com/go/storage"
- "cos.googlesource.com/cos/tools.git/src/pkg/gpuconfig"
-)
-
-func BuildPrecompiledDrivers(ctx context.Context, client *storage.Client, configs []gpuconfig.GPUPrecompilationConfig) error {
- for _, config := range configs {
- log.Printf("building precompiled GPU driver for: %s, driver version %s\n", config.Version, config.DriverVersion)
- }
- return nil
-}
diff --git a/src/cmd/cos_gpu_driver_builder/internal/builder/builder.go b/src/cmd/cos_gpu_driver_builder/internal/builder/builder.go
new file mode 100644
index 0000000..04c6a23
--- /dev/null
+++ b/src/cmd/cos_gpu_driver_builder/internal/builder/builder.go
@@ -0,0 +1,12 @@
+package builder
+
+import (
+ "context"
+
+ "cloud.google.com/go/storage"
+ "cos.googlesource.com/cos/tools.git/src/pkg/gpuconfig"
+)
+
+func BuildPrecompiledDriver(ctx context.Context, client *storage.Client, config gpuconfig.GPUPrecompilationConfig) (string, error) {
+ return "", nil
+}
diff --git a/src/cmd/cos_gpu_driver_builder/internal/config/configs.go b/src/cmd/cos_gpu_driver_builder/internal/config/configs.go
new file mode 100644
index 0000000..ba673c6
--- /dev/null
+++ b/src/cmd/cos_gpu_driver_builder/internal/config/configs.go
@@ -0,0 +1,38 @@
+package config
+
+import (
+ "context"
+ "fmt"
+ "log"
+
+ "cloud.google.com/go/storage"
+ "cos.googlesource.com/cos/tools.git/src/cmd/cos_gpu_driver_builder/internal/builder"
+ "cos.googlesource.com/cos/tools.git/src/pkg/gcs"
+ "cos.googlesource.com/cos/tools.git/src/pkg/gpuconfig"
+)
+
+func outputDriverFile(config gpuconfig.GPUPrecompilationConfig) string {
+ driverRunfile := fmt.Sprintf("NVIDIA-Linux-x86_64-%s-custom.run", config.DriverVersion)
+ return fmt.Sprintf("%s/%s", config.ProtoConfig.GetDriverOutputGcsDir(), driverRunfile)
+}
+
+func ProcessConfigs(ctx context.Context, client *storage.Client, configs []gpuconfig.GPUPrecompilationConfig) error {
+ for _, config := range configs {
+ log.Printf("building precompiled GPU driver for: %s, driver version %s\n", config.Version, config.DriverVersion)
+ if processed, _ := gcs.GCSObjectExists(ctx, client, outputDriverFile(config)); processed {
+ continue
+ }
+
+ precompiledDriver, err := builder.BuildPrecompiledDriver(ctx, client, config)
+ if err != nil {
+ log.Printf("precompilation failed for: %s, driver version %s\n", config.Version, config.DriverVersion)
+ continue
+ }
+
+ if err := gcs.UploadGCSObject(ctx, client, precompiledDriver, outputDriverFile(config)); err != nil {
+ log.Printf("export failed for: %s, driver version %s\n", config.Version, config.DriverVersion)
+ continue
+ }
+ }
+ return nil
+}
diff --git a/src/cmd/cos_gpu_driver_builder/main.go b/src/cmd/cos_gpu_driver_builder/main.go
index 96513eb..a478fca 100644
--- a/src/cmd/cos_gpu_driver_builder/main.go
+++ b/src/cmd/cos_gpu_driver_builder/main.go
@@ -6,8 +6,9 @@
"log"
"cloud.google.com/go/storage"
- "cos.googlesource.com/cos/tools.git/src/cmd/cos_gpu_driver_builder/builder"
"cos.googlesource.com/cos/tools.git/src/pkg/gpuconfig"
+
+ "cos.googlesource.com/cos/tools.git/src/cmd/cos_gpu_driver_builder/internal/config"
)
var (
@@ -32,5 +33,5 @@
log.Fatal("could not read configs: %v", err)
}
- _ = builder.BuildPrecompiledDrivers(ctx, client, configs)
+ config.ProcessConfigs(ctx, client, configs)
}
diff --git a/src/pkg/gcs/gcs_client.go b/src/pkg/gcs/gcs_client.go
index e2983ba..79570ad 100644
--- a/src/pkg/gcs/gcs_client.go
+++ b/src/pkg/gcs/gcs_client.go
@@ -71,6 +71,19 @@
return string(ret), nil
}
+// GCSObjectExists checks if an object exists at inputURL
+func GCSObjectExists(ctx context.Context, gcsClient *storage.Client, inputURL string) (bool, error) {
+ gcsBucket, name, err := getGCSVariables(inputURL)
+ _, err = gcsClient.Bucket(gcsBucket).Object(name).Attrs(ctx)
+ if err == storage.ErrObjectNotExist {
+ return false, nil
+ }
+ if err != nil {
+ return false, err
+ }
+ return true, nil
+}
+
// UploadGCSObject uploads an object at inputPath to destination URL
func UploadGCSObject(ctx context.Context, gcsClient *storage.Client, inputPath, destinationURL string) error {
fileReader, err := os.Open(inputPath)
