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,16 @@
 import atexit
 import socket
 import threading
+from collections import deque
 from collections.abc import Iterator
 from contextlib import contextmanager
 from typing import Optional
 
+# Size of the recently released ports queue
+# This prevents immediate reuse of ports that were just released
+# Increased to 1024 to reduce the chance of cycling back to TIME_WAIT ports
+_RECENTLY_RELEASED_PORTS_MAXLEN = 1024
+
 
 class PortManager:
     """Thread-safe port manager to prevent EADDRINUSE errors.
@@ -33,10 +39,12 @@ 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)
         # 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,27 @@ 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
+            # Let the OS choose a free port, but verify it's not in our tracking structures
+            # The OS naturally avoids ports in TIME_WAIT (without SO_REUSEADDR)
             for attempt in range(max_attempts):
                 port = self._find_free_port()
 
-                # Double-check it's not in our reserved set (shouldn't happen, but be safe)
-                if port not in self._allocated_ports:
+                # Skip if already allocated by us or recently released
+                # This prevents race conditions within our process
+                if port not in self._allocated_ports and port not in self._recently_released:
                     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 +94,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 +164,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 +184,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)
@@ -221,6 +221,56 @@ def leave_no_artifacts_behind():
     assert not difference, f"Test left artifacts behind: {difference}"
 
 
+def pytest_runtest_makereport(item: pytest.Item, call: pytest.CallInfo) -> None:
+    """Retry tests that fail with EADDRINUSE errors.
+
+    This handles the race condition where a port might enter TIME_WAIT between our port allocation check and actual
+    binding by TCPStore.
+
+    """
+    if call.excinfo is not None and call.when == "call":
+        exception_msg = str(call.excinfo.value)
+        # Check if this is an EADDRINUSE error from distributed training
+        if "EADDRINUSE" in exception_msg or "address already in use" in exception_msg.lower():
+            # Get the retry count from the test node
+            retry_count = getattr(item, "_port_retry_count", 0)
+            max_retries = 3
+
+            if retry_count < max_retries:
+                # Increment retry counter
+                item._port_retry_count = retry_count + 1
+
+                # Log the retry
+                if hasattr(item.config, "get_terminal_writer"):
+                    writer = item.config.get_terminal_writer()
+                    writer.write(
+                        f"\n[Port conflict detected] Retrying test {item.name} "
+                        f"(attempt {retry_count + 2}/{max_retries + 1})...\n",
+                        yellow=True,
+                    )
+
+                # Clear the port manager's state to get fresh ports
+                from lightning.fabric.utilities.port_manager import get_port_manager
+
+                manager = get_port_manager()
+                manager.release_all()
+
+                # Re-run the test by raising Rerun exception
+                # Note: This requires pytest-rerunfailures plugin
+                import time
+
+                time.sleep(0.5)  # Brief delay to let ports settle
+
+                # If pytest-rerunfailures is available, use it
+                try:
+                    from pytest_rerunfailures import Rerun
+
+                    raise Rerun(f"Port conflict (EADDRINUSE), retry {retry_count + 1}/{max_retries}")
+                except ImportError:
+                    # Plugin not available, just let the test fail
+                    pass
+
+
 def pytest_collection_modifyitems(items: list[pytest.Function], config: pytest.Config) -> None:
     """An adaptation of `tests/tests_pytorch/conftest.py::pytest_collection_modifyitems`"""
     initial_size = len(items)

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()