Merge pull request #1794 from Onder7994/worker_request_limit

Worker. Add option to set k8s resources.

functions/onprem/orborus/orborus.go

@@ -52,6 +52,7 @@ import (
 	appsv1 "k8s.io/api/apps/v1"
 	corev1 "k8s.io/api/core/v1"
 	rbacv1 "k8s.io/api/rbac/v1"
+	"k8s.io/apimachinery/pkg/api/resource"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/util/intstr"
 )
@@ -799,6 +800,48 @@ func buildEnvVars(envMap map[string]string) []corev1.EnvVar {
 	return envVars
 }
 
+func buildResourcesFromEnv() corev1.ResourceRequirements {
+	requests := corev1.ResourceList{}
+	limits := corev1.ResourceList{}
+
+	type item struct {
+		env          string
+		resourceName corev1.ResourceName
+		resourceList corev1.ResourceList
+	}
+
+	items := []item{
+		// kubernetes requests
+		{env: "SHUFFLE_WORKER_CPU_REQUEST", resourceName: corev1.ResourceCPU, resourceList: requests},
+		{env: "SHUFFLE_WORKER_MEMORY_REQUEST", resourceName: corev1.ResourceMemory, resourceList: requests},
+		{env: "SHUFFLE_WORKER_EPHEMERAL_STORAGE_REQUEST", resourceName: corev1.ResourceEphemeralStorage, resourceList: requests},
+		// kubernetes limits
+		{env: "SHUFFLE_WORKER_CPU_LIMIT", resourceName: corev1.ResourceCPU, resourceList: limits},
+		{env: "SHUFFLE_WORKER_MEMORY_LIMIT", resourceName: corev1.ResourceMemory, resourceList: limits},
+		{env: "SHUFFLE_WORKER_EPHEMERAL_STORAGE_LIMIT", resourceName: corev1.ResourceEphemeralStorage, resourceList: limits},
+	}
+
+	for _, it := range items {
+		if value := strings.TrimSpace(os.Getenv(it.env)); value != "" {
+			if quantity, err := resource.ParseQuantity(value); err == nil {
+				it.resourceList[it.resourceName] = quantity
+			} else {
+				log.Printf("[WARNING] Cannot parse %s=%q as resource quantity: %v", it.env, value, err)
+			}
+		}
+	}
+
+	rr := corev1.ResourceRequirements{}
+	if len(requests) > 0 {
+		rr.Requests = requests
+	}
+	if len(limits) > 0 {
+		rr.Limits = limits
+	}
+
+	return rr
+}
+
 func handleBackendImageDownload(ctx context.Context, images string) error {
 
 	// Replicate images with lowercase, as the name may be wrong
@@ -1042,6 +1085,20 @@ func deployK8sWorker(image string, identifier string, env []string) error {
 	env = append(env, fmt.Sprintf("IS_KUBERNETES=true"))
 	env = append(env, fmt.Sprintf("KUBERNETES_NAMESPACE=%s", os.Getenv("KUBERNETES_NAMESPACE")))
 
+	// app resource env
+	for _, k := range []string{
+		"SHUFFLE_APP_CPU_REQUEST",
+		"SHUFFLE_APP_MEMORY_REQUEST",
+		"SHUFFLE_APP_EPHEMERAL_STORAGE_REQUEST",
+		"SHUFFLE_APP_CPU_LIMIT",
+		"SHUFFLE_APP_MEMORY_LIMIT",
+		"SHUFFLE_APP_EPHEMERAL_STORAGE_LIMIT",
+	} {
+		if v := os.Getenv(k); v != "" {
+			env = append(env, fmt.Sprintf("%s=%s", k, v))
+		}
+	}
+
 	if len(os.Getenv("KUBERNETES_SERVICE_HOST")) > 0 {
 		env = append(env, fmt.Sprintf("KUBERNETES_SERVICE_HOST=%s", os.Getenv("KUBERNETES_SERVICE_HOST")))
 	}
@@ -1206,6 +1263,7 @@ func deployK8sWorker(image string, identifier string, env []string) error {
 		Image:           kubernetesImage,
 		Env:             buildEnvVars(envMap),
 		SecurityContext: containerSecurityContext,
+		Resources:       buildResourcesFromEnv(),
 
 		//ImagePullPolicy: "Never",
 		ImagePullPolicy: corev1.PullIfNotPresent,

functions/onprem/worker/worker.go

@@ -43,6 +43,7 @@ import (
 	//k8s deps
 	appsv1 "k8s.io/api/apps/v1"
 	corev1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/api/resource"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/util/intstr"
 	"k8s.io/client-go/kubernetes"
@@ -652,6 +653,7 @@ func deployk8sApp(image string, identifier string, env []string) error {
 								},
 							},
 							SecurityContext: containerSecurityContext,
+							Resources:       buildResourcesFromEnv(),
 						},
 					},
 					DNSPolicy:          corev1.DNSClusterFirst,
@@ -2436,6 +2438,49 @@ func buildEnvVars(envMap map[string]string) []corev1.EnvVar {
 	}
 	return envVars
 }
