extensions/gpu/gpu.go

// Package gpu provides the tools need to install a gpu extension.
package gpu

import (
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"slices"
	"strings"

	"cos-extensions/tools/osutils"

	deviceInfo "cos.googlesource.com/cos/tools.git/src/cmd/cos_gpu_installer/deviceinfo"

	"github.com/golang/glog"
)

const (
	osReleasePath          = "/etc/os-release"
	versionID              = "VERSION_ID"
	buildID                = "BUILD_ID"
	apiDomain              = "API_DOMAIN"
	artifactRegistryDomain = "ARTIFACT_REGISTRY_DOMAIN"
	projectPrefix          = "PROJECT_PREFIX"
	instanceURL            = "http://metadata.google.internal/computeMetadata/v1/instance/zone?alt=text"
	extensionCache         = "/var/lib/cos-extensions"
	gcrPrefix              = "gcr.io/cos-cloud"
	defaultHostDir         = "/var/lib/nvidia"
)

var (
	run                     = osutils.Run
	parse                   = osutils.Parse
	artifactRegions         = []string{"us", "eu", "asia"}
	apiDomainsPath          = "/etc/cloud-api-domains"
	tpcGPUInstallerRepo     = "docker.%s/%s/cos-cloud/cos-gpu-installer"
	persistencedServicePath = "/etc/systemd/system/nvidia-persistenced.service"
	// Note: Please update the app-admin/extensions-manager ebuild file.
	DefaultInstaller = "gcr.io/cos-cloud/cos-gpu-installer:v2.4.4"
)

// Installers represent the names of the possible installers to be used during the program execution.
type Installers struct {
	// Default installer represents the default gpu driver installer (cos-gpu-installer)
	DefaultInstaller string
	// Selected installer represents the installer set in the environmental variable [COS_GPU_INSTALLER]
	SelectedInstaller string
}

var writer io.Writer = os.Stdout

// List returns all available COS extensions and their versions.
func List(installers Installers, printInstaller bool, installerArgs []string) error {
	osReleaseFile, err := parse(osReleasePath)
	if err != nil {
		return fmt.Errorf("couldn't parse os release file: %v", err)
	}
	if printInstaller {
		// Using Fprint as this displays output to the stdout
		// Logging is set to stderr
		fmt.Fprintln(writer, installers.DefaultInstaller)
	} else {
		fmt.Fprintf(writer, "Available extensions for COS version %s-%s:\n\n", osReleaseFile[versionID], osReleaseFile[buildID])
		fmt.Fprintln(writer, "[gpu]")
		fmt.Fprintf(writer, "gpu installer: %s\n", installers.DefaultInstaller)

		err := runInstaller("list", true, extensionCache, installers, installerArgs...)
		if err != nil {
			return fmt.Errorf("couldn't retrieve available GPU driver versions: %v", err)
		}
	}
	return nil
}

// Install installs the GPU drivers according to specifications
// from the installerArgs.
func Install(installers Installers, installerArgs []string) error {
	var skipKernelModuleLoadingArg bool
	var hostDir string

	// Check for `no-verify` flag and parse the host directory.
	for i := 0; i < len(installerArgs); i++ {
		arg := installerArgs[i]
		if arg == "-no-verify" || arg == "--no-verify" {
			skipKernelModuleLoadingArg = true
		}
		if flagName, value, found := strings.Cut(arg, "="); found && (flagName == "-host-dir" || flagName == "--host-dir") {
			hostDir = value
		} else if arg == "-host-dir" || arg == "--host-dir" {
			if i+1 < len(installerArgs) {
				hostDir = installerArgs[i+1]
				i++
			}
		}
	}

	// If host dir was not passed by the user, then set host dir to default.
	var args []string
	if hostDir == "" {
		hostDir = defaultHostDir
		args = []string{fmt.Sprintf("-host-dir=%s", hostDir)}
	}

	if !skipKernelModuleLoadingArg {
		gpuType, _ := deviceInfo.GetGPUTypeInfo("/sys/bus/pci/devices")
		if gpuType == deviceInfo.GB200 || gpuType == deviceInfo.GB300 {
			glog.Infof("%s GPU detected. Enabling IMEX channel creation via NVreg_CreateImexChannel0=1", gpuType)
			installerArgs = append(installerArgs, "-module-arg", "nvidia.NVreg_CreateImexChannel0=1")
		}
	}
	args = append(args, installerArgs...)
	err := runInstaller("install", false, extensionCache, installers, args...)
	if err != nil {
		return fmt.Errorf("could not install GPU drivers: %v", err)
	}
	glog.Info("GPU drivers successfully installed")
	if skipKernelModuleLoadingArg {
		fmt.Fprintln(writer, "Skipping post_install due to no-verify flag.")
	} else {
		postInstall(hostDir)
	}
	return nil
}

