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 67df76601..a5b0560ce 100644 --- a/src/compute-plane-services/nvca/pkg/nvca/agent.go +++ b/src/compute-plane-services/nvca/pkg/nvca/agent.go @@ -364,8 +364,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. @@ -381,9 +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 + // 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 @@ -598,6 +600,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 +1130,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). @@ -1189,11 +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) + 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 { @@ -1206,8 +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.gpuRegistration.getRegistrationStatus)) } if a.FeatureFlagFetcher.IsAttributeEnabled(featureflag.AttrHostIsolation) { statusUpdaters = append(statusUpdaters, hostisolation.NewStatusGetter( @@ -1235,10 +1249,13 @@ 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") + 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, @@ -1251,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") @@ -1397,35 +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() } - // 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() - } - }) - - // 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 @@ -2156,24 +2153,26 @@ func (a *Agent) RenewICMSQueueCreds(ctx context.Context) error { return nil } - 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 529e1c5f9..b8c6d7935 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,593 @@ 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 + credentialResponse *types.ICMSCredentialResponse + credentialErr error + credentialStarted chan struct{} + credentialRelease chan struct{} + credentialReleaseOnce sync.Once +} + +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) 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() + 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: 365 * 24 * time.Hour, + 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.gpuRegistration.monitor) + require.NotNil(t, agent.queueManager) + 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) + + _, 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.gpuRegistration.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.gpuRegistration.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) { + logCtx, logHook := core.WithTestingLogger(newTestContext()) + ctx, cancel := context.WithCancel(core.WithRandomSeed(logCtx, 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) + 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) + 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 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.gpuRegistration.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.gpuRegistration.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) + + 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.gpuRegistration.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.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.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()) + 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..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,21 +74,23 @@ func (a *Agent) initICMSRegistrationSyncer(ctx context.Context) error { var lastRegTime = core.GetCurrentTime(ctx) a.syncICMSRegistration = func(ctx context.Context) error { - if timeSinceLastUpdate := core.GetCurrentTime(ctx).Sub(lastRegTime); timeSinceLastUpdate < time.Hour { - return nil - } + return a.gpuRegistration.withRegistrationOperation(ctx, func() error { + 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 } @@ -98,35 +100,37 @@ func (a *Agent) initICMSRegistrationSyncer(ctx context.Context) error { var lastRegTime = core.GetCurrentTime(ctx) a.syncICMSRegistration = func(ctx context.Context) error { - 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..d73e9e40b --- /dev/null +++ b/src/compute-plane-services/nvca/pkg/nvca/gpu_registration_manager.go @@ -0,0 +1,322 @@ +/* +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" +const maxGracefulNoGPURegistrationRetryInterval = 5 * time.Minute + +type gpuRegistrationMonitor interface { + HasGPUs() bool + SetOnGPUStateChange(GPUStateChangeCallback) + Start(context.Context) + GetComponentStatus(context.Context) (types.AgentHealth, error) +} + +type gpuRegistrationQueue interface { + Pause() + 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 + monitor gpuRegistrationMonitor + queueManager gpuRegistrationQueue + ready atomic.Bool + generation atomic.Uint64 + registrationRequests chan struct{} + retryInterval time.Duration + register func(context.Context) error + newRetryTimer func(time.Duration) gpuRegistrationRetryTimer +} + +// 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{} +} + +// 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 { + 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) + 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 + } +} + +// 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 + } + defer m.operationGate.unlock() + return operation() +} + +// configureGracefulNoGPU initializes state before the monitor and registration +// worker start. +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) +} + +// 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, + 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 +} + +// 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) + + 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: + } +} + +// 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, + ) + newRetryTimer := m.newRetryTimer + if newRetryTimer == nil { + newRetryTimer = func(delay time.Duration) gpuRegistrationRetryTimer { + return &realGPURegistrationRetryTimer{timer: time.NewTimer(delay)} + } + } + + for { + select { + case <-ctx.Done(): + return + case <-m.registrationRequests: + backoff.reset() + } + + for !m.ready.Load() && m.hasGPUs() { + if !m.tryRegistration(ctx) { + break + } + + retryTimer := newRetryTimer(backoff.next()) + select { + case <-ctx.Done(): + retryTimer.Stop() + return + case <-m.registrationRequests: + retryTimer.Stop() + backoff.reset() + 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..dbee64cb6 --- /dev/null +++ b/src/compute-plane-services/nvca/pkg/nvca/gpu_registration_manager_test.go @@ -0,0 +1,424 @@ +/* +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" + "errors" + "sync" + "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) +} + +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() +} + +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 + + 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 + + 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() { + waiterDone <- manager.withRegistrationOperation(ctx, func() error { + waiterRan.Store(true) + return nil + }) + }() + <-ctx.doneObserved + 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 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 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) + 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()) +} 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 {