add searching for free ports in unit tests

tests/helpers/utils.py

@@ -2,12 +2,25 @@
 import os
 import uuid
 from typing import Any, Dict, List, Optional, Tuple
+import random
+import socket
 
 import torch.cuda
 from nanotron.parallel import ParallelContext
 from torch.distributed.launcher import elastic_launch
 
 
+def find_free_port(min_port: int = 2000, max_port: int = 65000) -> int:
+    while True:
+        port = random.randint(min_port, max_port)
+        try:
+            with socket.socket() as sock:
+                sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+                sock.bind(("localhost", port))
+                return port
+        except OSError as e:
+            raise e
+
 def available_gpus():
     if not torch.cuda.is_available():
         return 0
@@ -92,6 +105,8 @@ def _init_distributed(func):
         """
         nb_gpus = tp * dp * pp
         run_id = uuid.uuid4()
+
+        port = find_free_port()
 
         config = torch.distributed.launcher.LaunchConfig(
             min_nodes=1,
@@ -101,7 +116,7 @@ def _init_distributed(func):
             rdzv_configs={"timeout": 60},
             # Setting port to `0` allows `torch` to randomly pick a port: https://pytorch.org/docs/stable/elastic/run.html#stacked-single-node-multi-worker
             # Works only for single node workload.
-            rdzv_endpoint="localhost:0",
+            rdzv_endpoint=f"localhost:{port}",
             run_id=str(run_id),
             max_restarts=0,
             # TODO @thomasw21: Tune as we increase the number of tests