Skip to content

Commit

Permalink
Merge pull request #6689 from wenxuan0923/wenx/add-vms-spec
Browse files Browse the repository at this point in the history 
[Azure VMs Pool] Support mixed agentpool types in Azure Cache
  • Loading branch information
@k8s-ci-robot
k8s-ci-robot authored Apr 12, 2024
2 parents 46f04ae + cca2f18 commit a7b4386
Show file tree
Hide file tree
Showing 7 changed files with 402 additions and 38 deletions.
79 changes: 57 additions & 22 deletions cluster-autoscaler/cloudprovider/azure/azure_cache.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,7 @@ type azureCache struct {
// Cache content.
resourceGroup string
vmType string
vmsPoolSet map[string]struct{} // track the nodepools that're vms pool
scaleSets map[string]compute.VirtualMachineScaleSet
virtualMachines map[string][]compute.VirtualMachine
registeredNodeGroups []cloudprovider.NodeGroup
Expand All @@ -67,6 +68,7 @@ func newAzureCache(client *azClient, cacheTTL time.Duration, resourceGroup, vmTy
refreshInterval: cacheTTL,
resourceGroup: resourceGroup,
vmType: vmType,
vmsPoolSet: make(map[string]struct{}),
scaleSets: make(map[string]compute.VirtualMachineScaleSet),
virtualMachines: make(map[string][]compute.VirtualMachine),
registeredNodeGroups: make([]cloudprovider.NodeGroup, 0),
Expand All @@ -87,6 +89,13 @@ func newAzureCache(client *azClient, cacheTTL time.Duration, resourceGroup, vmTy
return cache, nil
}

func (m *azureCache) getVMsPoolSet() map[string]struct{} {
m.mutex.Lock()
defer m.mutex.Unlock()

return m.vmsPoolSet
}

func (m *azureCache) getVirtualMachines() map[string][]compute.VirtualMachine {
m.mutex.Lock()
defer m.mutex.Unlock()
Expand Down Expand Up @@ -165,54 +174,78 @@ func (m *azureCache) fetchAzureResources() error {
m.mutex.Lock()
defer m.mutex.Unlock()

switch m.vmType {
case vmTypeVMSS:
// List all VMSS in the RG.
vmssResult, err := m.fetchScaleSets()
if err == nil {
m.scaleSets = vmssResult
} else {
return err
}
case vmTypeStandard:
// List all VMs in the RG.
vmResult, err := m.fetchVirtualMachines()
if err == nil {
m.virtualMachines = vmResult
} else {
return err
}
// fetch all the resources since CAS may be operating on mixed nodepools
// including both VMSS and VMs pools
vmssResult, err := m.fetchScaleSets()
if err == nil {
m.scaleSets = vmssResult
} else {
return err
}

vmResult, vmsPoolSet, err := m.fetchVirtualMachines()
if err == nil {
m.virtualMachines = vmResult
m.vmsPoolSet = vmsPoolSet
} else {
return err
}

return nil
}

const (
legacyAgentpoolNameTag = "poolName"
agentpoolNameTag = "aks-managed-poolName"
agentpoolTypeTag = "aks-managed-agentpool-type"
vmsPoolType = "VirtualMachines"
)

// fetchVirtualMachines returns the updated list of virtual machines in the config resource group using the Azure API.
func (m *azureCache) fetchVirtualMachines() (map[string][]compute.VirtualMachine, error) {
func (m *azureCache) fetchVirtualMachines() (map[string][]compute.VirtualMachine, map[string]struct{}, error) {
ctx, cancel := getContextWithCancel()
defer cancel()

result, err := m.azClient.virtualMachinesClient.List(ctx, m.resourceGroup)
if err != nil {
klog.Errorf("VirtualMachinesClient.List in resource group %q failed: %v", m.resourceGroup, err)
return nil, err.Error()
return nil, nil, err.Error()
}

instances := make(map[string][]compute.VirtualMachine)
// track the nodepools that're vms pools
vmsPoolSet := make(map[string]struct{})
for _, instance := range result {
if instance.Tags == nil {
continue
}

tags := instance.Tags
vmPoolName := tags["poolName"]
vmPoolName := tags[agentpoolNameTag]
// fall back to legacy tag name if not found
if vmPoolName == nil {
vmPoolName = tags[legacyAgentpoolNameTag]
}

if vmPoolName == nil {
continue
}

instances[to.String(vmPoolName)] = append(instances[to.String(vmPoolName)], instance)

// if the nodepool is already in the map, skip it
if _, ok := vmsPoolSet[to.String(vmPoolName)]; ok {
continue
}

// nodes from vms pool will have tag "aks-managed-agentpool-type" set to "VirtualMachines"
if agnetpoolType := tags[agentpoolTypeTag]; agnetpoolType != nil {
if strings.EqualFold(to.String(agnetpoolType), vmsPoolType) {
vmsPoolSet[to.String(vmPoolName)] = struct{}{}
}
}
}
return instances, nil
return instances, vmsPoolSet, nil
}

// fetchScaleSets returns the updated list of scale sets in the config resource group using the Azure API.
Expand Down Expand Up @@ -323,6 +356,7 @@ func (m *azureCache) getAutoscalingOptions(ref azureRef) map[string]string {

// FindForInstance returns node group of the given Instance
func (m *azureCache) FindForInstance(instance *azureRef, vmType string) (cloudprovider.NodeGroup, error) {
vmsPoolSet := m.getVMsPoolSet()
m.mutex.Lock()
defer m.mutex.Unlock()

Expand All @@ -340,7 +374,8 @@ func (m *azureCache) FindForInstance(instance *azureRef, vmType string) (cloudpr
return nil, nil
}

if vmType == vmTypeVMSS {
// cluster with vmss pool only
if vmType == vmTypeVMSS && len(vmsPoolSet) == 0 {
if m.areAllScaleSetsUniform() {
// Omit virtual machines not managed by vmss only in case of uniform scale set.
if ok := virtualMachineRE.Match([]byte(inst.Name)); ok {
Expand Down
73 changes: 68 additions & 5 deletions cluster-autoscaler/cloudprovider/azure/azure_cloud_provider_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,9 @@ func newTestAzureManager(t *testing.T) *AzureManager {
mockVMSSClient.EXPECT().List(gomock.Any(), "rg").Return(expectedScaleSets, nil).AnyTimes()
mockVMSSVMClient := mockvmssvmclient.NewMockInterface(ctrl)
mockVMSSVMClient.EXPECT().List(gomock.Any(), "rg", "test-vmss", gomock.Any()).Return(expectedVMSSVMs, nil).AnyTimes()
mockVMClient := mockvmclient.NewMockInterface(ctrl)
expectedVMs := newTestVMList(3)
mockVMClient.EXPECT().List(gomock.Any(), "rg").Return(expectedVMs, nil).AnyTimes()

manager := &AzureManager{
env: azure.PublicCloud,
Expand All @@ -57,6 +60,7 @@ func newTestAzureManager(t *testing.T) *AzureManager {
azClient: &azClient{
virtualMachineScaleSetsClient: mockVMSSClient,
virtualMachineScaleSetVMsClient: mockVMSSVMClient,
virtualMachinesClient: mockVMClient,
deploymentsClient: &DeploymentsClientMock{
FakeStore: map[string]resources.DeploymentExtended{
"deployment": {
Expand Down Expand Up @@ -115,9 +119,68 @@ func TestNodeGroups(t *testing.T) {
assert.Equal(t, len(provider.NodeGroups()), 0)

registered := provider.azureManager.RegisterNodeGroup(
newTestScaleSet(provider.azureManager, "test-asg"))
newTestScaleSet(provider.azureManager, "test-asg"),
)
assert.True(t, registered)
assert.Equal(t, len(provider.NodeGroups()), 1)
registered = provider.azureManager.RegisterNodeGroup(
newTestVMsPool(provider.azureManager, "test-vms-pool"),
)
assert.True(t, registered)
assert.Equal(t, len(provider.NodeGroups()), 2)
}

func TestMixedNodeGroups(t *testing.T) {
ctrl := gomock.NewController(t)
provider := newTestProvider(t)
mockVMSSClient := mockvmssclient.NewMockInterface(ctrl)
mockVMClient := mockvmclient.NewMockInterface(ctrl)
mockVMSSVMClient := mockvmssvmclient.NewMockInterface(ctrl)
provider.azureManager.azClient.virtualMachinesClient = mockVMClient
provider.azureManager.azClient.virtualMachineScaleSetsClient = mockVMSSClient
provider.azureManager.azClient.virtualMachineScaleSetVMsClient = mockVMSSVMClient

expectedScaleSets := newTestVMSSList(3, "test-asg", "eastus", compute.Uniform)
expectedVMsPoolVMs := newTestVMsPoolVMList(3)
expectedVMSSVMs := newTestVMSSVMList(3)

mockVMSSClient.EXPECT().List(gomock.Any(), provider.azureManager.config.ResourceGroup).Return(expectedScaleSets, nil).AnyTimes()
mockVMClient.EXPECT().List(gomock.Any(), provider.azureManager.config.ResourceGroup).Return(expectedVMsPoolVMs, nil).AnyTimes()
mockVMSSVMClient.EXPECT().List(gomock.Any(), provider.azureManager.config.ResourceGroup, "test-asg", gomock.Any()).Return(expectedVMSSVMs, nil).AnyTimes()

assert.Equal(t, len(provider.NodeGroups()), 0)
registered := provider.azureManager.RegisterNodeGroup(
newTestScaleSet(provider.azureManager, "test-asg"),
)
provider.azureManager.explicitlyConfigured["test-asg"] = true
assert.True(t, registered)

registered = provider.azureManager.RegisterNodeGroup(
newTestVMsPool(provider.azureManager, "test-vms-pool"),
)
provider.azureManager.explicitlyConfigured["test-vms-pool"] = true
assert.True(t, registered)
assert.Equal(t, len(provider.NodeGroups()), 2)

// refresh cache
provider.azureManager.forceRefresh()

// node from vmss pool
node := newApiNode(compute.Uniform, 0)
group, err := provider.NodeGroupForNode(node)
assert.NoError(t, err)
assert.NotNil(t, group, "Group should not be nil")
assert.Equal(t, group.Id(), "test-asg")
assert.Equal(t, group.MinSize(), 1)
assert.Equal(t, group.MaxSize(), 5)

// node from vms pool
vmsPoolNode := newVMsNode(0)
group, err = provider.NodeGroupForNode(vmsPoolNode)
assert.NoError(t, err)
assert.NotNil(t, group, "Group should not be nil")
assert.Equal(t, group.Id(), "test-vms-pool")
assert.Equal(t, group.MinSize(), 3)
assert.Equal(t, group.MaxSize(), 10)
}

func TestNodeGroupForNode(t *testing.T) {
Expand All @@ -135,6 +198,9 @@ func TestNodeGroupForNode(t *testing.T) {
mockVMSSClient := mockvmssclient.NewMockInterface(ctrl)
mockVMSSClient.EXPECT().List(gomock.Any(), provider.azureManager.config.ResourceGroup).Return(expectedScaleSets, nil)
provider.azureManager.azClient.virtualMachineScaleSetsClient = mockVMSSClient
mockVMClient := mockvmclient.NewMockInterface(ctrl)
provider.azureManager.azClient.virtualMachinesClient = mockVMClient
mockVMClient.EXPECT().List(gomock.Any(), provider.azureManager.config.ResourceGroup).Return(expectedVMs, nil).AnyTimes()

if orchMode == compute.Uniform {

Expand All @@ -143,11 +209,8 @@ func TestNodeGroupForNode(t *testing.T) {
provider.azureManager.azClient.virtualMachineScaleSetVMsClient = mockVMSSVMClient
} else {

mockVMClient := mockvmclient.NewMockInterface(ctrl)
provider.azureManager.config.EnableVmssFlex = true
mockVMClient.EXPECT().ListVmssFlexVMsWithoutInstanceView(gomock.Any(), "test-asg").Return(expectedVMs, nil).AnyTimes()
provider.azureManager.azClient.virtualMachinesClient = mockVMClient

}

registered := provider.azureManager.RegisterNodeGroup(
Expand Down
5 changes: 5 additions & 0 deletions cluster-autoscaler/cloudprovider/azure/azure_manager.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -147,6 +147,11 @@ func (m *AzureManager) buildNodeGroupFromSpec(spec string) (cloudprovider.NodeGr
return nil, fmt.Errorf("failed to parse node group spec: %v", err)
}

vmsPoolSet := m.azureCache.getVMsPoolSet()
if _, ok := vmsPoolSet[s.Name]; ok {
return NewVMsPool(s, m), nil
}

switch m.config.VMType {
case vmTypeStandard:
return NewAgentPool(s, m)
Expand Down
9 changes: 9 additions & 0 deletions cluster-autoscaler/cloudprovider/azure/azure_manager_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -144,6 +144,7 @@ func TestCreateAzureManagerValidConfig(t *testing.T) {
mockVMClient := mockvmclient.NewMockInterface(ctrl)
mockVMSSClient := mockvmssclient.NewMockInterface(ctrl)
mockVMSSClient.EXPECT().List(gomock.Any(), "fakeId").Return([]compute.VirtualMachineScaleSet{}, nil).Times(2)
mockVMClient.EXPECT().List(gomock.Any(), "fakeId").Return([]compute.VirtualMachine{}, nil).Times(2)
mockAzClient := &azClient{
virtualMachinesClient: mockVMClient,
virtualMachineScaleSetsClient: mockVMSSClient,
Expand Down Expand Up @@ -226,6 +227,7 @@ func TestCreateAzureManagerValidConfigForStandardVMType(t *testing.T) {
mockVMClient := mockvmclient.NewMockInterface(ctrl)
mockVMClient.EXPECT().List(gomock.Any(), "fakeId").Return([]compute.VirtualMachine{}, nil).Times(2)
mockVMSSClient := mockvmssclient.NewMockInterface(ctrl)
mockVMSSClient.EXPECT().List(gomock.Any(), "fakeId").Return([]compute.VirtualMachineScaleSet{}, nil).Times(2)
mockAzClient := &azClient{
virtualMachinesClient: mockVMClient,
virtualMachineScaleSetsClient: mockVMSSClient,
Expand Down Expand Up @@ -338,6 +340,7 @@ func TestCreateAzureManagerWithNilConfig(t *testing.T) {
mockVMClient := mockvmclient.NewMockInterface(ctrl)
mockVMSSClient := mockvmssclient.NewMockInterface(ctrl)
mockVMSSClient.EXPECT().List(gomock.Any(), "resourceGroup").Return([]compute.VirtualMachineScaleSet{}, nil).AnyTimes()
mockVMClient.EXPECT().List(gomock.Any(), "resourceGroup").Return([]compute.VirtualMachine{}, nil).AnyTimes()
mockAzClient := &azClient{
virtualMachinesClient: mockVMClient,
virtualMachineScaleSetsClient: mockVMSSClient,
Expand Down Expand Up @@ -663,7 +666,10 @@ func TestGetFilteredAutoscalingGroupsVmss(t *testing.T) {
expectedScaleSets := []compute.VirtualMachineScaleSet{fakeVMSSWithTags(vmssName, map[string]*string{vmssTag: &vmssTagValue, "min": &min, "max": &max})}
mockVMSSClient := mockvmssclient.NewMockInterface(ctrl)
mockVMSSClient.EXPECT().List(gomock.Any(), manager.config.ResourceGroup).Return(expectedScaleSets, nil).AnyTimes()
mockVMClient := mockvmclient.NewMockInterface(ctrl)
mockVMClient.EXPECT().List(gomock.Any(), manager.config.ResourceGroup).Return([]compute.VirtualMachine{}, nil).AnyTimes()
manager.azClient.virtualMachineScaleSetsClient = mockVMSSClient
manager.azClient.virtualMachinesClient = mockVMClient
err := manager.forceRefresh()
assert.NoError(t, err)

Expand Down Expand Up @@ -736,6 +742,9 @@ func TestFetchAutoAsgsVmss(t *testing.T) {
mockVMSSVMClient := mockvmssvmclient.NewMockInterface(ctrl)
mockVMSSVMClient.EXPECT().List(gomock.Any(), manager.config.ResourceGroup, vmssName, gomock.Any()).Return(expectedVMSSVMs, nil).AnyTimes()
manager.azClient.virtualMachineScaleSetVMsClient = mockVMSSVMClient
mockVMClient := mockvmclient.NewMockInterface(ctrl)
manager.azClient.virtualMachinesClient = mockVMClient
mockVMClient.EXPECT().List(gomock.Any(), manager.config.ResourceGroup).Return([]compute.VirtualMachine{}, nil).AnyTimes()
err := manager.forceRefresh()
assert.NoError(t, err)

Expand Down
Loading

0 comments on commit a7b4386

Please sign in to comment.