Remove Ubuntu install disk after install seems done

Author: nyanhp

AutomatedLabWorker/functions/VirtualMachines/New-LWHypervVM.ps1

@@ -556,36 +556,6 @@ rm -rf /etc/cron.d/postconf
 
     $vm = Hyper-V\New-VM @vmParameter
 
-    if ($vm.Generation -ge 2 -and $Machine.OperatingSystemType -eq 'Linux' -and $Machine.LinuxType -eq 'Ubuntu')
-    {
-        $systemDisk = $vm | Get-VMHardDiskDrive
-        $ubuntuSpecialDisk = $vm | Add-VMHardDiskDrive -Path "$vmPath\$($Machine.ResourceName)_INSTALL.vhdx" -Passthru
-        $vm | Set-VMFirmware -BootOrder $systemDisk, $ubuntuSpecialDisk
-
-        Start-Job -Name "Ensure Ubuntu Boot Order $($Machine.ResourceName)" -ScriptBlock {
-            param
-            (
-                $Name
-            )
-
-            # Boot order of Ubuntu system changes during installation. We need to ensure the OS
-            # is booted, lest it is stuck in an installation loop.
-            # No cancellation token, we'll remove the job if the lab installation has finished.
-            while ($true)
-            {
-                Write-Verbose "Reordering boot order for Ubuntu VM '$Name'"
-                $linvm = Get-LWHypervVM -Name $Name
-                $order = ($linvm | Get-VMFirmware).BootOrder
-                if ($order[0].BootType -eq 'Drive' -and $order[0].Device.Path -like "*_INSTALL*")
-                {
-                    $newOrder = $order[1..$order.Count] + $order[0]
-                    $linvm | Set-VMFirmware -BootOrder $newOrder
-                }
-                Start-Sleep -Seconds 1
-            }
-        } -ArgumentList $Machine.ResourceName
-    }
-
     Set-LWHypervVMDescription -ComputerName $Machine.ResourceName -Hashtable @{
         CreatedBy    = '{0} ({1})' -f $PSCmdlet.MyInvocation.MyCommand.Module.Name, $PSCmdlet.MyInvocation.MyCommand.Module.Version
         CreationTime = Get-Date

AutomatedLabWorker/functions/VirtualMachines/Start-LWHypervVM.ps1

@@ -24,10 +24,27 @@
     foreach ($Name in $(Get-LabVM -ComputerName $ComputerName -IncludeLinux | Where-Object SkipDeployment -eq $false))
     {
         $machine = Get-LabVM -ComputerName $Name -IncludeLinux
+        $hvMachine = Get-LWHypervVM -Name $Name.ResourceName
+
+        <#
+            Remove _INSTALL.vhdx on Ubuntu if the VM has been shut down once - indicating that the
+            cloudinit/subiquity phase was successfully finished.
+            We compare GuestStatePath LastWriteTime as a simple and quick way to check if
+            the VM's status has changed i.e. when it was stopped.
+            One minute seems like a sane interval, we might need to increase it in the future.
+        #>
+        if ($hvMachine.State -ne 'Running' -and ((Get-Item -Path $hvMachine.GuestStatePath).LastWriteTime - $hvMachine.CreationTime) -gt '00:01:00')
+        {
+            Write-ScreenInfo -Type Verbose "Removing installation disk '$Name'"
+            $disk = $hvMachine | Get-VMHardDiskDrive | Where-Object Path -like "*_INSTALL*"
+            $diskPath = $disk.Path # Otherwise $disk will be update after remove-vmharddiskdrive was called
+            $disk | Remove-VMHardDiskDrive
+            Remove-Item -Path $diskPath -Force
+        }
 
         try
         {
-            Get-LWHypervVM -Name $Name.ResourceName | Hyper-V\Start-VM -ErrorAction Stop
+            $hvMachine | Hyper-V\Start-VM -ErrorAction Stop
         }
         catch
         {