From b99d10c8b0ae1d97c226aa97159c64f99c598da6 Mon Sep 17 00:00:00 2001 From: TanishqDatabricks Date: Thu, 2 Jul 2026 16:33:08 +0200 Subject: [PATCH] experimental/ssh: wait full startup timeout for server health check The post-RUNNING health check against the driver-proxy /metadata endpoint used a fixed 30 x 2s = 60s window. After the bootstrap task reaches RUNNING the driver proxy still answers 503 until the container's HTTP endpoint is reachable; with a custom --base-environment that install happens post-RUNNING, so warmup routinely outlasts 60s and `ssh connect` fails with a driver-proxy 503 even though the server comes up fine. Poll with retries.Poll for the same accelerator-aware budget already used for the task to start (10m CPU / 45m GPU), backing off between attempts instead of hammering the proxy every 2s during a long wait. Extract formatServerStartError so a terminated bootstrap job returns its failure trace once, while only a timeout appends run diagnostics. Co-authored-by: Isaac --- NEXT_CHANGELOG.md | 2 + experimental/ssh/internal/client/client.go | 38 ++++++++++++------- .../internal/client/client_internal_test.go | 32 ++++++++++++++++ 3 files changed, 59 insertions(+), 13 deletions(-) diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index 30a64c54100..d633d477642 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -6,6 +6,8 @@ ### CLI +* Wait for the SSH server health check to pass for the full startup timeout (10 minutes, or 45 minutes with `--accelerator`) instead of a fixed 60 seconds, so `databricks ssh connect` no longer fails with a driver-proxy 503 while a custom `--base-environment` finishes installing ([#5807](https://github.com/databricks/cli/pull/5807)). + ### Bundles ### Dependency updates diff --git a/experimental/ssh/internal/client/client.go b/experimental/ssh/internal/client/client.go index 9fc20aa56df..844a3b96f69 100644 --- a/experimental/ssh/internal/client/client.go +++ b/experimental/ssh/internal/client/client.go @@ -1063,28 +1063,28 @@ func ensureSSHServerIsRunning(ctx context.Context, client *databricks.WorkspaceC sp := cmdio.NewSpinner(ctx, cmdio.WithElapsedTime()) defer sp.Close() sp.Update("Waiting for the SSH server to start...") - maxRetries := 30 - for retries := range maxRetries { - if ctx.Err() != nil { - return "", 0, "", ctx.Err() - } + // After the task reaches RUNNING the driver proxy still answers /metadata with 503 + // until the container's HTTP endpoint is reachable; with a custom base environment + // that install happens post-RUNNING, so warmup can outlast a short fixed window. Poll + // for the same budget we allowed for the task to start (accelerator-aware), backing + // off between attempts rather than hammering the proxy every 2s during a long wait. + _, pollErr := retries.Poll(ctx, opts.TaskStartupTimeout, func() (*struct{}, *retries.Err) { serverPort, userName, effectiveClusterID, err = getServerMetadata(ctx, client, sessionID, clusterID, version, opts.Liteswap) if err == nil { - cmdio.LogString(ctx, "Health check successful, starting ssh WebSocket connection...") - break + return &struct{}{}, nil } // The metadata never appears if the bootstrap job dies after reaching RUNNING. // Surface the job's actual error instead of waiting out the full timeout with a // generic "metadata.json doesn't exist" message. if failure, terminated := runFailureIfTerminated(ctx, client, runID); terminated { - return "", 0, "", fmt.Errorf("ssh server bootstrap job failed:\n%s", failure) - } - if retries < maxRetries-1 { - time.Sleep(2 * time.Second) - } else { - return "", 0, "", fmt.Errorf("failed to start the ssh server: %w\n%s", err, describeRunFailure(ctx, client, runID)) + return nil, retries.Halt(fmt.Errorf("ssh server bootstrap job failed:\n%s", failure)) } + return nil, retries.Continue(fmt.Errorf("waiting for ssh server health check: %w", err)) + }) + if pollErr != nil { + return "", 0, "", formatServerStartError(ctx, client, runID, pollErr) } + cmdio.LogString(ctx, "Health check successful, starting ssh WebSocket connection...") } else if err != nil { return "", 0, "", err } @@ -1092,6 +1092,18 @@ func ensureSSHServerIsRunning(ctx context.Context, client *databricks.WorkspaceC return userName, serverPort, effectiveClusterID, nil } +// formatServerStartError turns a failed health-check poll into the error returned to the user. +// A halted poll (the bootstrap job terminated) already carries the job's failure details, so it +// is returned as-is. A timeout carries only the last generic health-check error, so the run's +// diagnostics are appended. Splitting on the timeout case avoids printing the failure trace twice. +func formatServerStartError(ctx context.Context, client *databricks.WorkspaceClient, runID int64, pollErr error) error { + var timedOut *retries.ErrTimedOut + if errors.As(pollErr, &timedOut) { + return fmt.Errorf("failed to start the ssh server: %w\n%s", pollErr, describeRunFailure(ctx, client, runID)) + } + return pollErr +} + func logSshTunnelEvent(ctx context.Context, opts ClientOptions, isSuccess, isReconnect bool, serverStartTimeMs int64) { computeType := protos.SshTunnelComputeTypeDedicated if opts.IsServerlessMode() { diff --git a/experimental/ssh/internal/client/client_internal_test.go b/experimental/ssh/internal/client/client_internal_test.go index d2bfde57833..7017e9842ed 100644 --- a/experimental/ssh/internal/client/client_internal_test.go +++ b/experimental/ssh/internal/client/client_internal_test.go @@ -1,12 +1,14 @@ package client import ( + "fmt" "strings" "testing" "time" "github.com/databricks/cli/libs/cmdio" "github.com/databricks/databricks-sdk-go/experimental/mocks" + "github.com/databricks/databricks-sdk-go/retries" "github.com/databricks/databricks-sdk-go/service/environments" "github.com/databricks/databricks-sdk-go/service/jobs" "github.com/stretchr/testify/assert" @@ -193,6 +195,36 @@ func TestRunFailureIfTerminated(t *testing.T) { }) } +func TestFormatServerStartErrorTimeoutAppendsRunFailure(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + m := mocks.NewMockWorkspaceClient(t) + api := m.GetMockJobsAPI() + api.EXPECT().GetRun(mock.Anything, jobs.GetRunRequest{RunId: 1}).Return( + terminatedRun(1, 99, "Could not reach driver of cluster 0605-x.", "https://example.test/run/1"), nil) + api.EXPECT().GetRunOutput(mock.Anything, jobs.GetRunOutputRequest{RunId: 99}).Return( + &jobs.RunOutput{}, nil) + + pollErr := &retries.ErrTimedOut{} + err := formatServerStartError(ctx, m.WorkspaceClient, 1, pollErr) + require.Error(t, err) + assert.Contains(t, err.Error(), "failed to start the ssh server") + // The run diagnostics are fetched and appended for the timeout case. + assert.Contains(t, err.Error(), "Could not reach driver of cluster 0605-x.") +} + +func TestFormatServerStartErrorHaltReturnedAsIs(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + // A halted poll already carries the job failure details, so formatServerStartError must not + // fetch the run again (no GetRun expectation) nor re-append the trace. + m := mocks.NewMockWorkspaceClient(t) + pollErr := fmt.Errorf("ssh server bootstrap job failed:\n%s", "Could not reach driver of cluster 0605-x.") + err := formatServerStartError(ctx, m.WorkspaceClient, 1, pollErr) + require.Error(t, err) + assert.Equal(t, pollErr.Error(), err.Error()) + assert.Equal(t, 1, strings.Count(err.Error(), "Could not reach driver of cluster 0605-x.")) + assert.NotContains(t, err.Error(), "failed to start the ssh server") +} + func TestWaitForJobToStartSurfacesFailure(t *testing.T) { ctx := cmdio.MockDiscard(t.Context()) m := mocks.NewMockWorkspaceClient(t)