[tmpnet] Enable exclusive scheduling (#3988)

Author: maru-ava

.envrc

@@ -16,5 +16,8 @@ export AVALANCHEGO_PATH="${AVALANCHEGO_PATH:-$PWD/bin/avalanchego}"
 mkdir -p $PWD/build/plugins                                       # avalanchego will FATAL if the directory does not exist
 export AVAGO_PLUGIN_DIR="${AVAGO_PLUGIN_DIR:-$PWD/build/plugins}" # Use an existing value if set
 
+# Default to tmpnetctl targeting the last deployed tmpnet network
+export TMPNET_NETWORK_DIR="${TMPNET_NETWORK_DIR:-${HOME}/.tmpnet/networks/latest}"
+
 # Allow individuals to add their own customisation
 source_env_if_exists .envrc.local

Taskfile.yml

@@ -233,6 +233,11 @@ tasks:
       - task: build
       - cmd: go run ./tests/load/c/main --avalanchego-path=./build/avalanchego {{.CLI_ARGS}}
 
+  test-load-exclusive:
+    desc: Runs load tests against kube with exclusive scheduling
+    cmds:
+      - cmd: go run ./tests/load/c/main --runtime=kube --kube-use-exclusive-scheduling {{.CLI_ARGS}}
+
   test-unit:
     desc: Runs unit tests
     # Invoking with bash ensures compatibility with CI execution on Windows

tests/fixture/e2e/helpers.go

@@ -281,8 +281,16 @@ func StartNetwork(
 ) {
 	require := require.New(tc)
 
-	err := tmpnet.BootstrapNewNetwork(
-		tc.DefaultContext(),
+	nodeCount := len(network.Nodes)
+	timeout, err := network.DefaultRuntimeConfig.GetNetworkStartTimeout(nodeCount)
+	require.NoError(err)
+	tc.Log().Info("waiting for network to start",
+		zap.Float64("timeoutSeconds", timeout.Seconds()),
+	)
+	ctx := tc.ContextWithTimeout(timeout)
+
+	err = tmpnet.BootstrapNewNetwork(
+		ctx,
 		tc.Log(),
 		network,
 		rootNetworkDir,

tests/fixture/tmpnet/flags/kube_runtime.go

@@ -20,29 +20,33 @@ const (
 )
 
 var (
-	errKubeNamespaceRequired     = errors.New("--kube-namespace is required")
-	errKubeImageRequired         = errors.New("--kube-image is required")
-	errKubeMinVolumeSizeRequired = fmt.Errorf("--kube-volume-size must be >= %d", tmpnet.MinimumVolumeSizeGB)
+	errKubeNamespaceRequired       = errors.New("--kube-namespace is required")
+	errKubeImageRequired           = errors.New("--kube-image is required")
+	errKubeMinVolumeSizeRequired   = fmt.Errorf("--kube-volume-size must be >= %d", tmpnet.MinimumVolumeSizeGB)
+	errKubeSchedulingLabelRequired = errors.New("--kube-scheduling-label-key and --kube-scheduling-label-value are required when --kube-use-exclusive-scheduling is enabled")
 )
 
 type kubeRuntimeVars struct {
-	namespace    string
-	image        string
-	volumeSizeGB uint
-	config       *KubeconfigVars
+	namespace              string
+	image                  string
+	volumeSizeGB           uint
+	useExclusiveScheduling bool
+	schedulingLabelKey     string
+	schedulingLabelValue   string
+	config                 *KubeconfigVars
 }
 
 func (v *kubeRuntimeVars) registerWithFlag() {
 	v.config = newKubeconfigFlagVars(kubeDocPrefix)
-	v.register(flag.StringVar, flag.UintVar)
+	v.register(flag.StringVar, flag.UintVar, flag.BoolVar)
 }
 
 func (v *kubeRuntimeVars) registerWithFlagSet(flagSet *pflag.FlagSet) {
 	v.config = newKubeconfigFlagSetVars(flagSet, kubeDocPrefix)
-	v.register(flagSet.StringVar, flagSet.UintVar)
+	v.register(flagSet.StringVar, flagSet.UintVar, flagSet.BoolVar)
 }
 
-func (v *kubeRuntimeVars) register(stringVar varFunc[string], uintVar varFunc[uint]) {
+func (v *kubeRuntimeVars) register(stringVar varFunc[string], uintVar varFunc[uint], boolVar varFunc[bool]) {
 	stringVar(
 		&v.namespace,
 		"kube-namespace",
@@ -64,6 +68,24 @@ func (v *kubeRuntimeVars) register(stringVar varFunc[string], uintVar varFunc[ui
 			tmpnet.MinimumVolumeSizeGB,
 		),
 	)
+	boolVar(
+		&v.useExclusiveScheduling,
+		"kube-use-exclusive-scheduling",
+		false,
+		kubeDocPrefix+"Whether to schedule each AvalancheGo node to a dedicated Kubernetes node",
+	)
+	stringVar(
+		&v.schedulingLabelKey,
+		"kube-scheduling-label-key",
+		"purpose",
+		kubeDocPrefix+"The label key to use for exclusive scheduling for node selection and toleration",
+	)
+	stringVar(
+		&v.schedulingLabelValue,
+		"kube-scheduling-label-value",
+		"higher-spec",
+		kubeDocPrefix+"The label value to use for exclusive scheduling for node selection and toleration",
+	)
 }
 
 func (v *kubeRuntimeVars) getKubeRuntimeConfig() (*tmpnet.KubeRuntimeConfig, error) {
@@ -76,11 +98,17 @@ func (v *kubeRuntimeVars) getKubeRuntimeConfig() (*tmpnet.KubeRuntimeConfig, err
 	if v.volumeSizeGB < tmpnet.MinimumVolumeSizeGB {
 		return nil, errKubeMinVolumeSizeRequired
 	}
+	if v.useExclusiveScheduling && (len(v.schedulingLabelKey) == 0 || len(v.schedulingLabelValue) == 0) {
+		return nil, errKubeSchedulingLabelRequired
+	}
 	return &tmpnet.KubeRuntimeConfig{
-		ConfigPath:    v.config.Path,
-		ConfigContext: v.config.Context,
-		Namespace:     v.namespace,
-		Image:         v.image,
-		VolumeSizeGB:  v.volumeSizeGB,
+		ConfigPath:             v.config.Path,
+		ConfigContext:          v.config.Context,
+		Namespace:              v.namespace,
+		Image:                  v.image,
+		VolumeSizeGB:           v.volumeSizeGB,
+		UseExclusiveScheduling: v.useExclusiveScheduling,
+		SchedulingLabelKey:     v.schedulingLabelKey,
+		SchedulingLabelValue:   v.schedulingLabelValue,
 	}, nil
 }

tests/fixture/tmpnet/kube_runtime.go

@@ -43,6 +43,12 @@ const (
 	// - EBS volume sizes are in GB
 	//   - The minimum number greater than 1GB is 2GB
 	MinimumVolumeSizeGB = 2
+
+	// All statefulsets configured for exclusive scheduling will use
+	// anti-affinity with the following labeling to ensure their pods
+	// are never scheduled to the same nodes.
+	antiAffinityLabelKey   = "tmpnet-scheduling"
+	antiAffinityLabelValue = "exclusive"
 )
 
 type KubeRuntimeConfig struct {
@@ -58,6 +64,12 @@ type KubeRuntimeConfig struct {
 	Image string `json:"image,omitempty"`
 	// Size in gigabytes of the PersistentVolumeClaim  to allocate for the node
 	VolumeSizeGB uint `json:"volumeSizeGB,omitempty"`
+	// Whether to schedule each AvalancheGo node to a dedicated Kubernetes node
+	UseExclusiveScheduling bool `json:"useExclusiveScheduling,omitempty"`
+	// Label key to use for exclusive scheduling for node selection and toleration
+	SchedulingLabelKey string `json:"schedulingLabelKey,omitempty"`
+	// Label value to use for exclusive scheduling for node selection and toleration
+	SchedulingLabelValue string `json:"schedulingLabelValue,omitempty"`
 }
 
 type KubeRuntime struct {
@@ -265,6 +277,22 @@ func (p *KubeRuntime) Start(ctx context.Context) error {
 		p.node.getMonitoringLabels(),
 	)
 
+	if runtimeConfig.UseExclusiveScheduling {
+		labelKey := runtimeConfig.SchedulingLabelKey
+		labelValue := runtimeConfig.SchedulingLabelValue
+		log.Debug("configuring exclusive scheduling",
+			zap.String("nodeID", nodeID),
+			zap.String("namespace", namespace),
+			zap.String("statefulSet", statefulSetName),
+			zap.String("schedulingLabelKey", labelKey),
+			zap.String("schedulingLabelValue", labelValue),
+		)
+		if labelKey == "" || labelValue == "" {
+			return errors.New("scheduling label key and value must be non-empty when exclusive scheduling is enabled")
+		}
+		configureExclusiveScheduling(&statefulSet.Spec.Template, labelKey, labelValue)
+	}
+
 	_, err = clientset.AppsV1().StatefulSets(runtimeConfig.Namespace).Create(
 		ctx,
 		statefulSet,
@@ -764,3 +792,48 @@ func (p *KubeRuntime) getFlags() (FlagsMap, error) {
 	flags[config.HTTPHostKey] = "0.0.0.0"
 	return flags, nil
 }
+
+// configureExclusiveScheduling ensures that the provided template schedules only to nodes with the provided
+// labeling, tolerates a taint that matches the labeling, and uses anti-affinity to ensure only a single
+// avalanchego pod is scheduled to a given target node.
+func configureExclusiveScheduling(template *corev1.PodTemplateSpec, labelKey string, labelValue string) {
+	podSpec := &template.Spec
+
+	// Configure node selection
+	if podSpec.NodeSelector == nil {
+		podSpec.NodeSelector = make(map[string]string)
+	}
+	podSpec.NodeSelector[labelKey] = labelValue
+
+	// Configure toleration. Nodes are assumed to have a taint with the same
+	// key+value as the label used to select it.
+	podSpec.Tolerations = []corev1.Toleration{
+		{
+			Key:      labelKey,
+			Operator: corev1.TolerationOpEqual,
+			Value:    labelValue,
+			Effect:   corev1.TaintEffectNoExecute,
+		},
+	}
+
+	// Configure anti-affinity to ensure only one pod per node
+	templateMeta := &template.ObjectMeta
+	if templateMeta.Labels == nil {
+		templateMeta.Labels = make(map[string]string)
+	}
+	templateMeta.Labels[antiAffinityLabelKey] = antiAffinityLabelValue
+	podSpec.Affinity = &corev1.Affinity{
+		PodAntiAffinity: &corev1.PodAntiAffinity{
+			RequiredDuringSchedulingIgnoredDuringExecution: []corev1.PodAffinityTerm{
+				{
+					LabelSelector: &metav1.LabelSelector{
+						MatchLabels: map[string]string{
+							antiAffinityLabelKey: antiAffinityLabelValue,
+						},
+					},
+					TopologyKey: "kubernetes.io/hostname",
+				},
+			},
+		},
+	}
+}

tests/fixture/tmpnet/network.go

@@ -787,7 +787,7 @@ func (n *Network) GetNodeURIs(ctx context.Context, deferCleanupFunc func(func())
 // GetAvailableNodeIDs returns the node IDs of nodes in the network that are running and not ephemeral.
 func (n *Network) GetAvailableNodeIDs() []string {
 	availableNodes := FilterAvailableNodes(n.Nodes)
-	ids := make([]string, len(availableNodes))
+	ids := []string{}
 	for _, node := range availableNodes {
 		ids = append(ids, node.NodeID.String())
 	}

tests/fixture/tmpnet/node.go

@@ -56,6 +56,27 @@ type NodeRuntimeConfig struct {
 	Kube    *KubeRuntimeConfig    `json:"kube,omitempty"`
 }
 
+// GetNetworkStartTimeout returns the timeout to use when starting a network.
+func (c *NodeRuntimeConfig) GetNetworkStartTimeout(nodeCount int) (time.Duration, error) {
+	switch {
+	case c.Process != nil:
+		// Processes are expected to start quickly, nodeCount is ignored
+		return DefaultNetworkTimeout, nil
+	case c.Kube != nil:
+		// Ensure sufficient time for scheduling and image pull
+		timeout := time.Duration(nodeCount) * time.Minute
+
+		if c.Kube.UseExclusiveScheduling {
+			// Ensure sufficient time for the creation of autoscaled nodes
+			timeout *= 2
+		}
+
+		return timeout, nil
+	default:
+		return 0, errors.New("no runtime configuration set")
+	}
+}
+
 // Node supports configuring and running a node participating in a temporary network.
 type Node struct {
 	// Set by EnsureNodeID which is also called when the node is read.

tests/fixture/tmpnet/tmpnetctl/main.go

@@ -80,7 +80,15 @@ func main() {
 				DefaultRuntimeConfig: *nodeRuntimeConfig,
 			}
 
-			ctx, cancel := context.WithTimeout(context.Background(), tmpnet.DefaultNetworkTimeout)
+			timeout, err := nodeRuntimeConfig.GetNetworkStartTimeout(nodeCount)
+			if err != nil {
+				return err
+			}
+			log.Info("waiting for network to start",
+				zap.Float64("timeoutSeconds", timeout.Seconds()),
+			)
+
+			ctx, cancel := context.WithTimeout(context.Background(), timeout)
 			defer cancel()
 			if err := tmpnet.BootstrapNewNetwork(
 				ctx,

tests/fixture/tmpnet/utils.go

@@ -59,7 +59,7 @@ type NodeURI struct {
 // nodes are running as local processes or in a kube cluster.
 func GetNodeURIs(ctx context.Context, nodes []*Node, deferCleanupFunc func(func())) ([]NodeURI, error) {
 	availableNodes := FilterAvailableNodes(nodes)
-	uris := make([]NodeURI, 0, len(availableNodes))
+	uris := []NodeURI{}
 	for _, node := range availableNodes {
 		uri, cancel, err := node.GetLocalURI(ctx)
 		if err != nil {
@@ -77,7 +77,7 @@ func GetNodeURIs(ctx context.Context, nodes []*Node, deferCleanupFunc func(func(
 
 // FilteredAvailableNodes filters the provided nodes by whether they are running and not ephemeral.
 func FilterAvailableNodes(nodes []*Node) []*Node {
-	filteredNodes := make([]*Node, 0, len(nodes))
+	filteredNodes := []*Node{}
 	for _, node := range nodes {
 		if node.IsEphemeral {
 			// Avoid returning URIs for nodes whose lifespan is indeterminate