implement cluset functionality (#68)

Signed-off-by: Dmitry Shmulevich <[email protected]>

Author: dmitsh

internal/cluset/cluset.go

@@ -0,0 +1,134 @@
+/*
+ * Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package cluset
+
+import (
+	"fmt"
+	"regexp"
+	"sort"
+	"strconv"
+	"strings"
+)
+
+var compactRegexp, expandRegexp *regexp.Regexp
+
+func init() {
+	compactRegexp = regexp.MustCompile(`^(.*?)(0*\d+)$`)
+	expandRegexp = regexp.MustCompile(`^(.*)\[(.*?)\]$`)
+}
+
+// Compact compresses a list of nodes into a compact range format
+func Compact(nodes []string) []string {
+	groups := make(map[string][]string)
+
+	// Parse node names into groups
+	for _, node := range nodes {
+		matches := compactRegexp.FindStringSubmatch(node)
+		if matches != nil {
+			prefix, numStr := matches[1], matches[2]
+			groups[prefix] = append(groups[prefix], numStr)
+		} else {
+			groups[node] = nil
+		}
+	}
+
+	// Sort prefixes for consistent output
+	prefixes := make([]string, 0, len(groups))
+	for prefix := range groups {
+		prefixes = append(prefixes, prefix)
+	}
+	sort.Strings(prefixes)
+
+	// Format compressed ranges
+	var result []string
+	for _, prefix := range prefixes {
+		numbers := groups[prefix]
+		sort.Slice(numbers, func(i, j int) bool {
+			return atoi(numbers[i]) < atoi(numbers[j])
+		})
+		compressed := compressRange(numbers)
+		result = append(result, fmt.Sprintf("%s%s", prefix, compressed))
+	}
+
+	return result
+}
+
+// Expand decompresses a list of compacted node names back to individual entries
+func Expand(compressed []string) []string {
+	var result []string
+
+	for _, entry := range compressed {
+		matches := expandRegexp.FindStringSubmatch(entry)
+		if matches != nil {
+			prefix, ranges := matches[1], matches[2]
+			rangesParts := strings.Split(ranges, ",")
+			for _, part := range rangesParts {
+				if strings.Contains(part, "-") {
+					bounds := strings.Split(part, "-")
+					width := len(bounds[0]) // Preserve leading zeros
+					for i := atoi(bounds[0]); i <= atoi(bounds[1]); i++ {
+						result = append(result, fmt.Sprintf("%s%0*d", prefix, width, i))
+					}
+				} else {
+					result = append(result, fmt.Sprintf("%s%s", prefix, part))
+				}
+			}
+		} else {
+			result = append(result, entry)
+		}
+	}
+	return result
+}
+
+// compressRange converts a list of numbers into a compact range format
+func compressRange(numbers []string) string {
+	switch len(numbers) {
+	case 0:
+		return ""
+	case 1:
+		return numbers[0]
+	default:
+		var parts []string
+		start, end := numbers[0], numbers[0]
+
+		for i := 1; i < len(numbers); i++ {
+			if atoi(numbers[i]) == atoi(end)+1 {
+				end = numbers[i]
+			} else {
+				parts = append(parts, formatRange(start, end))
+				start, end = numbers[i], numbers[i]
+			}
+		}
+		parts = append(parts, formatRange(start, end))
+
+		return "[" + strings.Join(parts, ",") + "]"
+	}
+}
+
+// formatRange formats a range into a string
+func formatRange(start, end string) string {
+	if start == end {
+		return start
+	}
+	return fmt.Sprintf("%s-%s", start, end)
+}
+
+// atoi converts a zero-padded string number to an integer
+func atoi(s string) int {
+	num, _ := strconv.Atoi(s)
+	return num
+}

internal/cluset/cluset_test.go

@@ -0,0 +1,53 @@
+package cluset
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestCompactExpand(t *testing.T) {
+	testCases := []struct {
+		name                string
+		expanded, compacted []string
+	}{
+		{
+			name: "Case 1: empty list",
+		},
+		{
+			name:      "Case 2: ranges",
+			expanded:  []string{"abc0507", "abc0509", "abc0482", "abc0483", "abc0508", "abc0484"},
+			compacted: []string{"abc[0482-0484,0507-0509]"},
+		},
+		{
+			name:      "Case 3: singles",
+			expanded:  []string{"abc0507", "abc0509", "xyz0482"},
+			compacted: []string{"abc[0507,0509]", "xyz0482"},
+		},
+		{
+			name:      "Case 4: mix1",
+			expanded:  []string{"abc0507", "abc0509", "def", "abc0482", "abc0508"},
+			compacted: []string{"abc[0482,0507-0509]", "def"},
+		},
+		{
+			name:      "Case 5: mix2",
+			expanded:  []string{"abc0507", "abc0509", "abc0508", "abc0482", "xyz8", "xyz9", "xyz10"},
+			compacted: []string{"abc[0482,0507-0509]", "xyz[8-10]"},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			require.Equal(t, tc.compacted, Compact(tc.expanded))
+			require.Equal(t, toMap(tc.expanded), toMap(Expand(tc.compacted)))
+		})
+	}
+}
+
+func toMap(arr []string) map[string]struct{} {
+	m := make(map[string]struct{})
+	for _, str := range arr {
+		m[str] = struct{}{}
+	}
+	return m
+}

pkg/server/http_server_test.go

@@ -90,7 +90,7 @@ func TestServer(t *testing.T) {
 }
 `,
 			expected: `SwitchName=S1 Switches=S[2-3]
-SwitchName=S2 Nodes=Node[201-202],Node205
+SwitchName=S2 Nodes=Node[201-202,205]
 SwitchName=S3 Nodes=Node[304-306]
 `,
 		},
@@ -155,29 +155,29 @@ SwitchName=sw14 Nodes=n14-[1-2]
 }
 `,
 			expected: `# block001=nvl-1-1
-BlockName=block001 Nodes=n1-1-0[1-8]
+BlockName=block001 Nodes=n1-1-[01-08]
 # block002=nvl-1-2
-BlockName=block002 Nodes=n1-2-0[1-8]
+BlockName=block002 Nodes=n1-2-[01-08]
 # block003=nvl-2-1
-BlockName=block003 Nodes=n2-1-0[1-8]
+BlockName=block003 Nodes=n2-1-[01-08]
 # block004=nvl-2-2
-BlockName=block004 Nodes=n2-2-0[1-8]
+BlockName=block004 Nodes=n2-2-[01-08]
 # block005=nvl-3-1
-BlockName=block005 Nodes=n3-1-0[1-8]
+BlockName=block005 Nodes=n3-1-[01-08]
 # block006=nvl-3-2
-BlockName=block006 Nodes=n3-2-0[1-8]
+BlockName=block006 Nodes=n3-2-[01-08]
 # block007=nvl-4-1
-BlockName=block007 Nodes=n4-1-0[1-8]
+BlockName=block007 Nodes=n4-1-[01-08]
 # block008=nvl-4-2
-BlockName=block008 Nodes=n4-2-0[1-8]
+BlockName=block008 Nodes=n4-2-[01-08]
 # block009=nvl-5-1
-BlockName=block009 Nodes=n5-1-0[1-8]
+BlockName=block009 Nodes=n5-1-[01-08]
 # block010=nvl-5-2
-BlockName=block010 Nodes=n5-2-0[1-8]
+BlockName=block010 Nodes=n5-2-[01-08]
 # block011=nvl-6-1
-BlockName=block011 Nodes=n6-1-0[1-8]
+BlockName=block011 Nodes=n6-1-[01-08]
 # block012=nvl-6-2
-BlockName=block012 Nodes=n6-2-0[1-8]
+BlockName=block012 Nodes=n6-2-[01-08]
 BlockSizes=8,16,32
 `,
 		},

pkg/translate/output.go

@@ -26,6 +26,7 @@ import (
 
 	"k8s.io/klog/v2"
 
+	"github.com/NVIDIA/topograph/internal/cluset"
 	"github.com/NVIDIA/topograph/pkg/metrics"
 	"github.com/NVIDIA/topograph/pkg/topology"
 )
@@ -54,7 +55,7 @@ func printBlock(wr io.Writer, block *topology.Vertex, domainVisited map[string]i
 		if len(block.Name) != 0 {
 			comment = fmt.Sprintf("# %s=%s\n", block.ID, block.Name)
 		}
-		_, err := wr.Write([]byte(fmt.Sprintf("%sBlockName=%s Nodes=%s\n", comment, block.ID, strings.Join(compress(nodes), ","))))
+		_, err := wr.Write([]byte(fmt.Sprintf("%sBlockName=%s Nodes=%s\n", comment, block.ID, strings.Join(cluset.Compact(nodes), ","))))
 		if err != nil {
 			return err
 		}
@@ -266,7 +267,7 @@ func toTreeTopology(wr io.Writer, root *topology.Vertex) error {
 			comment = ""
 			switchName = sw
 		}
-		_, err := wr.Write([]byte(fmt.Sprintf("%sSwitchName=%s Nodes=%s\n", comment, switchName, strings.Join(compress(nodes), ","))))
+		_, err := wr.Write([]byte(fmt.Sprintf("%sSwitchName=%s Nodes=%s\n", comment, switchName, strings.Join(cluset.Compact(nodes), ","))))
 		if err != nil {
 			return err
 		}
@@ -295,103 +296,14 @@ func writeSwitch(wr io.Writer, v *topology.Vertex) error {
 	} else {
 		comment = fmt.Sprintf("# %s=%s\n", v.Name, v.ID)
 	}
-	_, err := wr.Write([]byte(fmt.Sprintf("%sSwitchName=%s Switches=%s\n", comment, v.Name, strings.Join(compress(arr), ","))))
+	_, err := wr.Write([]byte(fmt.Sprintf("%sSwitchName=%s Switches=%s\n", comment, v.Name, strings.Join(cluset.Compact(arr), ","))))
 	if err != nil {
 		return err
 	}
 
 	return nil
 }
 
-// compress finds contiguos numerical suffixes in names and presents then as ranges.
-// example: ["eos0507", "eos0509", "eos0508"] -> ["eos0[507-509"]
-func compress(input []string) []string {
-	ret := []string{}
-	keys := []string{}
-	m := make(map[string][]int) // map of prefix : array of numerical suffix
-	for _, name := range input {
-		if prefix, suffix := split(name); len(suffix) == 0 {
-			ret = append(ret, name)
-		} else {
-			num, _ := strconv.Atoi(suffix)
-			if arr, ok := m[prefix]; !ok {
-				m[prefix] = []int{num}
-			} else {
-				m[prefix] = append(arr, num)
-			}
-		}
-	}
-
-	// we sort the prefix to get consistent output for tests
-	for prefix := range m {
-		keys = append(keys, prefix)
-	}
-	sort.Strings(keys)
-
-	for _, prefix := range keys {
-		arr := m[prefix]
-		sort.Ints(arr)
-		var start, end int
-		for i, num := range arr {
-			if i == 0 {
-				start = num
-				end = num
-			} else if num == end+1 {
-				end = num
-			} else if start == end {
-				ret = append(ret, fmt.Sprintf("%s%d", prefix, start))
-				start = num
-				end = num
-			} else {
-				ret = append(ret, fmt.Sprintf("%s[%d-%d]", prefix, start, end))
-				start = num
-				end = num
-			}
-		}
-		if start == end {
-			ret = append(ret, fmt.Sprintf("%s%d", prefix, end))
-		} else {
-			ret = append(ret, fmt.Sprintf("%s[%d-%d]", prefix, start, end))
-		}
-	}
-
-	return ret
-}
-
-// split divides a string into a prefix and a numerical suffix
-func split(input string) (string, string) {
-	n := len(input)
-	if n == 0 {
-		return input, ""
-	}
-
-	// find numerical suffix
-	i := n - 1
-	for i >= 0 {
-		if input[i] >= '0' && input[i] <= '9' {
-			i--
-		} else {
-			break
-		}
-	}
-	i++
-
-	// ignore leading zeros
-	for i < n {
-		if input[i] == '0' {
-			i++
-		} else {
-			break
-		}
-	}
-
-	if i == n { // no suffix
-		return input, ""
-	}
-
-	return input[:i], input[i:]
-}
-
 func GetTreeTestSet(testForLongLabelName bool) (*topology.Vertex, map[string]string) {
 	//
 	//        S1