+
+func buildResourcesFromEnv() corev1.ResourceRequirements {
+	requests := corev1.ResourceList{}
+	limits := corev1.ResourceList{}
+
+	type item struct {
+		env          string
+		resourceName corev1.ResourceName
+		resourceList corev1.ResourceList
+	}
+
+	items := []item{
+		// kubernetes requests
+		{env: "SHUFFLE_APP_CPU_REQUEST", resourceName: corev1.ResourceCPU, resourceList: requests},
+		{env: "SHUFFLE_APP_MEMORY_REQUEST", resourceName: corev1.ResourceMemory, resourceList: requests},
+		{env: "SHUFFLE_APP_EPHEMERAL_STORAGE_REQUEST", resourceName: corev1.ResourceEphemeralStorage, resourceList: requests},
+		// kubernetes limits
+		{env: "SHUFFLE_APP_CPU_LIMIT", resourceName: corev1.ResourceCPU, resourceList: limits},
+		{env: "SHUFFLE_APP_MEMORY_LIMIT", resourceName: corev1.ResourceMemory, resourceList: limits},
+		{env: "SHUFFLE_APP_EPHEMERAL_STORAGE_LIMIT", resourceName: corev1.ResourceEphemeralStorage, resourceList: limits},
+	}
+
+	for _, it := range items {
+		if value := strings.TrimSpace(os.Getenv(it.env)); value != "" {
+			if quantity, err := resource.ParseQuantity(value); err == nil {
+				it.resourceList[it.resourceName] = quantity
+			} else {
+				log.Printf("[WARNING] Cannot parse %s=%q as resource quantity: %v", it.env, value, err)
+			}
+		}
+	}
+
+	rr := corev1.ResourceRequirements{}
+	if len(requests) > 0 {
+		rr.Requests = requests
+	}
+	if len(limits) > 0 {
+		rr.Limits = limits
+	}
+
+	return rr
+}
+
 func getWorkerBackendExecution(auth string, executionId string) (*shuffle.WorkflowExecution, error) {
 	backendUrl := os.Getenv("BASE_URL")
 	if len(backendUrl) == 0 {
@@ -2446,9 +2491,9 @@ func getWorkerBackendExecution(auth string, executionId string) (*shuffle.Workfl
 
 	streamResultUrl := fmt.Sprintf("%s/api/v1/streams/results", backendUrl)
 	topClient := shuffle.GetExternalClient(backendUrl)
-	requestData := shuffle.ActionResult {
+	requestData := shuffle.ActionResult{
 		Authorization: auth,
-		ExecutionId: executionId,
+		ExecutionId:   executionId,
 	}
 
 	data, err := json.Marshal(requestData)
@@ -2708,7 +2753,7 @@ func runWorkflowExecutionTransaction(ctx context.Context, attempts int64, workfl
 	}
 
 	if setExecution || workflowExecution.Status == "FINISHED" || workflowExecution.Status == "ABORTED" || workflowExecution.Status == "FAILURE" {
-		if debug { 
+		if debug {
 			log.Printf("[DEBUG][%s] Running setexec with status %s and %d/%d results", workflowExecution.ExecutionId, workflowExecution.Status, len(workflowExecution.Results), len(workflowExecution.Workflow.Actions))
 		}
 
@@ -2726,7 +2771,7 @@ func runWorkflowExecutionTransaction(ctx context.Context, attempts int64, workfl
 		if os.Getenv("SHUFFLE_SWARM_CONFIG") == "run" || os.Getenv("SHUFFLE_SWARM_CONFIG") == "swarm" {
 			finished := shuffle.ValidateFinished(ctx, -1, *workflowExecution)
 			if !finished {
-				if debug { 
+				if debug {
 					log.Printf("[DEBUG][%s] Handling next node since it's not finished!", workflowExecution.ExecutionId)
 				}
 
@@ -3926,7 +3971,7 @@ func sendAppRequest(ctx context.Context, incomingUrl, appName string, port int,
 		log.Printf("[ERROR] Failed reading app request body body: %s", err)
 		return err
 	} else {
-		if debug { 
+		if debug {
 			log.Printf("[DEBUG][%s] NEWRESP (from app): %s", workflowExecution.ExecutionId, string(body))
 		}
 	}