Merge #153653

153653: asim: allow configuring heterogeneous CPU r=tbg a=tbg

This PR lays the groundwork to allow for `gen_cluster nodes=3 node_cpu_rate_capacity(A,B,C)` where A, B,
C aren't identical. If only one value is specified, it's applied to all nodes. In other words,
existing tests still work without change.

Then it adds the `hetergeneous_cpu` test - the simple scenario outlined in #153516. We see that neither SMA nor MMA relieve the two 8vcpu nodes from their high CPU utilization, even though the 16vcpu node is below 50% utilization.

Closes #153516.

Epic: CRDB-49117

Co-authored-by: Tobias Grieger <[email protected]>

Author: craig[bot]

pkg/kv/kvserver/asim/gen/generator.go

@@ -229,14 +229,14 @@ type BasicCluster struct {
 	StoreByteCapacity   int64
 	Region              []string
 	NodesPerRegion      []int
-	NodeCPURateCapacity int64
+	NodeCPURateCapacity state.NodeCPURateCapacities
 }
 
 func (bc BasicCluster) String() string {
 	var b strings.Builder
 	_, _ = fmt.Fprintf(&b,
-		"[nodes: %d, stores_per_node:%d, store_disk_capacity: %dGiB, node_capacity: %dcpu-sec/sec",
-		bc.Nodes, bc.StoresPerNode, bc.StoreByteCapacity>>30, bc.NodeCPURateCapacity/time.Second.Nanoseconds())
+		"[nodes: %d, stores_per_node:%d, store_disk_capacity: %dGiB, node_capacity: %s",
+		bc.Nodes, bc.StoresPerNode, bc.StoreByteCapacity>>30, bc.NodeCPURateCapacity)
 	if len(bc.Region) != 0 {
 		_, _ = fmt.Fprintf(&b, ", region: %v, nodes_per_region: %v", bc.Region, bc.NodesPerRegion)
 	}

pkg/kv/kvserver/asim/state/BUILD.bazel

