Add retry policy client package and armadactl commands

cmd/armadactl/cmd/commands.go

@@ -27,6 +27,7 @@ func createCmd(a *armadactl.App) *cobra.Command {
 	}
 	cmd.Flags().Bool("dry-run", false, "Validate the input file and exit without making any changes.")
 	cmd.AddCommand(queueCreateCmd())
+	cmd.AddCommand(retryPolicyCreateCmd())
 	return cmd
 }
 
@@ -37,6 +38,7 @@ func deleteCmd() *cobra.Command {
 		Long:  "Delete Armada resource. Supported: queue",
 	}
 	cmd.AddCommand(queueDeleteCmd())
+	cmd.AddCommand(retryPolicyDeleteCmd())
 	return cmd
 }
 
@@ -47,6 +49,7 @@ func updateCmd() *cobra.Command {
 		Long:  "Update Armada resource. Supported: queue",
 	}
 	cmd.AddCommand(queueUpdateCmd())
+	cmd.AddCommand(retryPolicyUpdateCmd())
 	return cmd
 }
 
@@ -59,6 +62,8 @@ func getCmd() *cobra.Command {
 	cmd.AddCommand(
 		queueGetCmd(),
 		queuesGetCmd(),
+		retryPolicyGetCmd(),
+		retryPolicyGetAllCmd(),
 		getSchedulingReportCmd(armadactl.New()),
 		getQueueSchedulingReportCmd(armadactl.New()),
 		getJobSchedulingReportCmd(armadactl.New()),

cmd/armadactl/cmd/params.go

@@ -8,6 +8,7 @@ import (
 	ce "github.com/armadaproject/armada/pkg/client/executor"
 	cn "github.com/armadaproject/armada/pkg/client/node"
 	cq "github.com/armadaproject/armada/pkg/client/queue"
+	crp "github.com/armadaproject/armada/pkg/client/retrypolicy"
 )
 
 // initParams initialises the command parameters, flags, and a configuration file.
@@ -33,6 +34,12 @@ func initParams(cmd *cobra.Command, params *armadactl.Params) error {
 	params.QueueAPI.Preempt = cq.Preempt(client.ExtractCommandlineArmadaApiConnectionDetails)
 	params.QueueAPI.Cancel = cq.Cancel(client.ExtractCommandlineArmadaApiConnectionDetails)
 
+	params.RetryPolicyAPI.Create = crp.Create(client.ExtractCommandlineArmadaApiConnectionDetails)
+	params.RetryPolicyAPI.Delete = crp.Delete(client.ExtractCommandlineArmadaApiConnectionDetails)
+	params.RetryPolicyAPI.Get = crp.Get(client.ExtractCommandlineArmadaApiConnectionDetails)
+	params.RetryPolicyAPI.GetAll = crp.GetAll(client.ExtractCommandlineArmadaApiConnectionDetails)
+	params.RetryPolicyAPI.Update = crp.Update(client.ExtractCommandlineArmadaApiConnectionDetails)
+
 	params.ExecutorAPI.Cordon = ce.CordonExecutor(client.ExtractCommandlineArmadaApiConnectionDetails)
 	params.ExecutorAPI.Uncordon = ce.UncordonExecutor(client.ExtractCommandlineArmadaApiConnectionDetails)
 

cmd/armadactl/cmd/queue.go

@@ -12,6 +12,8 @@ import (
 	"github.com/armadaproject/armada/pkg/client/queue"
 )
 
+const retryPoliciesFlag = "retry-policies"
+
 func queueCreateCmd() *cobra.Command {
 	return queueCreateCmdWithApp(armadactl.New())
 }
@@ -63,13 +65,19 @@ Job priority is evaluated inside queue, queue has its own priority.  Any labels
 				return fmt.Errorf("error converting queue labels to map: %s", err)
 			}
 
+			retryPolicies, err := cmd.Flags().GetStringSlice(retryPoliciesFlag)
+			if err != nil {
+				return fmt.Errorf("error reading retry-policies: %s", err)
+			}
+
 			newQueue, err := queue.NewQueue(&api.Queue{
 				Name:           name,
 				PriorityFactor: priorityFactor,
 				UserOwners:     owners,
 				GroupOwners:    groups,
 				Cordoned:       cordoned,
 				Labels:         labelsAsMap,
+				RetryPolicies:  retryPolicies,
 			})
 			if err != nil {
 				return fmt.Errorf("invalid queue data: %s", err)
@@ -83,6 +91,7 @@ Job priority is evaluated inside queue, queue has its own priority.  Any labels
 	cmd.Flags().StringSlice("group-owners", []string{}, "Comma separated list of queue group owners, defaults to empty list.")
 	cmd.Flags().Bool("cordon", false, "Used to pause scheduling on specified queue. Defaults to false.")
 	cmd.Flags().StringSliceP("labels", "l", []string{}, "Comma separated list of key-value queue labels, for example: armadaproject.io/submitter=airflow. Defaults to empty list.")
+	cmd.Flags().StringSlice(retryPoliciesFlag, []string{}, "Comma separated list of retry policy names to assign to this queue, in evaluation order. Defaults to empty list.")
 	return cmd
 }
 
@@ -192,8 +201,14 @@ func queueUpdateCmdWithApp(a *armadactl.App) *cobra.Command {
 	cmd := &cobra.Command{
 		Use:   "queue <queue-name>",
 		Short: "Update an existing queue",
-		Long:  "Update settings of an existing queue",
-		Args:  cobra.ExactArgs(1),
+		Long: `Update settings of an existing queue.
+
+This is a full replace, not a partial patch. Every queue attribute is set from
+the flags on this command, and any flag you omit resets that attribute to its
+default. If the queue has retry policies attached, pass --retry-policies on
+every update, otherwise the attachment is cleared and the queue falls back to
+the default retry behaviour.`,
+		Args: cobra.ExactArgs(1),
 		PreRunE: func(cmd *cobra.Command, args []string) error {
 			return initParams(cmd, a.Params)
 		},
@@ -230,13 +245,19 @@ func queueUpdateCmdWithApp(a *armadactl.App) *cobra.Command {
 				return fmt.Errorf("error converting queue labels to map: %s", err)
 			}
 
+			retryPolicies, err := cmd.Flags().GetStringSlice(retryPoliciesFlag)
+			if err != nil {
+				return fmt.Errorf("error reading retry-policies: %s", err)
+			}
+
 			newQueue, err := queue.NewQueue(&api.Queue{
 				Name:           name,
 				PriorityFactor: priorityFactor,
 				UserOwners:     owners,
 				GroupOwners:    groups,
 				Cordoned:       cordoned,
 				Labels:         labelsAsMap,
+				RetryPolicies:  retryPolicies,
 			})
 			if err != nil {
 				return fmt.Errorf("invalid queue data: %s", err)
@@ -251,5 +272,6 @@ func queueUpdateCmdWithApp(a *armadactl.App) *cobra.Command {
 	cmd.Flags().StringSlice("group-owners", []string{}, "Comma separated list of queue group owners, defaults to empty list.")
 	cmd.Flags().Bool("cordon", false, "Used to pause scheduling on specified queue. Defaults to false.")
 	cmd.Flags().StringSliceP("labels", "l", []string{}, "Comma separated list of key-value queue labels, for example: armadaproject.io/submitter=airflow. Defaults to empty list.")
+	cmd.Flags().StringSlice(retryPoliciesFlag, []string{}, "Comma separated list of retry policy names to assign to this queue, in evaluation order. Defaults to empty list.")
 	return cmd
 }

cmd/armadactl/cmd/retrypolicy.go

@@ -0,0 +1,98 @@
+package cmd
+
+import (
+	"github.com/spf13/cobra"
+
+	"github.com/armadaproject/armada/internal/armadactl"
+)
+
+func retryPolicyCreateCmd() *cobra.Command {
+	a := armadactl.New()
+	return retryPolicyFileCmd(a,
+		"Create a retry policy from a YAML/JSON file",
+		"Create a retry policy that defines rules for whether failed jobs should be retried.",
+		a.CreateRetryPolicyFromFile)
+}
+
+func retryPolicyUpdateCmd() *cobra.Command {
+	a := armadactl.New()
+	return retryPolicyFileCmd(a,
+		"Update a retry policy from a YAML/JSON file",
+		"Update an existing retry policy with the definition from a YAML/JSON file.",
+		a.UpdateRetryPolicyFromFile)
+}
+
+func retryPolicyGetCmd() *cobra.Command {
+	a := armadactl.New()
+	return retryPolicyNameCmd(a,
+		"Get a retry policy by name",
+		"Get the definition of a retry policy by its name.",
+		a.GetRetryPolicy)
+}
+
+func retryPolicyDeleteCmd() *cobra.Command {
+	a := armadactl.New()
+	return retryPolicyNameCmd(a,
+		"Delete a retry policy by name",
+		"Delete an existing retry policy by its name.",
+		a.DeleteRetryPolicy)
+}
+
+func retryPolicyGetAllCmd() *cobra.Command {
+	a := armadactl.New()
+	return &cobra.Command{
+		Use:   "retry-policies",
+		Short: "List all retry policies",
+		Long:  "List all retry policies defined in the system.",
+		Args:  cobra.NoArgs,
+		PreRunE: func(cmd *cobra.Command, args []string) error {
+			return initParams(cmd, a.Params)
+		},
+		RunE: func(cmd *cobra.Command, args []string) error {
+			return a.GetAllRetryPolicies()
+		},
+	}
+}
+
+// retryPolicyFileCmd builds a command that reads a retry policy from a
+// YAML/JSON file and applies it via run.
+func retryPolicyFileCmd(a *armadactl.App, short, long string, run func(fileName string) error) *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "retry-policy",
+		Short: short,
+		Long:  long,
+		Args:  cobra.NoArgs,
+		PreRunE: func(cmd *cobra.Command, args []string) error {
+			return initParams(cmd, a.Params)
+		},
+		RunE: func(cmd *cobra.Command, args []string) error {
+			filePath, err := cmd.Flags().GetString("file")
+			if err != nil {
+				return err
+			}
+			return run(filePath)
+		},
+	}
+	cmd.Flags().StringP("file", "f", "", "Path to YAML/JSON file defining the retry policy.")
+	if err := cmd.MarkFlagRequired("file"); err != nil {
+		panic(err)
+	}
+	return cmd
+}
+
+// retryPolicyNameCmd builds a command that takes a single retry policy name
+// argument and applies it via run.
+func retryPolicyNameCmd(a *armadactl.App, short, long string, run func(name string) error) *cobra.Command {
+	return &cobra.Command{
+		Use:   "retry-policy <name>",
+		Short: short,
+		Long:  long,
+		Args:  cobra.ExactArgs(1),
+		PreRunE: func(cmd *cobra.Command, args []string) error {
+			return initParams(cmd, a.Params)
+		},
+		RunE: func(cmd *cobra.Command, args []string) error {
+			return run(args[0])
+		},
+	}
+}

internal/armadactl/app.go

@@ -21,6 +21,7 @@ import (
 	"github.com/armadaproject/armada/pkg/client/executor"
 	"github.com/armadaproject/armada/pkg/client/node"
 	"github.com/armadaproject/armada/pkg/client/queue"
+	"github.com/armadaproject/armada/pkg/client/retrypolicy"
 )
 
 type App struct {
@@ -41,6 +42,7 @@ type App struct {
 type Params struct {
 	ApiConnectionDetails *client.ApiConnectionDetails
 	QueueAPI             *QueueAPI
+	RetryPolicyAPI       *RetryPolicyAPI
 	ExecutorAPI          *ExecutorAPI
 	NodeAPI              *NodeAPI
 }
@@ -69,6 +71,14 @@ type ExecutorAPI struct {
 	PreemptOnExecutor executor.PreemptAPI
 }
 
+type RetryPolicyAPI struct {
+	Create retrypolicy.CreateAPI
+	Delete retrypolicy.DeleteAPI
+	Get    retrypolicy.GetAPI
+	GetAll retrypolicy.GetAllAPI
+	Update retrypolicy.UpdateAPI
+}
+
 type NodeAPI struct {
 	PreemptOnNode node.PreemptAPI
 	CancelOnNode  node.CancelAPI
@@ -79,9 +89,10 @@ type NodeAPI struct {
 func New() *App {
 	return &App{
 		Params: &Params{
-			QueueAPI:    &QueueAPI{},
-			ExecutorAPI: &ExecutorAPI{},
-			NodeAPI:     &NodeAPI{},
+			QueueAPI:       &QueueAPI{},
+			RetryPolicyAPI: &RetryPolicyAPI{},
+			ExecutorAPI:    &ExecutorAPI{},
+			NodeAPI:        &NodeAPI{},
 		},
 		Out:    os.Stdout,
 		Random: rand.Reader,

internal/armadactl/queue.go

@@ -55,6 +55,15 @@ func (a *App) CreateResource(fileName string, dryRun bool) error {
 		if !dryRun {
 			return a.Params.QueueAPI.Create(queue)
 		}
+	case client.ResourceKindRetryPolicy:
+		// Parse the file even on a dry run so a malformed policy is caught.
+		policy, err := retryPolicyFromFile(fileName)
+		if err != nil {
+			return err
+		}
+		if !dryRun {
+			return a.CreateRetryPolicy(policy)
+		}
 	default:
 		return errors.Errorf("invalid resource kind: %s", resource.Kind)
 	}