Fabric: reserve externally provided MASTER_PORT values

Author: littlebullGit

src/lightning/fabric/plugins/environments/lightning.py

@@ -104,24 +104,38 @@ def teardown(self) -> None:
         if "WORLD_SIZE" in os.environ:
             del os.environ["WORLD_SIZE"]
 
+        if self._main_port != -1:
+            get_port_manager().release_port(self._main_port)
+            self._main_port = -1
+
+        os.environ.pop("MASTER_PORT", None)
+        os.environ.pop("MASTER_ADDR", None)
+
 
 def find_free_network_port() -> int:
     """Finds a free port on localhost.
 
     It is useful in single-node training when we don't want to connect to a real main node but have to set the
     `MASTER_PORT` environment variable.
 
-    This function uses a global port manager to prevent internal race conditions within the test suite.
     The allocated port is reserved and won't be returned by subsequent calls until it's explicitly released.
 
-    Note:
-        While this prevents collisions between concurrent Lightning tests, external processes can still
-        claim the port between allocation and binding. For production use, explicitly set the MASTER_PORT
-        environment variable.
-
     Returns:
         A port number that is reserved and free at the time of allocation
 
     """
+    # If an external launcher already specified a MASTER_PORT (for example, torch.distributed.spawn or
+    # multiprocessing helpers), reserve it through the port manager so no other test reuses the same number.
+    if "MASTER_PORT" in os.environ:
+        master_port_str = os.environ["MASTER_PORT"]
+        try:
+            existing_port = int(master_port_str)
+        except ValueError:
+            pass
+        else:
+            port_manager = get_port_manager()
+            if port_manager.reserve_existing_port(existing_port):
+                return existing_port
+
     port_manager = get_port_manager()
     return port_manager.allocate_port()

src/lightning/fabric/utilities/port_manager.py

@@ -16,10 +16,14 @@
 import atexit
 import socket
 import threading
+from collections import deque
 from collections.abc import Iterator
 from contextlib import contextmanager
 from typing import Optional
 
+# Maximum number of recently released ports to track before reuse
+_RECENTLY_RELEASED_PORTS_MAXLEN = 256
+
 
 class PortManager:
     """Thread-safe port manager to prevent EADDRINUSE errors.
@@ -33,10 +37,14 @@ class PortManager:
     def __init__(self) -> None:
         self._lock = threading.Lock()
         self._allocated_ports: set[int] = set()
+        # Recently released ports are kept in a queue to avoid immediate reuse
+        self._recently_released: deque[int] = deque(maxlen=_RECENTLY_RELEASED_PORTS_MAXLEN)
+        # Counter to vary starting position on each allocation
+        self._allocation_counter = 0
         # Register cleanup to release all ports on exit
         atexit.register(self.release_all)
 
-    def allocate_port(self, preferred_port: Optional[int] = None, max_attempts: int = 100) -> int:
+    def allocate_port(self, preferred_port: Optional[int] = None, max_attempts: int = 1000) -> int:
         """Allocate a free port, ensuring it's not already reserved.
 
         Args:
