[misc] Find unused port in distributed init (#475)

Author: Edenzzzz

fastvideo/v1/tests/old_tests/distributed_example.py

@@ -29,14 +29,14 @@
 class Worker:
 
     def __init__(self, fastvideo_args: FastVideoArgs, local_rank: int,
-                 rank: int, pipe):
+                 rank: int, pipe, master_port: int):
         self.fastvideo_args = fastvideo_args
         self.local_rank = local_rank
         self.rank = rank
         # TODO(will): don't hardcode this
         self.distributed_init_method = "env://"
         self.pipe = pipe
-
+        self.master_port = master_port
         self.init_device()
 
         # Init request dispatcher
@@ -76,7 +76,7 @@ def init_device(self) -> None:
                 f"Unsupported device: {self.fastvideo_args.device_str}")
 
         os.environ["MASTER_ADDR"] = "localhost"
-        os.environ["MASTER_PORT"] = "29503"
+        os.environ["MASTER_PORT"] = str(self.master_port)
         os.environ["LOCAL_RANK"] = str(self.local_rank)
         os.environ["RANK"] = str(self.rank)
 
@@ -191,7 +191,7 @@ def init_worker_distributed_environment(
 
 
 def run_worker_process(fastvideo_args: FastVideoArgs, local_rank: int,
-                       rank: int, pipe):
+                       rank: int, pipe, master_port: int):
     # Add process-specific prefix to stdout and stderr
     process_name = mp.current_process().name
     pid = os.getpid()
@@ -206,8 +206,9 @@ def run_worker_process(fastvideo_args: FastVideoArgs, local_rank: int,
     logger.info("Worker %d initializing...",
                 rank,
                 local_main_process_only=False)
+
     try:
-        worker = Worker(fastvideo_args, local_rank, rank, pipe)
+        worker = Worker(fastvideo_args, local_rank, rank, pipe, master_port)
         logger.info("Worker %d sending ready", rank)
         pipe.send({
             "status": "ready",

fastvideo/v1/worker/gpu_worker.py

@@ -3,6 +3,7 @@
 import multiprocessing as mp
 import os
 import signal
+import socket
 import time
 from multiprocessing.process import BaseProcess
 from typing import Any, Callable, List, Optional, Union, cast
@@ -28,6 +29,15 @@ def _init_executor(self) -> None:
 
         self.workers: List[BaseProcess] = []
         self.worker_pipes = []
+        self.master_port = None
+
+        for port in range(29503, 65535):
+            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+                if s.connect_ex(('localhost', port)) != 0:
+                    self.master_port = port
+                    break
+        if self.master_port is None:
+            raise ValueError("No unused port found to use as master port")
 
         # Create pipes and start workers
         for rank in range(self.world_size):
@@ -39,7 +49,8 @@ def _init_executor(self) -> None:
                                 kwargs=dict(fastvideo_args=self.fastvideo_args,
                                             local_rank=rank,
                                             rank=rank,
-                                            pipe=worker_pipe))
+                                            pipe=worker_pipe,
+                                            master_port=self.master_port))
             worker.start()
             self.workers.append(worker)