some extra error docs

Author: apoorvkh

src/torchrunx/agent.py

@@ -169,8 +169,10 @@ def main(launcher_agent_group: LauncherAgentGroup, logger_hostname: str, logger_
         status = None
         while True:
             if status is None or status.state == "running":
-                status = AgentStatus.from_result(ctx.wait(5))
+                # status can contain ExceptionFromWorker or WorkerFailedError
+                status = AgentStatus.from_result(result=ctx.wait(5))
 
+            # can raise AgentFailedError in launcher and all agents
             agent_statuses = launcher_agent_group.sync_agent_statuses(status=status)
 
             all_done = all(s.state == "done" for s in agent_statuses)

src/torchrunx/utils.py

@@ -23,6 +23,7 @@ def get_open_port() -> int:
 class AgentFailedError(Exception):
     pass
 
+
 class WorkerFailedError(Exception):
     pass
 
@@ -108,7 +109,7 @@ class ExceptionFromWorker:
 @dataclass
 class AgentStatus:
     state: Literal["running", "failed", "done"]
-    return_values: list[Any | ExceptionFromWorker] = field(
+    return_values: list[Any | WorkerFailedError | ExceptionFromWorker] = field(
         default_factory=list
     )  # indexed by local rank