fix(ci): stabilize GPU E2E tests on shared RDP agents

.pipelines/Modelkit E2E Test.yml

@@ -178,6 +178,11 @@ parameters:
     displayName: 'OV: self-hosted agent name (Agent.Name)'
     type: string
     default: NPU-OV2
+  - name: dmlPinnedAgents
+    displayName: 'Agents whose DML perf eval pins the sole physical GPU (empty disables)'
+    type: object
+    default:
+      - NPU-OV2
   - name: ovPairs
     displayName: 'OV: EP/device pairs (any of: dml_gpu, mlas_cpu, ov_cpu, ov_gpu, ov_npu)'
     type: object
@@ -223,6 +228,7 @@ stages:
         parameters:
           agentName: ${{ parameters.ovAgentName }}
           agentSuffix: ov
+          dmlPinnedAgents: ${{ parameters.dmlPinnedAgents }}
           pairNames: ${{ parameters.ovPairs }}
           models: ${{ parameters.models }}
           modelTimeout: ${{ parameters.modelTimeout }}

.pipelines/templates/e2e-test-jobs.yml

@@ -21,6 +21,10 @@ parameters:
   #   mlas_cpu, ov_cpu, ov_gpu, ov_npu, trtrtx_gpu
   - name: pairNames
     type: object
+  # Opt-in for shared RDP development agents; affects DML perf eval and perf pytest.
+  - name: dmlPinnedAgents
+    type: object
+    default: []
   # Inline list of models to test. Each entry: { hf_id: <str>, task: <str> }.
   # `task` is optional — omit (or leave empty) to let winml infer.
   - name: models
@@ -228,6 +232,10 @@ jobs:
                       if ($task) {
                           $uvArgs += @("--task", $task)
                       }
+                      $pinDml = '${{ containsValue(parameters.dmlPinnedAgents, parameters.agentName) }}'
+                      if ($name -eq 'dml_gpu' -and $pinDml -eq 'True') {
+                          $uvArgs += "--pin-single-dml-gpu"
+                      }
                       if ($isLastPair) {
                           $uvArgs += "--clean-cache"
                       }
@@ -255,6 +263,12 @@ jobs:
       # ----------------------------------------------------------------------
       - ${{ each target in parameters.pytestTargets }}:
           - powershell: |
+              $env:WINML_E2E_PIN_SINGLE_DML_GPU = '0'
+              $pinDml = '${{ containsValue(parameters.dmlPinnedAgents, parameters.agentName) }}'
+              if ('${{ target }}' -eq 'perf' -and $pinDml -eq 'True') {
+                  $env:WINML_E2E_PIN_SINGLE_DML_GPU = '1'
+                  Write-Host "DML perf tests will exercise the physical GPU on this shared RDP agent."
+              }
               uv run --no-sync python -m pytest tests/e2e/test_${{ target }}_e2e.py -m e2e --timeout=${{ parameters.pytestTimeout }} --junitxml="$(Agent.TempDirectory)/junit-${{ parameters.agentSuffix }}-${{ target }}.xml" -v
               exit $LASTEXITCODE
             workingDirectory: $(Build.SourcesDirectory)

scripts/agent_setup/README.md

@@ -74,17 +74,25 @@ After registration, sign out and sign back in to verify the agent console launch
 
 **Only needed on GPU agents that run DirectML (`GPU-DML`) and are accessed over Remote Desktop.** Skip this on CPU/NPU/OpenVINO-only agents.
 
-When you disconnect an RDP session (instead of signing out), Windows detaches the physical display adapter from that session and falls back to the "Microsoft Remote Display Adapter," which has no Direct3D 12 support. DirectML can no longer create a device, so every DirectML CI step fails while the session is disconnected — even though OpenVINO and CPU runs keep working. The symptom in the pipeline is that only the `dml_gpu` eval steps and `test_perf_e2e.py` fail on that agent.
+RDP session transitions on a shared GPU agent can leave stale Remote Display Adapter nodes and duplicate GPU routes. DirectML may then select an unusable route even while the physical GPU remains available. Symptoms include DML device-creation errors and mismatches between ORT and DXCore adapter inventories.
 
