Separating single and multi-queue scheduling control commands in arma…

…dactl
armadaproject · Jul 25, 2024 · 6080eff · 6080eff
1 parent 09204b8
commit 6080eff
Show file tree

Hide file tree

Showing 3 changed files with 103 additions and 4 deletions.
diff --git a/cmd/armadactl/cmd/scheduling_control.go b/cmd/armadactl/cmd/scheduling_control.go
@@ -15,6 +15,7 @@ func cordon() *cobra.Command {
 		Short: "Pause scheduling by resource",
 		Long:  "Pause scheduling by resource. Supported: queue",
 	}
+	cmd.AddCommand(cordonQueues(a))
 	cmd.AddCommand(cordonQueue(a))
 	return cmd
 }
@@ -26,13 +27,14 @@ func uncordon() *cobra.Command {
 		Short: "Resume scheduling by resource",
 		Long:  "Resume scheduling by resource. Supported: queue",
 	}
+	cmd.AddCommand(uncordonQueues(a))
 	cmd.AddCommand(uncordonQueue(a))
 	return cmd
 }
 
-func cordonQueue(a *armadactl.App) *cobra.Command {
+func cordonQueues(a *armadactl.App) *cobra.Command {
 	cmd := &cobra.Command{
-		Use:   "queue",
+		Use:   "queues",
 		Short: "Pause scheduling for select queues",
 		Long:  "Pause scheduling for select queues. This can be achieved by queue name or by labels.",
 		PreRunE: func(cmd *cobra.Command, args []string) error {
@@ -73,9 +75,9 @@ func cordonQueue(a *armadactl.App) *cobra.Command {
 	return cmd
 }
 
-func uncordonQueue(a *armadactl.App) *cobra.Command {
+func uncordonQueues(a *armadactl.App) *cobra.Command {
 	cmd := &cobra.Command{
-		Use:   "queue",
+		Use:   "queues",
 		Short: "Resume scheduling for select queues",
 		Long:  "Resume scheduling for select queues. This can be achieved by queue name or by labels.",
 		PreRunE: func(cmd *cobra.Command, args []string) error {
@@ -115,3 +117,49 @@ func uncordonQueue(a *armadactl.App) *cobra.Command {
 	cmd.Flags().Bool("dry-run", false, "Show selection of queues that will be modified in this operation")
 	return cmd
 }
+
+func cordonQueue(a *armadactl.App) *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "queue <queue>",
+		Short: "Pause scheduling for a given queue",
+		PreRunE: func(cmd *cobra.Command, args []string) error {
+			return initParams(cmd, a.Params)
+		},
+		Args: cobra.ExactArgs(1),
+		RunE: func(cmd *cobra.Command, args []string) error {
+			queue := args[0]
+
+			dryRun, err := cmd.Flags().GetBool("dry-run")
+			if err != nil {
+				return fmt.Errorf("error reading dry-run flag: %s", err)
+			}
+
+			return a.CordonQueue(queue, dryRun)
+		},
+	}
+	cmd.Flags().Bool("dry-run", false, "Show selection of queues that will be modified in this operation")
+	return cmd
+}
+
+func uncordonQueue(a *armadactl.App) *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "queue <queue>",
+		Short: "Resume scheduling for a given queue",
+		PreRunE: func(cmd *cobra.Command, args []string) error {
+			return initParams(cmd, a.Params)
+		},
+		Args: cobra.ExactArgs(1),
+		RunE: func(cmd *cobra.Command, args []string) error {
+			queue := args[0]
+
+			dryRun, err := cmd.Flags().GetBool("dry-run")
+			if err != nil {
+				return fmt.Errorf("error reading dry-run flag: %s", err)
+			}
+
+			return a.UncordonQueue(queue, dryRun)
+		},
+	}
+	cmd.Flags().Bool("dry-run", false, "Show selection of queues that will be modified in this operation")
+	return cmd
+}
diff --git a/internal/armadactl/queue.go b/internal/armadactl/queue.go
@@ -58,6 +58,15 @@ func (a *App) DeleteQueue(name string) error {
 	return nil
 }
 
+func (a *App) getQueueAsAPIQueue(queue string) (*api.Queue, error) {
+	queueToReturn, err := a.Params.QueueAPI.Get(queue)
+	if err != nil {
+		return nil, errors.Errorf("[armadactl.getQueueAsAPIQueue] error getting queue %s: %s", queue, err)
+	}
+
+	return queueToReturn, nil
+}
+
 // GetQueue calls app.QueueAPI.Get with the provided parameters.
 func (a *App) GetQueue(name string) error {
 	queue, err := a.Params.QueueAPI.Get(name)

diff --git a/internal/armadactl/scheduling_control.go b/internal/armadactl/scheduling_control.go
@@ -17,6 +17,27 @@ func (a *App) cordonQueue(q *api.Queue) error {
 	return err
 }
 
+func (a *App) CordonQueue(queue string, dryRun bool) error {
+	selectedQueue, err := a.getQueueAsAPIQueue(queue)
+	if err != nil {
+		return fmt.Errorf("error retrieving queue %s: %s", queue, err)
+	}
+
+	if dryRun {
+		fmt.Printf("Pausing scheduling for the following queues: (DRY RUN)")
+		fmt.Println(queue)
+	} else {
+		fmt.Println("Pausing scheduling for the following queues:")
+		err = a.cordonQueue(selectedQueue)
+		if err != nil {
+			return fmt.Errorf("could not pause scheduling on queue %s: %s", queue, err)
+		} else {
+			fmt.Printf("%s paused\n", queue)
+		}
+	}
+	return nil
+}
+
 func (a *App) CordonQueues(matchQueues []string, matchLabels []string, dryRun bool, inverse bool) error {
 	selectedQueues, err := a.getAllQueuesAsAPIQueue(matchQueues, matchLabels, inverse)
 	if err != nil {
@@ -52,6 +73,27 @@ func (a *App) uncordonQueue(q *api.Queue) error {
 	return err
 }
 
+func (a *App) UncordonQueue(queue string, dryRun bool) error {
+	selectedQueue, err := a.getQueueAsAPIQueue(queue)
+	if err != nil {
+		return fmt.Errorf("error retrieving queue %s: %s", queue, err)
+	}
+
+	if dryRun {
+		fmt.Printf("Resuming scheduling for the following queues: (DRY RUN)")
+		fmt.Println(queue)
+	} else {
+		fmt.Println("Resuming scheduling for the following queues:")
+		err = a.uncordonQueue(selectedQueue)
+		if err != nil {
+			return fmt.Errorf("could not resume scheduling on queue %s: %s", queue, err)
+		} else {
+			fmt.Printf("%s resumed\n", queue)
+		}
+	}
+	return nil
+}
+
 func (a *App) UncordonQueues(matchQueues []string, matchLabels []string, dryRun bool, inverse bool) error {
 	selectedQueues, err := a.getAllQueuesAsAPIQueue(matchQueues, matchLabels, inverse)
 	if err != nil {