Implement resource limits for Tasks

api/v1alpha1/task_types.go

@@ -57,6 +57,9 @@ type TaskSpec struct {
 	// source defines where AME will pull the project from.
 	// This can either be AME's own object storage or a git repository.
 	Source ProjectSource `json:"source,omitempty"`
+
+	// Resources define what resources this Task requires.
+	Resources v1.ResourceList `json:"resources,omitempty"`
 }
 
 // A ProjectSource describes where AME will fetch the project from.
@@ -215,6 +218,7 @@ func NewTask(runCmd string, projectId string) *Task {
 			Env:        []TaskEnvVar{},
 			Secrets:    []TaskSecret{},
 			Pipeline:   []PipelineStep{},
+			Resources:  v1.ResourceList{},
 		},
 	}
 }

cmd/ame/cli_test.go

@@ -20,6 +20,7 @@ import (
 	"github.com/google/go-cmp/cmp"
 	"golang.org/x/sync/errgroup"
 
+	corev1 "k8s.io/api/core/v1"
 	v1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 
 	"github.com/ActiveState/vt10x"
@@ -763,8 +764,9 @@ func TestCreateTaskConfig(t *testing.T) {
 						Value: "val1",
 					},
 				},
-				Secrets:  []amev1alpha1.TaskSecret{},
-				Pipeline: []amev1alpha1.PipelineStep{},
+				Secrets:   []amev1alpha1.TaskSecret{},
+				Pipeline:  []amev1alpha1.PipelineStep{},
+				Resources: corev1.ResourceList{},
 			},
 		},
 	}

config/crd/bases/ame.teainspace.com_reccurringtasks.yaml

@@ -115,6 +115,15 @@ spec:
                     description: A unique identifier for the project wich the Task
                       will be running based on.
                     type: string
+                  resources:
+                    additionalProperties:
+                      anyOf:
+                      - type: integer
+                      - type: string
+                      pattern: ^(\+|-)?(([0-9]+(\.[0-9]*)?)|(\.[0-9]+))(([KMGTPE]i)|[numkMGTPE]|([eE](\+|-)?(([0-9]+(\.[0-9]*)?)|(\.[0-9]+))))?$
+                      x-kubernetes-int-or-string: true
+                    description: Resource define what resources this Task requires.
+                    type: object
                   runcommand:
                     description: 'The command AME will use to execute the Task. The
                       command must be runnable from a bash shell. TODO: define propper

config/crd/bases/ame.teainspace.com_tasks.yaml

@@ -109,6 +109,15 @@ spec:
                 description: A unique identifier for the project wich the Task will
                   be running based on.
                 type: string
+              resources:
+                additionalProperties:
+                  anyOf:
+                  - type: integer
+                  - type: string
+                  pattern: ^(\+|-)?(([0-9]+(\.[0-9]*)?)|(\.[0-9]+))(([KMGTPE]i)|[numkMGTPE]|([eE](\+|-)?(([0-9]+(\.[0-9]*)?)|(\.[0-9]+))))?$
+                  x-kubernetes-int-or-string: true
+                description: Resource define what resources this Task requires.
+                type: object
               runcommand:
                 description: 'The command AME will use to execute the Task. The command
                   must be runnable from a bash shell. TODO: define propper requirements

internal/clients/clients.go

@@ -10,6 +10,7 @@ import (
 	"teainspace.com/ame/internal/common"
 
 	argo "github.com/argoproj/argo-workflows/v3/pkg/client/clientset/versioned/typed/workflow/v1alpha1"
+	corev1 "k8s.io/api/core/v1"
 )
 
 // TODO: should we handle errors when creating clients instead of dying?
@@ -52,3 +53,7 @@ func RecTasksClientFromConfig(cfg *rest.Config, ns string) v1alpha1.ReccurringTa
 func GenericTaskClientFromConfig(cfg *rest.Config, ns string) common.AmeGenClient[*amev1alpha1.Task] {
 	return common.NewAmeGenClient[*amev1alpha1.Task](TasksClientFromConfig(cfg, ns))
 }
+
+func GenericPodClientFromConfig(cfg *rest.Config, ns string) common.AmeGenClient[*corev1.Pod] {
+	return common.NewAmeGenClient[*corev1.Pod](PodClientFromConfig(cfg, ns))
+}

internal/workflows/workflows.go

@@ -5,6 +5,7 @@ package workflows
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	"time"
 
@@ -32,7 +33,10 @@ func GenWorkflowSpec(task v1alpha1.Task) (argo.WorkflowSpec, error) {
 	}
 
 	setupWfTemplate := genWfSetupTemplate("setup", &task, taskVolName(&task))
-	mainWfTemplate := genWfTemplate("main", &task, taskVolName(&task))
+	mainWfTemplate, err := genWfTemplate("main", &task, taskVolName(&task))
+	if err != nil {
+		return argo.WorkflowSpec{}, err
+	}
 
 	wfSteps := []argo.ParallelSteps{
 		{
@@ -83,7 +87,11 @@ func GenPipelineWf(t *v1alpha1.Task) (*argo.WorkflowSpec, error) {
 
 	for _, s := range t.Spec.Pipeline {
 		tSpec := v1alpha1.WfSpecFromPipelineStep(t, s)
-		wfTemplate := genWfTemplate(s.TaskName, v1alpha1.NewTaskFromSpec(tSpec, t.GetName()+s.TaskName), taskVolName(t))
+		wfTemplate, err := genWfTemplate(s.TaskName, v1alpha1.NewTaskFromSpec(tSpec, t.GetName()+s.TaskName), taskVolName(t))
+		if err != nil {
+			return nil, err
+		}
+
 		wfSteps = append(wfSteps, argo.ParallelSteps{
 			Steps: []argo.WorkflowStep{
 				{
@@ -153,7 +161,12 @@ func taskVolName(t *v1alpha1.Task) string {
 	return fmt.Sprintf("%s-volume", t.GetName())
 }
 
-func genWfTemplate(templateName string, t *v1alpha1.Task, volName string) argo.Template {
+func genWfTemplate(templateName string, t *v1alpha1.Task, volName string) (argo.Template, error) {
+	podSpecPatch, err := podSpecPatch(t.Spec)
+	if err != nil {
+		return argo.Template{}, err
+	}
+
 	return argo.Template{
 		Name: templateName,
 
@@ -204,9 +217,21 @@ func genWfTemplate(templateName string, t *v1alpha1.Task, volName string) argo.T
 			},
 		},
 
-		PodSpecPatch: `{"containers":[{"name":"main", "resources":{"limits":{
-        "memory": "3Gi"   }}}]}`,
+		PodSpecPatch: podSpecPatch,
+	}, nil
+}
+
+func podSpecPatch(t v1alpha1.TaskSpec) (string, error) {
+	if len(t.Resources) == 0 {
+		return "", nil
 	}
+
+	requirements, err := json.Marshal(t.Resources)
+	if err != nil {
+		return "", err
+	}
+
+	return fmt.Sprintf("{\"containers\":[{\"name\":\"main\", \"resources\":{\"limits\":%s}}]}", requirements), nil
 }
 
 // TaskEnvToContainerEnv constructs an array of EnvVar from t's environment