From b3b2cf166f073edc8f29e9f5245fed3769eb80af Mon Sep 17 00:00:00 2001 From: Mike Camp Date: Tue, 25 Aug 2026 14:17:25 -0400 Subject: [PATCH 1/7] fix(nvca): make no-GPU recovery registration-safe Keep readiness unhealthy and creation queues paused until GPU discovery is followed by a successful ICMS registration and credential installation. Retry transient failures while preserving liveness. Serialize recovery and periodic registration across inventory capture and response application so stale credentials cannot overwrite recovered state. Tests: focused graceful-no-GPU lifecycle and response-ordering regressions; full pkg/nvca suite; focused race run. Signed-off-by: Mike Camp --- .../nvca/pkg/nvca/agent.go | 150 ++++++- .../nvca/pkg/nvca/agent_test.go | 418 +++++++++++++++++- .../nvca/pkg/nvca/agent_updates.go | 12 + 3 files changed, 556 insertions(+), 24 deletions(-) diff --git a/src/compute-plane-services/nvca/pkg/nvca/agent.go b/src/compute-plane-services/nvca/pkg/nvca/agent.go index 67df76601..2891737fa 100644 --- a/src/compute-plane-services/nvca/pkg/nvca/agent.go +++ b/src/compute-plane-services/nvca/pkg/nvca/agent.go @@ -24,6 +24,7 @@ import ( "net/http" "os" "strings" + "sync" "sync/atomic" "time" @@ -86,6 +87,8 @@ const ( ICMSRequestAckMaxGoroutines = 20 ICMSInstanceRequestStatusUpdatesMaxGoroutines = 20 + + gracefulNoGPURegistrationComponentName = "icmsregistration" ) // Controller tick types. @@ -364,8 +367,9 @@ type ICMSClientInterface interface { type Agent struct { *AgentOptions - metricsName string - newKubeClients func(ctx context.Context, path string) (*kubeclients.KubeClients, error) + metricsName string + newKubeClients func(ctx context.Context, path string) (*kubeclients.KubeClients, error) + newBackendK8sCacheBuilder func() *BackendK8sCacheBuilder // clientMetricsShutdown releases the OTel MeterProvider that backs outbound // client metrics. It is a no-op when client metrics are disabled. @@ -383,7 +387,13 @@ type Agent struct { // gpuMonitor monitors GPU availability and controls queue processing // when GracefulNoGPU feature flag is enabled. - gpuMonitor *GPUMonitor + gpuMonitor *GPUMonitor + gpuRegistrationMu sync.Mutex + gpuRegistrationReady atomic.Bool + gpuRegistrationGeneration atomic.Uint64 + gpuRegistrationRequests chan struct{} + // Serializes GPU inventory snapshots through ICMS response and queue credential application. + registrationOperationMu sync.Mutex startControllerManager func(context.Context, *kubeclients.KubeClients) error @@ -410,6 +420,109 @@ type Agent struct { selfDestruct *atomic.Bool } +func (a *Agent) getGracefulNoGPURegistrationStatus(context.Context) (types.AgentHealth, error) { + component := types.ComponentHealth{ + Status: types.HealthStatusHealthy, + StatusLevel: types.StatusLevelError, + } + if !a.gpuRegistrationReady.Load() { + component.Status = types.HealthStatusUnhealthy + component.Errors = []string{"waiting for successful ICMS registration after GPU discovery"} + } + + return types.AgentHealth{ + Components: map[string]types.ComponentHealth{ + gracefulNoGPURegistrationComponentName: component, + }, + }, nil +} + +func (a *Agent) handleGracefulNoGPUStateChange(ctx context.Context, hasGPUs bool) { + log := core.GetLogger(ctx) + a.gpuRegistrationGeneration.Add(1) + + a.gpuRegistrationMu.Lock() + a.gpuRegistrationReady.Store(false) + a.queueManager.Pause() + a.gpuRegistrationMu.Unlock() + + if !hasGPUs { + log.Warn("GPUs no longer available - pausing queue manager") + return + } + + log.Info("GPUs detected - waiting for successful ICMS registration before resuming queue manager") + select { + case a.gpuRegistrationRequests <- struct{}{}: + default: + } +} + +func (a *Agent) runGracefulNoGPURegistration(ctx context.Context) { + retryInterval := a.GPUPollInterval + if retryInterval <= 0 { + retryInterval = DefaultGPUPollInterval + } + + for { + select { + case <-ctx.Done(): + return + case <-a.gpuRegistrationRequests: + } + + for !a.gpuRegistrationReady.Load() && a.gpuMonitor.HasGPUs() { + if !a.tryGracefulNoGPURegistration(ctx) { + break + } + + retryTimer := time.NewTimer(retryInterval) + select { + case <-ctx.Done(): + retryTimer.Stop() + return + case <-a.gpuRegistrationRequests: + retryTimer.Stop() + case <-retryTimer.C: + } + } + } +} + +// tryGracefulNoGPURegistration returns true when registration should be retried. +func (a *Agent) tryGracefulNoGPURegistration(ctx context.Context) bool { + log := core.GetLogger(ctx) + a.registrationOperationMu.Lock() + defer a.registrationOperationMu.Unlock() + if ctx.Err() != nil { + return false + } + + generation := a.gpuRegistrationGeneration.Load() + if !a.gpuMonitor.HasGPUs() { + return false + } + + log.Info("Registering with ICMS after GPUs became available") + if _, err := a.RegisterWithICMS(ctx); err != nil { + log.WithError(err).Error("Failed to register with ICMS after GPUs became available; will retry") + return a.gpuMonitor.HasGPUs() + } + + a.gpuRegistrationMu.Lock() + defer a.gpuRegistrationMu.Unlock() + if !a.gpuMonitor.HasGPUs() || generation != a.gpuRegistrationGeneration.Load() { + log.Warn("GPU availability changed during ICMS registration - keeping queue manager paused") + return a.gpuMonitor.HasGPUs() + } + + // RegisterWithICMS installs the returned queue credentials before it returns. + a.gpuRegistrationReady.Store(true) + a.queueManager.Resume() + log.Info("Successfully registered with ICMS after GPUs became available") + return false +} + func (o *AgentOptions) sanitizedString() string { var sanitized AgentOptions sanitized.NCAId = o.NCAId @@ -598,6 +711,7 @@ func NewAgent(ctx context.Context, opts *AgentOptions) (*Agent, error) { } a.newKubeClients = defaultNewKubeClients + a.newBackendK8sCacheBuilder = NewBackendk8sCacheBuilder a.icmsClient = NewICMSClientWithHostHeaderOverride(ctx, opts.ClusterID, opts.EffectiveICMSURL(), opts.ICMSHostHeaderOverride, tokenFetcher, a.tracer, icmsHTTPOpts...) a.instStatusThreadPool = pool.New().WithMaxGoroutines(ICMSInstanceRequestStatusUpdatesMaxGoroutines) a.ackThreadPool = pool.New().WithMaxGoroutines(ICMSRequestAckMaxGoroutines) @@ -1127,7 +1241,7 @@ func (a *Agent) Start(ctx context.Context) error { infraOverheadGetter := enforce.NewInfraOverheadGetter(a.FeatureFlagFetcher, a.Config, enforce.GetRuntimeClassK8sClient(k8sclients.K8s)) log.Info("Configuring backendk8scache") - backendk8scache, _, err := NewBackendk8sCacheBuilder(). + backendk8scache, _, err := a.newBackendK8sCacheBuilder(). WithConfig(a.Config). WithClusterProvider(a.CloudProvider). WithClusterRegion(a.ClusterRegion). @@ -1194,6 +1308,8 @@ func (a *Agent) Start(ctx context.Context) error { gpus, gpuErr := nfClient.GetAllBackendGPUs(ctx) hasGPUs := gpuErr == nil && len(gpus) > 0 a.gpuMonitor.SetHasGPUs(hasGPUs) + a.gpuRegistrationReady.Store(hasGPUs) + a.gpuRegistrationRequests = make(chan struct{}, 1) if hasGPUs { log.Info("GPUs found during startup, proceeding normally") } else { @@ -1208,6 +1324,8 @@ func (a *Agent) Start(ctx context.Context) error { // Add GPU monitor to status updaters for readiness checks when GracefulNoGPU is enabled if a.gpuMonitor != nil { statusUpdaters = append(statusUpdaters, a.gpuMonitor) + statusUpdaters = append(statusUpdaters, + health.GetComponentStatusFunc(a.getGracefulNoGPURegistrationStatus)) } if a.FeatureFlagFetcher.IsAttributeEnabled(featureflag.AttrHostIsolation) { statusUpdaters = append(statusUpdaters, hostisolation.NewStatusGetter( @@ -1239,6 +1357,9 @@ func (a *Agent) Start(ctx context.Context) error { if skipHealthWait { log.Warn("GracefulNoGPU enabled with no GPUs - skipping health wait, readiness will report not-ready") + if _, refreshErr := a.backendHealthCache.RefreshStatus(ctx); refreshErr != nil { + log.WithError(refreshErr).Warn("Failed to prime health status while waiting for GPUs; continuing startup") + } } else { log.WithFields(logrus.Fields{ "interval": healthInterval, @@ -1403,25 +1524,10 @@ func (a *Agent) Start(ctx context.Context) error { a.queueManager.Pause() } + go a.runGracefulNoGPURegistration(ctx) + // Set up GPU state change callback - a.gpuMonitor.SetOnGPUStateChange(func(callbackCtx context.Context, hasGPUs bool) { - callbackLog := core.GetLogger(callbackCtx) - if hasGPUs { - callbackLog.Info("GPUs detected - resuming queue manager and registering with ICMS") - // Resume queue processing - a.queueManager.Resume() - // Register/re-register with ICMS to update GPU inventory. - if _, regErr := a.RegisterWithICMS(callbackCtx); regErr != nil { - callbackLog.WithError(regErr).Error("Failed to register with ICMS after GPUs became available") - } else { - callbackLog.Info("Successfully registered with ICMS after GPUs became available") - } - } else { - callbackLog.Warn("GPUs no longer available - pausing queue manager") - // Pause queue processing (allows termination messages, blocks creation) - a.queueManager.Pause() - } - }) + a.gpuMonitor.SetOnGPUStateChange(a.handleGracefulNoGPUStateChange) // Start the GPU monitor polling loop log.Info("Starting GPU monitor") diff --git a/src/compute-plane-services/nvca/pkg/nvca/agent_test.go b/src/compute-plane-services/nvca/pkg/nvca/agent_test.go index 529e1c5f9..502cb782d 100644 --- a/src/compute-plane-services/nvca/pkg/nvca/agent_test.go +++ b/src/compute-plane-services/nvca/pkg/nvca/agent_test.go @@ -65,6 +65,7 @@ import ( nvcaerrors "github.com/NVIDIA/nvcf/src/compute-plane-services/nvca/pkg/nvca/errors" "github.com/NVIDIA/nvcf/src/compute-plane-services/nvca/pkg/nvca/health" "github.com/NVIDIA/nvcf/src/compute-plane-services/nvca/pkg/queue" + mockqueue "github.com/NVIDIA/nvcf/src/compute-plane-services/nvca/pkg/queue/mock" natsqueue "github.com/NVIDIA/nvcf/src/compute-plane-services/nvca/pkg/queue/nats" "github.com/NVIDIA/nvcf/src/compute-plane-services/nvca/pkg/types" ) @@ -109,8 +110,8 @@ func TestAgentApis(t *testing.T) { NamespaceLabels: labels.Set{"foo": "bar"}, K8sVersion: "1.27.8", CredRenewInterval: DefaultCredRenewInterval, - HeartbeatInterval: DefaultHeartBeatInterval, - SyncQueueInterval: defaultSyncQueueInterval, + HeartbeatInterval: 365 * 24 * time.Hour, + SyncQueueInterval: 365 * 24 * time.Hour, SyncRequestStatusInterval: DefaultSyncRequestStatusInterval, PeriodicInstanceStatusInterval: DefaultPeriodicInstanceStatusInterval, SyncAcknowledgeRequestInterval: ackReqInterval, @@ -192,6 +193,419 @@ func TestAgentApis(t *testing.T) { assert.Equal(t, is.LastReportedStatus, string(types.ICMSInstanceTerminated)) } +type blockingRecordingICMSClient struct { + *mockICMSClient + + mu sync.Mutex + registrationRequests []types.ICMSRegistrationRequest + results []blockingRegistrationResult + registrationStarted chan int +} + +type registrationResult struct { + response *types.ICMSRegistrationResponse + err error +} + +type blockingRegistrationResult struct { + registrationResult + release chan struct{} + releaseOnce sync.Once +} + +func newBlockingRecordingICMSClient(results ...registrationResult) *blockingRecordingICMSClient { + blockingResults := make([]blockingRegistrationResult, len(results)) + for i, result := range results { + blockingResults[i] = blockingRegistrationResult{ + registrationResult: result, + release: make(chan struct{}), + } + } + return &blockingRecordingICMSClient{ + mockICMSClient: &mockICMSClient{}, + results: blockingResults, + registrationStarted: make(chan int, len(results)), + } +} + +func (m *blockingRecordingICMSClient) Register( + ctx context.Context, + req *types.ICMSRegistrationRequest, +) (*types.ICMSRegistrationResponse, error) { + requestCopy := *req + requestCopy.BackendGPUs = append([]types.RegistrationGPU(nil), req.BackendGPUs...) + for i := range requestCopy.BackendGPUs { + requestCopy.BackendGPUs[i].InstanceTypes = append( + []types.RegistrationInstanceType(nil), + req.BackendGPUs[i].InstanceTypes..., + ) + } + + m.mu.Lock() + attempt := len(m.registrationRequests) + m.registrationRequests = append(m.registrationRequests, requestCopy) + if attempt >= len(m.results) { + m.mu.Unlock() + return nil, fmt.Errorf("unexpected ICMS registration attempt %d", attempt) + } + result := &m.results[attempt] + m.mu.Unlock() + m.registrationStarted <- attempt + + select { + case <-result.release: + case <-ctx.Done(): + return nil, ctx.Err() + } + + return result.response, result.err +} + +func (m *blockingRecordingICMSClient) release(attempt int) { + m.mu.Lock() + defer m.mu.Unlock() + result := &m.results[attempt] + result.releaseOnce.Do(func() { close(result.release) }) +} + +func (m *blockingRecordingICMSClient) requests() []types.ICMSRegistrationRequest { + m.mu.Lock() + defer m.mu.Unlock() + return append([]types.ICMSRegistrationRequest(nil), m.registrationRequests...) +} + +func requireRegistrationAttempt(t *testing.T, client *blockingRecordingICMSClient, expected int) { + t.Helper() + select { + case attempt := <-client.registrationStarted: + require.Equal(t, expected, attempt) + case <-time.After(5 * time.Second): + t.Fatalf("timed out waiting for ICMS registration attempt %d", expected) + } +} + +func requireNoRegistrationAttempt(t *testing.T, client *blockingRecordingICMSClient, wait time.Duration) { + t.Helper() + select { + case attempt := <-client.registrationStarted: + t.Fatalf("unexpected ICMS registration attempt %d while another registration is in flight", attempt) + case <-time.After(wait): + } +} + +func newGracefulNoGPUTestAgent( + t *testing.T, + ctx context.Context, + icmsClient ICMSClientInterface, +) (*Agent, *kubeclients.KubeClients) { + t.Helper() + oldSyncInterval := syncICMSRegistrationInterval + // Callers use a fixed random seed, keeping periodic registration outside the test window. + syncICMSRegistrationInterval = 365 * 24 * time.Hour + t.Cleanup(func() { syncICMSRegistrationInterval = oldSyncInterval }) + + featureFlags := &featureflagmock.Fetcher{} + featureFlags.SetFeatureFlags(featureflag.GracefulNoGPU) + agentOpts := AgentOptions{ + TokenFetcherOptions: nvcaauth.TokenFetcherOptions{ + OAuthTokenScope: "byoc_registration", + OAuthClientID: "foo", + OAuthClientSecretKey: "bar", + }, + NCAId: "randomNCAId123", + ClusterName: "bartnvbackend", + ClusterID: "clusterid-1", + ClusterDescription: "this is a test cluster", + ClusterGroupName: "group of all A30", + ComputeBackend: "k8s", + CloudProvider: "on-prem", + NamespaceLabels: labels.Set{"foo": "bar"}, + K8sVersion: "1.27.8", + CredRenewInterval: DefaultCredRenewInterval, + HeartbeatInterval: DefaultHeartBeatInterval, + SyncQueueInterval: defaultSyncQueueInterval, + SyncRequestStatusInterval: DefaultSyncRequestStatusInterval, + PeriodicInstanceStatusInterval: DefaultPeriodicInstanceStatusInterval, + SyncAcknowledgeRequestInterval: ackReqInterval, + DynamicGPUDiscoveryEnabled: true, + MultipleGPUTypesAllowed: true, + UniformInstanceLabelsEnabled: true, + GPUPollInterval: 10 * time.Millisecond, + GPUDebounceTime: time.Millisecond, + FeatureFlagFetcher: featureFlags, + MetricsRegisterer: prometheus.NewRegistry(), + } + + agent := newMockAgent(t, ctx, agentOpts) + // Isolate startup readiness from the immediate heartbeat and queue-sync events. + delete(agent.resourceEventWorkerQueues, EventTickUpdateHeartbeat) + delete(agent.resourceEventWorkerQueues, EventTickSyncSQSQueue) + oldNewQueueClient := newQueueClient + newQueueClient = func(string) queue.Client { + return &mockqueue.Client{Use10MillisForWaits: true} + } + t.Cleanup(func() { newQueueClient = oldNewQueueClient }) + k8sClients := mockKubeClientsDynamicGPUs() + agent.newKubeClients = func(context.Context, string) (*kubeclients.KubeClients, error) { + return k8sClients, nil + } + agent.newBackendK8sCacheBuilder = func() *BackendK8sCacheBuilder { + builder := NewBackendk8sCacheBuilder() + builder.addSharedClusterNodePublisher = mockAddSharedClusterNodePublisherFunc + return builder + } + agent.icmsClient = icmsClient + return agent, k8sClients +} + +func requireHTTPStatusEventually(t *testing.T, address, path string, expected int) { + t.Helper() + require.EventuallyWithT(t, func(ct *assert.CollectT) { + resp, err := http.Get("http://" + address + path) + if !assert.NoError(ct, err) { + return + } + assert.NoError(ct, resp.Body.Close()) + assert.Equal(ct, expected, resp.StatusCode) + }, 5*time.Second, 10*time.Millisecond) +} + +func requireHTTPStatus(t *testing.T, address, path string, expected int) { + t.Helper() + resp, err := http.Get("http://" + address + path) + require.NoError(t, err) + require.NoError(t, resp.Body.Close()) + require.Equal(t, expected, resp.StatusCode) +} + +func requireRefreshedReadinessStatus( + t *testing.T, + ctx context.Context, + agent *Agent, + expected int, +) { + t.Helper() + _, err := agent.backendHealthCache.RefreshStatus(ctx) + require.NoError(t, err) + requireHTTPStatusEventually(t, agent.NVCASvcAddress, health.HTTPReadinessRoutePath, expected) +} + +func TestAgentStartGracefulNoGPURecoversWhenGPUAppears(t *testing.T) { + ctx, cancel := context.WithCancel(core.WithRandomSeed(newTestContext(), 42)) + t.Cleanup(cancel) + + recoveredCredentials := getTestQueueCreds(true) + icmsClient := newBlockingRecordingICMSClient(registrationResult{ + response: &types.ICMSRegistrationResponse{ + ClusterID: "registered-cluster-id", + ClusterGroupID: "registered-cluster-group-id", + Credentials: recoveredCredentials, + }, + }) + agent, k8sClients := newGracefulNoGPUTestAgent(t, ctx, icmsClient) + + require.NoError(t, agent.Start(ctx)) + requireHTTPStatus(t, agent.NVCASvcAddress, health.HTTPReadinessRoutePath, http.StatusServiceUnavailable) + requireHTTPStatus(t, agent.NVCASvcAddress, health.HTTPLivenessRoutePath, http.StatusOK) + require.NotNil(t, agent.gpuMonitor) + require.NotNil(t, agent.queueManager) + assert.False(t, agent.gpuMonitor.HasGPUs()) + assert.True(t, agent.queueManager.IsPaused()) + assert.Empty(t, icmsClient.requests()) + assert.Empty(t, agent.queueManager.getCreateQueue(testGPUNameDefault).QueueURL) + + _, err := k8sClients.K8s.CoreV1().Nodes().Create(ctx, functionNode.DeepCopy(), metav1.CreateOptions{}) + require.NoError(t, err) + requireRegistrationAttempt(t, icmsClient, 0) + + assert.True(t, agent.queueManager.IsPaused(), "queue must remain paused while registration is in flight") + assert.Empty(t, agent.queueManager.getCreateQueue(testGPUNameDefault).QueueURL) + requireRefreshedReadinessStatus(t, ctx, agent, http.StatusServiceUnavailable) + requireHTTPStatusEventually(t, agent.NVCASvcAddress, health.HTTPLivenessRoutePath, http.StatusOK) + requests := icmsClient.requests() + require.Len(t, requests, 1) + assert.Equal(t, []types.RegistrationGPU{{ + Name: "A100", + InstanceTypes: []types.RegistrationInstanceType{{ + Name: "ON-PREM.GPU.A100_1x", + Value: "ON-PREM.GPU.A100", + Description: "A100-SXM4-40GB (ampere family) on a Google-Compute-Engine machine", + Default: true, + CPUCores: 6, + CPU: "6", + SystemMemory: "32Gi", + GPUCount: 1, + GPUMemory: "40Gi", + Storage: "512Gi", + CPUArch: "unknown", + OS: "unknown", + DriverVersion: "unknown", + NodeType: types.RegistrationInstanceTypeNodeTypeSingle, + MaxInstances: 1, + }}, + }}, requests[0].BackendGPUs) + + icmsClient.release(0) + + require.EventuallyWithT(t, func(ct *assert.CollectT) { + assert.True(ct, agent.gpuMonitor.HasGPUs()) + assert.False(ct, agent.queueManager.IsPaused()) + assert.Equal(ct, recoveredCredentials.CreationQueues[testGPUNameDefault], + agent.queueManager.getCreateQueue(testGPUNameDefault)) + assert.Equal(ct, recoveredCredentials.TerminationQueue, agent.queueManager.getTermQueue()) + }, 5*time.Second, 10*time.Millisecond) + requireRefreshedReadinessStatus(t, ctx, agent, http.StatusOK) + requireHTTPStatusEventually(t, agent.NVCASvcAddress, health.HTTPLivenessRoutePath, http.StatusOK) + + require.NoError(t, k8sClients.K8s.CoreV1().Nodes().Delete(ctx, functionNode.Name, metav1.DeleteOptions{})) + require.EventuallyWithT(t, func(ct *assert.CollectT) { + assert.False(ct, agent.gpuMonitor.HasGPUs()) + assert.True(ct, agent.queueManager.IsPaused()) + }, 5*time.Second, 10*time.Millisecond) + requireRefreshedReadinessStatus(t, ctx, agent, http.StatusServiceUnavailable) + requireHTTPStatusEventually(t, agent.NVCASvcAddress, health.HTTPLivenessRoutePath, http.StatusOK) + require.Len(t, icmsClient.requests(), 1) +} + +func TestAgentStartGracefulNoGPURegistrationFailureRetriesBeforeResuming(t *testing.T) { + ctx, cancel := context.WithCancel(core.WithRandomSeed(newTestContext(), 42)) + t.Cleanup(cancel) + + recoveredCredentials := getTestQueueCreds(true) + icmsClient := newBlockingRecordingICMSClient( + registrationResult{err: fmt.Errorf("registration unavailable")}, + registrationResult{response: &types.ICMSRegistrationResponse{ + ClusterID: "registered-cluster-id", + ClusterGroupID: "registered-cluster-group-id", + Credentials: recoveredCredentials, + }}, + ) + agent, k8sClients := newGracefulNoGPUTestAgent(t, ctx, icmsClient) + require.NoError(t, agent.Start(ctx)) + requireHTTPStatus(t, agent.NVCASvcAddress, health.HTTPReadinessRoutePath, http.StatusServiceUnavailable) + requireHTTPStatus(t, agent.NVCASvcAddress, health.HTTPLivenessRoutePath, http.StatusOK) + + _, err := k8sClients.K8s.CoreV1().Nodes().Create(ctx, functionNode.DeepCopy(), metav1.CreateOptions{}) + require.NoError(t, err) + requireRegistrationAttempt(t, icmsClient, 0) + assert.True(t, agent.queueManager.IsPaused(), "queue must remain paused while registration is in flight") + + icmsClient.release(0) + requireRegistrationAttempt(t, icmsClient, 1) + assert.True(t, agent.queueManager.IsPaused(), "queue must remain paused while registration retry is in flight") + assert.Empty(t, agent.queueManager.getCreateQueue(testGPUNameDefault).QueueURL) + requireRefreshedReadinessStatus(t, ctx, agent, http.StatusServiceUnavailable) + requireHTTPStatusEventually(t, agent.NVCASvcAddress, health.HTTPLivenessRoutePath, http.StatusOK) + + icmsClient.release(1) + require.EventuallyWithT(t, func(ct *assert.CollectT) { + assert.False(ct, agent.queueManager.IsPaused()) + assert.Equal(ct, recoveredCredentials.CreationQueues[testGPUNameDefault], + agent.queueManager.getCreateQueue(testGPUNameDefault)) + assert.Equal(ct, recoveredCredentials.TerminationQueue, agent.queueManager.getTermQueue()) + }, 5*time.Second, 10*time.Millisecond) + requireRefreshedReadinessStatus(t, ctx, agent, http.StatusOK) + requireHTTPStatusEventually(t, agent.NVCASvcAddress, health.HTTPLivenessRoutePath, http.StatusOK) + require.Len(t, icmsClient.requests(), 2) +} + +func TestAgentStartGracefulNoGPUSerializesPeriodicRegistrationBeforeRecovery(t *testing.T) { + ctx, cancel := context.WithCancel(core.WithRandomSeed(newTestContext(), 42)) + t.Cleanup(cancel) + + initialCredentials := getTestQueueCreds(false) + stalePeriodicCredentials := getTestQueueCreds(false) + recoveredGPU := types.GPUName("AD102GL") + recoveredQueue := getTestCreationMessageQueueInfo(true) + recoveredQueue.GPU = string(recoveredGPU) + recoveredCredentials := getTestQueueCreds(true) + recoveredCredentials.CreationQueues = types.CreationQueueInfoSet{ + recoveredGPU: recoveredQueue, + } + icmsClient := newBlockingRecordingICMSClient( + registrationResult{response: &types.ICMSRegistrationResponse{ + ClusterID: "registered-cluster-id", + ClusterGroupID: "registered-cluster-group-id", + Credentials: initialCredentials, + }}, + registrationResult{response: &types.ICMSRegistrationResponse{ + ClusterID: "registered-cluster-id", + ClusterGroupID: "registered-cluster-group-id", + Credentials: stalePeriodicCredentials, + }}, + registrationResult{response: &types.ICMSRegistrationResponse{ + ClusterID: "registered-cluster-id", + ClusterGroupID: "registered-cluster-group-id", + Credentials: recoveredCredentials, + }}, + ) + agent, k8sClients := newGracefulNoGPUTestAgent(t, ctx, icmsClient) + _, err := k8sClients.K8s.CoreV1().Nodes().Create(ctx, functionNode.DeepCopy(), metav1.CreateOptions{}) + require.NoError(t, err) + icmsClient.release(0) + require.NoError(t, agent.Start(ctx)) + requireRegistrationAttempt(t, icmsClient, 0) + assert.False(t, agent.queueManager.IsPaused()) + + periodicDone := make(chan error, 1) + go func() { + periodicDone <- agent.syncICMSRegistration(ctx) + }() + requireRegistrationAttempt(t, icmsClient, 1) + requests := icmsClient.requests() + require.Len(t, requests, 2) + require.Len(t, requests[1].BackendGPUs, 1) + assert.Equal(t, "A100", requests[1].BackendGPUs[0].Name) + + require.NoError(t, k8sClients.K8s.CoreV1().Nodes().Delete(ctx, functionNode.Name, metav1.DeleteOptions{})) + require.EventuallyWithT(t, func(ct *assert.CollectT) { + assert.False(ct, agent.gpuMonitor.HasGPUs()) + assert.True(ct, agent.queueManager.IsPaused()) + }, 5*time.Second, 10*time.Millisecond) + + recoveryNode := functionNode.DeepCopy() + recoveryNode.Name = "node-2" + recoveryNode.Labels[nodefeatures.UniformInstanceTypeLabelKey] = "ON-PREM.GPU.AD102GL" + recoveryNode.Labels["nvidia.com/gpu.family"] = "volta" + recoveryNode.Labels["nvidia.com/gpu.memory"] = "32768" + recoveryNode.Labels["nvidia.com/gpu.product"] = "V100-SXM2-32GB" + recoveryNode.Labels["nvca.nvcf.nvidia.io/gpu.product"] = string(recoveredGPU) + lossGeneration := agent.gpuRegistrationGeneration.Load() + _, err = k8sClients.K8s.CoreV1().Nodes().Create(ctx, recoveryNode, metav1.CreateOptions{}) + require.NoError(t, err) + require.EventuallyWithT(t, func(ct *assert.CollectT) { + assert.True(ct, agent.gpuMonitor.HasGPUs()) + assert.Greater(ct, agent.gpuRegistrationGeneration.Load(), lossGeneration) + }, 5*time.Second, 10*time.Millisecond) + requireNoRegistrationAttempt(t, icmsClient, 250*time.Millisecond) + assert.True(t, agent.queueManager.IsPaused()) + assert.Empty(t, agent.queueManager.getCreateQueue(recoveredGPU).QueueURL) + + icmsClient.release(1) + select { + case periodicErr := <-periodicDone: + require.NoError(t, periodicErr) + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for periodic registration to complete") + } + requireRegistrationAttempt(t, icmsClient, 2) + requests = icmsClient.requests() + require.Len(t, requests, 3) + require.Len(t, requests[2].BackendGPUs, 1) + assert.Equal(t, recoveredGPU, types.GPUName(requests[2].BackendGPUs[0].Name)) + assert.True(t, agent.queueManager.IsPaused()) + assert.Empty(t, agent.queueManager.getCreateQueue(recoveredGPU).QueueURL) + + icmsClient.release(2) + require.EventuallyWithT(t, func(ct *assert.CollectT) { + assert.False(ct, agent.queueManager.IsPaused()) + assert.Equal(ct, recoveredQueue, agent.queueManager.getCreateQueue(recoveredGPU)) + assert.Equal(ct, recoveredCredentials.TerminationQueue, agent.queueManager.getTermQueue()) + }, 5*time.Second, 10*time.Millisecond) + requireRefreshedReadinessStatus(t, ctx, agent, http.StatusOK) + requireHTTPStatusEventually(t, agent.NVCASvcAddress, health.HTTPLivenessRoutePath, http.StatusOK) +} + func TestAgentRegisterWithICMSUpdatesQueueManagerCredentials(t *testing.T) { ctx, cancel := context.WithCancel(newTestContext()) t.Cleanup(cancel) diff --git a/src/compute-plane-services/nvca/pkg/nvca/agent_updates.go b/src/compute-plane-services/nvca/pkg/nvca/agent_updates.go index 9fb158b4f..1b7977dad 100644 --- a/src/compute-plane-services/nvca/pkg/nvca/agent_updates.go +++ b/src/compute-plane-services/nvca/pkg/nvca/agent_updates.go @@ -74,6 +74,12 @@ func (a *Agent) initICMSRegistrationSyncer(ctx context.Context) error { var lastRegTime = core.GetCurrentTime(ctx) a.syncICMSRegistration = func(ctx context.Context) error { + a.registrationOperationMu.Lock() + defer a.registrationOperationMu.Unlock() + if err := ctx.Err(); err != nil { + return err + } + if timeSinceLastUpdate := core.GetCurrentTime(ctx).Sub(lastRegTime); timeSinceLastUpdate < time.Hour { return nil } @@ -98,6 +104,12 @@ func (a *Agent) initICMSRegistrationSyncer(ctx context.Context) error { var lastRegTime = core.GetCurrentTime(ctx) a.syncICMSRegistration = func(ctx context.Context) error { + a.registrationOperationMu.Lock() + defer a.registrationOperationMu.Unlock() + if err := ctx.Err(); err != nil { + return err + } + regBackendGPUs, err := a.getRegistrationGPUs(ctx) if err != nil { return err From 4e0f23f4a25f147a11475187234931d32e844655 Mon Sep 17 00:00:00 2001 From: Mike Camp Date: Wed, 26 Aug 2026 14:25:22 -0400 Subject: [PATCH 2/7] fix(nvca): log recoverable registration retries as warnings Signed-off-by: Mike Camp --- src/compute-plane-services/nvca/pkg/nvca/agent.go | 2 +- .../nvca/pkg/nvca/agent_test.go | 11 ++++++++++- 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/src/compute-plane-services/nvca/pkg/nvca/agent.go b/src/compute-plane-services/nvca/pkg/nvca/agent.go index 2891737fa..cb4325b11 100644 --- a/src/compute-plane-services/nvca/pkg/nvca/agent.go +++ b/src/compute-plane-services/nvca/pkg/nvca/agent.go @@ -505,7 +505,7 @@ func (a *Agent) tryGracefulNoGPURegistration(ctx context.Context) bool { log.Info("Registering with ICMS after GPUs became available") if _, err := a.RegisterWithICMS(ctx); err != nil { - log.WithError(err).Error("Failed to register with ICMS after GPUs became available; will retry") + log.WithError(err).Warn("Failed to register with ICMS after GPUs became available; will retry") return a.gpuMonitor.HasGPUs() } diff --git a/src/compute-plane-services/nvca/pkg/nvca/agent_test.go b/src/compute-plane-services/nvca/pkg/nvca/agent_test.go index 502cb782d..1152ecfc8 100644 --- a/src/compute-plane-services/nvca/pkg/nvca/agent_test.go +++ b/src/compute-plane-services/nvca/pkg/nvca/agent_test.go @@ -468,7 +468,8 @@ func TestAgentStartGracefulNoGPURecoversWhenGPUAppears(t *testing.T) { } func TestAgentStartGracefulNoGPURegistrationFailureRetriesBeforeResuming(t *testing.T) { - ctx, cancel := context.WithCancel(core.WithRandomSeed(newTestContext(), 42)) + logCtx, logHook := core.WithTestingLogger(newTestContext()) + ctx, cancel := context.WithCancel(core.WithRandomSeed(logCtx, 42)) t.Cleanup(cancel) recoveredCredentials := getTestQueueCreds(true) @@ -492,6 +493,14 @@ func TestAgentStartGracefulNoGPURegistrationFailureRetriesBeforeResuming(t *test icmsClient.release(0) requireRegistrationAttempt(t, icmsClient, 1) + retryLogs := make([]*logrus.Entry, 0, 1) + for _, entry := range logHook.AllEntries() { + if entry.Message == "Failed to register with ICMS after GPUs became available; will retry" { + retryLogs = append(retryLogs, entry) + } + } + require.Len(t, retryLogs, 1) + assert.Equal(t, logrus.WarnLevel, retryLogs[0].Level) assert.True(t, agent.queueManager.IsPaused(), "queue must remain paused while registration retry is in flight") assert.Empty(t, agent.queueManager.getCreateQueue(testGPUNameDefault).QueueURL) requireRefreshedReadinessStatus(t, ctx, agent, http.StatusServiceUnavailable) From 8fdd6cf6dade11031e44bc00758318efb3347498 Mon Sep 17 00:00:00 2001 From: Mike Camp Date: Mon, 31 Aug 2026 16:44:26 +0200 Subject: [PATCH 3/7] fix(nvca): serialize credential refresh with registration Signed-off-by: Mike Camp --- .../nvca/pkg/nvca/agent.go | 8 +- .../nvca/pkg/nvca/agent_test.go | 175 +++++++++++++++++- 2 files changed, 177 insertions(+), 6 deletions(-) diff --git a/src/compute-plane-services/nvca/pkg/nvca/agent.go b/src/compute-plane-services/nvca/pkg/nvca/agent.go index cb4325b11..a024ff618 100644 --- a/src/compute-plane-services/nvca/pkg/nvca/agent.go +++ b/src/compute-plane-services/nvca/pkg/nvca/agent.go @@ -392,7 +392,7 @@ type Agent struct { gpuRegistrationReady atomic.Bool gpuRegistrationGeneration atomic.Uint64 gpuRegistrationRequests chan struct{} - // Serializes GPU inventory snapshots through ICMS response and queue credential application. + // Serializes ICMS registration and credential refresh through queue credential application. registrationOperationMu sync.Mutex startControllerManager func(context.Context, *kubeclients.KubeClients) error @@ -2262,6 +2262,12 @@ func (a *Agent) RenewICMSQueueCreds(ctx context.Context) error { return nil } + a.registrationOperationMu.Lock() + defer a.registrationOperationMu.Unlock() + if err := ctx.Err(); err != nil { + return err + } + credRes, err := a.icmsClient.GetCreds(ctx) nvcametrics.FromContext(ctx).RecordUpstreamRequest(nvcametrics.UpstreamOperationCredentials, err) if err != nil { diff --git a/src/compute-plane-services/nvca/pkg/nvca/agent_test.go b/src/compute-plane-services/nvca/pkg/nvca/agent_test.go index 1152ecfc8..bcfc3deed 100644 --- a/src/compute-plane-services/nvca/pkg/nvca/agent_test.go +++ b/src/compute-plane-services/nvca/pkg/nvca/agent_test.go @@ -196,10 +196,15 @@ func TestAgentApis(t *testing.T) { type blockingRecordingICMSClient struct { *mockICMSClient - mu sync.Mutex - registrationRequests []types.ICMSRegistrationRequest - results []blockingRegistrationResult - registrationStarted chan int + mu sync.Mutex + registrationRequests []types.ICMSRegistrationRequest + results []blockingRegistrationResult + registrationStarted chan int + credentialResponse *types.ICMSCredentialResponse + credentialErr error + credentialStarted chan struct{} + credentialRelease chan struct{} + credentialReleaseOnce sync.Once } type registrationResult struct { @@ -261,6 +266,47 @@ func (m *blockingRecordingICMSClient) Register( return result.response, result.err } +func (m *blockingRecordingICMSClient) GetCreds(ctx context.Context) (*types.ICMSCredentialResponse, error) { + m.mu.Lock() + credentialStarted := m.credentialStarted + credentialRelease := m.credentialRelease + credentialResponse := m.credentialResponse + credentialErr := m.credentialErr + m.mu.Unlock() + if credentialStarted == nil { + return m.mockICMSClient.GetCreds(ctx) + } + + select { + case credentialStarted <- struct{}{}: + default: + } + select { + case <-credentialRelease: + return credentialResponse, credentialErr + case <-ctx.Done(): + return nil, ctx.Err() + } +} + +func (m *blockingRecordingICMSClient) blockCredentialFetch( + response *types.ICMSCredentialResponse, + err error, +) { + m.mu.Lock() + defer m.mu.Unlock() + m.credentialResponse = response + m.credentialErr = err + m.credentialStarted = make(chan struct{}, 1) + m.credentialRelease = make(chan struct{}) +} + +func (m *blockingRecordingICMSClient) releaseCredentialFetch() { + m.mu.Lock() + defer m.mu.Unlock() + m.credentialReleaseOnce.Do(func() { close(m.credentialRelease) }) +} + func (m *blockingRecordingICMSClient) release(attempt int) { m.mu.Lock() defer m.mu.Unlock() @@ -321,7 +367,7 @@ func newGracefulNoGPUTestAgent( CloudProvider: "on-prem", NamespaceLabels: labels.Set{"foo": "bar"}, K8sVersion: "1.27.8", - CredRenewInterval: DefaultCredRenewInterval, + CredRenewInterval: 365 * 24 * time.Hour, HeartbeatInterval: DefaultHeartBeatInterval, SyncQueueInterval: defaultSyncQueueInterval, SyncRequestStatusInterval: DefaultSyncRequestStatusInterval, @@ -518,6 +564,125 @@ func TestAgentStartGracefulNoGPURegistrationFailureRetriesBeforeResuming(t *test require.Len(t, icmsClient.requests(), 2) } +func TestAgentStartGracefulNoGPUStaysPausedWhenGPUDisappearsDuringRegistration(t *testing.T) { + ctx, cancel := context.WithCancel(core.WithRandomSeed(newTestContext(), 42)) + t.Cleanup(cancel) + + recoveredCredentials := getTestQueueCreds(true) + icmsClient := newBlockingRecordingICMSClient( + registrationResult{response: &types.ICMSRegistrationResponse{ + ClusterID: "registered-cluster-id", + ClusterGroupID: "registered-cluster-group-id", + Credentials: recoveredCredentials, + }}, + registrationResult{response: &types.ICMSRegistrationResponse{ + ClusterID: "registered-cluster-id", + ClusterGroupID: "registered-cluster-group-id", + Credentials: recoveredCredentials, + }}, + ) + agent, k8sClients := newGracefulNoGPUTestAgent(t, ctx, icmsClient) + require.NoError(t, agent.Start(ctx)) + + _, err := k8sClients.K8s.CoreV1().Nodes().Create(ctx, functionNode.DeepCopy(), metav1.CreateOptions{}) + require.NoError(t, err) + requireRegistrationAttempt(t, icmsClient, 0) + require.NoError(t, k8sClients.K8s.CoreV1().Nodes().Delete(ctx, functionNode.Name, metav1.DeleteOptions{})) + require.EventuallyWithT(t, func(ct *assert.CollectT) { + assert.False(ct, agent.gpuMonitor.HasGPUs()) + assert.True(ct, agent.queueManager.IsPaused()) + }, 5*time.Second, 10*time.Millisecond) + + icmsClient.release(0) + requireNoRegistrationAttempt(t, icmsClient, 100*time.Millisecond) + assert.True(t, agent.queueManager.IsPaused()) + requireRefreshedReadinessStatus(t, ctx, agent, http.StatusServiceUnavailable) + requireHTTPStatusEventually(t, agent.NVCASvcAddress, health.HTTPLivenessRoutePath, http.StatusOK) + + _, err = k8sClients.K8s.CoreV1().Nodes().Create(ctx, functionNode.DeepCopy(), metav1.CreateOptions{}) + require.NoError(t, err) + requireRegistrationAttempt(t, icmsClient, 1) + assert.True(t, agent.queueManager.IsPaused()) + icmsClient.release(1) + require.EventuallyWithT(t, func(ct *assert.CollectT) { + assert.False(ct, agent.queueManager.IsPaused()) + assert.Equal(ct, recoveredCredentials.CreationQueues[testGPUNameDefault], + agent.queueManager.getCreateQueue(testGPUNameDefault)) + }, 5*time.Second, 10*time.Millisecond) + requireRefreshedReadinessStatus(t, ctx, agent, http.StatusOK) + requireHTTPStatusEventually(t, agent.NVCASvcAddress, health.HTTPLivenessRoutePath, http.StatusOK) +} + +func TestAgentStartGracefulNoGPUSerializesCredentialRenewalBeforeRecovery(t *testing.T) { + ctx, cancel := context.WithCancel(core.WithRandomSeed(newTestContext(), 42)) + t.Cleanup(cancel) + + initialCredentials := getTestQueueCreds(false) + staleCredentials := getTestQueueCreds(false) + recoveredCredentials := getTestQueueCreds(true) + icmsClient := newBlockingRecordingICMSClient( + registrationResult{response: &types.ICMSRegistrationResponse{ + ClusterID: "registered-cluster-id", + ClusterGroupID: "registered-cluster-group-id", + Credentials: initialCredentials, + }}, + registrationResult{response: &types.ICMSRegistrationResponse{ + ClusterID: "registered-cluster-id", + ClusterGroupID: "registered-cluster-group-id", + Credentials: recoveredCredentials, + }}, + ) + agent, k8sClients := newGracefulNoGPUTestAgent(t, ctx, icmsClient) + _, err := k8sClients.K8s.CoreV1().Nodes().Create(ctx, functionNode.DeepCopy(), metav1.CreateOptions{}) + require.NoError(t, err) + icmsClient.release(0) + require.NoError(t, agent.Start(ctx)) + requireRegistrationAttempt(t, icmsClient, 0) + assert.False(t, agent.queueManager.IsPaused()) + + require.NoError(t, k8sClients.K8s.CoreV1().Nodes().Delete(ctx, functionNode.Name, metav1.DeleteOptions{})) + require.EventuallyWithT(t, func(ct *assert.CollectT) { + assert.False(ct, agent.gpuMonitor.HasGPUs()) + assert.True(ct, agent.queueManager.IsPaused()) + }, 5*time.Second, 10*time.Millisecond) + + icmsClient.blockCredentialFetch(&types.ICMSCredentialResponse{QueueCredentials: staleCredentials}, nil) + credentialDone := make(chan error, 1) + go func() { + credentialDone <- agent.RenewICMSQueueCreds(ctx) + }() + select { + case <-icmsClient.credentialStarted: + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for credential renewal to start") + } + + _, err = k8sClients.K8s.CoreV1().Nodes().Create(ctx, functionNode.DeepCopy(), metav1.CreateOptions{}) + require.NoError(t, err) + requireNoRegistrationAttempt(t, icmsClient, 250*time.Millisecond) + assert.True(t, agent.queueManager.IsPaused()) + + icmsClient.releaseCredentialFetch() + select { + case credentialErr := <-credentialDone: + require.NoError(t, credentialErr) + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for credential renewal to finish") + } + requireRegistrationAttempt(t, icmsClient, 1) + assert.True(t, agent.queueManager.IsPaused()) + icmsClient.release(1) + + require.EventuallyWithT(t, func(ct *assert.CollectT) { + assert.False(ct, agent.queueManager.IsPaused()) + assert.Equal(ct, recoveredCredentials.CreationQueues[testGPUNameDefault], + agent.queueManager.getCreateQueue(testGPUNameDefault)) + assert.Equal(ct, recoveredCredentials.TerminationQueue, agent.queueManager.getTermQueue()) + }, 5*time.Second, 10*time.Millisecond) + requireRefreshedReadinessStatus(t, ctx, agent, http.StatusOK) + requireHTTPStatusEventually(t, agent.NVCASvcAddress, health.HTTPLivenessRoutePath, http.StatusOK) +} + func TestAgentStartGracefulNoGPUSerializesPeriodicRegistrationBeforeRecovery(t *testing.T) { ctx, cancel := context.WithCancel(core.WithRandomSeed(newTestContext(), 42)) t.Cleanup(cancel) From 0333bbb9cd29a0feabf3e4b7125dd9118e6eda6e Mon Sep 17 00:00:00 2001 From: Mike Camp Date: Mon, 31 Aug 2026 19:13:14 +0200 Subject: [PATCH 4/7] fix(nvca): keep no-GPU rollouts single-active Signed-off-by: Mike Camp --- .../nvca/pkg/operator/reconcile/nvcaagent_reconcile.go | 8 ++++++++ .../pkg/operator/reconcile/nvcaagent_reconcile_test.go | 4 ++++ 2 files changed, 12 insertions(+) diff --git a/src/compute-plane-services/nvca/pkg/operator/reconcile/nvcaagent_reconcile.go b/src/compute-plane-services/nvca/pkg/operator/reconcile/nvcaagent_reconcile.go index a99f3437c..5ddc0bcff 100644 --- a/src/compute-plane-services/nvca/pkg/operator/reconcile/nvcaagent_reconcile.go +++ b/src/compute-plane-services/nvca/pkg/operator/reconcile/nvcaagent_reconcile.go @@ -2171,6 +2171,14 @@ func (bc *BackendK8sCache) setupNVCADeployment(ctx context.Context, original *nv }, }, } + // GracefulNoGPU deliberately keeps the agent NotReady while the cluster has + // no GPUs. A rolling update would therefore surge a second singleton agent + // and retain the old replica indefinitely. Recreate keeps configuration + // rollouts single-active while preserving the default rollout behavior for + // clusters that do not opt in. + if slices.Contains(strings.Split(bc.getNVCAFeatureFlags(nb), ","), featureflag.GracefulNoGPU.Key) { + deployment.Spec.Strategy = appsv1.DeploymentStrategy{Type: appsv1.RecreateDeploymentStrategyType} + } if nb.Spec.VaultConfig.Enabled { deployment.Spec.Template.Annotations = mergeMaps(deployment.Spec.Template.Annotations, getVaultAnnotations(nb)) diff --git a/src/compute-plane-services/nvca/pkg/operator/reconcile/nvcaagent_reconcile_test.go b/src/compute-plane-services/nvca/pkg/operator/reconcile/nvcaagent_reconcile_test.go index 145510c09..3883aaa7f 100644 --- a/src/compute-plane-services/nvca/pkg/operator/reconcile/nvcaagent_reconcile_test.go +++ b/src/compute-plane-services/nvca/pkg/operator/reconcile/nvcaagent_reconcile_test.go @@ -273,6 +273,7 @@ func TestSetupNVCADeployment(t *testing.T) { Values: []string{ "LogPosting", "CachingSupport", + "GracefulNoGPU", "PeriodicInstanceStatusUpdate", "SharedCluster", }, @@ -322,6 +323,7 @@ func TestSetupNVCADeployment(t *testing.T) { LogLevel: "info", FeatureFlags: []string{ "CachingSupport", + "GracefulNoGPU", "LogPosting", "PeriodicInstanceStatusUpdate", "SharedCluster", @@ -538,6 +540,7 @@ func TestSetupNVCADeployment(t *testing.T) { assert.Equal(t, expectedAnnotations, gotSvc.Annotations) assert.Equal(t, expectedAnnotations, gotDep.Annotations) assert.Empty(t, gotDep.Spec.Template.Annotations) + assert.Equal(t, appsv1.RecreateDeploymentStrategyType, gotDep.Spec.Strategy.Type) // Try rollout with the same spec. err = bc.setupNVCADeployment(ctx, inNVCFBackend) @@ -631,6 +634,7 @@ func TestSetupNVCADeployment_OverrideEnvironmentVars(t *testing.T) { gotDep, getErr = depIface.Get(ctx, nvcaoptypes.NVCAModuleName, metav1.GetOptions{}) require.NoError(ct, getErr) }, 10*time.Second, 100*time.Millisecond) + assert.Empty(t, gotDep.Spec.Strategy.Type, "default rollout strategy must remain unchanged") var nvcaContainer *corev1.Container for i := range gotDep.Spec.Template.Spec.Containers { From 551d084fe4b2bf33456ba849f7a8d25d16af2a4c Mon Sep 17 00:00:00 2001 From: Mike Camp Date: Mon, 31 Aug 2026 21:24:48 +0200 Subject: [PATCH 5/7] refactor(nvca): isolate GPU registration coordination Signed-off-by: Mike Camp --- .../nvca/pkg/nvca/BUILD.bazel | 2 + .../nvca/pkg/nvca/agent.go | 193 +++------------ .../nvca/pkg/nvca/agent_test.go | 20 +- .../nvca/pkg/nvca/agent_updates.go | 88 +++---- .../nvca/pkg/nvca/gpu_registration_manager.go | 231 ++++++++++++++++++ .../pkg/nvca/gpu_registration_manager_test.go | 153 ++++++++++++ 6 files changed, 476 insertions(+), 211 deletions(-) create mode 100644 src/compute-plane-services/nvca/pkg/nvca/gpu_registration_manager.go create mode 100644 src/compute-plane-services/nvca/pkg/nvca/gpu_registration_manager_test.go diff --git a/src/compute-plane-services/nvca/pkg/nvca/BUILD.bazel b/src/compute-plane-services/nvca/pkg/nvca/BUILD.bazel index 23eeb1d7d..59d73f970 100644 --- a/src/compute-plane-services/nvca/pkg/nvca/BUILD.bazel +++ b/src/compute-plane-services/nvca/pkg/nvca/BUILD.bazel @@ -16,6 +16,7 @@ go_library( "cli.go", "computebackend.go", "gpumonitor.go", + "gpu_registration_manager.go", "icms_client.go", "jwks_updater.go", "k8scomputebackend.go", @@ -181,6 +182,7 @@ go_test( "cli_test.go", "encrypt_modelcache_test.go", "gpumonitor_test.go", + "gpu_registration_manager_test.go", "icms_client_test.go", "jwks_updater_test.go", "k8scomputebackend_miniservice_test.go", diff --git a/src/compute-plane-services/nvca/pkg/nvca/agent.go b/src/compute-plane-services/nvca/pkg/nvca/agent.go index a024ff618..a5b0560ce 100644 --- a/src/compute-plane-services/nvca/pkg/nvca/agent.go +++ b/src/compute-plane-services/nvca/pkg/nvca/agent.go @@ -24,7 +24,6 @@ import ( "net/http" "os" "strings" - "sync" "sync/atomic" "time" @@ -87,8 +86,6 @@ const ( ICMSRequestAckMaxGoroutines = 20 ICMSInstanceRequestStatusUpdatesMaxGoroutines = 20 - - gracefulNoGPURegistrationComponentName = "icmsregistration" ) // Controller tick types. @@ -385,15 +382,10 @@ type Agent struct { backendk8scache *BackendK8sCache backendHealthCache health.StatusCache - // gpuMonitor monitors GPU availability and controls queue processing - // when GracefulNoGPU feature flag is enabled. - gpuMonitor *GPUMonitor - gpuRegistrationMu sync.Mutex - gpuRegistrationReady atomic.Bool - gpuRegistrationGeneration atomic.Uint64 - gpuRegistrationRequests chan struct{} - // Serializes ICMS registration and credential refresh through queue credential application. - registrationOperationMu sync.Mutex + // gpuRegistration coordinates GPU availability, ICMS registration, and + // queue readiness when GracefulNoGPU is enabled. It also serializes all + // registration and credential-refresh operations. + gpuRegistration gpuRegistrationManager startControllerManager func(context.Context, *kubeclients.KubeClients) error @@ -420,109 +412,6 @@ type Agent struct { selfDestruct *atomic.Bool } -func (a *Agent) getGracefulNoGPURegistrationStatus(context.Context) (types.AgentHealth, error) { - component := types.ComponentHealth{ - Status: types.HealthStatusHealthy, - StatusLevel: types.StatusLevelError, - } - if !a.gpuRegistrationReady.Load() { - component.Status = types.HealthStatusUnhealthy - component.Errors = []string{"waiting for successful ICMS registration after GPU discovery"} - } - - return types.AgentHealth{ - Components: map[string]types.ComponentHealth{ - gracefulNoGPURegistrationComponentName: component, - }, - }, nil -} - -func (a *Agent) handleGracefulNoGPUStateChange(ctx context.Context, hasGPUs bool) { - log := core.GetLogger(ctx) - a.gpuRegistrationGeneration.Add(1) - - a.gpuRegistrationMu.Lock() - a.gpuRegistrationReady.Store(false) - a.queueManager.Pause() - a.gpuRegistrationMu.Unlock() - - if !hasGPUs { - log.Warn("GPUs no longer available - pausing queue manager") - return - } - - log.Info("GPUs detected - waiting for successful ICMS registration before resuming queue manager") - select { - case a.gpuRegistrationRequests <- struct{}{}: - default: - } -} - -func (a *Agent) runGracefulNoGPURegistration(ctx context.Context) { - retryInterval := a.GPUPollInterval - if retryInterval <= 0 { - retryInterval = DefaultGPUPollInterval - } - - for { - select { - case <-ctx.Done(): - return - case <-a.gpuRegistrationRequests: - } - - for !a.gpuRegistrationReady.Load() && a.gpuMonitor.HasGPUs() { - if !a.tryGracefulNoGPURegistration(ctx) { - break - } - - retryTimer := time.NewTimer(retryInterval) - select { - case <-ctx.Done(): - retryTimer.Stop() - return - case <-a.gpuRegistrationRequests: - retryTimer.Stop() - case <-retryTimer.C: - } - } - } -} - -// tryGracefulNoGPURegistration returns true when registration should be retried. -func (a *Agent) tryGracefulNoGPURegistration(ctx context.Context) bool { - log := core.GetLogger(ctx) - a.registrationOperationMu.Lock() - defer a.registrationOperationMu.Unlock() - if ctx.Err() != nil { - return false - } - - generation := a.gpuRegistrationGeneration.Load() - if !a.gpuMonitor.HasGPUs() { - return false - } - - log.Info("Registering with ICMS after GPUs became available") - if _, err := a.RegisterWithICMS(ctx); err != nil { - log.WithError(err).Warn("Failed to register with ICMS after GPUs became available; will retry") - return a.gpuMonitor.HasGPUs() - } - - a.gpuRegistrationMu.Lock() - defer a.gpuRegistrationMu.Unlock() - if !a.gpuMonitor.HasGPUs() || generation != a.gpuRegistrationGeneration.Load() { - log.Warn("GPU availability changed during ICMS registration - keeping queue manager paused") - return a.gpuMonitor.HasGPUs() - } - - // RegisterWithICMS installs the returned queue credentials before it returns. - a.gpuRegistrationReady.Store(true) - a.queueManager.Resume() - log.Info("Successfully registered with ICMS after GPUs became available") - return false -} - func (o *AgentOptions) sanitizedString() string { var sanitized AgentOptions sanitized.NCAId = o.NCAId @@ -1303,13 +1192,20 @@ func (a *Agent) Start(ctx context.Context) error { if nodeInformer := a.backendk8scache.GetNodeInformer(); nodeInformer != nil { gpuMonitorOpts = append(gpuMonitorOpts, WithNodeInformer(nodeInformer)) } - a.gpuMonitor = NewGPUMonitor(nfClient, gpuMonitorOpts...) + gpuMonitor := NewGPUMonitor(nfClient, gpuMonitorOpts...) // Check initial GPU availability gpus, gpuErr := nfClient.GetAllBackendGPUs(ctx) hasGPUs := gpuErr == nil && len(gpus) > 0 - a.gpuMonitor.SetHasGPUs(hasGPUs) - a.gpuRegistrationReady.Store(hasGPUs) - a.gpuRegistrationRequests = make(chan struct{}, 1) + gpuMonitor.SetHasGPUs(hasGPUs) + a.gpuRegistration.configureGracefulNoGPU( + gpuMonitor, + hasGPUs, + a.GPUPollInterval, + func(registrationCtx context.Context) error { + _, registrationErr := a.RegisterWithICMS(registrationCtx) + return registrationErr + }, + ) if hasGPUs { log.Info("GPUs found during startup, proceeding normally") } else { @@ -1322,10 +1218,10 @@ func (a *Agent) Start(ctx context.Context) error { statusUpdaters := []health.ComponentStatusGetter{a.backendk8scache} // Add GPU monitor to status updaters for readiness checks when GracefulNoGPU is enabled - if a.gpuMonitor != nil { - statusUpdaters = append(statusUpdaters, a.gpuMonitor) + if a.gpuRegistration.enabled() { + statusUpdaters = append(statusUpdaters, a.gpuRegistration.monitor) statusUpdaters = append(statusUpdaters, - health.GetComponentStatusFunc(a.getGracefulNoGPURegistrationStatus)) + health.GetComponentStatusFunc(a.gpuRegistration.getRegistrationStatus)) } if a.FeatureFlagFetcher.IsAttributeEnabled(featureflag.AttrHostIsolation) { statusUpdaters = append(statusUpdaters, hostisolation.NewStatusGetter( @@ -1353,7 +1249,7 @@ func (a *Agent) Start(ctx context.Context) error { // Skip waiting for healthy status when GracefulNoGPU is enabled with no GPUs. // In this case, readiness will report not-ready until GPUs appear, but liveness will pass. - skipHealthWait := a.gpuMonitor != nil && !a.gpuMonitor.HasGPUs() + skipHealthWait := a.gpuRegistration.enabled() && !a.gpuRegistration.hasGPUs() if skipHealthWait { log.Warn("GracefulNoGPU enabled with no GPUs - skipping health wait, readiness will report not-ready") @@ -1372,7 +1268,7 @@ func (a *Agent) Start(ctx context.Context) error { } // Check if we should defer ICMS registration (no GPUs with GracefulNoGPU enabled). - skipInitialRegistration := a.gpuMonitor != nil && !a.gpuMonitor.HasGPUs() + skipInitialRegistration := a.gpuRegistration.enabled() && !a.gpuRegistration.hasGPUs() var res *types.ICMSRegistrationResponse if skipInitialRegistration { log.Warn("No GPUs available - deferring ICMS registration until GPUs are detected") @@ -1518,20 +1414,15 @@ func (a *Agent) Start(ctx context.Context) error { // If GracefulNoGPU is enabled and we started without GPUs, pause the queue manager // and set up callbacks to handle GPU availability changes - if a.gpuMonitor != nil { - if !a.gpuMonitor.HasGPUs() { + if a.gpuRegistration.enabled() { + a.gpuRegistration.setQueueManager(a.queueManager) + if !a.gpuRegistration.hasGPUs() { log.Warn("Starting with queue manager paused due to no GPUs") a.queueManager.Pause() } - go a.runGracefulNoGPURegistration(ctx) - - // Set up GPU state change callback - a.gpuMonitor.SetOnGPUStateChange(a.handleGracefulNoGPUStateChange) - - // Start the GPU monitor polling loop log.Info("Starting GPU monitor") - a.gpuMonitor.Start(ctx) + a.gpuRegistration.start(ctx) } // Evict all workloads once during startup if in CordonAndDrainMaintenance mode @@ -2262,30 +2153,26 @@ func (a *Agent) RenewICMSQueueCreds(ctx context.Context) error { return nil } - a.registrationOperationMu.Lock() - defer a.registrationOperationMu.Unlock() - if err := ctx.Err(); err != nil { - return err - } - - credRes, err := a.icmsClient.GetCreds(ctx) - nvcametrics.FromContext(ctx).RecordUpstreamRequest(nvcametrics.UpstreamOperationCredentials, err) - if err != nil { - return fmt.Errorf("failed to GetCreds from ICMS, err: %v", err) - } + return a.gpuRegistration.withRegistrationOperation(ctx, func() error { + credRes, err := a.icmsClient.GetCreds(ctx) + nvcametrics.FromContext(ctx).RecordUpstreamRequest(nvcametrics.UpstreamOperationCredentials, err) + if err != nil { + return fmt.Errorf("failed to GetCreds from ICMS, err: %v", err) + } - // TODO: this is a hack remove this once ICMS properly sends back the queue credentials - queueCreds := a.postProcessQueueCredentials(ctx, credRes.QueueCredentials) + // TODO: this is a hack remove this once ICMS properly sends back the queue credentials + queueCreds := a.postProcessQueueCredentials(ctx, credRes.QueueCredentials) - err = a.backendk8scache.StoreUpdatedCredentials(ctx, queueCreds) - if err != nil { - return fmt.Errorf("failed to store renewed Queue Credentials, err: %v", err) - } + err = a.backendk8scache.StoreUpdatedCredentials(ctx, queueCreds) + if err != nil { + return fmt.Errorf("failed to store renewed Queue Credentials, err: %v", err) + } - a.queueManager.updateQueues(queueCreds) + a.queueManager.updateQueues(queueCreds) - log.Debugf("refreshed queueManager with new Creds") - return nil + log.Debugf("refreshed queueManager with new Creds") + return nil + }) } // evictAllWorkloads directly purges all workload instances and sends termination status updates to ICMS. diff --git a/src/compute-plane-services/nvca/pkg/nvca/agent_test.go b/src/compute-plane-services/nvca/pkg/nvca/agent_test.go index bcfc3deed..b8c6d7935 100644 --- a/src/compute-plane-services/nvca/pkg/nvca/agent_test.go +++ b/src/compute-plane-services/nvca/pkg/nvca/agent_test.go @@ -453,9 +453,9 @@ func TestAgentStartGracefulNoGPURecoversWhenGPUAppears(t *testing.T) { require.NoError(t, agent.Start(ctx)) requireHTTPStatus(t, agent.NVCASvcAddress, health.HTTPReadinessRoutePath, http.StatusServiceUnavailable) requireHTTPStatus(t, agent.NVCASvcAddress, health.HTTPLivenessRoutePath, http.StatusOK) - require.NotNil(t, agent.gpuMonitor) + require.NotNil(t, agent.gpuRegistration.monitor) require.NotNil(t, agent.queueManager) - assert.False(t, agent.gpuMonitor.HasGPUs()) + assert.False(t, agent.gpuRegistration.hasGPUs()) assert.True(t, agent.queueManager.IsPaused()) assert.Empty(t, icmsClient.requests()) assert.Empty(t, agent.queueManager.getCreateQueue(testGPUNameDefault).QueueURL) @@ -494,7 +494,7 @@ func TestAgentStartGracefulNoGPURecoversWhenGPUAppears(t *testing.T) { icmsClient.release(0) require.EventuallyWithT(t, func(ct *assert.CollectT) { - assert.True(ct, agent.gpuMonitor.HasGPUs()) + assert.True(ct, agent.gpuRegistration.hasGPUs()) assert.False(ct, agent.queueManager.IsPaused()) assert.Equal(ct, recoveredCredentials.CreationQueues[testGPUNameDefault], agent.queueManager.getCreateQueue(testGPUNameDefault)) @@ -505,7 +505,7 @@ func TestAgentStartGracefulNoGPURecoversWhenGPUAppears(t *testing.T) { require.NoError(t, k8sClients.K8s.CoreV1().Nodes().Delete(ctx, functionNode.Name, metav1.DeleteOptions{})) require.EventuallyWithT(t, func(ct *assert.CollectT) { - assert.False(ct, agent.gpuMonitor.HasGPUs()) + assert.False(ct, agent.gpuRegistration.hasGPUs()) assert.True(ct, agent.queueManager.IsPaused()) }, 5*time.Second, 10*time.Millisecond) requireRefreshedReadinessStatus(t, ctx, agent, http.StatusServiceUnavailable) @@ -589,7 +589,7 @@ func TestAgentStartGracefulNoGPUStaysPausedWhenGPUDisappearsDuringRegistration(t requireRegistrationAttempt(t, icmsClient, 0) require.NoError(t, k8sClients.K8s.CoreV1().Nodes().Delete(ctx, functionNode.Name, metav1.DeleteOptions{})) require.EventuallyWithT(t, func(ct *assert.CollectT) { - assert.False(ct, agent.gpuMonitor.HasGPUs()) + assert.False(ct, agent.gpuRegistration.hasGPUs()) assert.True(ct, agent.queueManager.IsPaused()) }, 5*time.Second, 10*time.Millisecond) @@ -642,7 +642,7 @@ func TestAgentStartGracefulNoGPUSerializesCredentialRenewalBeforeRecovery(t *tes require.NoError(t, k8sClients.K8s.CoreV1().Nodes().Delete(ctx, functionNode.Name, metav1.DeleteOptions{})) require.EventuallyWithT(t, func(ct *assert.CollectT) { - assert.False(ct, agent.gpuMonitor.HasGPUs()) + assert.False(ct, agent.gpuRegistration.hasGPUs()) assert.True(ct, agent.queueManager.IsPaused()) }, 5*time.Second, 10*time.Millisecond) @@ -733,7 +733,7 @@ func TestAgentStartGracefulNoGPUSerializesPeriodicRegistrationBeforeRecovery(t * require.NoError(t, k8sClients.K8s.CoreV1().Nodes().Delete(ctx, functionNode.Name, metav1.DeleteOptions{})) require.EventuallyWithT(t, func(ct *assert.CollectT) { - assert.False(ct, agent.gpuMonitor.HasGPUs()) + assert.False(ct, agent.gpuRegistration.hasGPUs()) assert.True(ct, agent.queueManager.IsPaused()) }, 5*time.Second, 10*time.Millisecond) @@ -744,12 +744,12 @@ func TestAgentStartGracefulNoGPUSerializesPeriodicRegistrationBeforeRecovery(t * recoveryNode.Labels["nvidia.com/gpu.memory"] = "32768" recoveryNode.Labels["nvidia.com/gpu.product"] = "V100-SXM2-32GB" recoveryNode.Labels["nvca.nvcf.nvidia.io/gpu.product"] = string(recoveredGPU) - lossGeneration := agent.gpuRegistrationGeneration.Load() + lossGeneration := agent.gpuRegistration.generation.Load() _, err = k8sClients.K8s.CoreV1().Nodes().Create(ctx, recoveryNode, metav1.CreateOptions{}) require.NoError(t, err) require.EventuallyWithT(t, func(ct *assert.CollectT) { - assert.True(ct, agent.gpuMonitor.HasGPUs()) - assert.Greater(ct, agent.gpuRegistrationGeneration.Load(), lossGeneration) + assert.True(ct, agent.gpuRegistration.hasGPUs()) + assert.Greater(ct, agent.gpuRegistration.generation.Load(), lossGeneration) }, 5*time.Second, 10*time.Millisecond) requireNoRegistrationAttempt(t, icmsClient, 250*time.Millisecond) assert.True(t, agent.queueManager.IsPaused()) diff --git a/src/compute-plane-services/nvca/pkg/nvca/agent_updates.go b/src/compute-plane-services/nvca/pkg/nvca/agent_updates.go index 1b7977dad..10da29b3e 100644 --- a/src/compute-plane-services/nvca/pkg/nvca/agent_updates.go +++ b/src/compute-plane-services/nvca/pkg/nvca/agent_updates.go @@ -74,27 +74,23 @@ func (a *Agent) initICMSRegistrationSyncer(ctx context.Context) error { var lastRegTime = core.GetCurrentTime(ctx) a.syncICMSRegistration = func(ctx context.Context) error { - a.registrationOperationMu.Lock() - defer a.registrationOperationMu.Unlock() - if err := ctx.Err(); err != nil { - return err - } + return a.gpuRegistration.withRegistrationOperation(ctx, func() error { + if timeSinceLastUpdate := core.GetCurrentTime(ctx).Sub(lastRegTime); timeSinceLastUpdate < time.Hour { + return nil + } - if timeSinceLastUpdate := core.GetCurrentTime(ctx).Sub(lastRegTime); timeSinceLastUpdate < time.Hour { - return nil - } + regBackendGPUs, err := a.getRegistrationGPUs(ctx) + if err != nil { + return err + } - regBackendGPUs, err := a.getRegistrationGPUs(ctx) - if err != nil { - return err - } + if err := a.doRegistrationUpdate(ctx, regBackendGPUs, "Performing hourly ICMS registration refresh"); err != nil { + return err + } - if err := a.doRegistrationUpdate(ctx, regBackendGPUs, "Performing hourly ICMS registration refresh"); err != nil { - return err - } - - lastRegTime = core.GetCurrentTime(ctx) - return nil + lastRegTime = core.GetCurrentTime(ctx) + return nil + }) } return nil } @@ -104,41 +100,37 @@ func (a *Agent) initICMSRegistrationSyncer(ctx context.Context) error { var lastRegTime = core.GetCurrentTime(ctx) a.syncICMSRegistration = func(ctx context.Context) error { - a.registrationOperationMu.Lock() - defer a.registrationOperationMu.Unlock() - if err := ctx.Err(); err != nil { - return err - } - - regBackendGPUs, err := a.getRegistrationGPUs(ctx) - if err != nil { - return err - } + return a.gpuRegistration.withRegistrationOperation(ctx, func() error { + regBackendGPUs, err := a.getRegistrationGPUs(ctx) + if err != nil { + return err + } - // Determine if we need to register - gpusChanged := len(lastRegBackendGPUs) == 0 || !cmp.Equal(lastRegBackendGPUs, regBackendGPUs, cmpopts.EquateEmpty()) - timeForHourlyRefresh := core.GetCurrentTime(ctx).Sub(lastRegTime) >= time.Hour + // Determine if we need to register + gpusChanged := len(lastRegBackendGPUs) == 0 || !cmp.Equal(lastRegBackendGPUs, regBackendGPUs, cmpopts.EquateEmpty()) + timeForHourlyRefresh := core.GetCurrentTime(ctx).Sub(lastRegTime) >= time.Hour - if !gpusChanged && !timeForHourlyRefresh { - log.Debug("Backend GPUs are up to date") - return nil - } + if !gpusChanged && !timeForHourlyRefresh { + log.Debug("Backend GPUs are up to date") + return nil + } - // Register with appropriate reason - var reason string - if timeForHourlyRefresh { - reason = "Performing hourly ICMS registration refresh" - } else { - reason = "Registering with ICMS due to GPU changes" - } + // Register with appropriate reason + var reason string + if timeForHourlyRefresh { + reason = "Performing hourly ICMS registration refresh" + } else { + reason = "Registering with ICMS due to GPU changes" + } - if err := a.doRegistrationUpdate(ctx, regBackendGPUs, reason); err != nil { - return err - } + if err := a.doRegistrationUpdate(ctx, regBackendGPUs, reason); err != nil { + return err + } - lastRegBackendGPUs = regBackendGPUs - lastRegTime = core.GetCurrentTime(ctx) - return nil + lastRegBackendGPUs = regBackendGPUs + lastRegTime = core.GetCurrentTime(ctx) + return nil + }) } return nil diff --git a/src/compute-plane-services/nvca/pkg/nvca/gpu_registration_manager.go b/src/compute-plane-services/nvca/pkg/nvca/gpu_registration_manager.go new file mode 100644 index 000000000..d7b7befc6 --- /dev/null +++ b/src/compute-plane-services/nvca/pkg/nvca/gpu_registration_manager.go @@ -0,0 +1,231 @@ +/* +SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +SPDX-License-Identifier: Apache-2.0 + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package nvca + +import ( + "context" + "sync" + "sync/atomic" + "time" + + "github.com/NVIDIA/nvcf/src/libraries/go/lib/pkg/core" + + "github.com/NVIDIA/nvcf/src/compute-plane-services/nvca/pkg/types" +) + +const gracefulNoGPURegistrationComponentName = "icmsregistration" + +type gpuRegistrationMonitor interface { + HasGPUs() bool + SetOnGPUStateChange(GPUStateChangeCallback) + Start(context.Context) + GetComponentStatus(context.Context) (types.AgentHealth, error) +} + +type gpuRegistrationQueue interface { + Pause() + Resume() +} + +type gpuRegistrationManager struct { + operationGate contextAwareRegistrationGate + stateMu sync.Mutex + monitor gpuRegistrationMonitor + queueManager gpuRegistrationQueue + ready atomic.Bool + generation atomic.Uint64 + registrationRequests chan struct{} + retryInterval time.Duration + register func(context.Context) error +} + +type contextAwareRegistrationGate struct { + once sync.Once + token chan struct{} +} + +func (g *contextAwareRegistrationGate) lock(ctx context.Context) error { + g.once.Do(func() { + g.token = make(chan struct{}, 1) + g.token <- struct{}{} + }) + + if err := ctx.Err(); err != nil { + return err + } + + select { + case <-ctx.Done(): + return ctx.Err() + case <-g.token: + if err := ctx.Err(); err != nil { + g.token <- struct{}{} + return err + } + return nil + } +} + +func (g *contextAwareRegistrationGate) unlock() { + g.token <- struct{}{} +} + +func (m *gpuRegistrationManager) withRegistrationOperation(ctx context.Context, operation func() error) error { + if err := m.operationGate.lock(ctx); err != nil { + return err + } + defer m.operationGate.unlock() + return operation() +} + +func (m *gpuRegistrationManager) configureGracefulNoGPU( + monitor gpuRegistrationMonitor, + initiallyReady bool, + retryInterval time.Duration, + register func(context.Context) error, +) { + m.monitor = monitor + m.ready.Store(initiallyReady) + m.retryInterval = retryInterval + m.register = register + m.registrationRequests = make(chan struct{}, 1) +} + +func (m *gpuRegistrationManager) setQueueManager(queueManager gpuRegistrationQueue) { + m.queueManager = queueManager +} + +func (m *gpuRegistrationManager) start(ctx context.Context) { + go m.run(ctx) + m.monitor.SetOnGPUStateChange(m.handleGPUStateChange) + m.monitor.Start(ctx) +} + +func (m *gpuRegistrationManager) enabled() bool { + return m.monitor != nil +} + +func (m *gpuRegistrationManager) hasGPUs() bool { + return m.monitor != nil && m.monitor.HasGPUs() +} + +func (m *gpuRegistrationManager) getRegistrationStatus(context.Context) (types.AgentHealth, error) { + component := types.ComponentHealth{ + Status: types.HealthStatusHealthy, + StatusLevel: types.StatusLevelError, + } + if !m.ready.Load() { + component.Status = types.HealthStatusUnhealthy + component.Errors = []string{"waiting for successful ICMS registration after GPU discovery"} + } + + return types.AgentHealth{ + Components: map[string]types.ComponentHealth{ + gracefulNoGPURegistrationComponentName: component, + }, + }, nil +} + +func (m *gpuRegistrationManager) handleGPUStateChange(ctx context.Context, hasGPUs bool) { + log := core.GetLogger(ctx) + m.generation.Add(1) + + m.stateMu.Lock() + m.ready.Store(false) + if m.queueManager != nil { + m.queueManager.Pause() + } + m.stateMu.Unlock() + + if !hasGPUs { + log.Warn("GPUs no longer available - pausing queue manager") + return + } + + log.Info("GPUs detected - waiting for successful ICMS registration before resuming queue manager") + select { + case m.registrationRequests <- struct{}{}: + default: + } +} + +func (m *gpuRegistrationManager) run(ctx context.Context) { + retryInterval := m.retryInterval + if retryInterval <= 0 { + retryInterval = DefaultGPUPollInterval + } + + for { + select { + case <-ctx.Done(): + return + case <-m.registrationRequests: + } + + for !m.ready.Load() && m.hasGPUs() { + if !m.tryRegistration(ctx) { + break + } + + retryTimer := time.NewTimer(retryInterval) + select { + case <-ctx.Done(): + retryTimer.Stop() + return + case <-m.registrationRequests: + retryTimer.Stop() + case <-retryTimer.C: + } + } + } +} + +// tryRegistration returns true when registration should be retried. +func (m *gpuRegistrationManager) tryRegistration(ctx context.Context) bool { + log := core.GetLogger(ctx) + shouldRetry := false + err := m.withRegistrationOperation(ctx, func() error { + generation := m.generation.Load() + if m.monitor == nil || !m.monitor.HasGPUs() { + return nil + } + + log.Info("Registering with ICMS after GPUs became available") + if err := m.register(ctx); err != nil { + log.WithError(err).Warn("Failed to register with ICMS after GPUs became available; will retry") + shouldRetry = m.monitor.HasGPUs() + return nil + } + + m.stateMu.Lock() + defer m.stateMu.Unlock() + if !m.monitor.HasGPUs() || generation != m.generation.Load() { + log.Warn("GPU availability changed during ICMS registration - keeping queue manager paused") + shouldRetry = m.monitor.HasGPUs() + return nil + } + + m.ready.Store(true) + if m.queueManager != nil { + m.queueManager.Resume() + } + log.Info("Successfully registered with ICMS after GPUs became available") + return nil + }) + return err == nil && shouldRetry +} diff --git a/src/compute-plane-services/nvca/pkg/nvca/gpu_registration_manager_test.go b/src/compute-plane-services/nvca/pkg/nvca/gpu_registration_manager_test.go new file mode 100644 index 000000000..8f7f16088 --- /dev/null +++ b/src/compute-plane-services/nvca/pkg/nvca/gpu_registration_manager_test.go @@ -0,0 +1,153 @@ +/* +SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +SPDX-License-Identifier: Apache-2.0 + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package nvca + +import ( + "context" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/NVIDIA/nvcf/src/compute-plane-services/nvca/pkg/types" +) + +type fakeGPUAvailability struct { + hasGPUs atomic.Bool + onStateChange GPUStateChangeCallback +} + +func (f *fakeGPUAvailability) HasGPUs() bool { + return f.hasGPUs.Load() +} + +func (f *fakeGPUAvailability) SetOnGPUStateChange(callback GPUStateChangeCallback) { + f.onStateChange = callback +} + +func (f *fakeGPUAvailability) Start(context.Context) {} + +func (f *fakeGPUAvailability) GetComponentStatus(context.Context) (types.AgentHealth, error) { + return types.AgentHealth{}, nil +} + +type fakeGPURegistrationQueue struct { + paused atomic.Bool +} + +func (f *fakeGPURegistrationQueue) Pause() { + f.paused.Store(true) +} + +func (f *fakeGPURegistrationQueue) Resume() { + f.paused.Store(false) +} + +func TestGPURegistrationManagerCancellationWhileWaiting(t *testing.T) { + var manager gpuRegistrationManager + + holderEntered := make(chan struct{}) + releaseHolder := make(chan struct{}) + holderDone := make(chan error, 1) + go func() { + holderDone <- manager.withRegistrationOperation(context.Background(), func() error { + close(holderEntered) + <-releaseHolder + return nil + }) + }() + <-holderEntered + + ctx, cancel := context.WithCancel(context.Background()) + waiterStarted := make(chan struct{}) + waiterDone := make(chan error, 1) + var waiterRan atomic.Bool + go func() { + close(waiterStarted) + waiterDone <- manager.withRegistrationOperation(ctx, func() error { + waiterRan.Store(true) + return nil + }) + }() + <-waiterStarted + cancel() + + select { + case err := <-waiterDone: + require.ErrorIs(t, err, context.Canceled) + assert.False(t, waiterRan.Load()) + case <-time.After(time.Second): + t.Fatal("canceled registration operation remained blocked behind the active operation") + } + + close(releaseHolder) + require.NoError(t, <-holderDone) +} + +func TestGPURegistrationManagerGPULossMarksUnreadyAndPausesQueue(t *testing.T) { + monitor := &fakeGPUAvailability{} + monitor.hasGPUs.Store(true) + queue := &fakeGPURegistrationQueue{} + var manager gpuRegistrationManager + manager.configureGracefulNoGPU(monitor, true, time.Second, nil) + manager.setQueueManager(queue) + + manager.handleGPUStateChange(context.Background(), false) + + status, err := manager.getRegistrationStatus(context.Background()) + require.NoError(t, err) + require.Contains(t, status.Components, gracefulNoGPURegistrationComponentName) + assert.Equal(t, types.HealthStatusUnhealthy, + status.Components[gracefulNoGPURegistrationComponentName].Status) + assert.True(t, queue.paused.Load()) +} + +func TestGPURegistrationManagerGPUChangeDuringRegistrationKeepsQueuePaused(t *testing.T) { + monitor := &fakeGPUAvailability{} + monitor.hasGPUs.Store(true) + queue := &fakeGPURegistrationQueue{} + queue.Pause() + registrationStarted := make(chan struct{}) + releaseRegistration := make(chan struct{}) + var manager gpuRegistrationManager + manager.configureGracefulNoGPU(monitor, false, time.Second, func(context.Context) error { + close(registrationStarted) + <-releaseRegistration + return nil + }) + manager.setQueueManager(queue) + + registrationDone := make(chan bool, 1) + go func() { + registrationDone <- manager.tryRegistration(context.Background()) + }() + <-registrationStarted + + monitor.hasGPUs.Store(false) + manager.handleGPUStateChange(context.Background(), false) + close(releaseRegistration) + + assert.False(t, <-registrationDone) + status, err := manager.getRegistrationStatus(context.Background()) + require.NoError(t, err) + assert.Equal(t, types.HealthStatusUnhealthy, + status.Components[gracefulNoGPURegistrationComponentName].Status) + assert.True(t, queue.paused.Load()) +} From 798145618ccf29fba593906a838170e9c432901d Mon Sep 17 00:00:00 2001 From: Mike Camp Date: Mon, 31 Aug 2026 21:48:55 +0200 Subject: [PATCH 6/7] fix(nvca): bound GPU registration retries Signed-off-by: Mike Camp --- .../nvca/pkg/nvca/gpu_registration_manager.go | 67 ++++++++++++++++++- .../pkg/nvca/gpu_registration_manager_test.go | 36 ++++++++-- 2 files changed, 98 insertions(+), 5 deletions(-) diff --git a/src/compute-plane-services/nvca/pkg/nvca/gpu_registration_manager.go b/src/compute-plane-services/nvca/pkg/nvca/gpu_registration_manager.go index d7b7befc6..f42105cce 100644 --- a/src/compute-plane-services/nvca/pkg/nvca/gpu_registration_manager.go +++ b/src/compute-plane-services/nvca/pkg/nvca/gpu_registration_manager.go @@ -29,6 +29,7 @@ import ( ) const gracefulNoGPURegistrationComponentName = "icmsregistration" +const maxGracefulNoGPURegistrationRetryInterval = 5 * time.Minute type gpuRegistrationMonitor interface { HasGPUs() bool @@ -42,6 +43,8 @@ type gpuRegistrationQueue interface { Resume() } +// gpuRegistrationManager owns GPU-registration readiness, retry, queue, and +// serialization state so all ICMS registration operations share one lifecycle. type gpuRegistrationManager struct { operationGate contextAwareRegistrationGate stateMu sync.Mutex @@ -54,11 +57,53 @@ type gpuRegistrationManager struct { register func(context.Context) error } +// contextAwareRegistrationGate serializes registration operations while +// allowing a queued caller to leave promptly when its context is canceled. type contextAwareRegistrationGate struct { once sync.Once token chan struct{} } +// gpuRegistrationRetryBackoff bounds repeated registration attempts while +// retaining the configured poll interval as the first retry delay. +type gpuRegistrationRetryBackoff struct { + initial time.Duration + current time.Duration + maximum time.Duration +} + +// newGPURegistrationRetryBackoff preserves intervals larger than the default +// ceiling so an operator-configured slow retry cadence is never shortened. +func newGPURegistrationRetryBackoff(initial, maximum time.Duration) gpuRegistrationRetryBackoff { + if maximum < initial { + maximum = initial + } + return gpuRegistrationRetryBackoff{ + initial: initial, + current: initial, + maximum: maximum, + } +} + +// next returns the current delay and advances the following delay to its cap. +func (b *gpuRegistrationRetryBackoff) next() time.Duration { + delay := b.current + if b.current < b.maximum { + if b.current > b.maximum/2 { + b.current = b.maximum + } else { + b.current *= 2 + } + } + return delay +} + +// reset restores the configured first-retry delay after a GPU state change. +func (b *gpuRegistrationRetryBackoff) reset() { + b.current = b.initial +} + +// lock waits for the single registration token or returns when ctx is canceled. func (g *contextAwareRegistrationGate) lock(ctx context.Context) error { g.once.Do(func() { g.token = make(chan struct{}, 1) @@ -81,10 +126,13 @@ func (g *contextAwareRegistrationGate) lock(ctx context.Context) error { } } +// unlock releases the registration token after a successful lock. func (g *contextAwareRegistrationGate) unlock() { g.token <- struct{}{} } +// withRegistrationOperation protects recovery, renewal, and periodic +// registration from overlapping while honoring cancellation before entry. func (m *gpuRegistrationManager) withRegistrationOperation(ctx context.Context, operation func() error) error { if err := m.operationGate.lock(ctx); err != nil { return err @@ -93,6 +141,8 @@ func (m *gpuRegistrationManager) withRegistrationOperation(ctx context.Context, return operation() } +// configureGracefulNoGPU initializes state before the monitor and registration +// worker start. func (m *gpuRegistrationManager) configureGracefulNoGPU( monitor gpuRegistrationMonitor, initiallyReady bool, @@ -106,24 +156,29 @@ func (m *gpuRegistrationManager) configureGracefulNoGPU( m.registrationRequests = make(chan struct{}, 1) } +// setQueueManager connects queue flow control once startup constructs it. func (m *gpuRegistrationManager) setQueueManager(queueManager gpuRegistrationQueue) { m.queueManager = queueManager } +// start launches recovery before the monitor can report its first transition. func (m *gpuRegistrationManager) start(ctx context.Context) { go m.run(ctx) m.monitor.SetOnGPUStateChange(m.handleGPUStateChange) m.monitor.Start(ctx) } +// enabled reports whether graceful no-GPU registration was configured. func (m *gpuRegistrationManager) enabled() bool { return m.monitor != nil } +// hasGPUs is nil-safe for agents that do not enable GPU monitoring. func (m *gpuRegistrationManager) hasGPUs() bool { return m.monitor != nil && m.monitor.HasGPUs() } +// getRegistrationStatus contributes recovery readiness to aggregate health. func (m *gpuRegistrationManager) getRegistrationStatus(context.Context) (types.AgentHealth, error) { component := types.ComponentHealth{ Status: types.HealthStatusHealthy, @@ -141,6 +196,8 @@ func (m *gpuRegistrationManager) getRegistrationStatus(context.Context) (types.A }, nil } +// handleGPUStateChange immediately marks registration unready and pauses queues; +// GPU arrival then schedules serialized recovery registration. func (m *gpuRegistrationManager) handleGPUStateChange(ctx context.Context, hasGPUs bool) { log := core.GetLogger(ctx) m.generation.Add(1) @@ -164,17 +221,24 @@ func (m *gpuRegistrationManager) handleGPUStateChange(ctx context.Context, hasGP } } +// run consumes coalesced GPU-arrival requests and retries registration with +// bounded backoff until readiness succeeds, GPUs disappear, or ctx ends. func (m *gpuRegistrationManager) run(ctx context.Context) { retryInterval := m.retryInterval if retryInterval <= 0 { retryInterval = DefaultGPUPollInterval } + backoff := newGPURegistrationRetryBackoff( + retryInterval, + maxGracefulNoGPURegistrationRetryInterval, + ) for { select { case <-ctx.Done(): return case <-m.registrationRequests: + backoff.reset() } for !m.ready.Load() && m.hasGPUs() { @@ -182,13 +246,14 @@ func (m *gpuRegistrationManager) run(ctx context.Context) { break } - retryTimer := time.NewTimer(retryInterval) + retryTimer := time.NewTimer(backoff.next()) select { case <-ctx.Done(): retryTimer.Stop() return case <-m.registrationRequests: retryTimer.Stop() + backoff.reset() case <-retryTimer.C: } } diff --git a/src/compute-plane-services/nvca/pkg/nvca/gpu_registration_manager_test.go b/src/compute-plane-services/nvca/pkg/nvca/gpu_registration_manager_test.go index 8f7f16088..cb2db7d7d 100644 --- a/src/compute-plane-services/nvca/pkg/nvca/gpu_registration_manager_test.go +++ b/src/compute-plane-services/nvca/pkg/nvca/gpu_registration_manager_test.go @@ -19,6 +19,7 @@ package nvca import ( "context" + "sync" "sync/atomic" "testing" "time" @@ -60,6 +61,19 @@ func (f *fakeGPURegistrationQueue) Resume() { f.paused.Store(false) } +type observedDoneContext struct { + context.Context + doneObserved chan struct{} + once sync.Once +} + +func (c *observedDoneContext) Done() <-chan struct{} { + c.once.Do(func() { + close(c.doneObserved) + }) + return c.Context.Done() +} + func TestGPURegistrationManagerCancellationWhileWaiting(t *testing.T) { var manager gpuRegistrationManager @@ -75,18 +89,20 @@ func TestGPURegistrationManagerCancellationWhileWaiting(t *testing.T) { }() <-holderEntered - ctx, cancel := context.WithCancel(context.Background()) - waiterStarted := make(chan struct{}) + baseContext, cancel := context.WithCancel(context.Background()) + ctx := &observedDoneContext{ + Context: baseContext, + doneObserved: make(chan struct{}), + } waiterDone := make(chan error, 1) var waiterRan atomic.Bool go func() { - close(waiterStarted) waiterDone <- manager.withRegistrationOperation(ctx, func() error { waiterRan.Store(true) return nil }) }() - <-waiterStarted + <-ctx.doneObserved cancel() select { @@ -101,6 +117,18 @@ func TestGPURegistrationManagerCancellationWhileWaiting(t *testing.T) { require.NoError(t, <-holderDone) } +func TestGPURegistrationRetryBackoffIsBoundedAndResettable(t *testing.T) { + backoff := newGPURegistrationRetryBackoff(5*time.Second, 20*time.Second) + + assert.Equal(t, 5*time.Second, backoff.next()) + assert.Equal(t, 10*time.Second, backoff.next()) + assert.Equal(t, 20*time.Second, backoff.next()) + assert.Equal(t, 20*time.Second, backoff.next()) + + backoff.reset() + assert.Equal(t, 5*time.Second, backoff.next()) +} + func TestGPURegistrationManagerGPULossMarksUnreadyAndPausesQueue(t *testing.T) { monitor := &fakeGPUAvailability{} monitor.hasGPUs.Store(true) From d5e4245b8f3104d7bddecf88761190aa4ddbae24 Mon Sep 17 00:00:00 2001 From: Mike Camp Date: Mon, 31 Aug 2026 22:00:47 +0200 Subject: [PATCH 7/7] test(nvca): verify GPU retry scheduling Signed-off-by: Mike Camp --- .../nvca/pkg/nvca/gpu_registration_manager.go | 30 ++- .../pkg/nvca/gpu_registration_manager_test.go | 243 ++++++++++++++++++ 2 files changed, 271 insertions(+), 2 deletions(-) diff --git a/src/compute-plane-services/nvca/pkg/nvca/gpu_registration_manager.go b/src/compute-plane-services/nvca/pkg/nvca/gpu_registration_manager.go index f42105cce..d73e9e40b 100644 --- a/src/compute-plane-services/nvca/pkg/nvca/gpu_registration_manager.go +++ b/src/compute-plane-services/nvca/pkg/nvca/gpu_registration_manager.go @@ -55,6 +55,7 @@ type gpuRegistrationManager struct { registrationRequests chan struct{} retryInterval time.Duration register func(context.Context) error + newRetryTimer func(time.Duration) gpuRegistrationRetryTimer } // contextAwareRegistrationGate serializes registration operations while @@ -64,6 +65,25 @@ type contextAwareRegistrationGate struct { token chan struct{} } +// gpuRegistrationRetryTimer is the minimal timer surface used by the retry +// loop, allowing deterministic manager-level tests without changing behavior. +type gpuRegistrationRetryTimer interface { + C() <-chan time.Time + Stop() bool +} + +type realGPURegistrationRetryTimer struct { + timer *time.Timer +} + +func (t *realGPURegistrationRetryTimer) C() <-chan time.Time { + return t.timer.C +} + +func (t *realGPURegistrationRetryTimer) Stop() bool { + return t.timer.Stop() +} + // gpuRegistrationRetryBackoff bounds repeated registration attempts while // retaining the configured poll interval as the first retry delay. type gpuRegistrationRetryBackoff struct { @@ -232,6 +252,12 @@ func (m *gpuRegistrationManager) run(ctx context.Context) { retryInterval, maxGracefulNoGPURegistrationRetryInterval, ) + newRetryTimer := m.newRetryTimer + if newRetryTimer == nil { + newRetryTimer = func(delay time.Duration) gpuRegistrationRetryTimer { + return &realGPURegistrationRetryTimer{timer: time.NewTimer(delay)} + } + } for { select { @@ -246,7 +272,7 @@ func (m *gpuRegistrationManager) run(ctx context.Context) { break } - retryTimer := time.NewTimer(backoff.next()) + retryTimer := newRetryTimer(backoff.next()) select { case <-ctx.Done(): retryTimer.Stop() @@ -254,7 +280,7 @@ func (m *gpuRegistrationManager) run(ctx context.Context) { case <-m.registrationRequests: retryTimer.Stop() backoff.reset() - case <-retryTimer.C: + case <-retryTimer.C(): } } } diff --git a/src/compute-plane-services/nvca/pkg/nvca/gpu_registration_manager_test.go b/src/compute-plane-services/nvca/pkg/nvca/gpu_registration_manager_test.go index cb2db7d7d..dbee64cb6 100644 --- a/src/compute-plane-services/nvca/pkg/nvca/gpu_registration_manager_test.go +++ b/src/compute-plane-services/nvca/pkg/nvca/gpu_registration_manager_test.go @@ -19,6 +19,7 @@ package nvca import ( "context" + "errors" "sync" "sync/atomic" "testing" @@ -74,6 +75,75 @@ func (c *observedDoneContext) Done() <-chan struct{} { return c.Context.Done() } +type fakeGPURegistrationRetryTimer struct { + ch chan time.Time + stopped chan struct{} + once sync.Once +} + +func newFakeGPURegistrationRetryTimer() *fakeGPURegistrationRetryTimer { + return &fakeGPURegistrationRetryTimer{ + ch: make(chan time.Time, 1), + stopped: make(chan struct{}), + } +} + +func (t *fakeGPURegistrationRetryTimer) C() <-chan time.Time { + return t.ch +} + +func (t *fakeGPURegistrationRetryTimer) Stop() bool { + t.once.Do(func() { + close(t.stopped) + }) + return true +} + +func (t *fakeGPURegistrationRetryTimer) fire() { + t.ch <- time.Time{} +} + +type gpuRegistrationRetryTimerRequest struct { + delay time.Duration + timer *fakeGPURegistrationRetryTimer +} + +type fakeGPURegistrationRetryTimerFactory struct { + requests chan gpuRegistrationRetryTimerRequest +} + +func newFakeGPURegistrationRetryTimerFactory() *fakeGPURegistrationRetryTimerFactory { + return &fakeGPURegistrationRetryTimerFactory{ + requests: make(chan gpuRegistrationRetryTimerRequest, 16), + } +} + +func (f *fakeGPURegistrationRetryTimerFactory) newTimer(delay time.Duration) gpuRegistrationRetryTimer { + timer := newFakeGPURegistrationRetryTimer() + f.requests <- gpuRegistrationRetryTimerRequest{delay: delay, timer: timer} + return timer +} + +func waitForGPURegistrationCall(t *testing.T, calls <-chan struct{}) { + t.Helper() + select { + case <-calls: + case <-time.After(time.Second): + t.Fatal("registration callback was not invoked") + } +} + +func waitForGPURegistrationRetryTimer(t *testing.T, factory *fakeGPURegistrationRetryTimerFactory) gpuRegistrationRetryTimerRequest { + t.Helper() + select { + case request := <-factory.requests: + return request + case <-time.After(time.Second): + t.Fatal("registration retry timer was not requested") + return gpuRegistrationRetryTimerRequest{} + } +} + func TestGPURegistrationManagerCancellationWhileWaiting(t *testing.T) { var manager gpuRegistrationManager @@ -129,6 +199,179 @@ func TestGPURegistrationRetryBackoffIsBoundedAndResettable(t *testing.T) { assert.Equal(t, 5*time.Second, backoff.next()) } +func TestGPURegistrationManagerRunUsesBoundedRetryBackoff(t *testing.T) { + monitor := &fakeGPUAvailability{} + monitor.hasGPUs.Store(true) + registrationCalls := make(chan struct{}, 16) + timers := newFakeGPURegistrationRetryTimerFactory() + var manager gpuRegistrationManager + manager.configureGracefulNoGPU(monitor, false, 5*time.Second, func(context.Context) error { + registrationCalls <- struct{}{} + return errors.New("registration unavailable") + }) + manager.newRetryTimer = timers.newTimer + + ctx, cancel := context.WithCancel(context.Background()) + runDone := make(chan struct{}) + go func() { + defer close(runDone) + manager.run(ctx) + }() + manager.handleGPUStateChange(ctx, true) + + expectedDelays := []time.Duration{ + 5 * time.Second, + 10 * time.Second, + 20 * time.Second, + 40 * time.Second, + 80 * time.Second, + 160 * time.Second, + 5 * time.Minute, + 5 * time.Minute, + } + var activeTimer *fakeGPURegistrationRetryTimer + for index, expectedDelay := range expectedDelays { + waitForGPURegistrationCall(t, registrationCalls) + request := waitForGPURegistrationRetryTimer(t, timers) + assert.Equal(t, expectedDelay, request.delay) + activeTimer = request.timer + if index < len(expectedDelays)-1 { + request.timer.fire() + } + } + + cancel() + select { + case <-activeTimer.stopped: + case <-time.After(time.Second): + t.Fatal("active retry timer was not stopped after cancellation") + } + select { + case <-runDone: + case <-time.After(time.Second): + t.Fatal("registration retry loop did not stop after cancellation") + } +} + +func TestGPURegistrationManagerRunResetsBackoffOnGPUArrival(t *testing.T) { + monitor := &fakeGPUAvailability{} + monitor.hasGPUs.Store(true) + registrationCalls := make(chan struct{}, 4) + timers := newFakeGPURegistrationRetryTimerFactory() + var manager gpuRegistrationManager + manager.configureGracefulNoGPU(monitor, false, 5*time.Second, func(context.Context) error { + registrationCalls <- struct{}{} + return errors.New("registration unavailable") + }) + manager.newRetryTimer = timers.newTimer + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + runDone := make(chan struct{}) + go func() { + defer close(runDone) + manager.run(ctx) + }() + manager.handleGPUStateChange(ctx, true) + + waitForGPURegistrationCall(t, registrationCalls) + firstTimer := waitForGPURegistrationRetryTimer(t, timers) + assert.Equal(t, 5*time.Second, firstTimer.delay) + firstTimer.timer.fire() + + waitForGPURegistrationCall(t, registrationCalls) + secondTimer := waitForGPURegistrationRetryTimer(t, timers) + assert.Equal(t, 10*time.Second, secondTimer.delay) + manager.handleGPUStateChange(ctx, true) + + select { + case <-secondTimer.timer.stopped: + case <-time.After(time.Second): + t.Fatal("GPU arrival did not wake and stop the active retry timer") + } + waitForGPURegistrationCall(t, registrationCalls) + resetTimer := waitForGPURegistrationRetryTimer(t, timers) + assert.Equal(t, 5*time.Second, resetTimer.delay) + + cancel() + select { + case <-runDone: + case <-time.After(time.Second): + t.Fatal("registration retry loop did not stop") + } +} + +func TestGPURegistrationManagerRunCancellationStopsBackoffTimer(t *testing.T) { + monitor := &fakeGPUAvailability{} + monitor.hasGPUs.Store(true) + registrationCalls := make(chan struct{}, 1) + timers := newFakeGPURegistrationRetryTimerFactory() + var manager gpuRegistrationManager + manager.configureGracefulNoGPU(monitor, false, time.Second, func(context.Context) error { + registrationCalls <- struct{}{} + return errors.New("registration unavailable") + }) + manager.newRetryTimer = timers.newTimer + + ctx, cancel := context.WithCancel(context.Background()) + runDone := make(chan struct{}) + go func() { + defer close(runDone) + manager.run(ctx) + }() + manager.handleGPUStateChange(ctx, true) + waitForGPURegistrationCall(t, registrationCalls) + activeTimer := waitForGPURegistrationRetryTimer(t, timers) + + cancel() + select { + case <-activeTimer.timer.stopped: + case <-time.After(time.Second): + t.Fatal("active retry timer was not stopped") + } + select { + case <-runDone: + case <-time.After(time.Second): + t.Fatal("registration retry loop did not exit") + } +} + +func TestGPURegistrationManagerRunPreservesConfiguredSlowRetry(t *testing.T) { + monitor := &fakeGPUAvailability{} + monitor.hasGPUs.Store(true) + registrationCalls := make(chan struct{}, 2) + timers := newFakeGPURegistrationRetryTimerFactory() + var manager gpuRegistrationManager + manager.configureGracefulNoGPU(monitor, false, 10*time.Minute, func(context.Context) error { + registrationCalls <- struct{}{} + return errors.New("registration unavailable") + }) + manager.newRetryTimer = timers.newTimer + + ctx, cancel := context.WithCancel(context.Background()) + runDone := make(chan struct{}) + go func() { + defer close(runDone) + manager.run(ctx) + }() + manager.handleGPUStateChange(ctx, true) + + waitForGPURegistrationCall(t, registrationCalls) + firstTimer := waitForGPURegistrationRetryTimer(t, timers) + assert.Equal(t, 10*time.Minute, firstTimer.delay) + firstTimer.timer.fire() + waitForGPURegistrationCall(t, registrationCalls) + secondTimer := waitForGPURegistrationRetryTimer(t, timers) + assert.Equal(t, 10*time.Minute, secondTimer.delay) + + cancel() + select { + case <-runDone: + case <-time.After(time.Second): + t.Fatal("registration retry loop did not stop") + } +} + func TestGPURegistrationManagerGPULossMarksUnreadyAndPausesQueue(t *testing.T) { monitor := &fakeGPUAvailability{} monitor.hasGPUs.Store(true)