ssh: add --keep-detached-processes so detached processes prevent the server shutdown

.nextchanges/cli/ssh-keep-detached-processes.md

@@ -0,0 +1 @@
+* `ssh connect` and `ssh setup` now accept a `--keep-detached-processes` flag to keep processes detached from the SSH session (`tmux`, `setsid`, `nohup`) running after the tunnel shuts down. Teardown then terminates only the tunnel's own process group, and the bootstrap job run is held open while any detached process is still running, so the survivors keep their `/Workspace` and `/Volumes` access. A held-open run also suppresses cluster autotermination, so the flag is off by default, is bounded by `--server-timeout`, and is dedicated-cluster only. Without it, the server now logs a warning naming the detached processes it is about to destroy, instead of sweeping them silently. ([#6387](https://github.com/databricks/cli/pull/6387))

acceptance/ssh/connect-serverless-cpu/output.txt

@@ -21,6 +21,7 @@
         "notebook_task": {
           "base_parameters": {
             "authorizedKeySecretName": "client-public-key",
+            "keepDetachedProcesses": "false",
             "maxClients": "10",
             "secretScopeName": "[USERNAME]-[CPU_CONN]-ssh-tunnel-keys",
             "serverless": "true",

acceptance/ssh/connect-serverless-gpu/output.txt

@@ -22,6 +22,7 @@
         "notebook_task": {
           "base_parameters": {
             "authorizedKeySecretName": "client-public-key",
+            "keepDetachedProcesses": "false",
             "maxClients": "10",
             "secretScopeName": "[USERNAME]-serverless-gpu-test-ssh-tunnel-keys",
             "serverless": "true",

acceptance/ssh/connection/output.txt

@@ -11,6 +11,7 @@
         "notebook_task": {
           "base_parameters": {
             "authorizedKeySecretName": "client-public-key",
+            "keepDetachedProcesses": "false",
             "maxClients": "10",
             "secretScopeName": "[USERNAME]-[TEST_DEFAULT_CLUSTER_ID]-ssh-tunnel-keys",
             "serverless": "false",

acceptance/ssh/setup/output.txt

@@ -5,6 +5,9 @@ ssh connect --proxy --cluster=[TEST_DEFAULT_CLUSTER_ID] --auto-start-cluster=tru
 === A shutdown delay beyond the default lifetime raises it, no --server-timeout needed
 ssh connect --proxy --cluster=[TEST_DEFAULT_CLUSTER_ID] --auto-start-cluster=true --shutdown-delay=48h0m0s --max-clients=10 --server-timeout=48h0m0s
 
+=== ProxyCommand written by setup --keep-detached-processes
+ssh connect --proxy --cluster=[TEST_DEFAULT_CLUSTER_ID] --auto-start-cluster=true --shutdown-delay=10m0s --keep-detached-processes --max-clients=10 --server-timeout=24h0m0s
+
 === Rejects a server that would refuse every connection
 >>> [CLI] ssh setup --name=broken --cluster=[TEST_DEFAULT_CLUSTER_ID] --max-clients=0
 Error: --max-clients must be at least 1, got 0
@@ -14,5 +17,6 @@ Error: --max-clients must be at least 1, got 0
 Error: --shutdown-delay (48h0m0s) cannot be longer than --server-timeout (24h0m0s)
 
 === No host config is written for the rejected setups
+home/.databricks/ssh-tunnel-configs/keep-detached
 home/.databricks/ssh-tunnel-configs/long-delay
 home/.databricks/ssh-tunnel-configs/my-cluster

acceptance/ssh/setup/script

@@ -17,11 +17,17 @@ title "A shutdown delay beyond the default lifetime raises it, no --server-timeo
 $CLI ssh setup --name=long-delay --cluster=$TEST_DEFAULT_CLUSTER_ID --shutdown-delay=48h &>LOG.long-delay
 sed -n 's/.*\(ssh connect --proxy.*\)/\1/p' "$HOME/.databricks/ssh-tunnel-configs/long-delay"
 
+# Holding the job run open for detached processes is fixed at submission too, so a host
+# configured through setup can only ask for it here.
+title "ProxyCommand written by setup --keep-detached-processes\n"
+$CLI ssh setup --name=keep-detached --cluster=$TEST_DEFAULT_CLUSTER_ID --keep-detached-processes &>LOG.keep-detached
+sed -n 's/.*\(ssh connect --proxy.*\)/\1/p' "$HOME/.databricks/ssh-tunnel-configs/keep-detached"
+
 title "Rejects a server that would refuse every connection"
 musterr trace $CLI ssh setup --name=broken --cluster=$TEST_DEFAULT_CLUSTER_ID --max-clients=0
 
 title "Rejects a shutdown delay the server can never reach"
 musterr trace $CLI ssh setup --name=broken --cluster=$TEST_DEFAULT_CLUSTER_ID --shutdown-delay=48h --server-timeout=24h
 
 title "No host config is written for the rejected setups\n"
-find.py 'ssh-tunnel-configs' --expect 2
+find.py 'ssh-tunnel-configs' --expect 3

experimental/ssh/README.md

@@ -70,6 +70,52 @@ See [filesystem troubleshooting](./FAILURE_MODES.md#filesystem-access-after-the-
 To reproduce and test the known `ssh connect` failure modes (container missing `sshd`, or a
 container that can't run the Python bootstrap), see [FAILURE_MODES.md](./FAILURE_MODES.md).
 
+## Keeping detached processes alive
+
+By default nothing outlives the session: when the last client disconnects, the server shuts
+down after `--shutdown-delay` and the bootstrap notebook sweeps every process it parents,
+including work that was deliberately detached with `tmux`, `setsid` or `nohup`.
+
+`--keep-detached-processes` prevents idle shutdown while detached work is still running.
+It works with dedicated clusters and serverless compute:
+
+```sh
+databricks ssh connect --cluster=<id> --keep-detached-processes
+databricks ssh connect --name=my-session --keep-detached-processes
+```
+
+When `--shutdown-delay` elapses with no SSH clients, the server checks for detached processes.
+If it finds any, it stays available for reconnection and checks again every 15 seconds.
+Once no detached work remains, it shuts down. A reconnect cancels the pending check; after
+the last client disconnects again, the full `--shutdown-delay` applies again. If the process
+tree cannot be read, the server postpones shutdown and retries rather than risking the work.
+
+- **It keeps compute running.** On dedicated clusters the active job also
+  suppresses autotermination. An idle `tmux` session counts as detached work even after the
+  command in its pane finishes; close the session when you no longer need it.
+- **Reconnect to the same session.** Use the same cluster ID or serverless connection name.
+  The existing server and notebook remain alive, so a reconnect can attach to the original
+  `tmux` session rather than creating a replacement run.
+- **The maximum lifetime still applies.** `--server-timeout` (24h by default) bounds the job
+  from its start, regardless of connected clients or detached work. The flag does not
+  survive a job cancellation, notebook restart, or compute termination. Multi-day work
+  belongs in Jobs/DABs.
+- **The notebook must remain alive too.** It anchors the detached processes' workspace
+  filesystem access. If the SSH server exits for another reason, the bootstrap still
+  preserves detached work and holds the run open, as before. This fallback preserves work,
+  not SSH access: it does not restart the server inside that run.
+
+`databricks ssh setup` takes the same flag and bakes it into the host's `ProxyCommand`, so
+`ssh <name>` sessions ask for it too. That is the only place a configured host can set it: the
+`ProxyCommand` is the invocation that submits the run, and the mode is fixed at submission.
+
+A reconnect that omits the flag reuses a running server that was started with it, hold
+included, so a session that never asked for it can end up holding the cluster open. Asking for
+it against a server that was started without it starts a fresh server instead.
+
+When the flag is *not* set and the server does find detached processes at teardown, it logs a
+warning naming them, so work that is about to be swept is no longer lost silently.
+
 ## Design
 
 High level:

experimental/ssh/cmd/connect.go

@@ -60,12 +60,14 @@ Connect to a dedicated cluster:
 	var baseEnvironment string
 	var autoApprove bool
 	var usagePolicyID string
+	var keepDetachedProcesses bool
 
 	cmd.Flags().StringVar(&clusterID, "cluster", "", "Databricks dedicated cluster ID")
 	cmd.Flags().DurationVar(&shutdownDelay, "shutdown-delay", defaultShutdownDelay, "Delay before shutting down the server after the last client disconnects")
 	cmd.Flags().IntVar(&maxClients, "max-clients", defaultMaxClients, "Maximum number of SSH clients")
 	cmd.Flags().DurationVar(&serverTimeout, "server-timeout", defaultServerTimeout, "Maximum lifetime of the SSH server; it is terminated after this duration even if clients are connected")
 	cmd.Flags().BoolVar(&autoStartCluster, "auto-start-cluster", true, "Automatically start the cluster if it is not running")
+	cmd.Flags().BoolVar(&keepDetachedProcesses, "keep-detached-processes", false, "Keep the SSH server and detached processes (tmux, setsid, nohup) running while detached work remains, bounded by --server-timeout")
 
 	cmd.Flags().StringVar(&connectionName, "name", "", "Connection name to reuse across sessions (serverless only)")
 	cmd.Flags().StringVar(&accelerator, "accelerator", "", "Serverless GPU accelerator type (GPU_1xA10 or GPU_8xH100)")
@@ -128,32 +130,33 @@ Connect to a dedicated cluster:
 			environmentVersion = 0
 		}
 		opts := client.ClientOptions{
-			Profile:              wsClient.Config.Profile,
-			ClusterID:            clusterID,
-			ConnectionName:       connectionName,
-			Accelerator:          accelerator,
-			ProxyMode:            proxyMode,
-			IDE:                  ide,
-			ServerMetadata:       serverMetadata,
-			ShutdownDelay:        shutdownDelay,
-			MaxClients:           maxClients,
-			HandoverTimeout:      handoverTimeout,
-			KeepaliveInterval:    defaultKeepaliveInterval,
-			ReleasesDir:          releasesDir,
-			ServerTimeout:        resolveServerTimeout(cmd.Flags(), serverTimeout, shutdownDelay),
-			TaskStartupTimeout:   startupTimeout,
-			AutoStartCluster:     autoStartCluster,
-			ClientPublicKeyName:  clientPublicKeyName,
-			ClientPrivateKeyName: clientPrivateKeyName,
-			ServerPublicKeyName:  serverPublicKeyName,
-			KnownHostsDir:        knownHostsDir,
-			Liteswap:             liteswap,
-			SkipSettingsCheck:    skipSettingsCheck,
-			EnvironmentVersion:   environmentVersion,
-			BaseEnvironment:      baseEnvironment,
-			AdditionalArgs:       args,
-			AutoApprove:          autoApprove,
-			UsagePolicyID:        usagePolicyID,
+			Profile:               wsClient.Config.Profile,
+			ClusterID:             clusterID,
+			ConnectionName:        connectionName,
+			Accelerator:           accelerator,
+			ProxyMode:             proxyMode,
+			IDE:                   ide,
+			ServerMetadata:        serverMetadata,
+			ShutdownDelay:         shutdownDelay,
+			MaxClients:            maxClients,
+			HandoverTimeout:       handoverTimeout,
+			KeepaliveInterval:     defaultKeepaliveInterval,
+			ReleasesDir:           releasesDir,
+			ServerTimeout:         resolveServerTimeout(cmd.Flags(), serverTimeout, shutdownDelay),
+			TaskStartupTimeout:    startupTimeout,
+			AutoStartCluster:      autoStartCluster,
+			ClientPublicKeyName:   clientPublicKeyName,
+			ClientPrivateKeyName:  clientPrivateKeyName,
+			ServerPublicKeyName:   serverPublicKeyName,
+			KnownHostsDir:         knownHostsDir,
+			Liteswap:              liteswap,
+			SkipSettingsCheck:     skipSettingsCheck,
+			EnvironmentVersion:    environmentVersion,
+			BaseEnvironment:       baseEnvironment,
+			AdditionalArgs:        args,
+			AutoApprove:           autoApprove,
+			UsagePolicyID:         usagePolicyID,
+			KeepDetachedProcesses: keepDetachedProcesses,
 		}
 		if err := opts.Validate(); err != nil {
 			return err

experimental/ssh/cmd/server.go

@@ -30,6 +30,7 @@ and proxies them to local SSH daemon processes.`,
 	var authorizedKeySecretName string
 	var serverless bool
 	var usagePolicyID string
+	var keepDetachedProcesses bool
 
 	cmd.Flags().StringVar(&clusterID, "cluster", "", "Databricks cluster ID")
 	cmd.MarkFlagRequired("cluster")
@@ -45,6 +46,7 @@ and proxies them to local SSH daemon processes.`,
 	cmd.Flags().StringVar(&version, "version", "", "Client version of the Databricks CLI")
 	cmd.Flags().BoolVar(&serverless, "serverless", false, "Enable serverless mode for Jupyter initialization")
 	cmd.Flags().StringVar(&usagePolicyID, "usage-policy-id", "", "Usage policy ID the job was submitted with")
+	cmd.Flags().BoolVar(&keepDetachedProcesses, "keep-detached-processes", false, "Keep the SSH server running while detached processes are alive")
 
 	cmd.PreRunE = func(cmd *cobra.Command, args []string) error {
 		// The server can be executed under a directory with an invalid bundle configuration.
@@ -74,6 +76,7 @@ and proxies them to local SSH daemon processes.`,
 			PortRange:               serverPortRange,
 			Serverless:              serverless,
 			UsagePolicyID:           usagePolicyID,
+			KeepDetachedProcesses:   keepDetachedProcesses,
 		}
 		return server.Run(ctx, wsc, opts)
 	}

experimental/ssh/cmd/setup.go

@@ -28,6 +28,7 @@ For serverless connections, use ` + "`databricks ssh connect`" + ` (no setup ste
 	var serverTimeout time.Duration
 	var autoStartCluster bool
 	var autoApprove bool
+	var keepDetachedProcesses bool
 
 	cmd.Flags().StringVar(&hostName, "name", "", "Host name to use in SSH config")
 	cmd.MarkFlagRequired("name")
@@ -37,6 +38,7 @@ For serverless connections, use ` + "`databricks ssh connect`" + ` (no setup ste
 	cmd.Flags().DurationVar(&shutdownDelay, "shutdown-delay", defaultShutdownDelay, "SSH server will terminate after this delay if there are no active connections")
 	cmd.Flags().IntVar(&maxClients, "max-clients", defaultMaxClients, "Maximum number of SSH clients")
 	cmd.Flags().DurationVar(&serverTimeout, "server-timeout", defaultServerTimeout, "Maximum lifetime of the SSH server; it is terminated after this duration even if clients are connected")
+	cmd.Flags().BoolVar(&keepDetachedProcesses, "keep-detached-processes", false, "Keep the SSH server and detached processes (tmux, setsid, nohup) running while detached work remains, bounded by --server-timeout")
 	cmd.Flags().BoolVar(&autoApprove, "auto-approve", false, "Skip confirmation prompts, recreating existing SSH host configs without asking")
 
 	cmd.PreRunE = func(cmd *cobra.Command, args []string) error {
@@ -50,15 +52,16 @@ For serverless connections, use ` + "`databricks ssh connect`" + ` (no setup ste
 		ctx := cmd.Context()
 		wsClient := cmdctx.WorkspaceClient(ctx)
 		setupOpts := setup.SetupOptions{
-			HostName:         hostName,
-			ClusterID:        clusterID,
-			AutoStartCluster: autoStartCluster,
-			SSHConfigPath:    sshConfigPath,
-			ShutdownDelay:    shutdownDelay,
-			MaxClients:       maxClients,
-			ServerTimeout:    resolveServerTimeout(cmd.Flags(), serverTimeout, shutdownDelay),
-			Profile:          wsClient.Config.Profile,
-			AutoApprove:      autoApprove,
+			HostName:              hostName,
+			ClusterID:             clusterID,
+			AutoStartCluster:      autoStartCluster,
+			SSHConfigPath:         sshConfigPath,
+			ShutdownDelay:         shutdownDelay,
+			MaxClients:            maxClients,
+			ServerTimeout:         resolveServerTimeout(cmd.Flags(), serverTimeout, shutdownDelay),
+			KeepDetachedProcesses: keepDetachedProcesses,
+			Profile:               wsClient.Config.Profile,
+			AutoApprove:           autoApprove,
 		}
 		return setup.Setup(ctx, wsClient, setupOpts)
 	}

experimental/ssh/internal/client/client.go

@@ -130,6 +130,8 @@ type ClientOptions struct {
 	AutoApprove bool
 	// Id of the usage policy to use for the serverless SSH server job. Serverless only.
 	UsagePolicyID string
+	// Whether detached processes prevent idle shutdown of the SSH server. Bounded by --server-timeout.
+	KeepDetachedProcesses bool
 }
 
 func (o *ClientOptions) Validate() error {
@@ -251,6 +253,9 @@ func (o *ClientOptions) ToProxyCommand() (string, error) {
 		proxyCommand = fmt.Sprintf("%q ssh connect --proxy --cluster=%s --auto-start-cluster=%t --shutdown-delay=%s",
 			executablePath, o.ClusterID, o.AutoStartCluster, o.ShutdownDelay.String())
 	}
+	if o.KeepDetachedProcesses {
+		proxyCommand += " --keep-detached-processes"
+	}
 
 	// Both of these are fixed when the server job is submitted, and for a host configured by
 	// `ssh setup` the submitting invocation is always the ProxyCommand, so they have to be
@@ -600,6 +605,8 @@ type serverMetadata struct {
 	ClusterID string
 	// UsagePolicyID the server was started with, used to decide whether a running server can be reused.
 	UsagePolicyID string
+	// KeepDetachedProcesses the server's run was submitted with, used the same way as UsagePolicyID.
+	KeepDetachedProcesses bool
 }
 
 // getServerMetadata retrieves the server metadata from the workspace and validates it via Driver Proxy.
@@ -647,10 +654,11 @@ func getServerMetadata(ctx context.Context, client *databricks.WorkspaceClient,
 	}
 
 	return serverMetadata{
-		Port:          wsMetadata.Port,
-		UserName:      string(bodyBytes),
-		ClusterID:     effectiveClusterID,
-		UsagePolicyID: wsMetadata.UsagePolicyID,
+		Port:                  wsMetadata.Port,
+		UserName:              string(bodyBytes),
+		ClusterID:             effectiveClusterID,
+		UsagePolicyID:         wsMetadata.UsagePolicyID,
+		KeepDetachedProcesses: wsMetadata.KeepDetachedProcesses,
 	}, nil
 }
 
@@ -719,6 +727,9 @@ func buildSSHServerSubmitRun(version, secretScopeName, jobNotebookPath, baseEnvi
 		// Recorded in the server's metadata.json so reconnects can tell which usage policy
 		// the running server was started under.
 		"usagePolicyId": opts.UsagePolicyID,
+		// The bootstrap only needs to know whether to hold the run open. How long it ends up
+		// holding it is decided by the work itself, bounded by the run's own timeout.
+		"keepDetachedProcesses": strconv.FormatBool(opts.KeepDetachedProcesses),
 	}
 
 	task := jobs.SubmitTask{
@@ -1302,6 +1313,14 @@ func usagePolicyMatches(storedPolicy, requestedPolicy string) bool {
 	return requestedPolicy == "" || storedPolicy == requestedPolicy
 }
 
+// keepDetachedMatches reports whether a running server holds its run open for detached
+// processes when this connection asked it to. The mode is fixed when the run is submitted, so
+// asking for it needs a server that has it; a connection that does not ask takes whatever is
+// already running.
+func keepDetachedMatches(stored, requested bool) bool {
+	return !requested || stored
+}
+
 func ensureSSHServerIsRunning(ctx context.Context, client *databricks.WorkspaceClient, version, secretScopeName string, opts ClientOptions) (string, int, string, error) {
 	sessionID := opts.SessionIdentifier()
 	// For dedicated clusters, use clusterID; for serverless, it will be read from metadata
@@ -1313,10 +1332,13 @@ func ensureSSHServerIsRunning(ctx context.Context, client *databricks.WorkspaceC
 	}
 
 	// Start a new server when none is running, or when the running one was started under a
-	// different usage policy. A job's usage policy is fixed at submission, so we can't retarget
-	// the existing server; the new server overwrites metadata.json and the old one idles out via
-	// shutdownDelay.
-	needNewServer := err != nil || !usagePolicyMatches(meta.UsagePolicyID, opts.UsagePolicyID)
+	// different usage policy or without keeping detached processes. Both are fixed at
+	// submission, so we can't retarget the existing server; the new server overwrites
+	// metadata.json, and its bootstrap terminates the running server on the cluster before
+	// starting (see cleanup() in ssh-server-bootstrap.py), which ends the previous run.
+	needNewServer := err != nil ||
+		!usagePolicyMatches(meta.UsagePolicyID, opts.UsagePolicyID) ||
+		!keepDetachedMatches(meta.KeepDetachedProcesses, opts.KeepDetachedProcesses)
 	if needNewServer {
 		cmdio.LogString(ctx, "Starting SSH server...")
 
@@ -1340,6 +1362,9 @@ func ensureSSHServerIsRunning(ctx context.Context, client *databricks.WorkspaceC
 			if err == nil && !usagePolicyMatches(meta.UsagePolicyID, opts.UsagePolicyID) {
 				err = fmt.Errorf("found a running SSH server with usage policy %q, waiting for the one with %q", meta.UsagePolicyID, opts.UsagePolicyID)
 			}
+			if err == nil && !keepDetachedMatches(meta.KeepDetachedProcesses, opts.KeepDetachedProcesses) {
+				err = errors.New("found a running SSH server that does not keep detached processes, waiting for the one that does")
+			}
 			if err == nil {
 				cmdio.LogString(ctx, "Health check successful, starting ssh WebSocket connection...")
 				break
@@ -1481,6 +1506,9 @@ func buildSshTunnelEvent(opts ClientOptions, outcome connectOutcome) *protos.Ssh
 		IsSuccess:          outcome.isSuccess,
 		HasBaseEnvironment: opts.BaseEnvironment != "",
 		HasUsagePolicy:     opts.UsagePolicyID != "",
-		ErrorCategory:      outcome.category(),
+		// The connect side can only report that the knob was asked for. Whether any detached
+		// process was there to keep is reported by the server, at teardown.
+		KeepDetachedRequested: opts.KeepDetachedProcesses,
+		ErrorCategory:         outcome.category(),
 	}
 }