set output topology format when forwarding requests (#25)

Signed-off-by: Dmitry Shmulevich <[email protected]>
NVIDIA · Nov 1, 2024 · cb092f9 · cb092f9
1 parent 3f9ea84
commit cb092f9
Show file tree

Hide file tree

Showing 5 changed files with 109 additions and 58 deletions.
diff --git a/pkg/engines/k8s/labeler_test.go b/pkg/engines/k8s/labeler_test.go
@@ -21,9 +21,9 @@ import (
 	"fmt"
 	"testing"
 
-	"github.com/NVIDIA/topograph/pkg/common"
-	"github.com/NVIDIA/topograph/pkg/translate"
 	"github.com/stretchr/testify/require"
+
+	"github.com/NVIDIA/topograph/pkg/translate"
 )
 
 type testLabeler struct {
@@ -50,7 +50,7 @@ func TestApplyNodeLabels(t *testing.T) {
 		"Node306": {"topology.kubernetes.io/network-level-1": "xf946c4acef2d5939", "topology.kubernetes.io/network-level-2": "S1"},
 	}
 
-	err := NewTopologyLabeler().ApplyNodeLabels(context.TODO(), root.Vertices[common.ValTopologyTree], labeler)
+	err := NewTopologyLabeler().ApplyNodeLabels(context.TODO(), root, labeler)
 	require.NoError(t, err)
 	require.Equal(t, data, labeler.data)
 }
diff --git a/pkg/server/grpc_client.go b/pkg/server/grpc_client.go
@@ -25,6 +25,7 @@ import (
 	"k8s.io/klog/v2"
 
 	"github.com/NVIDIA/topograph/pkg/common"
+	"github.com/NVIDIA/topograph/pkg/metrics"
 	pb "github.com/NVIDIA/topograph/pkg/protos"
 )
 
@@ -58,10 +59,20 @@ func forwardRequest(ctx context.Context, tr *common.TopologyRequest, url string,
 
 	klog.V(4).Infof("Response: %s", response.String())
 
-	return toGraph(response, cis), nil
+	return toGraph(response, cis, getTopologyFormat(tr.Engine.Params)), nil
 }
 
-func toGraph(response *pb.TopologyResponse, cis []common.ComputeInstances) *common.Vertex {
+// getTopologyFormat derives topology format from engine parameters: tree (default) or block
+func getTopologyFormat(params map[string]string) string {
+	if len(params) != 0 {
+		if format, ok := params[common.KeyPlugin]; ok && len(format) != 0 {
+			return format
+		}
+	}
+	return common.ValTopologyTree
+}
+
+func toGraph(response *pb.TopologyResponse, cis []common.ComputeInstances, format string) *common.Vertex {
 	i2n := make(map[string]string)
 	for _, ci := range cis {
 		for instance, node := range ci.Instances {
@@ -135,9 +146,10 @@ func toGraph(response *pb.TopologyResponse, cis []common.ComputeInstances) *comm
 	}
 
 	if len(i2n) != 0 {
-		klog.V(4).Infof("Adding unclaimed nodes %v", i2n)
+		klog.V(4).Infof("Adding nodes w/o topology: %v", i2n)
+		metrics.SetMissingTopology("GTS", len(i2n))
 		sw := &common.Vertex{
-			ID:       "extra",
+			ID:       common.NoTopology,
 			Vertices: make(map[string]*common.Vertex),
 		}
 		for instanceID, nodeName := range i2n {
@@ -146,7 +158,7 @@ func toGraph(response *pb.TopologyResponse, cis []common.ComputeInstances) *comm
 				ID:   instanceID,
 			}
 		}
-		forest["extra"] = sw
+		forest[common.NoTopology] = sw
 	}
 
 	treeRoot := &common.Vertex{
@@ -156,15 +168,25 @@ func toGraph(response *pb.TopologyResponse, cis []common.ComputeInstances) *comm
 		treeRoot.Vertices[name] = node
 	}
 
-	blockRoot := &common.Vertex{
-		Vertices: make(map[string]*common.Vertex),
-	}
-	for name, domain := range blocks {
-		blockRoot.Vertices[name] = domain
-	}
+	metadata := map[string]string{common.KeyPlugin: format}
+	if format == common.ValTopologyBlock {
+		blockRoot := &common.Vertex{
+			Vertices: make(map[string]*common.Vertex),
+		}
+		for name, domain := range blocks {
+			blockRoot.Vertices[name] = domain
+		}
+
+		return &common.Vertex{
+			Vertices: map[string]*common.Vertex{
+				common.ValTopologyBlock: blockRoot,
+				common.ValTopologyTree:  treeRoot,
+			},
+			Metadata: metadata,
+		}
+	} else {
+		treeRoot.Metadata = metadata
 
-	root := &common.Vertex{
-		Vertices: map[string]*common.Vertex{common.ValTopologyBlock: blockRoot, common.ValTopologyTree: treeRoot},
+		return treeRoot
 	}
-	return root
 }
diff --git a/pkg/server/grpc_client_test.go b/pkg/server/grpc_client_test.go
@@ -135,12 +135,53 @@ func TestToGraph(t *testing.T) {
 
 	nv1 := &common.Vertex{ID: "nvlink-nv1", Vertices: map[string]*common.Vertex{"n10-1": v101, "n10-2": v102, "n11-1": v111, "n11-2": v112}}
 
-	extra := &common.Vertex{ID: "extra", Vertices: map[string]*common.Vertex{"cpu1": cpu1}}
-	treeRoot := &common.Vertex{Vertices: map[string]*common.Vertex{"nvlink-nv1": nv1, "sw3": sw3, "extra": extra}}
+	extra := &common.Vertex{ID: common.NoTopology, Vertices: map[string]*common.Vertex{"cpu1": cpu1}}
+	treeRoot := &common.Vertex{Vertices: map[string]*common.Vertex{"nvlink-nv1": nv1, "sw3": sw3, common.NoTopology: extra}}
 	blockRoot := &common.Vertex{Vertices: map[string]*common.Vertex{"nvlink-nv1": nv1}}
 	root := &common.Vertex{
 		Vertices: map[string]*common.Vertex{common.ValTopologyBlock: blockRoot, common.ValTopologyTree: treeRoot},
+		Metadata: map[string]string{common.KeyPlugin: common.ValTopologyBlock},
 	}
 
-	require.Equal(t, root, toGraph(&pb.TopologyResponse{Instances: instances}, cis))
+	require.Equal(t, root, toGraph(&pb.TopologyResponse{Instances: instances}, cis, common.ValTopologyBlock))
+}
+
+func TestGetTopologyFormat(t *testing.T) {
+	testCases := []struct {
+		name   string
+		params map[string]string
+		format string
+	}{
+		{
+			name:   "Case 1: nil params",
+			params: nil,
+			format: common.ValTopologyTree,
+		},
+		{
+			name:   "Case 2: empty params",
+			params: make(map[string]string),
+			format: common.ValTopologyTree,
+		},
+		{
+			name:   "Case 3: missing key",
+			params: map[string]string{"a": "b"},
+			format: common.ValTopologyTree,
+		},
+		{
+			name:   "Case 4: block topology",
+			params: map[string]string{common.KeyPlugin: common.ValTopologyBlock},
+			format: common.ValTopologyBlock,
+		},
+		{
+			name:   "Case 5: tree topology",
+			params: map[string]string{common.KeyPlugin: common.ValTopologyTree},
+			format: common.ValTopologyTree,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			require.Equal(t, tc.format, getTopologyFormat(tc.params))
+		})
+	}
 }
diff --git a/pkg/translate/output.go b/pkg/translate/output.go
@@ -151,11 +151,10 @@ func toBlockTopology(wr io.Writer, root *common.Vertex) error {
 }
 
 func toTreeTopology(wr io.Writer, root *common.Vertex) error {
-	treeRoot := root.Vertices[common.ValTopologyTree]
 	visited := make(map[string]bool)
 	leaves := make(map[string][]string)
 	parents := []*common.Vertex{}
-	queue := []*common.Vertex{treeRoot}
+	queue := []*common.Vertex{root}
 	idToName := make(map[string]string)
 
 	for len(queue) > 0 {
@@ -367,14 +366,8 @@ func GetTreeTestSet(testForLongLabelName bool) (*common.Vertex, map[string]strin
 		ID:       "S1",
 		Vertices: map[string]*common.Vertex{"S2": sw2, s3name: sw3},
 	}
-	treeRoot := &common.Vertex{
-		Vertices: map[string]*common.Vertex{"S1": sw1},
-	}
-	blockRoot := &common.Vertex{
-		Vertices: map[string]*common.Vertex{},
-	}
 	root := &common.Vertex{
-		Vertices: map[string]*common.Vertex{common.ValTopologyBlock: blockRoot, common.ValTopologyTree: treeRoot},
+		Vertices: map[string]*common.Vertex{"S1": sw1},
 	}
 
 	return root, instance2node

diff --git a/pkg/translate/output_test.go b/pkg/translate/output_test.go
@@ -99,49 +99,44 @@ func TestToBlockIBTopology(t *testing.T) {
 func TestToSlurmNameShortener(t *testing.T) {
 	v := &common.Vertex{
 		Vertices: map[string]*common.Vertex{
-			common.ValTopologyTree: {
+			"hpcislandid-1": {
+				ID:   "hpcislandid-1",
+				Name: "switch.3.1",
 				Vertices: map[string]*common.Vertex{
-					"hpcislandid-1": {
-						ID:   "hpcislandid-1",
-						Name: "switch.3.1",
+					"network-block-1": {
+						ID:   "network-block-1",
+						Name: "switch.2.1",
 						Vertices: map[string]*common.Vertex{
-							"network-block-1": {
-								ID:   "network-block-1",
-								Name: "switch.2.1",
+							"local-block-1": {
+								ID:   "local-block-1",
+								Name: "switch.1.1",
 								Vertices: map[string]*common.Vertex{
-									"local-block-1": {
-										ID:   "local-block-1",
-										Name: "switch.1.1",
-										Vertices: map[string]*common.Vertex{
-											"node-1": {
-												ID:   "node-1-id",
-												Name: "node-1",
-											},
-										},
+									"node-1": {
+										ID:   "node-1-id",
+										Name: "node-1",
 									},
 								},
 							},
-							"network-block-2": {
-								ID:   "network-block-2",
-								Name: "switch.2.2",
+						},
+					},
+					"network-block-2": {
+						ID:   "network-block-2",
+						Name: "switch.2.2",
+						Vertices: map[string]*common.Vertex{
+							"local-block-2": {
+								ID:   "local-block-2",
+								Name: "switch.1.2",
 								Vertices: map[string]*common.Vertex{
-									"local-block-2": {
-										ID:   "local-block-2",
-										Name: "switch.1.2",
-										Vertices: map[string]*common.Vertex{
-											"node-2": {
-												ID:   "node-2-id",
-												Name: "node-2",
-											},
-										},
+									"node-2": {
+										ID:   "node-2-id",
+										Name: "node-2",
 									},
 								},
 							},
 						},
 					},
 				},
 			},
-			common.ValTopologyBlock: {Vertices: map[string]*common.Vertex{}},
 		},
 	}