Skip to content

Commit c0fb784

Browse files
Code2LifeCode2Life
committed
feat: add UpdateWorkerEnv method to HypervisorManager for in-place environment variable updates
1 parent 3f52516 commit c0fb784

5 files changed

Lines changed: 58 additions & 23 deletions

File tree

‎go.mod‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -7,15 +7,15 @@ require (
77
github.com/blang/semver/v4 v4.0.0
88
github.com/charmbracelet/lipgloss v1.1.0
99
github.com/go-resty/resty/v2 v2.17.1
10-
github.com/onsi/ginkgo/v2 v2.28.0
10+
github.com/onsi/ginkgo/v2 v2.28.1
1111
github.com/onsi/gomega v1.39.1
1212
github.com/spf13/cobra v1.10.2
1313
github.com/stretchr/testify v1.11.1
1414
k8s.io/klog/v2 v2.130.1
1515
)
1616

1717
// replace with the target commit id and run go mod tidy to get pseudo version
18-
replace github.com/NexusGPU/tensor-fusion => github.com/NexusGPU/tensor-fusion v1.54.7-0.20260207152551-11640a6a00c4
18+
replace github.com/NexusGPU/tensor-fusion => github.com/NexusGPU/tensor-fusion v1.54.7-0.20260212073720-55025d4cfb2c
1919

