fix(nvca): make no-GPU recovery registration-safe

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",

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.