src/cmd/cos_gpu_driver_builder/main.go

package main

import (
	"context"
	"flag"

	log "github.com/golang/glog"

	"cloud.google.com/go/storage"

	"cos.googlesource.com/cos/tools.git/src/cmd/cos_gpu_driver_builder/internal/config"
	"cos.googlesource.com/cos/tools.git/src/pkg/gpuconfig"
)

var (
	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")
	dryRun = flag.Bool("dry-run", false, "invoking the driver builder with -dry-run will not upload any build precompiled outputs")
)

func main() {
	flag.Parse()
	// always run driver builder with verbose logging.
	flag.Set("v", "2")
	if *bucket == "" && *configDir == "" {
		log.Fatal("empty watcher gcs dir and config file dir")
	}

	ctx := context.Background()
	client, err := storage.NewClient(ctx)
	if err != nil {
		log.Fatal("failed to setup client for GCS:", err)
	}

	var configs []gpuconfig.GPUPrecompilationConfig
	if *bucket != "" { // cos_gpu_driver_builder --watcher-gcs="cos-gpu-configs"
		configs, err = gpuconfig.ReadConfigs(ctx, client, *bucket, *lookBack, *mode)
		if err != nil {
			log.Fatal("could not read configs:", err)
		}
	} else { // cos_gpu_driver_builder --config="gs://cos-gpu-configs/2022-09-26T16:03:17-dc65ba40/"
		config, err := gpuconfig.ReadConfig(ctx, client, *configDir)
		if err != nil {
			log.Fatal("could not read config:", err)
		}
		configs = append(configs, config)
	}

	config.ProcessConfigs(ctx, client, configs, *dryRun)
}