-The provided script registers a SYSTEM, event-triggered Scheduled Task (`KeepSessionOnConsole`) that fires on RDP disconnect (TerminalServices-LocalSessionManager **Event ID 24**). After a short settle delay — so a quick reconnect is left alone — it redirects a still-disconnected session back to the physical console via `tscon`, keeping the GPU bound so DirectML continues to work headless.
+The script registers a SYSTEM Scheduled Task (`KeepSessionOnConsole`) triggered by TerminalServices-LocalSessionManager **Event ID 24**. Its enhanced worker matches the version deployed on the shared development agent:
+
+- A remote IP source triggers immediate handling; local/loopback transitions are ignored to avoid disrupting reconnects. Older tasks that only supply a SessionId use a matching Event 24 from the last minute to resolve the source address.
+- Session 1 is eligible; only service Session 0 is excluded. A reconnected session or another user's occupied console is left alone.
+- The worker attempts `tscon` up to three times, waiting two seconds between state checks, and verifies that the target reached console.
+- Only after verification, it removes non-present RDP display nodes (`SWD\REMOTEDISPLAYENUM\*`, phantom/error 45). Physical PCI GPUs and present RDP devices are excluded.
+- The optional `-CleanupOnly` mode runs that same cleanup only when a console session exists.
+
+Console recovery and PnP cleanup **do not prove DML readiness** or remove every stale ORT/DXGI route. The worker logs warnings and retains its deployed exit-zero behavior even on errors. Use the separate physical-GPU pin in the E2E pipeline to avoid known duplicate routes, and verify with real inference.
 
 **Register** (run once, elevated — the script reports an error if not elevated):
 
 ```powershell
 powershell -ExecutionPolicy Bypass -File .\scripts\agent_setup\setup_rdp_gpu_keepalive.ps1
 ```
 
-The script writes the worker to `C:\agent\tools\keep_console.ps1` and logs activity to `C:\agent\tools\keep_console.log`. To verify, disconnect and reconnect your RDP session, then check the log — you should reconnect on the first try, with no `tscon` errors.
+The installer writes the worker to `C:/agent/tools/keep_console.ps1` and logs activity to `C:/agent/tools/keep_console.log`. Re-registering passes both the event SessionId and source address. Its smoke test omits SessionId, so it performs no redirect or device cleanup. To verify recovery during a planned RDP disconnect, check for a verified console transition and review cleanup warnings, then run a DML inference check. Merely editing this branch does not update the installed worker or task.
 
 **Unregister:**
 

scripts/agent_setup/setup_rdp_gpu_keepalive.ps1

@@ -1,17 +1,17 @@
 # setup_rdp_gpu_keepalive.ps1
 #
-# Problem this solves:
-#   On this self-hosted ADO agent the GPU (Intel Arc) is only usable by DirectML
-#   while the user's RDP session is connected. When RDP disconnects, the session
-#   loses the physical display adapter (falls back to "Microsoft Remote Display
-#   Adapter", no D3D12) and every DirectML CI step fails. OpenVINO/CPU are
-#   unaffected, so only the `dml_gpu` eval steps and `test_perf_e2e.py` break.
+# Purpose:
+#   Recover the interactive agent session after an RDP disconnect. Session
+#   transitions can leave stale Remote Display Adapter nodes and duplicate GPU
+#   routes. Returning to console and cleaning non-present RDP nodes helps restore
+#   the environment, but does not guarantee that every ORT/DXGI route is usable.
 #
-# Fix:
-#   Register a SYSTEM, event-triggered scheduled task ("KeepSessionOnConsole")
-#   that fires on RDP disconnect (TerminalServices-LocalSessionManager Event ID
-#   24) and redirects the still-disconnected session back to the physical console
-#   via `tscon`, so the GPU stays bound and DirectML keeps working headless.
+# Enhanced worker:
+#   Handles network-backed Event 24 immediately, ignores local transitions,
+#   protects another user's active console, verifies/retries tscon, and removes
+#   only non-present RDP display nodes after a verified console transition.
+#   Keep this worker aligned with the deployed C:/agent/tools/keep_console.ps1.
+#   The separate CI physical-GPU pin is still needed for stale ORT routes.
 #
 # Usage (run once, elevated):
 #   powershell -NoProfile -ExecutionPolicy Bypass -File scripts\agent_setup\setup_rdp_gpu_keepalive.ps1
