Use find free port (#877)

Author: erikwijmans

habitat_baselines/rl/ddppo/ddp_utils.py

@@ -1,6 +1,8 @@
+import contextlib
 import functools
 import os
 import signal
+import socket
 import subprocess
 import threading
 from os import path as osp
@@ -267,3 +269,21 @@ def init_distrib_slurm(
     )
 
     return local_rank, tcp_store
+
+
+def find_free_port() -> int:
+    """
+    Returns a free port on the system.
+    Note that this can only be used to find a port for torch.distribted
+    if it's called by a process on the node that will have
+    world_rank == 0 and then all ranks are created. If you
+    just called `find_free_port()` on each rank independently, every
+    rank will have a different port!
+    """
+    with contextlib.closing(
+        socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    ) as sock:
+        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        sock.bind(("localhost", 0))
+        _, port = sock.getsockname()
+        return port

test/test_baseline_trainers.py

@@ -23,6 +23,7 @@
     from habitat_baselines.common.base_trainer import BaseRLTrainer
     from habitat_baselines.common.baseline_registry import baseline_registry
     from habitat_baselines.config.default import get_config
+    from habitat_baselines.rl.ddppo.ddp_utils import find_free_port
     from habitat_baselines.run import execute_exp, run_exp
     from habitat_baselines.utils.common import (
         ObservationBatchingCache,
@@ -77,8 +78,8 @@ def _powerset(s):
     ),
 )
 def test_trainers(test_cfg_path, mode, gpu2gpu, observation_transforms):
-    # For testing with world_size=1, -1 works as port in PyTorch
-    os.environ["MAIN_PORT"] = str(-1)
+    # For testing with world_size=1
+    os.environ["MAIN_PORT"] = str(find_free_port())
 
     if gpu2gpu:
         try: