Renamed timeout argument to worker_timeout

Author: apoorvkh

src/torchrunx/agent.py

@@ -105,7 +105,7 @@ def main(
                     local_world_size=num_workers,
                     world_size=worker_world_size,
                     hostname=launcher_payload.hostnames[agent_rank],
-                    timeout=launcher_payload.timeout,
+                    timeout=launcher_payload.worker_timeout,
                 ).serialize(),
             )
             for i in range(num_workers)

src/torchrunx/integrations/cli.py

@@ -46,7 +46,7 @@ def add_torchrunx_argument_group(parser: ArgumentParser) -> None:
     )
 
     group.add_argument(
-        "--timeout",
+        "--worker-timeout",
         type=int,
         default=600,
         help="Worker process group timeout in seconds. Default: 600.",
@@ -112,7 +112,7 @@ def launcher_from_args(args: Namespace) -> Launcher:
     else:
         backend = _backend  # pyright: ignore [reportAssignmentType]
 
-    timeout: int = args.timeout
+    worker_timeout: int = args.worker_timeout
     agent_timeout: int = args.agent_timeout
 
     copy_env_vars: tuple[str, ...] = tuple(args.copy_env_vars)
@@ -131,7 +131,7 @@ def launcher_from_args(args: Namespace) -> Launcher:
         workers_per_host=workers_per_host,
         ssh_config_file=ssh_config_file,
         backend=backend,
-        timeout=timeout,
+        worker_timeout=worker_timeout,
         agent_timeout=agent_timeout,
         copy_env_vars=copy_env_vars,
         extra_env_vars=extra_env_vars,

src/torchrunx/launcher.py

@@ -61,7 +61,7 @@ class Launcher:
     """`Backend <https://pytorch.org/docs/stable/distributed.html#torch.distributed.Backend>`_
         for worker process group. By default, NCCL (GPU backend).
         Use GLOO for CPU backend. ``None`` for no process group."""
-    timeout: int = 600
+    worker_timeout: int = 600
     """Worker process group timeout (seconds)."""
     agent_timeout: int = 30
     """Agent communication timeout (seconds)."""
@@ -119,7 +119,8 @@ def run(  # noqa: C901, PLR0912, PLR0915
         )
         ssh_config_file = self.ssh_config_file
         backend = self.backend
-        timeout = self.timeout
+        worker_timeout = self.worker_timeout
+        agent_timeout = self.agent_timeout
 
         env_vars = {
             k: v
@@ -161,7 +162,7 @@ def handler_factory() -> list[logging.Handler]:
             worker_global_ranks=worker_global_ranks,
             worker_world_size=sum(workers_per_host),
             backend=backend,
-            timeout=timeout,
+            worker_timeout=worker_timeout,
         )
         agent_payloads = None
 
@@ -201,7 +202,7 @@ def handler_factory() -> list[logging.Handler]:
                         env_vars=env_vars,
                         env_file=env_file,
                         hostname=hostname,
-                        agent_timeout=self.agent_timeout,
+                        agent_timeout=agent_timeout,
                     ),
                     hostname=hostname,
                     ssh_config_file=ssh_config_file,
@@ -217,7 +218,7 @@ def handler_factory() -> list[logging.Handler]:
                 launcher_port=launcher_port,
                 world_size=world_size,
                 rank=0,
-                agent_timeout=self.agent_timeout,
+                agent_timeout=agent_timeout,
             )
 
             # Sync initial payloads between launcher and agents

src/torchrunx/utils/comm.py

@@ -121,7 +121,7 @@ class LauncherPayload:
     worker_global_ranks: list[list[int]]
     worker_world_size: int
     backend: Literal["nccl", "gloo", "mpi", "ucc"] | None
-    timeout: int
+    worker_timeout: int
 
 
 @dataclass