Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
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
13 changes: 12 additions & 1 deletion experimental/ssh/internal/client/client.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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{
{
Expand Down Expand Up @@ -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.
Expand Down
30 changes: 21 additions & 9 deletions experimental/ssh/internal/client/client_internal_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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,
},
},
{
Expand Down
5 changes: 5 additions & 0 deletions libs/telemetry/protos/ssh_tunnel.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"`
Expand Down
Loading