rename to timeout

Author: pmcurtin

src/torchrunx/agent.py

@@ -34,7 +34,7 @@ class WorkerArgs:
     local_world_size: int
     world_size: int
     log_file: os.PathLike
-    pg_timeout: int
+    timeout: int
 
     def to_bytes(self) -> bytes:
         return cloudpickle.dumps(self)
@@ -87,7 +87,7 @@ def entrypoint(serialized_worker_args: bytes):
             world_size=worker_args.world_size,
             rank=worker_args.rank,
             store=store,
-            timeout=datetime.timedelta(seconds=worker_args.pg_timeout),
+            timeout=datetime.timedelta(seconds=worker_args.timeout),
         )
 
         os.environ["RANK"] = str(worker_args.rank)
@@ -136,7 +136,7 @@ def main(launcher_agent_group: LauncherAgentGroup):
                     local_world_size=num_workers,
                     world_size=worker_world_size,
                     log_file=worker_log_files[i],
-                    pg_timeout=launcher_payload.pg_timeout,
+                    timeout=launcher_payload.timeout,
                 ).to_bytes(),
             )
             for i in range(num_workers)

src/torchrunx/launcher.py

@@ -96,7 +96,7 @@ class Launcher:
         ]
     )
     env_file: str | os.PathLike | None = None
-    pg_timeout: int = 600
+    timeout: int = 600
 
     def run(
         self,
@@ -210,7 +210,7 @@ def run(
             worker_global_ranks=worker_global_ranks,
             worker_log_files=worker_log_files,
             backend=self.backend,
-            pg_timeout=self.pg_timeout,
+            timeout=self.timeout,
         )
 
         agent_payloads: list[AgentPayload] = launcher_agent_group.sync_payloads(payload=payload)[1:]  # pyright: ignore[reportAssignmentType]
@@ -272,7 +272,7 @@ def launch(
         "NCCL*",
     ],
     env_file: str | os.PathLike | None = None,
-    pg_timeout: int = 600,
+    timeout: int = 600,
 ) -> dict[int, Any]:
     """
     Launch a distributed PyTorch function on the specified nodes.
@@ -295,8 +295,8 @@ def launch(
     :type env_vars: list[str], optional
     :param env_file: An additional environment file that will be sourced prior to executing ``func``, defaults to None
     :type env_file: str | os.PathLike | None, optional
-    :param pg_timeout: Worker process group timeout, defaults to 600
-    :type pg_timeout: int, optional
+    :param timeout: Worker process group timeout, defaults to 600
+    :type timeout: int, optional
     :raises RuntimeError: May fail due to misconfiguration, or errors thrown by ``func``
     :return: A dictionary mapping worker ranks to their output
     :rtype: dict[int, Any]
@@ -309,5 +309,5 @@ def launch(
         log_dir=log_dir,
         env_vars=env_vars,
         env_file=env_file,
-        pg_timeout=pg_timeout,
+        timeout=timeout,
     ).run(func=func, func_kwargs=func_kwargs)

src/torchrunx/utils.py

@@ -29,7 +29,7 @@ class LauncherPayload:
     worker_global_ranks: list[list[int]]
     worker_log_files: list[list[Path]]
     backend: Literal["mpi", "gloo", "nccl", "ucc", None]
-    pg_timeout: int
+    timeout: int
 
 
 @dataclass