fix(controller): cancel in-flight spark-submit on controller shutdown

When the controller leader receives SIGTERM during spark-submit, the reconcile
context is cancelled immediately and the JVM subprocess may keep running until
the process exits or the kubelet sends SIGKILL. That can create a driver pod
the controller never records, contributing to duplicate submissions on restart.

- Add SparkSubmitter.runSparkSubmit using exec.CommandContext; send SIGTERM
  on cancellation and wait up to ShutdownGracePeriod before force-killing.
- Reuse --driver-pod-creation-grace-period for ShutdownGracePeriod and
  Manager.GracefulShutdownTimeout (no separate graceful-shutdown flag).
- Helm: controller pod terminationGracePeriodSeconds default 60.
- Add runSparkSubmit unit tests and chart deployment tests.

Refs #2788

Signed-off-by: Amir Alavi <amiralavi7@gmail.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

Author: a7i

charts/spark-operator-chart/README.md

@@ -102,7 +102,8 @@ See [helm uninstall](https://helm.sh/docs/helm/helm_uninstall) for command docum
 | controller.workers | int | `10` | Reconcile concurrency, higher values might increase memory usage. |
 | controller.logLevel | string | `"info"` | Configure the verbosity of logging, can be one of `debug`, `info`, `error`. |
 | controller.logEncoder | string | `"console"` | Configure the encoder of logging, can be one of `console` or `json`. |
-| controller.driverPodCreationGracePeriod | string | `"10s"` | Grace period after a successful spark-submit when driver pod not found errors will be retried. Useful if the driver pod can take some time to be created. |
+| controller.driverPodCreationGracePeriod | string | `"10s"` | Grace period after a successful spark-submit when driver pod not found errors will be retried. Also caps how long an in-flight spark-submit is given to exit after SIGTERM during controller shutdown (reconcile context is cancelled immediately on SIGTERM; this window is for reaping the JVM subprocess). Set lower than `terminationGracePeriodSeconds` so the manager exits before the kubelet sends SIGKILL. |
+| controller.terminationGracePeriodSeconds | int | `60` | Termination grace period for the controller pod. The kubelet sends SIGKILL when this elapses. Must be set higher than `controller.driverPodCreationGracePeriod` so the manager has time to interrupt and reap in-flight `spark-submit` subprocesses before the process is forcibly killed. |
 | controller.maxTrackedExecutorPerApp | int | `1000` | Specifies the maximum number of Executor pods that can be tracked by the controller per SparkApplication. |
 | controller.kubeAPIQPS | int | `20` | Maximum QPS to the API server from the controller client. |
 | controller.kubeAPIBurst | int | `30` | Maximum burst for throttle from the controller client. |

charts/spark-operator-chart/templates/controller/deployment.yaml

@@ -212,6 +212,9 @@ spec:
       {{- end }}
       serviceAccountName: {{ include "spark-operator.controller.serviceAccountName" . }}
       automountServiceAccountToken: {{ .Values.controller.serviceAccount.automountServiceAccountToken }}
+      {{- with .Values.controller.terminationGracePeriodSeconds }}
+      terminationGracePeriodSeconds: {{ . }}
+      {{- end }}
       {{- with .Values.controller.podSecurityContext }}
       securityContext:
         {{- toYaml . | nindent 8 }}

charts/spark-operator-chart/tests/controller/deployment_test.yaml

@@ -844,6 +844,21 @@ tests:
           path: spec.template.spec.containers[?(@.name=="spark-operator-controller")].args
           content: --max-tracked-executor-per-app=123
 
+  - it: Should set `terminationGracePeriodSeconds` on the pod by default
+    asserts:
+      - equal:
+          path: spec.template.spec.terminationGracePeriodSeconds
+          value: 60
+
+  - it: Should override `terminationGracePeriodSeconds` if `controller.terminationGracePeriodSeconds` is set
+    set:
+      controller:
+        terminationGracePeriodSeconds: 240
+    asserts:
+      - equal:
+          path: spec.template.spec.terminationGracePeriodSeconds
+          value: 240
+
 
   - it: Should add leader election parameters if `controller.leaderElection.leaseDuration`, `controller.leaderElection.renewDeadline` and `controller.leaderElection.retryPeriod` are set.
     set:

charts/spark-operator-chart/values.yaml

@@ -107,9 +107,19 @@ controller:
   # -- Configure the encoder of logging, can be one of `console` or `json`.
   logEncoder: console
 
-  # -- Grace period after a successful spark-submit when driver pod not found errors will be retried. Useful if the driver pod can take some time to be created.
+  # -- Grace period after a successful spark-submit when driver pod not found errors will be retried.
+  # Also caps how long an in-flight spark-submit is given to exit after SIGTERM during controller
+  # shutdown (reconcile context is cancelled immediately on SIGTERM; this window is for reaping the
+  # JVM subprocess). Set lower than `terminationGracePeriodSeconds` so the manager exits before the
+  # kubelet sends SIGKILL.
   driverPodCreationGracePeriod: 10s
 
+  # -- Termination grace period for the controller pod. The kubelet sends SIGKILL when this
+  # elapses. Must be set higher than `controller.driverPodCreationGracePeriod` so the manager
+  # has time to interrupt and reap in-flight `spark-submit` subprocesses before the process is
+  # forcibly killed.
+  terminationGracePeriodSeconds: 60
+
   # -- Specifies the maximum number of Executor pods that can be tracked by the controller per SparkApplication.
   maxTrackedExecutorPerApp: 1000
 

cmd/operator/controller/start.go

@@ -222,7 +222,7 @@ func NewStartCommand() *cobra.Command {
 	command.Flags().DurationVar(&leaderElectionRenewDeadline, "leader-election-renew-deadline", 10*time.Second, "Leader election renew deadline.")
 	command.Flags().DurationVar(&leaderElectionRetryPeriod, "leader-election-retry-period", 2*time.Second, "Leader election retry period.")
 
-	command.Flags().DurationVar(&driverPodCreationGracePeriod, "driver-pod-creation-grace-period", 10*time.Second, "Grace period after a successful spark-submit when driver pod not found errors will be retried. Useful if the driver pod can take some time to be created.")
+	command.Flags().DurationVar(&driverPodCreationGracePeriod, "driver-pod-creation-grace-period", 10*time.Second, "Grace period after a successful spark-submit when driver pod not found errors will be retried, and maximum time to wait for an in-flight spark-submit to exit after SIGTERM during controller shutdown. Useful if driver pod creation can take some time. Should be set lower than the pod's terminationGracePeriodSeconds.")
 
 	command.Flags().Float32Var(&kubeAPIQPS, "kube-api-qps", 20, "Maximum QPS to the API server from the controller client.")
 	command.Flags().IntVar(&kubeAPIBurst, "kube-api-burst", 30, "Maximum burst for throttle from the controller client.")
@@ -312,6 +312,7 @@ func start() {
 		LeaseDuration:           &leaderElectionLeaseDuration,
 		RenewDeadline:           &leaderElectionRenewDeadline,
 		RetryPeriod:             &leaderElectionRetryPeriod,
+		GracefulShutdownTimeout: &driverPodCreationGracePeriod,
 		// LeaderElectionReleaseOnCancel defines if the leader should step down voluntarily
 		// when the Manager ends. This requires the binary to immediately end when the
 		// Manager is stopped, otherwise, this setting is unsafe. Setting this significantly
@@ -539,7 +540,9 @@ func newSparkConnectReconcilerOptions() sparkconnect.Options {
 
 func newSparkSubmitter(ctx context.Context) (sparkapplication.SparkApplicationSubmitter, error) {
 	if !features.Enabled(features.RestSubmitter) {
-		return &sparkapplication.SparkSubmitter{}, nil
+		return &sparkapplication.SparkSubmitter{
+			ShutdownGracePeriod: driverPodCreationGracePeriod,
+		}, nil
 	}
 
 	var tlsCfg *sparkapplication.TLSConfig

internal/controller/sparkapplication/submission.go

@@ -25,6 +25,8 @@ import (
 	"reflect"
 	"slices"
 	"strings"
+	"syscall"
+	"time"
 
 	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
@@ -44,37 +46,53 @@ type SparkApplicationSubmitter interface {
 
 // SparkSubmitter submits a SparkApplication by calling spark-submit.
 type SparkSubmitter struct {
+	// ShutdownGracePeriod caps how long runSparkSubmit waits for spark-submit to exit
+	// after SIGTERM when the reconcile context is cancelled (e.g. controller shutdown).
+	ShutdownGracePeriod time.Duration
 }
 
 // SparkSubmitter implements SparkApplicationSubmitter interface.
 // This interface is highly experimental and may go under significant changes or removed in the future.
 var _ SparkApplicationSubmitter = &SparkSubmitter{}
 
 // Submit implements SparkApplicationSubmitter interface.
-func (*SparkSubmitter) Submit(ctx context.Context, app *v1beta2.SparkApplication) error {
+func (s *SparkSubmitter) Submit(ctx context.Context, app *v1beta2.SparkApplication) error {
 	logger := log.FromContext(ctx)
 
 	args, err := buildSparkSubmitArgs(app, false)
 	if err != nil {
 		return fmt.Errorf("failed to build spark-submit arguments: %v", err)
 	}
 
-	// Try submitting the application by running spark-submit.
+	// Try submitting the application by running spark-submit. The reconcile context is
+	// threaded into spark-submit so that controller shutdown (or any other context
+	// cancellation) signals the child process to exit. Cancellation can still arrive
+	// after spark-submit has already created the driver pod (e.g., between the create
+	// call and process exit), so this reduces, rather than eliminates, the chance of
+	// losing a status update for a freshly submitted SparkApplication.
 	logger.Info("Running spark-submit", "arguments", args)
-	if err := runSparkSubmit(args); err != nil {
+	if err := s.runSparkSubmit(ctx, args); err != nil {
 		return fmt.Errorf("failed to run spark-submit: %v", err)
 	}
 
 	return nil
 }
 
-func runSparkSubmit(args []string) error {
+func (s *SparkSubmitter) runSparkSubmit(ctx context.Context, args []string) error {
 	sparkHome, present := os.LookupEnv(common.EnvSparkHome)
-	if !present {
+	if !present || sparkHome == "" {
 		return fmt.Errorf("env %s is not specified", common.EnvSparkHome)
 	}
 	command := filepath.Join(sparkHome, "bin", "spark-submit")
-	cmd := exec.Command(command, args...)
+	cmd := exec.CommandContext(ctx, command, args...)
+	// exec.CommandContext defaults to SIGKILL on cancellation. Override so we send
+	// SIGTERM first, giving spark-submit and the underlying JVM a chance to shut
+	// down gracefully (release the API server connection, flush logs). WaitDelay
+	// caps how long we wait after the signal before the process is force-killed.
+	cmd.Cancel = func() error {
+		return cmd.Process.Signal(syscall.SIGTERM)
+	}
+	cmd.WaitDelay = s.ShutdownGracePeriod
 	_, err := cmd.Output()
 	if err != nil {
 		var errorMsg string

internal/controller/sparkapplication/submission_test.go

@@ -17,14 +17,19 @@ limitations under the License.
 package sparkapplication
 
 import (
+	"context"
 	"fmt"
 	"os"
+	"path/filepath"
+	"runtime"
 	"slices"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/google/uuid"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/resource"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
@@ -37,6 +42,104 @@ import (
 	"github.com/kubeflow/spark-operator/v2/pkg/util"
 )
 
+// writeFakeSparkHome creates a temporary SPARK_HOME directory containing a
+// `bin/spark-submit` script with the given contents and sets SPARK_HOME for the test.
+func writeFakeSparkHome(t *testing.T, script string) {
+	t.Helper()
+	if runtime.GOOS == "windows" {
+		t.Skip("fake spark-submit shell script is not supported on Windows")
+	}
+	root := t.TempDir()
+	binDir := filepath.Join(root, "bin")
+	require.NoError(t, os.MkdirAll(binDir, 0o755))
+	sparkSubmit := filepath.Join(binDir, "spark-submit")
+	require.NoError(t, os.WriteFile(sparkSubmit, []byte(script), 0o755))
+	t.Setenv(common.EnvSparkHome, root)
+}
+
+func TestRunSparkSubmit(t *testing.T) {
+	submitter := &SparkSubmitter{ShutdownGracePeriod: 10 * time.Second}
+
+	t.Run("returns error when SPARK_HOME is not set", func(t *testing.T) {
+		t.Setenv(common.EnvSparkHome, "")
+
+		err := submitter.runSparkSubmit(context.Background(), nil)
+		require.Error(t, err)
+		assert.Contains(t, err.Error(), common.EnvSparkHome)
+	})
+
+	t.Run("returns nil on successful spark-submit", func(t *testing.T) {
+		writeFakeSparkHome(t, "#!/bin/sh\nexit 0\n")
+
+		assert.NoError(t, submitter.runSparkSubmit(context.Background(), []string{"--foo", "bar"}))
+	})
+
+	t.Run("surfaces stderr when spark-submit fails", func(t *testing.T) {
+		writeFakeSparkHome(t, "#!/bin/sh\necho boom 1>&2\nexit 1\n")
+
+		err := submitter.runSparkSubmit(context.Background(), nil)
+		require.Error(t, err)
+		assert.Contains(t, err.Error(), "boom")
+	})
+
+	t.Run("returns driver pod already exists error", func(t *testing.T) {
+		script := fmt.Sprintf("#!/bin/sh\necho %q 1>&2\nexit 1\n", "got "+common.ErrorCodePodAlreadyExists+" from apiserver")
+		writeFakeSparkHome(t, script)
+
+		err := submitter.runSparkSubmit(context.Background(), nil)
+		require.Error(t, err)
+		assert.Contains(t, err.Error(), "driver pod already exist")
+	})
+
+	t.Run("context cancellation triggers SIGTERM rather than SIGKILL", func(t *testing.T) {
+		if _, err := os.Stat("/bin/bash"); err != nil {
+			t.Skip("/bin/bash is required for this test")
+		}
+		// The script traps SIGTERM, emits a marker to stderr, and exits non-zero.
+		// A foreground busy-loop is used so the trap fires between sleeps and so
+		// no background process keeps stderr open after the script exits (which
+		// would delay Wait by WaitDelay). We require bash (not /bin/sh) because
+		// some POSIX shells (notably bash 3.2 invoked as sh on macOS) do not
+		// deliver trapped signals during builtins like `wait`. If runSparkSubmit
+		// defaulted to SIGKILL (the os/exec default for context-cancelled
+		// commands), the trap would never run and stderr would be empty.
+		// The script writes a "ready" marker once the trap is registered so the
+		// test can wait until SIGTERM will be delivered after the handler is in
+		// place, avoiding a race on slow/loaded test machines.
+		readyMarker := filepath.Join(t.TempDir(), "ready")
+		script := fmt.Sprintf(`#!/bin/bash
+trap 'echo SIGTERM_RECEIVED 1>&2; exit 42' TERM
+touch %q
+while true; do
+  sleep 0.1
+done
+`, readyMarker)
+		writeFakeSparkHome(t, script)
+
+		ctx, cancel := context.WithCancel(context.Background())
+		// Cancel only after the script has installed the trap.
+		go func() {
+			deadline := time.Now().Add(5 * time.Second)
+			for time.Now().Before(deadline) {
+				if _, statErr := os.Stat(readyMarker); statErr == nil {
+					break
+				}
+				time.Sleep(20 * time.Millisecond)
+			}
+			cancel()
+		}()
+
+		start := time.Now()
+		err := submitter.runSparkSubmit(ctx, nil)
+		elapsed := time.Since(start)
+
+		require.Error(t, err)
+		assert.Less(t, elapsed, 8*time.Second, "runSparkSubmit should return promptly after SIGTERM")
+		assert.Contains(t, err.Error(), "SIGTERM_RECEIVED",
+			"spark-submit child should receive SIGTERM, not SIGKILL")
+	})
+}
+
 func TestExecutorConfOption(t *testing.T) {
 	tests := []struct {
 		name     string