Inject and Use values from Security Context

go.sum

@@ -72,7 +72,6 @@ github.com/Azure/go-autorest/logger v0.2.0/go.mod h1:T9E3cAhj2VqvPOtCYAvby9aBXkZ
 github.com/Azure/go-autorest/tracing v0.5.0/go.mod h1:r/s2XiOKccPW3HrqB+W0TQzfbtp2fGCgRFtBroKn4Dk=
 github.com/Azure/go-autorest/tracing v0.6.0 h1:TYi4+3m5t6K48TGI9AUdb+IzbnSxvnvUMfuitfgcfuo=
 github.com/Azure/go-autorest/tracing v0.6.0/go.mod h1:+vhtPC754Xsa23ID7GlGsrdKBpUA79WCAKPPZVC2DeU=
-github.com/BurntSushi/toml v0.3.1 h1:WXkYYl6Yr3qBf1K79EBnL4mak0OimBfB0XUf9Vl28OQ=
 github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
 github.com/BurntSushi/xgb v0.0.0-20160522181843-27f122750802/go.mod h1:IVnqGOEym/WlBOVXweHU+Q+/VP0lqqI8lqeDx9IjBqo=
 github.com/GoogleCloudPlatform/spark-on-k8s-operator v0.0.0-20200723154620-6f35a1152625 h1:cQyO5JQ2iuHnEcF3v24kdDMsgh04RjyFPDtuvD6PCE0=
@@ -502,8 +501,6 @@ github.com/kubeflow/tf-operator v0.5.3/go.mod h1:EBtz5LQoKaHUl/5fV5vD1qXVNVNyn3T
 github.com/lightstep/lightstep-tracer-common/golang/gogo v0.0.0-20190605223551-bc2310a04743/go.mod h1:qklhhLq1aX+mtWk9cPHPzaBjWImj5ULL6C7HFJtXQMM=
 github.com/lightstep/lightstep-tracer-go v0.18.1/go.mod h1:jlF1pusYV4pidLvZ+XD0UBX0ZE6WURAspgAczcDHrL4=
 github.com/lyft/protoc-gen-validate v0.0.13/go.mod h1:XbGvPuh87YZc5TdIa2/I4pLk0QoUACkjt2znoq26NVQ=
-github.com/lyft/spark-on-k8s-operator v0.1.4-0.20201027003055-c76b67e3b6d0 h1:1vSmc+Bo70X0JVYywQ9Hy/aet6p613ejacy9x5td0m4=
-github.com/lyft/spark-on-k8s-operator v0.1.4-0.20201027003055-c76b67e3b6d0/go.mod h1:hkRqdqAsdNnxT/Zst6MNMRbTAoiCZ0JRw7svRgAYb0A=
 github.com/magiconair/properties v1.8.0/go.mod h1:PppfXfuXeibc/6YijjN8zIbojt8czPbwD3XqdrwzmxQ=
 github.com/magiconair/properties v1.8.1/go.mod h1:PppfXfuXeibc/6YijjN8zIbojt8czPbwD3XqdrwzmxQ=
 github.com/magiconair/properties v1.8.4 h1:8KGKTcQQGm0Kv7vEbKFErAoAOFyyacLStRtQSeYtvkY=
@@ -756,7 +753,6 @@ go.uber.org/atomic v1.5.0/go.mod h1:sABNBOSYdrvTF6hTgEIbc7YasKWGhgEQZyfxyTvoXHQ=
 go.uber.org/atomic v1.6.0/go.mod h1:sABNBOSYdrvTF6hTgEIbc7YasKWGhgEQZyfxyTvoXHQ=
 go.uber.org/atomic v1.7.0 h1:ADUqmZGgLDDfbSL9ZmPxKTybcoEYHgpYfELNoN+7hsw=
 go.uber.org/atomic v1.7.0/go.mod h1:fEN4uk6kAWBTFdckzkM89CLk9XfWZrxpCo0nPH17wJc=
