Skip to content

Commit b22425d

Browse files
YanaXuYanaXu
authored
[AKS] support NodeTaints and AksCustomHeaders (#20319)
* [AKS] support NodeTaints * [AKS] support AksCustomHeader
1 parent a6ab979 commit b22425d

16 files changed

+3306
-4437
lines changed

‎src/Aks/Aks.Test/ScenarioTests/KubernetesTests.cs

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -93,5 +93,12 @@ public void TestNodeLabelsAndTags()
9393
{
9494
TestRunner.RunTestScript("Test-NodeLabels-Tags");
9595
}
96+
97+
[Fact]
98+
[Trait(Category.AcceptanceType, Category.CheckIn)]
99+
public void TestNodeTaints()
100+
{
101+
TestRunner.RunTestScript("Test-NodeTaints");
102+
}
96103
}
97104
}

‎src/Aks/Aks.Test/ScenarioTests/KubernetesTests.ps1

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -487,4 +487,54 @@ function Test-NodeLabels-Tags {
487487
finally {
488488
Remove-AzResourceGroup -Name $resourceGroupName -Force
489489
}
490+
}
491+
492+
function Test-NodeTaints {
493+
# Setup
494+
$resourceGroupName = Get-RandomResourceGroupName
495+
$kubeClusterName = Get-RandomClusterName
496+
$location = 'eastus'
497+
$nodeVmSize = "Standard_D2_v2"
498+
499+
try {
500+
New-AzResourceGroup -Name $resourceGroupName -Location $location
501+
502+
# create aks cluster with default nodepool
503+
New-AzAksCluster -ResourceGroupName $resourceGroupName -Name $kubeClusterName -NodeVmSize $nodeVmSize -NodeCount 1
504+
$cluster = Get-AzAksCluster -ResourceGroupName $resourceGroupName -Name $kubeClusterName
505+
Assert-AreEqual 1 $cluster.AgentPoolProfiles.Count
506+
Assert-AreEqual 0 $cluster.AgentPoolProfiles[0].NodeTaints.Count
507+
$pools = Get-AzAksNodePool -ResourceGroupName $resourceGroupName -ClusterName $kubeClusterName
508+
Assert-AreEqual 1 $pools.Count
509+
Assert-AreEqual 0 $pools.NodeTaints.Count
510+
511+
# create a 2nd nodepool
512+
$nodetains = @("sku=gpu:NoSchedule")
513+
New-AzAksNodePool -ResourceGroupName $resourceGroupName -ClusterName $kubeClusterName -Name "pool2" -Count 1 -NodeTaint $nodetains
514+
$cluster = Get-AzAksCluster -ResourceGroupName $resourceGroupName -Name $kubeClusterName
515+
Assert-AreEqual 2 $cluster.AgentPoolProfiles.Count
516+
Assert-AreEqual 0 ($cluster.AgentPoolProfiles | where {$_.Name -eq "default"}).NodeTaints.Count
517+
Assert-AreEqualArray $nodetains ($cluster.AgentPoolProfiles | where {$_.Name -eq "pool2"}).NodeTaints
518+
$pools = Get-AzAksNodePool -ResourceGroupName $resourceGroupName -ClusterName $kubeClusterName
519+
Assert-AreEqual 2 $pools.Count
520+
Assert-AreEqual 0 ($pools | where {$_.Name -eq "default"}).NodeTaints.Count
521+
Assert-AreEqualArray $nodetains ($pools | where {$_.Name -eq "pool2"}).NodeTaints
522+
523+
# update the 2nd nodepool
524+
$nodetains2 = @("CriticalAddonsOnly=true:NoSchedule")
525+
Update-AzAksNodePool -ResourceGroupName $resourceGroupName -ClusterName $kubeClusterName -Name "pool2" -NodeTaint $nodetains2
526+
$cluster = Get-AzAksCluster -ResourceGroupName $resourceGroupName -Name $kubeClusterName
527+
Assert-AreEqual 2 $cluster.AgentPoolProfiles.Count
528+
Assert-AreEqual 0 ($cluster.AgentPoolProfiles | where {$_.Name -eq "default"}).NodeTaints.Count
529+
Assert-AreEqualArray $nodetains2 ($cluster.AgentPoolProfiles | where {$_.Name -eq "pool2"}).NodeTaints
530+
$pools = Get-AzAksNodePool -ResourceGroupName $resourceGroupName -ClusterName $kubeClusterName
531+
Assert-AreEqual 2 $pools.Count
532+
Assert-AreEqual 0 ($pools | where {$_.Name -eq "default"}).NodeTaints.Count
533+
Assert-AreEqualArray $nodetains2 ($pools | where {$_.Name -eq "pool2"}).NodeTaints
534+
535+
$cluster | Remove-AzAksCluster -Force
536+
}
537+
finally {
538+
Remove-AzResourceGroup -Name $resourceGroupName -Force
539+
}
490540
}

