Repository navigation
[SSH] Check port_fwd_serverless_enabled before starting serverless compute - #7063
Conversation
Integration test reportCommit: a84b403
Top 6 slowest tests (at least 2 minutes):
|
rugpanov
left a comment
rugpanov
left a comment
There was a problem hiding this comment.
LGTM. I left three non-blocking suggestions inline.
| // environment. A read that fails therefore lets the connection go ahead, and the proxy's own 403 | ||
| // (errDriverProxyForbidden) remains the backstop. | ||
| func checkPortForwardingToServerless(ctx context.Context, client *databricks.WorkspaceClient) error { | ||
| setting, err := client.WorkspaceSettingsV2.GetPublicWorkspaceSetting(ctx, settingsv2.GetPublicWorkspaceSettingRequest{ |
There was a problem hiding this comment.
[nice to have] This read uses the default SDK client, which retries 429 and 503 responses until its retry budget runs out. If the settings service is slow, an optional check can block a working connection for minutes, and the ProxyCommand child runs it again. Use newProbeAPIClient from agentshim.go, or give the call a short context timeout and continue when it expires.
There was a problem hiding this comment.
IMO this should definitely be fixed, if not in this PR then in a followup
There was a problem hiding this comment.
Fixed in 1ea406f. The read now goes through newProbeAPIClient, so a 429 or 503 fails on the first attempt and the connection continues. I added a throttled case to TestCheckPortForwardingToServerless that asserts the setting is read exactly once.
| cmdio.LogString(ctx, fmt.Sprintf("Connecting to %s...", sessionID)) | ||
| } | ||
|
|
||
| if opts.IsServerlessMode() { |
There was a problem hiding this comment.
[nice to have] spawnSSHClient and ensureSSHConfigEntry write a ProxyCommand with --metadata, and that child process runs Run again. So one ssh connect sends the GET twice, and each IDE reconnect sends it again. With --metadata, nothing is uploaded or submitted, so the check prevents nothing. In production, the read always fails today. Run the check only when opts.IsServerlessMode() && opts.ServerMetadata == "", as the ValidateClusterAccess comment does for proxy mode.
There was a problem hiding this comment.
Fixed in 1ea406f. The check now runs only when opts.IsServerlessMode() && opts.ServerMetadata == "", so the ProxyCommand child and IDE reconnects no longer send the GET.
|
|
||
| if opts.IsServerlessMode() { | ||
| if err := checkPortForwardingToServerless(ctx, client); err != nil { | ||
| return err |
There was a problem hiding this comment.
[nice to have] This early return does not set outcome.errorCategory, so logSshTunnelEvent records UNKNOWN. The comment in Run says each failing step sets a category, and UNKNOWN is meant to signal CLI bugs. Add a category such as PORT_FORWARDING_DISABLED, which needs the matching proto change, and use it for the #7057 403 path too. If you add it in a follow-up PR, please say so here.
…mpute A serverless ssh connect now reads the port_fwd_serverless_enabled workspace setting first, in proxy mode too. When an administrator has turned it off, connect stops before it creates the secret scope, uploads binaries or submits the server run, and says which setting is off and who can change it. The driver proxy's gate also depends on an enforcement flag and a port allow-list the CLI cannot see, and the setting is not served through the public API in every environment (production today). A failed read therefore lets the connection go ahead, and the 403 handling from the previous change stays the backstop. The acceptance test for that path now stubs the settings read the way production answers it. DECO-28889 Co-authored-by: Isaac <no-reply@databricks.com>
The setting check used the default SDK client, which retries 429 and 503 until its retry budget runs out, so a slow settings service could hold up a working connection for minutes. It now reads through the no-retry probe client the AI Gateway check already uses. The check is also skipped when reconnecting with --metadata. The ProxyCommand child runs Run again with the server already running, so the read prevented nothing and was repeated on every IDE reconnect. Co-authored-by: Isaac <no-reply@databricks.com>
1ea406f to
a84b403
Compare
5 participants
Changes
A serverless
databricks ssh connect(CPU or--accelerator) now reads theport_fwd_serverless_enabledworkspace setting first, in ProxyCommand mode too. When an administrator has turned it off, connect stops before it creates the secret scope, uploads binaries or submits the server run:checkPortForwardingToServerlessreads the setting withWorkspaceSettingsV2.GetPublicWorkspaceSettingand refuses only when the read succeeds andeffective_boolean_valisfalse.libs/testserveranswers the settings GET, defaultingport_fwd_serverless_enabledtotruelike the real setting.Why
Stacked on #7057, PR 2 of DECO-28889. #7057 makes connect stop on the driver proxy's 403, but by then it has uploaded binaries and started serverless compute. The setting is what the proxy enforces, so checking it first avoids both.
The CLI can't see everything the proxy's gate (
PortForwardingToServerlessGate) checks. The gate also needs its enforcement flag on and skips allow-listed ports, and the setting is only served through the public API in dev and staging (public_api_stage: STAGING; production answersSetting port_fwd_serverless_enabled is not enabled for public API). So a failed read is never an error, and the 403 handling stays as the backstop. The acceptance test for that backstop now stubs the settings read the way production answers it.Tests
TestCheckPortForwardingToServerlesscovers disabled, enabled, no effective value, and a failed read. Plus newlibs/testserver/settings_test.gotests.acceptance/ssh/connect-serverless-port-forwarding-disabledchecks the error and that connect sent no requests besides reads.connect-serverless-feature-disabledkeeps covering the 403 path, with the settings read failing. Allacceptance/sshtests pass.--accelerator GPU_1xA10, and ProxyCommand mode (ssh -o ProxyCommand="databricks ssh connect --proxy --name …") all stop within 1 s with the error above.This pull request and its description were written by Isaac.