Standardize vertex structure with both tree and block children

Signed-off-by: Ritika Srivastava <[email protected]>
NVIDIA · Nov 26, 2024 · 52efde3 · 52efde3
1 parent 43a3ff3
commit 52efde3
Show file tree

Hide file tree

Showing 14 changed files with 57 additions and 75 deletions.
diff --git a/pkg/engines/k8s/labeler_test.go b/pkg/engines/k8s/labeler_test.go
@@ -23,7 +23,6 @@ import (
 
 	"github.com/stretchr/testify/require"
 
-	"github.com/NVIDIA/topograph/pkg/topology"
 	"github.com/NVIDIA/topograph/pkg/translate"
 )
 
@@ -41,7 +40,6 @@ func (l *testLabeler) AddNodeLabels(_ context.Context, nodeName string, labels m
 
 func TestApplyNodeLabels(t *testing.T) {
 	root, _ := translate.GetTreeTestSet(true)
-	require.Equal(t, root.Metadata[topology.KeyPlugin], topology.TopologyTree)
 	labeler := &testLabeler{data: make(map[string]map[string]string)}
 	data := map[string]map[string]string{
 		"Node201": {"topology.kubernetes.io/network-level-1": "S2", "topology.kubernetes.io/network-level-2": "S1"},

diff --git a/pkg/engines/slurm/slurm.go b/pkg/engines/slurm/slurm.go
@@ -148,8 +148,17 @@ func GenerateOutputParams(ctx context.Context, tree *topology.Vertex, params *Pa
 			plugin = tree.Metadata[topology.KeyPlugin]
 		}
 		if len(plugin) == 0 {
+			if tree.Vertices[topology.TopologyBlock] != nil {
+				plugin = topology.TopologyBlock
+			} else {
+				plugin = topology.TopologyTree
+			}
+		} else if plugin == topology.TopologyBlock && tree.Vertices[topology.TopologyBlock] == nil {
+			klog.Infof("Admin provided plugin %v not valid, overriding with %v", plugin, topology.TopologyTree)
 			plugin = topology.TopologyTree
+			tree.Metadata[topology.KeyPlugin] = plugin
 		}
+
 		if _, err := buf.WriteString(fmt.Sprintf(TopologyHeader, plugin)); err != nil {
 			return nil, err
 		}

diff --git a/pkg/engines/test/test.go b/pkg/engines/test/test.go
@@ -19,7 +19,6 @@ package test
 import (
 	"context"
 	"errors"
-	"fmt"
 
 	"github.com/NVIDIA/topograph/internal/config"
 	"github.com/NVIDIA/topograph/pkg/engines"
@@ -59,10 +58,6 @@ func (eng *TestEngine) GenerateOutput(ctx context.Context, tree *topology.Vertex
 		return nil, err
 	}
 
-	if len(tree.Metadata) == 0 {
-		return nil, fmt.Errorf("metadata for test engine not set")
-	}
-
 	if len(p.Plugin) != 0 {
 		tree.Metadata[topology.KeyPlugin] = p.Plugin
 	}

diff --git a/pkg/ib/ib.go b/pkg/ib/ib.go
@@ -64,13 +64,14 @@ func GenerateTopologyConfig(data []byte) (*topology.Vertex, error) {
 	}
 	seen = make(map[int]map[string]*Switch)
 	root.simplify(root.getHeight())
-	rootNode, err := root.toGraph()
+	treeNode, err := root.toGraph()
 	if err != nil {
 		return nil, err
 	}
-	rootNode.Metadata = map[string]string{
-		topology.KeyPlugin: topology.TopologyTree,
+	rootNode := &topology.Vertex{
+		Vertices: make(map[string]*topology.Vertex),
 	}
+	rootNode.Vertices[topology.TopologyTree] = treeNode
 	return rootNode, nil
 }
 

diff --git a/pkg/models/model.go b/pkg/models/model.go
@@ -213,13 +213,15 @@ func (model *Model) ToGraph() (*topology.Vertex, map[string]string) {
 	for k, v := range blockVertexMap {
 		blockRoot.Vertices[k] = v
 	}
+
+	rootNode := &topology.Vertex{
+		Vertices: make(map[string]*topology.Vertex),
+	}
+
 	if block_topology {
-		root := &topology.Vertex{
-			Vertices: map[string]*topology.Vertex{topology.TopologyBlock: blockRoot, topology.TopologyTree: treeRoot},
-			Metadata: map[string]string{topology.KeyPlugin: topology.TopologyBlock},
-		}
-		return root, instance2node
+		rootNode.Vertices[topology.TopologyBlock] = blockRoot
 	}
-	treeRoot.Metadata = map[string]string{topology.KeyPlugin: topology.TopologyTree}
-	return treeRoot, instance2node
+
+	rootNode.Vertices[topology.TopologyTree] = treeRoot
+	return rootNode, instance2node
 }
diff --git a/pkg/providers/aws/instance_topology.go b/pkg/providers/aws/instance_topology.go
@@ -179,15 +179,16 @@ func toGraph(top []types.InstanceTopology, cis []topology.ComputeInstances) (*to
 		forest[topology.NoTopology] = sw
 	}
 
-	root := &topology.Vertex{
+	treeRoot := &topology.Vertex{
 		Vertices: make(map[string]*topology.Vertex),
-		Metadata: map[string]string{
-			topology.KeyPlugin: topology.TopologyTree,
-		},
 	}
 	for name, node := range forest {
-		root.Vertices[name] = node
+		treeRoot.Vertices[name] = node
 	}
 
+	root := &topology.Vertex{
+		Vertices: make(map[string]*topology.Vertex),
+	}
+	root.Vertices[topology.TopologyTree] = treeRoot
 	return root, nil
 }
diff --git a/pkg/providers/aws/instance_topology_test.go b/pkg/providers/aws/instance_topology_test.go
@@ -106,11 +106,12 @@ func TestNewInstanceTopology(t *testing.T) {
 		Vertices: map[string]*topology.Vertex{"nn-224a2a4d9df61a975": v2},
 	}
 
-	expected := &topology.Vertex{
+	v0 := &topology.Vertex{
 		Vertices: map[string]*topology.Vertex{"nn-098f9e7674016cb1c": v1},
-		Metadata: map[string]string{
-			topology.KeyPlugin: topology.TopologyTree,
-		},
+	}
+
+	expected := &topology.Vertex{
+		Vertices: map[string]*topology.Vertex{topology.TopologyTree: v0},
 	}
 
 	tree, err := toGraph(top, []topology.ComputeInstances{{Instances: i2n}})

diff --git a/pkg/providers/baremetal/mnnvl.go b/pkg/providers/baremetal/mnnvl.go
@@ -8,7 +8,6 @@ import (
 	"strings"
 
 	"github.com/NVIDIA/topograph/internal/exec"
-	"github.com/NVIDIA/topograph/pkg/engines"
 	"github.com/NVIDIA/topograph/pkg/ib"
 	"github.com/NVIDIA/topograph/pkg/topology"
 )
@@ -208,9 +207,6 @@ func toGraph(domainMap map[string]domain, treeRoot *topology.Vertex) *topology.V
 		}
 		blockRoot.Vertices[domainName] = tree
 	}
-	// add root metadata
-	root.Metadata[topology.KeyEngine] = engines.EngineSLURM
-	root.Metadata[topology.KeyPlugin] = topology.TopologyBlock
 	root.Vertices[topology.TopologyBlock] = blockRoot
 	return root
 }

diff --git a/pkg/providers/gcp/instance_topology.go b/pkg/providers/gcp/instance_topology.go
@@ -141,15 +141,17 @@ func (cfg *InstanceTopology) toGraph() (*topology.Vertex, error) {
 		sw1.Vertices[id2] = sw2
 	}
 
-	root := &topology.Vertex{
+	treeRoot := &topology.Vertex{
 		Vertices: make(map[string]*topology.Vertex),
-		Metadata: map[string]string{
-			topology.KeyPlugin: topology.TopologyTree,
-		},
 	}
 	for name, node := range forest {
-		root.Vertices[name] = node
+		treeRoot.Vertices[name] = node
+	}
+
+	root := &topology.Vertex{
+		Vertices: make(map[string]*topology.Vertex),
 	}
+	root.Vertices[topology.TopologyTree] = treeRoot
 
 	return root, nil
 }

diff --git a/pkg/providers/oci/instance_topology.go b/pkg/providers/oci/instance_topology.go
@@ -220,16 +220,17 @@ func toGraph(bareMetalHostSummaries []*core.ComputeBareMetalHostSummary, cis []t
 		forest[topology.NoTopology] = sw
 	}
 
-	root := &topology.Vertex{
+	treeRoot := &topology.Vertex{
 		Vertices: make(map[string]*topology.Vertex),
-		Metadata: map[string]string{
-			topology.KeyPlugin: topology.TopologyTree,
-		},
 	}
 	for name, node := range forest {
-		root.Vertices[name] = node
+		treeRoot.Vertices[name] = node
 	}
 
+	root := &topology.Vertex{
+		Vertices: make(map[string]*topology.Vertex),
+	}
+	root.Vertices[topology.TopologyTree] = treeRoot
 	return root, nil
 
 }

diff --git a/pkg/server/grpc_client.go b/pkg/server/grpc_client.go
@@ -169,24 +169,19 @@ func toGraph(response *pb.TopologyResponse, cis []topology.ComputeInstances, for
 		treeRoot.Vertices[name] = node
 	}
 
-	metadata := map[string]string{topology.KeyPlugin: format}
+	root := &topology.Vertex{
+		Vertices: make(map[string]*topology.Vertex),
+	}
 	if format == topology.TopologyBlock {
 		blockRoot := &topology.Vertex{
 			Vertices: make(map[string]*topology.Vertex),
 		}
 		for name, domain := range blocks {
 			blockRoot.Vertices[name] = domain
 		}
-
-		return &topology.Vertex{
-			Vertices: map[string]*topology.Vertex{
-				topology.TopologyBlock: blockRoot,
-				topology.TopologyTree:  treeRoot,
-			},
-			Metadata: metadata,
-		}
-	} else {
-		treeRoot.Metadata = metadata
-		return treeRoot
+		root.Vertices[topology.TopologyBlock] = blockRoot
 	}
+	root.Vertices[topology.TopologyTree] = treeRoot
+	return root
+
 }
diff --git a/pkg/server/grpc_client_test.go b/pkg/server/grpc_client_test.go
@@ -140,7 +140,6 @@ func TestToGraph(t *testing.T) {
 	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},
-		Metadata: map[string]string{topology.KeyPlugin: topology.TopologyBlock},
 	}
 
 	require.Equal(t, root, toGraph(&pb.TopologyResponse{Instances: instances}, cis, topology.TopologyBlock))

diff --git a/pkg/translate/output.go b/pkg/translate/output.go
@@ -26,7 +26,6 @@ import (
 
 	"k8s.io/klog/v2"
 
-	"github.com/NVIDIA/topograph/pkg/engines"
 	"github.com/NVIDIA/topograph/pkg/metrics"
 	"github.com/NVIDIA/topograph/pkg/topology"
 )
@@ -39,7 +38,10 @@ func ToGraph(wr io.Writer, root *topology.Vertex) error {
 			return toTreeTopology(wr, root)
 		}
 	}
-	return fmt.Errorf("topology metadata not set in root node")
+	if root.Vertices[topology.TopologyBlock] != nil {
+		return toBlockTopology(wr, root)
+	}
+	return toTreeTopology(wr, root)
 }
 
 func printBlock(wr io.Writer, block *topology.Vertex, domainVisited map[string]int) error {
@@ -383,9 +385,6 @@ func GetTreeTestSet(testForLongLabelName bool) (*topology.Vertex, map[string]str
 	}
 	root := &topology.Vertex{
 		Vertices: map[string]*topology.Vertex{"S1": sw1},
-		Metadata: map[string]string{
-			topology.KeyPlugin: topology.TopologyTree,
-		},
 	}
 
 	return root, instance2node
@@ -477,8 +476,6 @@ func GetBlockWithMultiIBTestSet() (*topology.Vertex, map[string]string) {
 	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: "3",
 		},
 	}
