e2e: add SSH retry logic for reboot scenarios

Add configurable SSH connectivity retry mechanism for scenarios where nodes may
reboot during provisioning (e.g., MIG-enabled GPU nodes).

- Add WaitForSSHAfterReboot config option to enable retry behavior
- Implement retry logic with exponential backoff using
  wait.PollUntilContextTimeout
- Add reboot detection for common SSH error patterns
- Enable 5-minute retry window for Ubuntu 2404 MIG GPU test
- Maintain backward compatibility with existing single-attempt behavior

This resolves flaky test failures when GPU driver installation triggers
automatic reboots during node provisioning.

Signed-off-by: Suraj Deshmukh <suraj.deshmukh@microsoft.com>

Author: surajssd

e2e/scenario_gpu_managed_experience_test.go

@@ -3,6 +3,7 @@ package e2e
 import (
 	"context"
 	"testing"
+	"time"
 
 	"github.com/Azure/agentbaker/e2e/components"
 	"github.com/Azure/agentbaker/e2e/config"
@@ -198,8 +199,9 @@ func Test_Ubuntu2404_NvidiaDevicePluginRunning_MIG(t *testing.T) {
 			GPU: true,
 		},
 		Config: Config{
-			Cluster: ClusterKubenet,
-			VHD:     config.VHDUbuntu2404Gen2Containerd,
+			Cluster:               ClusterKubenet,
+			VHD:                   config.VHDUbuntu2404Gen2Containerd,
+			WaitForSSHAfterReboot: 5 * time.Minute,
 			BootstrapConfigMutator: func(nbc *datamodel.NodeBootstrappingConfiguration) {
 				nbc.AgentPoolProfile.VMSize = "Standard_NC24ads_A100_v4"
 				nbc.ConfigGPUDriverIfNeeded = true

e2e/test_helpers.go

@@ -22,6 +22,7 @@ import (
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/to"
 	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/compute/armcompute/v6"
 	"github.com/stretchr/testify/require"
+	"k8s.io/apimachinery/pkg/util/wait"
 	ctrruntimelog "sigs.k8s.io/controller-runtime/pkg/log"
 	"sigs.k8s.io/controller-runtime/pkg/log/zap"
 )
@@ -649,7 +650,81 @@ func CreateSIGImageVersionFromDisk(ctx context.Context, s *Scenario, version str
 	return &customVHD
 }
 
+// isRebootRelatedSSHError checks if the error is related to a system reboot
+func isRebootRelatedSSHError(err error, stderr string) bool {
+	if err == nil {
+		return false
+	}
+
+	rebootIndicators := []string{
+		"System is going down",
+		"pam_nologin",
+		"Connection closed by",
+		"Connection refused",
+		"Connection timed out",
+	}
+
+	errMsg := err.Error()
+	for _, indicator := range rebootIndicators {
+		if strings.Contains(errMsg, indicator) || strings.Contains(stderr, indicator) {
+			return true
+		}
+	}
+	return false
+}
+
 func validateSSHConnectivity(ctx context.Context, s *Scenario) error {
+	// If WaitForSSHAfterReboot is not set, use the original single-attempt behavior
+	if s.Config.WaitForSSHAfterReboot == 0 {
+		return attemptSSHConnection(ctx, s)
+	}
+
+	// Retry logic with exponential backoff for scenarios that may reboot
+	s.T.Logf("SSH connectivity validation will retry for up to %s if reboot-related errors are encountered", s.Config.WaitForSSHAfterReboot)
+	startTime := time.Now()
+	var lastSSHError error
+
+	err := wait.PollUntilContextTimeout(ctx, 5*time.Second, s.Config.WaitForSSHAfterReboot, true, func(ctx context.Context) (bool, error) {
+		err := attemptSSHConnection(ctx, s)
+		if err == nil {
+			elapsed := time.Since(startTime)
+			s.T.Logf("SSH connectivity established after %s", toolkit.FormatDuration(elapsed))
+			return true, nil
+		}
+
+		// Save the last error for better error messages
+		lastSSHError = err
+
+		// Extract stderr from the error
+		stderr := ""
+		if strings.Contains(err.Error(), "Stderr:") {
+			parts := strings.Split(err.Error(), "Stderr:")
+			if len(parts) > 1 {
+				stderr = parts[1]
+			}
+		}
+
+		// Check if this is a reboot-related error
+		if isRebootRelatedSSHError(err, stderr) {
+			s.T.Logf("Detected reboot-related SSH error, will retry: %v", err)
+			return false, nil // Continue polling
+		}
+
+		// Not a reboot error, fail immediately
+		return false, err
+	})
+
+	// If we timed out while retrying reboot-related errors, provide a better error message
+	if err != nil && lastSSHError != nil {
+		elapsed := time.Since(startTime)
+		return fmt.Errorf("SSH connection failed after waiting %s for node to reboot and come back up. Last SSH error: %w", toolkit.FormatDuration(elapsed), lastSSHError)
+	}
+
+	return err
+}
+
+// attemptSSHConnection performs a single SSH connectivity check
+func attemptSSHConnection(ctx context.Context, s *Scenario) error {
 	connectionTest := fmt.Sprintf("%s echo 'SSH_CONNECTION_OK'", sshString(s.Runtime.VMPrivateIP))
 	connectionResult, err := execOnPrivilegedPod(ctx, s.Runtime.Cluster.Kube, defaultNamespace, s.Runtime.Cluster.DebugPod.Name, connectionTest)
 

e2e/types.go

@@ -10,6 +10,7 @@ import (
 	"strconv"
 	"strings"
 	"testing"
+	"time"
 
 	aksnodeconfigv1 "github.com/Azure/agentbaker/aks-node-controller/pkg/gen/aksnodeconfig/v1"
 	"github.com/Azure/agentbaker/e2e/config"
@@ -166,6 +167,11 @@ type Config struct {
 	// It shouldn't be used for majority of scenarios, currently only used for scenarios where the node is not expected to be reachable via ssh
 	SkipSSHConnectivityValidation bool
 
+	// WaitForSSHAfterReboot if set to non-zero duration, SSH connectivity validation will retry with exponential backoff
+	// for up to this duration when encountering reboot-related errors. This is useful for scenarios where the node
+	// reboots during provisioning (e.g., MIG-enabled GPU nodes). Default (zero value) means no retry.
+	WaitForSSHAfterReboot time.Duration
+
 	// if VHDCaching is set then a VHD will be created first for the test scenario and then a VM will be created from that VHD.
 	// The main purpose is to validate VHD Caching logic and ensure a reboot step between basePrep and nodePrep doesn't break anything.
 	VHDCaching bool