Skip to content

Commit 37dad45

Browse files
DarrylWongDarrylWong
committed
roachtest: add unit test for arch selection
This test checks that our random arch selection produces architectures within a margin of error of the expected distribution.
1 parent d8a3da7 commit 37dad45

File tree

4 files changed

+133
-10
lines changed

4 files changed

+133
-10
lines changed

‎pkg/cmd/roachtest/cluster.go‎

Lines changed: 10 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -3189,7 +3189,7 @@ func archForTest(ctx context.Context, l *logger.Logger, testSpec registry.TestSp
31893189
validArchs = testSpec.Cluster.CompatibleArchs
31903190
}
31913191

3192-
arch := randomArch(ctx, l, validArchs, prng)
3192+
arch := randomArch(ctx, l, validArchs, prng, roachtestflags.ARM64Probability, roachtestflags.FIPSProbability)
31933193

31943194
if roachtestflags.Cloud == spec.GCE && arch == vm.ArchARM64 {
31953195
// N.B. T2A support is rather limited, both in terms of supported
@@ -3211,19 +3211,22 @@ func archForTest(ctx context.Context, l *logger.Logger, testSpec registry.TestSp
32113211
}
32123212

32133213
// randomArch chooses a random architecture, respecting the set of valid architectures
3214-
// specified by the test as well as the global architecture probability flags for the
3215-
// entire run.
3214+
// specified by the test as well as the provided architecture probability flags.
32163215
func randomArch(
3217-
ctx context.Context, l *logger.Logger, validArchs spec.ArchSet, prng *rand.Rand,
3216+
ctx context.Context,
3217+
l *logger.Logger,
3218+
validArchs spec.ArchSet,
3219+
prng *rand.Rand,
3220+
arm64Probability, fipsProbability float64,
32183221
) vm.CPUArch {
32193222
baseProbabilities := map[vm.CPUArch]float64{
3220-
vm.ArchAMD64: (1.0 - roachtestflags.ARM64Probability) * (1.0 - roachtestflags.FIPSProbability),
3221-
vm.ArchARM64: roachtestflags.ARM64Probability,
3223+
vm.ArchAMD64: (1.0 - arm64Probability) * (1.0 - fipsProbability),
3224+
vm.ArchARM64: arm64Probability,
32223225
// N.B. FIPS is only supported on 'amd64' at this time:
32233226
// FIPS is taken with probability
32243227
// (1 - arm64Probability) * fipsProbability
32253228
// which is P(fips & amd64)
3226-
vm.ArchFIPS: (1.0 - roachtestflags.ARM64Probability) * roachtestflags.FIPSProbability,
3229+
vm.ArchFIPS: (1.0 - arm64Probability) * fipsProbability,
32273230
}
32283231

32293232
// Calculate total weight for valid architectures only.

‎pkg/cmd/roachtest/cluster_test.go‎

Lines changed: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -909,3 +909,119 @@ func TestVerifyLibraries(t *testing.T) {
909909
})
910910
}
911911
}
912+
913+
func TestRandomArchProbabilities(t *testing.T) {
914+
ctx := context.Background()
915+
916+
tests := []struct {
917+
validArchs spec.ArchSet
918+
arm64Probability float64
919+
fipsProbability float64
920+
expectedDistribution map[vm.CPUArch]float64
921+
}{
922+
{
923+
validArchs: spec.AllArchs,
924+
arm64Probability: 0.3,
925+
fipsProbability: 0.2,
926+
expectedDistribution: map[vm.CPUArch]float64{
927+
vm.ArchAMD64: 0.56,
928+
vm.ArchARM64: 0.3,
929+
vm.ArchFIPS: 0.14,
930+
},
931+
},
932+
{
933+
validArchs: spec.AllExceptFIPS,
934+
arm64Probability: 0.4,
935+
fipsProbability: 0.1,
936+
expectedDistribution: map[vm.CPUArch]float64{
937+
vm.ArchAMD64: 0.57447,
938+
vm.ArchARM64: 0.42553,
939+
},
940+
},
941+
{
942+
validArchs: spec.OnlyAMD64,
943+
arm64Probability: 0.5,
944+
fipsProbability: 0.3,
945+
expectedDistribution: map[vm.CPUArch]float64{
946+
vm.ArchAMD64: 1.0, // Only valid architecture
947+
},
948+
},
949+
{
950+
validArchs: spec.OnlyARM64,
951+
arm64Probability: 0.5,
952+
fipsProbability: 0.3,
953+
expectedDistribution: map[vm.CPUArch]float64{
954+
vm.ArchARM64: 1.0, // Only valid architecture
955+
},
956+
},
957+
{
958+
validArchs: spec.OnlyFIPS,
959+
arm64Probability: 0.2,
960+
fipsProbability: 0.0,
961+
expectedDistribution: map[vm.CPUArch]float64{
962+
vm.ArchAMD64: 1.0, // Should fall back to AMD64
963+
},
964+
},
965+
{
966+
validArchs: spec.AllExceptFIPS,
967+
arm64Probability: 0.0,
968+
fipsProbability: 1.0,
969+
expectedDistribution: map[vm.CPUArch]float64{
970+
vm.ArchAMD64: 1.0, // Should fall back to AMD64
971+
},
972+
},
973+
{
974+
validArchs: spec.AllExceptFIPS,
975+
arm64Probability: 0.5,
976+
fipsProbability: 1.0,
977+
expectedDistribution: map[vm.CPUArch]float64{
978+
vm.ArchARM64: 1.0,
979+
},
980+
},
981+
{
982+
validArchs: spec.Archs(spec.ArchAMD64, spec.ArchFIPS),
983+
arm64Probability: 0.3,
984+
fipsProbability: 0.4,
985+
expectedDistribution: map[vm.CPUArch]float64{
986+
vm.ArchAMD64: 0.6,
987+
vm.ArchFIPS: 0.4,
988+
},
989+
},
990+
{
991+
validArchs: spec.Archs(spec.ArchARM64, spec.ArchFIPS),
992+
arm64Probability: 0.6,
993+
fipsProbability: 0.2,
994+
expectedDistribution: map[vm.CPUArch]float64{
995+
vm.ArchARM64: 0.88235,
996+
vm.ArchFIPS: 0.11765,
997+
},
998+
},
999+
}
1000+
1001+
// Since this is a statistical test, we want to use a fixed seed to avoid flakes,
1002+
// i.e. a 99% confidence interval would be expected to fail when stressed 100 times.
1003+
//
1004+
// We can run this manually with a random seed for more confidence in our distribution:
1005+
// prng, _ := randutil.NewTestRand()
1006+
prng := rand.New(rand.NewSource(12345))
1007+
for _, test := range tests {
1008+
t.Run(fmt.Sprintf("%s/arm=%d%%/fips=%d%%", test.validArchs, int(test.arm64Probability*100), int(test.fipsProbability*100)), func(t *testing.T) {
1009+
const numSamples = 10000
1010+
counts := make(map[vm.CPUArch]int)
1011+
1012+
// Generate samples
1013+
for i := 0; i < numSamples; i++ {
1014+
arch := randomArch(ctx, nilLogger(), test.validArchs, prng, test.arm64Probability, test.fipsProbability)
1015+
counts[arch]++
1016+
}
1017+
1018+
for expectedArch, expectedProb := range test.expectedDistribution {
1019+
actualCount := float64(counts[expectedArch])
1020+
actualProb := actualCount / float64(numSamples)
1021+
1022+
tolerance := 0.02
1023+
require.InDelta(t, expectedProb, actualProb, tolerance)
1024+
}
1025+
})
1026+
}
1027+
}

