improve testing fixtures

Author: alec-flowers

components/backends/vllm/src/dynamo/tests/test_ports.py

@@ -13,7 +13,6 @@
     DynamoPortRange,
     EtcdContext,
     PortAllocationRequest,
-    PortBinding,
     PortMetadata,
     allocate_and_reserve_port,
     allocate_and_reserve_port_block,
@@ -50,19 +49,6 @@ def test_invalid_port_range_min_greater_than_max(self):
             DynamoPortRange(min=3000, max=3000)
 
 
-class TestEtcdContext:
-    """Test ETCD context functionality."""
-
-    def test_make_port_key(self):
-        """Test port key generation."""
-        mock_client = Mock()
-        context = EtcdContext(client=mock_client, namespace="test-ns")
-
-        with patch("socket.gethostname", return_value="test-host"):
-            key = context.make_port_key(8080)
-            assert key == "dyn://test-ns/ports/test-host/8080"
-
-
 class TestPortMetadata:
     """Test port metadata functionality."""
 
@@ -80,77 +66,6 @@ def test_to_etcd_value_with_block_info(self):
         assert value["block_start"] == 8080
 
 
-class TestPortBinding:
-    """Test PortBinding context manager."""
-
-    def test_single_port_binding(self):
-        """Test binding a single port."""
-        # Find an available port
-        with socket.socket() as s:
-            s.bind(("", 0))
-            port = s.getsockname()[1]
-
-        binding = PortBinding(port)
-        assert binding.ports == [port]
-
-        # Test context manager
-        with binding:
-            # Port should be bound
-            assert len(binding.sockets) == 1
-            # Trying to bind again should fail
-            with pytest.raises(OSError):
-                with socket.socket() as s:
-                    s.bind(("", port))
-
-        # After context, sockets should be closed
-        # We should be able to bind again
-        with socket.socket() as s:
-            s.bind(("", port))
-
-    def test_multiple_port_binding(self):
-        """Test binding multiple ports."""
-        # Find available ports
-        ports = []
-        for _ in range(3):
-            with socket.socket() as s:
-                s.bind(("", 0))
-                ports.append(s.getsockname()[1])
-
-        binding = PortBinding(ports)
-        assert binding.ports == ports
-
-        with binding:
-            assert len(binding.sockets) == 3
-            # All ports should be bound
-            for port in ports:
-                with pytest.raises(OSError):
-                    with socket.socket() as s:
-                        s.bind(("", port))
-
-    def test_partial_binding_failure(self):
-        """Test cleanup when only some ports can be bound."""
-        # Bind a port that we'll keep occupied
-        blocker = socket.socket()
-        blocker.bind(("", 0))
-        blocked_port = blocker.getsockname()[1]
-
-        # Find another available port
-        with socket.socket() as s:
-            s.bind(("", 0))
-            good_port = s.getsockname()[1]
-
-        # Try to bind both ports (one will fail)
-        binding = PortBinding([good_port, blocked_port])
-
-        with pytest.raises(OSError):
-            binding.__enter__()
-
-        # All sockets should be cleaned up
-        assert all(sock._closed for sock in binding.sockets if hasattr(sock, "_closed"))
-
-        blocker.close()
-
-
 class TestHoldPorts:
     """Test hold_ports context manager."""
 
@@ -160,9 +75,11 @@ def test_hold_single_port(self):
             s.bind(("", 0))
             port = s.getsockname()[1]
 
-        with hold_ports(port) as binding:
-            assert isinstance(binding, PortBinding)
-            assert len(binding.sockets) == 1
+        with hold_ports(port):
+            assert not check_port_available(port)
+
+        # Port should be released after context exit
+        assert check_port_available(port)
 
     def test_hold_multiple_ports(self):
         """Test holding multiple ports."""
@@ -172,31 +89,13 @@ def test_hold_multiple_ports(self):
                 s.bind(("", 0))
                 ports.append(s.getsockname()[1])
 
