feat: new e2e uses bastion hosts for ssh

.pipelines/e2e-gpu.yaml

@@ -0,0 +1,40 @@
+name: $(Date:yyyyMMdd)$(Rev:.r)
+variables:
+  TAGS_TO_RUN: "gpu=true"
+  TAGS_TO_SKIP: "os=windows"
+  SKIP_E2E_TESTS: false
+trigger:
+  branches:
+    include:
+      - main
+pr:
+  branches:
+    include:
+      - official/*
+      - windows/*
+      - main
+  paths:
+    include:
+      - .pipelines/e2e.yaml
+      - .pipelines/templates/e2e-template.yaml
+      - .pipelines/scripts/e2e_run.sh
+      - e2e
+      - parts/linux
+      - parts/common/components.json
+      - pkg/agent
+      - go.mod
+      - go.sum
+    exclude:
+      - pkg/agent/datamodel/sig_config*.go # SIG config changes
+      - pkg/agent/datamodel/*.json # SIG version changes
+      - pkg/agent/testdata/AKSWindows* # Windows test data
+      - parts/common/components.json # centralized components management file
+      - staging/cse/windows/README
+      - /**/*.md
+      - e2e/scenario_win_test.go
+
+jobs:
+  - template: ./templates/e2e-template.yaml
+    parameters:
+      name: Linux GPU Tests
+      IgnoreScenariosWithMissingVhd: false

.pipelines/e2e.yaml

@@ -1,6 +1,6 @@
 name: $(Date:yyyyMMdd)$(Rev:.r)
 variables:
-  TAGS_TO_SKIP: "os=windows"
+  TAGS_TO_SKIP: "os=windows,gpu=true"
   SKIP_E2E_TESTS: false
 trigger:
   branches:

.pipelines/scripts/e2e_run.sh

@@ -64,16 +64,37 @@ if [ -n "${SIG_GALLERY_NAME}" ]; then
   export GALLERY_NAME=$SIG_GALLERY_NAME
 fi
 
+az extension add --name bastion
+
 # this software is used to take the output of "go test" and produce a junit report that we can upload to the pipeline
 # and see fancy test results.
 cd e2e
 mkdir -p bin
-GOBIN=`pwd`/bin/ go install gotest.tools/gotestsum@latest
+architecture=$(uname -m)
+
+case "$architecture" in
+  x86_64 | amd64) architecture="amd64" ;;
+  aarch64 | arm64) architecture="arm64" ;;
+  *)
+    echo "Unsupported architecture: $architecture"
+    exit 1
+    ;;
+esac
+
+gotestsum_version="1.13.0"
+gotestsum_archive="gotestsum_${gotestsum_version}_linux_${architecture}.tar.gz"
+gotestsum_url="https://github.com/gotestyourself/gotestsum/releases/download/v${gotestsum_version}/${gotestsum_archive}"
+
+temp_file="$(mktemp)"
+curl -fsSL "$gotestsum_url" -o "$temp_file"
+tar -xzf "$temp_file" -C bin
+chmod +x bin/gotestsum
+rm -f "$temp_file"
 
 # gotestsum configure to only show logs for failed tests, json file for detailed logs
 # Run the tests! Yey!
 test_exit_code=0
-./bin/gotestsum --format testdox --junitfile "${BUILD_SRC_DIR}/e2e/report.xml" --jsonfile "${BUILD_SRC_DIR}/e2e/test-log.json" -- -parallel 100 -timeout 90m || test_exit_code=$?
+./bin/gotestsum --format testdox --junitfile "${BUILD_SRC_DIR}/e2e/report.xml" --jsonfile "${BUILD_SRC_DIR}/e2e/test-log.json" -- -parallel 150 -timeout 90m || test_exit_code=$?
 
 # Upload test results as Azure DevOps artifacts
 echo "##vso[artifact.upload containerfolder=test-results;artifactname=e2e-test-log]${BUILD_SRC_DIR}/e2e/test-log.json"

.pipelines/templates/e2e-template.yaml

@@ -34,6 +34,8 @@ jobs:
         displayName: Run AgentBaker E2E
         env:
           E2E_SUBSCRIPTION_ID: $(E2E_SUBSCRIPTION_ID)
