blob: d7fd5e061cc5d2006a52fcff361e4bdf2c0cd3bf [file] [log] [blame]
// Package gpuconfig implements routines for manipulating proto based
// GPU build configuration files.
// It also implements the construction of these configs for
// the COS Image and the COS Kernel CI.
package gpuconfig
//go:generate protoc --go_out=:./pb -I. proto/config.proto
import (
const (
kernelGCSPrefix string = "gs://cos-kernel-artifacts/builds"
kernelSrcTarballPathTemplate string = "%s/%[2]s/cos-kernel-src-%[2]s.tgz"
kernelHeadersTarballPathTemplate string = "%s/%[2]s/cos-kernel-headers-%[2]s-x86_64.tgz"
toolchainTarballPathTemplate string = "builds/%s/toolchain_url.x86_64"
toolchainEnvPathTemplate string = "%s/%s/toolchain_env.x86_64"
driverOutputGcsDirTemplate string = "gs://nvidia-drivers-us-public/nvidia-cos-project/%s/"
nvidiaRunfileAddressTemplate string = "[1]s/NVIDIA-Linux-x86_64-%[1]"
timeFormatTemplate string = "2006-01-02-15:04:05"
type GPUPrecompilationConfig struct {
ProtoConfig *pb.COSGPUBuildRequest `json:"-"`
DriverVersion string `json:"driver_version"`
Milestone string `json:"milestone"`
Version string `json:"version"`
VersionType string `json:"version_type"`
// Generates and GPU precompilation build configs(and metadata) for a given
// tuple of kernelVersion and driver versions
func GenerateKernelCIConfigs(ctx context.Context, client *storage.Client, kernelVersion string, driverVersions []string) ([]GPUPrecompilationConfig, error) {
configs := []GPUPrecompilationConfig{}
for _, driverVersion := range driverVersions {
config, err := constructKernelCIConfig(ctx, client, kernelVersion, driverVersion)
if err != nil {
return nil, err
milestone := strings.Split(kernelVersion, "m")[1]
configs = append(configs, GPUPrecompilationConfig{config, driverVersion, milestone, kernelVersion, "Kernel"})
return configs, nil
func constructKernelCIConfig(ctx context.Context, client *storage.Client, kernelVersion, driverVersion string) (*pb.COSGPUBuildRequest, error) {
config := pb.COSGPUBuildRequest{
KernelSrcTarballGcs: stringPtr(fmt.Sprintf(kernelSrcTarballPathTemplate, kernelGCSPrefix, kernelVersion)),
KernelHeadersTarballGcs: stringPtr(fmt.Sprintf(kernelHeadersTarballPathTemplate, kernelGCSPrefix, kernelVersion)),
NvidiaRunfileAddress: stringPtr(fmt.Sprintf(nvidiaRunfileAddressTemplate, driverVersion)),
ToolchainEnvGcs: stringPtr(fmt.Sprintf(toolchainEnvPathTemplate, kernelGCSPrefix, kernelVersion)),
DriverOutputGcsDir: stringPtr(fmt.Sprintf(driverOutputGcsDirTemplate, kernelVersion)),
toolchainTarballPath, err := fetchToolchainTarballPath(ctx, client, kernelVersion)
if err != nil {
return nil, err
config.ToolchainTarballGcs = &toolchainTarballPath
return &config, nil
func fetchToolchainTarballPath(ctx context.Context, client *storage.Client, kernelVersion string) (string, error) {
toolchainTarballPathURL := fmt.Sprintf(toolchainTarballPathTemplate, kernelVersion)
reader, err := client.Bucket("cos-kernel-artifacts").Object(toolchainTarballPathURL).NewReader(ctx)
if err != nil {
return "", fmt.Errorf("Could not fetch the toolchain tarball path: %w", err)
var toolchainTarballPath []byte
if toolchainTarballPath, err = ioutil.ReadAll(reader); err != nil {
return "", fmt.Errorf("Could not read file contents of toolchain tarball path: %w", err)
return string(toolchainTarballPath), nil
// stringPtr returns a pointer to a string.
func stringPtr(s string) *string {
return &s