From 58906a748273f77e10d779790ce6830e26b7c261 Mon Sep 17 00:00:00 2001 From: Gothos Date: Wed, 11 Mar 2026 21:36:39 +0530 Subject: [PATCH 1/2] chore: make DistributedRunner use random ports by default --- projects/fal/src/fal/distributed/utils.py | 18 +++ projects/fal/src/fal/distributed/worker.py | 121 ++++++++++++++------- 2 files changed, 97 insertions(+), 42 deletions(-) diff --git a/projects/fal/src/fal/distributed/utils.py b/projects/fal/src/fal/distributed/utils.py index 8d56eb02f..9995fb202 100644 --- a/projects/fal/src/fal/distributed/utils.py +++ b/projects/fal/src/fal/distributed/utils.py @@ -3,6 +3,7 @@ import json import os import pickle +import socket import threading import warnings from collections.abc import Callable @@ -314,6 +315,23 @@ def wrap_distributed_worker( dist.destroy_process_group() +def _pick_random_free_port(addr: str) -> int: + """ + Picks a random free port on the given address. + :param addr: The address to pick a free port on. + :return: The free port. + """ + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind((addr, 0)) + return s.getsockname()[1] + + +def _check_port_in_use(errs: list[Exception], exc: Exception) -> bool: + txt = f"{exc}\n" + "\n".join(str(e) for e in errs) + txt = txt.lower() + return "address already in use" in txt or "eaddrinuse" in txt + + def launch_distributed_processes( func: Callable, world_size: int = 1, diff --git a/projects/fal/src/fal/distributed/worker.py b/projects/fal/src/fal/distributed/worker.py index ca3e05ec6..d18d2ccf0 100644 --- a/projects/fal/src/fal/distributed/worker.py +++ b/projects/fal/src/fal/distributed/worker.py @@ -18,6 +18,8 @@ from fal.distributed.utils import ( KeepAliveTimer, + _check_port_in_use, + _pick_random_free_port, distributed_deserialize, distributed_serialize, encode_text_event, @@ -231,9 +233,9 @@ def __init__( worker_cls: type[DistributedWorker] = DistributedWorker, world_size: int = 1, master_addr: str = "127.0.0.1", - master_port: int = 29500, + master_port: int = 0, worker_addr: str = "127.0.0.1", - worker_port: int = 54923, + worker_port: int = 0, timeout: int = 86400, # 24 hours keepalive_payload: dict[str, Any] = {}, keepalive_interval: Optional[Union[int, float]] = None, @@ -537,52 +539,87 @@ async def start(self, timeout: int = 1800, **kwargs: Any) -> None: if self.is_alive(): raise RuntimeError("Distributed processes are already running.") - self.context = launch_distributed_processes( - self.run, - world_size=self.world_size, - master_addr=self.master_addr, - master_port=self.master_port, - timeout=self.timeout, - cwd=self.cwd, - **kwargs, - ) + max_attempts = int(kwargs.pop("max_start_attempts", 3)) + if max_attempts < 1: + raise ValueError("max_start_attempts must be at least 1.") - try: - ready_workers: set[int] = set() - socket = self.get_zmq_socket() - start_time = time.perf_counter() + pick_master_port = self.master_port == 0 + pick_worker_port = self.worker_port == 0 - while len(ready_workers) < self.world_size: - try: - ident, msg = await socket.recv_multipart(flags=zmq.NOBLOCK) # type: ignore[misc] + if pick_worker_port: + # Pick a concrete worker port before spawn without creating ZMQ socket; + # spawned children pickle `self`, and pyzmq sockets are not picklable. + self.worker_port = _pick_random_free_port(self.worker_addr) - if msg != b"READY": - worker_id = ident.decode("utf-8") - worker_msg = msg.decode("utf-8") - raise RuntimeError( - f"Unexpected message from worker {worker_id}: {worker_msg}" - ) + for attempt in range(max_attempts): + try: + if pick_master_port: + self.master_port = _pick_random_free_port(self.master_addr) + if pick_worker_port and attempt > 0: + self.worker_port = _pick_random_free_port(self.worker_addr) + + self.context = launch_distributed_processes( + self.run, + world_size=self.world_size, + master_addr=self.master_addr, + master_port=self.master_port, + timeout=self.timeout, + cwd=self.cwd, + **kwargs, + ) - print(f"[debug] Worker {ident.decode('utf-8')} is ready.") - ready_workers.add(ident) - except zmq.Again: - total_wait_time = time.perf_counter() - start_time - if total_wait_time > timeout: - raise TimeoutError( - f"Timeout reached after {timeout} seconds while " - f"waiting for workers to be ready." - ) - await asyncio.sleep(0.5) - self.ensure_alive() - except Exception as e: - print(f"[debug] Error during startup: {e}\n{traceback.format_exc()}") - self.terminate(timeout=timeout) - raise RuntimeError("Failed to start distributed processes.") from e + ready_workers: set[bytes] = set() + socket = self.get_zmq_socket() + start_time = time.perf_counter() - print("[debug] All workers are ready and running.") + while len(ready_workers) < self.world_size: + try: + ident, msg = await socket.recv_multipart(flags=zmq.NOBLOCK) # type: ignore[misc] + + if msg != b"READY": + worker_id = ident.decode("utf-8") + worker_msg = msg.decode("utf-8") + raise RuntimeError( + f"Unexpected message from worker {worker_id}:" + f"{worker_msg}" + ) + + print(f"[debug] Worker {ident.decode('utf-8')} is ready.") + ready_workers.add(ident) + except zmq.Again: + total_wait_time = time.perf_counter() - start_time + if total_wait_time > timeout: + raise TimeoutError( + f"Timeout reached after {timeout} seconds" + "while waiting for workers to be ready." + ) + await asyncio.sleep(0.5) + self.ensure_alive() + except Exception as e: + errs = self.gather_errors() if self.context else [] + print( + f"[debug] Error during startup (attempt" + f" {attempt + 1}/{max_attempts}): {e}\n{traceback.format_exc()}" + ) + self.terminate(timeout=timeout) + self.context = None + self.close_zmq_socket() + + can_retry = ( + (pick_master_port or pick_worker_port) + and attempt < max_attempts - 1 + and _check_port_in_use(errs, e) + ) + if not can_retry: + raise RuntimeError("Failed to start distributed processes.") from e - # Start the keepalive timer - self.maybe_start_keepalive() + await asyncio.sleep(0.1) + continue + + print("[debug] All workers are ready and running.") + # Start the keepalive timer + self.maybe_start_keepalive() + return def keepalive(self, timeout: Optional[Union[int, float]] = 60.0) -> None: """ From 4a2fc39b2c2e2dfa553a403f8aac450c869a5b63 Mon Sep 17 00:00:00 2001 From: Gothos Date: Wed, 11 Mar 2026 21:56:15 +0530 Subject: [PATCH 2/2] misc: fix typing for py3.8 --- projects/fal/src/fal/distributed/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/projects/fal/src/fal/distributed/utils.py b/projects/fal/src/fal/distributed/utils.py index 9995fb202..57b0ebce7 100644 --- a/projects/fal/src/fal/distributed/utils.py +++ b/projects/fal/src/fal/distributed/utils.py @@ -326,7 +326,7 @@ def _pick_random_free_port(addr: str) -> int: return s.getsockname()[1] -def _check_port_in_use(errs: list[Exception], exc: Exception) -> bool: +def _check_port_in_use(errs: list, exc: Exception) -> bool: txt = f"{exc}\n" + "\n".join(str(e) for e in errs) txt = txt.lower() return "address already in use" in txt or "eaddrinuse" in txt