Deallocate VM on user logoff

Recently we've announced the public preview of start VM on connect which will allow your deallocated virtual machines getting started automatically when the assigned user tries to connect. Let us have a look how we can further optimize our cost by deallocating the VM when it is not used anymore.

Let's first have a high level view on the needed steps:

• Create a custom role to deallocate a virtual machine
• Create a managed identity for your virtual machines
• Allow each virtual machine to deallocate itself via a role assignment
• implement logoff script and policies around idle/disconnected sessions

So let's start by creating our custom role:

1. Open the Azure portal, go to Subscriptions and select the appropriate subscription

2. Go to Access control (IAM) and select Add a custom role.
   

3. Next, name the custom role and add a description. In this example I'll call it "Deallocate VM on logoff"

4. On the Permissions tab, add the following permission to the subscription you're assigning the role to:

   • Microsoft.Compute/virtualMachines/deallocate/action

5. When you're finished, select Ok.

If you prefer a JSON definition, please use the following template:

{
    "properties": {
        "roleName": "Deallocate VM on logoff",
        "description": "This custom role will allow your virtual machines to be deallocated when the user logs off.",
        "assignableScopes": [
            "/subscriptions/<<<SubscriptionID>>>"
        ],
        "permissions": [
            {
                "actions": [
                    "Microsoft.Compute/virtualMachines/deallocate/action"
                ],
                "notActions": [],
                "dataActions": [],
                "notDataActions": []
            }
        ]
    }
}

After we've created our custom role, we'll need to create a managed identity for our virtual machines. By this managed identities we don't need to store any credentials locally on the virtual machine or in an Azure KeyVault and can assign each virtual machine granular permission to shutdown only itself.

As this can be a bigger task, depending on the number of virtual machines you have in your personal host pools, I've prepared a script that will utilize the Azure PowerShell modules to assign that fine-grained permissions, so you may need to install those modules first:

Install-Module -Name Az.Account,Az.Compute,Az.DesktopVirtualization,Az.Resources

The script itself takes the host pool name, associated resource group and the role definition name selected above as parameters. It will then iterate through all virtual machines assigned to the specified host pool, create a managed identity when not already present and create a role assignment limited to the virtual machine itself:

$hostPoolName = "<<<HostPoolName>>>"
$resourceGroupName = "<<<ResourceGroupName>>>"
$roleDefinitionName = "<<<RoleDefinitionName>>>"

Connect-AzAccount
$sessionHosts = Get-AzWvdSessionHost -HostPoolName $hostPoolName -ResourceGroupName $resourceGroupName
foreach ($sessionHost in $sessionHosts) {
	<# get virtual machine by session host reference #>
	$resource =  Get-AzResource -ResourceId $sessionHost.ResourceId
	$vm = Get-AzVM -ResourceGroupName $resource.ResourceGroupName -Name $resource.Name
	
	<# create system-assigned managed identiy unless it already exists #>
	$managedIdentity = ($vm.Identity | where Type -eq "SystemAssigned").PrincipalId
	if ($managedIdentity -eq $Null) {
		Update-AzVM -ResourceGroupName $vm.ResourceGroupName -VM $vm -IdentityType SystemAssigned
		$managedIdentity = ((Get-AzVM -ResourceGroupName $vm.ResourceGroupName -VMName $vm.Name).Identity | where Type -eq "SystemAssigned").PrincipalId
	}
	
	<# create role-assignment unless it already exists #>
	if ((Get-AzRoleAssignment -RoleDefinitionName $roleDefinitionName -ObjectId $managedIdentity) -eq $Null) {
		New-AzRoleAssignment -ObjectId $managedIdentity -RoleDefinitionName $roleDefinitionName -Scope $vm.Id
	}
}

Next we'll configure our session host to disconnect idle sessions and logoff disconnected sessions after a certain period of time:

1. Connect remotely to the VM that you want to set the policy for.

2. Open the Group Policy Editor, then go to Local Computer Policy > Computer Configuration > Administrative Templates > Windows Components > Remote Desktop Services > Remote Desktop Session Host > Session Time Limits.

3. Find the policy that says Set time limit for disconnected sessions, then change its value to Enabled.

4. After you've enabled the policy, select your preferred time limit at End a disconnected session.

5. Find the policy that says Set time limit for active but idle Remote Desktop Services sessions, then change its value to Enabled.
6. After you've enabled the policy, select your preferred time limit at Idle session limit.

The above settings also ensure that an user will get a warning message two minutes before reaching the specified time limit so he can press a key or move the mouse to prevent getting disconnected.

In the last step we'll create the PowerShell script initiating the deallocation and configure it as a logoff script.

The PowerShell script will query the details of your virtual machine using the Azure instance metadata, connect to Azure using the created managed identity and initiate the actual deallocation via REST API:

$metadata = Invoke-RestMethod -Headers @{"Metadata"="true"} -Method GET -Proxy $Null -Uri "http://169.254.169.254/metadata/instance?api-version=2021-01-01"
$authorizationToken = Invoke-RestMethod -Headers @{"Metadata"="true"} -Method Get -Proxy $Null -Uri "http://169.254.169.254/metadata/identity/oauth2/token?api-version=2021-01-01&resource=https://management.azure.com/"

$subscriptionId = $metadata.compute.subscriptionId
$resourceGroupName = $metadata.compute.resourceGroupName
$vmName = $metadata.compute.name
$accessToken = $authorizationToken.access_token

$RestartEvents = Get-EventLog -LogName System -After (Get-Date).AddMinutes(-1) |? {($_.EventID -eq 1074) -and ($_.Message -match "restart" )}
$SessionCount = (query user | Measure-Object | select Count).count - 1 # remove headline

if (($SessionCount -gt 1) -or ($RestartEvents.count -ge 1))
{
	# skip deallocate because of user-sessions or initiated reboot
} else {
	Invoke-WebRequest -UseBasicParsing -Headers @{ Authorization ="Bearer $accessToken"} -Method POST -Proxy $Null -Uri https://management.azure.com/subscriptions/$subscriptionId/resourceGroups/$resourceGroupName/providers/Microsoft.Compute/virtualMachines/$vmName/deallocate?api-version=2021-03-01 -ContentType "application/json"
}

To have that script executed upon logoff, we need to configure it:

1. Connect remotely to the VM that you want to set the policy for.

2. Open the Group Policy Editor, then go to Local Computer Policy > User Configuration > Windows Settings > Scripts (Logon/Logoff).

3. Find the item that says Logoff.

4. Specify the script you've created on the PowerShell Scripts tab.

Place

When rolling out to bigger host-pools, ideally place the logoff script on a centrally accessible location, e.g. SYSVOL-share.

Hope this short tutorial will help you take full benefit of the start VM on connect feature. Happy to read your feedback and comments below.