func postInstall(hostDir string) {
	if _, err := os.Stat(filepath.Join(hostDir, "bin/nvidia-gridd")); err == nil {
		if _, err := os.Stat(persistencedServicePath); err == nil {
			glog.Info("NVIDIA Grid driver detected; daemon setup is handled by cos-gpu-installer.")
			return
		}
	}

	glog.Info("Making the GPU driver installation path executable by re-mounting it.")
	cmd := fmt.Sprintf(
		`sudo -n -- sh -c "mount --bind %[1]s %[1]s; mount -o remount,exec %[1]s"`,
		hostDir)
	errMessage := fmt.Sprintf(
		`Before deploying GPU workloads - please make the driver installation path executable by re-mounting it using:
		sudo mount --bind %[1]s %[1]s
		sudo mount -o remount,exec %[1]s`,
		hostDir)

	_, err := run("/bin/sh", []string{"-c", cmd}, false)
	if err != nil {
		glog.Errorf("Automatic remount failed: %v\n%v\n", errMessage, err)
	}

	glog.Info("Enabling NVIDIA persistence mode.")

	// Enable persistence mode - Whenever the NVIDIA device resources are no longer
	// in use, the NVIDIA kernel driver will tear down the device state.
	// When persistence mode is enabled, the daemon holds the NVIDIA character
	// device files open, preventing the NVIDIA kernel driver from tearing down
	// device state when no other process is using the device This utility does not
	// actually use any device resources itself - it will simply sleep while
	// maintaining a reference to the NVIDIA device state.
	nvidiaPersistencePath := filepath.Join(hostDir, "bin/nvidia-persistenced")
	cmd = fmt.Sprintf("sudo %s", nvidiaPersistencePath)
	_, err = run("/bin/sh", []string{"-c", cmd}, false)
	if err != nil {
		glog.Errorf("Failed to run nvidia-persistenced daemon.\n%v\n", err)
	}
}

// Client represents an http.Client interface
type Client interface {
	Do(r *http.Request) (*http.Response, error)
}

func retrieveInstanceZone(c Client, url string) (string, error) {
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return "", fmt.Errorf("failed to create request to retrieve instance zone: %v", err)
	}
	req.Header.Set("Metadata-Flavor", "Google")
	resp, err := c.Do(req)
	if err != nil {
		return "", fmt.Errorf("failed to retrieve instance zone: %v", err)
	}
	defer resp.Body.Close()
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return "", fmt.Errorf("failed to read response body: %v", err)
	}

	// Syntax of response: projects/<project-id>/zones/<region>
	// Example: projects/475556798229/zones/us-central1-a
	zone := strings.Split(string(body), "/")
	zone = strings.Split(zone[len(zone)-1], "-")
	return zone[0], nil
}

func checkArch() error {
	arch, err := run("uname", []string{"-m"}, false)
	if err != nil {
		return fmt.Errorf("failed to retrieve machine architecture: %v", err)
	}
	if !strings.Contains(arch, "x86_64") && !strings.Contains(arch, "aarch64") {
		return fmt.Errorf("GPU installation is only supported on X86 and ARM64 for now.\n Current architecture detected: %s", arch)
	}
	return nil
}

