Skip to content

Commit 55cbcbb

Browse files
jpascucci-nvjpascucci-nv
committed
CI: prevent display container restart during mode switching
Temporarily disable NVDisplay.ContainerLocalSystem while changing driver modes so queued recovery actions cannot restart it during the device cycle. Handle in-flight starts, restore the original startup type in a finally block, and report the requested driver mode by name.
1 parent bb9970e commit 55cbcbb

1 file changed

Lines changed: 48 additions & 36 deletions

File tree

‎ci/tools/configure_driver_mode.ps1‎

Lines changed: 48 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -12,28 +12,32 @@
1212
# DRIVER_MODE One of WDDM, TCC, MCDM.
1313

1414
function Stop-DisplayContainer {
15+
param (
16+
[System.ServiceProcess.ServiceController]$Service
17+
)
1518

1619
# A crash-looping display-container process can race the subsequent device
1720
# cycle and prevent the NVIDIA kernel driver from unloading.
1821
Write-Output "Stopping NVDisplay.ContainerLocalSystem..."
19-
$service = Get-Service -Name "NVDisplay.ContainerLocalSystem" -ErrorAction Stop
2022
$deadline = (Get-Date).AddSeconds(60)
21-
if ($service.Status -ne "Stopped" -and $service.Status -ne "StopPending") {
22-
Stop-Service -InputObject $service -Force -NoWait -ErrorAction Stop
23-
}
2423

2524
Write-Output "Waiting for NVDisplay.ContainerLocalSystem to stop and NVDisplay.Container processes to exit..."
2625
while ($true) {
27-
$service.Refresh()
26+
$Service.Refresh()
27+
# A recovery start may already be in flight when the service is
28+
# disabled. Stop it as soon as it reaches a stoppable state.
29+
if ($Service.Status -eq "Running" -or $Service.Status -eq "Paused") {
30+
Stop-Service -InputObject $Service -Force -NoWait -ErrorAction Stop
31+
$Service.Refresh()
32+
}
2833
$container_processes = @(Get-Process -Name "NVDisplay.Container" -ErrorAction SilentlyContinue)
29-
if ($service.Status -eq "Stopped" -and $container_processes.Count -eq 0) {
34+
if ($Service.Status -eq "Stopped" -and $container_processes.Count -eq 0) {
3035
break
3136
}
3237
if ((Get-Date) -ge $deadline) {
3338
$process_ids = $container_processes.Id -join ", "
3439
if (-not $process_ids) { $process_ids = "none" }
35-
Write-Error "Timed out after 60s waiting for the NVIDIA display container (service status: $($service.Status); remaining NVDisplay.Container PIDs: $process_ids)"
36-
exit 1
40+
throw "Timed out after 60s waiting for the NVIDIA display container (service status: $($Service.Status); remaining NVDisplay.Container PIDs: $process_ids)"
3741
}
3842
Start-Sleep -Milliseconds 500
3943
}
@@ -45,48 +49,56 @@ function Set-DriverMode {
4549
# This assumes we have the prior knowledge on which GPU can use which mode.
4650
$driver_mode = $env:DRIVER_MODE
4751
if ($driver_mode -eq "WDDM") {
48-
Write-Output "Setting driver mode to WDDM..."
4952
$driver_mode_code = 0
5053
} elseif ($driver_mode -eq "TCC") {
51-
Write-Output "Setting driver mode to TCC..."
5254
$driver_mode_code = 1
5355
} elseif ($driver_mode -eq "MCDM") {
54-
Write-Output "Setting driver mode to MCDM..."
5556
$driver_mode_code = 2
5657
} else {
5758
Write-Output "Unknown driver mode: $driver_mode"
5859
exit 1
5960
}
6061

61-
Stop-DisplayContainer
62+
$display_container_service = Get-Service -Name "NVDisplay.ContainerLocalSystem" -ErrorAction Stop
63+
$display_container_start_type = $display_container_service.StartType
64+
try {
65+
# A queued recovery restart cannot be canceled, so keep the service
66+
# disabled until the device cycle and NVML recovery are complete.
67+
Write-Output "Temporarily disabling NVDisplay.ContainerLocalSystem..."
68+
Set-Service -InputObject $display_container_service -StartupType Disabled -ErrorAction Stop
69+
Stop-DisplayContainer -Service $display_container_service
6270

63-
nvidia-smi -fdm $driver_mode_code
71+
Write-Output "Setting driver mode to $driver_mode..."
72+
nvidia-smi -fdm $driver_mode_code
6473

65-
# Only restart NVIDIA display adapters, not other display devices (e.g. QEMU VGA).
66-
$nvidia_devices = @(Get-PnpDevice -Class Display -FriendlyName "NVIDIA*")
67-
$gpu_count = $nvidia_devices.Count
68-
foreach ($device in $nvidia_devices) {
69-
Write-Output "Restarting device: $($device.FriendlyName) ($($device.InstanceId))"
70-
pnputil /disable-device "$($device.InstanceId)"
71-
pnputil /enable-device "$($device.InstanceId)"
72-
}
74+
# Only restart NVIDIA display adapters, not other display devices (e.g. QEMU VGA).
75+
$nvidia_devices = @(Get-PnpDevice -Class Display -FriendlyName "NVIDIA*")
76+
$gpu_count = $nvidia_devices.Count
77+
foreach ($device in $nvidia_devices) {
78+
Write-Output "Restarting device: $($device.FriendlyName) ($($device.InstanceId))"
79+
pnputil /disable-device "$($device.InstanceId)"
80+
pnputil /enable-device "$($device.InstanceId)"
81+
}
7382

74-
# Initial settle after the device cycle.
75-
Start-Sleep -Seconds 5
83+
# Initial settle after the device cycle.
84+
Start-Sleep -Seconds 5
7685

77-
# Poll nvidia-smi for N consecutive successes (N == cycled GPUs)
78-
# so a mid-init "ok" flap doesn't fool the loop; bail after ~60s.
79-
Write-Output "Waiting for nvidia-smi/NVML to come back up after device cycle..."
80-
$deadline = (Get-Date).AddSeconds(60)
81-
$consecutive_ok = 0
82-
do {
83-
Start-Sleep -Seconds 3
84-
& nvidia-smi.exe 2>&1 | Out-Null
85-
if ($LASTEXITCODE -eq 0) { $consecutive_ok++ } else { $consecutive_ok = 0 }
86-
} while ($consecutive_ok -lt $gpu_count -and (Get-Date) -lt $deadline)
87-
if ($consecutive_ok -lt $gpu_count) {
88-
Write-Error "nvidia-smi did not return cleanly $gpu_count times in a row within 60s of the device cycle"
89-
exit 1
86+
# Poll nvidia-smi for N consecutive successes (N == cycled GPUs)
87+
# so a mid-init "ok" flap doesn't fool the loop; bail after ~60s.
88+
Write-Output "Waiting for nvidia-smi/NVML to come back up after device cycle..."
89+
$deadline = (Get-Date).AddSeconds(60)
90+
$consecutive_ok = 0
91+
do {
92+
Start-Sleep -Seconds 3
93+
& nvidia-smi.exe 2>&1 | Out-Null
94+
if ($LASTEXITCODE -eq 0) { $consecutive_ok++ } else { $consecutive_ok = 0 }
95+
} while ($consecutive_ok -lt $gpu_count -and (Get-Date) -lt $deadline)
96+
if ($consecutive_ok -lt $gpu_count) {
97+
throw "nvidia-smi did not return cleanly $gpu_count times in a row within 60s of the device cycle"
98+
}
99+
} finally {
100+
Write-Output "Restoring NVDisplay.ContainerLocalSystem startup type to $display_container_start_type..."
101+
Set-Service -Name $display_container_service.Name -StartupType $display_container_start_type -ErrorAction Stop
90102
}
91103
}
92104

0 commit comments

Comments
 (0)