@@ -11,6 +11,7 @@ go_library(
         "load.go",
         "new_state.go",
         "new_state_test_helper.go",
+        "node_cpu_rate_capacities.go",
         "parser_replica_placement.go",
         "split_decider.go",
         "state.go",
@@ -59,6 +60,7 @@ go_test(
         "change_test.go",
         "config_loader_test.go",
         "liveness_test.go",
+        "node_cpu_rate_capacities_test.go",
         "parser_replica_placement_test.go",
         "split_decider_test.go",
         "state_test.go",

pkg/kv/kvserver/asim/state/config_loader.go

@@ -26,8 +26,8 @@ var AllClusterOptions = [...]string{"single_region", "single_region_multi_store"
 // SingleRegionConfig is a simple cluster config with a single region and 3
 // zones, all have the same number of nodes.
 var SingleRegionConfig = ClusterInfo{
-	NodeCPURateCapacityNanos: config.DefaultNodeCPURateCapacityNanos, // 8vpucs
-	StoreDiskCapacityBytes:   config.DefaultStoreDiskCapacityBytes,   // 1024 GiB
+	NodeCPURateCapacityNanos: []uint64{config.DefaultNodeCPURateCapacityNanos}, // 8vpucs
+	StoreDiskCapacityBytes:   config.DefaultStoreDiskCapacityBytes,             // 1024 GiB
 	Regions: []Region{
 		{
 			Name: "US",
@@ -43,8 +43,8 @@ var SingleRegionConfig = ClusterInfo{
 // SingleRegionMultiStoreConfig is a simple cluster config with a single region
 // and 3 zones, all zones have 1 node and 5 stores per node.
 var SingleRegionMultiStoreConfig = ClusterInfo{
-	NodeCPURateCapacityNanos: config.DefaultNodeCPURateCapacityNanos, // 8 vcpus
-	StoreDiskCapacityBytes:   config.DefaultStoreDiskCapacityBytes,   // 1024 GiB
+	NodeCPURateCapacityNanos: []uint64{config.DefaultNodeCPURateCapacityNanos}, // 8 vcpus
+	StoreDiskCapacityBytes:   config.DefaultStoreDiskCapacityBytes,             // 1024 GiB
 	Regions: []Region{
 		{
 			Name: "US",
@@ -59,8 +59,8 @@ var SingleRegionMultiStoreConfig = ClusterInfo{
 
 // MultiRegionConfig is a perfectly balanced cluster config with 3 regions.
 var MultiRegionConfig = ClusterInfo{
-	NodeCPURateCapacityNanos: config.DoubleDefaultNodeCPURateCapacityNanos, // 16 vcpus
-	StoreDiskCapacityBytes:   config.DoubleDefaultStoreDiskCapacityBytes,   // 2048 GiB
+	NodeCPURateCapacityNanos: []uint64{config.DoubleDefaultNodeCPURateCapacityNanos}, // 16 vcpus
+	StoreDiskCapacityBytes:   config.DoubleDefaultStoreDiskCapacityBytes,             // 2048 GiB
 	Regions: []Region{
 		{
 			Name: "US_East",
@@ -91,8 +91,8 @@ var MultiRegionConfig = ClusterInfo{
 
 // ComplexConfig is an imbalanced multi-region cluster config.
 var ComplexConfig = ClusterInfo{
-	NodeCPURateCapacityNanos: config.DoubleDefaultNodeCPURateCapacityNanos, // 16 vcpus
-	StoreDiskCapacityBytes:   config.DoubleDefaultStoreDiskCapacityBytes,   // 2048 GiB
+	NodeCPURateCapacityNanos: []uint64{config.DoubleDefaultNodeCPURateCapacityNanos}, // 16 vcpus
+	StoreDiskCapacityBytes:   config.DoubleDefaultStoreDiskCapacityBytes,             // 2048 GiB
 	Regions: []Region{
 		{
 			Name: "US_East",
@@ -282,7 +282,7 @@ type Region struct {
 type ClusterInfo struct {
 	Regions                  []Region
 	StoreDiskCapacityBytes   int64
-	NodeCPURateCapacityNanos int64
+	NodeCPURateCapacityNanos NodeCPURateCapacities
 }
 
 func (c ClusterInfo) String() (s string) {
@@ -300,7 +300,7 @@ func (c ClusterInfo) String() (s string) {
 		}
 		buf.WriteString("]\n")
 	}
-	buf.WriteString(fmt.Sprintf("store_disk_capacity=%d bytes, node_cpu_rate_capacity=%d cpu-ns/sec",
+	buf.WriteString(fmt.Sprintf("store_disk_capacity=%d bytes, node_cpu_rate_capacity=%s",
 		c.StoreDiskCapacityBytes, c.NodeCPURateCapacityNanos))
 	return buf.String()
 }
@@ -361,6 +361,7 @@ func LoadClusterInfo(c ClusterInfo, settings *config.SimulationSettings) State {
 	s := newState(settings)
 	// A new state has a single range - add the replica load for that range.
 	s.clusterinfo = c
+	var nodeIdx int
 	for _, r := range c.Regions {
 		regionTier := roachpb.Tier{
 			Key:   "region",
@@ -375,7 +376,17 @@ func LoadClusterInfo(c ClusterInfo, settings *config.SimulationSettings) State {
 				Tiers: []roachpb.Tier{regionTier, zoneTier},
 			}
 			for i := 0; i < z.NodeCount; i++ {
-				node := s.AddNode(c.NodeCPURateCapacityNanos, locality)
+				var cpuCap uint64
+				if len(c.NodeCPURateCapacityNanos) == 1 {
+					// As a special case, if only one CPU is specified, use it for all nodes.
+					cpuCap = c.NodeCPURateCapacityNanos[0]
+				} else {
+					// Otherwise, expect a CPU capacity for each node. Crash if this is
+					// not the case.
+					cpuCap = c.NodeCPURateCapacityNanos[nodeIdx]
+				}
+				nodeIdx += 1
+				node := s.AddNode(int64(cpuCap), locality)
 				storesRequired := z.StoresPerNode
 				if storesRequired < 1 {
 					panic(fmt.Sprintf("storesPerNode cannot be less than one but found %v", storesRequired))

pkg/kv/kvserver/asim/state/new_state.go

@@ -12,6 +12,7 @@ import (
 	"sort"
 	"strings"
 
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/asim/config"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 )
 
@@ -253,7 +254,9 @@ func RangesInfoWithDistribution(
 func ClusterInfoWithDistribution(
 	nodeCount int, storesPerNode int, regions []string, regionNodeWeights []float64,
 ) ClusterInfo {
-	ret := ClusterInfo{}
+	ret := ClusterInfo{
+		NodeCPURateCapacityNanos: []uint64{config.DefaultNodeCPURateCapacityNanos},
+	}
 
 	ret.Regions = make([]Region, len(regions))
 	availableNodes := nodeCount

pkg/kv/kvserver/asim/state/new_state_test_helper.go

@@ -27,7 +27,7 @@ func NewStateWithDistribution(
 	// Currently multi-store is not tested for correctness. Default to a single
 	// store per node.
 	clusterInfo := ClusterInfoWithStoreCount(numNodes, 1 /* storesPerNode */)
-	clusterInfo.NodeCPURateCapacityNanos = config.DefaultNodeCPURateCapacityNanos
+	clusterInfo.NodeCPURateCapacityNanos = []uint64{config.DefaultNodeCPURateCapacityNanos}
 	s := LoadClusterInfo(clusterInfo, settings)
 
 	stores := make([]StoreID, numNodes)

pkg/kv/kvserver/asim/state/node_cpu_rate_capacities.go

@@ -0,0 +1,40 @@
+// Copyright 2025 The Cockroach Authors.
+//
+// Use of this software is governed by the CockroachDB Software License
+// included in the /LICENSE file.
+
+package state
+
+import (
+	"fmt"
+	"strings"
+)
+
+type NodeCPURateCapacities []uint64
+
+func (sl NodeCPURateCapacities) String() string {
+	if len(sl) == 0 {
+		return "no cpus"
+	}
+	var buf strings.Builder
+	if len(sl) > 1 {
+		buf.WriteString("(")
+	}
+	for i, cc := range sl {
+		if i > 0 {
+			buf.WriteString(" ")
+		}
+
+		// Convert nanoseconds to seconds with up to 2 decimal places
+		if cc%1e9 == 0 {
+			fmt.Fprintf(&buf, "%d", cc/1e9)
+		} else {
+			fmt.Fprintf(&buf, "%.2f", float64(cc)/1e9)
+		}
+	}
+	if len(sl) > 1 {
+		buf.WriteString(")")
+	}
+	buf.WriteString(" cpu-sec/sec")
+	return buf.String()
+}

pkg/kv/kvserver/asim/state/node_cpu_rate_capacities_test.go

@@ -0,0 +1,59 @@
+// Copyright 2025 The Cockroach Authors.
+//
+// Use of this software is governed by the CockroachDB Software License
+// included in the /LICENSE file.
+
+package state
+
+import (
+	"path/filepath"
+	"testing"
+	"time"
+
+	"github.com/cockroachdb/cockroach/pkg/testutils/echotest"
+)
+
+func TestNodeCPURateCapacities_String(t *testing.T) {
+	testCases := []struct {
+		name     string
+		capacity NodeCPURateCapacities
+	}{
+		{
+			name:     "empty",
+			capacity: NodeCPURateCapacities{},
+		},
+		{
+			name:     "single_capacity",
+			capacity: NodeCPURateCapacities{uint64(time.Second.Nanoseconds())},
+		},
+		{
+			name: "multiple_capacities",
+			capacity: NodeCPURateCapacities{
+				uint64(2 * time.Second.Nanoseconds()),
+				uint64(4 * time.Second.Nanoseconds()),
+				uint64(8 * time.Second.Nanoseconds()),
+			},
+		},
+		{
+			name: "fractional_seconds",
+			capacity: NodeCPURateCapacities{
+				uint64(500 * time.Millisecond.Nanoseconds()),
+				uint64(1500 * time.Millisecond.Nanoseconds()),
+			},
+		},
+		{
+			name: "mixed_exact_and_fractional",
+			capacity: NodeCPURateCapacities{
+				uint64(time.Second.Nanoseconds()),            // 1.00
+				uint64(500 * time.Millisecond.Nanoseconds()), // 0.50
+				uint64(2 * time.Second.Nanoseconds()),        // 2.00
+			},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			echotest.Require(t, tc.capacity.String(), filepath.Join("testdata", "NodeCPURateCapacities_String", tc.name+".txt"))
+		})
+	}
+}

pkg/kv/kvserver/asim/state/testdata/NodeCPURateCapacities_String/empty.txt

@@ -0,0 +1,3 @@
+echo
+----
+no cpus

pkg/kv/kvserver/asim/state/testdata/NodeCPURateCapacities_String/fractional_seconds.txt

@@ -0,0 +1,3 @@
+echo
+----
+(0.50 1.50) cpu-sec/sec

pkg/kv/kvserver/asim/state/testdata/NodeCPURateCapacities_String/mixed_exact_and_fractional.txt

@@ -0,0 +1,3 @@
+echo
+----
+(1 0.50 2) cpu-sec/sec