‎pkg/cmd/roachtest/spec/cluster_spec.go‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,11 @@ var AllArchs = Archs(ArchAMD64, ArchARM64, ArchFIPS)
5757
// OnlyAMD64 contains only the AMD64 architecture.
5858
var OnlyAMD64 = Archs(ArchAMD64)
5959

60+
var OnlyARM64 = Archs(ArchARM64)
61+
62+
// OnlyFIPS contains only the FIPS architecture.
63+
var OnlyFIPS = Archs(ArchFIPS)
64+
6065
var AllExceptFIPS = AllArchs.remove(ArchFIPS)
6166

6267
// Archs creates an ArchSet for the given architectures.

‎pkg/cmd/roachtest/test_registry_test.go‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,6 @@ import (
1010

1111
"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/registry"
1212
"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/spec"
13-
"github.com/cockroachdb/cockroach/pkg/roachprod/vm"
1413
"github.com/prometheus/client_golang/prometheus"
1514
"github.com/stretchr/testify/require"
1615
)
@@ -37,10 +36,10 @@ func TestMakeTestRegistry(t *testing.T) {
3736
require.EqualValues(t, 4, s.CPUs)
3837
require.True(t, s.TerminateOnMigration)
3938

40-
s = r.MakeClusterSpec(10, spec.CPU(16), spec.Arch(registry.OnlyARM64))
39+
s = r.MakeClusterSpec(10, spec.CPU(16), spec.Arch(spec.OnlyARM64))
4140
require.EqualValues(t, 10, s.NodeCount)
4241
require.EqualValues(t, 16, s.CPUs)
43-
require.EqualValues(t, vm.ArchARM64, s.Arch)
42+
require.EqualValues(t, spec.OnlyARM64, s.CompatibleArchs)
4443
}
4544

4645
// TestPrometheusMetricParser tests that the registry.PromSub()

0 commit comments

Comments
 (0)