@@ -41,29 +41,136 @@ try {
 
     # ---- keep_console.ps1 (the redirect worker) ----
     $worker = @'
-param([string]$SessionId = '')
+    param(
+      [string]$SessionId = '',
+      [string]$SourceNetworkAddress = '',
+      [switch]$CleanupOnly
+    )
 
 $ErrorActionPreference = 'Stop'
 $log = 'C:\agent\tools\keep_console.log'
 function Write-Log($m) { "{0}  {1}" -f (Get-Date -Format 'yyyy-MM-dd HH:mm:ss'), $m | Out-File -FilePath $log -Append -Encoding utf8 }
 
-# Parse `qwinsta` into a map of sessionId -> state. Anchors on the state keyword
-# and the integer immediately before it, so an empty SESSIONNAME column is fine.
+# Parse `qwinsta` into a map of sessionId -> session facts. Anchors on the state
+# keyword and the integer immediately before it, so an empty USERNAME column is
+# fine. SESSIONNAME and USERNAME are only used for fail-closed safety checks.
 function Get-Sessions {
     $map = @{}
     foreach ($r in (qwinsta 2>$null) -split "`r?`n") {
         $t = ($r -split '\s+') | Where-Object { $_ -ne '' }
         for ($i = 0; $i -lt $t.Count; $i++) {
             if ($t[$i] -in @('Active','Disc','Conn','Listen','Down','Idle')) {
-                if ($i -ge 1 -and $t[$i-1] -match '^\d+$') { $map[[int]$t[$i-1]] = $t[$i] }
+        if ($i -ge 1 -and $t[$i-1] -match '^\d+$') {
+          $id = [int]$t[$i-1]
+          $name = if ($t.Count) { $t[0].TrimStart('>') } else { '' }
+          $username = if ($i -ge 3) { $t[$i-2] } else { '' }
+          $map[$id] = [pscustomobject]@{
+            Id = $id
+            Name = $name
+            Username = $username
+            State = $t[$i]
+          }
+        }
                 break
             }
         }
     }
     return $map
 }
 
+  # Backward-compatible fallback for an already-registered task whose action only
+  # passes SessionId. Re-registering the task makes Address available directly,
+  # but looking up the newest matching Event 24 lets a worker-only deployment take
+  # effect immediately without elevation. Limit the lookup window so a stale
+  # network disconnect can never be mistaken for the event being handled.
+  function Resolve-SourceNetworkAddress([int]$Target, [string]$Provided) {
+    $providedAddress = $Provided.Trim()
+    if ($providedAddress) { return $providedAddress }
+
+    try {
+      $events = Get-WinEvent -FilterHashtable @{
+        LogName = 'Microsoft-Windows-TerminalServices-LocalSessionManager/Operational'
+        Id = 24
+        StartTime = (Get-Date).AddMinutes(-1)
+      } -ErrorAction Stop
+      foreach ($event in $events) {
+        [xml]$eventXml = $event.ToXml()
+        $payload = $eventXml.Event.UserData.EventXML
+        $eventSessionId = 0
+        if (
+          [int]::TryParse([string]$payload.SessionID, [ref]$eventSessionId) -and
+          $eventSessionId -eq $Target
+        ) {
+          return [string]$payload.Address
+        }
+      }
+    }
+    catch {
+      Write-Log "WARN: cannot resolve Event 24 source address: $($_.Exception.Message)"
+    }
+    return ''
+  }
+
+# Remove only stale/non-present RDP indirect-display device nodes. On affected
+# machines a node can be CM_PROB_PHANTOM after `tscon`, yet its cloned GPU handle
+# remains visible through DXGI/ORT and may sort before the real adapter. Removing
+# that non-present node is safe for the physical GPU (which has a PCI instance ID)
+# and a future RDP connection recreates its own Remote Display Adapter as needed.
+function Remove-StaleRemoteDisplayAdapters {
+  try {
+    $stale = @(Get-PnpDevice -Class Display -ErrorAction Stop | Where-Object {
+      $_.InstanceId -like 'SWD\REMOTEDISPLAYENUM\*' -and
+      ($_.Problem -eq 'CM_PROB_PHANTOM' -or $_.ConfigManagerErrorCode -eq 45)
+    })
+  }
+  catch {
+    Write-Log "WARN: cannot enumerate display devices: $($_.Exception.Message)"
+    return $false
+  }
+
+  if (-not $stale) {
+    Write-Log "No stale Remote Display Adapter device nodes found."
+    return $true
+  }
+
+  $allRemoved = $true
+  $pnputil = Join-Path $env:SystemRoot 'System32\pnputil.exe'
+  foreach ($device in $stale) {
+    Write-Log "Removing stale Remote Display Adapter: $($device.InstanceId)"
+    # Windows PowerShell 5.1 turns redirected native stderr into errors.
+    # Capture it without bypassing exit-code and session-state handling.
+    $savedErrorActionPreference = $ErrorActionPreference
+    try {
+        $ErrorActionPreference = 'Continue'
+        $out = & $pnputil /remove-device $device.InstanceId 2>&1
+        $code = $LASTEXITCODE
+    }
+    finally {
+        $ErrorActionPreference = $savedErrorActionPreference
+    }
+    if ($out) { $out | ForEach-Object { Write-Log "  pnputil: $_" } }
+    if ($code -ne 0) {
+      $allRemoved = $false
+      Write-Log "WARN: pnputil remove-device exited $code for $($device.InstanceId)."
+    }
+  }
+  return $allRemoved
+}
+
 try {
+  if ($CleanupOnly) {
+    $sessions = Get-Sessions
+    $console = @($sessions.Values | Where-Object {
+      $_.Name -eq 'console' -and $_.State -in @('Active', 'Conn')
+    })
+    if (-not $console) {
+      Write-Log "WARN: cleanup requested without a console session; refusing."
+      exit 0
+    }
+    [void](Remove-StaleRemoteDisplayAdapters)
+    exit 0
+  }
+
     # Act ONLY on the session reported by the disconnect event. The broad
     # "find any Disc session" scan is intentionally gone: during a reconnect the
     # real session briefly flips Disc->Active and the scan would grab an unrelated
@@ -73,36 +180,83 @@ try {
         exit 0
     }
     $target = [int]$SessionId
-    if ($target -lt 2) {
-        Write-Log "Session $target is a system/console session; skipping."
+    if ($target -eq 0) {
+      Write-Log "Session 0 is the services session; skipping."
         exit 0
     }
 
-    # Settle delay: a plain reconnect produces a transient disconnect. Waiting a
-    # few seconds and re-checking lets a reconnecting session return to Active so
-    # we DON'T yank it back to console mid-reconnect (which caused the
-    # "connect twice" tug-of-war). Only a session that is STILL disconnected
-    # after the delay is treated as "user really left".
-    Start-Sleep -Seconds 10
+    # A network-backed Event 24 is the real RDP transport disconnect. LOCAL
+    # events are generated by tscon and reconnect transitions; ignoring them
+    # prevents the old "connect twice" tug-of-war without delaying the genuine
+    # redirect long enough for Windows to settle on the phantom adapter.
+    $sourceAddress = Resolve-SourceNetworkAddress $target $SourceNetworkAddress
+    $parsedAddress = $null
+    $isRemoteAddress = (
+      [System.Net.IPAddress]::TryParse($sourceAddress, [ref]$parsedAddress) -and
+      -not [System.Net.IPAddress]::IsLoopback($parsedAddress) -and
+      -not $parsedAddress.Equals([System.Net.IPAddress]::Any) -and
+      -not $parsedAddress.Equals([System.Net.IPAddress]::IPv6Any)
+    )
+    if (-not $isRemoteAddress) {
+      Write-Log "Ignoring Event 24 for session $target with source '$sourceAddress'."
+      exit 0
+    }
 
     $sessions = Get-Sessions
-    $state = $sessions[$target]
-    if ($state -ne 'Disc') {
-        Write-Log "Session $target is '$state' after settle (reconnected or gone); no redirect needed."
+    $session = $sessions[$target]
+    if ($null -eq $session -or $session.State -ne 'Disc') {
+      $state = if ($null -eq $session) { '' } else { $session.State }
+      Write-Log "Session $target is '$state' at dispatch (reconnected or gone); no redirect needed."
         exit 0
     }
 
-    Write-Log "Session $target still Disc after settle; redirecting to console."
-    $out = & tscon $target /dest:console 2>&1
-    $code = $LASTEXITCODE
-    if ($out) { $out | ForEach-Object { Write-Log "  tscon: $_" } }
-    if ($code -eq 0) {
-        Write-Log "Session $target redirected to console (GPU retained)."
+    # Never displace another user's active console session. A connected but
+    # unowned console placeholder is safe; tscon will attach the target to it.
+    $occupiedConsole = @($sessions.Values | Where-Object {
+      $_.Id -ne $target -and $_.Name -eq 'console' -and
+      $_.State -eq 'Active' -and $_.Username
+    }) | Select-Object -First 1
+    if ($occupiedConsole) {
+      Write-Log "WARN: console is occupied by '$($occupiedConsole.Username)' (session $($occupiedConsole.Id)); refusing to redirect session $target."
+      exit 0
     }
-    else {
-        # Non-fatal: the session state can change between the check and tscon.
-        Write-Log "WARN: tscon $target /dest:console exited $code (session likely changed state); ignoring."
+
+    Write-Log "Remote disconnect from $sourceAddress left session $target Disc; redirecting immediately."
+    $redirected = $false
+    for ($attempt = 1; $attempt -le 3; $attempt++) {
+      # Windows PowerShell 5.1 turns redirected native stderr into errors.
+      # Capture it without bypassing exit-code and session-state handling.
+      $savedErrorActionPreference = $ErrorActionPreference
+      try {
+          $ErrorActionPreference = 'Continue'
+          $out = & tscon $target /dest:console 2>&1
+          $code = $LASTEXITCODE
+      }
+      finally {
+          $ErrorActionPreference = $savedErrorActionPreference
+      }
+      if ($out) { $out | ForEach-Object { Write-Log "  tscon: $_" } }
+
+      Start-Sleep -Seconds 2
+      $current = (Get-Sessions)[$target]
+      if ($null -ne $current -and $current.Name -eq 'console' -and $current.State -in @('Active', 'Conn')) {
+        $redirected = $true
+        break
+      }
+      if ($null -eq $current -or $current.State -ne 'Disc') {
+        Write-Log "WARN: session $target changed state after tscon; not retrying."
+        break
+      }
+      Write-Log "WARN: tscon attempt ${attempt} exited $code without a verified console transition."
     }
+
+    if (-not $redirected) {
+      Write-Log "WARN: session $target was not verified on console; skipping display cleanup."
+      exit 0
+    }
+
+    Write-Log "Session $target verified on console; cleaning stale RDP display nodes."
+    [void](Remove-StaleRemoteDisplayAdapters)
     exit 0
 }
 catch {
@@ -126,6 +280,7 @@ catch {
       <Subscription>&lt;QueryList&gt;&lt;Query Id="0" Path="Microsoft-Windows-TerminalServices-LocalSessionManager/Operational"&gt;&lt;Select Path="Microsoft-Windows-TerminalServices-LocalSessionManager/Operational"&gt;*[System[(EventID=24)]]&lt;/Select&gt;&lt;/Query&gt;&lt;/QueryList&gt;</Subscription>
       <ValueQueries>
         <Value name="SessionID">Event/UserData/EventXML/SessionID</Value>
+        <Value name="SourceNetworkAddress">Event/UserData/EventXML/Address</Value>
       </ValueQueries>
     </EventTrigger>
   </Triggers>
@@ -156,12 +311,12 @@ catch {
   <Actions Context="Author">
     <Exec>
       <Command>powershell.exe</Command>
-      <Arguments>-NoProfile -NonInteractive -ExecutionPolicy Bypass -File "__WORKER__" -SessionId $(SessionID)</Arguments>
+      <Arguments>-NoProfile -NonInteractive -ExecutionPolicy Bypass -File "__WORKER__" -SessionId "$(SessionID)" -SourceNetworkAddress "$(SourceNetworkAddress)"</Arguments>
     </Exec>
   </Actions>
 </Task>
 '@
-    # Single-quoted here-string keeps $(SessionID) literal for Task Scheduler;
+    # Single-quoted here-string keeps event value substitutions literal for Task Scheduler;
     # inject the worker path via a placeholder so it stays parameterized.
     $xml = $xml.Replace('__WORKER__', $WorkerPath)
 
@@ -171,7 +326,7 @@ catch {
     $t = Get-ScheduledTask -TaskName $TaskName
     L ("task state={0} principal={1} runlevel={2}" -f $t.State, $t.Principal.UserId, $t.Principal.RunLevel)
 
-    # ---- safe smoke test: run worker once now (user is connected => no Disc => no redirect) ----
+    # ---- safe smoke test: no SessionId means no redirect or display cleanup ----
     $smoke = & powershell.exe -NoProfile -ExecutionPolicy Bypass -File $WorkerPath 2>&1
     L ("smoke worker exit={0}" -f $LASTEXITCODE)
     L ("smoke worker output: {0}" -f (($smoke | Out-String).Trim()))