-        with hold_ports(ports) as binding:
-            assert isinstance(binding, PortBinding)
-            assert len(binding.sockets) == 2
-
-
-class TestCheckPortAvailable:
-    """Test check_port_available function."""
-
-    def test_port_is_available(self):
-        """Test checking an available port."""
-        # Find an available port
-        with socket.socket() as s:
-            s.bind(("", 0))
-            port = s.getsockname()[1]
-
-        assert check_port_available(port) is True
-
-    def test_port_is_not_available(self):
-        """Test checking an occupied port."""
-        # Occupy a port
-        with socket.socket() as blocker:
-            blocker.bind(("", 0))
-            port = blocker.getsockname()[1]
+        with hold_ports(ports):
+            for port in ports:
+                assert not check_port_available(port)
 
-            assert check_port_available(port) is False
+        # All ports should be released after context exit
+        for port in ports:
+            assert check_port_available(port)
 
 
 class TestReservePortInEtcd:
@@ -206,20 +105,18 @@ class TestReservePortInEtcd:
     async def test_reserve_port_success(self):
         """Test successful port reservation in ETCD."""
         mock_client = AsyncMock()
-        mock_client.primary_lease_id = Mock(
-            return_value="test-lease-123"
-        )  # Regular Mock, not async
+        mock_client.primary_lease_id = Mock(return_value="test-lease-123")
 
         context = EtcdContext(client=mock_client, namespace="test-ns")
         metadata = PortMetadata(worker_id="test-worker", reason="test")
 
-        with patch("socket.gethostname", return_value="test-host"):
-            await reserve_port_in_etcd(context, 8080, metadata)
+        host_ip = get_host_ip()
+        await reserve_port_in_etcd(context, 8080, metadata)
 
         mock_client.kv_create.assert_called_once()
         call_args = mock_client.kv_create.call_args
 
-        assert call_args.kwargs["key"] == "dyn://test-ns/ports/test-host/8080"
+        assert call_args.kwargs["key"] == f"dyn://test-ns/ports/{host_ip}/8080"
         assert call_args.kwargs["lease_id"] == "test-lease-123"
 
         # Check the value is valid JSON
@@ -236,7 +133,7 @@ class TestAllocateAndReservePort:
     async def test_allocate_single_port_success(self):
         """Test successful single port allocation."""
         mock_client = AsyncMock()
-        mock_client.primary_lease_id.return_value = "test-lease"
+        mock_client.primary_lease_id = Mock(return_value="test-lease")
 
         context = EtcdContext(client=mock_client, namespace="test-ns")
         metadata = PortMetadata(worker_id="test-worker", reason="test")
@@ -264,7 +161,7 @@ class TestAllocateAndReservePortBlock:
     async def test_allocate_block_success(self):
         """Test successful port block allocation."""
         mock_client = AsyncMock()
-        mock_client.primary_lease_id.return_value = "test-lease"
+        mock_client.primary_lease_id = Mock(return_value="test-lease")
 
         context = EtcdContext(client=mock_client, namespace="test-ns")
         metadata = PortMetadata(worker_id="test-worker", reason="test")

components/backends/vllm/src/dynamo/vllm/ports.py

@@ -19,26 +19,6 @@
 DEFAULT_DYNAMO_PORT_MIN = 20000
 DEFAULT_DYNAMO_PORT_MAX = 30000
 
-__all__ = [
-    # Constants
-    "DEFAULT_DYNAMO_PORT_MIN",
-    "DEFAULT_DYNAMO_PORT_MAX",
-    # Classes
-    "DynamoPortRange",
-    "EtcdContext",
-    "PortMetadata",
-    "PortAllocationRequest",
-    "PortBinding",
-    # Functions
-    "check_port_available",
-    "reserve_port_in_etcd",
-    "allocate_and_reserve_port_block",
-    "allocate_and_reserve_port",
-    "get_host_ip",
-    # Context managers
-    "hold_ports",
-]
-
 
 @dataclass
 class DynamoPortRange:
