Fabric: reserve externally provided MASTER_PORT values

Author: littlebullGit

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

@@ -111,17 +111,24 @@ def find_free_network_port() -> int:
     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

@@ -89,6 +89,26 @@ def release_all(self) -> None:
         with self._lock:
             self._allocated_ports.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
+
+            self._allocated_ports.add(port)
+            return True
+
     @contextmanager
     def allocated_port(self, preferred_port: Optional[int] = None) -> Iterator[int]:
         """Context manager for automatic port cleanup.

tests/tests_fabric/conftest.py

@@ -97,6 +97,9 @@ def teardown_process_group():
     if port_to_release is not None:
         get_port_manager().release_port(port_to_release)
 
+    # Remove the MASTER_PORT so subsequent tests don't reuse the same value
+    os.environ.pop("MASTER_PORT", None)
+
 
 @pytest.fixture(autouse=True)
 def thread_police_duuu_daaa_duuu_daaa():

tests/tests_fabric/utilities/test_port_manager.py

@@ -18,9 +18,45 @@
 import threading
 from collections import Counter
 
+import pytest
+
 from lightning.fabric.plugins.environments.lightning import find_free_network_port
 from lightning.fabric.utilities.port_manager import PortManager, get_port_manager
 
+# =============================================================================
+# Fixtures
+# =============================================================================
+
+
+@pytest.fixture
+def with_master_port():
+    """Fixture that sets MASTER_PORT before test runs, for conftest coverage."""
+    port = find_free_network_port()
+    previous_value = os.environ.get("MASTER_PORT")
+    os.environ["MASTER_PORT"] = str(port)
+    try:
+        yield port
+    finally:
+        if previous_value is None:
+            os.environ.pop("MASTER_PORT", None)
+        else:
+            os.environ["MASTER_PORT"] = previous_value
+
+
+@pytest.fixture
+def with_invalid_master_port():
+    """Fixture that sets invalid MASTER_PORT to test error handling."""
+    previous_value = os.environ.get("MASTER_PORT")
+    os.environ["MASTER_PORT"] = "not_a_valid_port_number"
+    try:
+        yield
+    finally:
+        if previous_value is None:
+            os.environ.pop("MASTER_PORT", None)
+        else:
+            os.environ["MASTER_PORT"] = previous_value
+
+
 # =============================================================================
 # Unit Tests for PortManager
 # =============================================================================
@@ -135,12 +171,19 @@ def test_port_manager_allocation_failure():
     """Test that PortManager raises error when unable to allocate after max attempts."""
     manager = PortManager()
 
-    # This is hard to test without actually exhausting ports, but we can test
-    # the error path by mocking or just ensure the code path exists
-    # For now, just verify that max_attempts parameter exists
-    port = manager.allocate_port(max_attempts=1)
+    # Pre-allocate a large number of ports to make it harder to find a free one
+    # Then try with max_attempts=1 which should fail quickly
+    allocated_ports = [manager.allocate_port() for _ in range(50)]
+
+    # Test that it can still allocate with enough attempts
+    port = manager.allocate_port(max_attempts=100)
     assert port >= 1024
 
+    # Clean up
+    for p in allocated_ports:
+        manager.release_port(p)
+    manager.release_port(port)
+
 
 def test_port_manager_prevents_reallocation():
     """Test that a port won't be allocated twice until released."""
@@ -160,6 +203,10 @@ def test_port_manager_prevents_reallocation():
     manager.release_port(port1)
     assert port1 not in manager._allocated_ports
 
+    # Clean up
+    for port in more_ports:
+        manager.release_port(port)
+
 
 def test_get_port_manager_singleton():
     """Test that get_port_manager returns the same instance."""
@@ -173,6 +220,9 @@ def test_get_port_manager_singleton():
     port = manager1.allocate_port()
     assert port in manager2._allocated_ports
 
+    # Clean up
+    manager1.release_port(port)
+
 
 def test_get_port_manager_thread_safe_singleton():
     """Test that get_port_manager creates singleton safely across threads."""
@@ -280,6 +330,42 @@ def test_port_manager_atexit_cleanup():
     assert len(manager._allocated_ports) == 0
 
 
