test(fabric): bound port manager queue tests and update warning log

Author: littlebullGit

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

@@ -17,7 +17,12 @@
 from typing_extensions import override
 
 from lightning.fabric.plugins.environments.cluster_environment import ClusterEnvironment
-from lightning.fabric.utilities.port_manager import get_port_manager
+from lightning.fabric.utilities.port_manager import (
+    find_free_network_port as _pm_find_free_network_port,
+)
+from lightning.fabric.utilities.port_manager import (
+    get_port_manager,
+)
 from lightning.fabric.utilities.rank_zero import rank_zero_only
 
 
@@ -64,7 +69,7 @@ def main_address(self) -> str:
     def main_port(self) -> int:
         if self._main_port == -1:
             self._main_port = (
-                int(os.environ["MASTER_PORT"]) if "MASTER_PORT" in os.environ else find_free_network_port()
+                int(os.environ["MASTER_PORT"]) if "MASTER_PORT" in os.environ else _pm_find_free_network_port()
             )
         return self._main_port
 
@@ -115,27 +120,8 @@ def teardown(self) -> 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.
-
-    The allocated port is reserved and won't be returned by subsequent calls until it's explicitly released.
-
-    Returns:
-        A port number that is reserved and free at the time of allocation
+    Deprecated alias. Use :func:`lightning.fabric.utilities.port_manager.find_free_network_port` instead.
 
     """
-    # 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()
+
+    return _pm_find_free_network_port()

src/lightning/fabric/utilities/port_manager.py

@@ -233,8 +233,9 @@ def allocate_port(self, preferred_port: Optional[int] = None, max_attempts: int
 
                     queue_count = len(state.recently_released)
                     if queue_count > 800:  # >78% of typical 1024 capacity
+                        utilization_pct = (queue_count / _RECENTLY_RELEASED_PORTS_MAXLEN) * 100
                         log.warning(
-                            f"Port queue utilization high: {queue_count} entries. "
+                            f"Port queue utilization high: {queue_count} entries ({utilization_pct:.1f}% of capacity). "
                             f"Allocated port {port}. Active allocations: {len(state.allocated_ports)}"
                         )
 
@@ -550,3 +551,26 @@ def get_port_manager() -> PortManager:
             if _port_manager is None:
                 _port_manager = PortManager()
     return _port_manager
+
+
+def find_free_network_port() -> int:
+    """Find and reserve a free network port using the global port manager.
+
+    Returns:
+        A port number that is reserved and free at the time of allocation.
+
+    """
+
+    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()

tests/tests_fabric/utilities/test_port_manager.py

@@ -13,15 +13,23 @@
 # limitations under the License.
 """Tests for the PortManager utility and port allocation integration."""
 
+import collections
 import os
 import socket
 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
+import lightning.fabric.utilities.port_manager as port_manager_module
+import lightning.fabric.utilities.port_state as port_state_module
+from lightning.fabric.plugins.environments.lightning import (
+    find_free_network_port as env_find_free_network_port,
+)
+from lightning.fabric.utilities.port_manager import (
+    PortManager,
+    find_free_network_port,
+    get_port_manager,
+)
 
 # =============================================================================
 # Fixtures
@@ -143,7 +151,7 @@ def allocate_ports():
     assert len(set(ports)) == 100, f"Expected 100 unique ports, got {len(set(ports))}"
 
     # Check for any duplicates
-    counts = Counter(ports)
+    counts = collections.Counter(ports)
     duplicates = {port: count for port, count in counts.items() if count > 1}
     assert not duplicates, f"Found duplicate ports: {duplicates}"
 
@@ -495,6 +503,20 @@ def allocate():
         manager.release_port(port)
 
 
+def test_find_free_network_port_alias(monkeypatch):
+    """Legacy environment alias should reuse the port manager backed implementation."""
+
+    manager = get_port_manager()
+    manager.release_all()
+
+    port = env_find_free_network_port()
+
+    try:
+        assert port in manager._allocated_ports
+    finally:
+        manager.release_port(port)
+
+
 def test_port_allocation_simulates_distributed_test_lifecycle():
     """Simulate the lifecycle of a distributed test with port allocation and release."""
     manager = get_port_manager()
@@ -714,9 +736,14 @@ def test_port_manager_recently_released_prevents_immediate_reuse():
     manager.release_port(new_port)
 
 
-def test_port_manager_recently_released_queue_cycles():
+def _set_recently_released_limit(monkeypatch, value: int) -> None:
+    monkeypatch.setattr(port_manager_module, "_RECENTLY_RELEASED_PORTS_MAXLEN", value, raising=True)
+    monkeypatch.setattr(port_state_module, "_RECENTLY_RELEASED_MAX_LEN", value, raising=True)
+
+
+def test_port_manager_recently_released_queue_cycles(monkeypatch):
     """Test that recently_released queue cycles after maxlen allocations."""
-    from lightning.fabric.utilities.port_manager import _RECENTLY_RELEASED_PORTS_MAXLEN
+    _set_recently_released_limit(monkeypatch, 64)
 
     manager = PortManager()
 
@@ -727,8 +754,10 @@ def test_port_manager_recently_released_queue_cycles():
     # Port should be in recently_released queue
     assert first_port in manager._recently_released
 
+    queue_limit = port_manager_module._RECENTLY_RELEASED_PORTS_MAXLEN
+
     # Allocate and release many ports to fill the queue beyond maxlen
-    for _ in range(_RECENTLY_RELEASED_PORTS_MAXLEN + 10):
+    for _ in range(queue_limit + 10):
         port = manager.allocate_port()
         manager.release_port(port)
 
@@ -755,14 +784,20 @@ def test_port_manager_reserve_clears_recently_released():
     manager.release_port(port)
 
 
-def test_port_manager_high_queue_utilization_warning(caplog):
+def test_port_manager_high_queue_utilization_warning(monkeypatch, caplog):
     """Test that warning is logged when queue utilization exceeds 80%."""
     import logging
 
+    _set_recently_released_limit(monkeypatch, 64)
+
+    queue_limit = port_manager_module._RECENTLY_RELEASED_PORTS_MAXLEN
+    trigger_count = int(queue_limit * 0.8) + 1  # Just over 80%
+    expected_pct = (trigger_count / queue_limit) * 100
+
     manager = PortManager()
 
-    # Fill queue to >80% (821/1024 = 80.2%)
-    for _ in range(821):
+    # Fill queue to just over 80%
+    for _ in range(trigger_count):
         port = manager.allocate_port()
         manager.release_port(port)
 
@@ -773,4 +808,4 @@ def test_port_manager_high_queue_utilization_warning(caplog):
 
     # Verify warning was logged
     assert any("Port queue utilization high" in record.message for record in caplog.records)
-    assert any("80." in record.message for record in caplog.records)  # Should show 80.x%
+    assert any(f"{expected_pct:.1f}%" in record.message for record in caplog.records)

tests/tests_fabric/utilities/test_port_manager_process_safe.py

@@ -22,6 +22,7 @@
 
 import pytest
 
+import lightning.fabric.utilities.port_manager as port_manager_module
 from lightning.fabric.utilities.file_lock import UnixFileLock, WindowsFileLock, create_file_lock
 from lightning.fabric.utilities.port_manager import PortManager, _get_lock_dir, _get_lock_file
 from lightning.fabric.utilities.port_state import PortAllocation, PortState
@@ -106,6 +107,43 @@ def test_file_lock_context_manager_timeout(tmpdir):
     lock1.release()
 
 
+def test_get_lock_dir_handles_permission_error(monkeypatch, tmp_path):
+    """_get_lock_dir should tolerate probe unlink permission errors and register cleanup."""
+
+    monkeypatch.setenv("LIGHTNING_PORT_LOCK_DIR", str(tmp_path))
+
+    registered_calls = []
+
+    def fake_register(func, *args, **kwargs):
+        registered_calls.append((func, args, kwargs))
+        return func
+
+    monkeypatch.setattr(port_manager_module.atexit, "register", fake_register)
+
+    original_unlink = Path.unlink
+    call_state = {"count": 0}
+
+    def fake_unlink(self, *args, **kwargs):
+        if self.name.startswith(".lightning_port_manager_write_test_") and call_state["count"] == 0:
+            call_state["count"] += 1
+            raise PermissionError("locked")
+        return original_unlink(self, *args, **kwargs)
+
+    monkeypatch.setattr(Path, "unlink", fake_unlink)
+
+    lock_dir = _get_lock_dir()
+    assert Path(lock_dir) == tmp_path
+    assert registered_calls, "Cleanup should be registered when unlink fails"
+
+    cleanup_func, args, kwargs = registered_calls[0]
+    probe_path = args[0]
+    assert isinstance(probe_path, Path)
+    assert probe_path.exists()
+
+    cleanup_func(*args, **kwargs)
+    assert not probe_path.exists()
+
+
 # =============================================================================
 # Tests for PortState
 # =============================================================================