additional ruff rules

pyproject.toml

@@ -41,7 +41,7 @@ include = ["pyproject.toml", "src/**/*.py", "tests/**/*.py"]
 line-length = 100
 src = ["src", "tests"]
 [tool.ruff.lint]
-extend-select = ["I"]
+select = ["E", "F", "B", "UP", "I"]
 
 [tool.pyright]
 include = ["src", "tests"]

src/torchrunx/launcher.py

@@ -9,11 +9,11 @@
 import subprocess
 import sys
 from collections import ChainMap
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 from functools import partial
 from logging import Handler
 from multiprocessing import Process
-from typing import Any, Callable, Literal
+from typing import Any, Callable, Literal, Sequence
 
 import fabric
 import torch.distributed as dist
@@ -58,31 +58,29 @@ def execute_command(
 
 @dataclass
 class Launcher:
-    hostnames: list[str] | Literal["auto", "slurm"] = field(default_factory=lambda: ["localhost"])
-    workers_per_host: int | list[int] | Literal["auto", "slurm"] = 1
+    hostnames: list[str] | Literal["auto", "slurm"] = "auto"
+    workers_per_host: int | list[int] | Literal["auto", "slurm"] = "auto"
     ssh_config_file: str | os.PathLike | None = None
     backend: Literal["mpi", "gloo", "nccl", "ucc", None] = None
     log_handlers: list[Handler] | Literal["auto"] | None = "auto"
-    env_vars: list[str] = field(
-        default_factory=lambda: [
-            "PATH",
-            "LD_LIBRARY",
-            "LIBRARY_PATH",
-            "PYTHON*",
-            "CUDA*",
-            "TORCH*",
-            "PYTORCH*",
-            "NCCL*",
-        ]
+    env_vars: Sequence[str] = (
+        "PATH",
+        "LD_LIBRARY",
+        "LIBRARY_PATH",
+        "PYTHON*",
+        "CUDA*",
+        "TORCH*",
+        "PYTORCH*",
+        "NCCL*",
     )
     env_file: str | os.PathLike | None = None
     timeout: int = 600
 
     def run(
         self,
         func: Callable,
-        func_args: tuple[Any] = tuple(),
-        func_kwargs: dict[str, Any] = {},
+        func_args: tuple[Any] | None = None,
+        func_kwargs: dict[str, Any] | None = None,
     ) -> dict[int, Any]:
         """
         Launch a distributed PyTorch function on the specified nodes. See :mod:`torchrunx.launch`
@@ -204,6 +202,11 @@ def run(
             host_ranks = range(_cumulative_workers[n], _cumulative_workers[n + 1])
             worker_global_ranks.append(list(host_ranks))
 
+        if func_args is None:
+            func_args = tuple()
+        if func_kwargs is None:
+            func_kwargs = dict()
+
         payload = LauncherPayload(
             fn=partial(func, *func_args, **func_kwargs),
             hostnames=self.hostnames,
@@ -251,14 +254,14 @@ def run(
 
 def launch(
     func: Callable,
-    func_args: tuple[Any] = tuple(),
-    func_kwargs: dict[str, Any] = {},
-    hostnames: list[str] | Literal["auto", "slurm"] = ["localhost"],
-    workers_per_host: int | list[int] | Literal["auto", "slurm"] = 1,
+    func_args: tuple[Any] | None = None,
+    func_kwargs: dict[str, Any] | None = None,
+    hostnames: list[str] | Literal["auto", "slurm"] = "auto",
+    workers_per_host: int | list[int] | Literal["auto", "slurm"] = "auto",
     ssh_config_file: str | os.PathLike | None = None,
     backend: Literal["mpi", "gloo", "nccl", "ucc", None] = None,
     log_handlers: list[Handler] | Literal["auto"] = "auto",
-    env_vars: list[str] = [
+    env_vars: Sequence[str] = (
         "PATH",
         "LD_LIBRARY",
         "LIBRARY_PATH",
@@ -267,7 +270,7 @@ def launch(
         "TORCH*",
         "PYTORCH*",
         "NCCL*",
-    ],
+    ),
     env_file: str | os.PathLike | None = None,
     timeout: int = 600,
 ) -> dict[int, Any]:

src/torchrunx/utils.py

@@ -47,9 +47,7 @@ class AgentStatus:
     return_values: dict[int, Any | WorkerException] = field(default_factory=dict)
 
     @classmethod
-    def from_result(
-        cls, result: RunProcsResult | None, worker_global_ranks: list[int]
-    ) -> Self:
+    def from_result(cls, result: RunProcsResult | None, worker_global_ranks: list[int]) -> Self:
         if result is None:
             return cls(state="running")
 
@@ -96,10 +94,8 @@ def _deserialize(self, serialized: bytes) -> Any:
     def _all_gather(self, object: Any) -> list:
         """gather object from every rank to list on every rank"""
         object_bytes = self._serialize(object)
-        object_list = [bytes()] * self.world_size
-        dist.all_gather_object(
-            object_list=object_list, obj=object_bytes, group=self.group
-        )
+        object_list = [b""] * self.world_size
+        dist.all_gather_object(object_list=object_list, obj=object_bytes, group=self.group)
         object_list = [self._deserialize(o) for o in object_list]
         return object_list
 

tests/test_CI.py

@@ -61,7 +61,7 @@ def dist_func():
     assert len(log_files) == 3
 
     for file in log_files:
-        with open(f"{tmp}/{file}", "r") as f:
+        with open(f"{tmp}/{file}") as f:
             contents = f.read()
             print(contents)
             if file.endswith("[0].log"):