There will be times when you will want to start clearing down all active sessions on you Window Virtual Desktop hostpool hosts. This maybe because you want to perform some maintenance or migrate users to another pool.
The following script allows you to set and unset ‘Drain’ mode on all the Hosts or just a specific VM and below is an Automation Account runbook friendly version 😉
<#
.SYNOPSIS
Set the drain mode on VMs in a hostpool.
.DESCRIPTION
Used to enable or diable new sessions on a host
.PARAMETER HostPoolName
The name of the host pool that contains the VMs you want to remove th tag from
.PARAMETER HostPoolResourceGroupName
The name of the resource group that contains the Host Pool
.PARAMETER VMName
The name of a single VM to change the drain mode on. If left blank then all VMs in a hostpool will be set
.PARAMETER SetDrainModeOn
Sets the drain mode to On (not allow new sessions). If not set then the default is to set drain mode off (allow ne sessions)
.NOTES
Version: 1.0
Author: Nicholas Rogoff
Creation Date: 2020-09-17
Purpose/Change: Initial script development
.EXAMPLE
./SetDrainModeOnHostPoolVMs.ps1 -HostPoolName $HostPoolName -HostPoolResourceGroupName $HostPoolResourceGroupName -SetDrainModeOn
Sets drain mode on all hosts in the host pool that currently Allow New Sessions
.EXAMPLE
./SetDrainModeOnHostPoolVMs.ps1 -HostPoolName $HostPoolName -HostPoolResourceGroupName $HostPoolResourceGroupName
Sets 'Allow New Sessions' on all hosts in the host pool that currently Drain mode on
.EXAMPLE
./SetDrainModeOnHostPoolVMs.ps1 -HostPoolName $HostPoolName -HostPoolResourceGroupName $HostPoolResourceGroupName -VMName $VMName -SetDrainModeOn
Sets drain mode on a specific host in the host pool that currently Drain mode off
#>
#---------------------------------------------------------[Script Parameters]------------------------------------------------------
[CmdletBinding()]
Param (
[Parameter(mandatory = $true)]
[string]$HostPoolName,
[Parameter(mandatory = $true)]
[string]$HostPoolResourceGroupName,
[Parameter(mandatory = $false)]
[string]$VMName,
[Parameter(Mandatory = $False, HelpMessage = "Sets drain mode on. If not set then drain mode is turned off whereever it's found to be on")]
[switch]$SetDrainModeOn
)
#---------------------------------------------------------[Initialisations]--------------------------------------------------------
#Set Error Action to Silently Continue
$ErrorActionPreference = 'Continue'
#----------------------------------------------------------[Declarations]----------------------------------------------------------
#Any Global Declarations go here
#-----------------------------------------------------------[Functions]------------------------------------------------------------
function Get-VMNameFromSessionHost {
<#
.SYNOPSIS
Extracts the VM Name from the full SessionHost name returned in the SessionHost object
#>
[CmdletBinding()]
param (
$SessionHost
)
#Name is in the format 'uol-vwdgdev003-hp-dev-uks/vwdgdev003-0.ds.leeds.ac.uk' so need to split the last part
$VMName = $SessionHost.Name.Split("/")[1]
$VMName = $VMName.Split(".")[0]
return $VMName
}
#-----------------------------------------------------------[Execution]------------------------------------------------------------
Write-Output "Starting to set Drain Mode..."
$success = 0
$failed = 0
$skipped = 0
$sessionHosts = Get-AzWvdSessionHost -ResourceGroupName $HostPoolResourceGroupName -HostPoolName $HostPoolName
$sessionHosts | Select-Object Name, Status, Session, AllowNewSession
if ($SetDrainModeOn) {
Write-Output "Setting Drain Mode: OFF"
if ($VMName) {
#Get matching host with drain mode on
$HostPoolVMName = $HostPoolName + "/" + $VMName + ".ds.leeds.ac.uk"
$hostsToProcess = $sessionHosts | Where-Object Name -eq $HostPoolVMName | Where-Object AllowNewSession -eq $true
}
else {
#Get all hosts with drain mode on
$hostsToProcess = $sessionHosts | Where-Object AllowNewSession -eq $true
}
}
else {
if ($VMName) {
#Get matching host with drain mode on
$HostPoolVMName = $HostPoolName + "/" + $VMName + ".ds.leeds.ac.uk"
$hostsToProcess = $sessionHosts | Where-Object Name -eq $HostPoolVMName | Where-Object AllowNewSession -eq $false
}
else {
#Get all host with drain mode off
Write-Output "Setting Drain Mode: ON"
$hostsToProcess = $sessionHosts | Where-Object AllowNewSession -eq $false
}
}
foreach ($sh in $hostsToProcess) {
$Error.clear()
$VMName = Get-VMNameFromSessionHost($sh)
$thisAllowNewSessions = $sh.AllowNewSession
if ($SetDrainModeOn -and $thisAllowNewSessions) {
Write-Output "Setting " $VMName " to drain"
Update-AzWvdSessionHost -ResourceGroupName $HostPoolResourceGroupName `
-HostPoolName $HostPoolName `
-Name $sh.Name.Split("/")[1] `
-AllowNewSession:$false
}
elseif (! $SetDrainModeOn -and ! $thisAllowNewSessions) {
Write-Output "Setting " $VMName " to Allow New Sessions"
Update-AzWvdSessionHost -ResourceGroupName $HostPoolResourceGroupName `
-HostPoolName $HostPoolName `
-Name $sh.Name.Split("/")[1] `
-AllowNewSession:$true
}
else {
$skipped ++
}
if (!$Error[0]) {
$success += 1
}
else {
$failed += 1
Write-Error "!! Failed to set drain mode on $VMName"
}
}
Write-Output "================================="
Write-Output "Completed Drain Mode changes"
Write-Output "Succeeded: $success"
Write-Output "Failed: $failed"
Write-Output "================================="
And for a Azure Automation Account Runbook friendly version:
<#
.SYNOPSIS
Set the drain mode on VMs in a hostpool.
.DESCRIPTION
Used to enable or diable new sessions on a host
.PARAMETER SubscriptionName
The subscription name
.PARAMETER HostPoolName
The name of the host pool that contains the VMs you want to remove th tag from
.PARAMETER HostPoolResourceGroupName
The name of the resource group that contains the Host Pool
.PARAMETER VMName
The name of a single VM to change the drain mode on. If left blank then all VMs in a hostpool will be set
.PARAMETER SetDrainModeOn
Sets the drain mode to On (not allow new sessions). If not set then the default is to set drain mode off (allow ne sessions)
.NOTES
Version: 1.0
Author: Nicholas Rogoff
Creation Date: 2020-09-17
Purpose/Change: Initial script development
.EXAMPLE
./SetDrainModeOnHostPoolVMs.ps1 -HostPoolName $HostPoolName -HostPoolResourceGroupName $HostPoolResourceGroupName -SetDrainModeOn
Sets drain mode on all hosts in the host pool that currently Allow New Sessions
.EXAMPLE
./SetDrainModeOnHostPoolVMs.ps1 -HostPoolName $HostPoolName -HostPoolResourceGroupName $HostPoolResourceGroupName
Sets 'Allow New Sessions' on all hosts in the host pool that currently Drain mode on
.EXAMPLE
./SetDrainModeOnHostPoolVMs.ps1 -HostPoolName $HostPoolName -HostPoolResourceGroupName $HostPoolResourceGroupName -VMName $VMName -SetDrainModeOn
Sets drain mode on a specific host in the host pool that currently Drain mode off
#>
#---------------------------------------------------------[Script Parameters]------------------------------------------------------
[CmdletBinding()]
Param (
[Parameter(mandatory = $true)]
[string]$SubscriptionName,
[Parameter(mandatory = $true)]
[string]$HostPoolName,
[Parameter(mandatory = $true)]
[string]$HostPoolResourceGroupName,
[Parameter(mandatory = $false)]
[string]$VMName,
[Parameter(Mandatory = $False, HelpMessage = "Sets drain mode on. If not set then drain mode is turned off whereever it's found to be on")]
[switch]$SetDrainModeOn
)
#---------------------------------------------------------[Initialisations]--------------------------------------------------------
#Set Error Action to Silently Continue
$ErrorActionPreference = 'Continue'
Install-Module -Name Az.DesktopVirtualization
$connectionName = "AzureRunAsConnection"
try {
# Get the connection "AzureRunAsConnection "
$servicePrincipalConnection = Get-AutomationConnection -Name $connectionName
"Logging in to Azure..."
Add-AzAccount `
-ServicePrincipal `
-TenantId $servicePrincipalConnection.TenantId `
-ApplicationId $servicePrincipalConnection.ApplicationId `
-CertificateThumbprint $servicePrincipalConnection.CertificateThumbprint
}
catch {
if (!$servicePrincipalConnection) {
$ErrorMessage = "Connection $connectionName not found."
throw $ErrorMessage
}
else {
Write-Error -Message $_.Exception
throw $_.Exception
}
}
Select-AzSubscription $SubscriptionName
#----------------------------------------------------------[Declarations]----------------------------------------------------------
#Any Global Declarations go here
#-----------------------------------------------------------[Functions]------------------------------------------------------------
function Get-VMNameFromSessionHost {
<#
.SYNOPSIS
Extracts the VM Name from the full SessionHost name returned in the SessionHost object
#>
[CmdletBinding()]
param (
$SessionHost
)
#Name is in the format 'uol-vwdgdev003-hp-dev-uks/vwdgdev003-0.ds.leeds.ac.uk' so need to split the last part
$VMName = $SessionHost.Name.Split("/")[1]
$VMName = $VMName.Split(".")[0]
return $VMName
}
#-----------------------------------------------------------[Execution]------------------------------------------------------------
Write-Output "Starting to set Drain Mode..."
$success = 0
$failed = 0
$skipped = 0
$sessionHosts = Get-AzWvdSessionHost -ResourceGroupName $HostPoolResourceGroupName -HostPoolName $HostPoolName
$sessionHosts | Select-Object Name, Status, Session, AllowNewSession
if ($SetDrainModeOn) {
Write-Output "Setting Drain Mode: OFF"
if ($VMName) {
#Get matching host with drain mode on
$HostPoolVMName = $HostPoolName + "/" + $VMName + ".ds.leeds.ac.uk"
$hostsToProcess = $sessionHosts | Where-Object Name -eq $HostPoolVMName | Where-Object AllowNewSession -eq $true
}
else {
#Get all hosts with drain mode on
$hostsToProcess = $sessionHosts | Where-Object AllowNewSession -eq $true
}
}
else {
if ($VMName) {
#Get matching host with drain mode on
$HostPoolVMName = $HostPoolName + "/" + $VMName + ".ds.leeds.ac.uk"
$hostsToProcess = $sessionHosts | Where-Object Name -eq $HostPoolVMName | Where-Object AllowNewSession -eq $false
}
else {
#Get all host with drain mode off
Write-Output "Setting Drain Mode: ON"
$hostsToProcess = $sessionHosts | Where-Object AllowNewSession -eq $false
}
}
foreach ($sh in $hostsToProcess) {
$Error.clear()
$VMName = Get-VMNameFromSessionHost($sh)
$thisAllowNewSessions = $sh.AllowNewSession
if ($SetDrainModeOn -and $thisAllowNewSessions) {
Write-Output "Setting " $VMName " to drain"
Update-AzWvdSessionHost -ResourceGroupName $HostPoolResourceGroupName `
-HostPoolName $HostPoolName `
-Name $sh.Name.Split("/")[1] `
-AllowNewSession:$false
}
elseif (! $SetDrainModeOn -and ! $thisAllowNewSessions) {
Write-Output "Setting " $VMName " to Allow New Sessions"
Update-AzWvdSessionHost -ResourceGroupName $HostPoolResourceGroupName `
-HostPoolName $HostPoolName `
-Name $sh.Name.Split("/")[1] `
-AllowNewSession:$true
}
else {
$skipped ++
}
if (!$Error[0]) {
$success += 1
}
else {
$failed += 1
Write-Error "!! Failed to set drain mode on $VMName"
}
}
Write-Output "================================="
Write-Output "Completed Drain Mode changes"
Write-Output "Succeeded: $success"
Write-Output "Failed: $failed"
Write-Output "================================="