diff --git a/.nextchanges/cli/ssh-keep-detached-processes.md b/.nextchanges/cli/ssh-keep-detached-processes.md new file mode 100644 index 00000000000..975d7b7d311 --- /dev/null +++ b/.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)) diff --git a/acceptance/ssh/connect-serverless-cpu/output.txt b/acceptance/ssh/connect-serverless-cpu/output.txt index 994a35cbb88..1a2da09b42c 100644 --- a/acceptance/ssh/connect-serverless-cpu/output.txt +++ b/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", diff --git a/acceptance/ssh/connect-serverless-gpu/output.txt b/acceptance/ssh/connect-serverless-gpu/output.txt index 7c213823257..c89a29c8193 100644 --- a/acceptance/ssh/connect-serverless-gpu/output.txt +++ b/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", diff --git a/acceptance/ssh/connection/output.txt b/acceptance/ssh/connection/output.txt index 58babb6e4dc..1372be147a3 100644 --- a/acceptance/ssh/connection/output.txt +++ b/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", diff --git a/acceptance/ssh/setup/output.txt b/acceptance/ssh/setup/output.txt index 589680305ce..c07eb2cc138 100644 --- a/acceptance/ssh/setup/output.txt +++ b/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 diff --git a/acceptance/ssh/setup/script b/acceptance/ssh/setup/script index a406a2e4c80..a04c038ff3d 100644 --- a/acceptance/ssh/setup/script +++ b/acceptance/ssh/setup/script @@ -17,6 +17,12 @@ 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 @@ -24,4 +30,4 @@ 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 diff --git a/experimental/ssh/README.md b/experimental/ssh/README.md index 7d93b72d8e1..90800d8987e 100644 --- a/experimental/ssh/README.md +++ b/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= --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 ` 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: diff --git a/experimental/ssh/cmd/connect.go b/experimental/ssh/cmd/connect.go index bb965d8306b..9ac3af75464 100644 --- a/experimental/ssh/cmd/connect.go +++ b/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 diff --git a/experimental/ssh/cmd/server.go b/experimental/ssh/cmd/server.go index 3675a4a7fe8..6bb0d4db821 100644 --- a/experimental/ssh/cmd/server.go +++ b/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) } diff --git a/experimental/ssh/cmd/setup.go b/experimental/ssh/cmd/setup.go index a8848e741f8..261fd092d39 100644 --- a/experimental/ssh/cmd/setup.go +++ b/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) } diff --git a/experimental/ssh/internal/client/client.go b/experimental/ssh/internal/client/client.go index 1810d7f90ad..2e651e68d8a 100644 --- a/experimental/ssh/internal/client/client.go +++ b/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(), } } diff --git a/experimental/ssh/internal/client/client_internal_test.go b/experimental/ssh/internal/client/client_internal_test.go index ae70cb3e85a..1fc34858488 100644 --- a/experimental/ssh/internal/client/client_internal_test.go +++ b/experimental/ssh/internal/client/client_internal_test.go @@ -511,6 +511,15 @@ func TestBuildSshTunnelEvent(t *testing.T) { HasBaseEnvironment: true, }, }, + { + name: "keeping detached processes records the request", + opts: ClientOptions{ClusterID: "abc-123", KeepDetachedProcesses: true}, + want: protos.SshTunnelEvent{ + ComputeType: protos.SshTunnelComputeTypeDedicated, + ClientMode: protos.SshTunnelClientModeSSH, + KeepDetachedRequested: true, + }, + }, } for _, tt := range tests { diff --git a/experimental/ssh/internal/client/client_test.go b/experimental/ssh/internal/client/client_test.go index 5435bf96130..8363d2e5d0b 100644 --- a/experimental/ssh/internal/client/client_test.go +++ b/experimental/ssh/internal/client/client_test.go @@ -121,6 +121,14 @@ func TestValidate(t *testing.T) { name: "usage policy with connection name", opts: client.ClientOptions{ConnectionName: "my-conn", UsagePolicyID: "pol-1"}, }, + { + name: "keep detached processes with serverless", + opts: client.ClientOptions{ConnectionName: "my-conn", KeepDetachedProcesses: true}, + }, + { + name: "keep detached processes with cluster ID", + opts: client.ClientOptions{ClusterID: "abc-123", KeepDetachedProcesses: true}, + }, } for _, tt := range tests { @@ -334,6 +342,18 @@ func TestToProxyCommand(t *testing.T) { opts: client.ClientOptions{ConnectionName: "my-conn", ShutdownDelay: 2 * time.Minute, MaxClients: 25, ServerTimeout: 48 * time.Hour}, want: quoted + " ssh connect --proxy --name=my-conn --shutdown-delay=2m0s --max-clients=25 --server-timeout=48h0m0s", }, + { + // Carried into the ProxyCommand so a reconnect through ssh asks for the same + // mode, instead of starting a server that would sweep the detached work. + name: "dedicated cluster keeping detached processes", + opts: client.ClientOptions{ClusterID: "abc-123", KeepDetachedProcesses: true, ShutdownDelay: 5 * time.Minute}, + want: quoted + " ssh connect --proxy --cluster=abc-123 --auto-start-cluster=false --shutdown-delay=5m0s --keep-detached-processes", + }, + { + name: "serverless keeping detached processes", + opts: client.ClientOptions{ConnectionName: "my-conn", KeepDetachedProcesses: true, ShutdownDelay: 5 * time.Minute}, + want: quoted + " ssh connect --proxy --name=my-conn --shutdown-delay=5m0s --keep-detached-processes", + }, { name: "with metadata", opts: client.ClientOptions{ClusterID: "abc-123", ServerMetadata: "user,2222,abc-123"}, diff --git a/experimental/ssh/internal/client/policy_internal_test.go b/experimental/ssh/internal/client/policy_internal_test.go index f501f0ba6e6..dde22f46187 100644 --- a/experimental/ssh/internal/client/policy_internal_test.go +++ b/experimental/ssh/internal/client/policy_internal_test.go @@ -24,3 +24,25 @@ func TestUsagePolicyMatches(t *testing.T) { }) } } + +func TestKeepDetachedMatches(t *testing.T) { + tests := []struct { + name string + stored bool + requested bool + want bool + }{ + {name: "no request takes a server that holds the run open", stored: true, requested: false, want: true}, + {name: "no request takes a server that does not", stored: false, requested: false, want: true}, + {name: "request matches a server that holds the run open", stored: true, requested: true, want: true}, + {name: "request against a server that does not does not match", stored: false, requested: true, want: false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := keepDetachedMatches(tt.stored, tt.requested); got != tt.want { + t.Errorf("keepDetachedMatches(%v, %v) = %v, want %v", tt.stored, tt.requested, got, tt.want) + } + }) + } +} diff --git a/experimental/ssh/internal/client/ssh-server-bootstrap.py b/experimental/ssh/internal/client/ssh-server-bootstrap.py index 28a20f73688..a3f20ef0d4d 100644 --- a/experimental/ssh/internal/client/ssh-server-bootstrap.py +++ b/experimental/ssh/internal/client/ssh-server-bootstrap.py @@ -1,4 +1,3 @@ -import atexit import collections import ctypes import ctypes.util @@ -14,6 +13,14 @@ SSH_TUNNEL_BASENAME = "databricks_cli" +# How often the linger loop re-checks for detached processes still holding the run open. +LINGER_POLL_SECONDS = 15 + +# How often the linger loop repeats its "holding the run open" line while the set of survivors +# is unchanged. The hold ends with the work rather than at a deadline, so a line per poll would +# flood the run log for as long as the work lives. +LINGER_REPORT_SECONDS = 300 + # Exit statuses collected by the SIGCHLD subreaper handler, keyed by pid. The handler # can reap the server subprocess before Popen.wait() does, in which case Popen would # report exit code 0; this map preserves the real status. @@ -27,10 +34,14 @@ dbutils.widgets.text("sessionId", "") dbutils.widgets.text("serverless", "false") dbutils.widgets.text("usagePolicyId", "") +dbutils.widgets.text("keepDetachedProcesses", "false") def cleanup(): - subprocess.run(["pkill", "-f", SSH_TUNNEL_BASENAME], check=False) + # Terminate an SSH server left behind by a previous, hard-killed run. The pattern matches + # the server's own argv rather than the CLI binary name alone, so detached work that + # happens to run the CLI is not swept away with it. + subprocess.run(["pkill", "-f", f"{SSH_TUNNEL_BASENAME}.*ssh server --cluster="], check=False) def setup_subreaper(): @@ -76,9 +87,94 @@ def kill_all_children(): print(f"Error while killing child processes: {e}") -def setup_exit_handler(): - # Register the cleanup function to be called when the script exits - atexit.register(kill_all_children) +def kill_server_group(server_pgid): + """Terminate the SSH server's own process group. + + That group holds exactly what the tunnel started: the server and the sshd processes it + spawned per connection. A process that deliberately left the group - which is what tmux, + setsid and disown do - is not in it, so detached work survives this. Killing by parentage + instead (pkill -P) would sweep those too, because PR_SET_CHILD_SUBREAPER makes this + process adopt every orphan in the session. + """ + try: + os.killpg(server_pgid, signal.SIGTERM) + print(f"Terminated SSH server process group {server_pgid}") + except ProcessLookupError: + print(f"SSH server process group {server_pgid} is already gone") + + +def detached_descendants(server_pgid): + """Adopted children of this process that are outside the SSH server's process group. + + PR_SET_CHILD_SUBREAPER makes every orphan in the session reparent to this process, so + work that detached itself - tmux, setsid, disown, a plain background command - resurfaces + here as a direct child. The server's own sshd children stay in its group and are excluded. + + Reads /proc directly, mirroring detachedDescendants in internal/server/descendants.go. + Asking ps for this process's children cannot work: ps is one of them, and subprocess.run + leaves it in this process's group, so it matches its own query on every poll and the + survivor list is never empty. + """ + self_pid = os.getpid() + survivors = [] + for entry in os.listdir("/proc"): + if not entry.isdigit(): + continue + try: + with open(f"/proc/{entry}/stat") as stat_file: + # Split on the last ')' rather than from the left: the comm field before it + # can contain both spaces and parentheses. state, ppid and pgrp follow it. + state, ppid, pgrp = stat_file.read().rsplit(")", 1)[1].split()[:3] + except OSError: + # The process exited while we were walking /proc. + continue + if ppid != str(self_pid) or pgrp == str(server_pgid): + continue + # Zombies are already dead; the SIGCHLD handler collects them. + if state.startswith("Z"): + continue + survivors.append(entry) + return sorted(survivors, key=int) + + +def has_children(): + """Probe without reaping; WNOHANG returning None still means children exist.""" + try: + os.waitid(os.P_ALL, 0, os.WEXITED | os.WNOHANG | os.WNOWAIT) + except ChildProcessError: + return False + return True + + +def wait_for_detached_descendants(server_pgid): + """Hold the notebook open while detached work is still running. + + WSFS authorises an I/O by walking the live process tree for a registered ancestor, and + this process is the registered one. Returning while detached work is still alive would + reparent it to PID 1, outside the registered subtree, and silently strip its /Workspace + and /Volumes access - trading a visible failure for an invisible one. Holding the run + open instead keeps the cluster from auto-terminating, which is why it is opt-in. + + The work decides how long this takes. The only bound is the run's own timeout + (--server-timeout), which Jobs enforces and the client requires to be set, so this loop + cannot hold a cluster indefinitely. + """ + reported = None + reported_at = 0.0 + while True: + survivors = detached_descendants(server_pgid) + if not survivors and not has_children(): + print("No detached processes left, releasing the run", flush=True) + return + now = time.monotonic() + if survivors != reported or now - reported_at >= LINGER_REPORT_SECONDS: + print( + f"Holding the run open for {len(survivors)} detached process(es): {','.join(survivors)}", + flush=True, + ) + reported = survivors + reported_at = now + time.sleep(LINGER_POLL_SECONDS) def run_ssh_server(): @@ -128,6 +224,7 @@ def run_ssh_server(): raise RuntimeError("Session ID is required. Please provide it using the 'sessionId' widget.") serverless = dbutils.widgets.get("serverless") usage_policy_id = dbutils.widgets.get("usagePolicyId") + keep_detached_processes = dbutils.widgets.get("keepDetachedProcesses") == "true" # Mark this process's WSFS command origin so workspace-file activity from the # remote SSH session is attributable @@ -178,11 +275,27 @@ def run_ssh_server(): if usage_policy_id: server_args.append(f"--usage-policy-id={usage_policy_id}") + # The server uses this to defer idle shutdown while detached work is running. + if keep_detached_processes: + server_args.append("--keep-detached-processes") + # Tee the server output instead of inheriting stdout: the run-page logs remain the only # place to debug a RUNNING server, but on failure we attach the log tail to the exception # so "ssh connect" can print it (the Jobs run-output API has no stdout logs for notebook tasks). tail = collections.deque(maxlen=100) - proc = subprocess.Popen(server_args, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, errors="replace") + proc = subprocess.Popen( + server_args, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + errors="replace", + # Make the server a session and process group leader, so teardown can target exactly + # the processes the tunnel started. See kill_server_group. + start_new_session=True, + ) + # The server leads the new group, so the group id is its pid. Recorded here because the + # pid may already have been reaped by the time we tear the group down. + server_pgid = proc.pid try: for line in proc.stdout: # flush so the run-page logs stay live while the server is running @@ -192,15 +305,29 @@ def run_ssh_server(): # The SIGCHLD subreaper handler may have collected the server first; Popen reports that as 0. if proc.pid in reaped_statuses: returncode = os.waitstatus_to_exitcode(reaped_statuses[proc.pid]) - if returncode != 0: + if returncode == -signal.SIGTERM: + # A newer session's bootstrap terminates a server already running on this cluster + # (see cleanup), which a reconnect asking for --keep-detached-processes against a + # server without it now reaches on a normal path. That is a handover, not a failure + # of this run, so it must not mark the run FAILED - the hold below still runs. + print("SSH server was terminated, most likely by a newer session on this cluster", flush=True) + elif returncode != 0: # The tail size matches maxRunFailureTraceBytes, the cap the client prints to the terminal. raise RuntimeError(f"SSH server exited with code {returncode}. Last server logs:\n" + "".join(tail)[-2000:]) finally: - kill_all_children() + # Always reap the server and the sshd children it spawned; they are the only things + # in its process group. What happens to work that left that group depends on the mode: + # keep it and hold the run open as its WSFS anchor, or sweep it as we always have. + # Narrowing the sweep without holding the run open would leave survivors alive but cut + # off from /Workspace, which is a worse failure than the one it fixes. + kill_server_group(server_pgid) + if keep_detached_processes: + wait_for_detached_descendants(server_pgid) + else: + kill_all_children() if __name__ == "__main__": cleanup() setup_subreaper() - setup_exit_handler() run_ssh_server() diff --git a/experimental/ssh/internal/client/ssh_server_bootstrap_test.go b/experimental/ssh/internal/client/ssh_server_bootstrap_test.go new file mode 100644 index 00000000000..fe29b12aff2 --- /dev/null +++ b/experimental/ssh/internal/client/ssh_server_bootstrap_test.go @@ -0,0 +1,17 @@ +package client_test + +import ( + "os/exec" + "testing" + + "github.com/databricks/cli/libs/python" + "github.com/stretchr/testify/require" +) + +func TestSSHServerBootstrap(test *testing.T) { + test.Parallel() + + cmd := exec.CommandContext(test.Context(), python.GetExecutable(), "testdata/ssh_server_bootstrap_test.py") + output, err := cmd.CombinedOutput() + require.NoError(test, err, "%s", output) +} diff --git a/experimental/ssh/internal/client/submit_internal_test.go b/experimental/ssh/internal/client/submit_internal_test.go index fa1bd5b2dd1..d4247cbc2e1 100644 --- a/experimental/ssh/internal/client/submit_internal_test.go +++ b/experimental/ssh/internal/client/submit_internal_test.go @@ -72,6 +72,22 @@ func TestBuildSSHServerSubmitRun(t *testing.T) { assert.Equal(t, "abc-123", got.Tasks[0].ExistingClusterId) assert.Empty(t, got.Tasks[0].EnvironmentKey) assert.Empty(t, got.Environments) + // False is what tells the bootstrap to sweep detached work as it always has. + assert.Equal(t, "false", got.Tasks[0].NotebookTask.BaseParameters["keepDetachedProcesses"]) + }) + + t.Run("dedicated cluster keeping detached processes", func(t *testing.T) { + opts := ClientOptions{ + ClusterID: "abc-123", + ServerTimeout: 24 * time.Hour, + KeepDetachedProcesses: true, + } + got := buildSSHServerSubmitRun("v1", "scope", notebookPath, "", opts) + + assert.Equal(t, "true", got.Tasks[0].NotebookTask.BaseParameters["keepDetachedProcesses"]) + // The hold happens inside the run, so the run's own timeout is what bounds it. + assert.Equal(t, int(24*time.Hour.Seconds()), got.TimeoutSeconds) + assert.Equal(t, int(24*time.Hour.Seconds()), got.Tasks[0].TimeoutSeconds) }) t.Run("server lifecycle", func(t *testing.T) { diff --git a/experimental/ssh/internal/client/testdata/ssh_server_bootstrap_test.py b/experimental/ssh/internal/client/testdata/ssh_server_bootstrap_test.py new file mode 100644 index 00000000000..c3ebf8f728f --- /dev/null +++ b/experimental/ssh/internal/client/testdata/ssh_server_bootstrap_test.py @@ -0,0 +1,63 @@ +#!/usr/bin/env python3 +"""Test notebook linger helpers without importing Databricks runtime dependencies.""" + +import ast +import unittest +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import Mock, call + + +class LingerTest(unittest.TestCase): + def test_wait_for_detached_descendants(self): + source = Path(__file__).resolve().parents[1] / "ssh-server-bootstrap.py" + module = ast.parse(source.read_text()) + module.body = [ + node + for node in module.body + if isinstance(node, ast.FunctionDef) and node.name in {"has_children", "wait_for_detached_descendants"} + ] + code = compile(module, str(source), "exec") + + cases = { + "no children": ([[]], [ChildProcessError()], [], 0), + "child awaiting adoption": ([[], ["42"], []], [None, ChildProcessError()], [0, 1], 2), + "exited child awaiting reaping": ([[], []], [object(), ChildProcessError()], [0], 1), + "adopted child": ([["42"], []], [ChildProcessError()], [0], 1), + "pending adoption spans multiple polls": ( + [[], [], ["42"], []], + [None, None, ChildProcessError()], + [0, 1, 2], + 2, + ), + "pending adoption report interval": ([[], [], []], [None, None, ChildProcessError()], [0, 300], 2), + } + for name, (survivors, children, timestamps, report_count) in cases.items(): + with self.subTest(name=name): + mock_os = SimpleNamespace(P_ALL=0, WEXITED=1, WNOHANG=2, WNOWAIT=4, waitid=Mock(side_effect=children)) + mock_time = SimpleNamespace(monotonic=Mock(side_effect=timestamps), sleep=Mock()) + namespace = { + "os": mock_os, + "time": mock_time, + "detached_descendants": Mock(side_effect=survivors), + "LINGER_POLL_SECONDS": 1, + "LINGER_REPORT_SECONDS": 300, + "print": Mock(), + } + exec(code, namespace) + + namespace["wait_for_detached_descendants"](123) + + self.assertEqual(namespace["detached_descendants"].call_args_list, [call(123)] * len(survivors)) + self.assertEqual( + mock_os.waitid.call_args_list, + [call(mock_os.P_ALL, 0, mock_os.WEXITED | mock_os.WNOHANG | mock_os.WNOWAIT)] * len(children), + ) + self.assertEqual(mock_time.sleep.call_args_list, [call(1)] * (len(survivors) - 1)) + self.assertEqual(mock_time.monotonic.call_count, len(timestamps)) + self.assertEqual(namespace["print"].call_count, report_count + 1) + self.assertIn("No detached processes left", namespace["print"].call_args[0][0]) + + +if __name__ == "__main__": + unittest.main() diff --git a/experimental/ssh/internal/proxy/connections.go b/experimental/ssh/internal/proxy/connections.go index 96194fa430e..3e3991571dc 100644 --- a/experimental/ssh/internal/proxy/connections.go +++ b/experimental/ssh/internal/proxy/connections.go @@ -24,7 +24,7 @@ func NewConnectionsManager(maxClients int, shutdownDelay time.Duration) *Connect connections: make(map[string]*proxyConnection), TimedOut: make(chan bool), } - cm.startShutdownTimer() + cm.startShutdownTimer(shutdownDelay) return cm } @@ -61,7 +61,7 @@ func (cm *ConnectionsManager) Remove(id string) { cm.removeConnection(id) count := cm.Count() if count <= 0 { - cm.startShutdownTimer() + cm.startShutdownTimer(cm.shutdownDelay) } } @@ -71,13 +71,21 @@ func (cm *ConnectionsManager) removeConnection(id string) { delete(cm.connections, id) } -func (cm *ConnectionsManager) startShutdownTimer() { +func (cm *ConnectionsManager) ExtendIdleTimeout(delay time.Duration) { + cm.connectionsMu.Lock() + defer cm.connectionsMu.Unlock() + if len(cm.connections) == 0 { + cm.startShutdownTimer(delay) + } +} + +func (cm *ConnectionsManager) startShutdownTimer(delay time.Duration) { cm.shutdownTimerMu.Lock() defer cm.shutdownTimerMu.Unlock() if cm.shutdownTimer != nil { cm.shutdownTimer.Stop() } - cm.shutdownTimer = time.AfterFunc(cm.shutdownDelay, func() { + cm.shutdownTimer = time.AfterFunc(delay, func() { cm.TimedOut <- true }) } diff --git a/experimental/ssh/internal/proxy/connections_idle_test.go b/experimental/ssh/internal/proxy/connections_idle_test.go new file mode 100644 index 00000000000..6c1c0b7c122 --- /dev/null +++ b/experimental/ssh/internal/proxy/connections_idle_test.go @@ -0,0 +1,29 @@ +package proxy_test + +import ( + "testing" + "testing/synctest" + "time" + + "github.com/databricks/cli/experimental/ssh/internal/proxy" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestConnectionsManagerExtendIdleTimeout(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + connections := proxy.NewConnectionsManager(1, time.Minute) + connections.ExtendIdleTimeout(time.Second) + time.Sleep(time.Second) + assert.True(t, <-connections.TimedOut) + require.True(t, connections.TryAdd("connected", nil)) + connections.ExtendIdleTimeout(time.Second) + time.Sleep(time.Minute) + synctest.Wait() + select { + case <-connections.TimedOut: + t.Fatal("extended timeout fired while a client was connected") + default: + } + }) +} diff --git a/experimental/ssh/internal/server/descendants.go b/experimental/ssh/internal/server/descendants.go new file mode 100644 index 00000000000..89163f56465 --- /dev/null +++ b/experimental/ssh/internal/server/descendants.go @@ -0,0 +1,112 @@ +package server + +import ( + "fmt" + "os" + "path/filepath" + "slices" + "strconv" + "strings" +) + +// procRoot is the procfs mount point. The functions below take the root as an argument +// so tests can run against a fixture tree instead of the live one. +const procRoot = "/proc" + +// procStat is the part of /proc//stat this package reads. +type procStat struct { + state string + ppid int + pgrp int +} + +// parseProcStat reads state, ppid and pgrp out of the contents of /proc//stat. +func parseProcStat(content string) (procStat, error) { + // Anchor on the last ')' instead of splitting from the left: the second field is the + // executable name, and it can contain both spaces and parentheses. + comm := strings.LastIndex(content, ")") + if comm < 0 { + return procStat{}, fmt.Errorf("no comm field in %q", content) + } + // state, ppid and pgrp are the three fields that follow comm. + fields := strings.Fields(content[comm+1:]) + if len(fields) < 3 { + return procStat{}, fmt.Errorf("expected state, ppid and pgrp after comm in %q", content) + } + ppid, err := strconv.Atoi(fields[1]) + if err != nil { + return procStat{}, fmt.Errorf("failed to parse ppid in %q: %w", content, err) + } + pgrp, err := strconv.Atoi(fields[2]) + if err != nil { + return procStat{}, fmt.Errorf("failed to parse pgrp in %q: %w", content, err) + } + return procStat{state: fields[0], ppid: ppid, pgrp: pgrp}, nil +} + +func readProcStat(root string, pid int) (procStat, error) { + content, err := os.ReadFile(filepath.Join(root, strconv.Itoa(pid), "stat")) + if err != nil { + return procStat{}, err + } + return parseProcStat(string(content)) +} + +// detachedDescendants returns the pids of processes the SSH session started that left +// the server's process group - what tmux, setsid, disown and a plain background command +// all do. +// +// They surface as siblings of the server: the bootstrap notebook sets +// PR_SET_CHILD_SUBREAPER, so a process that orphans itself is reparented to the notebook +// rather than to PID 1. Taking the notebook's children and excluding the server's own +// process group therefore leaves exactly the detached work - the server's sshd children +// stay in its group, and the notebook starts nothing else. +func detachedDescendants(root string, selfPid int) ([]int, error) { + self, err := readProcStat(root, selfPid) + if err != nil { + return nil, fmt.Errorf("failed to read own process stat: %w", err) + } + // The notebook is gone and the server has been reparented to PID 1: there is no + // anchor left to enumerate against, and nothing left to keep alive. + if self.ppid <= 1 { + return nil, nil + } + + entries, err := os.ReadDir(root) + if err != nil { + return nil, fmt.Errorf("failed to read %s: %w", root, err) + } + + var pids []int + for _, entry := range entries { + pid, err := strconv.Atoi(entry.Name()) + if err != nil { + // Not a process directory. + continue + } + stat, err := readProcStat(root, pid) + if err != nil { + // The process exited while we were walking, or its stat is unreadable. + continue + } + if stat.ppid != self.ppid || stat.pgrp == self.pgrp { + continue + } + // Zombies are already dead; the notebook's SIGCHLD handler collects them. + if strings.HasPrefix(stat.state, "Z") { + continue + } + pids = append(pids, pid) + } + slices.Sort(pids) + return pids, nil +} + +// formatPids renders pids for a log line, e.g. "1234, 1235". +func formatPids(pids []int) string { + parts := make([]string, len(pids)) + for i, pid := range pids { + parts[i] = strconv.Itoa(pid) + } + return strings.Join(parts, ", ") +} diff --git a/experimental/ssh/internal/server/descendants_test.go b/experimental/ssh/internal/server/descendants_test.go new file mode 100644 index 00000000000..6fe22fb50fe --- /dev/null +++ b/experimental/ssh/internal/server/descendants_test.go @@ -0,0 +1,161 @@ +package server + +import ( + "os" + "path/filepath" + "runtime" + "strconv" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// fakeProc writes a /proc-like tree. Each process is described by its ppid, pgrp and state. +type fakeProcess struct { + comm string + ppid int + pgrp int + state string +} + +func fakeProc(t *testing.T, processes map[int]fakeProcess) string { + t.Helper() + root := t.TempDir() + for pid, p := range processes { + dir := filepath.Join(root, strconv.Itoa(pid)) + require.NoError(t, os.MkdirAll(dir, 0o755)) + // The real format has 50+ fields; only the four leading ones are read. + line := strconv.Itoa(pid) + " (" + p.comm + ") " + p.state + " " + + strconv.Itoa(p.ppid) + " " + strconv.Itoa(p.pgrp) + " 0 0 -1 4194304\n" + require.NoError(t, os.WriteFile(filepath.Join(dir, "stat"), []byte(line), 0o644)) + } + return root +} + +func TestParseProcStat(t *testing.T) { + t.Run("parses state, ppid and pgrp", func(t *testing.T) { + got, err := parseProcStat("4242 (databricks) S 4200 4242 4242 0 -1 4194304 1 0") + require.NoError(t, err) + assert.Equal(t, procStat{state: "S", ppid: 4200, pgrp: 4242}, got) + }) + + // The comm field is unquoted and may contain spaces and parentheses, so the fields + // after it can only be located from the last ')'. + t.Run("handles a comm with spaces and parentheses", func(t *testing.T) { + got, err := parseProcStat("7 (weird ) name) R 3 9 9 0 -1 0") + require.NoError(t, err) + assert.Equal(t, procStat{state: "R", ppid: 3, pgrp: 9}, got) + }) + + t.Run("rejects a line without comm", func(t *testing.T) { + _, err := parseProcStat("4242 S 4200 4242") + assert.ErrorContains(t, err, "no comm field") + }) + + t.Run("rejects a truncated line", func(t *testing.T) { + _, err := parseProcStat("4242 (databricks) S 4200") + assert.ErrorContains(t, err, "expected state, ppid and pgrp") + }) +} + +func TestDetachedDescendants(t *testing.T) { + // The shape the server sees at teardown: the notebook (100) is the subreaper, the + // server (200) leads its own group, sshd (300) is in the server's group, and the + // detached work (400, 500) has been reparented onto the notebook with its own groups. + const notebook, server, sshd = 100, 200, 300 + + t.Run("returns detached work only", func(t *testing.T) { + root := fakeProc(t, map[int]fakeProcess{ + notebook: {comm: "python", ppid: 1, pgrp: notebook, state: "S"}, + server: {comm: "databricks", ppid: notebook, pgrp: server, state: "S"}, + sshd: {comm: "sshd", ppid: server, pgrp: server, state: "S"}, + 400: {comm: "tmux: server", ppid: notebook, pgrp: 400, state: "S"}, + 500: {comm: "train.py", ppid: notebook, pgrp: 500, state: "R"}, + }) + + pids, err := detachedDescendants(root, server) + require.NoError(t, err) + assert.Equal(t, []int{400, 500}, pids) + }) + + t.Run("excludes the server's own process group", func(t *testing.T) { + // A login shell that sshd put in its own group, but that is still parented by + // sshd rather than adopted by the notebook, is not detached work. + root := fakeProc(t, map[int]fakeProcess{ + notebook: {comm: "python", ppid: 1, pgrp: notebook, state: "S"}, + server: {comm: "databricks", ppid: notebook, pgrp: server, state: "S"}, + sshd: {comm: "sshd", ppid: server, pgrp: server, state: "S"}, + 400: {comm: "bash", ppid: sshd, pgrp: 400, state: "S"}, + }) + + pids, err := detachedDescendants(root, server) + require.NoError(t, err) + assert.Empty(t, pids) + }) + + t.Run("skips zombies", func(t *testing.T) { + root := fakeProc(t, map[int]fakeProcess{ + notebook: {comm: "python", ppid: 1, pgrp: notebook, state: "S"}, + server: {comm: "databricks", ppid: notebook, pgrp: server, state: "S"}, + 400: {comm: "gone", ppid: notebook, pgrp: 400, state: "Z"}, + }) + + pids, err := detachedDescendants(root, server) + require.NoError(t, err) + assert.Empty(t, pids) + }) + + // The notebook died first, so the server was reparented to PID 1. Nothing is anchored + // any more and there is no sibling set to enumerate. + t.Run("returns nothing once the notebook is gone", func(t *testing.T) { + root := fakeProc(t, map[int]fakeProcess{ + 1: {comm: "systemd", ppid: 0, pgrp: 1, state: "S"}, + server: {comm: "databricks", ppid: 1, pgrp: server, state: "S"}, + 400: {comm: "tmux: server", ppid: 1, pgrp: 400, state: "S"}, + }) + + pids, err := detachedDescendants(root, server) + require.NoError(t, err) + assert.Empty(t, pids) + }) + + t.Run("ignores unreadable and non-process entries", func(t *testing.T) { + root := fakeProc(t, map[int]fakeProcess{ + notebook: {comm: "python", ppid: 1, pgrp: notebook, state: "S"}, + server: {comm: "databricks", ppid: notebook, pgrp: server, state: "S"}, + 400: {comm: "tmux: server", ppid: notebook, pgrp: 400, state: "S"}, + }) + require.NoError(t, os.MkdirAll(filepath.Join(root, "self"), 0o755)) + // A process that exited between the readdir and the stat read. + require.NoError(t, os.MkdirAll(filepath.Join(root, "999"), 0o755)) + + pids, err := detachedDescendants(root, server) + require.NoError(t, err) + assert.Equal(t, []int{400}, pids) + }) + + t.Run("fails when own stat is missing", func(t *testing.T) { + _, err := detachedDescendants(t.TempDir(), server) + assert.ErrorContains(t, err, "failed to read own process stat") + }) +} + +// The fixtures above are hand-written, so this pins the field offsets against a real +// /proc//stat. The server only ever runs on Linux compute. +func TestParseProcStatAgainstRealProcfs(t *testing.T) { + if runtime.GOOS != "linux" { + t.Skip("procfs is Linux-only") + } + + got, err := readProcStat(procRoot, os.Getpid()) + require.NoError(t, err) + assert.Equal(t, os.Getppid(), got.ppid) + assert.NotZero(t, got.pgrp) + assert.NotEmpty(t, got.state) +} + +func TestFormatPids(t *testing.T) { + assert.Equal(t, "1, 22, 333", formatPids([]int{1, 22, 333})) + assert.Empty(t, formatPids(nil)) +} diff --git a/experimental/ssh/internal/server/export_test.go b/experimental/ssh/internal/server/export_test.go index b5d6290b9df..7313c43c446 100644 --- a/experimental/ssh/internal/server/export_test.go +++ b/experimental/ssh/internal/server/export_test.go @@ -1,6 +1,13 @@ package server var ( - WorkspaceToken = workspaceToken - FuseUserInfo = fuseUserInfo + WorkspaceToken = workspaceToken + FuseUserInfo = fuseUserInfo + WaitForIdleShutdown = waitForIdleShutdown + ProcWithDetachedWork = procWithDetachedWork +) + +const ( + DetachedProcessCheckInterval = detachedProcessCheckInterval + TestServerPid = testServerPid ) diff --git a/experimental/ssh/internal/server/idle.go b/experimental/ssh/internal/server/idle.go new file mode 100644 index 00000000000..2935d63e452 --- /dev/null +++ b/experimental/ssh/internal/server/idle.go @@ -0,0 +1,37 @@ +package server + +import ( + "context" + "time" + + "github.com/databricks/cli/experimental/ssh/internal/proxy" + "github.com/databricks/cli/libs/log" +) + +const detachedProcessCheckInterval = 15 * time.Second + +func waitForIdleShutdown(ctx context.Context, connections *proxy.ConnectionsManager, keepDetached bool, root string, selfPid int) error { + for { + select { + case <-ctx.Done(): + return ctx.Err() + case <-connections.TimedOut: + } + if connections.Count() > 0 { + continue + } + if !keepDetached { + return nil + } + pids, err := detachedDescendants(root, selfPid) + if err == nil && len(pids) == 0 { + return nil + } + if err != nil { + log.Warnf(ctx, "Cannot check detached processes; postponing SSH idle shutdown: %v", err) + } else { + log.Infof(ctx, "Keeping SSH server running for %d detached process(es): %s", len(pids), formatPids(pids)) + } + connections.ExtendIdleTimeout(detachedProcessCheckInterval) + } +} diff --git a/experimental/ssh/internal/server/idle_test.go b/experimental/ssh/internal/server/idle_test.go new file mode 100644 index 00000000000..591ff884372 --- /dev/null +++ b/experimental/ssh/internal/server/idle_test.go @@ -0,0 +1,121 @@ +package server_test + +import ( + "context" + "os" + "path/filepath" + "testing" + "testing/synctest" + "time" + + "github.com/databricks/cli/experimental/ssh/internal/proxy" + "github.com/databricks/cli/experimental/ssh/internal/server" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +const detachedProcessDirectory = "400" + +func TestWaitForIdleShutdown(t *testing.T) { + for _, test := range []struct { + name string + keepDetached bool + hasDetached bool + unreadable bool + wantHold bool + }{ + {name: "flag off with detached work", hasDetached: true}, + {name: "flag off without detached work"}, + {name: "flag on without detached work", keepDetached: true}, + {name: "flag on with detached work", keepDetached: true, hasDetached: true, wantHold: true}, + {name: "flag on with unreadable process tree", keepDetached: true, unreadable: true, wantHold: true}, + {name: "flag off does not need process tree", unreadable: true}, + } { + t.Run(test.name, func(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + root := server.ProcWithDetachedWork(t) + if !test.hasDetached { + require.NoError(t, os.RemoveAll(filepath.Join(root, detachedProcessDirectory))) + } + if test.unreadable { + root = t.TempDir() + } + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + connections := proxy.NewConnectionsManager(1, time.Second) + shutdown := make(chan error, 1) + go func() { + shutdown <- server.WaitForIdleShutdown(ctx, connections, test.keepDetached, root, server.TestServerPid) + }() + time.Sleep(time.Second + 2*server.DetachedProcessCheckInterval) + synctest.Wait() + assert.Equal(t, !test.wantHold, len(shutdown) > 0) + require.True(t, connections.TryAdd("cleanup", nil)) + cancel() + synctest.Wait() + require.Len(t, shutdown, 1) + shutdownErr := <-shutdown + if test.wantHold { + assert.ErrorIs(t, shutdownErr, context.Canceled) + } else { + assert.NoError(t, shutdownErr) + } + }) + }) + } +} + +func TestWaitForIdleShutdownWhenDetachedWorkFinishes(t *testing.T) { + for _, shutdownDelay := range []time.Duration{0, time.Second} { + t.Run(shutdownDelay.String(), func(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + root := server.ProcWithDetachedWork(t) + connections := proxy.NewConnectionsManager(1, shutdownDelay) + shutdown := make(chan error, 1) + go func() { + shutdown <- server.WaitForIdleShutdown(t.Context(), connections, true, root, server.TestServerPid) + }() + time.Sleep(shutdownDelay) + synctest.Wait() + require.Empty(t, shutdown) + require.NoError(t, os.RemoveAll(filepath.Join(root, detachedProcessDirectory))) + time.Sleep(server.DetachedProcessCheckInterval) + synctest.Wait() + require.Len(t, shutdown, 1) + assert.NoError(t, <-shutdown) + }) + }) + } +} + +func TestWaitForIdleShutdownReconnect(t *testing.T) { + for _, connectedFor := range []time.Duration{time.Second, time.Minute} { + t.Run(connectedFor.String(), func(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + root := server.ProcWithDetachedWork(t) + shutdownDelay := time.Minute + connections := proxy.NewConnectionsManager(1, shutdownDelay) + shutdown := make(chan error, 1) + go func() { + shutdown <- server.WaitForIdleShutdown(t.Context(), connections, true, root, server.TestServerPid) + }() + time.Sleep(shutdownDelay) + synctest.Wait() + require.Empty(t, shutdown) + require.True(t, connections.TryAdd("reconnected", nil)) + require.NoError(t, os.RemoveAll(filepath.Join(root, detachedProcessDirectory))) + time.Sleep(connectedFor) + synctest.Wait() + require.Empty(t, shutdown) + connections.Remove("reconnected") + time.Sleep(shutdownDelay - time.Second) + synctest.Wait() + require.Empty(t, shutdown) + time.Sleep(time.Second) + synctest.Wait() + require.Len(t, shutdown, 1) + assert.NoError(t, <-shutdown) + }) + }) + } +} diff --git a/experimental/ssh/internal/server/server.go b/experimental/ssh/internal/server/server.go index e5454a60ed4..7ee1797d237 100644 --- a/experimental/ssh/internal/server/server.go +++ b/experimental/ssh/internal/server/server.go @@ -22,6 +22,8 @@ import ( "github.com/databricks/cli/experimental/ssh/internal/workspace" "github.com/databricks/cli/libs/env" "github.com/databricks/cli/libs/log" + "github.com/databricks/cli/libs/telemetry" + "github.com/databricks/cli/libs/telemetry/protos" "github.com/databricks/databricks-sdk-go" ) @@ -45,6 +47,8 @@ type ServerOptions struct { // UsagePolicyID the job was submitted with. Persisted to metadata.json so reconnects // can tell which usage policy the running server was started under. UsagePolicyID string + // KeepDetachedProcesses prevents idle shutdown while detached processes are running. + KeepDetachedProcesses bool // The directory to store sshd configuration ConfigDir string // The name of the secrets scope to use for client and server keys @@ -81,9 +85,10 @@ func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ServerOpt // Save metadata including ClusterID (required for Driver Proxy connections in serverless mode) metadata := &workspace.WorkspaceMetadata{ - Port: port, - ClusterID: opts.ClusterID, - UsagePolicyID: opts.UsagePolicyID, + Port: port, + ClusterID: opts.ClusterID, + UsagePolicyID: opts.UsagePolicyID, + KeepDetachedProcesses: opts.KeepDetachedProcesses, } err = workspace.SaveWorkspaceMetadata(ctx, client, opts.Version, opts.SessionID, metadata) if err != nil { @@ -122,13 +127,66 @@ func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ServerOpt http.HandleFunc("/driver-proxy-http/logs", logBuf.serveHTTP) http.HandleFunc("/driver-proxy-http/capabilities", serveCapabilities) - go handleTimeout(ctx, connections.TimedOut, opts.ShutdownDelay) + listenErr := make(chan error, 1) + go func() { + listenErr <- http.ListenAndServe(listenAddr, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // Normalize double slashes from the driver proxy (e.g. //metadata -> /metadata) + r.URL.Path = path.Clean(r.URL.Path) + http.DefaultServeMux.ServeHTTP(w, r) + })) + }() - return http.ListenAndServe(listenAddr, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - // Normalize double slashes from the driver proxy (e.g. //metadata -> /metadata) - r.URL.Path = path.Clean(r.URL.Path) - http.DefaultServeMux.ServeHTTP(w, r) - })) + idleErr := make(chan error, 1) + go func() { + idleErr <- waitForIdleShutdown(ctx, connections, opts.KeepDetachedProcesses, procRoot, os.Getpid()) + }() + + select { + case err := <-listenErr: + return err + case err := <-idleErr: + if err != nil { + return err + } + // Return rather than exiting in place, so the notebook that started us gets to run + // its teardown and this process reports the shutdown through the CLI's normal path. + log.Info(ctx, fmt.Sprintf("No SSH clients for %v, shutting down...", opts.ShutdownDelay)) + reportDetachedDescendants(ctx, opts, procRoot, os.Getpid()) + return nil + } +} + +// reportDetachedDescendants records what the session leaves behind when the server shuts +// down. Only the server can see this: the client that started the session is long gone by +// the time the idle timer fires, and the notebook's teardown runs after this process exits. +func reportDetachedDescendants(ctx context.Context, opts ServerOptions, root string, selfPid int) { + pids, err := detachedDescendants(root, selfPid) + if err != nil { + log.Debugf(ctx, "Failed to look for detached processes: %v", err) + return + } + + // Warning, not info: without --keep-detached-processes these processes do not outlive the + // run, and until now they vanished with no explanation anywhere. The client reads this + // back through /logs. + if len(pids) > 0 && !opts.KeepDetachedProcesses { + log.Warnf(ctx, "Shutting down with %d detached process(es) still running (pids %s). "+ + "They do not survive the end of this run. To keep them, reconnect with "+ + "\"databricks ssh connect --keep-detached-processes\", which keeps the SSH server running while they run.", + len(pids), formatPids(pids)) + } + + computeType := protos.SshTunnelComputeTypeDedicated + if opts.Serverless { + computeType = protos.SshTunnelComputeTypeServerless + } + telemetry.Log(ctx, protos.DatabricksCliLog{ + SshTunnelTeardownEvent: &protos.SshTunnelTeardownEvent{ + ComputeType: computeType, + KeepDetachedRequested: opts.KeepDetachedProcesses, + HadDetachedDescendantsAtTeardown: len(pids) > 0, + }, + }) } // serveCapabilities tells the client which optional parts of the tunnel protocol this server @@ -153,12 +211,6 @@ func serveMetadata(w http.ResponseWriter, r *http.Request) { } } -func handleTimeout(ctx context.Context, timedOutChannel chan bool, shutdownDelay time.Duration) { - <-timedOutChannel - log.Info(ctx, fmt.Sprintf("No SSH clients for %v, shutting down...", shutdownDelay)) - os.Exit(0) -} - func findAvailablePort(startPort, maxAttempts int) (int, error) { for i := range maxAttempts { port := startPort + i diff --git a/experimental/ssh/internal/server/teardown_test.go b/experimental/ssh/internal/server/teardown_test.go new file mode 100644 index 00000000000..6e61b364ebf --- /dev/null +++ b/experimental/ssh/internal/server/teardown_test.go @@ -0,0 +1,141 @@ +package server + +import ( + "encoding/json" + "testing" + + "github.com/databricks/cli/libs/cmdctx" + "github.com/databricks/cli/libs/telemetry" + "github.com/databricks/cli/libs/telemetry/protos" + "github.com/databricks/cli/libs/testserver" + "github.com/databricks/databricks-sdk-go/config" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +const ( + testNotebookPid = 100 + testServerPid = 200 +) + +// procWithDetachedWork is the teardown shape the warning is about: the notebook anchors the +// server, and one process has been adopted by the notebook with a process group of its own. +func procWithDetachedWork(t *testing.T) string { + return fakeProc(t, map[int]fakeProcess{ + testNotebookPid: {comm: "python", ppid: 1, pgrp: testNotebookPid, state: "S"}, + testServerPid: {comm: "databricks", ppid: testNotebookPid, pgrp: testServerPid, state: "S"}, + 400: {comm: "tmux: server", ppid: testNotebookPid, pgrp: 400, state: "S"}, + }) +} + +func TestReportDetachedDescendantsWarning(t *testing.T) { + t.Run("warns and names the flag when the work is about to be swept", func(t *testing.T) { + ctx, logs := captureWarnLogs(t.Context()) + reportDetachedDescendants(ctx, ServerOptions{}, procWithDetachedWork(t), testServerPid) + + assert.Contains(t, logs.String(), "1 detached process(es) still running (pids 400)") + assert.Contains(t, logs.String(), "--keep-detached-processes") + }) + + t.Run("stays quiet when the run is held open for them", func(t *testing.T) { + ctx, logs := captureWarnLogs(t.Context()) + opts := ServerOptions{KeepDetachedProcesses: true} + reportDetachedDescendants(ctx, opts, procWithDetachedWork(t), testServerPid) + + assert.Empty(t, logs.String()) + }) + + t.Run("warns on serverless", func(t *testing.T) { + ctx, logs := captureWarnLogs(t.Context()) + reportDetachedDescendants(ctx, ServerOptions{Serverless: true}, procWithDetachedWork(t), testServerPid) + + assert.Contains(t, logs.String(), "--keep-detached-processes") + }) + + t.Run("stays quiet when nothing was left behind", func(t *testing.T) { + root := fakeProc(t, map[int]fakeProcess{ + testNotebookPid: {comm: "python", ppid: 1, pgrp: testNotebookPid, state: "S"}, + testServerPid: {comm: "databricks", ppid: testNotebookPid, pgrp: testServerPid, state: "S"}, + }) + ctx, logs := captureWarnLogs(t.Context()) + reportDetachedDescendants(ctx, ServerOptions{}, root, testServerPid) + + assert.Empty(t, logs.String()) + }) + + t.Run("does not warn when the process tree cannot be read", func(t *testing.T) { + ctx, logs := captureWarnLogs(t.Context()) + reportDetachedDescendants(ctx, ServerOptions{}, t.TempDir(), testServerPid) + + assert.Empty(t, logs.String()) + }) +} + +// The teardown event is the only measurement of how often the tunnel is about to destroy +// detached work, so this pins that it reaches the wire with the fields a query needs. +func TestReportDetachedDescendantsTelemetry(t *testing.T) { + tests := []struct { + name string + opts ServerOptions + root func(t *testing.T) string + want protos.SshTunnelTeardownEvent + }{ + { + name: "detached work left behind on a dedicated cluster", + opts: ServerOptions{}, + root: procWithDetachedWork, + want: protos.SshTunnelTeardownEvent{ + ComputeType: protos.SshTunnelComputeTypeDedicated, + HadDetachedDescendantsAtTeardown: true, + }, + }, + { + name: "the run was held open for it", + opts: ServerOptions{KeepDetachedProcesses: true}, + root: procWithDetachedWork, + want: protos.SshTunnelTeardownEvent{ + ComputeType: protos.SshTunnelComputeTypeDedicated, + KeepDetachedRequested: true, + HadDetachedDescendantsAtTeardown: true, + }, + }, + { + name: "nothing detached, on serverless", + opts: ServerOptions{Serverless: true}, + root: func(t *testing.T) string { + return fakeProc(t, map[int]fakeProcess{ + testNotebookPid: {comm: "python", ppid: 1, pgrp: testNotebookPid, state: "S"}, + testServerPid: {comm: "databricks", ppid: testNotebookPid, pgrp: testServerPid, state: "S"}, + }) + }, + want: protos.SshTunnelTeardownEvent{ComputeType: protos.SshTunnelComputeTypeServerless}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + server := testserver.New(t) + t.Cleanup(server.Close) + + var body telemetry.RequestBody + server.Handle("POST", "/telemetry-ext", func(req testserver.Request) any { + require.NoError(t, json.Unmarshal(req.Body, &body)) + return telemetry.ResponseBody{NumProtoSuccess: 1} + }) + + ctx := telemetry.WithNewLogger(t.Context()) + ctx = cmdctx.SetConfigUsed(ctx, &config.Config{Host: server.URL, Token: "token"}) + + reportDetachedDescendants(ctx, tt.opts, tt.root(t), testServerPid) + require.NoError(t, telemetry.Upload(ctx, protos.ExecutionContext{})) + + require.Len(t, body.ProtoLogs, 1) + var logged protos.FrontendLog + require.NoError(t, json.Unmarshal([]byte(body.ProtoLogs[0]), &logged)) + require.NotNil(t, logged.Entry.DatabricksCliLog.SshTunnelTeardownEvent) + assert.Equal(t, tt.want, *logged.Entry.DatabricksCliLog.SshTunnelTeardownEvent) + // The connect event stays untouched, so is_success queries keep counting connections. + assert.Nil(t, logged.Entry.DatabricksCliLog.SshTunnelEvent) + }) + } +} diff --git a/experimental/ssh/internal/setup/setup.go b/experimental/ssh/internal/setup/setup.go index 2f55da1ca46..bdfd88cc3c3 100644 --- a/experimental/ssh/internal/setup/setup.go +++ b/experimental/ssh/internal/setup/setup.go @@ -30,6 +30,11 @@ type SetupOptions struct { // Maximum lifetime of the SSH server, will be added as a --server-timeout flag to the ProxyCommand. // Also fixed at submission time. ServerTimeout time.Duration + // Whether detached work prevents idle shutdown, added as --keep-detached-processes + // to the ProxyCommand. Fixed at + // submission time like the two above, so this is the only place a host configured through + // setup can ask for it. + KeepDetachedProcesses bool // Optional path to the local ssh config. Defaults to ~/.ssh/config SSHConfigPath string // Optional path to the local directory to store SSH keys. Defaults to ~/.databricks/ssh-tunnel-keys @@ -119,12 +124,13 @@ func Setup(ctx context.Context, client *databricks.WorkspaceClient, opts SetupOp // omits --cluster, the ID is only known after the interactive picker above, // so building it earlier would serialize an empty --cluster= flag. clientOpts := sshclient.ClientOptions{ - ClusterID: opts.ClusterID, - AutoStartCluster: opts.AutoStartCluster, - ShutdownDelay: opts.ShutdownDelay, - MaxClients: opts.MaxClients, - ServerTimeout: opts.ServerTimeout, - Profile: opts.Profile, + ClusterID: opts.ClusterID, + AutoStartCluster: opts.AutoStartCluster, + ShutdownDelay: opts.ShutdownDelay, + MaxClients: opts.MaxClients, + ServerTimeout: opts.ServerTimeout, + KeepDetachedProcesses: opts.KeepDetachedProcesses, + Profile: opts.Profile, } // The ProxyCommand is persisted in the SSH config, so reject values that would produce a // tunnel that can never work (e.g. --max-clients=0) here rather than at first `ssh `. diff --git a/experimental/ssh/internal/setup/setup_test.go b/experimental/ssh/internal/setup/setup_test.go index 5dc04e1b366..0846b6caae6 100644 --- a/experimental/ssh/internal/setup/setup_test.go +++ b/experimental/ssh/internal/setup/setup_test.go @@ -280,23 +280,25 @@ func TestSetup_SerializesServerLifecycleFlags(t *testing.T) { }, nil) opts := SetupOptions{ - HostName: "test-host", - ClusterID: "cluster-123", - SSHConfigPath: filepath.Join(tmpDir, "ssh_config"), - SSHKeysDir: tmpDir, - ShutdownDelay: 30 * time.Second, - MaxClients: 25, - ServerTimeout: 48 * time.Hour, + HostName: "test-host", + ClusterID: "cluster-123", + SSHConfigPath: filepath.Join(tmpDir, "ssh_config"), + SSHKeysDir: tmpDir, + ShutdownDelay: 30 * time.Second, + MaxClients: 25, + ServerTimeout: 48 * time.Hour, + KeepDetachedProcesses: true, } require.NoError(t, Setup(ctx, m.WorkspaceClient, opts)) - // The ProxyCommand is the invocation that submits the server job, so both values have to - // reach the persisted host config or the user's choice is silently dropped. + // The ProxyCommand is the invocation that submits the server job, so every value fixed at + // submission has to reach the persisted host config or the user's choice is silently dropped. hostContent, err := os.ReadFile(filepath.Join(tmpDir, ".databricks", "ssh-tunnel-configs", "test-host")) require.NoError(t, err) assert.Contains(t, string(hostContent), "--max-clients=25") assert.Contains(t, string(hostContent), "--server-timeout=48h0m0s") + assert.Contains(t, string(hostContent), "--keep-detached-processes") } func TestSetup_RejectsUnusableServerLifecycleFlags(t *testing.T) { diff --git a/experimental/ssh/internal/workspace/workspace.go b/experimental/ssh/internal/workspace/workspace.go index 576e8a6df9f..c67adb6bac4 100644 --- a/experimental/ssh/internal/workspace/workspace.go +++ b/experimental/ssh/internal/workspace/workspace.go @@ -22,6 +22,11 @@ type WorkspaceMetadata struct { // UsagePolicyID records the usage policy the server's job was submitted with, so a // reconnect can tell whether a running server matches the requested usage policy. UsagePolicyID string `json:"usage_policy_id,omitempty"` + // KeepDetachedProcesses records whether the server's bootstrap notebook will hold the job + // run open for detached processes after the server exits (--keep-detached-processes), so a + // reconnect can tell whether a running server honours the requested mode. False, and so + // absent, when the session did not ask for it. + KeepDetachedProcesses bool `json:"keep_detached_processes,omitempty"` } func getWorkspaceRootDir(ctx context.Context, client *databricks.WorkspaceClient) (string, error) { diff --git a/libs/telemetry/protos/frontend_log.go b/libs/telemetry/protos/frontend_log.go index 75cc50c7057..5ba72b6bf8b 100644 --- a/libs/telemetry/protos/frontend_log.go +++ b/libs/telemetry/protos/frontend_log.go @@ -20,6 +20,7 @@ type DatabricksCliLog struct { BundleInitEvent *BundleInitEvent `json:"bundle_init_event,omitempty"` BundleDeployEvent *BundleDeployEvent `json:"bundle_deploy_event,omitempty"` SshTunnelEvent *SshTunnelEvent `json:"ssh_tunnel_event,omitempty"` + SshTunnelTeardownEvent *SshTunnelTeardownEvent `json:"ssh_tunnel_teardown_event,omitempty"` BundleConfigRemoteSyncEvent *BundleConfigRemoteSyncEvent `json:"bundle_config_remote_sync_event,omitempty"` AitoolsInstallEvent *AitoolsInstallEvent `json:"aitools_install_event,omitempty"` SetupLocalEvent *SetupLocalEvent `json:"setup_local_event,omitempty"` diff --git a/libs/telemetry/protos/ssh_tunnel.go b/libs/telemetry/protos/ssh_tunnel.go index 3240c7e9efe..55d8ca30fda 100644 --- a/libs/telemetry/protos/ssh_tunnel.go +++ b/libs/telemetry/protos/ssh_tunnel.go @@ -150,6 +150,12 @@ type SshTunnelEvent struct { // Only the presence is recorded, not the policy ID itself. HasUsagePolicy bool `json:"has_usage_policy"` + // Whether the connection asked for detached processes (tmux, setsid, nohup) to + // outlive the server via --keep-detached-processes. Only the request is recorded, not + // the duration. Whether any such process actually existed at teardown is reported + // separately by the server, in SshTunnelTeardownEvent. + KeepDetachedRequested bool `json:"keep_detached_requested"` + // Why the connection attempt failed, how an established session ended, or // TYPE_UNSPECIFIED when neither applies. Deliberately without omitempty: the field is // what identifies a failure's cause, so an empty value must not be silently dropped into diff --git a/libs/telemetry/protos/ssh_tunnel_teardown.go b/libs/telemetry/protos/ssh_tunnel_teardown.go new file mode 100644 index 00000000000..1c5c4fb48ea --- /dev/null +++ b/libs/telemetry/protos/ssh_tunnel_teardown.go @@ -0,0 +1,28 @@ +package protos + +// SshTunnelTeardownEvent is emitted by the SSH tunnel server on the compute when it +// shuts down after its idle timeout. It is a separate event from SshTunnelEvent +// because it is not a connection attempt: folding it into that event would add rows +// that every existing is_success query would count as connections. +// +// It exists to size the problem the --keep-detached-processes flag addresses: only the +// server, running on the compute at teardown, can see whether the session left +// detached processes behind, and by then the client that started it is long gone. +// +// The linger itself is deliberately not reported here. The bootstrap notebook is what +// holds the run open, and it outlives every Go process in the session, so no CLI +// process can observe when the linger ends; the run's own duration carries that. +type SshTunnelTeardownEvent struct { + // Type of compute: dedicated cluster or serverless. + ComputeType SshTunnelComputeType `json:"compute_type,omitempty"` + + // Whether the session asked for detached processes to be kept via + // --keep-detached-processes. + KeepDetachedRequested bool `json:"keep_detached_requested"` + + // Whether processes the tunnel started, but that left its process group (tmux, + // setsid, nohup), were still running when the server shut down. Without + // --keep-detached-processes those processes do not survive the run, so this counts how + // often the tunnel destroys work a user meant to keep. + HadDetachedDescendantsAtTeardown bool `json:"had_detached_descendants_at_teardown"` +} diff --git a/libs/telemetry/protos/ssh_tunnel_test.go b/libs/telemetry/protos/ssh_tunnel_test.go index d4ac094f716..503d961283b 100644 --- a/libs/telemetry/protos/ssh_tunnel_test.go +++ b/libs/telemetry/protos/ssh_tunnel_test.go @@ -25,6 +25,24 @@ func TestSshTunnelEventEncodesFailureExplicitly(t *testing.T) { "auto_start_cluster", "has_base_environment", "has_usage_policy", + "keep_detached_requested", + } { + assert.Equal(t, false, got[field], "%s must be sent as false, not omitted", field) + } +} + +// The teardown event's whole purpose is counting how often detached work is destroyed, so a +// "nothing was left behind" teardown has to arrive as false rather than as an absent field. +func TestSshTunnelTeardownEventEncodesFalseExplicitly(t *testing.T) { + b, err := json.Marshal(SshTunnelTeardownEvent{}) + require.NoError(t, err) + + var got map[string]any + require.NoError(t, json.Unmarshal(b, &got)) + + for _, field := range []string{ + "keep_detached_requested", + "had_detached_descendants_at_teardown", } { assert.Equal(t, false, got[field], "%s must be sent as false, not omitted", field) } @@ -33,13 +51,17 @@ func TestSshTunnelEventEncodesFailureExplicitly(t *testing.T) { // Guards fields added later: a bool that can legitimately be false must not // carry omitempty, or its false case arrives as NULL and cannot be counted. func TestSshTunnelEventBoolFieldsOmitOmitempty(t *testing.T) { - typ := reflect.TypeFor[SshTunnelEvent]() - for field := range typ.Fields() { - if field.Type.Kind() != reflect.Bool { - continue + for _, typ := range []reflect.Type{ + reflect.TypeFor[SshTunnelEvent](), + reflect.TypeFor[SshTunnelTeardownEvent](), + } { + for field := range typ.Fields() { + if field.Type.Kind() != reflect.Bool { + continue + } + tag := field.Tag.Get("json") + assert.NotContains(t, tag, "omitempty", + "%s.%s has omitempty; a false value would be indistinguishable from not reported", typ.Name(), field.Name) } - tag := field.Tag.Get("json") - assert.NotContains(t, tag, "omitempty", - "%s has omitempty; a false value would be indistinguishable from not reported", field.Name) } }