@@ -531,8 +528,6 @@ func GetBlockWithIBTestSet() (*topology.Vertex, map[string]string) {
 	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: "3",
 		},
 	}
@@ -605,8 +600,6 @@ func GetBlockWithDFSIBTestSet() (*topology.Vertex, map[string]string) {
 	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",
 		},
 	}
@@ -643,8 +636,6 @@ func GetBlockTestSet() (*topology.Vertex, map[string]string) {
 	root := &topology.Vertex{
 		Vertices: map[string]*topology.Vertex{topology.TopologyBlock: blockRoot},
 		Metadata: map[string]string{
-			topology.KeyEngine:     engines.EngineSLURM,
-			topology.KeyPlugin:     topology.TopologyBlock,
 			topology.KeyBlockSizes: "3",
 		},
 	}

diff --git a/pkg/translate/output_test.go b/pkg/translate/output_test.go
@@ -63,7 +63,6 @@ SwitchName=switch.1.2 Nodes=node-2
 
 func TestToTreeTopology(t *testing.T) {
 	v, _ := GetTreeTestSet(false)
-	require.Equal(t, v.Metadata[topology.KeyPlugin], topology.TopologyTree)
 	buf := &bytes.Buffer{}
 	err := ToGraph(buf, v)
 	require.NoError(t, err)
@@ -72,7 +71,6 @@ func TestToTreeTopology(t *testing.T) {
 
 func TestToBlockTopology(t *testing.T) {
 	v, _ := GetBlockTestSet()
-	require.Equal(t, v.Metadata[topology.KeyPlugin], topology.TopologyBlock)
 	buf := &bytes.Buffer{}
 	err := ToGraph(buf, v)
 	require.NoError(t, err)
@@ -81,7 +79,6 @@ func TestToBlockTopology(t *testing.T) {
 
 func TestToBlockMultiIBTopology(t *testing.T) {
 	v, _ := GetBlockWithMultiIBTestSet()
-	require.Equal(t, v.Metadata[topology.KeyPlugin], topology.TopologyBlock)
 	buf := &bytes.Buffer{}
 	err := ToGraph(buf, v)
 	require.NoError(t, err)
@@ -95,7 +92,6 @@ func TestToBlockMultiIBTopology(t *testing.T) {
 
 func TestToBlockIBTopology(t *testing.T) {
 	v, _ := GetBlockWithIBTestSet()
-	require.Equal(t, v.Metadata[topology.KeyPlugin], topology.TopologyBlock)
 	buf := &bytes.Buffer{}
 	err := ToGraph(buf, v)
 	require.NoError(t, err)
@@ -109,7 +105,6 @@ 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)
@@ -163,12 +158,8 @@ func TestToSlurmNameShortener(t *testing.T) {
 				},
 			},
 		},
-		Metadata: map[string]string{
-			topology.KeyPlugin: topology.TopologyTree,
-		},
 	}
 
-	require.Equal(t, v.Metadata[topology.KeyPlugin], topology.TopologyTree)
 	buf := &bytes.Buffer{}
 	err := ToGraph(buf, v)
 	require.NoError(t, err)