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,6 +37,8 @@ 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)
 
@@ -55,6 +61,7 @@ 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)
@@ -64,7 +71,10 @@ def allocate_port(self, preferred_port: Optional[int] = None, max_attempts: int
             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)
+                # Skip ports that were recently released to avoid TIME_WAIT conflicts
+                if port in self._recently_released:
+                    continue
+
                 if port not in self._allocated_ports:
                     self._allocated_ports.add(port)
                     return port
@@ -82,12 +92,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]:

tests/tests_fabric/conftest.py

@@ -94,8 +94,16 @@ def teardown_process_group():
     _destroy_dist_connection()
 
     # Release the port from the manager so it can be reused
+    manager = get_port_manager()
+
     if port_to_release is not None:
-        get_port_manager().release_port(port_to_release)
+        manager.release_port(port_to_release)
+
+    # If a process group created or updated MASTER_PORT during the test, make sure we track it and then clear it.
+    if "MASTER_PORT" in os.environ:
+        with contextlib.suppress(ValueError):
+            manager.reserve_existing_port(int(os.environ["MASTER_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()