From cb0937ae59f98486bb2c83e39980d8d99e9e0a6d Mon Sep 17 00:00:00 2001 From: Rada Kamysheva Date: Thu, 9 Jul 2026 12:45:17 +0000 Subject: [PATCH 1/2] acc: run ssh/connect-serverless-gpu locally Flip acceptance/ssh/connect-serverless-gpu to Local = true so it runs against the in-process testserver (it was disabled on both cloud and local in #4838). A full SSH handshake can't complete locally (no real compute/sshd + a wss driver-proxy), so the local run drives the connect flow up to the connection attempt and asserts the serverless GPU bootstrap job the CLI submits (accelerator, environment, secret scope). Teach libs/testserver enough of the SSH tunnel backend for the flow to get there deterministically and fast: - secrets/list (ListSecretsByScope), returning RESOURCE_DOES_NOT_EXIST for a missing scope so CreateKeysSecretScope creates it - synthesize metadata.json on the ssh-server-bootstrap run submit, as a real tunnel server would publish it, so getServerMetadata succeeds - driver-proxy metadata/logs endpoints - workspace/export now 404s a missing object instead of returning a nil body (which the response normalizer rejects) The connect attempt itself is expected to fail locally; its noisy output goes to LOG.stderr (excluded from the golden files). --- .../ssh/connect-serverless-gpu/out.stdout.txt | 1 - .../ssh/connect-serverless-gpu/out.test.toml | 2 +- .../ssh/connect-serverless-gpu/output.txt | 40 ++++++++++++++++++ acceptance/ssh/connect-serverless-gpu/script | 28 +++++++++++-- .../ssh/connect-serverless-gpu/test.toml | 10 ++++- libs/testserver/handlers.go | 26 ++++++++++++ libs/testserver/jobs.go | 41 +++++++++++++++++++ libs/testserver/secrets.go | 35 ++++++++++++++++ 8 files changed, 177 insertions(+), 6 deletions(-) diff --git a/acceptance/ssh/connect-serverless-gpu/out.stdout.txt b/acceptance/ssh/connect-serverless-gpu/out.stdout.txt index 41cae5e7d16..e69de29bb2d 100644 --- a/acceptance/ssh/connect-serverless-gpu/out.stdout.txt +++ b/acceptance/ssh/connect-serverless-gpu/out.stdout.txt @@ -1 +0,0 @@ -Connection successful diff --git a/acceptance/ssh/connect-serverless-gpu/out.test.toml b/acceptance/ssh/connect-serverless-gpu/out.test.toml index ab8691ddb9b..c777e3ce206 100644 --- a/acceptance/ssh/connect-serverless-gpu/out.test.toml +++ b/acceptance/ssh/connect-serverless-gpu/out.test.toml @@ -1,4 +1,4 @@ -Local = false +Local = true Cloud = false RequiresUnityCatalog = true CloudEnvs.gcp = false diff --git a/acceptance/ssh/connect-serverless-gpu/output.txt b/acceptance/ssh/connect-serverless-gpu/output.txt index e69de29bb2d..b7013545a8b 100644 --- a/acceptance/ssh/connect-serverless-gpu/output.txt +++ b/acceptance/ssh/connect-serverless-gpu/output.txt @@ -0,0 +1,40 @@ + +>>> print_requests.py //api/2.2/jobs/runs/submit +{ + "method": "POST", + "path": "/api/2.2/jobs/runs/submit", + "body": { + "environments": [ + { + "environment_key": "ssh_tunnel_serverless", + "spec": { + "environment_version": "4" + } + } + ], + "run_name": "ssh-server-bootstrap-serverless-gpu-test", + "tasks": [ + { + "compute": { + "hardware_accelerator": "GPU_1xA10" + }, + "environment_key": "ssh_tunnel_serverless", + "notebook_task": { + "base_parameters": { + "authorizedKeySecretName": "client-public-key", + "maxClients": "10", + "secretScopeName": "[USERNAME]-serverless-gpu-test-ssh-tunnel-keys", + "serverless": "true", + "sessionId": "serverless-gpu-test", + "shutdownDelay": "10m0s", + "version": "[CLI_VERSION]" + }, + "notebook_path": "/Workspace/Users/[USERNAME]/.databricks/ssh-tunnel/[CLI_VERSION]/serverless-gpu-test/ssh-server-bootstrap" + }, + "task_key": "start_ssh_server", + "timeout_seconds": 86400 + } + ], + "timeout_seconds": 86400 + } +} diff --git a/acceptance/ssh/connect-serverless-gpu/script b/acceptance/ssh/connect-serverless-gpu/script index 8c8538f6e32..7835de376ad 100644 --- a/acceptance/ssh/connect-serverless-gpu/script +++ b/acceptance/ssh/connect-serverless-gpu/script @@ -1,6 +1,28 @@ +if [ -z "${CLOUD_ENV:-}" ]; then + # Local runs have no cloud-built release artifacts and cannot download them + # from GitHub. The uploaded tunnel binaries are never executed by a real + # compute in a local run, so stand in empty release archives (named for the + # dev build, which omits the version from the archive name) to satisfy + # --releases-dir and keep the upload step on the local filesystem. + mkdir -p releases + : >releases/databricks_cli_linux_amd64.zip + : >releases/databricks_cli_linux_arm64.zip + CLI_RELEASES_DIR="$PWD/releases" +fi + errcode $CLI ssh connect --name serverless-gpu-test --accelerator=GPU_1xA10 --releases-dir=$CLI_RELEASES_DIR -- "echo 'Connection successful'" >out.stdout.txt 2>LOG.stderr -if ! grep -q "Connection successful" out.stdout.txt; then - run_id=$(cat LOG.stderr | grep -o "Job submitted successfully with run ID: [0-9]*" | grep -o "[0-9]*$") - trace $CLI jobs get-run "$run_id" > LOG.job +if [ -n "${CLOUD_ENV:-}" ]; then + # On cloud a real serverless GPU connection runs the remote command and prints + # its output; on failure, dump the bootstrap run to help diagnose. + if ! grep -q "Connection successful" out.stdout.txt; then + run_id=$(cat LOG.stderr | grep -o "Job submitted successfully with run ID: [0-9]*" | grep -o "[0-9]*$") + trace $CLI jobs get-run "$run_id" > LOG.job + fi +else + # A local run has no real compute, so the SSH handshake cannot complete: the + # connect attempt fails after the bootstrap job is submitted (see LOG.stderr). + # The local test's job is to verify the CLI submits the right serverless GPU + # bootstrap job, so assert that from the recorded request. + trace print_requests.py //api/2.2/jobs/runs/submit fi diff --git a/acceptance/ssh/connect-serverless-gpu/test.toml b/acceptance/ssh/connect-serverless-gpu/test.toml index fa125d9bb5d..24b0870236f 100644 --- a/acceptance/ssh/connect-serverless-gpu/test.toml +++ b/acceptance/ssh/connect-serverless-gpu/test.toml @@ -1,9 +1,17 @@ -Local = false +Local = true Cloud = false # Serverless GPU is only available in newer environments RequiresUnityCatalog = true +# Assert the serverless GPU bootstrap job the CLI submits on a local run. +RecordRequests = true + +# Local runs stage stand-in release archives here; not part of the golden output. +Ignore = [ + "releases", +] + # Serverless GPU is not available in GCP yet [CloudEnvs] gcp = false diff --git a/libs/testserver/handlers.go b/libs/testserver/handlers.go index 802baa0118a..6fb2a210994 100644 --- a/libs/testserver/handlers.go +++ b/libs/testserver/handlers.go @@ -105,6 +105,17 @@ func AddDefaultHandlers(server *Server) { path := req.URL.Query().Get("path") data := req.Workspace.WorkspaceExport(path) + // A missing object returns 404, matching the real API. Without this the + // handler would return a nil body (the zero FileEntry's Data), which the + // response normalizer rejects. `ssh connect` relies on this 404 to detect + // that the server's metadata.json has not been published yet. + if data == nil { + return Response{ + StatusCode: 404, + Body: map[string]string{"message": fmt.Sprintf("Path (%s) doesn't exist.", path)}, + } + } + // The filer reads the raw object body via ?direct_download=true, while // the SDK's Workspace.Export (used by `databricks workspace export`) // requests JSON and expects the base64-encoded content field. @@ -730,6 +741,21 @@ func AddDefaultHandlers(server *Server) { return req.Workspace.SecretsGet(req) }) + server.Handle("GET", "/api/2.0/secrets/list", func(req Request) any { + return req.Workspace.SecretsList(req) + }) + + // The SSH tunnel server sits behind the workspace driver proxy. `ssh connect` + // GETs .../metadata to read the remote login user before opening the tunnel, + // and best-effort GETs .../logs to surface recent server errors on failure. + server.Handle("GET", "/driver-proxy-api/o/{workspace_id}/{cluster_id}/{port}/metadata", func(req Request) any { + return Response{Body: sshTunnelRemoteUser} + }) + + server.Handle("GET", "/driver-proxy-api/o/{workspace_id}/{cluster_id}/{port}/logs", func(req Request) any { + return Response{Body: ""} + }) + // Secrets ACLs: server.Handle("GET", "/api/2.0/secrets/acls/get", func(req Request) any { return req.Workspace.SecretsAclsGet(req) diff --git a/libs/testserver/jobs.go b/libs/testserver/jobs.go index c3537d00187..615074226c9 100644 --- a/libs/testserver/jobs.go +++ b/libs/testserver/jobs.go @@ -15,6 +15,7 @@ import ( "github.com/databricks/databricks-sdk-go/service/compute" "github.com/databricks/databricks-sdk-go/service/jobs" + "github.com/databricks/databricks-sdk-go/service/workspace" ) const missingJobGitProviderMessage = "git_source.git_provider must be one of: github,gitlab,bitbucketcloud,gitlabenterpriseedition,bitbucketserver,azuredevopsservices,githubenterprise,awscodecommit" @@ -519,9 +520,49 @@ func (s *FakeWorkspace) JobsSubmit(req Request) Response { Tasks: tasks, } + // The SSH tunnel bootstrap run stands in for the server the CLI starts on + // compute. That server publishes a metadata.json (its port and cluster ID) + // to the workspace, which `ssh connect` polls before connecting. No server + // runs in a local test, so synthesize that file so the connect flow proceeds. + if strings.HasPrefix(runName, sshTunnelBootstrapRunPrefix) { + s.writeSSHTunnelMetadata(request) + } + return Response{Body: jobs.SubmitRunResponse{RunId: runId}} } +const ( + sshTunnelBootstrapRunPrefix = "ssh-server-bootstrap-" + sshTunnelBootstrapNotebook = "ssh-server-bootstrap" + sshTunnelServerPort = 7772 + sshTunnelClusterID = "1234-567890-serverless" + sshTunnelRemoteUser = "spark" +) + +// writeSSHTunnelMetadata publishes the metadata.json that a real SSH tunnel +// server writes to the workspace on startup, next to the bootstrap notebook. +// Callers must already hold the workspace lock. +func (s *FakeWorkspace) writeSSHTunnelMetadata(request jobs.SubmitRun) { + for _, t := range request.Tasks { + if t.NotebookTask == nil { + continue + } + // The bootstrap notebook and metadata.json share the session directory. + metadataPath := strings.TrimSuffix(t.NotebookTask.NotebookPath, sshTunnelBootstrapNotebook) + "metadata.json" + metadata, err := json.Marshal(map[string]any{ + "port": sshTunnelServerPort, + "cluster_id": sshTunnelClusterID, + }) + if err != nil { + continue + } + s.files[metadataPath] = FileEntry{ + Info: workspace.ObjectInfo{ObjectType: "FILE", Path: metadataPath}, + Data: metadata, + } + } +} + // executePythonWheelTask runs a python wheel task locally using uv. // For tasks using existing_cluster_id, the venv is cached per cluster to match // cloud behavior where libraries are cached on running clusters. diff --git a/libs/testserver/secrets.go b/libs/testserver/secrets.go index 1f298ea9f0c..31e8d398629 100644 --- a/libs/testserver/secrets.go +++ b/libs/testserver/secrets.go @@ -4,6 +4,7 @@ import ( "encoding/base64" "encoding/json" "fmt" + "slices" "github.com/databricks/databricks-sdk-go/service/workspace" ) @@ -40,6 +41,40 @@ func (s *FakeWorkspace) SecretsPut(req Request) Response { return Response{} } +// SecretsList models GET /api/2.0/secrets/list (ListSecretsByScope). A missing +// scope must return RESOURCE_DOES_NOT_EXIST so the SDK maps it to +// apierr.ErrResourceDoesNotExist; `ssh connect`'s CreateKeysSecretScope relies on +// that sentinel to decide whether to create the scope rather than fail. +func (s *FakeWorkspace) SecretsList(req Request) Response { + defer s.LockUnlock()() + + scope := req.URL.Query().Get("scope") + + if _, exists := s.SecretScopes[scope]; !exists { + return Response{ + StatusCode: 404, + Body: map[string]string{"error_code": "RESOURCE_DOES_NOT_EXIST", "message": fmt.Sprintf("Scope %s does not exist", scope)}, + } + } + + keys := make([]string, 0, len(s.Secrets[scope])) + for key := range s.Secrets[scope] { + keys = append(keys, key) + } + slices.Sort(keys) + + secrets := make([]workspace.SecretMetadata, 0, len(keys)) + for _, key := range keys { + secrets = append(secrets, workspace.SecretMetadata{Key: key}) + } + + return Response{ + Body: workspace.ListSecretsResponse{ + Secrets: secrets, + }, + } +} + func (s *FakeWorkspace) SecretsGet(req Request) Response { defer s.LockUnlock()() From 50973572c3f91363a4805f88eca695c8619b3f05 Mon Sep 17 00:00:00 2001 From: Rada Kamysheva Date: Thu, 9 Jul 2026 13:30:37 +0000 Subject: [PATCH 2/2] acc: tighten testserver/ssh comments --- acceptance/ssh/connect-serverless-gpu/script | 16 +++++----------- libs/testserver/handlers.go | 11 ++++------- libs/testserver/jobs.go | 12 ++++-------- libs/testserver/secrets.go | 6 ++---- 4 files changed, 15 insertions(+), 30 deletions(-) diff --git a/acceptance/ssh/connect-serverless-gpu/script b/acceptance/ssh/connect-serverless-gpu/script index 7835de376ad..e7e95e443d0 100644 --- a/acceptance/ssh/connect-serverless-gpu/script +++ b/acceptance/ssh/connect-serverless-gpu/script @@ -1,9 +1,6 @@ if [ -z "${CLOUD_ENV:-}" ]; then - # Local runs have no cloud-built release artifacts and cannot download them - # from GitHub. The uploaded tunnel binaries are never executed by a real - # compute in a local run, so stand in empty release archives (named for the - # dev build, which omits the version from the archive name) to satisfy - # --releases-dir and keep the upload step on the local filesystem. + # Local runs have no release artifacts and never execute the uploaded + # binaries, so stand in empty archives (dev name omits the version). mkdir -p releases : >releases/databricks_cli_linux_amd64.zip : >releases/databricks_cli_linux_arm64.zip @@ -13,16 +10,13 @@ fi errcode $CLI ssh connect --name serverless-gpu-test --accelerator=GPU_1xA10 --releases-dir=$CLI_RELEASES_DIR -- "echo 'Connection successful'" >out.stdout.txt 2>LOG.stderr if [ -n "${CLOUD_ENV:-}" ]; then - # On cloud a real serverless GPU connection runs the remote command and prints - # its output; on failure, dump the bootstrap run to help diagnose. + # On cloud the connection runs the remote command; dump the run on failure. if ! grep -q "Connection successful" out.stdout.txt; then run_id=$(cat LOG.stderr | grep -o "Job submitted successfully with run ID: [0-9]*" | grep -o "[0-9]*$") trace $CLI jobs get-run "$run_id" > LOG.job fi else - # A local run has no real compute, so the SSH handshake cannot complete: the - # connect attempt fails after the bootstrap job is submitted (see LOG.stderr). - # The local test's job is to verify the CLI submits the right serverless GPU - # bootstrap job, so assert that from the recorded request. + # Locally the handshake can't complete (no compute); just assert the CLI + # submitted the right serverless GPU bootstrap job. trace print_requests.py //api/2.2/jobs/runs/submit fi diff --git a/libs/testserver/handlers.go b/libs/testserver/handlers.go index 6fb2a210994..338eb7e3e50 100644 --- a/libs/testserver/handlers.go +++ b/libs/testserver/handlers.go @@ -105,10 +105,8 @@ func AddDefaultHandlers(server *Server) { path := req.URL.Query().Get("path") data := req.Workspace.WorkspaceExport(path) - // A missing object returns 404, matching the real API. Without this the - // handler would return a nil body (the zero FileEntry's Data), which the - // response normalizer rejects. `ssh connect` relies on this 404 to detect - // that the server's metadata.json has not been published yet. + // A missing object returns 404, matching the real API; returning the nil + // body otherwise trips the response normalizer. if data == nil { return Response{ StatusCode: 404, @@ -745,9 +743,8 @@ func AddDefaultHandlers(server *Server) { return req.Workspace.SecretsList(req) }) - // The SSH tunnel server sits behind the workspace driver proxy. `ssh connect` - // GETs .../metadata to read the remote login user before opening the tunnel, - // and best-effort GETs .../logs to surface recent server errors on failure. + // SSH tunnel server behind the driver proxy: /metadata returns the remote + // login user, /logs is a best-effort error tail fetched on failure. server.Handle("GET", "/driver-proxy-api/o/{workspace_id}/{cluster_id}/{port}/metadata", func(req Request) any { return Response{Body: sshTunnelRemoteUser} }) diff --git a/libs/testserver/jobs.go b/libs/testserver/jobs.go index 615074226c9..2d019c3e496 100644 --- a/libs/testserver/jobs.go +++ b/libs/testserver/jobs.go @@ -520,10 +520,8 @@ func (s *FakeWorkspace) JobsSubmit(req Request) Response { Tasks: tasks, } - // The SSH tunnel bootstrap run stands in for the server the CLI starts on - // compute. That server publishes a metadata.json (its port and cluster ID) - // to the workspace, which `ssh connect` polls before connecting. No server - // runs in a local test, so synthesize that file so the connect flow proceeds. + // No tunnel server runs locally, so synthesize the metadata.json it would + // publish; `ssh connect` polls for it before connecting. if strings.HasPrefix(runName, sshTunnelBootstrapRunPrefix) { s.writeSSHTunnelMetadata(request) } @@ -539,15 +537,13 @@ const ( sshTunnelRemoteUser = "spark" ) -// writeSSHTunnelMetadata publishes the metadata.json that a real SSH tunnel -// server writes to the workspace on startup, next to the bootstrap notebook. -// Callers must already hold the workspace lock. +// writeSSHTunnelMetadata publishes the metadata.json a real tunnel server would +// write next to the bootstrap notebook. Callers must hold the workspace lock. func (s *FakeWorkspace) writeSSHTunnelMetadata(request jobs.SubmitRun) { for _, t := range request.Tasks { if t.NotebookTask == nil { continue } - // The bootstrap notebook and metadata.json share the session directory. metadataPath := strings.TrimSuffix(t.NotebookTask.NotebookPath, sshTunnelBootstrapNotebook) + "metadata.json" metadata, err := json.Marshal(map[string]any{ "port": sshTunnelServerPort, diff --git a/libs/testserver/secrets.go b/libs/testserver/secrets.go index 31e8d398629..1b737be0d61 100644 --- a/libs/testserver/secrets.go +++ b/libs/testserver/secrets.go @@ -41,10 +41,8 @@ func (s *FakeWorkspace) SecretsPut(req Request) Response { return Response{} } -// SecretsList models GET /api/2.0/secrets/list (ListSecretsByScope). A missing -// scope must return RESOURCE_DOES_NOT_EXIST so the SDK maps it to -// apierr.ErrResourceDoesNotExist; `ssh connect`'s CreateKeysSecretScope relies on -// that sentinel to decide whether to create the scope rather than fail. +// SecretsList models GET /api/2.0/secrets/list. A missing scope must return +// RESOURCE_DOES_NOT_EXIST so callers can branch on apierr.ErrResourceDoesNotExist. func (s *FakeWorkspace) SecretsList(req Request) Response { defer s.LockUnlock()()