@@ -67,8 +47,8 @@ class EtcdContext:
 
     def make_port_key(self, port: int) -> str:
         """Generate ETCD key for a port reservation"""
-        node_name = socket.gethostname()
-        return f"dyn://{self.namespace}/ports/{node_name}/{port}"
+        node_ip = get_host_ip()
+        return f"dyn://{self.namespace}/ports/{node_ip}/{port}"
 
 
 @dataclass
@@ -103,64 +83,38 @@ class PortAllocationRequest:
     max_attempts: int = 100
 
 
-class PortBinding:
-    """Holds socket bindings to ensure exclusive access to ports during reservation.
-
-    Can handle a single port or multiple ports.
-    """
-
-    def __init__(self, ports: int | list[int]):
-        """Initialize with a single port or list of ports."""
-        if isinstance(ports, int):
-            self.ports = [ports]
-        else:
-            self.ports = ports
-        self.sockets = []
-
-    def __enter__(self):
-        """Bind to all ports."""
-        try:
-            for port in self.ports:
-                sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-                sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
-                sock.bind(("", port))
-                self.sockets.append(sock)
-            return self
-        except OSError:
-            # Clean up any sockets we managed to bind
-            for sock in self.sockets:
-                sock.close()
-            raise
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        """Release all port bindings."""
-        for sock in self.sockets:
-            sock.close()
-
-
 @contextmanager
 def hold_ports(ports: int | list[int]):
     """Context manager to hold port binding(s).
 
+    Holds socket bindings to ensure exclusive access to ports during reservation.
+    Can handle a single port or multiple ports.
+
     Args:
         ports: Single port number or list of port numbers to hold
-
-    Yields:
-        PortBinding: The binding object holding the port(s)
     """
-    binding = PortBinding(ports)
+    if isinstance(ports, int):
+        ports = [ports]
+
+    sockets = []
     try:
-        binding.__enter__()
-        yield binding
+        for port in ports:
+            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+            sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+            sock.bind(("", port))
+            sockets.append(sock)
+
+        yield
+
     finally:
-        binding.__exit__(None, None, None)
+        for sock in sockets:
+            sock.close()
 
 
 def check_port_available(port: int) -> bool:
     """Check if a specific port is available for binding."""
     try:
         with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
-            sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
             sock.bind(("", port))
             return True
     except OSError:
@@ -211,7 +165,6 @@ async def allocate_and_reserve_port_block(request: PortAllocationRequest) -> lis
     available_start_ports = list(range(request.port_range.min, max_start_port + 1))
     random.shuffle(available_start_ports)
 
-    # Limit attempts to the number of valid starting ports or max_attempts
     actual_max_attempts = min(len(available_start_ports), request.max_attempts)
 
     for attempt in range(1, actual_max_attempts + 1):
@@ -227,7 +180,6 @@ async def allocate_and_reserve_port_block(request: PortAllocationRequest) -> lis
 
                 # We have exclusive access to these ports, now reserve them in ETCD
                 for i, port in enumerate(ports_to_reserve):
-                    # Create metadata for this specific port in the block
                     port_metadata = PortMetadata(
                         worker_id=f"{request.metadata.worker_id}-{i}"
                         if request.block_size > 1
@@ -255,12 +207,10 @@ async def allocate_and_reserve_port_block(request: PortAllocationRequest) -> lis
                 return ports_to_reserve
 
         except OSError as e:
-            # Failed to bind to ports, they're in use
             logger.debug(
                 f"Failed to bind to port block starting at {start_port} (attempt {attempt}): {e}"
             )
         except Exception as e:
-            # Failed to reserve in ETCD (but we had the ports bound)
             logger.debug(
                 f"Failed to reserve port block starting at {start_port} in ETCD (attempt {attempt}): {e}"
             )