Skip to content

Change launch to Launcher dataclass #40

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 8 commits into from
Jul 15, 2024
Merged

Change launch to Launcher dataclass #40

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
0e4db09
Change `launch` to `Launcher` dataclass
apoorvkh Jul 13, 2024
1d7b3be
moved slurm utils to slurm.py
apoorvkh Jul 13, 2024
165d4a0
`launch` alias function for `Launcher.run`
apoorvkh Jul 13, 2024
fb62a94
Merge branch 'main' of github.com:apoorvkh/torchrunx into argument-class
apoorvkh Jul 14, 2024
44dc2a5
Update api.rst
apoorvkh Jul 14, 2024
4b4cf13
Update docs-preview.yml
apoorvkh Jul 14, 2024
765e802
fix typing, formatting
pmcurtin Jul 15, 2024
cc9c1d0
mutable default dataclass values
pmcurtin Jul 15, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 19 additions & 21 deletions src/torchrunx/launcher.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -78,24 +78,22 @@ def monitor_log(log_file: Path):

@dataclass
class Launcher:
hostnames: list[str] = (["localhost"],)
workers_per_host: int | list[int] = (1,)
ssh_config_file: str | os.PathLike | None = (None,)
backend: Literal["mpi", "gloo", "nccl", "ucc", None] = (None,)
log_dir: os.PathLike | str = ("./logs",)
env_vars: list[str] = (
[
"PATH",
"LD_LIBRARY",
"LIBRARY_PATH",
"PYTHON*",
"CUDA*",
"TORCH*",
"PYTORCH*",
"NCCL*",
],
)
env_file: str | os.PathLike | None = (None,)
hostnames: list[str] = ["localhost"]
workers_per_host: int | list[int] = 1
ssh_config_file: str | os.PathLike | None = None
backend: Literal["mpi", "gloo", "nccl", "ucc", None] = None
log_dir: os.PathLike | str = "./logs"
env_vars: list[str] = [
"PATH",
"LD_LIBRARY",
"LIBRARY_PATH",
"PYTHON*",
"CUDA*",
"TORCH*",
"PYTORCH*",
"NCCL*",
]
env_file: str | os.PathLike | None = None

def run(
self,
Expand Down Expand Up @@ -137,7 +135,7 @@ def run(
workers_per_host = [workers_per_host] * num_hosts

assert workers_per_host is not None
assert len(workers_per_host) == num_hosts
assert len(workers_per_host) == num_hosts # type: ignore

# launch command

Expand Down Expand Up @@ -199,7 +197,7 @@ def run(

# build and sync payloads between launcher and agents

_cumulative_workers = [0] + list(itertools.accumulate(workers_per_host))
_cumulative_workers = [0] + list(itertools.accumulate(workers_per_host)) # type: ignore

worker_world_size = _cumulative_workers[-1]

Expand All @@ -211,7 +209,7 @@ def run(
worker_log_files = [
[
log_dir / f"{timestamp}_{hostname}_{local_rank}.log"
for local_rank in range(workers_per_host[i])
for local_rank in range(workers_per_host[i]) # type: ignore
]
for i, hostname in enumerate(self.hostnames)
]
Expand Down
2 changes: 1 addition & 1 deletion src/torchrunx/slurm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
import subprocess


def slurm_hosts() -> list[str]:
def slurm_hosts() -> "list[str]":
Copy link
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Wait, why this change?

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pyright was complaining about it.

"""Retrieves hostnames of Slurm-allocated nodes.

:return: Hostnames of nodes in current Slurm allocation
Expand Down
Loading