@@ -55,23 +63,42 @@ def allocate_port(self, preferred_port: Optional[int] = None, max_attempts: int
             if (
                 preferred_port is not None
                 and preferred_port not in self._allocated_ports
+                and preferred_port not in self._recently_released
                 and self._is_port_free(preferred_port)
             ):
                 self._allocated_ports.add(preferred_port)
                 return preferred_port
 
-            # Try to find a free port
+            # Ephemeral port range (49152-65535)
+            # We'll search through this range to find a port that's:
+            # 1. Not in our allocated set
+            # 2. Not in our recently_released queue
+            # 3. Actually free according to the OS
+
+            # Use combination of thread ID and allocation counter to vary starting point
+            # This distributes allocation across the port range and avoids hotspots
+            port_range_size = 65535 - 49152 + 1
+            thread_id = threading.get_ident()
+            self._allocation_counter += 1
+            start_offset = (hash(thread_id) + self._allocation_counter) % port_range_size
+
             for attempt in range(max_attempts):
-                port = self._find_free_port()
+                # Cycle through port range with varied offset
+                port = 49152 + ((start_offset + attempt) % port_range_size)
 
-                # Double-check it's not in our reserved set (shouldn't happen, but be safe)
-                if port not in self._allocated_ports:
+                # Skip if in our tracking structures
+                if port in self._allocated_ports or port in self._recently_released:
+                    continue
+
+                # Check if actually free
+                if self._is_port_free(port):
                     self._allocated_ports.add(port)
                     return port
 
             raise RuntimeError(
                 f"Failed to allocate a free port after {max_attempts} attempts. "
-                f"Currently allocated ports: {len(self._allocated_ports)}"
+                f"Currently allocated: {len(self._allocated_ports)}, "
+                f"recently released: {len(self._recently_released)}"
             )
 
     def release_port(self, port: int) -> None:
@@ -82,12 +109,43 @@ def release_port(self, port: int) -> None:
 
         """
         with self._lock:
-            self._allocated_ports.discard(port)
+            if port in self._allocated_ports:
+                self._allocated_ports.remove(port)
+                # Add to the back of the queue; oldest will be evicted when queue is full
+                self._recently_released.append(port)
 
     def release_all(self) -> None:
         """Release all allocated ports."""
         with self._lock:
             self._allocated_ports.clear()
+            self._recently_released.clear()
+
+    def reserve_existing_port(self, port: int) -> bool:
+        """Reserve a port that was allocated externally.
+
+        Args:
+            port: The externally assigned port to reserve.
+
+        Returns:
+            True if the port was reserved (or already reserved), False if the port value is invalid.
+
+        """
+        if port <= 0 or port > 65535:
+            return False
+
+        with self._lock:
+            if port in self._allocated_ports:
+                return True
+
+            # Remove from recently released queue if present (we're explicitly reserving it)
+            if port in self._recently_released:
+                # Create a new deque without this port
+                self._recently_released = deque(
+                    (p for p in self._recently_released if p != port), maxlen=_RECENTLY_RELEASED_PORTS_MAXLEN
+                )
+
+            self._allocated_ports.add(port)
+            return True
 
     @contextmanager
     def allocated_port(self, preferred_port: Optional[int] = None) -> Iterator[int]:
@@ -121,7 +179,8 @@ def _find_free_port() -> int:
 
         """
         s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        # Don't use SO_REUSEADDR - we need to match the behavior of TCPStore
+        # which binds without it, so ports in TIME_WAIT will be rejected
         s.bind(("", 0))
         port = s.getsockname()[1]
         s.close()
@@ -140,7 +199,8 @@ def _is_port_free(port: int) -> bool:
         """
         try:
             s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-            s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+            # Don't use SO_REUSEADDR - we need to match the behavior of TCPStore
+            # which binds without it, so ports in TIME_WAIT will be rejected
             s.bind(("", port))
             s.close()
             return True

tests/tests_fabric/conftest.py

@@ -82,20 +82,20 @@ def teardown_process_group():
 
     from lightning.fabric.utilities.port_manager import get_port_manager
 
-    # Record the port used in this test (if any)
-    port_to_release = None
-    if "MASTER_PORT" in os.environ:
-        with contextlib.suppress(ValueError, KeyError):
-            port_to_release = int(os.environ["MASTER_PORT"])
-
     yield
 
     # Clean up distributed connection
     _destroy_dist_connection()
 
-    # Release the port from the manager so it can be reused
-    if port_to_release is not None:
-        get_port_manager().release_port(port_to_release)
+    manager = get_port_manager()
+
+    # If a process group created or updated MASTER_PORT during the test, reserve it and then clear it
+    if "MASTER_PORT" in os.environ:
+        with contextlib.suppress(ValueError):
+            port = int(os.environ["MASTER_PORT"])
+            manager.reserve_existing_port(port)
+            manager.release_port(port)
+        os.environ.pop("MASTER_PORT", None)
 
 
 @pytest.fixture(autouse=True)

tests/tests_fabric/plugins/environments/test_lightning.py

@@ -17,6 +17,7 @@
 import pytest
 
 from lightning.fabric.plugins.environments import LightningEnvironment
+from lightning.fabric.utilities.port_manager import get_port_manager
 
 
 @mock.patch.dict(os.environ, {}, clear=True)
@@ -84,5 +85,18 @@ def test_teardown():
     assert "WORLD_SIZE" not in os.environ
 
 
+@mock.patch.dict(os.environ, {}, clear=True)
+def test_teardown_releases_port_and_env():
+    env = LightningEnvironment()
+    port = env.main_port
+    assert port in get_port_manager()._allocated_ports
+
+    env.teardown()
+
+    assert port not in get_port_manager()._allocated_ports
+    assert "MASTER_PORT" not in os.environ
+    assert "MASTER_ADDR" not in os.environ
+
+
 def test_detect():
     assert LightningEnvironment.detect()