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 "================================="