From 8b714698955f3bf90191f239de82d79ca1225d67 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Tue, 25 Jun 2024 00:43:18 +0800 Subject: [PATCH] fix(ray): Use default svc account if not set in task metadata (#5499) Signed-off-by: Kevin Su --- flyteplugins/go/tasks/plugins/k8s/ray/ray.go | 3 +++ .../go/tasks/plugins/k8s/ray/ray_test.go | 22 ++++++++++++++----- 2 files changed, 19 insertions(+), 6 deletions(-) diff --git a/flyteplugins/go/tasks/plugins/k8s/ray/ray.go b/flyteplugins/go/tasks/plugins/k8s/ray/ray.go index 7774c50376..90388b46a5 100644 --- a/flyteplugins/go/tasks/plugins/k8s/ray/ray.go +++ b/flyteplugins/go/tasks/plugins/k8s/ray/ray.go @@ -190,6 +190,9 @@ func constructRayJob(taskCtx pluginsCore.TaskExecutionContext, rayJob plugins.Ra } serviceAccountName := flytek8s.GetServiceAccountNameFromTaskExecutionMetadata(taskCtx.TaskExecutionMetadata()) + if len(serviceAccountName) == 0 { + serviceAccountName = cfg.ServiceAccount + } rayClusterSpec.HeadGroupSpec.Template.Spec.ServiceAccountName = serviceAccountName for index := range rayClusterSpec.WorkerGroupSpecs { diff --git a/flyteplugins/go/tasks/plugins/k8s/ray/ray_test.go b/flyteplugins/go/tasks/plugins/k8s/ray/ray_test.go index 65ccfac643..7b555e9f23 100644 --- a/flyteplugins/go/tasks/plugins/k8s/ray/ray_test.go +++ b/flyteplugins/go/tasks/plugins/k8s/ray/ray_test.go @@ -106,7 +106,7 @@ func dummyRayTaskTemplate(id string, rayJob *plugins.RayJob) *core.TaskTemplate } } -func dummyRayTaskContext(taskTemplate *core.TaskTemplate, resources *corev1.ResourceRequirements, extendedResources *core.ExtendedResources, containerImage string) pluginsCore.TaskExecutionContext { +func dummyRayTaskContext(taskTemplate *core.TaskTemplate, resources *corev1.ResourceRequirements, extendedResources *core.ExtendedResources, containerImage, serviceAccount string) pluginsCore.TaskExecutionContext { taskCtx := &mocks.TaskExecutionContext{} inputReader := &pluginIOMocks.InputReader{} inputReader.OnGetInputPrefixPath().Return("/input/prefix") @@ -177,7 +177,8 @@ func TestBuildResourceRay(t *testing.T) { err := config.SetK8sPluginConfig(&config.K8sPluginConfig{DefaultTolerations: toleration}) assert.Nil(t, err) - RayResource, err := rayJobResourceHandler.BuildResource(context.TODO(), dummyRayTaskContext(taskTemplate, resourceRequirements, nil, "")) + rayCtx := dummyRayTaskContext(taskTemplate, resourceRequirements, nil, "", serviceAccount) + RayResource, err := rayJobResourceHandler.BuildResource(context.TODO(), rayCtx) assert.Nil(t, err) assert.NotNil(t, RayResource) @@ -208,6 +209,15 @@ func TestBuildResourceRay(t *testing.T) { assert.Equal(t, ray.Spec.RayClusterSpec.WorkerGroupSpecs[0].Template.Annotations, map[string]string{"annotation-1": "val1"}) assert.Equal(t, ray.Spec.RayClusterSpec.WorkerGroupSpecs[0].Template.Labels, map[string]string{"label-1": "val1"}) assert.Equal(t, ray.Spec.RayClusterSpec.WorkerGroupSpecs[0].Template.Spec.Tolerations, toleration) + + // Make sure the default service account is being used if SA is not provided in the task context + rayCtx = dummyRayTaskContext(taskTemplate, resourceRequirements, nil, "", "") + RayResource, err = rayJobResourceHandler.BuildResource(context.TODO(), rayCtx) + assert.Nil(t, err) + assert.NotNil(t, RayResource) + ray, ok = RayResource.(*rayv1.RayJob) + assert.True(t, ok) + assert.Equal(t, ray.Spec.RayClusterSpec.HeadGroupSpec.Template.Spec.ServiceAccountName, GetConfig().ServiceAccount) } func TestBuildResourceRayContainerImage(t *testing.T) { @@ -241,7 +251,7 @@ func TestBuildResourceRayContainerImage(t *testing.T) { for _, f := range fixtures { t.Run(f.name, func(t *testing.T) { taskTemplate := dummyRayTaskTemplate("id", dummyRayCustomObj()) - taskContext := dummyRayTaskContext(taskTemplate, f.resources, nil, f.containerImageOverride) + taskContext := dummyRayTaskContext(taskTemplate, f.resources, nil, f.containerImageOverride, serviceAccount) rayJobResourceHandler := rayJobResourceHandler{} r, err := rayJobResourceHandler.BuildResource(context.TODO(), taskContext) assert.Nil(t, err) @@ -372,7 +382,7 @@ func TestBuildResourceRayExtendedResources(t *testing.T) { t.Run(p.name, func(t *testing.T) { taskTemplate := dummyRayTaskTemplate("ray-id", dummyRayCustomObj()) taskTemplate.ExtendedResources = p.extendedResourcesBase - taskContext := dummyRayTaskContext(taskTemplate, p.resources, p.extendedResourcesOverride, "") + taskContext := dummyRayTaskContext(taskTemplate, p.resources, p.extendedResourcesOverride, "", serviceAccount) rayJobResourceHandler := rayJobResourceHandler{} r, err := rayJobResourceHandler.BuildResource(context.TODO(), taskContext) assert.Nil(t, err) @@ -431,7 +441,7 @@ func TestDefaultStartParameters(t *testing.T) { err := config.SetK8sPluginConfig(&config.K8sPluginConfig{DefaultTolerations: toleration}) assert.Nil(t, err) - RayResource, err := rayJobResourceHandler.BuildResource(context.TODO(), dummyRayTaskContext(taskTemplate, resourceRequirements, nil, "")) + RayResource, err := rayJobResourceHandler.BuildResource(context.TODO(), dummyRayTaskContext(taskTemplate, resourceRequirements, nil, "", serviceAccount)) assert.Nil(t, err) assert.NotNil(t, RayResource) @@ -637,7 +647,7 @@ func TestInjectLogsSidecar(t *testing.T) { assert.NoError(t, SetConfig(&Config{ LogsSidecar: p.logsSidecarCfg, })) - taskContext := dummyRayTaskContext(&p.taskTemplate, resourceRequirements, nil, "") + taskContext := dummyRayTaskContext(&p.taskTemplate, resourceRequirements, nil, "", serviceAccount) rayJobResourceHandler := rayJobResourceHandler{} r, err := rayJobResourceHandler.BuildResource(context.TODO(), taskContext) assert.Nil(t, err)