commands/rootCmds.go - cos/cos-extensions - Git at Google


            
            package commands


            import (

            
            	"fmt"

            
            	"net/http"

            
            	"os"

            
            	"slices"


            	"cos-extensions/extensions/gpu"


            	"github.com/spf13/cobra"

            
            )


            const (

            
            	helpTemplate  = `{{if or .Runnable .HasSubCommands}}{{.UsageString}}{{end}}`

            
            	usageTemplate = `Usage:

            
            	{{.UseLine}}{{with (or .Long .Short)}}


            Description:

            
            	{{. | trimTrailingWhitespaces}}{{end}}{{if gt (len .Aliases) 0}}


            Aliases:

            
            	{{.NameAndAliases}}{{end}}{{if .HasExample}}


            Examples:

            
              {{.Example}}{{end}}{{if .HasAvailableSubCommands}}{{$cmds := .Commands}}{{if eq (len .Groups) 0}}


            {{if eq .Name "cos-extensions"}}Available Commands:{{else}}Available Extensions:{{end}}{{range $cmds}}{{if (or .IsAvailableCommand (eq .Name "help"))}}

            
            	{{rpad .Name .NamePadding }} {{.Short}}{{end}}{{end}}{{else}}{{range $group := .Groups}}


              {{.Title}}{{range $cmds}}{{if (and (eq .GroupID $group.ID) (or .IsAvailableCommand (eq .Name "help")))}}

            
            	{{rpad .Name .NamePadding }} {{.Short}}{{end}}{{end}}{{end}}{{if not .AllChildCommandsHaveGroup}}


            Additional Commands:{{range $cmds}}{{if (and (eq .GroupID "") (or .IsAvailableCommand (eq .Name "help")))}}

            
            	{{rpad .Name .NamePadding }} {{.Short}}{{end}}{{end}}{{end}}{{end}}{{end}}{{if .HasAvailableLocalFlags}}


            Flags:

            
              {{.LocalFlags.FlagUsages | trimTrailingWhitespaces}}{{end}}{{if .HasAvailableInheritedFlags}}


            Global Flags:

            
              {{.InheritedFlags.FlagUsages | trimTrailingWhitespaces}}{{end}}{{if .HasHelpSubCommands}}


            Additional help topics:{{range .Commands}}{{if .IsAdditionalHelpTopicCommand}}

            
            	{{rpad .CommandPath .CommandPathPadding}} {{.Short}}{{end}}{{end}}{{end}}{{if (and .HasAvailableSubCommands (eq .Name "cos-extensions"))}}


            Use "{{.CommandPath}} [command] --help" for more information about a command or

            
            use "{{.CommandPath}} [command] <extension> --help" to learn how to use an extension with a command. {{else}} {{if .HasAvailableSubCommands}}


            Use "{{.CommandPath}} <extension> --help" to learn how to use an extension with {{.Name}} command{{end}}{{end}}


            `

            
            )


            var (

            
            	rootCmd = &cobra.Command{

            
            		Use:                   "cos-extensions [command] <extension>",

            
            		DisableFlagsInUseLine: true,

            
            		Long: `cos-extensions        Utility to maintain cos extensions.


              Available Extensions:

            
                [gpu]                 Installs gpu drivers and lists the available versions.


                    Extension usage:

            
                    cos-extensions list gpu [--gpu-installer]   lists all available gpu driver verions.

            
                                                                When --gpu-installer is set, lists the default

            
                                                                cos-gpu-installer.


                    cos-extensions install gpu                  installs gpu drivers.


                    Additional Description:

            
                    cos-extensions install gpu -- [options]


                    --version

            
                    The gpu extension can be invoked with --version. The possible values are 'latest', 'default',

            
                    'R<major-version>' eg. 'R470', 'R535'. or precise driver versions retrievable by running

            
                    cos-extensions list gpu.

            
                    If unspecified, the default driver version is installed.


                    --force-fallback

            
                    This flag indicates whether to use fallback mechanism when specified GPU driver is not compatible

            
                    with GPU devices. In case unspecified, the fallback behavior of the installer is not applicable

            
                    for --version=R<major-version> eg. 'R470', 'R525' or --version=<precise-version> eg.

            
                    '535.129.03', '525.147.05'.

            
                    The fallback behavior of the installer is active for --version is unset or --version=default or

            
                    --version=latest.

            
                    When fallback behavior is active, the installer will find a compatible driver detected GPU on

            
                    the VM.


                    --prepare-build-tools

            
                    The gpu extension can be invoked with a --prepare-build-tools optional argument that can be used

            
                    to cache the toolchain for the installer.

            
                    Caching the toolchain carries the overhead of ~1GB disk space on the stateful partition.

            
                    Using this command only populates the cache and does NOT install the GPU drivers thus may save

            
                    time on downloading the toolchain during subsequent installations.


                    --clean-build-tools

            
                    Use this optional command to delete the cache for the toolchain present on the stateful

            
                    partition.


                    -test

            
                    The gpu extension can be used to install drivers on a dev channel image with this.


                    -no-verify

            
                    Skip kernel module loading, installation verification,

            
                    and the enabling of NVIDIA persistence mode.

            
                    Useful for preloading drivers without attached GPU.


                    -debug

            
                    The granularity of logging of the gpu extension can be increased by specifying this.


                    cos-extensions list gpu -- [options]


                    --target-gpu

            
                    This flag specifies the GPU device to display its compatible drivers.

            
                    If specified, it must be one of: NVIDIA_TESLA_K80, NVIDIA_TESLA_P4, NVIDIA_TESLA_P100,

            
                    NVIDIA_TESLA_V100, NVIDIA_L4, NVIDIA_H100_80GB, NVIDIA_TESLA_A100, NVIDIA_A100_80GB,

            
                    NVIDIA_TESLA_T4.

            
                    If not specified, the GPU device will be auto-detected by the installer.


                    --gpu-proto-cache-dir

            
                    The GPU proto cache directory that GPU driver versions proto file is stored into.

            
                    If unspecified, the GPU driver versions proto file will not be cached.


                [kernel-mods]        Installs out-of-tree kernel mods from GCS bucket to the COS kernel.


                    Extension usage:

            
                    cos-extensions list kernel-mods [flags]      lists all the available compiled kernel modules in a

            
                                                                 gcs directory.


                    cos-extension install kernel-mods            installs compiled kernel modules from a GCS bucket

            
                    [module files] [flags] -- [module args]      directory.

            
                                                                 Note that module arguments should be avoided during

            
                                                                 multiple module installations as it might lead to

            
                                                                 unexpected results.


                    Additional Description:

            
                    cos-extensions install kernel-mods [module files] [flags] -- [module arguments]


                    --gcs_path

            
                    This flag can be used to specify the GCS path which stores the available kernel modules to be

            
                    installed. The expected syntax is gs://<bucket name>/<directory path>.

            
                    If not specified, the default gcs path is used.


                    --local_dir

            
                    This flag can be used to specify the local directory where the kernel modules would be installed.

            
                    If not specified, the default local directory is used.


                    --download-only, -d

            
                    This flag indicates whether to download the module to the local directory without installing.

            
                    This flag cannot be used with the local_only flag.


                    --local_only, -l

            
                    This flag indicates whether to install modules directly from local directory without downloading.

            
                    This flag cannot be used with the download_only flag.


                    --recursive, -r

            
                    This flag indicates whether to recursively installs modules from a gcs bucket directory and its

            
                    subdirectories.

            
                    This flag cannot be used with the local_only flag nor module files.

            
                    Note that the modules are downloaded with their subdirectory path.


                    cos-extensions list kernel-mods [flags]


                    --gcs_path

            
                    This flag can be used to specify the GCS path from which the available kernel modules would be

            
                    listed. The expected syntax is gs://<bucket name>/<directory path>.

            
                    If not specified, the default path is used.


                    --local_dir

            
                    This flag can be used to specify the local directory from which the downloaded or installed

            
                    kernel modules would be listed.

            
                    If not specified, the default path is used.


                    --downloaded, -d

            
                    This flag indicates whether to list only the downloaded compiled modules in a local directory.

            
                    This can be used with local_dir flag, else it uses the default directory. It can be used with

            
                    the recursive flag to include downloaded modules in the subdirectories.


                    --installed, -i

            
                    This flag indicates whether to list only the installed kernel modules in a local directory.

            
                    This can be used with local_dir flag, else it uses the default directory. It can be used with

            
                    the recursive flag to include installed modules in the subdirectories.


                    --recursive, -r

            
                    This flag indicates whether to recursively list all available modules in the gcs subdirectories.

            
                    This can be used with gcs_path flag, else it uses default gcs path. It can be used with installed

            
                    flag to list all installed kernels in the directory and its subdirectories. It can also be used

            
                    with the downloaded flag.`,

            
            	}


            	installCmd = &cobra.Command{

            
            		Use:                   "install <extension>",

            
            		DisableFlagsInUseLine: true,

            
            		Short:                 "Installs a COS extension.",

            
            	}


            	listCmd = &cobra.Command{

            
            		Use:                   "list <extension>",

            
            		DisableFlagsInUseLine: true,

            
            		Short:                 "Lists all available COS extensions and their versions.",

            
            		RunE: func(cmd *cobra.Command, args []string) error {

            
            			printInstaller, _ := cmd.Flags().GetBool("gpu-installer")

            
            			if slices.Contains(args, "--gpu-installer") {

            
            				printInstaller = true

            
            			}

            
            			installers, err := gpu.GetCosInstaller(&http.Client{})

            
            			if err != nil {

            
            				return fmt.Errorf("Failed to set gpu installer: %v", err)

            
            			}

            
            			err = gpu.List(installers, printInstaller, args)

            
            			if err != nil {

            
            				return fmt.Errorf("Failed to run list command: %v", err)

            
            			}

            
            			return nil

            
            		},

            
            	}

            
            )


            func Execute() {

            
            	if err := rootCmd.Execute(); err != nil {

            
            		// Did not return error as it is handled by the RunE field.

            
            		os.Exit(1)

            
            	}

            
            }


            func init() {

            
            	rootCmd.CompletionOptions.DisableDefaultCmd = true

            
            	rootCmd.SetUsageTemplate(usageTemplate)

            
            	rootCmd.SetHelpTemplate(helpTemplate)


            	// Registering subcommands for cos-extensions

            
            	rootCmd.AddCommand(installCmd)

            
            	rootCmd.AddCommand(listCmd)

            
            }
	package commands

	import (
	"fmt"
	"net/http"
	"os"
	"slices"

	"cos-extensions/extensions/gpu"

	"github.com/spf13/cobra"
	)

	const (
	helpTemplate = `{{if or .Runnable .HasSubCommands}}{{.UsageString}}{{end}}`
	usageTemplate = `Usage:
	{{.UseLine}}{{with (or .Long .Short)}}

	Description:
	{{. \| trimTrailingWhitespaces}}{{end}}{{if gt (len .Aliases) 0}}

	Aliases:
	{{.NameAndAliases}}{{end}}{{if .HasExample}}

	Examples:
	{{.Example}}{{end}}{{if .HasAvailableSubCommands}}{{$cmds := .Commands}}{{if eq (len .Groups) 0}}

	{{if eq .Name "cos-extensions"}}Available Commands:{{else}}Available Extensions:{{end}}{{range $cmds}}{{if (or .IsAvailableCommand (eq .Name "help"))}}
	{{rpad .Name .NamePadding }} {{.Short}}{{end}}{{end}}{{else}}{{range $group := .Groups}}

	{{.Title}}{{range $cmds}}{{if (and (eq .GroupID $group.ID) (or .IsAvailableCommand (eq .Name "help")))}}
	{{rpad .Name .NamePadding }} {{.Short}}{{end}}{{end}}{{end}}{{if not .AllChildCommandsHaveGroup}}

	Additional Commands:{{range $cmds}}{{if (and (eq .GroupID "") (or .IsAvailableCommand (eq .Name "help")))}}
	{{rpad .Name .NamePadding }} {{.Short}}{{end}}{{end}}{{end}}{{end}}{{end}}{{if .HasAvailableLocalFlags}}

	Flags:
	{{.LocalFlags.FlagUsages \| trimTrailingWhitespaces}}{{end}}{{if .HasAvailableInheritedFlags}}

	Global Flags:
	{{.InheritedFlags.FlagUsages \| trimTrailingWhitespaces}}{{end}}{{if .HasHelpSubCommands}}

	Additional help topics:{{range .Commands}}{{if .IsAdditionalHelpTopicCommand}}
	{{rpad .CommandPath .CommandPathPadding}} {{.Short}}{{end}}{{end}}{{end}}{{if (and .HasAvailableSubCommands (eq .Name "cos-extensions"))}}

	Use "{{.CommandPath}} [command] --help" for more information about a command or
	use "{{.CommandPath}} [command] <extension> --help" to learn how to use an extension with a command. {{else}} {{if .HasAvailableSubCommands}}

	Use "{{.CommandPath}} <extension> --help" to learn how to use an extension with {{.Name}} command{{end}}{{end}}

	`
	)

	var (
	rootCmd = &cobra.Command{
	Use: "cos-extensions [command] <extension>",
	DisableFlagsInUseLine: true,
	Long: `cos-extensions Utility to maintain cos extensions.

	Available Extensions:
	[gpu] Installs gpu drivers and lists the available versions.

	Extension usage:
	cos-extensions list gpu [--gpu-installer] lists all available gpu driver verions.
	When --gpu-installer is set, lists the default
	cos-gpu-installer.

	cos-extensions install gpu installs gpu drivers.

	Additional Description:
	cos-extensions install gpu -- [options]

	--version
	The gpu extension can be invoked with --version. The possible values are 'latest', 'default',
	'R<major-version>' eg. 'R470', 'R535'. or precise driver versions retrievable by running
	cos-extensions list gpu.
	If unspecified, the default driver version is installed.

	--force-fallback
	This flag indicates whether to use fallback mechanism when specified GPU driver is not compatible
	with GPU devices. In case unspecified, the fallback behavior of the installer is not applicable
	for --version=R<major-version> eg. 'R470', 'R525' or --version=<precise-version> eg.
	'535.129.03', '525.147.05'.
	The fallback behavior of the installer is active for --version is unset or --version=default or
	--version=latest.
	When fallback behavior is active, the installer will find a compatible driver detected GPU on
	the VM.

	--prepare-build-tools
	The gpu extension can be invoked with a --prepare-build-tools optional argument that can be used
	to cache the toolchain for the installer.
	Caching the toolchain carries the overhead of ~1GB disk space on the stateful partition.
	Using this command only populates the cache and does NOT install the GPU drivers thus may save
	time on downloading the toolchain during subsequent installations.

	--clean-build-tools
	Use this optional command to delete the cache for the toolchain present on the stateful
	partition.

	-test
	The gpu extension can be used to install drivers on a dev channel image with this.

	-no-verify
	Skip kernel module loading, installation verification,
	and the enabling of NVIDIA persistence mode.
	Useful for preloading drivers without attached GPU.

	-debug
	The granularity of logging of the gpu extension can be increased by specifying this.

	cos-extensions list gpu -- [options]

	--target-gpu
	This flag specifies the GPU device to display its compatible drivers.
	If specified, it must be one of: NVIDIA_TESLA_K80, NVIDIA_TESLA_P4, NVIDIA_TESLA_P100,
	NVIDIA_TESLA_V100, NVIDIA_L4, NVIDIA_H100_80GB, NVIDIA_TESLA_A100, NVIDIA_A100_80GB,
	NVIDIA_TESLA_T4.
	If not specified, the GPU device will be auto-detected by the installer.

	--gpu-proto-cache-dir
	The GPU proto cache directory that GPU driver versions proto file is stored into.
	If unspecified, the GPU driver versions proto file will not be cached.

	[kernel-mods] Installs out-of-tree kernel mods from GCS bucket to the COS kernel.

	Extension usage:
	cos-extensions list kernel-mods [flags] lists all the available compiled kernel modules in a
	gcs directory.

	cos-extension install kernel-mods installs compiled kernel modules from a GCS bucket
	[module files] [flags] -- [module args] directory.
	Note that module arguments should be avoided during
	multiple module installations as it might lead to
	unexpected results.

	Additional Description:
	cos-extensions install kernel-mods [module files] [flags] -- [module arguments]

	--gcs_path
	This flag can be used to specify the GCS path which stores the available kernel modules to be
	installed. The expected syntax is gs://<bucket name>/<directory path>.
	If not specified, the default gcs path is used.

	--local_dir
	This flag can be used to specify the local directory where the kernel modules would be installed.
	If not specified, the default local directory is used.

	--download-only, -d
	This flag indicates whether to download the module to the local directory without installing.
	This flag cannot be used with the local_only flag.

	--local_only, -l
	This flag indicates whether to install modules directly from local directory without downloading.
	This flag cannot be used with the download_only flag.

	--recursive, -r
	This flag indicates whether to recursively installs modules from a gcs bucket directory and its
	subdirectories.
	This flag cannot be used with the local_only flag nor module files.
	Note that the modules are downloaded with their subdirectory path.

	cos-extensions list kernel-mods [flags]

	--gcs_path
	This flag can be used to specify the GCS path from which the available kernel modules would be
	listed. The expected syntax is gs://<bucket name>/<directory path>.
	If not specified, the default path is used.

	--local_dir
	This flag can be used to specify the local directory from which the downloaded or installed
	kernel modules would be listed.
	If not specified, the default path is used.

	--downloaded, -d
	This flag indicates whether to list only the downloaded compiled modules in a local directory.
	This can be used with local_dir flag, else it uses the default directory. It can be used with
	the recursive flag to include downloaded modules in the subdirectories.

	--installed, -i
	This flag indicates whether to list only the installed kernel modules in a local directory.
	This can be used with local_dir flag, else it uses the default directory. It can be used with
	the recursive flag to include installed modules in the subdirectories.

	--recursive, -r
	This flag indicates whether to recursively list all available modules in the gcs subdirectories.
	This can be used with gcs_path flag, else it uses default gcs path. It can be used with installed
	flag to list all installed kernels in the directory and its subdirectories. It can also be used
	with the downloaded flag.`,
	}

	installCmd = &cobra.Command{
	Use: "install <extension>",
	DisableFlagsInUseLine: true,
	Short: "Installs a COS extension.",
	}

	listCmd = &cobra.Command{
	Use: "list <extension>",
	DisableFlagsInUseLine: true,
	Short: "Lists all available COS extensions and their versions.",
	RunE: func(cmd *cobra.Command, args []string) error {
	printInstaller, _ := cmd.Flags().GetBool("gpu-installer")
	if slices.Contains(args, "--gpu-installer") {
	printInstaller = true
	}
	installers, err := gpu.GetCosInstaller(&http.Client{})
	if err != nil {
	return fmt.Errorf("Failed to set gpu installer: %v", err)
	}
	err = gpu.List(installers, printInstaller, args)
	if err != nil {
	return fmt.Errorf("Failed to run list command: %v", err)
	}
	return nil
	},
	}
	)

	func Execute() {
	if err := rootCmd.Execute(); err != nil {
	// Did not return error as it is handled by the RunE field.
	os.Exit(1)
	}
	}

	func init() {
	rootCmd.CompletionOptions.DisableDefaultCmd = true
	rootCmd.SetUsageTemplate(usageTemplate)
	rootCmd.SetHelpTemplate(helpTemplate)

	// Registering subcommands for cos-extensions
	rootCmd.AddCommand(installCmd)
	rootCmd.AddCommand(listCmd)
	}