-go.uber.org/goleak v1.1.10 h1:z+mqJhf6ss6BSfSM671tgKyZBFPTTJM+HLxnhPC3wu0=
 go.uber.org/goleak v1.1.10/go.mod h1:8a7PlsEVH3e/a/GLqe5IIrQx6GzcnRmZEufDUTk4A7A=
 go.uber.org/multierr v1.1.0/go.mod h1:wR5kodmAFQ0UK8QlbwjlSNy0Z68gJhDJUG5sjR94q/0=
 go.uber.org/multierr v1.3.0/go.mod h1:VgVr7evmIr6uPjLBxg28wmKNXyqE9akIJ5XnfpiKl+4=
@@ -1218,7 +1214,6 @@ honnef.co/go/tools v0.0.0-20190418001031-e561f6794a2a/go.mod h1:rf3lG4BRIbNafJWh
 honnef.co/go/tools v0.0.0-20190523083050-ea95bdfd59fc/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4=
 honnef.co/go/tools v0.0.1-2019.2.3/go.mod h1:a3bituU0lyd329TUQxRnasdCoJDkEUEAqEt0JzvZhAg=
 honnef.co/go/tools v0.0.1-2020.1.3/go.mod h1:X/FiERA/W4tHapMX5mGpAtMSVEeEUOyHaw9vFzvIQ3k=
-honnef.co/go/tools v0.0.1-2020.1.4 h1:UoveltGrhghAA7ePc+e+QYDHXrBps2PqFZiHkGR/xK8=
 honnef.co/go/tools v0.0.1-2020.1.4/go.mod h1:X/FiERA/W4tHapMX5mGpAtMSVEeEUOyHaw9vFzvIQ3k=
 k8s.io/api v0.0.0-20210217171935-8e2decd92398/go.mod h1:60tmSUpHxGPFerNHbo/ayI2lKxvtrhbxFyXuEIWJd78=
 k8s.io/api v0.18.2/go.mod h1:SJCWI7OLzhZSvbY7U8zwNl9UA4o1fizoug34OV/2r78=

go/tasks/pluginmachinery/core/exec_metadata.go

@@ -33,5 +33,6 @@ type TaskExecutionMetadata interface {
 	GetMaxAttempts() uint32
 	GetAnnotations() map[string]string
 	GetK8sServiceAccount() string
+	GetSecurityContext() core.SecurityContext
 	IsInterruptible() bool
 }

go/tasks/pluginmachinery/flytek8s/utils.go

@@ -2,6 +2,7 @@ package flytek8s
 
 import (
 	"github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/core"
+	pluginmachinery_core "github.com/flyteorg/flyteplugins/go/tasks/pluginmachinery/core"
 	v1 "k8s.io/api/core/v1"
 )
 
@@ -12,3 +13,18 @@ func ToK8sEnvVar(env []*core.KeyValuePair) []v1.EnvVar {
 	}
 	return envVars
 }
+
+func GetServiceAccountNameFromTaskExecutionMetadata(taskExecutionMetadata pluginmachinery_core.TaskExecutionMetadata) string {
+	var serviceAccount string
+	securityContext := taskExecutionMetadata.GetSecurityContext()
+	if securityContext.GetRunAs() != nil {
+		serviceAccount = securityContext.GetRunAs().GetK8SServiceAccount()
+	}
+
+	// TO BE DEPRECATED
+	if len(serviceAccount) == 0 {
+		serviceAccount = taskExecutionMetadata.GetK8sServiceAccount()
+	}
+
+	return serviceAccount
+}

go/tasks/pluginmachinery/flytek8s/utils_test.go

@@ -1 +1,27 @@
 package flytek8s
