diff --git a/experimental/ssh/internal/client/client.go b/experimental/ssh/internal/client/client.go index 80142d54d4c..a31a5a075a9 100644 --- a/experimental/ssh/internal/client/client.go +++ b/experimental/ssh/internal/client/client.go @@ -226,6 +226,16 @@ func (o *ClientOptions) IsServerlessMode() bool { return o.ClusterID == "" && o.ConnectionName != "" } +// submittedEnvironmentVersion returns the serverless environment version the server job is +// submitted with, or 0 when it sets none: dedicated clusters have no environment, and a custom +// base environment carries its own version. +func (o *ClientOptions) submittedEnvironmentVersion() int { + if !o.IsServerlessMode() || o.BaseEnvironment != "" { + return 0 + } + return max(o.EnvironmentVersion, minEnvironmentVersion) +} + // SessionIdentifier returns the unique identifier for the session. // For dedicated clusters, this is the cluster ID. For serverless, this is the connection name. func (o *ClientOptions) SessionIdentifier() string { @@ -842,7 +852,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(opts.submittedEnvironmentVersion()) } submitRequest.Environments = []jobs.JobEnvironment{ { @@ -1595,6 +1605,7 @@ func buildSshTunnelEvent(opts ClientOptions, outcome connectOutcome) *protos.Ssh ServerStartTimeMs: outcome.serverStartTimeMs, IsSuccess: outcome.isSuccess, HasBaseEnvironment: opts.BaseEnvironment != "", + EnvironmentVersion: opts.submittedEnvironmentVersion(), HasUsagePolicy: opts.UsagePolicyID != "", // The connect side can only report that the knob was asked for. Whether any detached // process was there to keep is reported by the server, at teardown. diff --git a/experimental/ssh/internal/client/client_internal_test.go b/experimental/ssh/internal/client/client_internal_test.go index 1db7adbd88a..452a8162b26 100644 --- a/experimental/ssh/internal/client/client_internal_test.go +++ b/experimental/ssh/internal/client/client_internal_test.go @@ -483,27 +483,39 @@ func TestBuildSshTunnelEvent(t *testing.T) { name: "serverless with accelerator", opts: ClientOptions{ConnectionName: "my-conn", Accelerator: "GPU_1xA10"}, want: protos.SshTunnelEvent{ - ComputeType: protos.SshTunnelComputeTypeServerless, - AcceleratorType: "GPU_1xA10", - ClientMode: protos.SshTunnelClientModeSSH, + ComputeType: protos.SshTunnelComputeTypeServerless, + AcceleratorType: "GPU_1xA10", + ClientMode: protos.SshTunnelClientModeSSH, + EnvironmentVersion: 4, }, }, { name: "proxy mode takes precedence over IDE", opts: ClientOptions{ConnectionName: "my-conn", ProxyMode: true, IDE: "vscode"}, want: protos.SshTunnelEvent{ - ComputeType: protos.SshTunnelComputeTypeServerless, - IdeType: "vscode", - ClientMode: protos.SshTunnelClientModeProxy, + ComputeType: protos.SshTunnelComputeTypeServerless, + IdeType: "vscode", + ClientMode: protos.SshTunnelClientModeProxy, + EnvironmentVersion: 4, }, }, { name: "IDE mode", opts: ClientOptions{ConnectionName: "my-conn", IDE: "cursor"}, want: protos.SshTunnelEvent{ - ComputeType: protos.SshTunnelComputeTypeServerless, - IdeType: "cursor", - ClientMode: protos.SshTunnelClientModeIDE, + ComputeType: protos.SshTunnelComputeTypeServerless, + IdeType: "cursor", + ClientMode: protos.SshTunnelClientModeIDE, + EnvironmentVersion: 4, + }, + }, + { + name: "explicit serverless environment version", + opts: ClientOptions{ConnectionName: "my-conn", EnvironmentVersion: 6}, + want: protos.SshTunnelEvent{ + ComputeType: protos.SshTunnelComputeTypeServerless, + ClientMode: protos.SshTunnelClientModeSSH, + EnvironmentVersion: 6, }, }, { diff --git a/libs/telemetry/protos/ssh_tunnel.go b/libs/telemetry/protos/ssh_tunnel.go index 708a453ac63..e4f77f91056 100644 --- a/libs/telemetry/protos/ssh_tunnel.go +++ b/libs/telemetry/protos/ssh_tunnel.go @@ -156,6 +156,11 @@ type SshTunnelEvent struct { // Whether the connection was successful. IsSuccess bool `json:"is_success"` + // Serverless environment version the SSH server job is submitted with: the + // --environment-version value or the CLI default. Unset for dedicated clusters and + // for a custom base environment, which carries its own version. + EnvironmentVersion int `json:"environment_version,omitempty"` + // Whether a serverless usage policy was set via --usage-policy-id. // Only the presence is recorded, not the policy ID itself. HasUsagePolicy bool `json:"has_usage_policy"`