Skip to content

Make --accelerator flag optional for ssh connect #4947

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
anton-107 merged 2 commits into from
Apr 14, 2026
+10 −10
Merged

Make --accelerator flag optional for ssh connect #4947

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
368e1e6
Make `--accelerator` flag optional for `ssh connect` in serverless mode
anton-107 Apr 13, 2026
253f6ab
Fix gofmt import ordering
anton-107 Apr 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 8 additions & 7 deletions experimental/ssh/internal/client/client.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,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"
"github.com/fatih/color"
"github.com/gorilla/websocket"
)

Expand Down Expand Up @@ -105,16 +106,9 @@ func (o *ClientOptions) Validate() error {
if o.Accelerator != "" && o.ConnectionName == "" {
return errors.New("--accelerator flag can only be used with serverless compute (--name flag)")
}
// Consider removing this check when we enable serverless CPU connections. Ideally Jobs API should do the validation
// for us, but they don't plan on doing it in the nearest future. For now we should not forget to check if there are
// any other possible values that can be here.
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")
}
// TODO: Remove when we add support for serverless CPU
if o.ConnectionName != "" && o.Accelerator == "" {
return errors.New("--name flag requires --accelerator to be set (for now we only support serverless GPU compute)")
}
if o.ConnectionName != "" && !connectionNameRegex.MatchString(o.ConnectionName) {
return fmt.Errorf("connection name %q must consist of letters, numbers, dashes, and underscores", o.ConnectionName)
}
Expand Down Expand Up @@ -215,6 +209,9 @@ func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ClientOpt

if !opts.ProxyMode {
cmdio.LogString(ctx, fmt.Sprintf("Connecting to %s...", sessionID))
if opts.IsServerlessMode() && opts.Accelerator == "" {
cmdio.LogString(ctx, color.YellowString("WARNING: serverless compute without an accelerator is in private preview. If you are not enrolled, this command will likely time out with an error. Contact your Databricks account team to enroll."))
}
}

if opts.IDE != "" && !opts.ProxyMode {
Expand Down Expand Up @@ -294,6 +291,10 @@ func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ClientOpt
}
userName, serverPort, clusterID, err = ensureSSHServerIsRunning(ctx, client, version, secretScopeName, opts)
if err != nil {
if opts.IsServerlessMode() && opts.Accelerator == "" && errors.Is(err, errServerMetadata) {
return fmt.Errorf("failed to ensure that ssh server is running: %w\n\n"+
color.YellowString("This may be because serverless compute without an accelerator is in private preview.\nContact your Databricks account team to enroll."), err)
}
return fmt.Errorf("failed to ensure that ssh server is running: %w", err)
}
} else {
Expand Down
5 changes: 2 additions & 3 deletions experimental/ssh/internal/client/client_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,9 +36,8 @@ func TestValidate(t *testing.T) {
wantErr: "--accelerator flag can only be used with serverless compute (--name flag)",
},
{
name: "connection name without accelerator",
opts: client.ClientOptions{ConnectionName: "my-conn"},
wantErr: "--name flag requires --accelerator to be set (for now we only support serverless GPU compute)",
name: "connection name without accelerator",
opts: client.ClientOptions{ConnectionName: "my-conn"},
},
{
name: "invalid connection name characters",
Expand Down
Loading