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

 package commands

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

 	"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 {
 			installers, err := gpu.GetCosInstaller(&http.Client{})
 			if err != nil {
 				return fmt.Errorf("Failed to set gpu installer: %v", err)
 			}
 			err = gpu.List(installers, false, 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"

	"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 {
	installers, err := gpu.GetCosInstaller(&http.Client{})
	if err != nil {
	return fmt.Errorf("Failed to set gpu installer: %v", err)
	}
	err = gpu.List(installers, false, 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)
	}