Traverse Topology in DFS format in block topology output (#32)

Signed-off-by: Ritika Srivastava <[email protected]>
NVIDIA · Nov 22, 2024 · b32d034 · b32d034
1 parent b5fb02b
commit b32d034
Show file tree

Hide file tree

Showing 4 changed files with 125 additions and 24 deletions.
diff --git a/pkg/server/grpc_client.go b/pkg/server/grpc_client.go
@@ -107,13 +107,12 @@ func toGraph(response *pb.TopologyResponse, cis []topology.ComputeInstances, for
 		if len(ins.NvlinkDomain) != 0 {
 			klog.V(4).Infof("Adding node %q to NVLink domain %q", nodeName, ins.NvlinkDomain)
 			switchName := fmt.Sprintf("nvlink-%s", ins.NvlinkDomain)
-			sw, ok := forest[switchName]
+			sw, ok := blocks[switchName]
 			if !ok {
 				sw = &topology.Vertex{
 					ID:       switchName,
 					Vertices: map[string]*topology.Vertex{id: vertex},
 				}
-				forest[switchName] = sw
 				blocks[switchName] = sw
 			} else {
 				sw.Vertices[id] = vertex

diff --git a/pkg/server/grpc_client_test.go b/pkg/server/grpc_client_test.go
@@ -136,7 +136,7 @@ func TestToGraph(t *testing.T) {
 	nv1 := &topology.Vertex{ID: "nvlink-nv1", Vertices: map[string]*topology.Vertex{"n10-1": v101, "n10-2": v102, "n11-1": v111, "n11-2": v112}}
 
 	extra := &topology.Vertex{ID: topology.NoTopology, Vertices: map[string]*topology.Vertex{"cpu1": cpu1}}
-	treeRoot := &topology.Vertex{Vertices: map[string]*topology.Vertex{"nvlink-nv1": nv1, "sw3": sw3, topology.NoTopology: extra}}
+	treeRoot := &topology.Vertex{Vertices: map[string]*topology.Vertex{"sw3": sw3, topology.NoTopology: extra}}
 	blockRoot := &topology.Vertex{Vertices: map[string]*topology.Vertex{"nvlink-nv1": nv1}}
 	root := &topology.Vertex{
 		Vertices: map[string]*topology.Vertex{topology.TopologyBlock: blockRoot, topology.TopologyTree: treeRoot},

diff --git a/pkg/translate/output.go b/pkg/translate/output.go
@@ -112,40 +112,23 @@ func getBlockSize(domainVisited map[string]int, adminBlockSize string) string {
 }
 
 func toBlockTopology(wr io.Writer, root *topology.Vertex) error {
-	// traverse tree topology and when a node is reached, check within blockRoot for domain and print that domain.
+	// traverse tree topology in DFS manner and when a node is reached, check within blockRoot for domain and print that domain.
 	// keep a map of which domain has been printed
 	treeRoot := root.Vertices[topology.TopologyTree]
 	blockRoot := root.Vertices[topology.TopologyBlock]
 	visited := make(map[string]bool)
-	queue := []*topology.Vertex{treeRoot}
 	domainVisited := make(map[string]int)
 
 	if treeRoot != nil {
-		for len(queue) > 0 {
-			v := queue[0]
-			queue = queue[1:]
-
-			// sort the IDs
-			keys := sortVertices(v)
-			for _, key := range keys {
-				w := v.Vertices[key]
-				if len(w.Vertices) == 0 { // it's a leaf; don't add to queue
-					err := findBlock(wr, w.ID, blockRoot, domainVisited)
-					if err != nil {
-						return err
-					}
-				} else if !visited[w.ID] {
-					queue = append(queue, w)
-					visited[w.ID] = true
-				}
-			}
+		err := dfsTraversal(wr, treeRoot, blockRoot, visited, domainVisited)
+		if err != nil {
+			return err
 		}
 	}
 	err := printDisconnectedBlocks(wr, blockRoot, domainVisited)
 	if err != nil {
 		return err
 	}
-
 	blockSize := ""
 	if _, exists := root.Metadata[topology.KeyBlockSizes]; exists {
 		blockSize = root.Metadata[topology.KeyBlockSizes]
@@ -155,6 +138,28 @@ func toBlockTopology(wr io.Writer, root *topology.Vertex) error {
 	return err
 }
 
+func dfsTraversal(wr io.Writer, curVertex *topology.Vertex, blockRoot *topology.Vertex, visited map[string]bool, domainVisited map[string]int) error {
+	visited[curVertex.ID] = true
+	keys := sortVertices(curVertex)
+	for _, key := range keys {
+		w := curVertex.Vertices[key]
+		if len(w.Vertices) == 0 { // it's a leaf; don't add to queue
+			err := findBlock(wr, w.ID, blockRoot, domainVisited)
+			if err != nil {
+				return err
+			}
+		} else {
+			if !visited[w.ID] {
+				err := dfsTraversal(wr, w, blockRoot, visited, domainVisited)
+				if err != nil {
+					return err
+				}
+			}
+		}
+	}
+	return nil
+}
+
 func toTreeTopology(wr io.Writer, root *topology.Vertex) error {
 	visited := make(map[string]bool)
 	leaves := make(map[string][]string)
@@ -418,6 +423,7 @@ func GetBlockWithMultiIBTestSet() (*topology.Vertex, map[string]string) {
 		Vertices: map[string]*topology.Vertex{"S5": sw5, "S6": sw6},
 	}
 	ibRoot1 := &topology.Vertex{
+		ID:       "ibRoot1",
 		Vertices: map[string]*topology.Vertex{"S4": sw4},
 	}
 
@@ -434,6 +440,7 @@ func GetBlockWithMultiIBTestSet() (*topology.Vertex, map[string]string) {
 		Vertices: map[string]*topology.Vertex{"S2": sw2, "S3": sw3},
 	}
 	ibRoot2 := &topology.Vertex{
+		ID:       "ibRoot2",
 		Vertices: map[string]*topology.Vertex{"S1": sw1},
 	}
 
@@ -527,6 +534,80 @@ func GetBlockWithIBTestSet() (*topology.Vertex, map[string]string) {
 	return root, instance2node
 }
 
+func GetBlockWithDFSIBTestSet() (*topology.Vertex, map[string]string) {
+	instance2node := map[string]string{
+		"I14": "Node104", "I15": "Node105",
+		"I22": "Node202", "I25": "Node205",
+	}
+
+	n14 := &topology.Vertex{ID: "I14", Name: "Node104"}
+	n15 := &topology.Vertex{ID: "I15", Name: "Node105"}
+
+	n22 := &topology.Vertex{ID: "I22", Name: "Node202"}
+	n25 := &topology.Vertex{ID: "I25", Name: "Node205"}
+
+	sw2 := &topology.Vertex{
+		ID:       "S2",
+		Vertices: map[string]*topology.Vertex{"I14": n14, "I15": n15},
+	}
+
+	sw4 := &topology.Vertex{
+		ID:       "S4",
+		Vertices: map[string]*topology.Vertex{"I22": n22},
+	}
+
+	sw5 := &topology.Vertex{
+		ID:       "S5",
+		Vertices: map[string]*topology.Vertex{"I25": n25},
+	}
+
+	sw3 := &topology.Vertex{
+		ID:       "S3",
+		Vertices: map[string]*topology.Vertex{"S5": sw5},
+	}
+	sw1 := &topology.Vertex{
+		ID:       "S1",
+		Vertices: map[string]*topology.Vertex{"S4": sw4},
+	}
+
+	sw0 := &topology.Vertex{
+		ID:       "S0",
+		Vertices: map[string]*topology.Vertex{"S1": sw1, "S2": sw2, "S3": sw3},
+	}
+
+	treeRoot := &topology.Vertex{
+		Vertices: map[string]*topology.Vertex{"S0": sw0},
+	}
+
+	block2 := &topology.Vertex{
+		ID:       "B2",
+		Vertices: map[string]*topology.Vertex{"I14": n14, "I15": n15},
+	}
+	block1 := &topology.Vertex{
+		ID:       "B1",
+		Vertices: map[string]*topology.Vertex{"I22": n22},
+	}
+
+	block3 := &topology.Vertex{
+		ID:       "B3",
+		Vertices: map[string]*topology.Vertex{"I25": n25},
+	}
+
+	blockRoot := &topology.Vertex{
+		Vertices: map[string]*topology.Vertex{"B1": block1, "B2": block2, "B3": block3},
+	}
+
+	root := &topology.Vertex{
+		Vertices: map[string]*topology.Vertex{topology.TopologyBlock: blockRoot, topology.TopologyTree: treeRoot},
+		Metadata: map[string]string{
+			topology.KeyEngine:     engines.EngineSLURM,
+			topology.KeyPlugin:     topology.TopologyBlock,
+			topology.KeyBlockSizes: "1",
+		},
+	}
+	return root, instance2node
+}
+
 func GetBlockTestSet() (*topology.Vertex, map[string]string) {
 	instance2node := map[string]string{
 		"I14": "Node104", "I15": "Node105", "I16": "Node106",

diff --git a/pkg/translate/output_test.go b/pkg/translate/output_test.go
@@ -41,6 +41,13 @@ BlockName=B1 Nodes=Node[104-106]
 BlockName=B2 Nodes=Node[201-202],Node205
 BlockSizes=3
 `
+
+	testBlockConfigDFS = `BlockName=B1 Nodes=Node202
+BlockName=B2 Nodes=Node[104-105]
+BlockName=B3 Nodes=Node205
+BlockSizes=1
+`
+
 	shortNameExpectedResult = `# switch.3.1=hpcislandid-1
 SwitchName=switch.3.1 Switches=switch.2.[1-2]
 # switch.2.1=network-block-1
@@ -100,6 +107,20 @@ func TestToBlockIBTopology(t *testing.T) {
 	}
 }
 
+func TestToBlockDFSIBTopology(t *testing.T) {
+	v, _ := GetBlockWithDFSIBTestSet()
+	require.Equal(t, v.Metadata[topology.KeyPlugin], topology.TopologyBlock)
+	buf := &bytes.Buffer{}
+	err := ToGraph(buf, v)
+	require.NoError(t, err)
+	switch buf.String() {
+	case testBlockConfigDFS:
+		// nop
+	default:
+		t.Errorf("unexpected result %s", buf.String())
+	}
+}
+
 func TestToSlurmNameShortener(t *testing.T) {
 	v := &topology.Vertex{
 		Vertices: map[string]*topology.Vertex{