fix: auto-assign free port in spawn_parallel_fn to avoid EADDRINUSE
This commit is contained in:
parent
9aca62c26c
commit
4c35d36146
|
|
@ -1,14 +1,21 @@
|
||||||
import os
|
import os
|
||||||
|
import socket
|
||||||
from abc import ABC, abstractmethod
|
from abc import ABC, abstractmethod
|
||||||
from contextlib import contextmanager
|
from contextlib import contextmanager
|
||||||
from functools import wraps
|
from functools import wraps
|
||||||
from typing import Callable
|
from typing import Callable, Optional
|
||||||
|
|
||||||
import torch
|
import torch
|
||||||
import torch.distributed as dist
|
import torch.distributed as dist
|
||||||
import torch.multiprocessing as mp
|
import torch.multiprocessing as mp
|
||||||
|
|
||||||
|
|
||||||
|
def find_free_port() -> str:
|
||||||
|
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
||||||
|
s.bind(("", 0))
|
||||||
|
return str(s.getsockname()[1])
|
||||||
|
|
||||||
|
|
||||||
def get_current_device():
|
def get_current_device():
|
||||||
return os.environ["LOCAL_DEVICE"]
|
return os.environ["LOCAL_DEVICE"]
|
||||||
|
|
||||||
|
|
@ -217,11 +224,13 @@ def spawn_parallel_fn(
|
||||||
world_size: int,
|
world_size: int,
|
||||||
backend: str = "nccl",
|
backend: str = "nccl",
|
||||||
master_addr: str = "localhost",
|
master_addr: str = "localhost",
|
||||||
master_port: str = "29500",
|
master_port: Optional[str] = None,
|
||||||
device_type: str = "cuda",
|
device_type: str = "cuda",
|
||||||
start_method: str = "spawn",
|
start_method: str = "spawn",
|
||||||
**kwargs,
|
**kwargs,
|
||||||
):
|
):
|
||||||
|
if master_port is None:
|
||||||
|
master_port = find_free_port()
|
||||||
launcher = _detect_launcher()
|
launcher = _detect_launcher()
|
||||||
if launcher in ("torchelastic", "torchrun", "external"):
|
if launcher in ("torchelastic", "torchrun", "external"):
|
||||||
strategy = TorchrunStrategy(
|
strategy = TorchrunStrategy(
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue