From 7acf45ba6bcc102a2b1d7ddfd497dd42ba494323 Mon Sep 17 00:00:00 2001 From: Anton Nekipelov <226657+anton-107@users.noreply.github.com> Date: Wed, 26 Aug 2026 19:45:41 +0000 Subject: [PATCH 1/6] ssh: add --keep-detached-for so detached processes survive the tunnel Today the SSH tunnel destroys work that was deliberately detached. When the server exits on its idle timer, the bootstrap notebook sweeps every process it parents with `pkill -P`, and because it sets PR_SET_CHILD_SUBREAPER it parents every orphan in the session - which is exactly what tmux, setsid and nohup make of themselves. Scoping that sweep to the server's own process group is not enough on its own, and that is the part measured on real compute (DECO-28187): survivors do live through the teardown, but the notebook is the WSFS-registered process, and WSFS authorises an I/O by walking the live process tree for a registered ancestor. The moment the notebook returns, the survivors reparent to PID 1 and every /Workspace and /Volumes call fails with EPERM. That turns a visible failure into an invisible one. So the two halves ship together, behind one flag: databricks ssh connect --cluster= --keep-detached-for= With it set, teardown terminates only the server's process group and the notebook then holds the job run open while any detached process is still alive, so the anchor stays in place. With it unset - the default - the notebook sweeps its children exactly as before, and the server logs a warning naming the processes it is about to destroy, so the original complaint (work vanished with no explanation) becomes actionable even for users who never set the flag. A duration rather than a boolean, because what the feature spends is time on a cluster: a held-open run suppresses autotermination. It is capped at the run's own timeout, rejected for serverless (the container teardown takes survivors regardless), and plumbed like --usage-policy-id: notebook base parameter, persisted in metadata.json, and a reconnect asking for a different value starts a new server. Telemetry: the connect event records whether the flag was asked for. Whether the session actually left detached work behind is only visible on the compute at teardown, so the server emits a new SshTunnelTeardownEvent for it - separate from SshTunnelEvent so it cannot be counted as a connection. That number is what tells us how often we destroy work today, and so whether the default should ever flip. Also narrows the startup sweep, which matched any command line containing `databricks_cli` and would otherwise let a new session destroy a previous session's detached work. Co-authored-by: Isaac --- .nextchanges/cli/ssh-keep-detached-for.md | 1 + .../ssh/connect-serverless-cpu/output.txt | 1 + .../ssh/connect-serverless-gpu/output.txt | 1 + acceptance/ssh/connection/output.txt | 1 + experimental/ssh/README.md | 30 ++++ experimental/ssh/cmd/connect.go | 3 + experimental/ssh/cmd/constants.go | 3 + experimental/ssh/cmd/server.go | 3 + experimental/ssh/internal/client/client.go | 58 ++++++- .../internal/client/client_internal_test.go | 9 + .../ssh/internal/client/client_test.go | 30 ++++ .../internal/client/policy_internal_test.go | 23 +++ .../internal/client/ssh-server-bootstrap.py | 117 ++++++++++++- .../internal/client/submit_internal_test.go | 16 ++ .../ssh/internal/server/descendants.go | 112 ++++++++++++ .../ssh/internal/server/descendants_test.go | 161 ++++++++++++++++++ experimental/ssh/internal/server/server.go | 79 +++++++-- .../ssh/internal/server/teardown_test.go | 143 ++++++++++++++++ .../ssh/internal/workspace/workspace.go | 5 + libs/telemetry/protos/frontend_log.go | 1 + libs/telemetry/protos/ssh_tunnel.go | 6 + libs/telemetry/protos/ssh_tunnel_teardown.go | 28 +++ libs/telemetry/protos/ssh_tunnel_test.go | 36 +++- 23 files changed, 831 insertions(+), 36 deletions(-) create mode 100644 .nextchanges/cli/ssh-keep-detached-for.md create mode 100644 experimental/ssh/internal/server/descendants.go create mode 100644 experimental/ssh/internal/server/descendants_test.go create mode 100644 experimental/ssh/internal/server/teardown_test.go create mode 100644 libs/telemetry/protos/ssh_tunnel_teardown.go diff --git a/.nextchanges/cli/ssh-keep-detached-for.md b/.nextchanges/cli/ssh-keep-detached-for.md new file mode 100644 index 00000000000..05eb21422e0 --- /dev/null +++ b/.nextchanges/cli/ssh-keep-detached-for.md @@ -0,0 +1 @@ +* `ssh connect` now accepts a `--keep-detached-for` 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 for up to the given duration so the survivors keep their `/Workspace` and `/Volumes` access. A held-open run also suppresses cluster autotermination, so the flag is off by default 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..40abd6c57c4 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", + "keepDetachedForSeconds": "0", "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..e3a2914d954 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", + "keepDetachedForSeconds": "0", "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..f400471bb8a 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", + "keepDetachedForSeconds": "0", "maxClients": "10", "secretScopeName": "[USERNAME]-[TEST_DEFAULT_CLUSTER_ID]-ssh-tunnel-keys", "serverless": "false", diff --git a/experimental/ssh/README.md b/experimental/ssh/README.md index 7d93b72d8e1..6399b249b48 100644 --- a/experimental/ssh/README.md +++ b/experimental/ssh/README.md @@ -70,6 +70,36 @@ 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`. + +`databricks ssh connect --cluster= --keep-detached-for=` changes that. On +teardown the tunnel terminates only its own process group - the server and its `sshd` +children - and then holds the job run open for up to `` while any detached process +is still running. Two things to know before using it: + +- **It holds the cluster up.** A `RUNNING` job run suppresses autotermination, so the + cluster keeps accruing DBUs until the work finishes or the duration runs out. That is why + the flag takes a duration rather than a boolean, and why it is off by default: the unit of + the knob is the thing being spent. `--keep-detached-for` cannot exceed the job's own 24h + timeout, and multi-day work still belongs in Jobs/DABs. Note also that reconnecting starts + a new run rather than rejoining the lingering one, so each session with live detached work + leaves its own run behind. +- **The notebook has to stay alive, not just the process.** Workspace filesystem access is + authorized by walking the live process tree for a registered ancestor, and the bootstrap + notebook is that ancestor. A detached process that outlives it keeps `/dbfs` and REST API + access but loses `/Workspace` and `/Volumes` with `EPERM` - which is why the group-scoped + teardown is tied to the linger and not enabled on its own. + +Dedicated clusters only. On serverless the container is torn down with the run, so survivors +die regardless and the flag is rejected. + +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..79fb274b9c4 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 keepDetachedFor time.Duration 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().DurationVar(&keepDetachedFor, "keep-detached-for", defaultKeepDetachedFor, "Keep processes detached from the SSH session (tmux, setsid, nohup) running for up to this long after the server shuts down, at the cost of holding the cluster up (dedicated clusters only)") 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)") @@ -154,6 +156,7 @@ Connect to a dedicated cluster: AdditionalArgs: args, AutoApprove: autoApprove, UsagePolicyID: usagePolicyID, + KeepDetachedFor: keepDetachedFor, } if err := opts.Validate(); err != nil { return err diff --git a/experimental/ssh/cmd/constants.go b/experimental/ssh/cmd/constants.go index 001ee7f23a6..9b279148c3a 100644 --- a/experimental/ssh/cmd/constants.go +++ b/experimental/ssh/cmd/constants.go @@ -16,6 +16,9 @@ const ( // Default cap on how long an SSH tunnel server is allowed to live. Fixed when the // server job is submitted, so it is only settable by the invocation that starts it. defaultServerTimeout = 24 * time.Hour + // Off by default: holding the run open for detached processes keeps the cluster from + // auto-terminating, so it is the caller who decides to spend that time. + defaultKeepDetachedFor = time.Duration(0) taskStartupTimeout = 10 * time.Minute gpuTaskStartupTimeout = 45 * time.Minute diff --git a/experimental/ssh/cmd/server.go b/experimental/ssh/cmd/server.go index 3675a4a7fe8..5a0765be92a 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 keepDetachedFor time.Duration 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().DurationVar(&keepDetachedFor, "keep-detached-for", defaultKeepDetachedFor, "How long the bootstrap notebook holds the job run open for detached processes after the server exits") 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, + KeepDetachedFor: keepDetachedFor, } return server.Run(ctx, wsc, opts) } diff --git a/experimental/ssh/internal/client/client.go b/experimental/ssh/internal/client/client.go index 1810d7f90ad..71a640c53f3 100644 --- a/experimental/ssh/internal/client/client.go +++ b/experimental/ssh/internal/client/client.go @@ -130,6 +130,11 @@ type ClientOptions struct { AutoApprove bool // Id of the usage policy to use for the serverless SSH server job. Serverless only. UsagePolicyID string + // How long the bootstrap notebook holds the job run open after the SSH server shuts + // down, so processes detached from the session (tmux, setsid, nohup) keep running with + // their workspace filesystem access intact. Zero, the default, keeps today's behaviour: + // the run ends with the server and nothing outlives it. Dedicated clusters only. + KeepDetachedFor time.Duration } func (o *ClientOptions) Validate() error { @@ -142,6 +147,14 @@ func (o *ClientOptions) Validate() error { if o.UsagePolicyID != "" && o.ClusterID != "" { return errors.New("--usage-policy-id flag can only be used with serverless compute (--name flag)") } + if o.KeepDetachedFor < 0 { + return fmt.Errorf("--keep-detached-for must not be negative, got %v", o.KeepDetachedFor) + } + // On serverless the container goes away with the run, so nothing survives the server + // however long the notebook lingers. + if o.KeepDetachedFor > 0 && o.ClusterID == "" { + return errors.New("--keep-detached-for flag can only be used with a dedicated cluster (--cluster flag)") + } if o.Accelerator != "" && o.Accelerator != "GPU_1xA10" && o.Accelerator != "GPU_8xH100" { return fmt.Errorf("invalid accelerator value: %q, expected %q or %q", o.Accelerator, "GPU_1xA10", "GPU_8xH100") } @@ -178,6 +191,11 @@ func (o *ClientOptions) Validate() error { if o.ShutdownDelay > o.ServerTimeout { return fmt.Errorf("--shutdown-delay (%s) cannot be longer than --server-timeout (%s)", o.ShutdownDelay, o.ServerTimeout) } + // The run's own timeout is the hard ceiling: Jobs terminates it regardless of what the + // notebook is waiting for, so a larger value would promise time we cannot deliver. + if o.KeepDetachedFor > o.ServerTimeout { + return fmt.Errorf("--keep-detached-for must not exceed %v, the maximum lifetime of the SSH server job, got %v", o.ServerTimeout, o.KeepDetachedFor) + } return nil } @@ -250,6 +268,9 @@ func (o *ClientOptions) ToProxyCommand() (string, error) { } else { 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.KeepDetachedFor > 0 { + proxyCommand += " --keep-detached-for=" + o.KeepDetachedFor.String() + } } // Both of these are fixed when the server job is submitted, and for a host configured by @@ -600,6 +621,8 @@ type serverMetadata struct { ClusterID string // UsagePolicyID the server was started with, used to decide whether a running server can be reused. UsagePolicyID string + // KeepDetachedForMs the server's run was submitted with, used the same way as UsagePolicyID. + KeepDetachedForMs int64 } // getServerMetadata retrieves the server metadata from the workspace and validates it via Driver Proxy. @@ -647,10 +670,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, + KeepDetachedForMs: wsMetadata.KeepDetachedForMs, }, nil } @@ -719,6 +743,10 @@ 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, + // Seconds rather than a duration string, because the bootstrap counts in seconds. + // Validate caps this at ServerTimeout, the run's own timeout, so the notebook can + // never wait longer than Jobs keeps the run alive. + "keepDetachedForSeconds": strconv.Itoa(int(opts.KeepDetachedFor.Seconds())), } task := jobs.SubmitTask{ @@ -1302,6 +1330,14 @@ func usagePolicyMatches(storedPolicy, requestedPolicy string) bool { return requestedPolicy == "" || storedPolicy == requestedPolicy } +// keepDetachedMatches reports whether a running server holds its run open for exactly as +// long as this connection asked for. The linger is fixed when the run is submitted, so a +// different duration needs a new server; a connection that asks for nothing takes whatever +// is already running. +func keepDetachedMatches(storedMs, requestedMs int64) bool { + return requestedMs == 0 || storedMs == requestedMs +} + 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 @@ -1316,7 +1352,10 @@ func ensureSSHServerIsRunning(ctx context.Context, client *databricks.WorkspaceC // 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) + keepDetachedMs := opts.KeepDetachedFor.Milliseconds() + needNewServer := err != nil || + !usagePolicyMatches(meta.UsagePolicyID, opts.UsagePolicyID) || + !keepDetachedMatches(meta.KeepDetachedForMs, keepDetachedMs) if needNewServer { cmdio.LogString(ctx, "Starting SSH server...") @@ -1340,6 +1379,10 @@ 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.KeepDetachedForMs, keepDetachedMs) { + err = fmt.Errorf("found a running SSH server keeping detached processes for %v, waiting for the one keeping them for %v", + time.Duration(meta.KeepDetachedForMs)*time.Millisecond, opts.KeepDetachedFor) + } if err == nil { cmdio.LogString(ctx, "Health check successful, starting ssh WebSocket connection...") break @@ -1481,6 +1524,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.KeepDetachedFor > 0, + 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..3108b110a6c 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 presence only", + opts: ClientOptions{ClusterID: "abc-123", KeepDetachedFor: 2 * time.Hour}, + 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..22d9c1996e4 100644 --- a/experimental/ssh/internal/client/client_test.go +++ b/experimental/ssh/internal/client/client_test.go @@ -121,6 +121,29 @@ func TestValidate(t *testing.T) { name: "usage policy with connection name", opts: client.ClientOptions{ConnectionName: "my-conn", UsagePolicyID: "pol-1"}, }, + { + name: "keep detached with serverless", + opts: client.ClientOptions{ConnectionName: "my-conn", KeepDetachedFor: time.Hour}, + wantErr: "--keep-detached-for flag can only be used with a dedicated cluster (--cluster flag)", + }, + { + name: "keep detached with cluster ID", + opts: client.ClientOptions{ClusterID: "abc-123", KeepDetachedFor: time.Hour}, + }, + { + name: "keep detached beyond the server timeout", + opts: client.ClientOptions{ClusterID: "abc-123", KeepDetachedFor: 25 * time.Hour}, + wantErr: "--keep-detached-for must not exceed 24h0m0s, the maximum lifetime of the SSH server job, got 25h0m0s", + }, + { + name: "keep detached exactly at the server timeout", + opts: client.ClientOptions{ClusterID: "abc-123", KeepDetachedFor: 24 * time.Hour}, + }, + { + name: "negative keep detached", + opts: client.ClientOptions{ClusterID: "abc-123", KeepDetachedFor: -time.Minute}, + wantErr: "--keep-detached-for must not be negative, got -1m0s", + }, } for _, tt := range tests { @@ -334,6 +357,13 @@ 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 + // linger, instead of starting a server that would sweep the detached work. + name: "dedicated cluster keeping detached processes", + opts: client.ClientOptions{ClusterID: "abc-123", KeepDetachedFor: 2 * time.Hour, ShutdownDelay: 5 * time.Minute}, + want: quoted + " ssh connect --proxy --cluster=abc-123 --auto-start-cluster=false --shutdown-delay=5m0s --keep-detached-for=2h0m0s", + }, { 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..cfe48ffa777 100644 --- a/experimental/ssh/internal/client/policy_internal_test.go +++ b/experimental/ssh/internal/client/policy_internal_test.go @@ -24,3 +24,26 @@ func TestUsagePolicyMatches(t *testing.T) { }) } } + +func TestKeepDetachedMatches(t *testing.T) { + tests := []struct { + name string + stored int64 + requested int64 + want bool + }{ + {name: "empty request takes a lingering server", stored: 3600000, requested: 0, want: true}, + {name: "empty request takes a server that does not linger", stored: 0, requested: 0, want: true}, + {name: "equal durations match", stored: 3600000, requested: 3600000, want: true}, + {name: "different durations do not match", stored: 3600000, requested: 7200000, want: false}, + {name: "request against a server that does not linger does not match", stored: 0, requested: 3600000, 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(%d, %d) = %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..fd74414fb0c 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,9 @@ SSH_TUNNEL_BASENAME = "databricks_cli" +# How often the linger loop re-checks for detached processes still holding the run open. +LINGER_POLL_SECONDS = 15 + # 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 +29,14 @@ dbutils.widgets.text("sessionId", "") dbutils.widgets.text("serverless", "false") dbutils.widgets.text("usagePolicyId", "") +dbutils.widgets.text("keepDetachedForSeconds", "0") 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 +82,78 @@ 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. + """ + result = subprocess.run( + ["ps", "-o", "pid=,pgid=,stat=", "--ppid", str(os.getpid())], + capture_output=True, + text=True, + check=False, + ) + survivors = [] + for line in result.stdout.splitlines(): + fields = line.split() + if len(fields) < 3: + continue + pid, pgid, stat = fields + # Zombies are already dead; the SIGCHLD handler collects them. + if stat.startswith("Z"): + continue + if pgid != str(server_pgid): + survivors.append(pid) + return survivors + + +def wait_for_detached_descendants(server_pgid, timeout_seconds): + """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 and + bounded by --keep-detached-for. + """ + deadline = time.monotonic() + timeout_seconds + while True: + survivors = detached_descendants(server_pgid) + if not survivors: + print("No detached processes left, releasing the run", flush=True) + return + if time.monotonic() > deadline: + print( + f"Reached the --keep-detached-for limit of {timeout_seconds}s with " + f"{len(survivors)} detached process(es) still running: {','.join(survivors)}. " + "Releasing the run; they lose /Workspace and /Volumes access from here.", + flush=True, + ) + return + print( + f"Holding the run open for {len(survivors)} detached process(es): {','.join(survivors)}", + flush=True, + ) + time.sleep(LINGER_POLL_SECONDS) def run_ssh_server(): @@ -128,6 +203,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_for_seconds = int(dbutils.widgets.get("keepDetachedForSeconds") or 0) # Mark this process's WSFS command origin so workspace-file activity from the # remote SSH session is attributable @@ -178,11 +254,28 @@ def run_ssh_server(): if usage_policy_id: server_args.append(f"--usage-policy-id={usage_policy_id}") + # The server does not linger itself; it uses this to persist the mode for reconnects and + # to warn about detached work it is about to leave behind when the mode is off. + if keep_detached_for_seconds > 0: + server_args.append(f"--keep-detached-for={keep_detached_for_seconds}s") + # 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 @@ -196,11 +289,19 @@ def run_ssh_server(): # 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 lingering 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_for_seconds > 0: + wait_for_detached_descendants(server_pgid, keep_detached_for_seconds) + else: + kill_all_children() if __name__ == "__main__": cleanup() setup_subreaper() - setup_exit_handler() run_ssh_server() diff --git a/experimental/ssh/internal/client/submit_internal_test.go b/experimental/ssh/internal/client/submit_internal_test.go index fa1bd5b2dd1..4aec4ba11a2 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) + // Zero is what tells the bootstrap to sweep detached work as it always has. + assert.Equal(t, "0", got.Tasks[0].NotebookTask.BaseParameters["keepDetachedForSeconds"]) + }) + + t.Run("dedicated cluster keeping detached processes", func(t *testing.T) { + opts := ClientOptions{ + ClusterID: "abc-123", + ServerTimeout: 24 * time.Hour, + KeepDetachedFor: 90 * time.Minute, + } + got := buildSSHServerSubmitRun("v1", "scope", notebookPath, "", opts) + + assert.Equal(t, "5400", got.Tasks[0].NotebookTask.BaseParameters["keepDetachedForSeconds"]) + // The linger happens inside the run, so the run's own timeout still 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/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/server.go b/experimental/ssh/internal/server/server.go index e5454a60ed4..28c19f983c3 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,12 @@ 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 + // KeepDetachedFor is how long the bootstrap notebook holds the job run open for + // detached processes after this server exits. Zero means it does not: the notebook + // sweeps them as it always has. The server does not linger itself; it only needs the + // value to persist it for reconnects and to decide whether to warn about work it is + // about to destroy. + KeepDetachedFor time.Duration // The directory to store sshd configuration ConfigDir string // The name of the secrets scope to use for client and server keys @@ -81,9 +89,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, + KeepDetachedForMs: opts.KeepDetachedFor.Milliseconds(), } err = workspace.SaveWorkspaceMetadata(ctx, client, opts.Version, opts.SessionID, metadata) if err != nil { @@ -122,13 +131,59 @@ 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) - })) + select { + case err := <-listenErr: + return err + case <-connections.TimedOut: + // 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-for these processes do not outlive the + // run, and until now they vanished with no explanation anywhere. The client reads this + // back through /logs. Serverless is excluded because the container teardown takes them + // regardless, so the flag cannot help there and is rejected for it. + if len(pids) > 0 && opts.KeepDetachedFor == 0 && !opts.Serverless { + 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-for=\", which holds the run open for them.", + 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.KeepDetachedFor > 0, + HadDetachedDescendantsAtTeardown: len(pids) > 0, + }, + }) } // serveCapabilities tells the client which optional parts of the tunnel protocol this server @@ -153,12 +208,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..9ee64b730cc --- /dev/null +++ b/experimental/ssh/internal/server/teardown_test.go @@ -0,0 +1,143 @@ +package server + +import ( + "encoding/json" + "testing" + "time" + + "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-for") + }) + + t.Run("stays quiet when the run is held open for them", func(t *testing.T) { + ctx, logs := captureWarnLogs(t.Context()) + opts := ServerOptions{KeepDetachedFor: time.Hour} + reportDetachedDescendants(ctx, opts, procWithDetachedWork(t), testServerPid) + + assert.Empty(t, logs.String()) + }) + + // The flag is rejected for serverless, so pointing at it there would be misleading. + t.Run("stays quiet on serverless", func(t *testing.T) { + ctx, logs := captureWarnLogs(t.Context()) + reportDetachedDescendants(ctx, ServerOptions{Serverless: true}, procWithDetachedWork(t), testServerPid) + + assert.Empty(t, logs.String()) + }) + + 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{KeepDetachedFor: 2 * time.Hour}, + 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/workspace/workspace.go b/experimental/ssh/internal/workspace/workspace.go index 576e8a6df9f..5da0938f86e 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"` + // KeepDetachedForMs records how long the server's bootstrap notebook will hold the job + // run open for detached processes after the server exits (--keep-detached-for), so a + // reconnect can tell whether a running server honours the requested duration. Zero, and + // so absent, when the session did not ask for it. + KeepDetachedForMs int64 `json:"keep_detached_for_ms,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..c35165de657 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-for. Only the presence 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..82de9adb06d --- /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-for 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-for. Only the presence is recorded, not the duration. + 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-for 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) } } From 5df516e0f69c017731bc984b62241d7b4e80a2ce Mon Sep 17 00:00:00 2001 From: Anton Nekipelov <226657+anton-107@users.noreply.github.com> Date: Thu, 27 Aug 2026 14:55:47 +0000 Subject: [PATCH 2/6] ssh: release the keep-detached linger early, and survive a server handover Fixes the two defects found verifying #6387 on real compute (DECO-28187). 1. The linger never released early. `detached_descendants()` asked `ps` for this process's children, but `ps` is itself one of them and `subprocess.run` leaves it in the notebook's process group, so it passed the "outside the server's group" filter on every poll. The survivor list was never empty: `--keep-detached-for` held the run - and suppressed cluster autotermination - for its full duration even with zero detached work, and "No detached processes left, releasing the run" was unreachable. It now reads `/proc` directly, mirroring `detachedDescendants` in `internal/server/descendants.go`, which counted correctly for exactly this reason. 2. A server handover marked the previous run FAILED. A new session's bootstrap terminates a server already running on the cluster, which a reconnect asking for a different `--keep-detached-for` now reaches on a normal path; the previous notebook saw -15 and raised, so its run ended INTERNAL_ERROR/FAILED. A SIGTERM exit is now treated as a handover rather than a failure. The comment beside `needNewServer` claimed the displaced server "idles out via shutdownDelay" - it is killed; corrected. Also documents what the verification measured about the feature's real ceiling: the work lives at most the linger plus the cluster's autotermination window, reconnecting after the linger leaves survivors with EPERM on every workspace path, and a reconnect that omits the flag inherits the linger of the server it reuses. Verified locally under a PR_SET_CHILD_SUBREAPER parent with a real `setsid` child: the old code reports a phantom survivor with a fresh pid each call and never empties, the new code reports the one real detached process and empties as soon as it exits. `./task test-exp-ssh` (327 unit + 7 acceptance), `./task lint-q`, ruff format and the whitespace check pass. Co-authored-by: Isaac --- experimental/ssh/README.md | 17 +++++++- experimental/ssh/internal/client/client.go | 9 ++-- .../internal/client/ssh-server-bootstrap.py | 43 ++++++++++++------- 3 files changed, 48 insertions(+), 21 deletions(-) diff --git a/experimental/ssh/README.md b/experimental/ssh/README.md index 6399b249b48..d4f505f8a24 100644 --- a/experimental/ssh/README.md +++ b/experimental/ssh/README.md @@ -79,7 +79,7 @@ including work that was deliberately detached with `tmux`, `setsid` or `nohup`. `databricks ssh connect --cluster= --keep-detached-for=` changes that. On teardown the tunnel terminates only its own process group - the server and its `sshd` children - and then holds the job run open for up to `` while any detached process -is still running. Two things to know before using it: +is still running. Things to know before using it: - **It holds the cluster up.** A `RUNNING` job run suppresses autotermination, so the cluster keeps accruing DBUs until the work finishes or the duration runs out. That is why @@ -88,15 +88,28 @@ is still running. Two things to know before using it: timeout, and multi-day work still belongs in Jobs/DABs. Note also that reconnecting starts a new run rather than rejoining the lingering one, so each session with live detached work leaves its own run behind. +- **The duration is not the whole lifetime the work gets.** Once the linger ends the run is + released and the cluster starts its own autotermination countdown, which is the last thing + that reaps survivors - detached work does not count as cluster activity, however busy it + is. The ceiling is therefore `` plus the cluster's `autotermination_minutes`, and + only the first part is yours to set: `--keep-detached-for=5m` does not carry a three-hour + job. - **The notebook has to stay alive, not just the process.** Workspace filesystem access is authorized by walking the live process tree for a registered ancestor, and the bootstrap notebook is that ancestor. A detached process that outlives it keeps `/dbfs` and REST API access but loses `/Workspace` and `/Volumes` with `EPERM` - which is why the group-scoped - teardown is tied to the linger and not enabled on its own. + teardown is tied to the linger and not enabled on its own. So: **reconnect before the + linger expires.** After it, the work is still running but every workspace path fails, and a + new window opened inside a surviving `tmux` inherits that failure, because the `tmux` + server - not the shell - is what lost its registered ancestor. Dedicated clusters only. On serverless the container is torn down with the run, so survivors die regardless and the flag is rejected. +A reconnect that omits the flag reuses a running server that was started with it, linger +included, so a session that never asked for it can end up holding the cluster open. Passing a +different duration 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. diff --git a/experimental/ssh/internal/client/client.go b/experimental/ssh/internal/client/client.go index 71a640c53f3..7499cd24c1b 100644 --- a/experimental/ssh/internal/client/client.go +++ b/experimental/ssh/internal/client/client.go @@ -1348,10 +1348,11 @@ func ensureSSHServerIsRunning(ctx context.Context, client *databricks.WorkspaceC return "", 0, "", err } - // 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. + // Start a new server when none is running, or when the running one was started with a + // different usage policy or linger duration. 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. keepDetachedMs := opts.KeepDetachedFor.Milliseconds() needNewServer := err != nil || !usagePolicyMatches(meta.UsagePolicyID, opts.UsagePolicyID) || diff --git a/experimental/ssh/internal/client/ssh-server-bootstrap.py b/experimental/ssh/internal/client/ssh-server-bootstrap.py index fd74414fb0c..d1e74d85c4e 100644 --- a/experimental/ssh/internal/client/ssh-server-bootstrap.py +++ b/experimental/ssh/internal/client/ssh-server-bootstrap.py @@ -104,25 +104,32 @@ def detached_descendants(server_pgid): 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. """ - result = subprocess.run( - ["ps", "-o", "pid=,pgid=,stat=", "--ppid", str(os.getpid())], - capture_output=True, - text=True, - check=False, - ) + self_pid = os.getpid() survivors = [] - for line in result.stdout.splitlines(): - fields = line.split() - if len(fields) < 3: + 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 - pid, pgid, stat = fields # Zombies are already dead; the SIGCHLD handler collects them. - if stat.startswith("Z"): + if state.startswith("Z"): continue - if pgid != str(server_pgid): - survivors.append(pid) - return survivors + survivors.append(entry) + return sorted(survivors, key=int) def wait_for_detached_descendants(server_pgid, timeout_seconds): @@ -285,7 +292,13 @@ 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 a different --keep-detached-for 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 linger 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: From e1daec5cdcf9e3f683a6a9f6cb486e594c9721e0 Mon Sep 17 00:00:00 2001 From: Anton Nekipelov <226657+anton-107@users.noreply.github.com> Date: Thu, 17 Sep 2026 21:07:24 +0000 Subject: [PATCH 3/6] ssh: make keeping detached processes a boolean flag Replaces --keep-detached-for= with --keep-detached-processes, a boolean bounded by the run's own timeout (--server-timeout, 24h by default). The duration was a second budget for the same run, measured from a different origin: the run timeout runs from run start, while the linger clock only started once the server exited. --keep-detached-for=24h under the default --server-timeout therefore validated but could never deliver 24h, and since Validate compared the two in isolation, the ceiling it accepted silently depended on how long the session had already run. Presence is the condition the feature wants - hold the run while detached work is alive - and --server-timeout already means "how long this run may live", so it can carry the cost control on its own. Dropped along with the duration: the negative-value check, the ceiling check, the seconds-based notebook widget, and duration matching on reconnect. keepDetachedMatches is now "asking for it needs a server that has it". The notebook's hold has no deadline any more, so its per-poll log line is throttled to one every five minutes while the set of survivors is unchanged, instead of every 15 seconds for as long as the work lives. Co-authored-by: Isaac --- .nextchanges/cli/ssh-keep-detached-for.md | 1 - .../cli/ssh-keep-detached-processes.md | 1 + .../ssh/connect-serverless-cpu/output.txt | 2 +- .../ssh/connect-serverless-gpu/output.txt | 2 +- acceptance/ssh/connection/output.txt | 2 +- experimental/ssh/README.md | 47 +++++------ experimental/ssh/cmd/connect.go | 58 ++++++------- experimental/ssh/cmd/constants.go | 3 - experimental/ssh/cmd/server.go | 6 +- experimental/ssh/internal/client/client.go | 82 ++++++++----------- .../internal/client/client_internal_test.go | 4 +- .../ssh/internal/client/client_test.go | 30 ++----- .../internal/client/policy_internal_test.go | 15 ++-- .../internal/client/ssh-server-bootstrap.py | 57 +++++++------ .../internal/client/submit_internal_test.go | 14 ++-- experimental/ssh/internal/server/server.go | 28 +++---- .../ssh/internal/server/teardown_test.go | 7 +- .../ssh/internal/workspace/workspace.go | 10 +-- libs/telemetry/protos/ssh_tunnel.go | 2 +- libs/telemetry/protos/ssh_tunnel_teardown.go | 6 +- 20 files changed, 175 insertions(+), 202 deletions(-) delete mode 100644 .nextchanges/cli/ssh-keep-detached-for.md create mode 100644 .nextchanges/cli/ssh-keep-detached-processes.md diff --git a/.nextchanges/cli/ssh-keep-detached-for.md b/.nextchanges/cli/ssh-keep-detached-for.md deleted file mode 100644 index 05eb21422e0..00000000000 --- a/.nextchanges/cli/ssh-keep-detached-for.md +++ /dev/null @@ -1 +0,0 @@ -* `ssh connect` now accepts a `--keep-detached-for` 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 for up to the given duration so the survivors keep their `/Workspace` and `/Volumes` access. A held-open run also suppresses cluster autotermination, so the flag is off by default 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/.nextchanges/cli/ssh-keep-detached-processes.md b/.nextchanges/cli/ssh-keep-detached-processes.md new file mode 100644 index 00000000000..2e130ac0d9b --- /dev/null +++ b/.nextchanges/cli/ssh-keep-detached-processes.md @@ -0,0 +1 @@ +* `ssh connect` now accepts 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 40abd6c57c4..1a2da09b42c 100644 --- a/acceptance/ssh/connect-serverless-cpu/output.txt +++ b/acceptance/ssh/connect-serverless-cpu/output.txt @@ -21,7 +21,7 @@ "notebook_task": { "base_parameters": { "authorizedKeySecretName": "client-public-key", - "keepDetachedForSeconds": "0", + "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 e3a2914d954..c89a29c8193 100644 --- a/acceptance/ssh/connect-serverless-gpu/output.txt +++ b/acceptance/ssh/connect-serverless-gpu/output.txt @@ -22,7 +22,7 @@ "notebook_task": { "base_parameters": { "authorizedKeySecretName": "client-public-key", - "keepDetachedForSeconds": "0", + "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 f400471bb8a..1372be147a3 100644 --- a/acceptance/ssh/connection/output.txt +++ b/acceptance/ssh/connection/output.txt @@ -11,7 +11,7 @@ "notebook_task": { "base_parameters": { "authorizedKeySecretName": "client-public-key", - "keepDetachedForSeconds": "0", + "keepDetachedProcesses": "false", "maxClients": "10", "secretScopeName": "[USERNAME]-[TEST_DEFAULT_CLUSTER_ID]-ssh-tunnel-keys", "serverless": "false", diff --git a/experimental/ssh/README.md b/experimental/ssh/README.md index d4f505f8a24..855fbbda623 100644 --- a/experimental/ssh/README.md +++ b/experimental/ssh/README.md @@ -76,39 +76,36 @@ By default nothing outlives the session: when the last client disconnects, the s down after `--shutdown-delay` and the bootstrap notebook sweeps every process it parents, including work that was deliberately detached with `tmux`, `setsid` or `nohup`. -`databricks ssh connect --cluster= --keep-detached-for=` changes that. On -teardown the tunnel terminates only its own process group - the server and its `sshd` -children - and then holds the job run open for up to `` while any detached process -is still running. Things to know before using it: - -- **It holds the cluster up.** A `RUNNING` job run suppresses autotermination, so the - cluster keeps accruing DBUs until the work finishes or the duration runs out. That is why - the flag takes a duration rather than a boolean, and why it is off by default: the unit of - the knob is the thing being spent. `--keep-detached-for` cannot exceed the job's own 24h - timeout, and multi-day work still belongs in Jobs/DABs. Note also that reconnecting starts - a new run rather than rejoining the lingering one, so each session with live detached work - leaves its own run behind. -- **The duration is not the whole lifetime the work gets.** Once the linger ends the run is - released and the cluster starts its own autotermination countdown, which is the last thing - that reaps survivors - detached work does not count as cluster activity, however busy it - is. The ceiling is therefore `` plus the cluster's `autotermination_minutes`, and - only the first part is yours to set: `--keep-detached-for=5m` does not carry a three-hour - job. +`databricks ssh connect --cluster= --keep-detached-processes` changes that. On teardown +the tunnel terminates only its own process group - the server and its `sshd` children - and +then holds the job run open for as long as any detached process is still running. Things to +know before using it: + +- **It holds the cluster up.** A `RUNNING` job run suppresses autotermination, so the cluster + keeps accruing DBUs until the last detached process exits. The bound is the run's own + lifetime, `--server-timeout` (24h by default), which is therefore the knob to reach for when + the cost is what matters; multi-day work still belongs in Jobs/DABs. Note also that + reconnecting starts a new run rather than rejoining the one being held open, so each session + with live detached work leaves its own run behind. +- **Releasing the run is not the end of the work.** Once the hold is released the cluster + starts its own autotermination countdown, which is the last thing that reaps survivors - + detached work does not count as cluster activity, however busy it is. - **The notebook has to stay alive, not just the process.** Workspace filesystem access is authorized by walking the live process tree for a registered ancestor, and the bootstrap notebook is that ancestor. A detached process that outlives it keeps `/dbfs` and REST API access but loses `/Workspace` and `/Volumes` with `EPERM` - which is why the group-scoped - teardown is tied to the linger and not enabled on its own. So: **reconnect before the - linger expires.** After it, the work is still running but every workspace path fails, and a - new window opened inside a surviving `tmux` inherits that failure, because the `tmux` - server - not the shell - is what lost its registered ancestor. + teardown is tied to holding the run open and not enabled on its own. Work that finishes + while the run is held never sees this; work still running when `--server-timeout` expires + does, and from there every workspace path fails, including in a new window opened inside a + surviving `tmux`, because the `tmux` server - not the shell - is what lost its registered + ancestor. So size `--server-timeout` to the work you intend to leave behind. Dedicated clusters only. On serverless the container is torn down with the run, so survivors die regardless and the flag is rejected. -A reconnect that omits the flag reuses a running server that was started with it, linger -included, so a session that never asked for it can end up holding the cluster open. Passing a -different duration starts a fresh server instead. +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. diff --git a/experimental/ssh/cmd/connect.go b/experimental/ssh/cmd/connect.go index 79fb274b9c4..278b0fc2016 100644 --- a/experimental/ssh/cmd/connect.go +++ b/experimental/ssh/cmd/connect.go @@ -60,14 +60,14 @@ Connect to a dedicated cluster: var baseEnvironment string var autoApprove bool var usagePolicyID string - var keepDetachedFor time.Duration + 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().DurationVar(&keepDetachedFor, "keep-detached-for", defaultKeepDetachedFor, "Keep processes detached from the SSH session (tmux, setsid, nohup) running for up to this long after the server shuts down, at the cost of holding the cluster up (dedicated clusters only)") + cmd.Flags().BoolVar(&keepDetachedProcesses, "keep-detached-processes", false, "Keep processes detached from the SSH session (tmux, setsid, nohup) running after the tunnel shuts down. Holds the cluster up until they exit or --server-timeout elapses (dedicated clusters only)") 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)") @@ -130,33 +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, - KeepDetachedFor: keepDetachedFor, + 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/constants.go b/experimental/ssh/cmd/constants.go index 9b279148c3a..001ee7f23a6 100644 --- a/experimental/ssh/cmd/constants.go +++ b/experimental/ssh/cmd/constants.go @@ -16,9 +16,6 @@ const ( // Default cap on how long an SSH tunnel server is allowed to live. Fixed when the // server job is submitted, so it is only settable by the invocation that starts it. defaultServerTimeout = 24 * time.Hour - // Off by default: holding the run open for detached processes keeps the cluster from - // auto-terminating, so it is the caller who decides to spend that time. - defaultKeepDetachedFor = time.Duration(0) taskStartupTimeout = 10 * time.Minute gpuTaskStartupTimeout = 45 * time.Minute diff --git a/experimental/ssh/cmd/server.go b/experimental/ssh/cmd/server.go index 5a0765be92a..7dc1b555dc4 100644 --- a/experimental/ssh/cmd/server.go +++ b/experimental/ssh/cmd/server.go @@ -30,7 +30,7 @@ and proxies them to local SSH daemon processes.`, var authorizedKeySecretName string var serverless bool var usagePolicyID string - var keepDetachedFor time.Duration + var keepDetachedProcesses bool cmd.Flags().StringVar(&clusterID, "cluster", "", "Databricks cluster ID") cmd.MarkFlagRequired("cluster") @@ -46,7 +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().DurationVar(&keepDetachedFor, "keep-detached-for", defaultKeepDetachedFor, "How long the bootstrap notebook holds the job run open for detached processes after the server exits") + cmd.Flags().BoolVar(&keepDetachedProcesses, "keep-detached-processes", false, "Whether the bootstrap notebook holds the job run open for detached processes after the server exits") cmd.PreRunE = func(cmd *cobra.Command, args []string) error { // The server can be executed under a directory with an invalid bundle configuration. @@ -76,7 +76,7 @@ and proxies them to local SSH daemon processes.`, PortRange: serverPortRange, Serverless: serverless, UsagePolicyID: usagePolicyID, - KeepDetachedFor: keepDetachedFor, + KeepDetachedProcesses: keepDetachedProcesses, } return server.Run(ctx, wsc, opts) } diff --git a/experimental/ssh/internal/client/client.go b/experimental/ssh/internal/client/client.go index 7499cd24c1b..1fd5035bf4a 100644 --- a/experimental/ssh/internal/client/client.go +++ b/experimental/ssh/internal/client/client.go @@ -130,11 +130,12 @@ type ClientOptions struct { AutoApprove bool // Id of the usage policy to use for the serverless SSH server job. Serverless only. UsagePolicyID string - // How long the bootstrap notebook holds the job run open after the SSH server shuts - // down, so processes detached from the session (tmux, setsid, nohup) keep running with - // their workspace filesystem access intact. Zero, the default, keeps today's behaviour: - // the run ends with the server and nothing outlives it. Dedicated clusters only. - KeepDetachedFor time.Duration + // Whether the bootstrap notebook holds the job run open after the SSH server shuts down, + // for as long as processes detached from the session (tmux, setsid, nohup) keep running, + // so they keep their workspace filesystem access intact. False, the default, keeps + // today's behaviour: the run ends with the server and nothing outlives it. The run's own + // timeout (--server-timeout) is what bounds the hold. Dedicated clusters only. + KeepDetachedProcesses bool } func (o *ClientOptions) Validate() error { @@ -147,13 +148,10 @@ func (o *ClientOptions) Validate() error { if o.UsagePolicyID != "" && o.ClusterID != "" { return errors.New("--usage-policy-id flag can only be used with serverless compute (--name flag)") } - if o.KeepDetachedFor < 0 { - return fmt.Errorf("--keep-detached-for must not be negative, got %v", o.KeepDetachedFor) - } // On serverless the container goes away with the run, so nothing survives the server - // however long the notebook lingers. - if o.KeepDetachedFor > 0 && o.ClusterID == "" { - return errors.New("--keep-detached-for flag can only be used with a dedicated cluster (--cluster flag)") + // however long the notebook holds the run open. + if o.KeepDetachedProcesses && o.ClusterID == "" { + return errors.New("--keep-detached-processes flag can only be used with a dedicated cluster (--cluster flag)") } if o.Accelerator != "" && o.Accelerator != "GPU_1xA10" && o.Accelerator != "GPU_8xH100" { return fmt.Errorf("invalid accelerator value: %q, expected %q or %q", o.Accelerator, "GPU_1xA10", "GPU_8xH100") @@ -191,11 +189,6 @@ func (o *ClientOptions) Validate() error { if o.ShutdownDelay > o.ServerTimeout { return fmt.Errorf("--shutdown-delay (%s) cannot be longer than --server-timeout (%s)", o.ShutdownDelay, o.ServerTimeout) } - // The run's own timeout is the hard ceiling: Jobs terminates it regardless of what the - // notebook is waiting for, so a larger value would promise time we cannot deliver. - if o.KeepDetachedFor > o.ServerTimeout { - return fmt.Errorf("--keep-detached-for must not exceed %v, the maximum lifetime of the SSH server job, got %v", o.ServerTimeout, o.KeepDetachedFor) - } return nil } @@ -268,8 +261,8 @@ func (o *ClientOptions) ToProxyCommand() (string, error) { } else { 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.KeepDetachedFor > 0 { - proxyCommand += " --keep-detached-for=" + o.KeepDetachedFor.String() + if o.KeepDetachedProcesses { + proxyCommand += " --keep-detached-processes" } } @@ -621,8 +614,8 @@ type serverMetadata struct { ClusterID string // UsagePolicyID the server was started with, used to decide whether a running server can be reused. UsagePolicyID string - // KeepDetachedForMs the server's run was submitted with, used the same way as UsagePolicyID. - KeepDetachedForMs int64 + // 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. @@ -670,11 +663,11 @@ func getServerMetadata(ctx context.Context, client *databricks.WorkspaceClient, } return serverMetadata{ - Port: wsMetadata.Port, - UserName: string(bodyBytes), - ClusterID: effectiveClusterID, - UsagePolicyID: wsMetadata.UsagePolicyID, - KeepDetachedForMs: wsMetadata.KeepDetachedForMs, + Port: wsMetadata.Port, + UserName: string(bodyBytes), + ClusterID: effectiveClusterID, + UsagePolicyID: wsMetadata.UsagePolicyID, + KeepDetachedProcesses: wsMetadata.KeepDetachedProcesses, }, nil } @@ -743,10 +736,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, - // Seconds rather than a duration string, because the bootstrap counts in seconds. - // Validate caps this at ServerTimeout, the run's own timeout, so the notebook can - // never wait longer than Jobs keeps the run alive. - "keepDetachedForSeconds": strconv.Itoa(int(opts.KeepDetachedFor.Seconds())), + // 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{ @@ -1330,12 +1322,12 @@ func usagePolicyMatches(storedPolicy, requestedPolicy string) bool { return requestedPolicy == "" || storedPolicy == requestedPolicy } -// keepDetachedMatches reports whether a running server holds its run open for exactly as -// long as this connection asked for. The linger is fixed when the run is submitted, so a -// different duration needs a new server; a connection that asks for nothing takes whatever -// is already running. -func keepDetachedMatches(storedMs, requestedMs int64) bool { - return requestedMs == 0 || storedMs == requestedMs +// 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) { @@ -1348,15 +1340,14 @@ func ensureSSHServerIsRunning(ctx context.Context, client *databricks.WorkspaceC return "", 0, "", err } - // Start a new server when none is running, or when the running one was started with a - // different usage policy or linger duration. 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. - keepDetachedMs := opts.KeepDetachedFor.Milliseconds() + // Start a new server when none is running, or when the running one was started under a + // 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.KeepDetachedForMs, keepDetachedMs) + !keepDetachedMatches(meta.KeepDetachedProcesses, opts.KeepDetachedProcesses) if needNewServer { cmdio.LogString(ctx, "Starting SSH server...") @@ -1380,9 +1371,8 @@ 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.KeepDetachedForMs, keepDetachedMs) { - err = fmt.Errorf("found a running SSH server keeping detached processes for %v, waiting for the one keeping them for %v", - time.Duration(meta.KeepDetachedForMs)*time.Millisecond, opts.KeepDetachedFor) + 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...") @@ -1527,7 +1517,7 @@ func buildSshTunnelEvent(opts ClientOptions, outcome connectOutcome) *protos.Ssh HasUsagePolicy: opts.UsagePolicyID != "", // 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.KeepDetachedFor > 0, + 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 3108b110a6c..1fc34858488 100644 --- a/experimental/ssh/internal/client/client_internal_test.go +++ b/experimental/ssh/internal/client/client_internal_test.go @@ -512,8 +512,8 @@ func TestBuildSshTunnelEvent(t *testing.T) { }, }, { - name: "keeping detached processes records presence only", - opts: ClientOptions{ClusterID: "abc-123", KeepDetachedFor: 2 * time.Hour}, + name: "keeping detached processes records the request", + opts: ClientOptions{ClusterID: "abc-123", KeepDetachedProcesses: true}, want: protos.SshTunnelEvent{ ComputeType: protos.SshTunnelComputeTypeDedicated, ClientMode: protos.SshTunnelClientModeSSH, diff --git a/experimental/ssh/internal/client/client_test.go b/experimental/ssh/internal/client/client_test.go index 22d9c1996e4..ceb88898087 100644 --- a/experimental/ssh/internal/client/client_test.go +++ b/experimental/ssh/internal/client/client_test.go @@ -122,27 +122,13 @@ func TestValidate(t *testing.T) { opts: client.ClientOptions{ConnectionName: "my-conn", UsagePolicyID: "pol-1"}, }, { - name: "keep detached with serverless", - opts: client.ClientOptions{ConnectionName: "my-conn", KeepDetachedFor: time.Hour}, - wantErr: "--keep-detached-for flag can only be used with a dedicated cluster (--cluster flag)", + name: "keep detached processes with serverless", + opts: client.ClientOptions{ConnectionName: "my-conn", KeepDetachedProcesses: true}, + wantErr: "--keep-detached-processes flag can only be used with a dedicated cluster (--cluster flag)", }, { - name: "keep detached with cluster ID", - opts: client.ClientOptions{ClusterID: "abc-123", KeepDetachedFor: time.Hour}, - }, - { - name: "keep detached beyond the server timeout", - opts: client.ClientOptions{ClusterID: "abc-123", KeepDetachedFor: 25 * time.Hour}, - wantErr: "--keep-detached-for must not exceed 24h0m0s, the maximum lifetime of the SSH server job, got 25h0m0s", - }, - { - name: "keep detached exactly at the server timeout", - opts: client.ClientOptions{ClusterID: "abc-123", KeepDetachedFor: 24 * time.Hour}, - }, - { - name: "negative keep detached", - opts: client.ClientOptions{ClusterID: "abc-123", KeepDetachedFor: -time.Minute}, - wantErr: "--keep-detached-for must not be negative, got -1m0s", + name: "keep detached processes with cluster ID", + opts: client.ClientOptions{ClusterID: "abc-123", KeepDetachedProcesses: true}, }, } @@ -359,10 +345,10 @@ func TestToProxyCommand(t *testing.T) { }, { // Carried into the ProxyCommand so a reconnect through ssh asks for the same - // linger, instead of starting a server that would sweep the detached work. + // mode, instead of starting a server that would sweep the detached work. name: "dedicated cluster keeping detached processes", - opts: client.ClientOptions{ClusterID: "abc-123", KeepDetachedFor: 2 * time.Hour, ShutdownDelay: 5 * time.Minute}, - want: quoted + " ssh connect --proxy --cluster=abc-123 --auto-start-cluster=false --shutdown-delay=5m0s --keep-detached-for=2h0m0s", + 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: "with metadata", diff --git a/experimental/ssh/internal/client/policy_internal_test.go b/experimental/ssh/internal/client/policy_internal_test.go index cfe48ffa777..dde22f46187 100644 --- a/experimental/ssh/internal/client/policy_internal_test.go +++ b/experimental/ssh/internal/client/policy_internal_test.go @@ -28,21 +28,20 @@ func TestUsagePolicyMatches(t *testing.T) { func TestKeepDetachedMatches(t *testing.T) { tests := []struct { name string - stored int64 - requested int64 + stored bool + requested bool want bool }{ - {name: "empty request takes a lingering server", stored: 3600000, requested: 0, want: true}, - {name: "empty request takes a server that does not linger", stored: 0, requested: 0, want: true}, - {name: "equal durations match", stored: 3600000, requested: 3600000, want: true}, - {name: "different durations do not match", stored: 3600000, requested: 7200000, want: false}, - {name: "request against a server that does not linger does not match", stored: 0, requested: 3600000, want: false}, + {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(%d, %d) = %v, want %v", 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 d1e74d85c4e..b8812688388 100644 --- a/experimental/ssh/internal/client/ssh-server-bootstrap.py +++ b/experimental/ssh/internal/client/ssh-server-bootstrap.py @@ -16,6 +16,11 @@ # 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. @@ -29,7 +34,7 @@ dbutils.widgets.text("sessionId", "") dbutils.widgets.text("serverless", "false") dbutils.widgets.text("usagePolicyId", "") -dbutils.widgets.text("keepDetachedForSeconds", "0") +dbutils.widgets.text("keepDetachedProcesses", "false") def cleanup(): @@ -132,34 +137,34 @@ def detached_descendants(server_pgid): return sorted(survivors, key=int) -def wait_for_detached_descendants(server_pgid, timeout_seconds): +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 and - bounded by --keep-detached-for. + 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. """ - deadline = time.monotonic() + timeout_seconds + reported = None + reported_at = 0.0 while True: survivors = detached_descendants(server_pgid) if not survivors: print("No detached processes left, releasing the run", flush=True) return - if time.monotonic() > deadline: + now = time.monotonic() + if survivors != reported or now - reported_at >= LINGER_REPORT_SECONDS: print( - f"Reached the --keep-detached-for limit of {timeout_seconds}s with " - f"{len(survivors)} detached process(es) still running: {','.join(survivors)}. " - "Releasing the run; they lose /Workspace and /Volumes access from here.", + f"Holding the run open for {len(survivors)} detached process(es): {','.join(survivors)}", flush=True, ) - return - 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) @@ -210,7 +215,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_for_seconds = int(dbutils.widgets.get("keepDetachedForSeconds") or 0) + 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 @@ -261,10 +266,10 @@ def run_ssh_server(): if usage_policy_id: server_args.append(f"--usage-policy-id={usage_policy_id}") - # The server does not linger itself; it uses this to persist the mode for reconnects and - # to warn about detached work it is about to leave behind when the mode is off. - if keep_detached_for_seconds > 0: - server_args.append(f"--keep-detached-for={keep_detached_for_seconds}s") + # The server does not hold the run open itself; it uses this to persist the mode for + # reconnects and to warn about detached work it is about to leave behind when the mode is off. + 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 @@ -294,9 +299,9 @@ def run_ssh_server(): returncode = os.waitstatus_to_exitcode(reaped_statuses[proc.pid]) if returncode == -signal.SIGTERM: # A newer session's bootstrap terminates a server already running on this cluster - # (see cleanup), which a reconnect asking for a different --keep-detached-for 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 linger below still runs. + # (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. @@ -305,11 +310,11 @@ def run_ssh_server(): # 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 lingering would leave survivors alive but cut off from - # /Workspace, which is a worse failure than the one it fixes. + # 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_for_seconds > 0: - wait_for_detached_descendants(server_pgid, keep_detached_for_seconds) + if keep_detached_processes: + wait_for_detached_descendants(server_pgid) else: kill_all_children() diff --git a/experimental/ssh/internal/client/submit_internal_test.go b/experimental/ssh/internal/client/submit_internal_test.go index 4aec4ba11a2..d4247cbc2e1 100644 --- a/experimental/ssh/internal/client/submit_internal_test.go +++ b/experimental/ssh/internal/client/submit_internal_test.go @@ -72,20 +72,20 @@ 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) - // Zero is what tells the bootstrap to sweep detached work as it always has. - assert.Equal(t, "0", got.Tasks[0].NotebookTask.BaseParameters["keepDetachedForSeconds"]) + // 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, - KeepDetachedFor: 90 * time.Minute, + ClusterID: "abc-123", + ServerTimeout: 24 * time.Hour, + KeepDetachedProcesses: true, } got := buildSSHServerSubmitRun("v1", "scope", notebookPath, "", opts) - assert.Equal(t, "5400", got.Tasks[0].NotebookTask.BaseParameters["keepDetachedForSeconds"]) - // The linger happens inside the run, so the run's own timeout still bounds it. + 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) }) diff --git a/experimental/ssh/internal/server/server.go b/experimental/ssh/internal/server/server.go index 28c19f983c3..254884a72df 100644 --- a/experimental/ssh/internal/server/server.go +++ b/experimental/ssh/internal/server/server.go @@ -47,12 +47,12 @@ 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 - // KeepDetachedFor is how long the bootstrap notebook holds the job run open for - // detached processes after this server exits. Zero means it does not: the notebook - // sweeps them as it always has. The server does not linger itself; it only needs the - // value to persist it for reconnects and to decide whether to warn about work it is - // about to destroy. - KeepDetachedFor time.Duration + // KeepDetachedProcesses is whether the bootstrap notebook holds the job run open for + // detached processes after this server exits. False means it does not: the notebook + // sweeps them as it always has. The server does not hold the run open itself; it only + // needs the value to persist it for reconnects and to decide whether to warn about work + // it is about to destroy. + KeepDetachedProcesses bool // The directory to store sshd configuration ConfigDir string // The name of the secrets scope to use for client and server keys @@ -89,10 +89,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, - KeepDetachedForMs: opts.KeepDetachedFor.Milliseconds(), + Port: port, + ClusterID: opts.ClusterID, + UsagePolicyID: opts.UsagePolicyID, + KeepDetachedProcesses: opts.KeepDetachedProcesses, } err = workspace.SaveWorkspaceMetadata(ctx, client, opts.Version, opts.SessionID, metadata) if err != nil { @@ -162,14 +162,14 @@ func reportDetachedDescendants(ctx context.Context, opts ServerOptions, root str return } - // Warning, not info: without --keep-detached-for these processes do not outlive the + // 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. Serverless is excluded because the container teardown takes them // regardless, so the flag cannot help there and is rejected for it. - if len(pids) > 0 && opts.KeepDetachedFor == 0 && !opts.Serverless { + if len(pids) > 0 && !opts.KeepDetachedProcesses && !opts.Serverless { 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-for=\", which holds the run open for them.", + "\"databricks ssh connect --keep-detached-processes\", which holds the run open while they run.", len(pids), formatPids(pids)) } @@ -180,7 +180,7 @@ func reportDetachedDescendants(ctx context.Context, opts ServerOptions, root str telemetry.Log(ctx, protos.DatabricksCliLog{ SshTunnelTeardownEvent: &protos.SshTunnelTeardownEvent{ ComputeType: computeType, - KeepDetachedRequested: opts.KeepDetachedFor > 0, + KeepDetachedRequested: opts.KeepDetachedProcesses, HadDetachedDescendantsAtTeardown: len(pids) > 0, }, }) diff --git a/experimental/ssh/internal/server/teardown_test.go b/experimental/ssh/internal/server/teardown_test.go index 9ee64b730cc..758b30f3bbf 100644 --- a/experimental/ssh/internal/server/teardown_test.go +++ b/experimental/ssh/internal/server/teardown_test.go @@ -3,7 +3,6 @@ package server import ( "encoding/json" "testing" - "time" "github.com/databricks/cli/libs/cmdctx" "github.com/databricks/cli/libs/telemetry" @@ -35,12 +34,12 @@ func TestReportDetachedDescendantsWarning(t *testing.T) { 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-for") + 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{KeepDetachedFor: time.Hour} + opts := ServerOptions{KeepDetachedProcesses: true} reportDetachedDescendants(ctx, opts, procWithDetachedWork(t), testServerPid) assert.Empty(t, logs.String()) @@ -93,7 +92,7 @@ func TestReportDetachedDescendantsTelemetry(t *testing.T) { }, { name: "the run was held open for it", - opts: ServerOptions{KeepDetachedFor: 2 * time.Hour}, + opts: ServerOptions{KeepDetachedProcesses: true}, root: procWithDetachedWork, want: protos.SshTunnelTeardownEvent{ ComputeType: protos.SshTunnelComputeTypeDedicated, diff --git a/experimental/ssh/internal/workspace/workspace.go b/experimental/ssh/internal/workspace/workspace.go index 5da0938f86e..c67adb6bac4 100644 --- a/experimental/ssh/internal/workspace/workspace.go +++ b/experimental/ssh/internal/workspace/workspace.go @@ -22,11 +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"` - // KeepDetachedForMs records how long the server's bootstrap notebook will hold the job - // run open for detached processes after the server exits (--keep-detached-for), so a - // reconnect can tell whether a running server honours the requested duration. Zero, and - // so absent, when the session did not ask for it. - KeepDetachedForMs int64 `json:"keep_detached_for_ms,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/ssh_tunnel.go b/libs/telemetry/protos/ssh_tunnel.go index c35165de657..55d8ca30fda 100644 --- a/libs/telemetry/protos/ssh_tunnel.go +++ b/libs/telemetry/protos/ssh_tunnel.go @@ -151,7 +151,7 @@ type SshTunnelEvent struct { HasUsagePolicy bool `json:"has_usage_policy"` // Whether the connection asked for detached processes (tmux, setsid, nohup) to - // outlive the server via --keep-detached-for. Only the presence is recorded, not + // 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"` diff --git a/libs/telemetry/protos/ssh_tunnel_teardown.go b/libs/telemetry/protos/ssh_tunnel_teardown.go index 82de9adb06d..1c5c4fb48ea 100644 --- a/libs/telemetry/protos/ssh_tunnel_teardown.go +++ b/libs/telemetry/protos/ssh_tunnel_teardown.go @@ -5,7 +5,7 @@ package protos // 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-for flag addresses: only the +// 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. // @@ -17,12 +17,12 @@ type SshTunnelTeardownEvent struct { ComputeType SshTunnelComputeType `json:"compute_type,omitempty"` // Whether the session asked for detached processes to be kept via - // --keep-detached-for. Only the presence is recorded, not the duration. + // --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-for those processes do not survive the run, so this counts how + // --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"` } From b3c4b6c3c3b35974934708dcccb954cdbfa3e241 Mon Sep 17 00:00:00 2001 From: Anton Nekipelov <226657+anton-107@users.noreply.github.com> Date: Thu, 17 Sep 2026 21:09:36 +0000 Subject: [PATCH 4/6] ssh: support --keep-detached-processes in `ssh setup` The mode is fixed when the server job is submitted, and for a host configured by `ssh setup` the submitting invocation is always the persisted ProxyCommand - the same reason --max-clients and --server-timeout are carried there. Without this, an `ssh ` session could never ask to keep its detached work: it either inherited the mode from a server someone had already started with `ssh connect --keep-detached-processes`, or started one that swept the work. Setup already validates the ClientOptions it serializes, so the dedicated-only check applies here too; setup requires a cluster, so it cannot trip. Co-authored-by: Isaac --- .../cli/ssh-keep-detached-processes.md | 2 +- acceptance/ssh/setup/output.txt | 4 ++++ acceptance/ssh/setup/script | 8 ++++++- experimental/ssh/README.md | 4 ++++ experimental/ssh/cmd/setup.go | 21 +++++++++++-------- experimental/ssh/internal/setup/setup.go | 18 ++++++++++------ experimental/ssh/internal/setup/setup_test.go | 20 ++++++++++-------- 7 files changed, 51 insertions(+), 26 deletions(-) diff --git a/.nextchanges/cli/ssh-keep-detached-processes.md b/.nextchanges/cli/ssh-keep-detached-processes.md index 2e130ac0d9b..975d7b7d311 100644 --- a/.nextchanges/cli/ssh-keep-detached-processes.md +++ b/.nextchanges/cli/ssh-keep-detached-processes.md @@ -1 +1 @@ -* `ssh connect` now accepts 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)) +* `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/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 855fbbda623..2340000cfc3 100644 --- a/experimental/ssh/README.md +++ b/experimental/ssh/README.md @@ -103,6 +103,10 @@ know before using it: Dedicated clusters only. On serverless the container is torn down with the run, so survivors die regardless and the flag is rejected. +`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. diff --git a/experimental/ssh/cmd/setup.go b/experimental/ssh/cmd/setup.go index a8848e741f8..238d82832b0 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 processes detached from the SSH session (tmux, setsid, nohup) running after the tunnel shuts down. Holds the cluster up until they exit or --server-timeout elapses") 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/setup/setup.go b/experimental/ssh/internal/setup/setup.go index 2f55da1ca46..a54e94b8eda 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 the tunnel keeps processes detached from the SSH session running after it shuts + // down, will be added as a --keep-detached-processes flag 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) { From 54f45151b6e87f1ae78bf2a22944143d1e5e0333 Mon Sep 17 00:00:00 2001 From: Anton Nekipelov <226657+anton-107@users.noreply.github.com> Date: Fri, 18 Sep 2026 12:15:18 +0200 Subject: [PATCH 5/6] ssh: wait for child adoption before releasing the run (#6737) ## Changes Add a non-reaping `waitid` check before releasing the SSH bootstrap run, plus six regression scenarios wired into the SSH Go test suite. ## Why Targets #6387. After SIGTERM, server-group helpers can still parent detached work that the notebook has not adopted yet. An empty survivor scan must not release the run while children remain. The existing SIGCHLD handler retains ownership of exit statuses, and polling/report throttling stay unchanged. ## Tests - Regression cases fail on the unpatched parent (`b3c4b6c3c`) and pass with this fix. An isolated Linux process reproduction also verifies delayed adoption and exit-status preservation. - `./task fmt`, `./task checks`, `./task lint`, and `./task test-exp-ssh` pass. - `./task test` passes with `OMNIGENT` unset, cloud testing disabled, and an isolated `TMPDIR`; two unrelated Terraform acceptance timeouts pass on automatic retry. _This PR was written by Codex._ --- _This PR was created with [GitHub MCP](http://go/mcps)._ --- .../internal/client/ssh-server-bootstrap.py | 11 +++- .../client/ssh_server_bootstrap_test.go | 17 +++++ .../testdata/ssh_server_bootstrap_test.py | 63 +++++++++++++++++++ 3 files changed, 90 insertions(+), 1 deletion(-) create mode 100644 experimental/ssh/internal/client/ssh_server_bootstrap_test.go create mode 100644 experimental/ssh/internal/client/testdata/ssh_server_bootstrap_test.py diff --git a/experimental/ssh/internal/client/ssh-server-bootstrap.py b/experimental/ssh/internal/client/ssh-server-bootstrap.py index b8812688388..bef6b809fe3 100644 --- a/experimental/ssh/internal/client/ssh-server-bootstrap.py +++ b/experimental/ssh/internal/client/ssh-server-bootstrap.py @@ -137,6 +137,15 @@ def detached_descendants(server_pgid): 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. @@ -154,7 +163,7 @@ def wait_for_detached_descendants(server_pgid): reported_at = 0.0 while True: survivors = detached_descendants(server_pgid) - if not survivors: + if not survivors and not has_children(): print("No detached processes left, releasing the run", flush=True) return now = time.monotonic() 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/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() From 8e3b76f44d82470210c90df18b5fe0399116034f Mon Sep 17 00:00:00 2001 From: Anton Nekipelov <226657+anton-107@users.noreply.github.com> Date: Sat, 19 Sep 2026 12:10:33 +0000 Subject: [PATCH 6/6] ssh: keep server running while detached work remains --- experimental/ssh/README.md | 54 ++++---- experimental/ssh/cmd/connect.go | 2 +- experimental/ssh/cmd/server.go | 2 +- experimental/ssh/cmd/setup.go | 2 +- experimental/ssh/internal/client/client.go | 17 +-- .../ssh/internal/client/client_test.go | 10 +- .../internal/client/ssh-server-bootstrap.py | 3 +- .../ssh/internal/proxy/connections.go | 16 ++- .../internal/proxy/connections_idle_test.go | 29 +++++ .../ssh/internal/server/export_test.go | 11 +- experimental/ssh/internal/server/idle.go | 37 ++++++ experimental/ssh/internal/server/idle_test.go | 121 ++++++++++++++++++ experimental/ssh/internal/server/server.go | 23 ++-- .../ssh/internal/server/teardown_test.go | 5 +- experimental/ssh/internal/setup/setup.go | 4 +- 15 files changed, 268 insertions(+), 68 deletions(-) create mode 100644 experimental/ssh/internal/proxy/connections_idle_test.go create mode 100644 experimental/ssh/internal/server/idle.go create mode 100644 experimental/ssh/internal/server/idle_test.go diff --git a/experimental/ssh/README.md b/experimental/ssh/README.md index 2340000cfc3..90800d8987e 100644 --- a/experimental/ssh/README.md +++ b/experimental/ssh/README.md @@ -76,32 +76,34 @@ By default nothing outlives the session: when the last client disconnects, the s down after `--shutdown-delay` and the bootstrap notebook sweeps every process it parents, including work that was deliberately detached with `tmux`, `setsid` or `nohup`. -`databricks ssh connect --cluster= --keep-detached-processes` changes that. On teardown -the tunnel terminates only its own process group - the server and its `sshd` children - and -then holds the job run open for as long as any detached process is still running. Things to -know before using it: - -- **It holds the cluster up.** A `RUNNING` job run suppresses autotermination, so the cluster - keeps accruing DBUs until the last detached process exits. The bound is the run's own - lifetime, `--server-timeout` (24h by default), which is therefore the knob to reach for when - the cost is what matters; multi-day work still belongs in Jobs/DABs. Note also that - reconnecting starts a new run rather than rejoining the one being held open, so each session - with live detached work leaves its own run behind. -- **Releasing the run is not the end of the work.** Once the hold is released the cluster - starts its own autotermination countdown, which is the last thing that reaps survivors - - detached work does not count as cluster activity, however busy it is. -- **The notebook has to stay alive, not just the process.** Workspace filesystem access is - authorized by walking the live process tree for a registered ancestor, and the bootstrap - notebook is that ancestor. A detached process that outlives it keeps `/dbfs` and REST API - access but loses `/Workspace` and `/Volumes` with `EPERM` - which is why the group-scoped - teardown is tied to holding the run open and not enabled on its own. Work that finishes - while the run is held never sees this; work still running when `--server-timeout` expires - does, and from there every workspace path fails, including in a new window opened inside a - surviving `tmux`, because the `tmux` server - not the shell - is what lost its registered - ancestor. So size `--server-timeout` to the work you intend to leave behind. - -Dedicated clusters only. On serverless the container is torn down with the run, so survivors -die regardless and the flag is rejected. +`--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 diff --git a/experimental/ssh/cmd/connect.go b/experimental/ssh/cmd/connect.go index 278b0fc2016..9ac3af75464 100644 --- a/experimental/ssh/cmd/connect.go +++ b/experimental/ssh/cmd/connect.go @@ -67,7 +67,7 @@ Connect to a dedicated cluster: 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 processes detached from the SSH session (tmux, setsid, nohup) running after the tunnel shuts down. Holds the cluster up until they exit or --server-timeout elapses (dedicated clusters only)") + 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)") diff --git a/experimental/ssh/cmd/server.go b/experimental/ssh/cmd/server.go index 7dc1b555dc4..6bb0d4db821 100644 --- a/experimental/ssh/cmd/server.go +++ b/experimental/ssh/cmd/server.go @@ -46,7 +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, "Whether the bootstrap notebook holds the job run open for detached processes after the server exits") + 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. diff --git a/experimental/ssh/cmd/setup.go b/experimental/ssh/cmd/setup.go index 238d82832b0..261fd092d39 100644 --- a/experimental/ssh/cmd/setup.go +++ b/experimental/ssh/cmd/setup.go @@ -38,7 +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 processes detached from the SSH session (tmux, setsid, nohup) running after the tunnel shuts down. Holds the cluster up until they exit or --server-timeout elapses") + 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 { diff --git a/experimental/ssh/internal/client/client.go b/experimental/ssh/internal/client/client.go index 1fd5035bf4a..2e651e68d8a 100644 --- a/experimental/ssh/internal/client/client.go +++ b/experimental/ssh/internal/client/client.go @@ -130,11 +130,7 @@ type ClientOptions struct { AutoApprove bool // Id of the usage policy to use for the serverless SSH server job. Serverless only. UsagePolicyID string - // Whether the bootstrap notebook holds the job run open after the SSH server shuts down, - // for as long as processes detached from the session (tmux, setsid, nohup) keep running, - // so they keep their workspace filesystem access intact. False, the default, keeps - // today's behaviour: the run ends with the server and nothing outlives it. The run's own - // timeout (--server-timeout) is what bounds the hold. Dedicated clusters only. + // Whether detached processes prevent idle shutdown of the SSH server. Bounded by --server-timeout. KeepDetachedProcesses bool } @@ -148,11 +144,6 @@ func (o *ClientOptions) Validate() error { if o.UsagePolicyID != "" && o.ClusterID != "" { return errors.New("--usage-policy-id flag can only be used with serverless compute (--name flag)") } - // On serverless the container goes away with the run, so nothing survives the server - // however long the notebook holds the run open. - if o.KeepDetachedProcesses && o.ClusterID == "" { - return errors.New("--keep-detached-processes flag can only be used with a dedicated cluster (--cluster flag)") - } if o.Accelerator != "" && o.Accelerator != "GPU_1xA10" && o.Accelerator != "GPU_8xH100" { return fmt.Errorf("invalid accelerator value: %q, expected %q or %q", o.Accelerator, "GPU_1xA10", "GPU_8xH100") } @@ -261,9 +252,9 @@ func (o *ClientOptions) ToProxyCommand() (string, error) { } else { 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" - } + } + if o.KeepDetachedProcesses { + proxyCommand += " --keep-detached-processes" } // Both of these are fixed when the server job is submitted, and for a host configured by diff --git a/experimental/ssh/internal/client/client_test.go b/experimental/ssh/internal/client/client_test.go index ceb88898087..8363d2e5d0b 100644 --- a/experimental/ssh/internal/client/client_test.go +++ b/experimental/ssh/internal/client/client_test.go @@ -122,9 +122,8 @@ func TestValidate(t *testing.T) { opts: client.ClientOptions{ConnectionName: "my-conn", UsagePolicyID: "pol-1"}, }, { - name: "keep detached processes with serverless", - opts: client.ClientOptions{ConnectionName: "my-conn", KeepDetachedProcesses: true}, - wantErr: "--keep-detached-processes flag can only be used with a dedicated cluster (--cluster flag)", + name: "keep detached processes with serverless", + opts: client.ClientOptions{ConnectionName: "my-conn", KeepDetachedProcesses: true}, }, { name: "keep detached processes with cluster ID", @@ -350,6 +349,11 @@ func TestToProxyCommand(t *testing.T) { 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/ssh-server-bootstrap.py b/experimental/ssh/internal/client/ssh-server-bootstrap.py index bef6b809fe3..a3f20ef0d4d 100644 --- a/experimental/ssh/internal/client/ssh-server-bootstrap.py +++ b/experimental/ssh/internal/client/ssh-server-bootstrap.py @@ -275,8 +275,7 @@ def run_ssh_server(): if usage_policy_id: server_args.append(f"--usage-policy-id={usage_policy_id}") - # The server does not hold the run open itself; it uses this to persist the mode for - # reconnects and to warn about detached work it is about to leave behind when the mode is off. + # The server uses this to defer idle shutdown while detached work is running. if keep_detached_processes: server_args.append("--keep-detached-processes") 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/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 254884a72df..7ee1797d237 100644 --- a/experimental/ssh/internal/server/server.go +++ b/experimental/ssh/internal/server/server.go @@ -47,11 +47,7 @@ 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 is whether the bootstrap notebook holds the job run open for - // detached processes after this server exits. False means it does not: the notebook - // sweeps them as it always has. The server does not hold the run open itself; it only - // needs the value to persist it for reconnects and to decide whether to warn about work - // it is about to destroy. + // KeepDetachedProcesses prevents idle shutdown while detached processes are running. KeepDetachedProcesses bool // The directory to store sshd configuration ConfigDir string @@ -140,10 +136,18 @@ func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ServerOpt })) }() + idleErr := make(chan error, 1) + go func() { + idleErr <- waitForIdleShutdown(ctx, connections, opts.KeepDetachedProcesses, procRoot, os.Getpid()) + }() + select { case err := <-listenErr: return err - case <-connections.TimedOut: + 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)) @@ -164,12 +168,11 @@ func reportDetachedDescendants(ctx context.Context, opts ServerOptions, root str // 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. Serverless is excluded because the container teardown takes them - // regardless, so the flag cannot help there and is rejected for it. - if len(pids) > 0 && !opts.KeepDetachedProcesses && !opts.Serverless { + // 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 holds the run open while they run.", + "\"databricks ssh connect --keep-detached-processes\", which keeps the SSH server running while they run.", len(pids), formatPids(pids)) } diff --git a/experimental/ssh/internal/server/teardown_test.go b/experimental/ssh/internal/server/teardown_test.go index 758b30f3bbf..6e61b364ebf 100644 --- a/experimental/ssh/internal/server/teardown_test.go +++ b/experimental/ssh/internal/server/teardown_test.go @@ -45,12 +45,11 @@ func TestReportDetachedDescendantsWarning(t *testing.T) { assert.Empty(t, logs.String()) }) - // The flag is rejected for serverless, so pointing at it there would be misleading. - t.Run("stays quiet on serverless", func(t *testing.T) { + t.Run("warns on serverless", func(t *testing.T) { ctx, logs := captureWarnLogs(t.Context()) reportDetachedDescendants(ctx, ServerOptions{Serverless: true}, procWithDetachedWork(t), testServerPid) - assert.Empty(t, logs.String()) + assert.Contains(t, logs.String(), "--keep-detached-processes") }) t.Run("stays quiet when nothing was left behind", func(t *testing.T) { diff --git a/experimental/ssh/internal/setup/setup.go b/experimental/ssh/internal/setup/setup.go index a54e94b8eda..bdfd88cc3c3 100644 --- a/experimental/ssh/internal/setup/setup.go +++ b/experimental/ssh/internal/setup/setup.go @@ -30,8 +30,8 @@ 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 the tunnel keeps processes detached from the SSH session running after it shuts - // down, will be added as a --keep-detached-processes flag to the ProxyCommand. Fixed at + // 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