+          SYS_SSH_PUBLIC_KEY: $(SYS_SSH_PUBLIC_KEY)
+          SYS_SSH_PRIVATE_KEY_B64: $(SYS_SSH_PRIVATE_KEY_B64)
           BUILD_SRC_DIR: $(System.DefaultWorkingDirectory)
           DefaultWorkingDirectory: $(Build.SourcesDirectory)
           VHD_BUILD_ID: $(VHD_BUILD_ID)

aks-node-controller/go.mod

@@ -1,6 +1,6 @@
 module github.com/Azure/agentbaker/aks-node-controller
 
-go 1.23.7
+go 1.24.0
 
 require (
 	github.com/Azure/agentbaker v0.20240503.0
@@ -28,7 +28,7 @@ require (
 	github.com/pkg/errors v0.9.1 // indirect
 	github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect
 	github.com/vincent-petithory/dataurl v1.0.0 // indirect
-	golang.org/x/sys v0.35.0 // indirect
+	golang.org/x/sys v0.39.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )
 

aks-node-controller/go.sum

@@ -55,12 +55,12 @@ github.com/vincent-petithory/dataurl v1.0.0 h1:cXw+kPto8NLuJtlMsI152irrVw9fRDX8A
 github.com/vincent-petithory/dataurl v1.0.0/go.mod h1:FHafX5vmDzyP+1CQATJn7WFKc9CvnvxyvZy6I1MrG/U=
 go.yaml.in/yaml/v3 v3.0.4 h1:tfq32ie2Jv2UxXFdLJdh3jXuOzWiL1fo0bu/FbuKpbc=
 go.yaml.in/yaml/v3 v3.0.4/go.mod h1:DhzuOOF2ATzADvBadXxruRBLzYTpT36CKvDb3+aBEFg=
-golang.org/x/net v0.43.0 h1:lat02VYK2j4aLzMzecihNvTlJNQUq316m2Mr9rnM6YE=
-golang.org/x/net v0.43.0/go.mod h1:vhO1fvI4dGsIjh73sWfUVjj3N7CA9WkKJNQm2svM6Jg=
-golang.org/x/sys v0.35.0 h1:vz1N37gP5bs89s7He8XuIYXpyY0+QlsKmzipCbUtyxI=
-golang.org/x/sys v0.35.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k=
-golang.org/x/text v0.28.0 h1:rhazDwis8INMIwQ4tpjLDzUhx6RlXqZNPEM0huQojng=
-golang.org/x/text v0.28.0/go.mod h1:U8nCwOR8jO/marOQ0QbDiOngZVEBB7MAiitBuMjXiNU=
+golang.org/x/net v0.48.0 h1:zyQRTTrjc33Lhh0fBgT/H3oZq9WuvRR5gPC70xpDiQU=
+golang.org/x/net v0.48.0/go.mod h1:+ndRgGjkh8FGtu1w1FGbEC31if4VrNVMuKTgcAAnQRY=
+golang.org/x/sys v0.39.0 h1:CvCKL8MeisomCi6qNZ+wbb0DN9E5AATixKsvNtMoMFk=
+golang.org/x/sys v0.39.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks=
+golang.org/x/text v0.32.0 h1:ZD01bjUt1FQ9WJ0ClOL5vxgxOI/sVCNgX1YtKwcY0mU=
+golang.org/x/text v0.32.0/go.mod h1:o/rUWzghvpD5TXrTIBuJU77MTaN0ljMWE47kxGJQ7jY=
 google.golang.org/protobuf v1.36.6 h1:z1NpPI8ku2WgiWnf+t9wTPsn6eP1L7ksHUlkfLvd9xY=
 google.golang.org/protobuf v1.36.6/go.mod h1:jduwjTPXsFjZGTmRluh+L6NjiWu7pchiJ2/5YcXBHnY=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=

e2e/aks_model.go

@@ -13,9 +13,9 @@ import (
 	"github.com/Azure/agentbaker/pkg/agent"
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/to"
-	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/containerregistry/armcontainerregistry"
-	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/containerservice/armcontainerservice/v6"
-	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/network/armnetwork/v6"
+	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/containerregistry/armcontainerregistry/v2"
+	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/containerservice/armcontainerservice/v8"
+	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/network/armnetwork/v7"
 	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/privatedns/armprivatedns"
 )
 
@@ -175,6 +175,16 @@ func getBaseClusterModel(clusterName, location, k8sSystemPoolSKU string) *armcon
 					Enabled: to.Ptr(false),
 				},
 			},
