Update for RoleSpec

Signed-off-by: Andrew Dye <[email protected]>
flyteorg · Oct 8, 2023 · ad57897 · ad57897
1 parent 583808b
commit ad57897
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 8 deletions.
diff --git a/flyteplugins/go/tasks/plugins/k8s/spark/spark.go b/flyteplugins/go/tasks/plugins/k8s/spark/spark.go
@@ -110,7 +110,7 @@ func (sparkResourceHandler) BuildResource(ctx context.Context, taskCtx pluginsCo
 		serviceAccountName = sparkTaskType
 	}
 
-	tolerations, err := getTolerations(sparkJob.GetDriverPod().GetPodSpec())
+	tolerations, err := getTolerations(sparkJob.GetDriverSpec().GetPod().GetPodSpec())
 	if err != nil {
 		return nil, err
 	}
@@ -131,7 +131,7 @@ func (sparkResourceHandler) BuildResource(ctx context.Context, taskCtx pluginsCo
 		ServiceAccount: &serviceAccountName,
 	}
 
-	tolerations, err = getTolerations(sparkJob.GetExecutorPod().GetPodSpec())
+	tolerations, err = getTolerations(sparkJob.GetExecutorSpec().GetPod().GetPodSpec())
 	if err != nil {
 		return nil, err
 	}

diff --git a/flyteplugins/go/tasks/plugins/k8s/spark/spark_test.go b/flyteplugins/go/tasks/plugins/k8s/spark/spark_test.go
@@ -701,9 +701,11 @@ func TestBuildResourcePodTemplate(t *testing.T) {
 	driverPodSpecPb := structpb.Struct{}
 	err = utils.MarshalStruct(&podSpec, &driverPodSpecPb)
 	assert.NoError(t, err)
-	sparkJob.DriverPodValue = &plugins.SparkJob_DriverPod{
-		DriverPod: &core.K8SPod{
-			PodSpec: &driverPodSpecPb,
+	sparkJob.DriverSpec = &plugins.RoleSpec{
+		PodValue: &plugins.RoleSpec_Pod{
+			Pod: &core.K8SPod{
+				PodSpec: &driverPodSpecPb,
+			},
 		},
 	}
 	extraExecutorToleration := corev1.Toleration{
@@ -719,9 +721,11 @@ func TestBuildResourcePodTemplate(t *testing.T) {
 	execPodSpecPb := structpb.Struct{}
 	err = utils.MarshalStruct(&podSpec, &execPodSpecPb)
 	assert.NoError(t, err)
-	sparkJob.ExecutorPodValue = &plugins.SparkJob_ExecutorPod{
-		ExecutorPod: &core.K8SPod{
-			PodSpec: &execPodSpecPb,
+	sparkJob.ExecutorSpec = &plugins.RoleSpec{
+		PodValue: &plugins.RoleSpec_Pod{
+			Pod: &core.K8SPod{
+				PodSpec: &execPodSpecPb,
+			},
 		},
 	}
 	taskTemplate := dummySparkTaskTemplate("blah-1", sparkJob)