Merge pull request NVIDIA#391 from klueska/fix-bug-termination

Ensure we exit cleanly when compute-domain-daemon shutdown prematurely

Author: klueska

cmd/compute-domain-daemon/computedomain.go

@@ -197,7 +197,11 @@ func (m *ComputeDomainManager) UpdateComputeDomainNodeInfo(ctx context.Context,
 
 // BlockUntilAllNodesJoinComputeDomain waits until all nodes have joined the compute domain
 // and returns the list of nodes in the compute domain.
-func (m *ComputeDomainManager) BlockUntilAllNodesJoinComputeDomain() []*nvapi.ComputeDomainNode {
-	<-m.nodesChan
-	return m.nodes
+func (m *ComputeDomainManager) BlockUntilAllNodesJoinComputeDomain(ctx context.Context) ([]*nvapi.ComputeDomainNode, error) {
+	select {
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	case <-m.nodesChan:
+		return m.nodes, nil
+	}
 }

cmd/compute-domain-daemon/controller.go

@@ -104,6 +104,6 @@ func (c *Controller) Run(ctx context.Context) error {
 
 // BlockUntilAllNodesJoinComputeDomain waits until all nodes have joined the compute domain
 // and returns the list of nodes in the compute domain.
-func (c *Controller) BlockUntilAllNodesJoinComputeDomain() []*nvapi.ComputeDomainNode {
-	return c.computeDomainManager.BlockUntilAllNodesJoinComputeDomain()
+func (c *Controller) BlockUntilAllNodesJoinComputeDomain(ctx context.Context) ([]*nvapi.ComputeDomainNode, error) {
+	return c.computeDomainManager.BlockUntilAllNodesJoinComputeDomain(ctx)
 }

cmd/compute-domain-daemon/main.go

@@ -173,7 +173,10 @@ func run(ctx context.Context, cancel context.CancelFunc, flags *Flags) error {
 	}()
 
 	// Wait until all nodes have joined the compute domain
-	nodes := controller.BlockUntilAllNodesJoinComputeDomain()
+	nodes, err := controller.BlockUntilAllNodesJoinComputeDomain(ctx)
+	if err != nil {
+		return fmt.Errorf("error waiting for all nodes to join ComputeDomain: %w", err)
+	}
 
 	if flags.cliqueID == "" {
 		fmt.Println("ClusterUUID and CliqueId are NOT set for GPUs on this node.")