dcu-tracker.go

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 <dcu_id> <operation>", 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
}