+			LinuxProfile: &armcontainerservice.LinuxProfile{
+				AdminUsername: to.Ptr("azureuser"),
+				SSH: &armcontainerservice.SSHConfiguration{
+					PublicKeys: []*armcontainerservice.SSHPublicKey{
+						{
+							KeyData: to.Ptr(string(config.SysSSHPublicKey)),
+						},
+					},
+				},
+			},
 		},
 		Identity: &armcontainerservice.ManagedClusterIdentity{
 			Type: to.Ptr(armcontainerservice.ResourceIdentityTypeSystemAssigned),
@@ -275,6 +285,17 @@ func addFirewallRules(
 	ctx context.Context, clusterModel *armcontainerservice.ManagedCluster,
 	location string,
 ) error {
+	routeTableName := "abe2e-fw-rt"
+	rtGetResp, err := config.Azure.RouteTables.Get(
+		ctx,
+		*clusterModel.Properties.NodeResourceGroup,
+		routeTableName,
+		nil,
+	)
+	if err == nil && len(rtGetResp.Properties.Subnets) != 0 {
+		// already associated with aks subnet
+		return nil
+	}
 
 	vnet, err := getClusterVNet(ctx, *clusterModel.Properties.NodeResourceGroup)
 	if err != nil {
@@ -366,7 +387,6 @@ func addFirewallRules(
 		return fmt.Errorf("failed to get firewall private IP address")
 	}
 
-	routeTableName := "abe2e-fw-rt"
 	routeTableParams := armnetwork.RouteTable{
 		Location: to.Ptr(location),
 		Properties: &armnetwork.RouteTablePropertiesFormat{
@@ -535,55 +555,33 @@ func airGapSecurityGroup(location, clusterFQDN string) (armnetwork.SecurityGroup
 
 func addPrivateEndpointForACR(ctx context.Context, nodeResourceGroup, privateACRName string, vnet VNet, location string) error {
 	logf(ctx, "Checking if private endpoint for private container registry is in rg %s", nodeResourceGroup)
-
 	var err error
-	var exists bool
+	var privateEndpoint *armnetwork.PrivateEndpoint
 	privateEndpointName := "PE-for-ABE2ETests"
-	if exists, err = privateEndpointExists(ctx, nodeResourceGroup, privateEndpointName); err != nil {
-		return err
-	}
-	if exists {
-		logf(ctx, "Private Endpoint already exists, skipping creation")
-		return nil
-	}
-
-	var peResp armnetwork.PrivateEndpointsClientCreateOrUpdateResponse
-	if peResp, err = createPrivateEndpoint(ctx, nodeResourceGroup, privateEndpointName, privateACRName, vnet, location); err != nil {
+	if privateEndpoint, err = createPrivateEndpoint(ctx, nodeResourceGroup, privateEndpointName, privateACRName, vnet, location); err != nil {
 		return err
 	}
 
 	privateZoneName := "privatelink.azurecr.io"
-	var pzResp armprivatedns.PrivateZonesClientCreateOrUpdateResponse
-	if pzResp, err = createPrivateZone(ctx, nodeResourceGroup, privateZoneName); err != nil {
+	var privateZone *armprivatedns.PrivateZone
+	if privateZone, err = createPrivateZone(ctx, nodeResourceGroup, privateZoneName); err != nil {
 		return err
 	}
 
 	if err = createPrivateDNSLink(ctx, vnet, nodeResourceGroup, privateZoneName); err != nil {
 		return err
 	}
 
-	if err = addRecordSetToPrivateDNSZone(ctx, peResp, nodeResourceGroup, privateZoneName); err != nil {
+	if err = addRecordSetToPrivateDNSZone(ctx, privateEndpoint, nodeResourceGroup, privateZoneName); err != nil {
 		return err
 	}
 
-	if err = addDNSZoneGroup(ctx, pzResp, nodeResourceGroup, privateZoneName, *peResp.Name); err != nil {
+	if err = addDNSZoneGroup(ctx, privateZone, nodeResourceGroup, privateZoneName, *privateEndpoint.Name); err != nil {
 		return err
 	}
 	return nil
 }
 
-func privateEndpointExists(ctx context.Context, nodeResourceGroup, privateEndpointName string) (bool, error) {
-	existingPE, err := config.Azure.PrivateEndpointClient.Get(ctx, nodeResourceGroup, privateEndpointName, nil)
-	if err == nil && existingPE.ID != nil {
-		logf(ctx, "Private Endpoint already exists with ID: %s", *existingPE.ID)
-		return true, nil
-	}
-	if err != nil && !strings.Contains(err.Error(), "ResourceNotFound") {
-		return false, fmt.Errorf("failed to get private endpoint: %w", err)
-	}
-	return false, nil
-}
-
 func createPrivateAzureContainerRegistryPullSecret(ctx context.Context, cluster *armcontainerservice.ManagedCluster, kubeconfig *Kubeclient, resourceGroup string, isNonAnonymousPull bool) error {
 	privateACRName := config.GetPrivateACRName(isNonAnonymousPull, *cluster.Location)
 	if isNonAnonymousPull {
@@ -768,7 +766,15 @@ func addCacheRulesToPrivateAzureContainerRegistry(ctx context.Context, resourceG
 	return nil
 }
 
-func createPrivateEndpoint(ctx context.Context, nodeResourceGroup, privateEndpointName, privateACRName string, vnet VNet, location string) (armnetwork.PrivateEndpointsClientCreateOrUpdateResponse, error) {
+func createPrivateEndpoint(ctx context.Context, nodeResourceGroup, privateEndpointName, privateACRName string, vnet VNet, location string) (*armnetwork.PrivateEndpoint, error) {
+	existingPE, err := config.Azure.PrivateEndpointClient.Get(ctx, nodeResourceGroup, privateEndpointName, nil)
+	if err == nil && existingPE.ID != nil {
+		logf(ctx, "Private Endpoint already exists with ID: %s", *existingPE.ID)
+		return &existingPE.PrivateEndpoint, nil
+	}
+	if err != nil && !strings.Contains(err.Error(), "ResourceNotFound") {
+		return nil, fmt.Errorf("failed to get private endpoint: %w", err)
+	}
 	logf(ctx, "Creating Private Endpoint in rg %s", nodeResourceGroup)
 	acrID := fmt.Sprintf("/subscriptions/%s/resourceGroups/%s/providers/Microsoft.ContainerRegistry/registries/%s", config.Config.SubscriptionID, config.ResourceGroupName(location), privateACRName)
 
@@ -798,18 +804,27 @@ func createPrivateEndpoint(ctx context.Context, nodeResourceGroup, privateEndpoi
 		nil,
 	)
 	if err != nil {
-		return armnetwork.PrivateEndpointsClientCreateOrUpdateResponse{}, fmt.Errorf("failed to create private endpoint in BeginCreateOrUpdate: %w", err)
+		return nil, fmt.Errorf("failed to create private endpoint in BeginCreateOrUpdate: %w", err)
 	}
 	resp, err := poller.PollUntilDone(ctx, nil)
 	if err != nil {
-		return armnetwork.PrivateEndpointsClientCreateOrUpdateResponse{}, fmt.Errorf("failed to create private endpoint in polling: %w", err)
+		return nil, fmt.Errorf("failed to create private endpoint in polling: %w", err)
 	}
 
 	logf(ctx, "Private Endpoint created or updated with ID: %s", *resp.ID)
-	return resp, nil
+	return &resp.PrivateEndpoint, nil
 }
 
-func createPrivateZone(ctx context.Context, nodeResourceGroup, privateZoneName string) (armprivatedns.PrivateZonesClientCreateOrUpdateResponse, error) {
+func createPrivateZone(ctx context.Context, nodeResourceGroup, privateZoneName string) (*armprivatedns.PrivateZone, error) {
+	pzResp, err := config.Azure.PrivateZonesClient.Get(
+		ctx,
+		nodeResourceGroup,
+		privateZoneName,
+		nil,
+	)
+	if err == nil {
+		return &pzResp.PrivateZone, nil
+	}
 	dnsZoneParams := armprivatedns.PrivateZone{
 		Location: to.Ptr("global"),
 	}
@@ -821,23 +836,36 @@ func createPrivateZone(ctx context.Context, nodeResourceGroup, privateZoneName s
 		nil,
 	)
 	if err != nil {
-		return armprivatedns.PrivateZonesClientCreateOrUpdateResponse{}, fmt.Errorf("failed to create private dns zone in BeginCreateOrUpdate: %w", err)
+		return nil, fmt.Errorf("failed to create private dns zone in BeginCreateOrUpdate: %w", err)
 	}
 	resp, err := poller.PollUntilDone(ctx, nil)
 	if err != nil {
-		return armprivatedns.PrivateZonesClientCreateOrUpdateResponse{}, fmt.Errorf("failed to create private dns zone in polling: %w", err)
+		return nil, fmt.Errorf("failed to create private dns zone in polling: %w", err)
 	}
 
 	logf(ctx, "Private DNS Zone created or updated with ID: %s", *resp.ID)
-	return resp, nil
+	return &resp.PrivateZone, nil
 }
 
 func createPrivateDNSLink(ctx context.Context, vnet VNet, nodeResourceGroup, privateZoneName string) error {
+	networkLinkName := "link-ABE2ETests"
+	_, err := config.Azure.VirutalNetworkLinksClient.Get(
+		ctx,
+		nodeResourceGroup,
+		privateZoneName,
+		networkLinkName,
+		nil,
+	)
+
+	if err == nil {
+		// private dns link already created
+		return nil
+	}
+
 	vnetForId, err := config.Azure.VNet.Get(ctx, nodeResourceGroup, vnet.name, nil)
 	if err != nil {
 		return fmt.Errorf("failed to get vnet: %w", err)
 	}
-	networkLinkName := "link-ABE2ETests"
 	linkParams := armprivatedns.VirtualNetworkLink{
 		Location: to.Ptr("global"),
 		Properties: &armprivatedns.VirtualNetworkLinkProperties{
@@ -867,16 +895,16 @@ func createPrivateDNSLink(ctx context.Context, vnet VNet, nodeResourceGroup, pri
 	return nil
 }
 
-func addRecordSetToPrivateDNSZone(ctx context.Context, peResp armnetwork.PrivateEndpointsClientCreateOrUpdateResponse, nodeResourceGroup, privateZoneName string) error {
-	for i, dnsConfigPtr := range peResp.Properties.CustomDNSConfigs {
+func addRecordSetToPrivateDNSZone(ctx context.Context, privateEndpoint *armnetwork.PrivateEndpoint, nodeResourceGroup, privateZoneName string) error {
+	for i, dnsConfigPtr := range privateEndpoint.Properties.CustomDNSConfigs {
 		var ipAddresses []string
 		if dnsConfigPtr == nil {
 			return fmt.Errorf("CustomDNSConfigs[%d] is nil", i)
 		}
 
 		// get the ip addresses
 		dnsConfig := *dnsConfigPtr
-		if dnsConfig.IPAddresses == nil || len(dnsConfig.IPAddresses) == 0 {
+		if len(dnsConfig.IPAddresses) == 0 {
 			return fmt.Errorf("CustomDNSConfigs[%d].IPAddresses is nil or empty", i)
 		}
 		for _, ipPtr := range dnsConfig.IPAddresses {
@@ -907,15 +935,19 @@ func addRecordSetToPrivateDNSZone(ctx context.Context, peResp armnetwork.Private
 	return nil
 }
 
-func addDNSZoneGroup(ctx context.Context, pzResp armprivatedns.PrivateZonesClientCreateOrUpdateResponse, nodeResourceGroup, privateZoneName, endpointName string) error {
+func addDNSZoneGroup(ctx context.Context, privateZone *armprivatedns.PrivateZone, nodeResourceGroup, privateZoneName, endpointName string) error {
 	groupName := strings.Replace(privateZoneName, ".", "-", -1) // replace . with -
+	_, err := config.Azure.PrivateDNSZoneGroup.Get(ctx, nodeResourceGroup, endpointName, groupName, nil)
+	if err == nil {
+		return nil
+	}
 	dnsZonegroup := armnetwork.PrivateDNSZoneGroup{
 		Name: to.Ptr(fmt.Sprintf("%s/default", privateZoneName)),
 		Properties: &armnetwork.PrivateDNSZoneGroupPropertiesFormat{
 			PrivateDNSZoneConfigs: []*armnetwork.PrivateDNSZoneConfig{{
 				Name: to.Ptr(groupName),
 				Properties: &armnetwork.PrivateDNSZonePropertiesFormat{
-					PrivateDNSZoneID: pzResp.ID,
+					PrivateDNSZoneID: privateZone.ID,
 				},
 			}},
 		},