@@ -16,19 +16,26 @@ function Stop-DisplayContainer {
     # A crash-looping display-container process can race the subsequent device
     # cycle and prevent the NVIDIA kernel driver from unloading.
     Write-Output "Stopping NVDisplay.ContainerLocalSystem..."
-    Stop-Service -Name "NVDisplay.ContainerLocalSystem" -Force -ErrorAction Stop
-
-    Write-Output "Waiting for NVDisplay.Container processes to exit..."
+    $service = Get-Service -Name "NVDisplay.ContainerLocalSystem" -ErrorAction Stop
     $deadline = (Get-Date).AddSeconds(60)
-    $container_processes = @(Get-Process -Name "NVDisplay.Container" -ErrorAction SilentlyContinue)
-    while ($container_processes.Count -gt 0 -and (Get-Date) -lt $deadline) {
-        Start-Sleep -Milliseconds 500
-        $container_processes = @(Get-Process -Name "NVDisplay.Container" -ErrorAction SilentlyContinue)
+    if ($service.Status -ne "Stopped" -and $service.Status -ne "StopPending") {
+        Stop-Service -InputObject $service -Force -NoWait -ErrorAction Stop
     }
-    if ($container_processes.Count -gt 0) {
-        $process_ids = $container_processes.Id -join ", "
-        Write-Error "NVDisplay.Container processes ($process_ids) did not exit within 60s of stopping NVDisplay.ContainerLocalSystem"
-        exit 1
+
+    Write-Output "Waiting for NVDisplay.ContainerLocalSystem to stop and NVDisplay.Container processes to exit..."
+    while ($true) {
+        $service.Refresh()
+        $container_processes = @(Get-Process -Name "NVDisplay.Container" -ErrorAction SilentlyContinue)
+        if ($service.Status -eq "Stopped" -and $container_processes.Count -eq 0) {
+            break
+        }
+        if ((Get-Date) -ge $deadline) {
+            $process_ids = $container_processes.Id -join ", "
+            if (-not $process_ids) { $process_ids = "none" }
+            Write-Error "Timed out after 60s waiting for the NVIDIA display container (service status: $($service.Status); remaining NVDisplay.Container PIDs: $process_ids)"
+            exit 1
+        }
+        Start-Sleep -Milliseconds 500
     }
 }