[SSH] Make --environment-version flag public

.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))

experimental/ssh/cmd/connect.go

@@ -36,6 +36,7 @@ Connect to serverless:
   databricks ssh connect
   databricks ssh connect --accelerator=<GPU_type>   # AI Runtime
   databricks ssh connect --base-environment=<name>  # custom base environment
+  databricks ssh connect --environment-version=5    # serverless environment version
 
 Connect to a dedicated cluster:
   databricks ssh connect --cluster=<cluster-id>`,
@@ -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,21 +114,23 @@ 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.
 		startupTimeout := taskStartupTimeout
 		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,

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

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{
 			{

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)
 			}
 		}

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",