2020
require (
2121
cel.dev/expr v0.24.0 // indirect
@@ -79,7 +79,7 @@ require (
7979
github.com/modern-go/reflect2 v1.0.3-0.20250322232337-35a7c28c31ee // indirect
8080
github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect
8181
github.com/pelletier/go-toml/v2 v2.2.4 // indirect
82-
github.com/posthog/posthog-go v1.9.1 // indirect
82+
github.com/posthog/posthog-go v1.10.0 // indirect
8383
github.com/prometheus/client_golang v1.23.2 // indirect
8484
github.com/prometheus/client_model v0.6.2 // indirect
8585
github.com/prometheus/common v0.66.1 // indirect
@@ -93,8 +93,8 @@ require (
9393
github.com/ugorji/go/codec v1.3.0 // indirect
9494
github.com/x448/float16 v0.8.4 // indirect
9595
github.com/yusufpapurcu/wmi v1.2.4 // indirect
96-
go.opentelemetry.io/otel v1.38.0 // indirect
97-
go.opentelemetry.io/otel/trace v1.38.0 // indirect
96+
go.opentelemetry.io/otel v1.40.0 // indirect
97+
go.opentelemetry.io/otel/trace v1.40.0 // indirect
9898
go.yaml.in/yaml/v2 v2.4.3 // indirect
9999
go.yaml.in/yaml/v3 v3.0.4 // indirect
100100
golang.org/x/arch v0.22.0 // indirect

‎go.sum‎

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,8 @@ filippo.io/edwards25519 v1.1.0 h1:FNf4tywRC1HmFuKW5xopWpigGjJKiJSV0Cqo0cJWDaA=
44
filippo.io/edwards25519 v1.1.0/go.mod h1:BxyFTGdWcka3PhytdK4V28tE5sGfRvvvRV7EaN4VDT4=
55
github.com/Masterminds/semver/v3 v3.4.0 h1:Zog+i5UMtVoCU8oKka5P7i9q9HgrJeGzI9SA1Xbatp0=
66
github.com/Masterminds/semver/v3 v3.4.0/go.mod h1:4V+yj/TJE1HU9XfppCwVMZq3I84lprf4nC11bSS5beM=
7-
github.com/NexusGPU/tensor-fusion v1.54.7-0.20260207152551-11640a6a00c4 h1:o7o/HwQVAM/zWcoTRHUI+7WWGSV+KNhreZfisMAhA9w=
8-
github.com/NexusGPU/tensor-fusion v1.54.7-0.20260207152551-11640a6a00c4/go.mod h1:v5InWOMmDN4yQe2K7YAQXz7wR7LFZn5Ebl4nD68On9I=
7+
github.com/NexusGPU/tensor-fusion v1.54.7-0.20260212073720-55025d4cfb2c h1:fw92lrfj1g0M1t00aPdsOSSMYfdsgZZlhpdcupmhWpk=
8+
github.com/NexusGPU/tensor-fusion v1.54.7-0.20260212073720-55025d4cfb2c/go.mod h1:iId/YtFuib6CxqR+PNEIIFMM+ZvP/8Z0iEN4bN5LptA=
99
github.com/antlr4-go/antlr/v4 v4.13.1 h1:SqQKkuVZ+zWkMMNkjy5FZe5mr5WURWnlpmOuzYWrPrQ=
1010
github.com/antlr4-go/antlr/v4 v4.13.1/go.mod h1:GKmUxMtwp6ZgGwZSva4eWPC5mS6vUAmOABFgjdkM7Nw=
1111
github.com/aymanbagabas/go-osc52/v2 v2.0.1 h1:HwpRHbFMcZLEVr42D4p7XBqjyuxQH5SMiErDT4WkJ2k=
@@ -209,8 +209,8 @@ github.com/muesli/termenv v0.16.0/go.mod h1:ZRfOIKPFDYQoDFF4Olj7/QJbW60Ol/kL1pU3
209209
github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 h1:C3w9PqII01/Oq1c1nUAm88MOHcQC9l5mIlSMApZMrHA=
210210
github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822/go.mod h1:+n7T8mK8HuQTcFwEeznm/DIxMOiR9yIdICNftLE1DvQ=
211211
github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e/go.mod h1:zD1mROLANZcx1PVRCS0qkT7pwLkGfwJo4zjcN/Tysno=
212-
github.com/onsi/ginkgo/v2 v2.28.0 h1:Rrf+lVLmtlBIKv6KrIGJCjyY8N36vDVcutbGJkyqjJc=
213-
github.com/onsi/ginkgo/v2 v2.28.0/go.mod h1:ArE1D/XhNXBXCBkKOLkbsb2c81dQHCRcF5zwn/ykDRo=
212+
github.com/onsi/ginkgo/v2 v2.28.1 h1:S4hj+HbZp40fNKuLUQOYLDgZLwNUVn19N3Atb98NCyI=
213+
github.com/onsi/ginkgo/v2 v2.28.1/go.mod h1:CLtbVInNckU3/+gC8LzkGUb9oF+e8W8TdUsxPwvdOgE=
214214
github.com/onsi/gomega v1.39.1 h1:1IJLAad4zjPn2PsnhH70V4DKRFlrCzGBNrNaru+Vf28=
215215
github.com/onsi/gomega v1.39.1/go.mod h1:hL6yVALoTOxeWudERyfppUcZXjMwIMLnuSfruD2lcfg=
216216
github.com/pelletier/go-toml/v2 v2.2.4 h1:mye9XuhQ6gvn5h28+VilKrrPoQVanw5PMw/TB0t5Ec4=
@@ -220,8 +220,8 @@ github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINE
220220
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
221221
github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRIccs7FGNTlIRMkT8wgtp5eCXdBlqhYGL6U=
222222
github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
223-
github.com/posthog/posthog-go v1.9.1 h1:9bkcRnYSvcgMxL2s9QlCnd1DVnm2qWXxWu5o0HSF0xM=
224-
github.com/posthog/posthog-go v1.9.1/go.mod h1:wB3/9Q7d9gGb1P/yf/Wri9VBlbP8oA8z++prRzL5OcY=
223+
github.com/posthog/posthog-go v1.10.0 h1:wfoy7Jfb4LigCoHYyMZoiJmmEoCLOkSaYfDxM/NtCqY=
224+
github.com/posthog/posthog-go v1.10.0/go.mod h1:wB3/9Q7d9gGb1P/yf/Wri9VBlbP8oA8z++prRzL5OcY=
225225
github.com/prometheus/client_golang v1.23.2 h1:Je96obch5RDVy3FDMndoUsjAhG5Edi49h0RJWRi/o0o=
226226
github.com/prometheus/client_golang v1.23.2/go.mod h1:Tb1a6LWHB3/SPIzCoaDXI4I8UHKeFTEQ1YCr+0Gyqmg=
227227
github.com/prometheus/client_model v0.6.2 h1:oBsgwpGs7iVziMvrGhE53c/GrLUsZdHnqNwqPLxwZyk=
@@ -279,10 +279,10 @@ github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e h1:JVG44RsyaB9T2KIHavM
279279
github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e/go.mod h1:RbqR21r5mrJuqunuUZ/Dhy/avygyECGrLceyNeo4LiM=
280280
github.com/yusufpapurcu/wmi v1.2.4 h1:zFUKzehAFReQwLys1b/iSMl+JQGSCSjtVqQn9bBrPo0=
281281
github.com/yusufpapurcu/wmi v1.2.4/go.mod h1:SBZ9tNy3G9/m5Oi98Zks0QjeHVDvuK0qfxQmPyzfmi0=
282-
go.opentelemetry.io/otel v1.38.0 h1:RkfdswUDRimDg0m2Az18RKOsnI8UDzppJAtj01/Ymk8=
283-
go.opentelemetry.io/otel v1.38.0/go.mod h1:zcmtmQ1+YmQM9wrNsTGV/q/uyusom3P8RxwExxkZhjM=
284-
go.opentelemetry.io/otel/trace v1.38.0 h1:Fxk5bKrDZJUH+AMyyIXGcFAPah0oRcT+LuNtJrmcNLE=
285-
go.opentelemetry.io/otel/trace v1.38.0/go.mod h1:j1P9ivuFsTceSWe1oY+EeW3sc+Pp42sO++GHkg4wwhs=
282+
go.opentelemetry.io/otel v1.40.0 h1:oA5YeOcpRTXq6NN7frwmwFR0Cn3RhTVZvXsP4duvCms=
283+
go.opentelemetry.io/otel v1.40.0/go.mod h1:IMb+uXZUKkMXdPddhwAHm6UfOwJyh4ct1ybIlV14J0g=
284+
go.opentelemetry.io/otel/trace v1.40.0 h1:WA4etStDttCSYuhwvEa8OP8I5EWu24lkOzp+ZYblVjw=
285+
go.opentelemetry.io/otel/trace v1.40.0/go.mod h1:zeAhriXecNGP/s2SEG3+Y8X9ujcJOTqQ5RgdEJcawiA=
286286
go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto=
287287
go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE=
288288
go.uber.org/mock v0.5.2 h1:LbtPTcP8A5k9WPXj54PPPbjcI4Y6lhyOZXn+VS7wNko=

‎internal/agent/agent_test.go‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -369,6 +369,10 @@ func (m *mockHypervisorManager) StopWorker(workerUID string) error {
369369
return nil
370370
}
371371

372+
func (m *mockHypervisorManager) UpdateWorkerEnv(workerUID string, env map[string]string) error {
373+
return nil
374+
}
375+
372376
func (m *mockHypervisorManager) GetDeviceMetrics() (map[string]*hvApi.GPUUsageMetrics, error) {
373377
return nil, nil
374378
}

‎internal/hypervisor/manager.go‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ type HypervisorManager interface {
2929
ListWorkers() []*api.WorkerInfo
3030
StartWorker(workerInfo *api.WorkerInfo) error
3131
StopWorker(workerUID string) error
32+
UpdateWorkerEnv(workerUID string, env map[string]string) error
3233
GetDeviceMetrics() (map[string]*api.GPUUsageMetrics, error)
3334
GetWorkerAllocation(workerUID string) (*api.WorkerAllocation, bool)
3435
RegisterWorkerHandler(handler framework.WorkerChangeHandler) error
@@ -400,6 +401,19 @@ func (m *Manager) StopWorker(workerUID string) error {
400401
return nil
401402
}
402403

404+
// UpdateWorkerEnv updates environment variables for a worker without restarting its process.
405+
// The new env vars take effect on next process restart (crash recovery).
406+
func (m *Manager) UpdateWorkerEnv(workerUID string, env map[string]string) error {
407+
m.mu.RLock()
408+
defer m.mu.RUnlock()
409+
410+
if !m.started {
411+
return ErrNotStarted
412+
}
413+
414+
return m.backend.UpdateWorkerEnv(workerUID, env)
415+
}
416+
403417
// GetWorkerAllocation returns the allocation for a specific worker
404418
func (m *Manager) GetWorkerAllocation(workerUID string) (*api.WorkerAllocation, bool) {
405419
m.mu.RLock()

‎internal/hypervisor/reconciler.go‎

Lines changed: 25 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -137,13 +137,21 @@ func (r *Reconciler) reconcile() {
137137
} else {
138138
added++
139139
}
140-
} else if r.needsUpdate(desiredInfo, actualWorker) {
141-
// Worker exists but config changed (non-status fields), restart it
140+
} else if r.needsRestart(desiredInfo, actualWorker) {
141+
// Structural change (GPU allocation, executable, args) requires restart
142142
if err := r.restartWorker(desiredInfo); err != nil {
143143
klog.Errorf("Failed to restart worker: worker_id=%s error=%v", workerID, err)
144144
} else {
145145
updated++
146146
}
147+
} else if r.needsEnvUpdate(desiredInfo, actualWorker) {
148+
// Env-only change: update in place without restarting the running process.
149+
// New env vars take effect on next process restart (crash recovery).
150+
if err := r.manager.UpdateWorkerEnv(workerID, desiredInfo.WorkerRunningInfo.Env); err != nil {
151+
klog.Errorf("Failed to update worker env: worker_id=%s error=%v", workerID, err)
152+
} else {
153+
updated++
154+
}
147155
}
148156
}
149157

@@ -204,8 +212,9 @@ func (r *Reconciler) restartWorker(info *api.WorkerInfo) error {
204212
return r.startWorker(info)
205213
}
206214

207-
// needsUpdate checks if worker config changed (non-status fields only)
208-
func (r *Reconciler) needsUpdate(desired, actual *api.WorkerInfo) bool {
215+
// needsRestart checks if structural config changed (GPU allocation, executable, args)
216+
// that requires stopping and restarting the worker process.
217+
func (r *Reconciler) needsRestart(desired, actual *api.WorkerInfo) bool {
209218
// Check if GPU allocation changed
210219
if len(desired.AllocatedDevices) != len(actual.AllocatedDevices) {
211220
return true
@@ -238,10 +247,6 @@ func (r *Reconciler) needsUpdate(desired, actual *api.WorkerInfo) bool {
238247
return true
239248
}
240249
}
241-
// Check environment variables (including fractional GPU config)
242-
if !maps.Equal(desired.WorkerRunningInfo.Env, actual.WorkerRunningInfo.Env) {
243-
return true
244-
}
245250
} else if desired.WorkerRunningInfo != nil || actual.WorkerRunningInfo != nil {
246251
// One is nil, the other is not - they differ
247252
return true
@@ -250,6 +255,18 @@ func (r *Reconciler) needsUpdate(desired, actual *api.WorkerInfo) bool {
250255
return false
251256
}
252257

258+
// needsEnvUpdate checks if only environment variables changed.
259+
// Env changes are applied in place without restarting the running process;
260+
// the new values take effect on next process restart (crash recovery).
261+
func (r *Reconciler) needsEnvUpdate(desired, actual *api.WorkerInfo) bool {
262+
if desired.WorkerRunningInfo != nil && actual.WorkerRunningInfo != nil {
263+
if !maps.Equal(desired.WorkerRunningInfo.Env, actual.WorkerRunningInfo.Env) {
264+
return true
265+
}
266+
}
267+
return false
268+
}
269+
253270
// GetStatus returns current reconciler status
254271
func (r *Reconciler) GetStatus() ReconcilerStatus {
255272
r.mu.RLock()

0 commit comments

Comments
 (0)