// GetCosInstaller retrieves the installers to be used in execution of the program's commands.
func GetCosInstaller(c Client) (Installers, error) {
	var installers Installers
	defaultInstaller := DefaultInstaller

	// Modify DefaultInstaller based on APIDomainsPath. If
	// APIDomainsPath does not exist, just use DefaultInstaller. This
	// adjustment is required when running in a TPC environment,
	// where the installer needs to be fetched from an
	// internal repository.
	if _, err := os.Stat(apiDomainsPath); err == nil {
		domains, err := parse(apiDomainsPath)
		if err != nil {
			return installers, fmt.Errorf("failed to parse %q, err: %v", apiDomainsPath, err)
		}
		if domains[apiDomain] != "googleapis.com" {
			// We are running in TPC, set default installer to in-universe
			// repo.
			universeRepo := fmt.Sprintf(tpcGPUInstallerRepo, domains[artifactRegistryDomain], strings.TrimSuffix(domains[projectPrefix], ":"))
			defaultInstaller = strings.Replace(defaultInstaller, gcrPrefix, universeRepo, -1)
			glog.Infof("Using TPC Default installer: %q", defaultInstaller)
		}
	} else if !errors.Is(err, os.ErrNotExist) {
		return installers, fmt.Errorf("failed to stat %q, err: %v", apiDomainsPath, err)
	}

	// Checks if it is in the environmental variables and returns the value.
	if installer := os.Getenv("COS_GPU_INSTALLER"); installer != "" {
		installers = Installers{
			DefaultInstaller:  defaultInstaller,
			SelectedInstaller: installer,
		}
		return installers, nil
	}

	// Change gpu installer container image name if location of the vm is available
	// and it's from cos-cloud project.
	zone, err := retrieveInstanceZone(c, instanceURL)
	if err != nil {
		// It's OK if metadata server returns error.
		glog.Errorf("couldn't retrieve instance zone: %v", err)
		return installers, nil
	}
	if zone != "" && slices.Contains(artifactRegions, zone) {
		installers = Installers{
			DefaultInstaller: fmt.Sprintf("%s.%s", zone, defaultInstaller),
		}
		return installers, nil
	}
	installers = Installers{
		DefaultInstaller: defaultInstaller,
	}
	glog.Infof("Using installers: %v", installers)
	return installers, nil
}

// runInstaller runs the gpu installer with the cmd and installerArgs.
func runInstaller(cmd string, hideStderr bool, cacheDir string, installers Installers, installerArgs ...string) error {
	err := checkArch()
	if err != nil {
		return err
	}
	var cosInstaller string

	// Checks if selectedInstaller is set and uses that, else it defaults to the
	// default installer.
	if installers.SelectedInstaller != "" {
		cosInstaller = installers.SelectedInstaller
	} else {
		cosInstaller = installers.DefaultInstaller
	}

	args := []string{
		"run",
		"--rm",
		"--name=cos-gpu-installer",
		"--privileged",
		"--net=host",
		"--pid=host",
		"--volume", "/dev:/dev",
		"--volume", "/:/root",
		"--log-driver", "journald",
	}

	if slices.Contains(installerArgs, "--clean-build-tools") {
		glog.Infof("Cleaning cache present at: %s", cacheDir)

		if err := os.RemoveAll(cacheDir); err != nil {
			return fmt.Errorf("failed to clean cache present at %s: %v", cacheDir, err)
		}
	} else {
		_, err := os.Stat(cacheDir)

		// Use extension cache if exists by default or if build flag passed in
		if slices.Contains(installerArgs, "--prepare-build-tools") || !os.IsNotExist(err) {
			buildDir := fmt.Sprintf("%s/:/build/", cacheDir)
			args = append(args, []string{"--volume", buildDir}...)
		}
		args = append(args, cosInstaller)
		args = append(args, cmd)
		args = append(args, installerArgs...)
		res, err := run("/usr/bin/docker", args, hideStderr)
		if !hideStderr && err != nil {
			return fmt.Errorf("failed to complete installation using installer '%s': %v", cosInstaller, err)
		}
		fmt.Fprintln(writer, res)
	}
	return nil
}