ciflow/torchtitan/185414: [torchelastic] handle d-state process (#185414)

Summary:

The torchrun elastic agent logged ChildFailedError, then went silent

— never exited — until a manual kill, at which point twinfra spent ~1min

sending stop signals every second that the container ignored.

Root cause inside torch elastic: when workers wedge on a kernel resource

(NCCL/GPU/RDMA — Linux D-state) SIGKILL cannot reap them, and both

MultiprocessContext._close and SubprocessContext._close end in an unbounded

proc.join() / proc.wait(). The agent itself wedges on that join, so the

supervising launcher can never exit and MAST keeps the slot pinned.

Two changes here:

1. 
   

   Bound the final proc.join() (api.py:900) and proc.wait() (api.py:1054)
   
   with the same timeout the rest of _close already honors. On timeout, log
   
   the unkillable PID (with the SIGKILL signal name and advice that the host
   
   may need recycling) and continue. The agent process can now exit even when
   
   workers are unkillable.

   
   


2. 
   

   Make the D-state detection added in the prior change actually escalate:
   
   when _check_d_state_timeout fires, set _remaining_restarts = 0 before
   
   returning UNHEALTHY. A D-state worker still holds GPU/NIC on the host, so
   
   a fresh worker group on the same host would conflict; we want _invoke_run
   
   to take the _stop_workers + return branch and exit immediately.

   
   

Caveat: this diff lets the supervisor exit cleanly. It does not free the

wedged worker processes on the host — those still need host recycling for the

GPUs to be reusable, since no userspace signal can unblock a process in true

D-state. The new error log surfaces the offending PIDs so oncall can recycle

the host explicitly.

Authored by Claude.

Test Plan:

New unit tests:

• 
  

  BoundedCloseTest in test/distributed/elastic/multiprocessing/api_test.py
  
  covers both MultiprocessContext._close and SubprocessContext._close.
  
  Each test mocks an unkillable process whose is_alive/poll always reports
  
  alive and whose join/wait honor the timeout but never resolve. Asserts
  
  _close returns in <5s (two bounded joins of 1s each + epsilon) instead of
  
  hanging forever.

  
  


• 
  

  Extended LocalElasticAgentDStateTest::test_check_d_state_timeout_fires in
  
  test/distributed/elastic/agent/server/test/local_elastic_agent_test.py to
  
  assert _remaining_restarts is set to 0 when the D-state timeout fires.

  
  

Reviewed By: d4l3k

Differential Revision: D105652547