‎src/Aks/Aks.Test/SessionRecords/Commands.Aks.Test.ScenarioTests.KubernetesTests/TestNodeLabels.json

Lines changed: 0 additions & 4366 deletions
This file was deleted.

‎src/Aks/Aks.Test/SessionRecords/Commands.Aks.Test.ScenarioTests.KubernetesTests/TestNodeTaints.json

Lines changed: 3058 additions & 0 deletions
Large diffs are not rendered by default.

‎src/Aks/Aks/ChangeLog.md

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,8 @@
2424
* Added property `-NodeImageVersion` for the output of `Get-AzAksNodePool`[#19893]
2525
* Added parameter `-NodePoolLabel` for `Set-AzAksCluster`, `-NodeLabel` for `New-AzAksNodePool` and `Update-AzAksNodePool`
2626
* Added parameter `-NodePoolTag` for `New-AzAksCluster` and `Set-AzAksCluster`, `-Tag` for `New-AzAksNodePool` and `Update-AzAksNodePool`
27+
* Added parameter `-NodeTaint` for `New-AzAksNodePool` and `Update-AzAksNodePool`
28+
* Added parameter `-AksCustomHeader` for `Set-AzAksCluster`, `New-AzAksNodePool` and `Update-AzAksNodePool`
2729

2830
## Version 5.0.1
2931
* Upgraded AutoMapper to Microsoft.Azure.PowerShell.AutoMapper 6.2.2 with fix [#18721]

‎src/Aks/Aks/Commands/CreateOrUpdateKubeBase.cs

Lines changed: 24 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@
3939
using Microsoft.Azure.Commands.Common.MSGraph.Version1_0.Applications;
4040
using Microsoft.Azure.Commands.Common.MSGraph.Version1_0;
4141
using ResourceIdentityType = Microsoft.Azure.Management.ContainerService.Models.ResourceIdentityType;
42+
using Microsoft.Azure.Commands.Aks.Commands;
4243

4344
namespace Microsoft.Azure.Commands.Aks
4445
{
@@ -190,6 +191,9 @@ public abstract class CreateOrUpdateKubeBase : KubeCmdletBase
190191
[Parameter(Mandatory = false, HelpMessage = "Alternative CA cert to use for connecting to proxy servers.")]
191192
public string HttpProxyConfigTrustedCa { get; set; }
192193

194+
[Parameter(Mandatory = false, HelpMessage = "Aks custom headers used for building Kubernetes network.")]
195+
public Hashtable AksCustomHeader { get; set; }
196+
193197
protected void BeforeBuildNewCluster()
194198
{
195199
if (!string.IsNullOrEmpty(ResourceGroupName) && string.IsNullOrEmpty(Location))
@@ -269,7 +273,7 @@ protected string GetSshKey(string sshKeyOrFile)
269273
protected AcsServicePrincipal EnsureServicePrincipal(string spId = null, string clientSecret = null)
270274
{
271275
//If user specifies service principal, just use it directly and no need to save to disk
272-
if(!string.IsNullOrEmpty(spId) && !string.IsNullOrEmpty(clientSecret))
276+
if (!string.IsNullOrEmpty(spId) && !string.IsNullOrEmpty(clientSecret))
273277
{
274278
return new AcsServicePrincipal()
275279
{
@@ -289,7 +293,7 @@ protected AcsServicePrincipal EnsureServicePrincipal(string spId = null, string
289293
if (clientSecret == null)
290294
{
291295
clientSecret = RandomBase64String(16);
292-
}
296+
}
293297

294298
acsServicePrincipal = BuildServicePrincipal(Name, clientSecret);
295299
WriteVerbose(Resources.CreatedANewServicePrincipalAndAssignedTheContributorRole);
@@ -365,7 +369,7 @@ protected void AddAcrRoleAssignment(string acrName, string acrParameterName, Acs
365369
var acrObjects = RmClient.Resources.List(acrQuery);
366370
acrResourceId = acrObjects.First().Id;
367371
}
368-
catch(Exception)
372+
catch (Exception)
369373
{
370374
throw new AzPSArgumentException(
371375
string.Format(Resources.CouldNotFindSpecifiedAcr, acrName),
@@ -389,12 +393,12 @@ protected void AddAcrRoleAssignment(string acrName, string acrParameterName, Acs
389393
var servicePrincipal = GraphClient.FilterServicePrincipals(oDataQuery).First();
390394
spObjectId = servicePrincipal.Id;
391395
}
392-
catch(Exception ex)
396+
catch (Exception ex)
393397
{
394398
throw new AzPSInvalidOperationException(
395399
string.Format(Resources.CouldNotFindObjectIdForServicePrincipal, acsServicePrincipal.SpId),
396400
ex,
397-
string.Format(Resources.CouldNotFindObjectIdForServicePrincipal,"*"));
401+
string.Format(Resources.CouldNotFindObjectIdForServicePrincipal, "*"));
398402
}
399403
}
400404
var success = RetryAction(() =>
@@ -472,7 +476,7 @@ protected bool RetryAction(Action action, string actionName = null)
472476
protected AcsServicePrincipal LoadServicePrincipal()
473477
{
474478
var config = LoadServicePrincipals();
475-
if(config?.ContainsKey(DefaultContext.Subscription.Id) == true)
479+
if (config?.ContainsKey(DefaultContext.Subscription.Id) == true)
476480
{
477481
return config[DefaultContext.Subscription.Id];
478482
}
@@ -518,7 +522,7 @@ protected string DefaultDnsPrefix()
518522
var subPart = string.Join("", DefaultContext.Subscription.Id.Take(4));
519523
return $"{namePart}{subPart}";
520524
}
521-
525+
522526
protected ManagedClusterLoadBalancerProfile CreateOrUpdateLoadBalancerProfile(ManagedClusterLoadBalancerProfile loadBalancerProfile)
523527
{
524528
if ((this.IsParameterBound(c => c.LoadBalancerManagedOutboundIpCount) ||
@@ -666,5 +670,18 @@ protected ManagedCluster SetIdentity(ManagedCluster cluster)
666670

667671
return cluster;
668672
}
673+
674+
private protected ManagedCluster CreateOrUpdate(string resourceGroupName, string resourceName, ManagedCluster parameters)
675+
{
676+
if (this.IsParameterBound(c => c.AksCustomHeader))
677+
{
678+
Dictionary<string, List<string>> customHeaders = Utilities.HashtableToDictionary(AksCustomHeader);
679+
return Client.ManagedClusters.CreateOrUpdateWithHttpMessagesAsync(resourceGroupName, resourceName, parameters, customHeaders).GetAwaiter().GetResult().Body;
680+
}
681+
else
682+
{
683+
return Client.ManagedClusters.CreateOrUpdate(resourceGroupName, resourceName, parameters);
684+
}
685+
}
669686
}
670687
}

‎src/Aks/Aks/Commands/NewAzureRmAks.cs

Lines changed: 2 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,6 @@
1313
// ----------------------------------------------------------------------------------
1414

1515
using System;
16-
using System.Collections;
1716
using System.Collections.Generic;
1817
using System.ComponentModel;
1918
using System.Diagnostics;
@@ -28,7 +27,6 @@
2827
using Microsoft.Azure.Commands.Common.Exceptions;
2928
using Microsoft.Azure.Commands.ResourceManager.Common.ArgumentCompleters;
3029
using Microsoft.Azure.Commands.ResourceManager.Common.Tags;
31-
using Microsoft.Azure.Management.ContainerService;
3230
using Microsoft.Azure.Management.ContainerService.Models;
3331
using Microsoft.Rest;
3432
using Microsoft.WindowsAzure.Commands.Common;
@@ -113,10 +111,6 @@ public class NewAzureRmAks : CreateOrUpdateKubeBase
113111
[Parameter(Mandatory = false, HelpMessage = "Docker bridge cidr used for building Kubernetes network.")]
114112
public string DockerBridgeCidr { get; set; }
115113

116-
[Parameter(Mandatory = false, HelpMessage = "Aks custom headers used for building Kubernetes network.")]
117-
118-
public Hashtable AksCustomHeader { get; set; }
119-
120114
[Parameter(Mandatory = false, HelpMessage = "The load balancer sku for the managed cluster.")]
121115
[PSArgumentCompleter("basic", "standard")]
122116
public string LoadBalancerSku { get; set; }
@@ -157,26 +151,8 @@ public override void ExecuteCmdlet()
157151
var managedCluster = BuildNewCluster();
158152
try
159153
{
160-
ManagedCluster cluster;
161-
if (this.IsParameterBound(c => c.AksCustomHeader))
162-
{
163-
Dictionary<string, List<string>> customHeaders = new Dictionary<string, List<string>>();
164-
foreach (var key in AksCustomHeader.Keys)
165-
{
166-
List<string> values = new List<string>();
167-
foreach (var value in (object[])AksCustomHeader[key])
168-
{
169-
values.Add(value.ToString());
170-
}
171-
customHeaders.Add(key.ToString(), values);
172-
}
173-
174-
cluster = Client.ManagedClusters.CreateOrUpdateWithHttpMessagesAsync(ResourceGroupName, Name, managedCluster, customHeaders).GetAwaiter().GetResult().Body;
175-
}
176-
else
177-
{
178-
cluster = Client.ManagedClusters.CreateOrUpdate(ResourceGroupName, Name, managedCluster);
179-
}
154+
155+
var cluster = this.CreateOrUpdate(ResourceGroupName, Name, managedCluster);
180156
var psObj = PSMapper.Instance.Map<PSKubernetesCluster>(cluster);
181157

182158
if (this.IsParameterBound(c => c.AcrNameToAttach))

‎src/Aks/Aks/Commands/NewAzureRmAksNodePool.cs

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,6 @@
1515
using System;
1616
using System.Collections.Generic;
1717
using System.Management.Automation;
18-
1918
using Microsoft.Azure.Commands.Aks.Models;
2019
using Microsoft.Azure.Commands.Aks.Properties;
2120
using Microsoft.Azure.Commands.Common.Exceptions;
@@ -113,7 +112,7 @@ public override void ExecuteCmdlet()
113112
ClusterName = ClusterObject.Name;
114113
}
115114
var agentPool = GetAgentPool();
116-
var pool = Client.AgentPools.CreateOrUpdate(ResourceGroupName, ClusterName, Name, agentPool);
115+
var pool = this.CreateOrUpdate(ResourceGroupName, ClusterName, Name, agentPool);
117116
var psPool = PSMapper.Instance.Map<PSNodePool>(pool);
118117
WriteObject(psPool);
119118
};
@@ -211,6 +210,10 @@ private AgentPool GetAgentPool()
211210
agentPool.Tags.Add(key.ToString(), Tag[key].ToString());
212211
}
213212
}
213+
if (this.IsParameterBound(c => c.NodeTaint))
214+
{
215+
agentPool.NodeTaints = NodeTaint;
216+
}
214217

215218
return agentPool;
216219
}

‎src/Aks/Aks/Commands/NewOrUpdateAgentPoolBase.cs

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,11 @@
1212
// limitations under the License.
1313
// ----------------------------------------------------------------------------------
1414

15+
using Microsoft.Azure.Commands.Aks.Commands;
1516
using Microsoft.Azure.Commands.ResourceManager.Common.ArgumentCompleters;
17+
using Microsoft.Azure.Management.ContainerService;
18+
using Microsoft.Azure.Management.ContainerService.Models;
19+
using Microsoft.WindowsAzure.Commands.Utilities.Common;
1620
using Newtonsoft.Json;
1721
using System.Collections;
1822
using System.Collections.Generic;
@@ -47,5 +51,24 @@ public class NewOrUpdateAgentPoolBase : KubeCmdletBase
4751

4852
[Parameter(Mandatory = false, HelpMessage = "The tags to be persisted on the agent pool virtual machine scale set.")]
4953
public Hashtable Tag { get; set; }
54+
55+
[Parameter(Mandatory = false, HelpMessage = "The node taints added to new nodes during node pool create and scale")]
56+
public string[] NodeTaint { get; set; }
57+
58+
[Parameter(Mandatory = false, HelpMessage = "Aks custom headers")]
59+
public Hashtable AksCustomHeader { get; set; }
60+
61+
private protected AgentPool CreateOrUpdate(string resourceGroupName, string resourceName, string agentPoolName, AgentPool parameters)
62+
{
63+
if (this.IsParameterBound(c => c.AksCustomHeader))
64+
{
65+
Dictionary<string, List<string>> customHeaders = Utilities.HashtableToDictionary(AksCustomHeader);
66+
return Client.AgentPools.CreateOrUpdateWithHttpMessagesAsync(resourceGroupName, resourceName, agentPoolName, parameters, customHeaders).GetAwaiter().GetResult().Body;
67+
}
68+
else
69+
{
70+
return Client.AgentPools.CreateOrUpdate(resourceGroupName, resourceName, agentPoolName, parameters);
71+
}
72+
}
5073
}
5174
}

‎src/Aks/Aks/Commands/SetAzureRmAks.cs

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@
1616
using System.Collections.Generic;
1717
using System.Linq;
1818
using System.Management.Automation;
19-
2019
using Microsoft.Azure.Commands.Aks.Models;
2120
using Microsoft.Azure.Commands.Aks.Properties;
2221
using Microsoft.Azure.Commands.Common.Exceptions;
@@ -31,11 +30,8 @@
3130
using Microsoft.Azure.Management.Internal.Resources.Utilities.Models;
3231
using Microsoft.Rest.Azure.OData;
3332
using Microsoft.WindowsAzure.Commands.Common;
34-
using Microsoft.WindowsAzure.Commands.Common.CustomAttributes;
3533
using Microsoft.WindowsAzure.Commands.Utilities.Common;
3634

37-
using ResourceIdentityType = Microsoft.Azure.Management.ContainerService.Models.ResourceIdentityType;
38-
3935
namespace Microsoft.Azure.Commands.Aks
4036
{
4137
[Cmdlet("Set", ResourceManager.Common.AzureRMConstants.AzureRMPrefix + "AksCluster", DefaultParameterSetName = DefaultParamSet, SupportsShouldProcess = true)]
@@ -401,7 +397,7 @@ public override void ExecuteCmdlet()
401397
}
402398
SetIdentity(cluster);
403399

404-
var kubeCluster = Client.ManagedClusters.CreateOrUpdate(ResourceGroupName, Name, cluster);
400+
var kubeCluster = this.CreateOrUpdate(ResourceGroupName, Name, cluster);
405401

406402
if (this.IsParameterBound(c => c.DiskEncryptionSetID))
407403
{

‎src/Aks/Aks/Commands/UpdateAzureRmAksNodePool.cs

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -180,8 +180,12 @@ public override void ExecuteCmdlet()
180180
pool.Tags.Add(key.ToString(), Tag[key].ToString());
181181
}
182182
}
183+
if (this.IsParameterBound(c => c.NodeTaint))
184+
{
185+
pool.NodeTaints = NodeTaint;
186+
}
183187

184-
var updatedPool = Client.AgentPools.CreateOrUpdate(ResourceGroupName, ClusterName, Name, pool);
188+
var updatedPool = this.CreateOrUpdate(ResourceGroupName, ClusterName, Name, pool);
185189
WriteObject(PSMapper.Instance.Map<PSNodePool>(updatedPool));
186190
});
187191
}

‎src/Aks/Aks/Commands/Utilities.cs

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,8 @@
1313
// ----------------------------------------------------------------------------------
1414

1515
using System;
16-
16+
using System.Collections;
17+
using System.Collections.Generic;
1718
using Microsoft.Azure.Commands.Common.Exceptions;
1819

1920
namespace Microsoft.Azure.Commands.Aks.Commands
@@ -32,5 +33,24 @@ public static string GetParentResourceName(string parentResource, string paramet
3233

3334
return items.Length > 0 ? items[items.Length - 1] : null;
3435
}
36+
37+
public static Dictionary<string, List<string>> HashtableToDictionary(Hashtable table)
38+
{
39+
Dictionary<string, List<string>> dict = new Dictionary<string, List<string>>();
40+
if (table == null)
41+
{
42+
return dict;
43+
}
44+
foreach (var key in table.Keys)
45+
{
46+
List<string> values = new List<string>();
47+
foreach (var value in (object[])table[key])
48+
{
49+
values.Add(value.ToString());
50+
}
51+
dict.Add(key.ToString(), values);
52+
}
53+
return dict;
54+
}
3555
}
3656
}

0 commit comments

Comments
 (0)