basic WorkerKilledError

src/torchrunx/launcher.py

@@ -22,7 +22,14 @@
 
 from .environment import auto_hosts, auto_workers, slurm_hosts, slurm_workers
 from .logging_utils import LogRecordSocketReceiver, default_handlers
-from .utils import AgentStatus, LauncherAgentGroup, LauncherPayload, WorkerException, get_open_port
+from .utils import (
+    AgentStatus,
+    LauncherAgentGroup,
+    LauncherPayload,
+    WorkerException,
+    WorkerKilledError,
+    get_open_port,
+)
 
 
 class AgentKilledError(Exception):
@@ -152,6 +159,8 @@ def run(  # noqa: C901, PLR0912
                     for value in s.return_values:
                         if isinstance(value, WorkerException):
                             raise value.exception
+                        if isinstance(value, WorkerKilledError):
+                            raise value
 
                 if all(s.state == "done" for s in agent_statuses):
                     break

src/torchrunx/utils.py

@@ -94,6 +94,11 @@ class WorkerException:
     exception: Exception
 
 
+@dataclass
+class WorkerKilledError(Exception):
+    failure: str
+
+
 @dataclass
 class AgentStatus:
     state: Literal["running", "failed", "done"]
@@ -105,9 +110,9 @@ class AgentStatus:
     def from_result(cls, result: RunProcsResult | None) -> Self:
         if result is None:
             return cls(state="running")
-
+        for local_rank, failure in result.failures.items():
+            result.return_values[local_rank] = WorkerKilledError(failure.message)
         return_values = list(result.return_values.values())
-
         failed = any(isinstance(v, WorkerException) for v in return_values)
         state = "failed" if failed else "done"
 

tests/test_func.py

@@ -1,4 +1,5 @@
 import os
+import time
 
 import torch
 import torch.distributed as dist