Skip to content

Commit a72085b

Browse files
gyuhogyuho
authored
feat(login, helm): re-register on machine-id change and refresh session token (#1255)
## Problem gpud persists its login identity — **machine id**, **session token**, endpoint — in `/var/lib/gpud`, a node-scoped `hostPath` that survives reboots. In the container/DaemonSet pattern that cache can outlive what it describes, and gpud's "skip login if a machine id is already persisted" fast path then runs the node on **stale** state: 1. **Stale machine id** — a node deleted from its node group and rejoined binds a **NEW machine object** (new `machine-id` label) while the OLD id is still on disk → gpud either silently skips login on the stale id (runs under a mismatched identity, never joins) or fails with a terse mismatch error. 2. **Stale session token** — the persisted session token can be invalidated server-side (gpud-manager resets the BYOK session on every login; the workspace session SAK can rotate) → skip-login reuses a dead token → keepalive breaks. This PR adds two opt-in flags (both **default `false`** in gpud for backward-compatibility, both **default `true`** in the chart, matching how the chart already enables container-pattern behaviors like `rebootCommands`/`findmntCommands`). They wire into the **same** skip/refresh fork in `login.Login`. ## 1. `--machine-id-overwrite` When a supplied `--machine-id` differs from the persisted one: - **false (default)** → reject with an actionable error (safe for host/systemd). - **true** → discard the persisted login identity (machine id + session token, **metadata table only**) and re-register as the new machine. Reboot/event history is **preserved**. Unified in `login.reconcileMachineID`, which runs *before* the skip/refresh fork so behavior is consistent whether or not node labels changed. The now-redundant deeper guard (`stored machine ID … does not match`) is removed (unreachable). The `--machine-id` flag docs now describe this. ## 2. `--refresh-session-token` When set, gpud **always re-logins on start** even with a persisted machine id, instead of reusing the cached session token. The control plane re-issues the (workspace-scoped) session token on every successful login, so re-login reliably re-fetches a current one. Costs one extra login round-trip per start; keeps the cached token from going stale. > Verified against gpud-manager: `server/handler_login.go` returns `sessionToken` on every success (`loginExistingMachine`/`loginNewMachine`), sourced from the workspace Secret via `tokenReader.GetToken(WorkspaceID, TokenPermissionSession)` — and `clearStaleSessionIfBYOK` resets the BYOK session on every login, so the server *expects* re-login. A `TODO` is noted on both flags: once validated in production, make always-refresh the default and retire the toggle. ## Helm `gpud.machineIdOverwrite` (default `true`) and `gpud.refreshSessionToken` (default `true`) → `GPUD_MACHINE_ID_OVERWRITE` / `GPUD_REFRESH_SESSION_TOKEN` env → the startup wrapper appends `--machine-id-overwrite` / `--refresh-session-token`. Set either to `false` to restore the strict/skip-login behavior. ## Backward compatibility **Strict invariant preserved**: with both flags unset (`false`) and no `--machine-id` supplied — i.e. every existing host/systemd install — behavior is byte-for-byte identical. `reconcileMachineID` is a no-op when there's no persisted id, no supplied id, or they match; and without `--refresh-session-token` the skip-login fast path is unchanged. ## Testing - `pkg/login/login_test.go` — `TestReconcileMachineID` (no-persisted / empty-requested / match / mismatch-fail / mismatch-overwrite). - `pkg/login/mock_login_test.go` — full `Login` flow: `…MismatchedOverrideFails` (updated to the new error), `TestLogin_MachineIDOverwriteReRegisters`, and `TestLogin_RefreshSessionTokenRelogsIn` (forces relogin + re-persists the token when a machine id is already assigned and labels are unchanged; contrast with `TestLogin_MachineIDAlreadyAssigned`, which skips). - `go build ./...`, `go vet`, `go test -gcflags="all=-N -l" ./pkg/login/` all green. - `helm lint` + render verified for both flags: default → env + flag present; `--set …=false` → both omitted. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Signed-off-by: Gyuho Lee <gyuhol@nvidia.com>
1 parent 4e754b1 commit a72085b

14 files changed

Lines changed: 844 additions & 143 deletions

‎cmd/gpud/command/command.go‎

Lines changed: 26 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -98,7 +98,19 @@ nohup sudo gpud run &>> <your log file path> &
9898
cli.StringFlag{
9999
Name: "machine-id",
100100
Hidden: true,
101-
Usage: "(optional) for override default machine id",
101+
Usage: "(optional) machine id to register/run as. If it matches the machine id already persisted in the state DB, login is skipped (reused across reboots). If it DIFFERS from the persisted one, login fails unless --machine-id-overwrite is set -- this prevents silently retargeting an already-registered node.",
102+
},
103+
cli.BoolFlag{
104+
Name: "machine-id-overwrite",
105+
Hidden: true,
106+
Usage: "(optional) when --machine-id differs from the persisted machine id, discard the persisted login identity (machine id + session token) and check in with the requested machine instead of failing. Health/system state (reboot history, events) is preserved. Use for the container/DaemonSet pattern where a node can be deleted from its node group and rejoined with a new machine object.",
107+
},
108+
// TODO: once validated in production, make session-token refresh the default
109+
// behavior (always re-login on start) and remove this flag.
110+
cli.BoolFlag{
111+
Name: "refresh-session-token",
112+
Hidden: true,
113+
Usage: "(optional) on every start, re-run login to re-fetch the current session token from the control plane instead of reusing the persisted one (skips the 'machine id already assigned' fast path). Use for the container/DaemonSet pattern after a workspace token rotation.",
102114
},
103115
cli.StringFlag{
104116
Name: "node-group",
@@ -166,7 +178,19 @@ sudo rm /etc/systemd/system/gpud.service
166178
&cli.StringFlag{
167179
Name: "machine-id",
168180
Hidden: true,
169-
Usage: "(optional) for override default machine id",
181+
Usage: "(optional) machine id to register/run as. If it matches the machine id already persisted in the state DB, login is skipped (reused across reboots). If it DIFFERS from the persisted one, login fails unless --machine-id-overwrite is set -- this prevents silently retargeting an already-registered node.",
182+
},
183+
&cli.BoolFlag{
184+
Name: "machine-id-overwrite",
185+
Hidden: true,
186+
Usage: "(optional) when --machine-id differs from the persisted machine id, discard the persisted login identity (machine id + session token) and check in with the requested machine instead of failing. Health/system state (reboot history, events) is preserved. Use for the container/DaemonSet pattern where a node can be deleted from its node group and rejoined with a new machine object.",
187+
},
188+
// TODO: once validated in production, make session-token refresh the default
189+
// behavior (always re-login on start) and remove this flag.
190+
&cli.BoolFlag{
191+
Name: "refresh-session-token",
192+
Hidden: true,
193+
Usage: "(optional) on every start, re-run login to re-fetch the current session token from the control plane instead of reusing the persisted one (skips the 'machine id already assigned' fast path). Use for the container/DaemonSet pattern after a workspace token rotation.",
170194
},
171195
&cli.StringFlag{
172196
Name: "token",

‎cmd/gpud/run/command.go‎

Lines changed: 72 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,9 @@ func Command(cliContext *cli.Context) error {
8080
controlPlaneLoginRegistrationToken := cliContext.String("token")
8181

8282
machineIDForOverride := cliContext.String("machine-id")
83+
shouldOverwriteMachineID := cliContext.Bool("machine-id-overwrite")
84+
refreshSessionToken := cliContext.Bool("refresh-session-token")
85+
controlPlaneLoginSucceeded := false
8386

8487
// Note: login.Login() ALWAYS writes to the persistent state file (via dataDir),
8588
// regardless of --db-in-memory flag. The login package doesn't know about in-memory mode.
@@ -92,10 +95,12 @@ func Command(cliContext *cli.Context) error {
9295
defer loginCancel()
9396

9497
loginCfg := login.LoginConfig{
95-
Token: controlPlaneLoginRegistrationToken,
96-
Endpoint: controlPlaneEndpoint,
97-
MachineID: machineIDForOverride,
98-
DataDir: dataDir,
98+
Token: controlPlaneLoginRegistrationToken,
99+
Endpoint: controlPlaneEndpoint,
100+
MachineID: machineIDForOverride,
101+
MachineIDOverwrite: shouldOverwriteMachineID,
102+
RefreshSessionToken: refreshSessionToken,
103+
DataDir: dataDir,
99104

100105
GPUCount: gpuCountStr,
101106
}
@@ -104,6 +109,7 @@ func Command(cliContext *cli.Context) error {
104109
if lerr := login.Login(loginCtx, loginCfg); lerr != nil {
105110
return lerr
106111
}
112+
controlPlaneLoginSucceeded = true
107113
log.Logger.Infow("successfully logged in in gpud run")
108114

109115
if err := recordLoginSuccessState(loginCtx, dataDir); err != nil {
@@ -375,34 +381,8 @@ func Command(cliContext *cli.Context) error {
375381
mctx, mcancel := context.WithTimeout(context.Background(), 10*time.Second)
376382
defer mcancel()
377383

378-
dbRW, err := pkgsqlite.Open(cfg.State)
379-
if err != nil {
380-
return fmt.Errorf("failed to open state for metadata overrides: %w", err)
381-
}
382-
defer func() {
383-
_ = dbRW.Close()
384-
}()
385-
386-
if err := pkgmetadata.CreateTableMetadata(mctx, dbRW); err != nil {
387-
return fmt.Errorf("failed to ensure metadata table: %w", err)
388-
}
389-
390-
if controlPlaneEndpoint != "" {
391-
if err := pkgmetadata.SetMetadata(mctx, dbRW, pkgmetadata.MetadataKeyEndpoint, controlPlaneEndpoint); err != nil {
392-
return fmt.Errorf("failed to set endpoint metadata: %w", err)
393-
}
394-
log.Logger.Infow("overriding endpoint from flag", "endpoint", controlPlaneEndpoint)
395-
}
396-
397-
// DO NOT overwrite "pkgmetadata.MetadataKeyToken"
398-
// because successful login operation will persist the session token in the metadata
399-
// NOT the registration token
400-
401-
if machineIDForOverride != "" {
402-
if err := pkgmetadata.SetMetadata(mctx, dbRW, pkgmetadata.MetadataKeyMachineID, machineIDForOverride); err != nil {
403-
return fmt.Errorf("failed to set machine-id metadata: %w", err)
404-
}
405-
log.Logger.Infow("overriding machine id from flag", "machineID", machineIDForOverride)
384+
if err := persistMetadataOverrides(mctx, cfg.State, controlPlaneEndpoint, machineIDForOverride, shouldOverwriteMachineID, controlPlaneLoginSucceeded); err != nil {
385+
return err
406386
}
407387
}
408388

@@ -456,6 +436,66 @@ func Command(cliContext *cli.Context) error {
456436
return nil
457437
}
458438

439+
func validateMachineIDOverride(prevMachineID, requestedMachineID string, overwrite, controlPlaneLoginSucceeded bool) error {
440+
if prevMachineID == "" || prevMachineID == requestedMachineID {
441+
return nil
442+
}
443+
if !overwrite {
444+
return fmt.Errorf("persisted machine ID %q differs from --machine-id %q; pass --machine-id-overwrite to replace it", prevMachineID, requestedMachineID)
445+
}
446+
if !controlPlaneLoginSucceeded {
447+
return fmt.Errorf("cannot overwrite persisted machine ID %q with --machine-id %q without a successful control-plane login; pass --token so gpud can check in to the requested machine", prevMachineID, requestedMachineID)
448+
}
449+
return nil
450+
}
451+
452+
func persistMetadataOverrides(ctx context.Context, stateFile, controlPlaneEndpoint, machineIDForOverride string, machineIDOverwrite, controlPlaneLoginSucceeded bool) error {
453+
dbRW, err := pkgsqlite.Open(stateFile)
454+
if err != nil {
455+
return fmt.Errorf("failed to open state for metadata overrides: %w", err)
456+
}
457+
defer func() {
458+
_ = dbRW.Close()
459+
}()
460+
461+
if err := pkgmetadata.CreateTableMetadata(ctx, dbRW); err != nil {
462+
return fmt.Errorf("failed to ensure metadata table: %w", err)
463+
}
464+
465+
if controlPlaneEndpoint != "" {
466+
if err := pkgmetadata.SetMetadata(ctx, dbRW, pkgmetadata.MetadataKeyEndpoint, controlPlaneEndpoint); err != nil {
467+
return fmt.Errorf("failed to set endpoint metadata: %w", err)
468+
}
469+
log.Logger.Infow("overriding endpoint from flag", "endpoint", controlPlaneEndpoint)
470+
}
471+
472+
if machineIDForOverride == "" {
473+
return nil
474+
}
475+
476+
prevMachineID, err := pkgmetadata.ReadMetadata(ctx, dbRW, pkgmetadata.MetadataKeyMachineID)
477+
if err != nil {
478+
return fmt.Errorf("failed to read persisted machine-id: %w", err)
479+
}
480+
if err := validateMachineIDOverride(prevMachineID, machineIDForOverride, machineIDOverwrite, controlPlaneLoginSucceeded); err != nil {
481+
return err
482+
}
483+
484+
// A successful control-plane login persists the machine ID returned by the
485+
// control plane. Do not overwrite that authoritative result with the CLI
486+
// value. The no-login path keeps the historical behavior of recording an
487+
// initial machine-id override locally.
488+
if prevMachineID == machineIDForOverride || controlPlaneLoginSucceeded {
489+
return nil
490+
}
491+
492+
if err := pkgmetadata.SetMetadata(ctx, dbRW, pkgmetadata.MetadataKeyMachineID, machineIDForOverride); err != nil {
493+
return fmt.Errorf("failed to set machine-id metadata: %w", err)
494+
}
495+
log.Logger.Infow("overriding machine id from flag", "machineID", machineIDForOverride)
496+
return nil
497+
}
498+
459499
func parseRetentionPeriods(cliContext *cli.Context) (metricsRetentionPeriod, eventsRetentionPeriod time.Duration) {
460500
metricsRetentionPeriod = cliContext.Duration("metrics-retention-period")
461501
deprecatedRetentionPeriod := cliContext.Duration("retention-period")

‎cmd/gpud/run/command_test.go‎

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package run
22

33
import (
4+
"context"
45
"flag"
56
"testing"
67
"time"
@@ -10,6 +11,8 @@ import (
1011
"github.com/urfave/cli"
1112

1213
"github.com/leptonai/gpud/pkg/config"
14+
pkgmetadata "github.com/leptonai/gpud/pkg/metadata"
15+
pkgsqlite "github.com/leptonai/gpud/pkg/sqlite"
1316
)
1417

1518
func TestParseInfinibandExcludeDevices(t *testing.T) {
@@ -128,3 +131,110 @@ func TestParseRetentionPeriods(t *testing.T) {
128131
})
129132
}
130133
}
134+
135+
func TestPersistMetadataOverrides(t *testing.T) {
136+
t.Parallel()
137+
138+
tests := []struct {
139+
name string
140+
persistedMachineID string
141+
persistedToken string
142+
controlPlaneEndpoint string
143+
requestedMachineID string
144+
machineIDOverwrite bool
145+
controlPlaneLoginSucceeded bool
146+
wantMachineID string
147+
wantToken string
148+
wantEndpoint string
149+
errContains string
150+
}{
151+
{
152+
name: "endpoint override without machine id is persisted",
153+
controlPlaneEndpoint: "gpud-manager.example.com",
154+
wantEndpoint: "gpud-manager.example.com",
155+
},
156+
{
157+
name: "initial machine id without login is allowed",
158+
requestedMachineID: "new-machine-id",
159+
wantMachineID: "new-machine-id",
160+
},
161+
{
162+
name: "matching machine id without login is allowed",
163+
persistedMachineID: "machine-id",
164+
persistedToken: "session-token",
165+
requestedMachineID: "machine-id",
166+
wantMachineID: "machine-id",
167+
wantToken: "session-token",
168+
},
169+
{
170+
name: "mismatch without overwrite is rejected",
171+
persistedMachineID: "old-machine-id",
172+
persistedToken: "old-session-token",
173+
requestedMachineID: "new-machine-id",
174+
wantMachineID: "old-machine-id",
175+
wantToken: "old-session-token",
176+
errContains: "pass --machine-id-overwrite",
177+
},
178+
{
179+
name: "mismatch without a successful login is rejected",
180+
persistedMachineID: "old-machine-id",
181+
persistedToken: "old-session-token",
182+
requestedMachineID: "new-machine-id",
183+
machineIDOverwrite: true,
184+
wantMachineID: "old-machine-id",
185+
wantToken: "old-session-token",
186+
errContains: "pass --token",
187+
},
188+
{
189+
name: "successful login keeps control plane machine id",
190+
persistedMachineID: "manager-machine-id",
191+
persistedToken: "fresh-session-token",
192+
requestedMachineID: "requested-machine-id",
193+
machineIDOverwrite: true,
194+
controlPlaneLoginSucceeded: true,
195+
wantMachineID: "manager-machine-id",
196+
wantToken: "fresh-session-token",
197+
},
198+
}
199+
200+
for _, tt := range tests {
201+
t.Run(tt.name, func(t *testing.T) {
202+
t.Parallel()
203+
204+
ctx := context.Background()
205+
stateFile := config.StateFilePath(t.TempDir())
206+
dbRW, err := pkgsqlite.Open(stateFile)
207+
require.NoError(t, err)
208+
require.NoError(t, pkgmetadata.CreateTableMetadata(ctx, dbRW))
209+
if tt.persistedMachineID != "" {
210+
require.NoError(t, pkgmetadata.SetMetadata(ctx, dbRW, pkgmetadata.MetadataKeyMachineID, tt.persistedMachineID))
211+
}
212+
if tt.persistedToken != "" {
213+
require.NoError(t, pkgmetadata.SetMetadata(ctx, dbRW, pkgmetadata.MetadataKeyToken, tt.persistedToken))
214+
}
215+
require.NoError(t, dbRW.Close())
216+
217+
err = persistMetadataOverrides(ctx, stateFile, tt.controlPlaneEndpoint, tt.requestedMachineID, tt.machineIDOverwrite, tt.controlPlaneLoginSucceeded)
218+
if tt.errContains == "" {
219+
require.NoError(t, err)
220+
} else {
221+
require.Error(t, err)
222+
assert.Contains(t, err.Error(), tt.errContains)
223+
}
224+
225+
dbRO, err := pkgsqlite.Open(stateFile, pkgsqlite.WithReadOnly(true))
226+
require.NoError(t, err)
227+
defer func() { _ = dbRO.Close() }()
228+
229+
machineID, err := pkgmetadata.ReadMachineID(ctx, dbRO)
230+
require.NoError(t, err)
231+
token, err := pkgmetadata.ReadToken(ctx, dbRO)
232+
require.NoError(t, err)
233+
endpoint, err := pkgmetadata.ReadMetadata(ctx, dbRO, pkgmetadata.MetadataKeyEndpoint)
234+
require.NoError(t, err)
235+
assert.Equal(t, tt.wantMachineID, machineID)
236+
assert.Equal(t, tt.wantToken, token)
237+
assert.Equal(t, tt.wantEndpoint, endpoint)
238+
})
239+
}
240+
}

0 commit comments

Comments
 (0)