+
+import (
+	"testing"
+
+	"github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/core"
+	"github.com/flyteorg/flyteplugins/go/tasks/pluginmachinery/core/mocks"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestGetServiceAccountNameFromTaskExecutionMetadata(t *testing.T) {
+	mockTaskExecMetadata := mocks.TaskExecutionMetadata{}
+	mockTaskExecMetadata.OnGetSecurityContext().Return(core.SecurityContext{
+		RunAs: &core.Identity{K8SServiceAccount: "service-account"},
+	})
+	result := GetServiceAccountNameFromTaskExecutionMetadata(&mockTaskExecMetadata)
+	assert.Equal(t, "service-account", result)
+}
+
+func TestGetServiceAccountNameFromServiceAccount(t *testing.T) {
+	mockTaskExecMetadata := mocks.TaskExecutionMetadata{}
+	mockTaskExecMetadata.OnGetSecurityContext().Return(core.SecurityContext{})
+	mockTaskExecMetadata.OnGetK8sServiceAccount().Return("service-account")
+	result := GetServiceAccountNameFromTaskExecutionMetadata(&mockTaskExecMetadata)
+	assert.Equal(t, "service-account", result)
+}

go/tasks/plugins/array/awsbatch/config/config.go

@@ -22,9 +22,10 @@ type Config struct {
 	// Provide additional environment variable pairs that plugin authors will provide to containers
 	DefaultEnvVars       map[string]string `json:"defaultEnvVars" pflag:"-,Additional environment variable that should be injected into every resource"`
 	MaxErrorStringLength int               `json:"maxErrLength" pflag:",Determines the maximum length of the error string returned for the array."`
-	RoleAnnotationKey    string            `json:"roleAnnotationKey" pflag:",Map key to use to lookup role from task annotations."`
-	OutputAssembler      workqueue.Config  `json:"outputAssembler"`
-	ErrorAssembler       workqueue.Config  `json:"errorAssembler"`
+	// This can be deprecated. Just having it for backward compatibility
+	RoleAnnotationKey string           `json:"roleAnnotationKey" pflag:",Map key to use to lookup role from task annotations."`
+	OutputAssembler   workqueue.Config `json:"outputAssembler"`
+	ErrorAssembler    workqueue.Config `json:"errorAssembler"`
 }
 
 type JobStoreConfig struct {

go/tasks/plugins/array/awsbatch/job_definition.go

@@ -50,7 +50,7 @@ func EnsureJobDefinition(ctx context.Context, tCtx pluginCore.TaskExecutionConte
 		return nil, errors.Errorf(pluginErrors.BadTaskSpecification, "Tasktemplate does not contain a container image.")
 	}
 
-	role := awsUtils.GetRole(ctx, cfg.RoleAnnotationKey, tCtx.TaskExecutionMetadata().GetAnnotations())
+	role := awsUtils.GetRoleFromSecurityContext(cfg.RoleAnnotationKey, tCtx.TaskExecutionMetadata())
 
 	cacheKey := definition.NewCacheKey(role, containerImage)
 	if existingArn, found := definitionCache.Get(cacheKey); found {

go/tasks/plugins/array/awsbatch/job_definition_test.go

@@ -64,7 +64,7 @@ func TestEnsureJobDefinition(t *testing.T) {
 	tMeta.OnGetTaskExecutionID().Return(tID)
 	tMeta.OnGetOverrides().Return(overrides)
 	tMeta.OnGetAnnotations().Return(map[string]string{})
-
+	tMeta.OnGetSecurityContext().Return(core.SecurityContext{})
 	tCtx := &mocks.TaskExecutionContext{}
 	tCtx.OnTaskReader().Return(tReader)
 	tCtx.OnTaskExecutionMetadata().Return(tMeta)
@@ -101,3 +101,70 @@ func TestEnsureJobDefinition(t *testing.T) {
 		assert.Equal(t, "their-arn", nextState.JobDefinitionArn)
 	})
 }
+
+func TestEnsureJobDefinitionWithSecurityContext(t *testing.T) {
+	ctx := context.Background()
+
+	tReader := &mocks.TaskReader{}
+	tReader.OnReadMatch(mock.Anything).Return(&core.TaskTemplate{
+		Interface: &core.TypedInterface{
+			Outputs: &core.VariableMap{
+				Variables: map[string]*core.Variable{"var1": {Type: &core.LiteralType{Type: &core.LiteralType_Simple{Simple: core.SimpleType_INTEGER}}}},
+			},
+		},
+		Target: &core.TaskTemplate_Container{
+			Container: createSampleContainerTask(),
+		},
+	}, nil)
+
+	overrides := &mocks.TaskOverrides{}
+	overrides.OnGetConfig().Return(&v1.ConfigMap{Data: map[string]string{
+		DynamicTaskQueueKey: "queue1",
+	}})
+
+	tID := &mocks.TaskExecutionID{}
+	tID.OnGetGeneratedName().Return("found")
+
+	tMeta := &mocks.TaskExecutionMetadata{}
+	tMeta.OnGetTaskExecutionID().Return(tID)
+	tMeta.OnGetOverrides().Return(overrides)
+	tMeta.OnGetAnnotations().Return(map[string]string{})
+	tMeta.OnGetSecurityContext().Return(core.SecurityContext{
+		RunAs: &core.Identity{IamRole: "new-role"},
+	})
+	tCtx := &mocks.TaskExecutionContext{}
+	tCtx.OnTaskReader().Return(tReader)
+	tCtx.OnTaskExecutionMetadata().Return(tMeta)
+
+	cfg := &config.Config{}
+	batchClient := NewCustomBatchClient(batchMocks.NewMockAwsBatchClient(), "", "",
+		utils.NewRateLimiter("", 10, 20),
+		utils.NewRateLimiter("", 10, 20))
+
+	t.Run("Not Found", func(t *testing.T) {
+		dCache := definition.NewCache(10)
+
+		nextState, err := EnsureJobDefinition(ctx, tCtx, cfg, batchClient, dCache, &State{
+			State: &arrayCore.State{},
+		})
+
+		assert.NoError(t, err)
+		assert.NotNil(t, nextState)
+		assert.Equal(t, "my-arn", nextState.JobDefinitionArn)
+		p, v := nextState.GetPhase()
+		assert.Equal(t, arrayCore.PhaseLaunch, p)
+		assert.Zero(t, v)
+	})
+
+	t.Run("Found", func(t *testing.T) {
+		dCache := definition.NewCache(10)
+		assert.NoError(t, dCache.Put(definition.NewCacheKey("new-role", "img1"), "their-arn"))
+
+		nextState, err := EnsureJobDefinition(ctx, tCtx, cfg, batchClient, dCache, &State{
+			State: &arrayCore.State{},
+		})
+		assert.NoError(t, err)
+		assert.NotNil(t, nextState)
+		assert.Equal(t, "their-arn", nextState.JobDefinitionArn)
+	})
+}

go/tasks/plugins/awsutils/awsutils.go

@@ -1,10 +1,26 @@
 package awsutils
 
-import "context"
+import (
+	core2 "github.com/flyteorg/flyteplugins/go/tasks/pluginmachinery/core"
+)
 
-func GetRole(_ context.Context, roleAnnotationKey string, annotations map[string]string) string {
-	if len(roleAnnotationKey) > 0 {
-		return annotations[roleAnnotationKey]
+func GetRoleFromSecurityContext(roleKey string, taskExecutionMetadata core2.TaskExecutionMetadata) string {
+	var role string
+	securityContext := taskExecutionMetadata.GetSecurityContext()
+	if securityContext.GetRunAs() != nil {
+		role = securityContext.GetRunAs().GetIamRole()
+	}
+
+	// Continue this for backward compatibility
+	if len(role) == 0 {
+		role = getRole(roleKey, taskExecutionMetadata.GetAnnotations())
+	}
+	return role
+}
+
+func getRole(roleKey string, keyValueMap map[string]string) string {
+	if len(roleKey) > 0 {
+		return keyValueMap[roleKey]
 	}
 
 	return ""

go/tasks/plugins/k8s/container/container.go

@@ -64,8 +64,7 @@ func (Plugin) BuildResource(ctx context.Context, taskCtx pluginsCore.TaskExecuti
 
 	pod := flytek8s.BuildPodWithSpec(podSpec)
 
-	// We want to Also update the serviceAccount to the serviceaccount of the workflow
-	pod.Spec.ServiceAccountName = taskCtx.TaskExecutionMetadata().GetK8sServiceAccount()
+	pod.Spec.ServiceAccountName = flytek8s.GetServiceAccountNameFromTaskExecutionMetadata(taskCtx.TaskExecutionMetadata())
 
 	return pod, nil
 }

go/tasks/plugins/k8s/container/container_test.go

@@ -39,6 +39,9 @@ func dummyContainerTaskMetadata(resources *v1.ResourceRequirements) pluginsCore.
 		Name: "blah",
 	})
 	taskMetadata.On("GetK8sServiceAccount").Return("service-account")
+	taskMetadata.On("GetSecurityContext").Return(core.SecurityContext{
+		RunAs: &core.Identity{K8SServiceAccount: "service-account"},
+	})
 	taskMetadata.On("GetOwnerID").Return(types.NamespacedName{
 		Namespace: "test-namespace",
 		Name:      "test-owner-name",

go/tasks/plugins/k8s/sagemaker/builtin_training.go

@@ -116,8 +116,9 @@ func (m awsSagemakerPlugin) buildResourceForTrainingJob(
 
 	inputModeString := strings.Title(strings.ToLower(sagemakerTrainingJob.GetAlgorithmSpecification().GetInputMode().String()))
 
-	role := awsUtils.GetRole(ctx, cfg.RoleAnnotationKey, taskCtx.TaskExecutionMetadata().GetAnnotations())
-	if role == "" {
+	role := awsUtils.GetRoleFromSecurityContext(cfg.RoleAnnotationKey, taskCtx.TaskExecutionMetadata())
+
+	if len(role) == 0 {
 		role = cfg.RoleArn
 	}
 

go/tasks/plugins/k8s/sagemaker/hyperparameter_tuning.go

@@ -137,8 +137,9 @@ func (m awsSagemakerPlugin) buildResourceForHyperparameterTuningJob(
 	tuningObjectiveTypeString := strings.Title(strings.ToLower(hpoJobConfig.GetTuningObjective().GetObjectiveType().String()))
 	trainingJobEarlyStoppingTypeString := strings.Title(strings.ToLower(hpoJobConfig.TrainingJobEarlyStoppingType.String()))
 
-	role := awsUtils.GetRole(ctx, cfg.RoleAnnotationKey, taskCtx.TaskExecutionMetadata().GetAnnotations())
-	if role == "" {
+	role := awsUtils.GetRoleFromSecurityContext(cfg.RoleAnnotationKey, taskCtx.TaskExecutionMetadata())
+
+	if len(role) == 0 {
 		role = cfg.RoleArn
 	}
 

go/tasks/plugins/k8s/sagemaker/plugin_test_utils.go

@@ -195,6 +195,10 @@ func generateMockCustomTrainingJobTaskContext(taskTemplate *flyteIdlCore.TaskTem
 	taskExecutionMetadata.OnGetTaskExecutionID().Return(tID)
 	taskExecutionMetadata.OnGetNamespace().Return("test-namespace")
 	taskExecutionMetadata.OnGetAnnotations().Return(map[string]string{"iam.amazonaws.com/role": "metadata_role"})
+	taskExecutionMetadata.OnGetSecurityContext().Return(flyteIdlCore.SecurityContext{
+		RunAs: &flyteIdlCore.Identity{IamRole: "new-role"},
+	})
+
 	taskExecutionMetadata.OnGetLabels().Return(map[string]string{"label-1": "val1"})
 	taskExecutionMetadata.OnGetOwnerReference().Return(v1.OwnerReference{
 		Kind: "node",
@@ -270,6 +274,7 @@ func generateMockTrainingJobTaskContext(taskTemplate *flyteIdlCore.TaskTemplate,
 	taskExecutionMetadata.OnGetTaskExecutionID().Return(tID)
 	taskExecutionMetadata.OnGetNamespace().Return("test-namespace")
 	taskExecutionMetadata.OnGetAnnotations().Return(map[string]string{"iam.amazonaws.com/role": "metadata_role"})
+	taskExecutionMetadata.OnGetSecurityContext().Return(flyteIdlCore.SecurityContext{})
 	taskExecutionMetadata.OnGetLabels().Return(map[string]string{"label-1": "val1"})
 	taskExecutionMetadata.OnGetOwnerReference().Return(v1.OwnerReference{
 		Kind: "node",
@@ -353,6 +358,7 @@ func generateMockHyperparameterTuningJobTaskContext(taskTemplate *flyteIdlCore.T
 	outputReader.OnGetOutputPath().Return(storage.DataReference("/data/outputs.pb"))
 	outputReader.OnGetOutputPrefixPath().Return(storage.DataReference("/data/"))
 	outputReader.OnGetRawOutputPrefix().Return(storage.DataReference("/raw/"))
+
 	taskCtx.OnOutputWriter().Return(outputReader)
 
 	taskReader := &mocks.TaskReader{}
@@ -384,6 +390,9 @@ func genMockTaskExecutionMetadata() *mocks.TaskExecutionMetadata {
 	taskExecutionMetadata.OnGetTaskExecutionID().Return(tID)
 	taskExecutionMetadata.OnGetNamespace().Return("test-namespace")
 	taskExecutionMetadata.OnGetAnnotations().Return(map[string]string{"iam.amazonaws.com/role": "metadata_role"})
+	taskExecutionMetadata.OnGetSecurityContext().Return(flyteIdlCore.SecurityContext{
+		RunAs: &flyteIdlCore.Identity{IamRole: "default_role"},
+	})
 	taskExecutionMetadata.OnGetLabels().Return(map[string]string{"label-1": "val1"})
 	taskExecutionMetadata.OnGetOwnerReference().Return(v1.OwnerReference{
 		Kind: "node",

go/tasks/plugins/k8s/sidecar/sidecar.go

@@ -118,8 +118,7 @@ func (sidecarResourceHandler) BuildResource(ctx context.Context, taskCtx plugins
 	// CrashLoopBackoff after the initial job completion.
 	pod.Spec.RestartPolicy = k8sv1.RestartPolicyNever
 
-	// We want to also update the serviceAccount to the serviceaccount of the workflow
-	pod.Spec.ServiceAccountName = taskCtx.TaskExecutionMetadata().GetK8sServiceAccount()
+	pod.Spec.ServiceAccountName = flytek8s.GetServiceAccountNameFromTaskExecutionMetadata(taskCtx.TaskExecutionMetadata())
 
 	pod, err = validateAndFinalizePod(ctx, taskCtx, primaryContainerName, *pod)
 	if err != nil {

go/tasks/plugins/k8s/sidecar/sidecar_test.go

@@ -56,12 +56,14 @@ func dummyContainerTaskMetadata(resources *v1.ResourceRequirements) pluginsCore.
 	taskMetadata := &pluginsCoreMock.TaskExecutionMetadata{}
 	taskMetadata.On("GetNamespace").Return("test-namespace")
 	taskMetadata.On("GetAnnotations").Return(map[string]string{"annotation-1": "val1"})
+
 	taskMetadata.On("GetLabels").Return(map[string]string{"label-1": "val1"})
 	taskMetadata.On("GetOwnerReference").Return(metav1.OwnerReference{
 		Kind: "node",
 		Name: "blah",
 	})
 	taskMetadata.On("IsInterruptible").Return(true)
+	taskMetadata.On("GetSecurityContext").Return(core.SecurityContext{})
 	taskMetadata.On("GetK8sServiceAccount").Return("service-account")
 	taskMetadata.On("GetOwnerID").Return(types.NamespacedName{
 		Namespace: "test-namespace",
@@ -319,6 +321,7 @@ func TestBuildSidecarResource(t *testing.T) {
 	assert.Len(t, res.(*v1.Pod).Spec.Containers[0].VolumeMounts, 1)
 	assert.Equal(t, "volume mount", res.(*v1.Pod).Spec.Containers[0].VolumeMounts[0].Name)
 
+	assert.Equal(t, "service-account", res.(*v1.Pod).Spec.ServiceAccountName)
 	// Assert user-specified tolerations don't get overridden
 	assert.Len(t, res.(*v1.Pod).Spec.Tolerations, 1)
 	for _, tol := range res.(*v1.Pod).Spec.Tolerations {