fix(spark): pass driver/executor pod specs through as the pod template

[pingsutw]

Why are the changes needed?

A Spark task that customizes its driver or executor pod cannot do so today, and in the
common case it fails outright:

flyte.TaskEnvironment(
    name="spark_exec_disk",
    plugin_config=Spark(
        spark_conf=base_conf,
        executor_pod=flyte.PodTemplate(
            primary_container_name="primary",
            pod_spec=V1PodSpec(containers=[V1Container(
                name="primary",
                resources=V1ResourceRequirements(requests={"ephemeral-storage": "9Gi"}),
            )]),
        ),
    ),
)

[BadTaskSpecification] invalid TaskSpecification, container [primary] not defined

createDriverSpec/createExecutorSpec merged the custom pod spec onto the pod spec Flyte
builds with MergeOverlayPodSpecOntoBase. That merge matches containers by name, and
iterates the base containers only — the container Flyte generates is named after the task
execution ID (<TaskAction name>-<attempt> in v2, container_helper.go BuildRawContainer),
which no user can predict. So:

1. any container in the custom pod spec is silently dropped (the ephemeral-storage request
   above never reaches the pod), and
2. primary_container_name then overrode the primary container name, and the following
   GetContainer(podSpec, primaryContainerName) failed the task.

Naming the container after the operator's spark-kubernetes-driver /
spark-kubernetes-executor — the names the Spark operator itself documents for pod
templates — hits exactly the same wall.

What changes were proposed in this pull request?

Hand the custom pod spec to the Spark operator verbatim as
spec.{driver,executor}.template instead of merging it into Flyte's pod spec, and drop the
primary_container_name override. The operator patches the template onto the container it
generates, which is the container users are naming in the first place.

• createSparkPodSpec now takes the custom pod spec separately and uses it for
  spec.Template, falling back to the pod spec Flyte builds when the task sets no custom
  driver/executor pod (unchanged behavior for the vast majority of tasks).
• The explicit SparkPodSpec fields (affinity, tolerations, node selector, image, env, …)
  are still derived from the pod spec Flyte builds, so platform scheduling constraints stay
  intact — the operator treats explicit fields as overrides of the template, so a custom
  affinity in the template can no longer dilute the non-interruptible node-selector
  requirement.

Tradeoff, called out explicitly: the template is now the only path a custom
driver/executor pod spec travels. On a cluster whose SparkApplication CRD predates
spec.driver.template (or with the enable-pod-template: false kill switch), the custom
pod spec no longer applies at all, where previously its pod-level fields were merged into
the legacy SparkPodSpec fields. Container-level fields never worked on that path anyway.
TestBuildResourceCustomPodSpecDroppedWithoutTemplateSupport pins this behavior.

How was this patch tested?

go test ./go/tasks/plugins/k8s/spark/ — green. Tests added/updated:

• TestBuildResourceCustomPodSpecPassthrough — driver pod named spark-kubernetes-driver
  and executor pod named primary, both with an ephemeral-storage request and a matching
  primary_container_name. Asserts the build succeeds and both containers reach
  spec.{driver,executor}.template verbatim, while image/env still come from Flyte's pod
  spec. Fails on main with invalid TaskSpecification, container [spark-kubernetes-driver] not defined (verified).
• TestBuildResourceCustomPodSpecDroppedWithoutTemplateSupport — documents the tradeoff
  above on a legacy CRD.
• TestBuildResourceSparkExecutorAffinityDilution — updated: the custom affinity now lands
  on the template, and every node selector term on the explicit Affinity field still
  carries the non-interruptible requirement.
• TestBuildResourcePodTemplateGating (existing) covers the no-custom-pod-spec fallback.

Labels

fixed

Check all the applicable boxes

• I updated the documentation accordingly. (n/a — no user-facing config change)
• All new and existing tests passed.
• All commits are signed-off.

https://claude.ai/code/session_01SAATDRHXajHapuojwQ4SMR

[Copilot]

Pull request overview

This PR fixes Spark driver/executor pod customization by passing user-provided driver/executor PodSpecs directly to the Spark Operator via spec.{driver,executor}.template, instead of merging them into Flyte’s generated pod spec (which dropped container-level overrides due to unpredictable Flyte container naming).

Changes:

• Updated SparkApplication construction so createSparkPodSpec can use a separate custom pod spec for the pod template (falling back to Flyte’s generated pod spec when none is provided).
• Removed the primary_container_name override path so user container naming no longer breaks task validation.
• Added/updated unit tests to validate custom pod spec passthrough, legacy-CRD behavior when template support is absent, and affinity “dilution” behavior.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
flyteplugins/go/tasks/plugins/k8s/spark/spark.go	Pass custom driver/executor pod specs through as Spark Operator pod templates; stop merging into Flyte’s pod spec.
flyteplugins/go/tasks/plugins/k8s/spark/spark_test.go	Add/adjust tests covering custom pod spec template passthrough and legacy/template-gating behaviors.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated no new comments.

Suppressed comments (3)

flyteplugins/go/tasks/plugins/k8s/spark/spark.go:270

• This uses pluginmachinery/utils.UnmarshalStructToObj, which is marked deprecated in flyteplugins/go/tasks/pluginmachinery/utils/marshal_utils.go. Consider switching to flytestdlib/utils.UnmarshalStructToObj instead and dropping the staticcheck suppression, so new code paths aren't built on deprecated helpers.

		if executorPod.GetPodSpec() != nil {
			err = utils.UnmarshalStructToObj(executorPod.GetPodSpec(), &customPodSpec) //nolint: staticcheck
			if err != nil {
				return nil, errors.Errorf(errors.BadTaskSpecification,
					"Unable to unmarshal executor pod spec [%v], Err: [%v]", executorPod.GetPodSpec(), err.Error())
			}

flyteplugins/go/tasks/plugins/k8s/spark/spark.go:186

• The comment says the custom driver/executor pod spec is passed through "verbatim", but the implementation intentionally applies some defaulting (e.g. EnableServiceLinks when unset). Consider tightening the wording to avoid implying a strict byte-for-byte passthrough.

	// the operator treats explicit fields as overrides of the template. The user's driver/executor
	// pod spec is passed through verbatim so the operator can patch it onto the container it
	// generates (`spark-kubernetes-driver`/`spark-kubernetes-executor`); Flyte's own container
	// names never match those, so merging it here would drop it on the floor.

flyteplugins/go/tasks/plugins/k8s/spark/spark.go:223

• This uses pluginmachinery/utils.UnmarshalStructToObj, which is marked deprecated in flyteplugins/go/tasks/pluginmachinery/utils/marshal_utils.go. Consider switching to flytestdlib/utils.UnmarshalStructToObj instead and dropping the staticcheck suppression, so new code paths aren't built on deprecated helpers.

This issue also appears on line 265 of the same file.

		if driverPod.GetPodSpec() != nil {
			err = utils.UnmarshalStructToObj(driverPod.GetPodSpec(), &customPodSpec) //nolint: staticcheck
			if err != nil {
				return nil, errors.Errorf(errors.BadTaskSpecification,
					"Unable to unmarshal driver pod spec [%v], Err: [%v]", driverPod.GetPodSpec(), err.Error())
			}