Skip to content

Commit a7e0dbe

Browse files
cardyokcodex
cardyok
and
codex
committed
feat(session): add single-connection transport
Co-Authored-By: Codex <noreply@openai.com>
1 parent f0e21ed commit a7e0dbe

31 files changed

Lines changed: 1983 additions & 32 deletions

‎cmd/gpud/command/command.go‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,15 @@ gpud scan
3939
sudo gpud up
4040
`
4141

42+
func sessionProtocolFlag() cli.Flag {
43+
return &cli.StringFlag{
44+
Name: "session-protocol",
45+
Usage: "control-plane session transport: v1, v2, or auto",
46+
Value: pkgconfig.DefaultSessionProtocol,
47+
EnvVar: "GPUD_SESSION_PROTOCOL",
48+
}
49+
}
50+
4251
func App() *cli.App {
4352
app := cli.NewApp()
4453

@@ -125,6 +134,7 @@ nohup sudo gpud run &>> <your log file path> &
125134
Usage: "(optional) endpoint for checking in",
126135
Value: "gpud-manager-prod01.dgxc-lepton.nvidia.com",
127136
},
137+
sessionProtocolFlag(),
128138
cli.StringFlag{
129139
Name: "gpu-count",
130140
Usage: "(optional) specify count of gpu (leave empty to auto-detect)",
@@ -268,6 +278,7 @@ sudo rm /etc/systemd/system/gpud.service
268278
Usage: "skips processing session updateConfig requests (testing only)",
269279
Hidden: true,
270280
},
281+
sessionProtocolFlag(),
271282

272283
&cli.DurationFlag{
273284
Name: "metrics-retention-period, retention-period",

‎cmd/gpud/command/command_test.go‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,46 @@
11
package command
22

33
import (
4+
"flag"
45
"strings"
56
"testing"
67

78
"github.com/stretchr/testify/require"
89
"github.com/urfave/cli"
910
)
1011

12+
func TestAppRunAndUpHaveSessionProtocolFlag(t *testing.T) {
13+
t.Parallel()
14+
15+
app := App()
16+
wantCommands := map[string]bool{"run": false, "up": false}
17+
for _, cmd := range app.Commands {
18+
if _, ok := wantCommands[cmd.Name]; !ok {
19+
continue
20+
}
21+
wantCommands[cmd.Name] = true
22+
23+
var protocolFlag cli.Flag
24+
for _, candidate := range cmd.Flags {
25+
if candidate.GetName() == "session-protocol" {
26+
protocolFlag = candidate
27+
break
28+
}
29+
}
30+
require.NotNilf(t, protocolFlag, "command %q is missing --session-protocol", cmd.Name)
31+
32+
set := flag.NewFlagSet(cmd.Name, flag.ContinueOnError)
33+
protocolFlag.Apply(set)
34+
require.NoError(t, set.Parse([]string{"--session-protocol", "v2"}))
35+
ctx := cli.NewContext(app, set, nil)
36+
require.Equal(t, "v2", ctx.String("session-protocol"))
37+
}
38+
39+
for cmdName, found := range wantCommands {
40+
require.Truef(t, found, "expected command %q to exist", cmdName)
41+
}
42+
}
43+
1144
func TestAppHasInfinibandExcludeDevicesFlag(t *testing.T) {
1245
t.Parallel()
1346

‎cmd/gpud/run/command.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -144,6 +144,7 @@ func Command(cliContext *cli.Context) error {
144144
versionFileSet := cliContext.IsSet("version-file")
145145
pluginSpecsFile := cliContext.String("plugin-specs-file")
146146
skipSessionUpdateConfig := cliContext.Bool("skip-session-update-config")
147+
sessionProtocol := cliContext.String("session-protocol")
147148

148149
ibClassRootDir := cliContext.String("infiniband-class-root-dir")
149150
ibExcludeDevicesStr := cliContext.String("infiniband-exclude-devices")
@@ -360,6 +361,7 @@ func Command(cliContext *cli.Context) error {
360361

361362
cfg.PluginSpecsFile = pluginSpecsFile
362363
cfg.SkipSessionUpdateConfig = skipSessionUpdateConfig
364+
cfg.SessionProtocol = sessionProtocol
363365

364366
if components != "" {
365367
cfg.Components = strings.Split(components, ",")

‎cmd/gpud/up/command.go‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -130,7 +130,7 @@ func Command(cliContext *cli.Context) (retErr error) {
130130
// The env file only contains: --log-level, --log-file, --endpoint, --data-dir, --db-in-memory
131131
// Session credentials are stored in the persistent state file, not the systemd env file.
132132

133-
if err := systemdInit(endpoint, dataDir, dbInMemory); err != nil {
133+
if err := systemdInit(endpoint, dataDir, dbInMemory, cliContext.String("session-protocol")); err != nil {
134134
return err
135135
}
136136
log.Logger.Debugw("successfully started systemd init")
@@ -178,11 +178,11 @@ func recordLoginSuccessState(ctx context.Context, dataDir string) error {
178178
return nil
179179
}
180180

181-
func systemdInit(endpoint string, dataDir string, dbInMemory bool) error {
181+
func systemdInit(endpoint string, dataDir string, dbInMemory bool, sessionProtocol ...string) error {
182182
// Always create/overwrite env file (consistent with v0.8.0 behavior).
183183
// IMPORTANT: The --token flag is NEVER written to the env file.
184184
// Only runtime configuration flags are written: --log-level, --log-file, --endpoint, --data-dir, --db-in-memory
185-
if err := systemd.CreateDefaultEnvFile(endpoint, dataDir, dbInMemory); err != nil {
185+
if err := systemd.CreateDefaultEnvFile(endpoint, dataDir, dbInMemory, sessionProtocol...); err != nil {
186186
return err
187187
}
188188
systemdUnitFileData := systemd.GPUdServiceUnitFileContents()

‎cmd/gpud/up/mock_command_additional_test.go‎

Lines changed: 11 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@ func newMockeyCLIContext(t *testing.T, args []string) *cli.Context {
3939
_ = flags.String("data-dir", "", "")
4040
_ = flags.String("token", "", "")
4141
_ = flags.String("endpoint", "", "")
42+
_ = flags.String("session-protocol", config.DefaultSessionProtocol, "")
4243
_ = flags.String("machine-id", "", "")
4344
_ = flags.Bool("machine-id-overwrite", false, "")
4445
_ = flags.Bool("refresh-session-token", false, "")
@@ -113,7 +114,7 @@ func TestCommand_LoginSuccessWithRecordStateError(t *testing.T) {
113114
}).Build()
114115
mockey.Mock(pkdsystemd.SystemctlExists).To(func() bool { return true }).Build()
115116
mockey.Mock(systemd.DefaultBinExists).To(func() bool { return true }).Build()
116-
mockey.Mock(systemd.CreateDefaultEnvFile).To(func(endpoint, dataDir string, dbInMemory bool) error {
117+
mockey.Mock(systemd.CreateDefaultEnvFile).To(func(endpoint, dataDir string, dbInMemory bool, sessionProtocol ...string) error {
117118
return nil
118119
}).Build()
119120
mockey.Mock(systemd.GPUdServiceUnitFileContents).To(func() string {
@@ -165,7 +166,7 @@ func TestCommand_WithAllFlags(t *testing.T) {
165166
}).Build()
166167
mockey.Mock(pkdsystemd.SystemctlExists).To(func() bool { return true }).Build()
167168
mockey.Mock(systemd.DefaultBinExists).To(func() bool { return true }).Build()
168-
mockey.Mock(systemd.CreateDefaultEnvFile).To(func(endpoint, dataDir string, dbInMemory bool) error {
169+
mockey.Mock(systemd.CreateDefaultEnvFile).To(func(endpoint, dataDir string, dbInMemory bool, sessionProtocol ...string) error {
169170
return nil
170171
}).Build()
171172
mockey.Mock(systemd.GPUdServiceUnitFileContents).To(func() string {
@@ -275,7 +276,7 @@ func TestCommand_WithEmptyToken(t *testing.T) {
275276
}).Build()
276277
mockey.Mock(pkdsystemd.SystemctlExists).To(func() bool { return true }).Build()
277278
mockey.Mock(systemd.DefaultBinExists).To(func() bool { return true }).Build()
278-
mockey.Mock(systemd.CreateDefaultEnvFile).To(func(endpoint, dataDir string, dbInMemory bool) error {
279+
mockey.Mock(systemd.CreateDefaultEnvFile).To(func(endpoint, dataDir string, dbInMemory bool, sessionProtocol ...string) error {
279280
return nil
280281
}).Build()
281282
mockey.Mock(systemd.GPUdServiceUnitFileContents).To(func() string {
@@ -299,7 +300,7 @@ func TestCommand_WithEmptyToken(t *testing.T) {
299300
func TestSystemdInit_WriteFileErrorWithDifferentPaths(t *testing.T) {
300301
mockey.PatchConvey("systemdInit write file error with path check", t, func() {
301302
writtenPath := ""
302-
mockey.Mock(systemd.CreateDefaultEnvFile).To(func(endpoint, dataDir string, dbInMemory bool) error {
303+
mockey.Mock(systemd.CreateDefaultEnvFile).To(func(endpoint, dataDir string, dbInMemory bool, sessionProtocol ...string) error {
303304
return nil
304305
}).Build()
305306
mockey.Mock(systemd.GPUdServiceUnitFileContents).To(func() string {
@@ -321,7 +322,7 @@ func TestSystemdInit_WriteFileErrorWithDifferentPaths(t *testing.T) {
321322
func TestSystemdInit_WithDbInMemory(t *testing.T) {
322323
mockey.PatchConvey("systemdInit with db in memory", t, func() {
323324
capturedDbInMemory := false
324-
mockey.Mock(systemd.CreateDefaultEnvFile).To(func(endpoint, dataDir string, dbInMemory bool) error {
325+
mockey.Mock(systemd.CreateDefaultEnvFile).To(func(endpoint, dataDir string, dbInMemory bool, sessionProtocol ...string) error {
325326
capturedDbInMemory = dbInMemory
326327
return nil
327328
}).Build()
@@ -471,6 +472,7 @@ func TestCommand_SystemdEnvFileWithEndpoint(t *testing.T) {
471472
tmpDir := t.TempDir()
472473
var capturedEndpoint string
473474
var capturedDataDir string
475+
var capturedSessionProtocol string
474476

475477
mockey.Mock(common.ResolveDataDir).To(func(cliContext *cli.Context) (string, error) {
476478
return tmpDir, nil
@@ -481,9 +483,10 @@ func TestCommand_SystemdEnvFileWithEndpoint(t *testing.T) {
481483
}).Build()
482484
mockey.Mock(pkdsystemd.SystemctlExists).To(func() bool { return true }).Build()
483485
mockey.Mock(systemd.DefaultBinExists).To(func() bool { return true }).Build()
484-
mockey.Mock(systemd.CreateDefaultEnvFile).To(func(endpoint, dataDir string, dbInMemory bool) error {
486+
mockey.Mock(systemd.CreateDefaultEnvFile).To(func(endpoint, dataDir string, dbInMemory bool, sessionProtocol ...string) error {
485487
capturedEndpoint = endpoint
486488
capturedDataDir = dataDir
489+
capturedSessionProtocol = sessionProtocol[0]
487490
return nil
488491
}).Build()
489492
mockey.Mock(systemd.GPUdServiceUnitFileContents).To(func() string {
@@ -495,11 +498,12 @@ func TestCommand_SystemdEnvFileWithEndpoint(t *testing.T) {
495498
mockey.Mock(pkgupdate.EnableGPUdSystemdUnit).To(func() error { return nil }).Build()
496499
mockey.Mock(pkgupdate.RestartGPUdSystemdUnit).To(func() error { return nil }).Build()
497500

498-
cliContext := newMockeyCLIContext(t, []string{"--endpoint", "https://custom.endpoint.com"})
501+
cliContext := newMockeyCLIContext(t, []string{"--endpoint", "https://custom.endpoint.com", "--session-protocol", "v2"})
499502
err := Command(cliContext)
500503
require.NoError(t, err)
501504
assert.Equal(t, "https://custom.endpoint.com", capturedEndpoint)
502505
assert.Equal(t, tmpDir, capturedDataDir)
506+
assert.Equal(t, "v2", capturedSessionProtocol)
503507
})
504508
}
505509

‎cmd/gpud/up/mock_command_test.go‎

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@ func newCLIContext(t *testing.T, args []string) *cli.Context {
3838
_ = flags.String("data-dir", "", "")
3939
_ = flags.String("token", "", "")
4040
_ = flags.String("endpoint", "", "")
41+
_ = flags.String("session-protocol", config.DefaultSessionProtocol, "")
4142
_ = flags.String("machine-id", "", "")
4243
_ = flags.String("node-group", "", "")
4344
_ = flags.String("gpu-count", "", "")
@@ -167,7 +168,7 @@ func TestCommand_SystemdInitError(t *testing.T) {
167168
}).Build()
168169
mockey.Mock(pkdsystemd.SystemctlExists).To(func() bool { return true }).Build()
169170
mockey.Mock(systemd.DefaultBinExists).To(func() bool { return true }).Build()
170-
mockey.Mock(systemd.CreateDefaultEnvFile).To(func(endpoint, dataDir string, dbInMemory bool) error {
171+
mockey.Mock(systemd.CreateDefaultEnvFile).To(func(endpoint, dataDir string, dbInMemory bool, sessionProtocol ...string) error {
171172
return errors.New("failed to create env file")
172173
}).Build()
173174

@@ -192,7 +193,7 @@ func TestCommand_EnableSystemdUnitError(t *testing.T) {
192193
}).Build()
193194
mockey.Mock(pkdsystemd.SystemctlExists).To(func() bool { return true }).Build()
194195
mockey.Mock(systemd.DefaultBinExists).To(func() bool { return true }).Build()
195-
mockey.Mock(systemd.CreateDefaultEnvFile).To(func(endpoint, dataDir string, dbInMemory bool) error {
196+
mockey.Mock(systemd.CreateDefaultEnvFile).To(func(endpoint, dataDir string, dbInMemory bool, sessionProtocol ...string) error {
196197
return nil
197198
}).Build()
198199
mockey.Mock(systemd.GPUdServiceUnitFileContents).To(func() string {
@@ -226,7 +227,7 @@ func TestCommand_RestartSystemdUnitError(t *testing.T) {
226227
}).Build()
227228
mockey.Mock(pkdsystemd.SystemctlExists).To(func() bool { return true }).Build()
228229
mockey.Mock(systemd.DefaultBinExists).To(func() bool { return true }).Build()
229-
mockey.Mock(systemd.CreateDefaultEnvFile).To(func(endpoint, dataDir string, dbInMemory bool) error {
230+
mockey.Mock(systemd.CreateDefaultEnvFile).To(func(endpoint, dataDir string, dbInMemory bool, sessionProtocol ...string) error {
230231
return nil
231232
}).Build()
232233
mockey.Mock(systemd.GPUdServiceUnitFileContents).To(func() string {
@@ -261,7 +262,7 @@ func TestCommand_SuccessWithoutToken(t *testing.T) {
261262
}).Build()
262263
mockey.Mock(pkdsystemd.SystemctlExists).To(func() bool { return true }).Build()
263264
mockey.Mock(systemd.DefaultBinExists).To(func() bool { return true }).Build()
264-
mockey.Mock(systemd.CreateDefaultEnvFile).To(func(endpoint, dataDir string, dbInMemory bool) error {
265+
mockey.Mock(systemd.CreateDefaultEnvFile).To(func(endpoint, dataDir string, dbInMemory bool, sessionProtocol ...string) error {
265266
return nil
266267
}).Build()
267268
mockey.Mock(systemd.GPUdServiceUnitFileContents).To(func() string {
@@ -317,7 +318,7 @@ func TestCommand_SuccessWithToken(t *testing.T) {
317318
}).Build()
318319
mockey.Mock(pkdsystemd.SystemctlExists).To(func() bool { return true }).Build()
319320
mockey.Mock(systemd.DefaultBinExists).To(func() bool { return true }).Build()
320-
mockey.Mock(systemd.CreateDefaultEnvFile).To(func(endpoint, dataDir string, dbInMemory bool) error {
321+
mockey.Mock(systemd.CreateDefaultEnvFile).To(func(endpoint, dataDir string, dbInMemory bool, sessionProtocol ...string) error {
321322
return nil
322323
}).Build()
323324
mockey.Mock(systemd.GPUdServiceUnitFileContents).To(func() string {
@@ -447,7 +448,7 @@ func TestRecordLoginSuccessState_Success(t *testing.T) {
447448
// TestSystemdInit_CreateEnvFileError tests systemdInit when creating env file fails.
448449
func TestSystemdInit_CreateEnvFileError(t *testing.T) {
449450
mockey.PatchConvey("systemdInit create env file error", t, func() {
450-
mockey.Mock(systemd.CreateDefaultEnvFile).To(func(endpoint, dataDir string, dbInMemory bool) error {
451+
mockey.Mock(systemd.CreateDefaultEnvFile).To(func(endpoint, dataDir string, dbInMemory bool, sessionProtocol ...string) error {
451452
return errors.New("failed to create env file")
452453
}).Build()
453454

@@ -460,7 +461,7 @@ func TestSystemdInit_CreateEnvFileError(t *testing.T) {
460461
// TestSystemdInit_WriteFileError tests systemdInit when writing unit file fails.
461462
func TestSystemdInit_WriteFileError(t *testing.T) {
462463
mockey.PatchConvey("systemdInit write file error", t, func() {
463-
mockey.Mock(systemd.CreateDefaultEnvFile).To(func(endpoint, dataDir string, dbInMemory bool) error {
464+
mockey.Mock(systemd.CreateDefaultEnvFile).To(func(endpoint, dataDir string, dbInMemory bool, sessionProtocol ...string) error {
464465
return nil
465466
}).Build()
466467
mockey.Mock(systemd.GPUdServiceUnitFileContents).To(func() string {
@@ -479,7 +480,7 @@ func TestSystemdInit_WriteFileError(t *testing.T) {
479480
// TestSystemdInit_Success tests successful systemdInit.
480481
func TestSystemdInit_Success(t *testing.T) {
481482
mockey.PatchConvey("systemdInit success", t, func() {
482-
mockey.Mock(systemd.CreateDefaultEnvFile).To(func(endpoint, dataDir string, dbInMemory bool) error {
483+
mockey.Mock(systemd.CreateDefaultEnvFile).To(func(endpoint, dataDir string, dbInMemory bool, sessionProtocol ...string) error {
483484
return nil
484485
}).Build()
485486
mockey.Mock(systemd.GPUdServiceUnitFileContents).To(func() string {

‎go.mod‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,9 @@ require (
3535
go.uber.org/zap v1.27.0
3636
golang.org/x/crypto v0.52.0
3737
golang.org/x/sys v0.45.0
38+
google.golang.org/genproto/googleapis/rpc v0.0.0-20251202230838-ff82c1b0f217
3839
google.golang.org/grpc v1.79.3
40+
google.golang.org/protobuf v1.36.10
3941
gopkg.in/natefinch/lumberjack.v2 v2.2.1
4042
k8s.io/api v0.32.0
4143
k8s.io/apimachinery v0.32.0
@@ -149,8 +151,6 @@ require (
149151
golang.org/x/time v0.11.0 // indirect
150152
golang.org/x/tools v0.44.0 // indirect
151153
golang.zx2c4.com/wireguard/windows v0.5.3 // indirect
152-
google.golang.org/genproto/googleapis/rpc v0.0.0-20251202230838-ff82c1b0f217 // indirect
153-
google.golang.org/protobuf v1.36.10 // indirect
154154
gopkg.in/inf.v0 v0.9.1 // indirect
155155
gopkg.in/yaml.v3 v3.0.1 // indirect
156156
k8s.io/klog/v2 v2.130.1 // indirect

‎pkg/config/config.go‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,9 @@ type Config struct {
103103
// SkipSessionUpdateConfig skips processing of updateConfig session commands. Intended for testing.
104104
SkipSessionUpdateConfig bool `json:"skip_session_update_config"`
105105

106+
// SessionProtocol selects the control-plane session transport: v1, v2, or auto.
107+
SessionProtocol string `json:"session_protocol"`
108+
106109
// DBInMemory enables in-memory SQLite database mode.
107110
// When true, the database is opened as a shared in-memory database (file::memory:?cache=shared)
108111
// instead of using the State file path. Data will not persist across restarts.
@@ -139,6 +142,11 @@ func (config *Config) Validate() error {
139142
if config.EventsRetentionPeriod.Duration > 0 && config.EventsRetentionPeriod.Duration < time.Minute {
140143
return fmt.Errorf("events_retention_period must be at least 1 minute, got %d", config.EventsRetentionPeriod.Duration)
141144
}
145+
switch config.SessionProtocol {
146+
case "", "v1", "v2", "auto":
147+
default:
148+
return fmt.Errorf("session_protocol must be one of v1, v2, or auto, got %q", config.SessionProtocol)
149+
}
142150

143151
return nil
144152
}

‎pkg/config/config_test.go‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,19 @@ import (
77
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
88
)
99

10+
func TestConfigValidateSessionProtocol(t *testing.T) {
11+
for _, protocol := range []string{"", "v1", "v2", "auto"} {
12+
cfg := &Config{Address: ":15132", MetricsRetentionPeriod: metav1.Duration{Duration: time.Minute}, SessionProtocol: protocol}
13+
if err := cfg.Validate(); err != nil {
14+
t.Fatalf("Validate() with protocol %q returned %v", protocol, err)
15+
}
16+
}
17+
cfg := &Config{Address: ":15132", MetricsRetentionPeriod: metav1.Duration{Duration: time.Minute}, SessionProtocol: "future"}
18+
if err := cfg.Validate(); err == nil {
19+
t.Fatal("Validate() accepted an unsupported session protocol")
20+
}
21+
}
22+
1023
func TestConfigValidate_AutoUpdateExitCode(t *testing.T) {
1124
tests := []struct {
1225
name string

‎pkg/config/default.go‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -13,9 +13,10 @@ import (
1313
)
1414

1515
const (
16-
DefaultAPIVersion = "v1"
17-
DefaultGPUdPort = 15132
18-
DefaultDataDir = "/var/lib/gpud"
16+
DefaultAPIVersion = "v1"
17+
DefaultGPUdPort = 15132
18+
DefaultDataDir = "/var/lib/gpud"
19+
DefaultSessionProtocol = "v1"
1920
)
2021

2122
var (
@@ -52,6 +53,7 @@ func DefaultConfig(ctx context.Context, opts ...OpOption) (*Config, error) {
5253
CompactPeriod: DefaultCompactPeriod,
5354
Pprof: false,
5455
EnableAutoUpdate: true,
56+
SessionProtocol: DefaultSessionProtocol,
5557
NvidiaToolOverwrites: nvidiacommon.ToolOverwrites{
5658
InfinibandClassRootDir: options.InfinibandClassRootDir,
5759
ExcludedInfinibandDevices: options.ExcludedInfinibandDevices,

0 commit comments

Comments
 (0)