fix: check if podSpec is not nil

Signed-off-by: machichima <nary12321@gmail.com>
flyteorg · Feb 12, 2025 · c7dbcfb · c7dbcfb
1 parent a783bdd
commit c7dbcfb
Showing 1 changed file with 24 additions and 21 deletions.
diff --git a/flyteplugins/go/tasks/plugins/k8s/spark/spark.go b/flyteplugins/go/tasks/plugins/k8s/spark/spark.go
@@ -200,22 +200,24 @@ func createDriverSpec(ctx context.Context, taskCtx pluginsCore.TaskExecutionCont
 
 	driverPod := sparkJob.GetDriverPod()
 	if driverPod != nil {
-		var customPodSpec *v1.PodSpec
+		if driverPod.GetPodSpec() != nil {
+			var customPodSpec *v1.PodSpec
 
-		err = utils.UnmarshalStructToObj(driverPod.GetPodSpec(), &customPodSpec)
-		if err != nil {
-			return nil, errors.Errorf(errors.BadTaskSpecification,
-				"Unable to unmarshal driver pod spec [%v], Err: [%v]", driverPod.GetPodSpec(), err.Error())
+			err = utils.UnmarshalStructToObj(driverPod.GetPodSpec(), &customPodSpec)
+			if err != nil {
+				return nil, errors.Errorf(errors.BadTaskSpecification,
+					"Unable to unmarshal driver pod spec [%v], Err: [%v]", driverPod.GetPodSpec(), err.Error())
+			}
+
+			podSpec, err = flytek8s.MergePodSpecs(podSpec, customPodSpec, primaryContainerName, "")
+			if err != nil {
+				return nil, err
+			}
 		}
 
 		if driverPod.GetPrimaryContainerName() != "" {
 			primaryContainerName = driverPod.GetPrimaryContainerName()
 		}
-
-		podSpec, err = flytek8s.MergePodSpecs(podSpec, customPodSpec, primaryContainerName, "")
-		if err != nil {
-			return nil, err
-		}
 	}
 
 	primaryContainer, err := flytek8s.GetContainer(podSpec, primaryContainerName)
@@ -251,22 +253,23 @@ func createExecutorSpec(ctx context.Context, taskCtx pluginsCore.TaskExecutionCo
 
 	executorPod := sparkJob.GetExecutorPod()
 	if executorPod != nil {
-		var customPodSpec *v1.PodSpec
+		if executorPod.GetPodSpec() != nil {
+			var customPodSpec *v1.PodSpec
 
-		err = utils.UnmarshalStructToObj(executorPod.GetPodSpec(), &customPodSpec)
-		if err != nil {
-			return nil, errors.Errorf(errors.BadTaskSpecification,
-				"Unable to unmarshal executor pod spec [%v], Err: [%v]", executorPod.GetPodSpec(), err.Error())
-		}
+			err = utils.UnmarshalStructToObj(executorPod.GetPodSpec(), &customPodSpec)
+			if err != nil {
+				return nil, errors.Errorf(errors.BadTaskSpecification,
+					"Unable to unmarshal executor pod spec [%v], Err: [%v]", executorPod.GetPodSpec(), err.Error())
+			}
 
+			podSpec, err = flytek8s.MergePodSpecs(podSpec, customPodSpec, primaryContainerName, "")
+			if err != nil {
+				return nil, err
+			}
+		}
 		if executorPod.GetPrimaryContainerName() != "" {
 			primaryContainerName = executorPod.GetPrimaryContainerName()
 		}
-
-		podSpec, err = flytek8s.MergePodSpecs(podSpec, customPodSpec, primaryContainerName, "")
-		if err != nil {
-			return nil, err
-		}
 	}
 
 	primaryContainer, err := flytek8s.GetContainer(podSpec, primaryContainerName)