Inject container resource during BuildRawContainer (#335)

* Inject container resource during BuildRawContainer Signed-off-by: byhsu <[email protected]> * address comment Signed-off-by: byhsu <[email protected]> --------- Signed-off-by: byhsu <[email protected]> Co-authored-by: byhsu <[email protected]>
flyteorg · Apr 7, 2023 · a3f7334 · a3f7334
1 parent e84d27a
commit a3f7334
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/go/tasks/pluginmachinery/flytek8s/container_helper.go b/go/tasks/pluginmachinery/flytek8s/container_helper.go
@@ -204,13 +204,19 @@ func BuildRawContainer(ctx context.Context, taskContainer *core.Container, taskE
 		containerName = rand.String(4)
 	}
 
+	res, err := ToK8sResourceRequirements(taskContainer.Resources)
+	if err != nil {
+		return nil, err
+	}
+
 	container := &v1.Container{
 		Name:                     containerName,
 		Image:                    taskContainer.GetImage(),
 		Args:                     taskContainer.GetArgs(),
 		Command:                  taskContainer.GetCommand(),
 		Env:                      ToK8sEnvVar(taskContainer.GetEnv()),
 		TerminationMessagePolicy: v1.TerminationMessageFallbackToLogsOnError,
+		Resources:                *res,
 	}
 
 	return container, nil
@@ -251,7 +257,7 @@ func ToK8sContainer(ctx context.Context, tCtx pluginscore.TaskExecutionContext)
 		Task:             tCtx.TaskReader(),
 	}
 
-	if err := AddFlyteCustomizationsToContainer(ctx, templateParameters, ResourceCustomizationModeAssignResources, container); err != nil {
+	if err := AddFlyteCustomizationsToContainer(ctx, templateParameters, ResourceCustomizationModeMergeExistingResources, container); err != nil {
 		return nil, err
 	}