blob: 6874500ca351c2634be66fbecf3d7e6cf1904078 [file] [log] [blame]
// Copyright 2018 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package main
import (
"context"
"encoding/json"
"errors"
"flag"
"fmt"
"io/ioutil"
"log"
"os"
"path/filepath"
"strings"
"text/template"
"cos.googlesource.com/cos/tools.git/src/pkg/config"
"cos.googlesource.com/cos/tools.git/src/pkg/fs"
"cos.googlesource.com/cos/tools.git/src/pkg/provisioner"
"cos.googlesource.com/cos/tools.git/src/pkg/utils"
"cloud.google.com/go/storage"
"github.com/google/subcommands"
"google.golang.org/api/iterator"
)
const (
gpuScript = "install_gpu.sh"
installerContainer = "gcr.io/cos-cloud/cos-gpu-installer:v20230117"
)
// TODO(b/121332360): Move most GPU functionality to cos-gpu-installer
var (
validGPUs = []string{"nvidia-tesla-k80", "nvidia-tesla-p100", "nvidia-tesla-v100", "nvidia-tesla-t4"}
)
// InstallGPU implements subcommands.Command for the "install-gpu" command.
// This command configures the current image build process to customize the result image
// with GPU drivers.
type InstallGPU struct {
NvidiaDriverVersion string
NvidiaDriverMd5sum string
NvidiaInstallDirHost string
gpuType string
getValidDrivers bool
gpuDataDir string
}
// Name implements subcommands.Command.Name.
func (*InstallGPU) Name() string {
return "install-gpu"
}
// Synopsis implements subcommands.Command.Synopsis.
func (*InstallGPU) Synopsis() string {
return "Configure the image build with GPU drivers."
}
// Usage implements subcommands.Command.Usage.
func (*InstallGPU) Usage() string {
return `install-gpu [flags]
`
}
// SetFlags implements subcommands.Command.SetFlags.
func (i *InstallGPU) SetFlags(f *flag.FlagSet) {
f.StringVar(&i.NvidiaDriverVersion, "version", "", "Driver version to install. Can also be the name of an nvidia installer present in the "+
"directory specified by '-deps-dir'; e.g., NVIDIA-Linux-x86_64-450.51.06.run.")
f.StringVar(&i.NvidiaDriverMd5sum, "md5sum", "", "Md5sum of the driver to install.")
f.StringVar(&i.NvidiaInstallDirHost, "install-dir", "/var/lib/nvidia",
"Location to install drivers on the image.")
f.StringVar(
&i.gpuType, "gpu-type", "nvidia-tesla-p100", "This flag is deprecated, please use the same flag in step 'finish-image-build'.")
f.BoolVar(
&i.getValidDrivers, "get-valid-drivers", false,
"Print the list of supported GPU driver versions. If this flag is given, no other actions will be taken.")
f.StringVar(&i.gpuDataDir, "deps-dir", "", "If provided, the local directory to search for cos-gpu-installer data dependencies. "+
"The exact data dependencies that must be present in this directory depends on the version of cos-gpu-installer "+
"used by cos-customizer. Do not expect this flag to be stable; it exists for compatibility with pre-release COS images.")
}
func validDriverVersions(ctx context.Context, gcsClient *storage.Client) (map[string]bool, error) {
// We gather the set of valid drivers from the set of drivers provided by Nvidia in their GCS bucket.
// Nominally, paths we care about in this bucket look like 'tesla/<version>/<binaries>'. Version 390.46 has
// a deprecated path structure, and since it's supported by cos-gpu-installer, we special case that here.
validDrivers := map[string]bool{"390.46": true}
query := &storage.Query{Prefix: "tesla/"}
it := gcsClient.Bucket("nvidia-drivers-us-public").Objects(ctx, query)
for {
objAttrs, err := it.Next()
if err == iterator.Done {
break
}
if err != nil {
return nil, err
}
// Example object: tesla/396.26/NVIDIA-Linux-x86_64-396.26-diagnostic.run
if splitPath := strings.SplitN(objAttrs.Name, "/", 3); len(splitPath) > 1 {
validDrivers[splitPath[1]] = true
}
}
return validDrivers, nil
}
func (i *InstallGPU) validate(ctx context.Context, gcsClient *storage.Client, files *fs.Files, provConfig *provisioner.Config) error {
isValidGPU := false
for _, g := range validGPUs {
if i.gpuType == g {
isValidGPU = true
break
}
}
if !isValidGPU {
return fmt.Errorf("%q is an invalid GPU type. Must be one of: %v", i.gpuType, validGPUs)
}
if i.NvidiaDriverVersion == "" {
return fmt.Errorf("version must be set")
}
var gpuAlreadyConf bool
for _, s := range provConfig.Steps {
if s.Type == "InstallGPU" {
gpuAlreadyConf = true
break
}
}
if gpuAlreadyConf {
return fmt.Errorf("install-gpu can only be invoked once in an image build process. Only one driver version can be installed on the image")
}
if strings.HasSuffix(i.NvidiaDriverVersion, ".run") {
log.Printf("driver version is set to %q, which looks like an nvidia installer file", i.NvidiaDriverVersion)
if i.gpuDataDir == "" {
return errors.New(`"-deps-dir" must be set when the version is specified as an nvidia installer file`)
}
fileName := filepath.Join(i.gpuDataDir, i.NvidiaDriverVersion)
info, err := os.Stat(fileName)
if os.IsNotExist(err) {
return fmt.Errorf("nvidia installer file at %q does not exist", fileName)
}
if info.IsDir() {
return fmt.Errorf("nvidia installer file at %q is a directory", fileName)
}
} else {
validDrivers, err := validDriverVersions(ctx, gcsClient)
if err != nil {
return err
}
if !validDrivers[i.NvidiaDriverVersion] {
var drivers []string
for d := range validDrivers {
drivers = append(drivers, d)
}
return fmt.Errorf("driver version %s is not valid; valid driver versions are: %v", i.NvidiaDriverVersion, drivers)
}
}
return nil
}
func (i *InstallGPU) templateScript(scriptPath string) error {
setCOSDownloadGCS := ""
if i.gpuDataDir != "" {
setCOSDownloadGCS = "true"
}
data := struct {
NvidiaDriverVersion string
NvidiaDriverMd5sum string
NvidiaInstallDirHost string
SetCOSDownloadGCS string
}{
NvidiaDriverVersion: utils.QuoteForShell(i.NvidiaDriverVersion),
NvidiaDriverMd5sum: utils.QuoteForShell(i.NvidiaDriverMd5sum),
NvidiaInstallDirHost: utils.QuoteForShell(i.NvidiaInstallDirHost),
SetCOSDownloadGCS: utils.QuoteForShell(setCOSDownloadGCS),
}
tmpl, err := template.New(filepath.Base(scriptPath)).ParseFiles(scriptPath)
if err != nil {
return err
}
w, err := os.Create(scriptPath)
if err != nil {
return err
}
defer w.Close()
return tmpl.Execute(w, data)
}
func (i *InstallGPU) updateBuildConfig(configPath string) error {
buildConfig := &config.Build{}
configFile, err := os.OpenFile(configPath, os.O_RDWR, 0666)
if err != nil {
return err
}
defer configFile.Close()
if err := config.Load(configFile, buildConfig); err != nil {
return err
}
buildConfig.GPUType = i.gpuType
if i.gpuDataDir != "" {
files, err := ioutil.ReadDir(i.gpuDataDir)
if err != nil {
return fmt.Errorf("error reading dir %q: %v", i.gpuDataDir, err)
}
for _, f := range files {
if f.Mode().IsRegular() {
buildConfig.GCSFiles = append(buildConfig.GCSFiles, filepath.Join(i.gpuDataDir, f.Name()))
}
}
}
if _, err := configFile.Seek(0, 0); err != nil {
return err
}
return config.SaveConfigToFile(configFile, buildConfig)
}
func (i *InstallGPU) updateProvConfig(provConfig *provisioner.Config) error {
buf, err := json.Marshal(&provisioner.InstallGPUStep{
NvidiaDriverVersion: i.NvidiaDriverVersion,
NvidiaDriverMD5Sum: i.NvidiaDriverMd5sum,
NvidiaInstallDirHost: i.NvidiaInstallDirHost,
NvidiaInstallerContainer: installerContainer,
// GCSDepsPrefix will be converted into a gs:// address by the preloader
// package.
GCSDepsPrefix: i.gpuDataDir,
})
if err != nil {
return err
}
provConfig.Steps = append(provConfig.Steps, provisioner.StepConfig{
Type: "InstallGPU",
Args: json.RawMessage(buf),
})
return nil
}
// Execute implements subcommands.Command.Execute. It configures the current image build process to
// customize the result image with GPU drivers.
func (i *InstallGPU) Execute(ctx context.Context, f *flag.FlagSet, args ...interface{}) subcommands.ExitStatus {
if f.NArg() != 0 {
f.Usage()
return subcommands.ExitUsageError
}
if len(args) < 2 {
log.Panic("InstallGPU expects two arguments; *fs.Files and ServiceClients")
}
files, ok := args[0].(*fs.Files)
if !ok {
log.Panic("InstallGPU expects two arguments; *fs.Files and ServiceClients")
}
serviceClients, ok := args[1].(ServiceClients)
if !ok {
log.Panic("InstallGPU expects two arguments; *fs.Files and ServiceClients")
}
_, gcsClient, err := serviceClients(ctx, false)
if err != nil {
log.Println(err)
return subcommands.ExitFailure
}
defer gcsClient.Close()
if i.getValidDrivers {
validDrivers, err := validDriverVersions(ctx, gcsClient)
if err != nil {
log.Println(err)
return subcommands.ExitFailure
}
var drivers []string
for d := range validDrivers {
drivers = append(drivers, d)
}
log.Printf("Valid driver versions are: %v\n", drivers)
return subcommands.ExitSuccess
}
var provConfig provisioner.Config
if err := config.LoadFromFile(files.ProvConfig, &provConfig); err != nil {
log.Println(err)
return subcommands.ExitFailure
}
if err := i.validate(ctx, gcsClient, files, &provConfig); err != nil {
log.Println(err)
return subcommands.ExitFailure
}
if err := i.updateBuildConfig(files.BuildConfig); err != nil {
log.Println(err)
return subcommands.ExitFailure
}
if err := i.updateProvConfig(&provConfig); err != nil {
log.Println(err)
return subcommands.ExitFailure
}
if err := config.SaveConfigToPath(files.ProvConfig, &provConfig); err != nil {
log.Println(err)
return subcommands.ExitFailure
}
return subcommands.ExitSuccess
}