From 3c837f112ce51fb834407bdb76b2035fc8adf068 Mon Sep 17 00:00:00 2001 From: Paul Dittamo Date: Tue, 5 Dec 2023 19:26:06 -0800 Subject: [PATCH 1/5] don't override service account from security context if already set Signed-off-by: Paul Dittamo --- .../tasks/plugins/k8s/pod/container_test.go | 61 +++++++++++++++++++ .../go/tasks/plugins/k8s/pod/plugin.go | 4 +- 2 files changed, 64 insertions(+), 1 deletion(-) diff --git a/flyteplugins/go/tasks/plugins/k8s/pod/container_test.go b/flyteplugins/go/tasks/plugins/k8s/pod/container_test.go index 19000e0c72..a01caca2ef 100644 --- a/flyteplugins/go/tasks/plugins/k8s/pod/container_test.go +++ b/flyteplugins/go/tasks/plugins/k8s/pod/container_test.go @@ -19,6 +19,7 @@ import ( flytek8sConfig "github.com/flyteorg/flyte/flyteplugins/go/tasks/pluginmachinery/flytek8s/config" pluginsIOMock "github.com/flyteorg/flyte/flyteplugins/go/tasks/pluginmachinery/io/mocks" "github.com/flyteorg/flyte/flyteplugins/go/tasks/pluginmachinery/k8s" + "github.com/flyteorg/flyte/flyteplugins/go/tasks/pluginmachinery/utils" ) var containerResourceRequirements = &v1.ResourceRequirements{ @@ -28,6 +29,8 @@ var containerResourceRequirements = &v1.ResourceRequirements{ }, } +var podTemplateServiceAccount = "test-service-account" + func dummyContainerTaskTemplate(command []string, args []string) *core.TaskTemplate { return &core.TaskTemplate{ Type: "test", @@ -40,6 +43,39 @@ func dummyContainerTaskTemplate(command []string, args []string) *core.TaskTempl } } +func dummyContainerTaskTemplateWithPodSpec(command []string, args []string) *core.TaskTemplate { + + podSpec := v1.PodSpec{ + Containers: []v1.Container{ + v1.Container{ + Name: "test-image", + Command: command, + Args: args, + }, + }, + ServiceAccountName: podTemplateServiceAccount, + } + + podSpecPb, err := utils.MarshalObjToStruct(podSpec) + if err != nil { + panic(err) + } + + taskTemplate := &core.TaskTemplate{ + Type: "test", + Target: &core.TaskTemplate_K8SPod{ + K8SPod: &core.K8SPod{ + PodSpec: podSpecPb, + }, + }, + Config: map[string]string{ + "primary_container_name": "test-image", + }, + } + + return taskTemplate +} + func dummyContainerTaskMetadata(resources *v1.ResourceRequirements, extendedResources *core.ExtendedResources) pluginsCore.TaskExecutionMetadata { taskMetadata := &pluginsCoreMock.TaskExecutionMetadata{} taskMetadata.On("GetNamespace").Return("test-namespace") @@ -147,6 +183,31 @@ func TestContainerTaskExecutor_BuildResource(t *testing.T) { assert.Equal(t, "service-account", j.Spec.ServiceAccountName) } +func TestContainerTaskExecutor_BuildResource_PodTemplate(t *testing.T) { + command := []string{"command"} + args := []string{"{{.Input}}"} + taskTemplate := dummyContainerTaskTemplateWithPodSpec(command, args) + taskCtx := dummyContainerTaskContext(taskTemplate, containerResourceRequirements, nil) + + r, err := DefaultPodPlugin.BuildResource(context.TODO(), taskCtx) + assert.NoError(t, err) + assert.NotNil(t, r) + j, ok := r.(*v1.Pod) + assert.True(t, ok) + + assert.NotEmpty(t, j.Spec.Containers) + assert.Equal(t, containerResourceRequirements.Limits[v1.ResourceCPU], j.Spec.Containers[0].Resources.Limits[v1.ResourceCPU]) + + // TODO: Once configurable, test when setting storage is supported on the cluster vs not. + storageRes := j.Spec.Containers[0].Resources.Limits[v1.ResourceStorage] + assert.Equal(t, int64(0), (&storageRes).Value()) + + assert.Equal(t, command, j.Spec.Containers[0].Command) + assert.Equal(t, []string{"test-data-reference"}, j.Spec.Containers[0].Args) + + assert.Equal(t, podTemplateServiceAccount, j.Spec.ServiceAccountName) +} + func TestContainerTaskExecutor_BuildResource_ExtendedResources(t *testing.T) { assert.NoError(t, flytek8sConfig.SetK8sPluginConfig(&flytek8sConfig.K8sPluginConfig{ GpuDeviceNodeLabel: "gpu-node-label", diff --git a/flyteplugins/go/tasks/plugins/k8s/pod/plugin.go b/flyteplugins/go/tasks/plugins/k8s/pod/plugin.go index 11de877021..d31c9a4b24 100644 --- a/flyteplugins/go/tasks/plugins/k8s/pod/plugin.go +++ b/flyteplugins/go/tasks/plugins/k8s/pod/plugin.go @@ -126,7 +126,9 @@ func (p plugin) BuildResource(ctx context.Context, taskCtx pluginsCore.TaskExecu objectMeta.Annotations[flytek8s.PrimaryContainerKey] = primaryContainerName } - podSpec.ServiceAccountName = flytek8s.GetServiceAccountNameFromTaskExecutionMetadata(taskCtx.TaskExecutionMetadata()) + if podSpec.ServiceAccountName == "" { + podSpec.ServiceAccountName = flytek8s.GetServiceAccountNameFromTaskExecutionMetadata(taskCtx.TaskExecutionMetadata()) + } pod := flytek8s.BuildIdentityPod() pod.ObjectMeta = *objectMeta From 1167153d6d0496f0e5c7d02d50941c16d64f2cab Mon Sep 17 00:00:00 2001 From: Paul Dittamo Date: Tue, 5 Dec 2023 20:40:42 -0800 Subject: [PATCH 2/5] update unit test Signed-off-by: Paul Dittamo --- .../tasks/plugins/k8s/pod/container_test.go | 110 ++++++++++-------- 1 file changed, 60 insertions(+), 50 deletions(-) diff --git a/flyteplugins/go/tasks/plugins/k8s/pod/container_test.go b/flyteplugins/go/tasks/plugins/k8s/pod/container_test.go index a01caca2ef..2fa7bffbdb 100644 --- a/flyteplugins/go/tasks/plugins/k8s/pod/container_test.go +++ b/flyteplugins/go/tasks/plugins/k8s/pod/container_test.go @@ -30,6 +30,7 @@ var containerResourceRequirements = &v1.ResourceRequirements{ } var podTemplateServiceAccount = "test-service-account" +var securityContextServiceAccount = "security-context-service-account" func dummyContainerTaskTemplate(command []string, args []string) *core.TaskTemplate { return &core.TaskTemplate{ @@ -76,7 +77,7 @@ func dummyContainerTaskTemplateWithPodSpec(command []string, args []string) *cor return taskTemplate } -func dummyContainerTaskMetadata(resources *v1.ResourceRequirements, extendedResources *core.ExtendedResources) pluginsCore.TaskExecutionMetadata { +func dummyContainerTaskMetadata(resources *v1.ResourceRequirements, extendedResources *core.ExtendedResources, returnsServiceAccount bool) pluginsCore.TaskExecutionMetadata { taskMetadata := &pluginsCoreMock.TaskExecutionMetadata{} taskMetadata.On("GetNamespace").Return("test-namespace") taskMetadata.On("GetAnnotations").Return(map[string]string{"annotation-1": "val1"}) @@ -85,9 +86,13 @@ func dummyContainerTaskMetadata(resources *v1.ResourceRequirements, extendedReso Kind: "node", Name: "blah", }) - taskMetadata.On("GetK8sServiceAccount").Return("service-account") + if returnsServiceAccount { + taskMetadata.On("GetK8sServiceAccount").Return("service-account") + } else { + taskMetadata.On("GetK8sServiceAccount").Return("") + } taskMetadata.On("GetSecurityContext").Return(core.SecurityContext{ - RunAs: &core.Identity{K8SServiceAccount: "service-account"}, + RunAs: &core.Identity{K8SServiceAccount: securityContextServiceAccount}, }) taskMetadata.On("GetOwnerID").Return(types.NamespacedName{ Namespace: "test-namespace", @@ -117,8 +122,7 @@ func dummyContainerTaskMetadata(resources *v1.ResourceRequirements, extendedReso return taskMetadata } -func dummyContainerTaskContext(taskTemplate *core.TaskTemplate, resources *v1.ResourceRequirements, extendedResources *core.ExtendedResources) pluginsCore.TaskExecutionContext { - dummyTaskMetadata := dummyContainerTaskMetadata(resources, extendedResources) +func dummyContainerTaskContext(taskTemplate *core.TaskTemplate, taskMetadata pluginsCore.TaskExecutionMetadata) pluginsCore.TaskExecutionContext { taskCtx := &pluginsCoreMock.TaskExecutionContext{} inputReader := &pluginsIOMock.InputReader{} inputReader.OnGetInputPrefixPath().Return("test-data-reference") @@ -139,7 +143,7 @@ func dummyContainerTaskContext(taskTemplate *core.TaskTemplate, resources *v1.Re taskReader.OnReadMatch(mock.Anything).Return(taskTemplate, nil) taskCtx.OnTaskReader().Return(taskReader) - taskCtx.OnTaskExecutionMetadata().Return(dummyTaskMetadata) + taskCtx.OnTaskExecutionMetadata().Return(taskMetadata) pluginStateReader := &pluginsCoreMock.PluginStateReader{} pluginStateReader.OnGetMatch(mock.Anything).Return(0, nil) @@ -158,54 +162,57 @@ func TestContainerTaskExecutor_BuildIdentityResource(t *testing.T) { assert.Equal(t, flytek8s.PodKind, r.GetObjectKind().GroupVersionKind().Kind) } -func TestContainerTaskExecutor_BuildResource(t *testing.T) { +func TestContainerTaskExecutor_BuildResource_U(t *testing.T) { command := []string{"command"} args := []string{"{{.Input}}"} - taskTemplate := dummyContainerTaskTemplate(command, args) - taskCtx := dummyContainerTaskContext(taskTemplate, containerResourceRequirements, nil) - - r, err := DefaultPodPlugin.BuildResource(context.TODO(), taskCtx) - assert.NoError(t, err) - assert.NotNil(t, r) - j, ok := r.(*v1.Pod) - assert.True(t, ok) - - assert.NotEmpty(t, j.Spec.Containers) - assert.Equal(t, containerResourceRequirements.Limits[v1.ResourceCPU], j.Spec.Containers[0].Resources.Limits[v1.ResourceCPU]) - - // TODO: Once configurable, test when setting storage is supported on the cluster vs not. - storageRes := j.Spec.Containers[0].Resources.Limits[v1.ResourceStorage] - assert.Equal(t, int64(0), (&storageRes).Value()) - - assert.Equal(t, command, j.Spec.Containers[0].Command) - assert.Equal(t, []string{"test-data-reference"}, j.Spec.Containers[0].Args) - - assert.Equal(t, "service-account", j.Spec.ServiceAccountName) -} - -func TestContainerTaskExecutor_BuildResource_PodTemplate(t *testing.T) { - command := []string{"command"} - args := []string{"{{.Input}}"} - taskTemplate := dummyContainerTaskTemplateWithPodSpec(command, args) - taskCtx := dummyContainerTaskContext(taskTemplate, containerResourceRequirements, nil) + testCases := []struct { + name string + taskTemplate *core.TaskTemplate + taskMetadata pluginsCore.TaskExecutionMetadata + expectServiceAccount string + }{ + { + name: "BuildResource", + taskTemplate: dummyContainerTaskTemplate(command, args), + taskMetadata: dummyContainerTaskMetadata(containerResourceRequirements, nil, true), + expectServiceAccount: "service-account", + }, + { + name: "BuildResource_PodTemplate", + taskTemplate: dummyContainerTaskTemplateWithPodSpec(command, args), + taskMetadata: dummyContainerTaskMetadata(containerResourceRequirements, nil, true), + expectServiceAccount: podTemplateServiceAccount, + }, + { + name: "BuildResource_SecurityContext", + taskTemplate: dummyContainerTaskTemplate(command, args), + taskMetadata: dummyContainerTaskMetadata(containerResourceRequirements, nil, false), + expectServiceAccount: securityContextServiceAccount, + }, + } + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + taskCtx := dummyContainerTaskContext(tc.taskTemplate, tc.taskMetadata) - r, err := DefaultPodPlugin.BuildResource(context.TODO(), taskCtx) - assert.NoError(t, err) - assert.NotNil(t, r) - j, ok := r.(*v1.Pod) - assert.True(t, ok) + r, err := DefaultPodPlugin.BuildResource(context.TODO(), taskCtx) + assert.NoError(t, err) + assert.NotNil(t, r) + j, ok := r.(*v1.Pod) + assert.True(t, ok) - assert.NotEmpty(t, j.Spec.Containers) - assert.Equal(t, containerResourceRequirements.Limits[v1.ResourceCPU], j.Spec.Containers[0].Resources.Limits[v1.ResourceCPU]) + assert.NotEmpty(t, j.Spec.Containers) + assert.Equal(t, containerResourceRequirements.Limits[v1.ResourceCPU], j.Spec.Containers[0].Resources.Limits[v1.ResourceCPU]) - // TODO: Once configurable, test when setting storage is supported on the cluster vs not. - storageRes := j.Spec.Containers[0].Resources.Limits[v1.ResourceStorage] - assert.Equal(t, int64(0), (&storageRes).Value()) + // TODO: Once configurable, test when setting storage is supported on the cluster vs not. + storageRes := j.Spec.Containers[0].Resources.Limits[v1.ResourceStorage] + assert.Equal(t, int64(0), (&storageRes).Value()) - assert.Equal(t, command, j.Spec.Containers[0].Command) - assert.Equal(t, []string{"test-data-reference"}, j.Spec.Containers[0].Args) + assert.Equal(t, command, j.Spec.Containers[0].Command) + assert.Equal(t, []string{"test-data-reference"}, j.Spec.Containers[0].Args) - assert.Equal(t, podTemplateServiceAccount, j.Spec.ServiceAccountName) + assert.Equal(t, tc.expectServiceAccount, j.Spec.ServiceAccountName) + }) + } } func TestContainerTaskExecutor_BuildResource_ExtendedResources(t *testing.T) { @@ -313,7 +320,8 @@ func TestContainerTaskExecutor_BuildResource_ExtendedResources(t *testing.T) { t.Run(f.name, func(t *testing.T) { taskTemplate := dummyContainerTaskTemplate([]string{"command"}, []string{"{{.Input}}"}) taskTemplate.ExtendedResources = f.extendedResourcesBase - taskContext := dummyContainerTaskContext(taskTemplate, f.resources, f.extendedResourcesOverride) + taskMetadata := dummyContainerTaskMetadata(f.resources, f.extendedResourcesOverride, true) + taskContext := dummyContainerTaskContext(taskTemplate, taskMetadata) r, err := DefaultPodPlugin.BuildResource(context.TODO(), taskContext) assert.Nil(t, err) assert.NotNil(t, r) @@ -338,7 +346,8 @@ func TestContainerTaskExecutor_GetTaskStatus(t *testing.T) { command := []string{"command"} args := []string{"{{.Input}}"} taskTemplate := dummyContainerTaskTemplate(command, args) - taskCtx := dummyContainerTaskContext(taskTemplate, containerResourceRequirements, nil) + taskMetadata := dummyContainerTaskMetadata(containerResourceRequirements, nil, true) + taskCtx := dummyContainerTaskContext(taskTemplate, taskMetadata) j := &v1.Pod{ Status: v1.PodStatus{}, @@ -427,7 +436,8 @@ func TestContainerTaskExecutor_GetTaskStatus_InvalidImageName(t *testing.T) { command := []string{"command"} args := []string{"{{.Input}}"} taskTemplate := dummyContainerTaskTemplate(command, args) - taskCtx := dummyContainerTaskContext(taskTemplate, containerResourceRequirements, nil) + taskMetadata := dummyContainerTaskMetadata(containerResourceRequirements, nil, true) + taskCtx := dummyContainerTaskContext(taskTemplate, taskMetadata) ctx := context.TODO() reason := "InvalidImageName" From ab38bf7dc00e5e7d8823170bce48511761f56a05 Mon Sep 17 00:00:00 2001 From: Paul Dittamo Date: Tue, 5 Dec 2023 20:49:42 -0800 Subject: [PATCH 3/5] cleanup Signed-off-by: Paul Dittamo --- flyteplugins/go/tasks/plugins/k8s/pod/container_test.go | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/flyteplugins/go/tasks/plugins/k8s/pod/container_test.go b/flyteplugins/go/tasks/plugins/k8s/pod/container_test.go index 2fa7bffbdb..8abef87ecf 100644 --- a/flyteplugins/go/tasks/plugins/k8s/pod/container_test.go +++ b/flyteplugins/go/tasks/plugins/k8s/pod/container_test.go @@ -29,6 +29,7 @@ var containerResourceRequirements = &v1.ResourceRequirements{ }, } +var serviceAccount = "service-account" var podTemplateServiceAccount = "test-service-account" var securityContextServiceAccount = "security-context-service-account" @@ -87,7 +88,7 @@ func dummyContainerTaskMetadata(resources *v1.ResourceRequirements, extendedReso Name: "blah", }) if returnsServiceAccount { - taskMetadata.On("GetK8sServiceAccount").Return("service-account") + taskMetadata.On("GetK8sServiceAccount").Return(serviceAccount) } else { taskMetadata.On("GetK8sServiceAccount").Return("") } @@ -175,7 +176,7 @@ func TestContainerTaskExecutor_BuildResource_U(t *testing.T) { name: "BuildResource", taskTemplate: dummyContainerTaskTemplate(command, args), taskMetadata: dummyContainerTaskMetadata(containerResourceRequirements, nil, true), - expectServiceAccount: "service-account", + expectServiceAccount: serviceAccount, }, { name: "BuildResource_PodTemplate", From 8ce1a8bdce19b35f609d61e06410b1ee3f43f298 Mon Sep 17 00:00:00 2001 From: Paul Dittamo Date: Tue, 5 Dec 2023 20:56:47 -0800 Subject: [PATCH 4/5] typo Signed-off-by: Paul Dittamo --- flyteplugins/go/tasks/plugins/k8s/pod/container_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/flyteplugins/go/tasks/plugins/k8s/pod/container_test.go b/flyteplugins/go/tasks/plugins/k8s/pod/container_test.go index 8abef87ecf..06970cf3e2 100644 --- a/flyteplugins/go/tasks/plugins/k8s/pod/container_test.go +++ b/flyteplugins/go/tasks/plugins/k8s/pod/container_test.go @@ -163,7 +163,7 @@ func TestContainerTaskExecutor_BuildIdentityResource(t *testing.T) { assert.Equal(t, flytek8s.PodKind, r.GetObjectKind().GroupVersionKind().Kind) } -func TestContainerTaskExecutor_BuildResource_U(t *testing.T) { +func TestContainerTaskExecutor_BuildResource(t *testing.T) { command := []string{"command"} args := []string{"{{.Input}}"} testCases := []struct { From 2817999ec495c19bcdad4a907060203239d528ec Mon Sep 17 00:00:00 2001 From: Paul Dittamo Date: Wed, 6 Dec 2023 09:52:33 -0800 Subject: [PATCH 5/5] clean up sytling Signed-off-by: Paul Dittamo --- flyteplugins/go/tasks/plugins/k8s/pod/container_test.go | 8 +++++--- flyteplugins/go/tasks/plugins/k8s/pod/plugin.go | 2 +- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/flyteplugins/go/tasks/plugins/k8s/pod/container_test.go b/flyteplugins/go/tasks/plugins/k8s/pod/container_test.go index 06970cf3e2..9a70f906b9 100644 --- a/flyteplugins/go/tasks/plugins/k8s/pod/container_test.go +++ b/flyteplugins/go/tasks/plugins/k8s/pod/container_test.go @@ -29,9 +29,11 @@ var containerResourceRequirements = &v1.ResourceRequirements{ }, } -var serviceAccount = "service-account" -var podTemplateServiceAccount = "test-service-account" -var securityContextServiceAccount = "security-context-service-account" +var ( + serviceAccount = "service-account" + podTemplateServiceAccount = "test-service-account" + securityContextServiceAccount = "security-context-service-account" +) func dummyContainerTaskTemplate(command []string, args []string) *core.TaskTemplate { return &core.TaskTemplate{ diff --git a/flyteplugins/go/tasks/plugins/k8s/pod/plugin.go b/flyteplugins/go/tasks/plugins/k8s/pod/plugin.go index d31c9a4b24..b266a6f5e8 100644 --- a/flyteplugins/go/tasks/plugins/k8s/pod/plugin.go +++ b/flyteplugins/go/tasks/plugins/k8s/pod/plugin.go @@ -126,7 +126,7 @@ func (p plugin) BuildResource(ctx context.Context, taskCtx pluginsCore.TaskExecu objectMeta.Annotations[flytek8s.PrimaryContainerKey] = primaryContainerName } - if podSpec.ServiceAccountName == "" { + if len(podSpec.ServiceAccountName) == 0 { podSpec.ServiceAccountName = flytek8s.GetServiceAccountNameFromTaskExecutionMetadata(taskCtx.TaskExecutionMetadata()) }