commit d72d557590c807a36c3398fda1ba02864f37f029
author Arnav Kansal <rnv@google.com> Fri Dec 09 20:05:36 2022 +0000
committer Arnav Kansal <rnv@google.com> Thu Dec 15 19:07:46 2022 +0000
tree ecc42c093035ce215811c12a864429c9c5994a7f
parent b9a0c2c22e0c7ce6d2f1bf27789c2b5d9c92bee9

cos_gpu_driver_builder: Disable kernel CI config processing by default

only process image CI configs by default unless kernel is specified.

BUG=b/261639461

Change-Id: I64072c3a0d94c328101d1362a377ab6f26149114
Reviewed-on: https://cos-review.googlesource.com/c/cos/tools/+/39852
Cloud-Build: GCB Service account <228075978874@cloudbuild.gserviceaccount.com>
Reviewed-by: Meena Shanmugam <meenashanmugam@google.com>
Tested-by: Arnav Kansal <rnv@google.com>

src/cmd/cos_gpu_driver_builder/main.go

diff --git a/src/cmd/cos_gpu_driver_builder/main.go b/src/cmd/cos_gpu_driver_builder/main.go
index 286bb55..0d64774 100644
--- a/src/cmd/cos_gpu_driver_builder/main.go
+++ b/src/cmd/cos_gpu_driver_builder/main.go
@@ -17,6 +17,8 @@
 	configDir = flag.String("config-dir", "", "Directory containing config.textproto and metadata file that needs to be processed.")
 	bucket    = flag.String("watcher-gcs", "", "GCS bucket to watch for unprocessed configs.")
 	lookBack  = flag.Int("lookBackDays", 7, "read configs produced within the past <lookBack> days.")
+	// default to only building image CI precompiled drivers
+	mode = flag.String("mode", "image", "image, kernel, or both for processing image CI/kernel CI configs. Works only with watcher-gcs arg")
 )
 
 func main() {
@@ -34,7 +36,7 @@
 
 	var configs []gpuconfig.GPUPrecompilationConfig
 	if *bucket != "" { // cos_gpu_driver_builder --watcher-gcs="cos-gpu-configs"
-		configs, err = gpuconfig.ReadConfigs(ctx, client, *bucket, *lookBack)
+		configs, err = gpuconfig.ReadConfigs(ctx, client, *bucket, *lookBack, *mode)
 		if err != nil {
 			log.Fatal("could not read configs:", err)
 		}

src/pkg/gpuconfig/read_configs.go

diff --git a/src/pkg/gpuconfig/read_configs.go b/src/pkg/gpuconfig/read_configs.go
index c87b94a..4cb0c25 100644
--- a/src/pkg/gpuconfig/read_configs.go
+++ b/src/pkg/gpuconfig/read_configs.go
@@ -41,6 +41,9 @@
 // Reads precompilation config from GCS bucket into GPUPrecompilationConfig struct.
 func ReadConfig(ctx context.Context, client *storage.Client, dirName string) (GPUPrecompilationConfig, error) {
 	var config GPUPrecompilationConfig
+	if len(dirName) > 1 && dirName[len(dirName)-1] != '/' {
+		dirName += "/"
+	}
 	metadata, err := gcs.DownloadGCSObjectString(ctx, client, dirName+"metadata")
 	if err != nil {
 		return config, err
@@ -64,7 +67,7 @@
 }
 
 // Reads all config dirs published within <lookBackDays> of current date into a list of GPUPrecompilationConfig struct
-func ReadConfigs(ctx context.Context, client *storage.Client, bucketName string, lookBackDays int) ([]GPUPrecompilationConfig, error) {
+func ReadConfigs(ctx context.Context, client *storage.Client, bucketName string, lookBackDays int, versionType string) ([]GPUPrecompilationConfig, error) {
 	startDay := strings.TrimSuffix(timeNow().AddDate(0, 0, -lookBackDays).Format(time.RFC3339), "Z")
 	dirNames, err := listConfigDirs(ctx, client, bucketName, startDay)
 	if err != nil {
@@ -77,7 +80,16 @@
 		if err != nil {
 			return nil, err
 		}
-		configs = append(configs, config)
+		if matchVersionType(versionType, config.VersionType) {
+			configs = append(configs, config)
+		}
 	}
 	return configs, nil
 }
+
+func matchVersionType(mode string, versionType string) bool {
+	if strings.EqualFold(mode, "both") {
+		return true
+	}
+	return strings.EqualFold(mode, versionType)
+}