src/pkg/gpuconfig/read_configs.go

package gpuconfig

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"
	"time"

	"cloud.google.com/go/storage"
	"cos.googlesource.com/cos/tools.git/src/pkg/gcs"
	"cos.googlesource.com/cos/tools.git/src/pkg/gpuconfig/pb"
	"github.com/golang/protobuf/proto"
	"google.golang.org/api/iterator"
)

func listConfigDirs(ctx context.Context, client *storage.Client, bucketName string, start string) ([]string, error) {
	query := &storage.Query{
		StartOffset: start, // Only list objects lexicographically >=
		Delimiter:   "/",   // Only list dirs
	}
	query.SetAttrSelection([]string{"Prefix"})

	bkt := client.Bucket(bucketName)
	var dirNames []string
	it := bkt.Objects(ctx, query)
	for {
		attrs, err := it.Next()
		if err == iterator.Done {
			break
		}
		if err != nil {
			return nil, err
		}
		dirNames = append(dirNames, fmt.Sprintf("gs://%s/%s", bucketName, attrs.Prefix))
	}

	return dirNames, nil
}

// 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
	}

	if err := json.Unmarshal([]byte(metadata), &config); err != nil {
		return config, err
	}

	textproto, err := gcs.DownloadGCSObjectString(ctx, client, dirName+"config.textproto")
	if err != nil {
		return config, err
	}

	config.ProtoConfig = &pb.COSGPUBuildRequest{}
	if err := proto.UnmarshalText(textproto, config.ProtoConfig); err != nil {
		return config, err
	}

	return config, nil
}

// 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, 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 {
		return nil, err
	}

	configs := []GPUPrecompilationConfig{}
	for _, dir := range dirNames {
		config, err := ReadConfig(ctx, client, dir)
		if err != nil {
			return nil, err
		}
		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)
}