diff --git a/.nextchanges/cli/ssh-environment-version.md b/.nextchanges/cli/ssh-environment-version.md new file mode 100644 index 00000000000..7dbf46a2f8d --- /dev/null +++ b/.nextchanges/cli/ssh-environment-version.md @@ -0,0 +1 @@ +* `databricks ssh connect` now documents `--environment-version` for choosing the serverless environment version, starts a separate default-named server for each version, and rejects the flag together with `--cluster`. Default connection names also include the environment version in their hash, so the default name when `--environment-version` is not specified changes with this release. Configuration associated with the old connection name (secret scope, SSH keys and known hosts in your local `.databricks` folder) will no longer be used and can be safely removed. ([#7310](https://github.com/databricks/cli/pull/7310)) diff --git a/experimental/ssh/cmd/connect.go b/experimental/ssh/cmd/connect.go index 9ac3af75464..d2040ad760f 100644 --- a/experimental/ssh/cmd/connect.go +++ b/experimental/ssh/cmd/connect.go @@ -36,6 +36,7 @@ Connect to serverless: databricks ssh connect databricks ssh connect --accelerator= # AI Runtime databricks ssh connect --base-environment= # custom base environment + databricks ssh connect --environment-version=5 # serverless environment version Connect to a dedicated cluster: databricks ssh connect --cluster=`, @@ -93,8 +94,7 @@ Connect to a dedicated cluster: cmd.Flags().BoolVar(&skipSettingsCheck, "skip-settings-check", false, "Skip checking and updating IDE settings") cmd.Flags().MarkHidden("skip-settings-check") - cmd.Flags().IntVar(&environmentVersion, "environment-version", defaultEnvironmentVersion, "Environment version for AI Runtime") - cmd.Flags().MarkHidden("environment-version") + cmd.Flags().IntVar(&environmentVersion, "environment-version", client.DefaultEnvironmentVersion, "Serverless environment version for the SSH server job (serverless only)") cmd.Flags().StringVar(&baseEnvironment, "base-environment", "", "Custom base environment for serverless compute: an env.yaml path, a workspace-base-environments resource ID, or a display name") @@ -114,8 +114,16 @@ Connect to a dedicated cluster: cmd.RunE = func(cmd *cobra.Command, args []string) error { ctx := cmd.Context() wsClient := cmdctx.WorkspaceClient(ctx) + // Generate the name before resetting the version so that the default version + // maps to the same name whether it is set explicitly or implicitly. if connectionName == "" && clusterID == "" && !proxyMode { - connectionName = client.GenerateDefaultConnectionName(wsClient.Config.Host, accelerator, baseEnvironment) + connectionName = client.GenerateDefaultConnectionName(wsClient.Config.Host, accelerator, baseEnvironment, environmentVersion) + } + // Only carry an explicitly-set environment version. Leaving it at 0 otherwise + // lets the submit path default to DefaultEnvironmentVersion and lets Validate + // detect a real --environment-version + --base-environment conflict. + if !cmd.Flags().Changed("environment-version") { + environmentVersion = 0 } // Serverless GPU compute can take much longer to provision than CPU compute, // so allow extra time for the SSH server job to start. @@ -123,12 +131,6 @@ Connect to a dedicated cluster: if accelerator != "" { startupTimeout = gpuTaskStartupTimeout } - // Only carry an explicitly-set environment version. Leaving it at 0 otherwise - // lets the submit path default to minEnvironmentVersion and lets Validate - // detect a real --environment-version + --base-environment conflict. - if !cmd.Flags().Changed("environment-version") { - environmentVersion = 0 - } opts := client.ClientOptions{ Profile: wsClient.Config.Profile, ClusterID: clusterID, diff --git a/experimental/ssh/cmd/constants.go b/experimental/ssh/cmd/constants.go index 001ee7f23a6..6216b7ca3e0 100644 --- a/experimental/ssh/cmd/constants.go +++ b/experimental/ssh/cmd/constants.go @@ -11,8 +11,7 @@ const ( // alive. Matches the keepalive interval of the vite bridge (libs/apps/vite/bridge.go), and sits // well under both the ~9 minutes after which idle sessions were observed to drop and the // 30 second SSH-level keepalive that was verified to prevent it. - defaultKeepaliveInterval = 20 * time.Second - defaultEnvironmentVersion = 4 + defaultKeepaliveInterval = 20 * time.Second // 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 diff --git a/experimental/ssh/internal/client/client.go b/experimental/ssh/internal/client/client.go index 80142d54d4c..0d482eef8a2 100644 --- a/experimental/ssh/internal/client/client.go +++ b/experimental/ssh/internal/client/client.go @@ -1,6 +1,7 @@ package client import ( + "cmp" "context" "crypto/md5" _ "embed" @@ -63,6 +64,9 @@ var errDriverProxyForbidden = errors.New("the workspace refused the SSH tunnel c var connectionNameRegex = regexp.MustCompile(`^[a-zA-Z0-9][a-zA-Z0-9_-]*$`) +// DefaultEnvironmentVersion is the serverless environment version used when --environment-version is not set. +const DefaultEnvironmentVersion = 4 + const ( sshServerTaskKey = "start_ssh_server" serverlessEnvironmentKey = "ssh_tunnel_serverless" @@ -170,7 +174,7 @@ func (o *ClientOptions) Validate() error { if o.IDE != "" && o.IDE != vscode.VSCodeOption && o.IDE != vscode.CursorOption { return fmt.Errorf("invalid IDE value: %q, expected %q or %q", o.IDE, vscode.VSCodeOption, vscode.CursorOption) } - if o.EnvironmentVersion > 0 && o.EnvironmentVersion < minEnvironmentVersion { + if o.EnvironmentVersion != 0 && o.EnvironmentVersion < minEnvironmentVersion { return fmt.Errorf("environment version must be >= %d, got %d", minEnvironmentVersion, o.EnvironmentVersion) } // base_environment and environment_version are mutually exclusive in the SDK, @@ -181,6 +185,9 @@ func (o *ClientOptions) Validate() error { if o.BaseEnvironment != "" && o.ClusterID != "" { return errors.New("--base-environment can only be used with serverless compute") } + if o.EnvironmentVersion > 0 && o.ClusterID != "" { + return errors.New("--environment-version can only be used with serverless compute") + } // A server started with fewer than one client slot rejects every connection with an // opaque websocket handshake failure, so catch it here instead. if o.MaxClients < 1 { @@ -201,17 +208,22 @@ func (o *ClientOptions) Validate() error { } // GenerateDefaultConnectionName creates a deterministic connection name from -// the workspace host, accelerator type, and base environment. The name includes +// the workspace host, accelerator type, and environment. The name includes // a hash so that different workspaces produce different names (avoiding SSH // known_hosts conflicts). The environment is folded into the hash because a // serverless server bakes in its environment at startup: distinct environments // must map to distinct connection names so they don't reuse each other's server. -func GenerateDefaultConnectionName(host, accelerator, baseEnvironment string) string { - // Keep the hash host-only when no base environment is set so existing default - // connection names are preserved. - hashInput := host - if baseEnvironment != "" { +func GenerateDefaultConnectionName(host, accelerator, baseEnvironment string, environmentVersion int) string { + // The version is hashed even when it is the default, so bumping + // DefaultEnvironmentVersion intentionally moves users to a new server. + var hashInput string + switch { + case baseEnvironment != "": hashInput = host + "\x00" + baseEnvironment + default: + // The key keeps a version from hashing like a base environment + // whose name is that number. + hashInput = host + "\x00environment_version=" + strconv.Itoa(environmentVersion) } h := md5.Sum([]byte(hashInput)) hashStr := hex.EncodeToString(h[:4]) @@ -842,7 +854,7 @@ func buildSSHServerSubmitRun(version, secretScopeName, jobNotebookPath, baseEnvi if baseEnvironment != "" { spec.BaseEnvironment = baseEnvironment } else { - spec.EnvironmentVersion = strconv.Itoa(max(opts.EnvironmentVersion, minEnvironmentVersion)) + spec.EnvironmentVersion = strconv.Itoa(cmp.Or(opts.EnvironmentVersion, DefaultEnvironmentVersion)) } submitRequest.Environments = []jobs.JobEnvironment{ { diff --git a/experimental/ssh/internal/client/client_test.go b/experimental/ssh/internal/client/client_test.go index 8363d2e5d0b..20bc56d0115 100644 --- a/experimental/ssh/internal/client/client_test.go +++ b/experimental/ssh/internal/client/client_test.go @@ -90,9 +90,19 @@ func TestValidate(t *testing.T) { opts: client.ClientOptions{ClusterID: "abc-123", EnvironmentVersion: 3}, wantErr: "environment version must be >= 4, got 3", }, + { + name: "negative environment version", + opts: client.ClientOptions{ConnectionName: "my-conn", EnvironmentVersion: -1}, + wantErr: "environment version must be >= 4, got -1", + }, { name: "valid environment version", - opts: client.ClientOptions{ClusterID: "abc-123", EnvironmentVersion: 4}, + opts: client.ClientOptions{ConnectionName: "my-conn", EnvironmentVersion: 4}, + }, + { + name: "environment version with cluster", + opts: client.ClientOptions{ClusterID: "abc-123", EnvironmentVersion: 5}, + wantErr: "--environment-version can only be used with serverless compute", }, { name: "base environment with environment version", @@ -228,35 +238,35 @@ func TestGenerateDefaultConnectionName(t *testing.T) { { name: "no accelerator", host: "https://my-workspace.cloud.databricks.com", - want: "databricks-cpu-961dabbd", + want: "databricks-cpu-2db656b0", }, { name: "GPU_1xA10 accelerator", host: "https://my-workspace.cloud.databricks.com", accelerator: "GPU_1xA10", - want: "databricks-gpu-1xa10-961dabbd", + want: "databricks-gpu-1xa10-2db656b0", }, { name: "GPU_8xH100 accelerator", host: "https://my-workspace.cloud.databricks.com", accelerator: "GPU_8xH100", - want: "databricks-gpu-8xh100-961dabbd", + want: "databricks-gpu-8xh100-2db656b0", }, { name: "different host produces different name", host: "https://other-workspace.cloud.databricks.com", - want: "databricks-cpu-e8a8ec19", + want: "databricks-cpu-ebd5d001", }, { name: "deterministic for same input", host: "https://my-workspace.cloud.databricks.com", - want: "databricks-cpu-961dabbd", + want: "databricks-cpu-2db656b0", }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - got := client.GenerateDefaultConnectionName(tt.host, tt.accelerator, "") + got := client.GenerateDefaultConnectionName(tt.host, tt.accelerator, "", client.DefaultEnvironmentVersion) assert.Equal(t, tt.want, got) }) } @@ -266,12 +276,32 @@ func TestGenerateDefaultConnectionName(t *testing.T) { // ignored when an existing server for the default name is reused). t.Run("base environment differentiates the name", func(t *testing.T) { const host = "https://my-workspace.cloud.databricks.com" - base := client.GenerateDefaultConnectionName(host, "", "") - withEnv := client.GenerateDefaultConnectionName(host, "", "my-env") - otherEnv := client.GenerateDefaultConnectionName(host, "", "other-env") + base := client.GenerateDefaultConnectionName(host, "", "", client.DefaultEnvironmentVersion) + withEnv := client.GenerateDefaultConnectionName(host, "", "my-env", client.DefaultEnvironmentVersion) + otherEnv := client.GenerateDefaultConnectionName(host, "", "other-env", client.DefaultEnvironmentVersion) assert.NotEqual(t, base, withEnv, "setting --base-environment must change the default name") assert.NotEqual(t, withEnv, otherEnv, "different environments must produce different names") - assert.Equal(t, withEnv, client.GenerateDefaultConnectionName(host, "", "my-env"), "must be deterministic") + assert.Equal(t, withEnv, client.GenerateDefaultConnectionName(host, "", "my-env", client.DefaultEnvironmentVersion), "must be deterministic") + }) + + // Same for an explicit environment version: without it, a connection asking for + // version 5 would reuse a running server started on the default version. + t.Run("environment version differentiates the name", func(t *testing.T) { + const host = "https://my-workspace.cloud.databricks.com" + base := client.GenerateDefaultConnectionName(host, "", "", client.DefaultEnvironmentVersion) + v5 := client.GenerateDefaultConnectionName(host, "", "", 5) + v6 := client.GenerateDefaultConnectionName(host, "", "", 6) + assert.NotEqual(t, base, v5, "setting --environment-version must change the default name") + assert.NotEqual(t, v5, v6, "different versions must produce different names") + assert.NotEqual(t, v5, client.GenerateDefaultConnectionName(host, "", "5", client.DefaultEnvironmentVersion), "a version must not collide with a base environment named like it") + assert.Equal(t, v5, client.GenerateDefaultConnectionName(host, "", "", 5), "must be deterministic") + }) + + t.Run("base environment takes precedence over environment version", func(t *testing.T) { + const host = "https://my-workspace.cloud.databricks.com" + withEnv := client.GenerateDefaultConnectionName(host, "", "my-env", 4) + assert.Equal(t, withEnv, client.GenerateDefaultConnectionName(host, "", "my-env", 5)) + assert.NotEqual(t, withEnv, client.GenerateDefaultConnectionName(host, "", "other-env", 4)) }) } @@ -288,7 +318,7 @@ func TestGenerateDefaultConnectionNameMatchesRegex(t *testing.T) { for _, host := range hosts { for _, acc := range accelerators { for _, env := range environments { - name := client.GenerateDefaultConnectionName(host, acc, env) + name := client.GenerateDefaultConnectionName(host, acc, env, client.DefaultEnvironmentVersion) assert.Regexp(t, nameRegex, name, "host=%q accelerator=%q environment=%q name=%q", host, acc, env, name) } } diff --git a/experimental/ssh/internal/client/submit_internal_test.go b/experimental/ssh/internal/client/submit_internal_test.go index d4247cbc2e1..f56043bfbb7 100644 --- a/experimental/ssh/internal/client/submit_internal_test.go +++ b/experimental/ssh/internal/client/submit_internal_test.go @@ -33,6 +33,22 @@ func TestBuildSSHServerSubmitRun(t *testing.T) { assert.Nil(t, got.Tasks[0].Compute) }) + t.Run("serverless environment version", func(t *testing.T) { + for _, tt := range []struct { + version int + want string + }{ + {0, "4"}, + {5, "5"}, + } { + opts := ClientOptions{ConnectionName: "conn", ServerTimeout: time.Hour, EnvironmentVersion: tt.version} + got := buildSSHServerSubmitRun("v1", "scope", notebookPath, "", opts) + + require.Len(t, got.Environments, 1) + assert.Equal(t, tt.want, got.Environments[0].Spec.EnvironmentVersion) + } + }) + t.Run("serverless with accelerator", func(t *testing.T) { opts := ClientOptions{ ConnectionName: "conn",