package dcuTracker import ( "dtk-container-toolkit/cmd/dtk-ctk/dcu-tracker/disable" "dtk-container-toolkit/cmd/dtk-ctk/dcu-tracker/enable" "dtk-container-toolkit/cmd/dtk-ctk/dcu-tracker/initialize" "dtk-container-toolkit/cmd/dtk-ctk/dcu-tracker/release" "dtk-container-toolkit/cmd/dtk-ctk/dcu-tracker/reset" "dtk-container-toolkit/cmd/dtk-ctk/dcu-tracker/status" "dtk-container-toolkit/internal/dcu-tracker" "dtk-container-toolkit/internal/logger" "fmt" "os/user" "github.com/urfave/cli/v2" ) type command struct { logger logger.Interface } func NewCommand(logger logger.Interface) *cli.Command { c := command{ logger: logger, } return c.build() } func (m command) build() *cli.Command { dcuTrackerCmd := cli.Command{ Name: "dcu-tracker", Usage: "DCU Tracker related commands", UsageText: `dtk-ctk dcu-tracker [dcu-ids] [accessibility] Arguments: dcu-ids Comma-separated list of DCU IDs (comma separated list, range operator, all) accessibility Must be either 'exclusive' or 'shared' Examples: dtk-ctk dcu-tracker 0,1,2 exclusive dtk-ctk dcu-tracker 0,1-2 shared dtk-ctk dcu-tracker all shared OR dtk-ctk dcu-tracker [command] [options]`, Before: func(c *cli.Context) error { return m.validateGenOptions(c) }, Action: func(c *cli.Context) error { return m.performAction(c) }, } dcuTrackerCmd.Subcommands = []*cli.Command{ disable.NewCommand(m.logger), enable.NewCommand(m.logger), initialize.NewCommand(m.logger), reset.NewCommand(m.logger), release.NewCommand(m.logger), status.NewCommand(m.logger), } return &dcuTrackerCmd } func (m command) validateGenOptions(c *cli.Context) error { curUser, err := user.Current() if err != nil || curUser.Uid != "0" { return fmt.Errorf("Permission denied: Not running as root") } return nil } func (m command) performAction(c *cli.Context) error { if c.Args().Len() == 0 { return cli.ShowAppHelp(c) } if c.Args().Len() > 2 { return cli.Exit("Error: Missing arguments. Usgae: dcu-tracker ", 1) } gpuIDs := c.Args().Get(0) operation := c.Args().Get(1) dcuTracker, err := dcuTracker.New() if err != nil { return fmt.Errorf("Failed to create DCU tracker, Error: %v", err) } switch operation { case "exclusive": dcuTracker.MakeDCUsExclusive(gpuIDs) case "shared": dcuTracker.MakeDCUsShared(gpuIDs) default: return cli.Exit("Error: Invalid operation. Must be 'exclusive' or 'shared", 1) } return nil }