diff --git a/cmd/kill.go b/cmd/kill.go index 88d1d9f..47140ba 100644 --- a/cmd/kill.go +++ b/cmd/kill.go @@ -15,6 +15,9 @@ package cmd import ( + "strconv" + "strings" + realis "github.com/aurora-scheduler/gorealis/v2" "github.com/spf13/cobra" ) @@ -26,6 +29,7 @@ func init() { // Kill Job killCmd.AddCommand(killJobCmd) + killCmd.AddCommand(killTasksCmd) killJobCmd.Flags().StringVarP(env, "environment", "e", "", "Aurora Environment") killJobCmd.Flags().StringVarP(role, "role", "r", "", "Aurora Role") @@ -34,6 +38,17 @@ func init() { killJobCmd.MarkFlagRequired("environment") killJobCmd.MarkFlagRequired("role") killJobCmd.MarkFlagRequired("name") + + //Set flags for killTask sub-command + killTasksCmd.Flags().StringVarP(env, "environment", "e", "", "Aurora Environment") + killTasksCmd.Flags().StringVarP(role, "role", "r", "", "Aurora Role") + killTasksCmd.Flags().StringVarP(name, "name", "n", "", "Aurora Name") + killTasksCmd.Flags().StringVarP(instances, "instances", "i", "", "Instances e.g. 1, 2, 5") + killTasksCmd.Flags().BoolVarP(&monitor, "monitor", "m", true, "monitor the result after sending the command") + killTasksCmd.MarkFlagRequired("environment") + killTasksCmd.MarkFlagRequired("role") + killTasksCmd.MarkFlagRequired("name") + killTasksCmd.MarkFlagRequired("instance") } var killCmd = &cobra.Command{ @@ -47,6 +62,24 @@ var killJobCmd = &cobra.Command{ Run: killJob, } +/* +* The killTasks command allows the user to kill a specific task of a job. +* The command also allows the user to kill multiple tasks of the same job. To do so the user needs to pass a list of instance numbers as comma separated values. +* Pass the instance number of the job to be killed after the --instances or -i flag +* Please note that all the instances passed must belong to the same job. +* +* example : australis kill tasks -e "environment" -r "role" -n "job_name" -i "1" +* The above example kills instance number 1. +* +* example 2 : australis kill tasks -e "environment" -r "role" -n "job_name" -i "1, 5, 9" +* The above example kills tasks 1, 5 and 9, which are part of the same job + */ +var killTasksCmd = &cobra.Command{ + Use: "tasks", + Short: "Kill Aurora Tasks", + Run: killTasks, +} + func killJob(cmd *cobra.Command, args []string) { log.Infof("Killing job [Env:%s Role:%s Name:%s]\n", *env, *role, *name) @@ -64,3 +97,48 @@ func killJob(cmd *cobra.Command, args []string) { } } } + +func killTasks(cmd *cobra.Command, args []string) { + log.Infof("Killing task [Env:%s Role:%s Name:%s Instance:%s]\n", *env, *role, *name, *instances) + + //Set jobKey for the tasks to be killed. + task := realis.NewTask(). + Environment(*env). + Role(*role). + Name(*name) + + //Check that instance number is passed + if instances == nil { + log.Fatalln("Instance number not found. Please pass a valid instance number. If you want to pass multiple instances, please pass them as comma separated integer values") + } + + /* + * In the following block, we convert instance numbers, which were passed as strings, to integer values + * After converting them to integers, we add them to a slice of type int32. + */ + + splitString := strings.Split(*instances, ",") + instanceList := make([]int32, len(splitString)) + + for i := range instanceList { + splitString[i] = strings.TrimSpace(splitString[i]) + instanceNumber, intErr := strconv.Atoi(splitString[i]) + if intErr != nil { + log.Fatalln("Instance passed should be a number. Error: " + intErr.Error()) + return + } else { + instanceList[i] = int32(instanceNumber) + } + } + + //Call the killtasks function, passing the instanceList as the list of instances to be killed. + if _, err := client.KillInstances(task.JobKey(), instanceList...); err != nil { + log.Fatalln(err) + } + + if monitor { + if ok, err := client.MonitorInstances(task.JobKey(), 0, 5, 50); !ok || err != nil { + log.Fatalln("Unable to kill the given task") + } + } +} diff --git a/cmd/root.go b/cmd/root.go index cc72ac3..a2084e9 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -51,6 +51,7 @@ var monitor bool var timeout time.Duration var log = logrus.New() var taskStatus = new(string) +var instances = new(string) const australisVer = "v1.0.4" diff --git a/internal/job.go b/internal/job.go index 2f5e09b..b2ffc8b 100644 --- a/internal/job.go +++ b/internal/job.go @@ -93,7 +93,6 @@ func (j *Job) ToRealis() (*realis.AuroraJob, error) { RAM(j.RAM). Disk(j.Disk). AddPorts(int(j.Port)). - GPU(j.GPU). IsService(j.Service). Tier(j.Tier). Priority(j.Priority). @@ -101,6 +100,10 @@ func (j *Job) ToRealis() (*realis.AuroraJob, error) { InstanceCount(j.Instances). MaxFailure(j.MaxFailures) + if j.GPU > 0 { + auroraJob.GPU(j.GPU) + } + if j.CronSchedule != nil { auroraJob.CronSchedule(*j.CronSchedule) }