+def test_port_manager_reserve_existing_port_free():
+    """reserve_existing_port should succeed for free ports and track them."""
+    manager = PortManager()
+
+    port = manager._find_free_port()
+    assert manager.reserve_existing_port(port)
+    assert port in manager._allocated_ports
+
+    # Second call should succeed but not duplicate
+    assert manager.reserve_existing_port(port)
+    assert len(manager._allocated_ports) == 1
+
+
+def test_port_manager_reserve_existing_port_invalid_value():
+    """reserve_existing_port should reject invalid port numbers."""
+    manager = PortManager()
+
+    assert not manager.reserve_existing_port(0)
+    assert not manager.reserve_existing_port(-1)
+    assert not manager.reserve_existing_port(70000)
+
+
+def test_port_manager_reserve_existing_port_after_release():
+    """Ports released from sockets should become reservable."""
+    manager = PortManager()
+
+    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+    s.bind(("", 0))
+    reusable_port = s.getsockname()[1]
+    s.close()
+
+    assert manager.reserve_existing_port(reusable_port)
+    assert reusable_port in manager._allocated_ports
+
+
 def test_port_manager_context_manager():
     """Test that context manager automatically releases ports."""
     manager = PortManager()
@@ -430,6 +516,44 @@ def test_port_allocation_simulates_distributed_test_lifecycle():
     assert len(manager._allocated_ports) == initial_count
 
 
+def test_conftest_cleanup_with_master_port_set(with_master_port):
+    """Test conftest cleanup when MASTER_PORT is set before test starts.
+
+    This test uses a fixture to set MASTER_PORT before the test runs, allowing the conftest teardown_process_group
+    fixture to capture and clean it up. This ensures the conftest cleanup code is covered.
+
+    """
+    manager = get_port_manager()
+    port = with_master_port  # Port was set by fixture
+
+    # Verify port is allocated
+    assert port in manager._allocated_ports
+    assert os.environ.get("MASTER_PORT") == str(port)
+
+    # Leave MASTER_PORT set - conftest teardown will clean it up
+    # After this test, teardown_process_group will:
+    # 1. Detect MASTER_PORT in os.environ (line captured before yield)
+    # 2. Call get_port_manager().release_port(port)
+    # 3. Port gets released back to manager
+
+
+def test_conftest_handles_invalid_master_port(with_invalid_master_port):
+    """Test conftest handles invalid MASTER_PORT gracefully.
+
+    This exercises the contextlib.suppress(ValueError, KeyError) path in the conftest teardown_process_group fixture.
+
+    """
+    # Fixture set MASTER_PORT to "not_a_valid_port_number"
+    # The conftest will try to parse it: int(os.environ["MASTER_PORT"])
+    # This will raise ValueError, which should be caught by contextlib.suppress
+
+    # Verify the invalid value is set
+    assert os.environ.get("MASTER_PORT") == "not_a_valid_port_number"
+
+    # This test just needs to complete without crashing
+    # The conftest teardown will handle the ValueError gracefully
+
+
 def test_multiple_tests_can_reuse_ports_after_release():
     """Test that ports can be reused after being released."""
     manager = get_port_manager()
@@ -521,3 +645,49 @@ def test_port_manager_survives_multiple_test_sessions():
     # Clean up
     for port in session2_ports + session3_ports:
         manager.release_port(port)
+
+
+def test_port_manager_allocation_runtime_error():
+    """Test that allocation fails gracefully when max_attempts is exhausted."""
+    manager = PortManager()
+
+    # Mock the _find_free_port to always return a port that's already allocated
+    # This will cause max_attempts to be exhausted
+    allocated_port = manager.allocate_port()
+
+    # Save original method
+    original_find = manager._find_free_port
+
+    # Make _find_free_port always return the already-allocated port
+    def always_return_allocated():
+        return allocated_port
+
+    manager._find_free_port = always_return_allocated
+
+    # This should raise RuntimeError after max_attempts
+    with pytest.raises(RuntimeError, match="Failed to allocate a free port after .* attempts"):
+        manager.allocate_port(max_attempts=5)
+
+    # Restore original method and clean up
+    manager._find_free_port = original_find
+    manager.release_port(allocated_port)
+
+
+def test_find_free_network_port_respects_existing_master_port(with_master_port):
+    """find_free_network_port should reuse externally provided MASTER_PORT."""
+    manager = get_port_manager()
+    port = with_master_port
+
+    returned_port = find_free_network_port()
+    assert returned_port == port
+    assert port in manager._allocated_ports
+
+
+def test_find_free_network_port_handles_invalid_master_port(with_invalid_master_port):
+    """Invalid MASTER_PORT values should fall back to allocating a fresh port."""
+    manager = get_port_manager()
+
+    returned_port = find_free_network_port()
+    assert isinstance(returned_port, int)
+    assert returned_port in manager._allocated_ports
+    assert returned_port != "not_a_valid_port_number"