chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,348 @@
|
||||
import heapq
|
||||
import logging
|
||||
import time
|
||||
from typing import Dict, List, Optional, Set, Tuple
|
||||
|
||||
from ray.serve._private.common import RequestProtocol
|
||||
from ray.serve._private.constants import (
|
||||
RAY_SERVE_DIRECT_INGRESS_MAX_GRPC_PORT,
|
||||
RAY_SERVE_DIRECT_INGRESS_MAX_HTTP_PORT,
|
||||
RAY_SERVE_DIRECT_INGRESS_MIN_GRPC_PORT,
|
||||
RAY_SERVE_DIRECT_INGRESS_MIN_HTTP_PORT,
|
||||
RAY_SERVE_PORT_QUARANTINE_S,
|
||||
SERVE_LOGGER_NAME,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(SERVE_LOGGER_NAME)
|
||||
|
||||
|
||||
class NoAvailablePortError(Exception):
|
||||
def __init__(self, protocol: str, node_id: str):
|
||||
message = f"No available ports on node {node_id} for {protocol} protocol."
|
||||
super().__init__(message)
|
||||
|
||||
|
||||
class PortAllocator:
|
||||
"""Manages a pool of ports for a specific protocol (e.g., HTTP or gRPC)."""
|
||||
|
||||
def __init__(self, min_port: int, max_port: int, protocol: str, node_id: str):
|
||||
self._protocol = protocol
|
||||
self._node_id = node_id
|
||||
# TODO(abrar): add a validation here to ensure min_port and max_port dont overlap with
|
||||
# ray params min_worker_port and max_worker_port.
|
||||
self._available_ports = list(range(min_port, max_port))
|
||||
heapq.heapify(self._available_ports)
|
||||
|
||||
self._allocated_ports: Dict[str, int] = {}
|
||||
self._blocked_ports: Set[int] = set()
|
||||
# port -> monotonic deadline at which it leaves quarantine.
|
||||
self._quarantined_ports: Dict[int, float] = {}
|
||||
# Min-heap of (deadline, port) mirroring _quarantined_ports so a no-op drain
|
||||
# is O(1) (peek) instead of O(quarantined). Lazy-deleted (P2).
|
||||
self._quarantine_heap: List[Tuple[float, int]] = []
|
||||
|
||||
def _drain_expired_quarantine(self) -> None:
|
||||
"""Return quarantined ports past their expiry to the available pool.
|
||||
|
||||
Drains via the min-heap keyed on expiry, so when nothing is due this is a
|
||||
single O(1) peek instead of an O(quarantined) dict scan. Lazy deletion: a
|
||||
popped heap entry is honored only if it still matches the live deadline in
|
||||
_quarantined_ports (entries left stale by update_port_if_missing's reclaim
|
||||
pop, or by a re-quarantine, are skipped).
|
||||
"""
|
||||
if not self._quarantine_heap:
|
||||
return
|
||||
now = time.monotonic()
|
||||
while self._quarantine_heap and self._quarantine_heap[0][0] <= now:
|
||||
deadline, port = heapq.heappop(self._quarantine_heap)
|
||||
if self._quarantined_ports.get(port) == deadline:
|
||||
heapq.heappush(self._available_ports, port)
|
||||
del self._quarantined_ports[port]
|
||||
logger.info(
|
||||
f"Released {self._protocol} port {port} from quarantine on "
|
||||
f"node {self._node_id}; returning it to the available pool."
|
||||
)
|
||||
|
||||
def has_pending_quarantine(self) -> bool:
|
||||
"""True if any port is still quarantined (drains expired entries first)."""
|
||||
self._drain_expired_quarantine()
|
||||
return bool(self._quarantined_ports)
|
||||
|
||||
def update_port_if_missing(self, replica_id: str, port: Optional[int]):
|
||||
"""Update port value for a replica."""
|
||||
if replica_id in self._allocated_ports:
|
||||
return
|
||||
# Reclaim trumps quarantine: the port is back in active use.
|
||||
if port is not None:
|
||||
self._quarantined_ports.pop(port, None)
|
||||
assert (
|
||||
port is not None
|
||||
), f"Port is None for {self._protocol} protocol on replica {replica_id} on node {self._node_id}"
|
||||
if self._protocol == RequestProtocol.HTTP:
|
||||
if not (
|
||||
RAY_SERVE_DIRECT_INGRESS_MIN_HTTP_PORT
|
||||
<= port
|
||||
<= RAY_SERVE_DIRECT_INGRESS_MAX_HTTP_PORT
|
||||
):
|
||||
logger.warning(f"HTTP port out of range: {port}")
|
||||
elif self._protocol == RequestProtocol.GRPC:
|
||||
if not (
|
||||
RAY_SERVE_DIRECT_INGRESS_MIN_GRPC_PORT
|
||||
<= port
|
||||
<= RAY_SERVE_DIRECT_INGRESS_MAX_GRPC_PORT
|
||||
):
|
||||
logger.warning(f"GRPC port out of range: {port}")
|
||||
self._allocated_ports[replica_id] = port
|
||||
|
||||
logger.info(
|
||||
f"Recovered {self._protocol} port {port} for replica {replica_id} on node {self._node_id}"
|
||||
)
|
||||
return port
|
||||
|
||||
def allocate(self, replica_id: str) -> int:
|
||||
if replica_id in self._allocated_ports:
|
||||
logger.warning(
|
||||
f"{self._protocol} port already allocated for replica {replica_id}"
|
||||
)
|
||||
return self._allocated_ports[replica_id]
|
||||
|
||||
self._drain_expired_quarantine()
|
||||
|
||||
# Recovered replicas live in _allocated_ports but their ports are
|
||||
# still in _available_ports, so guard against re-handing them out.
|
||||
# A recovered port released into quarantine is also still in the heap
|
||||
# (update_port_if_missing never pops it), so guard against handing out
|
||||
# a still-quarantined port too.
|
||||
allocated = set(self._allocated_ports.values())
|
||||
while self._available_ports:
|
||||
port = heapq.heappop(self._available_ports)
|
||||
if (
|
||||
port not in self._blocked_ports
|
||||
and port not in allocated
|
||||
and port not in self._quarantined_ports
|
||||
):
|
||||
self._allocated_ports[replica_id] = port
|
||||
logger.info(
|
||||
f"Allocated {self._protocol} port {port} to replica {replica_id} on node {self._node_id}"
|
||||
)
|
||||
return port
|
||||
|
||||
raise NoAvailablePortError(self._protocol, self._node_id)
|
||||
|
||||
def release(self, replica_id: str, port: int, block_port: bool = False):
|
||||
"""
|
||||
Releases a port for a replica.
|
||||
|
||||
Args:
|
||||
replica_id: The ID of the replica to release the port for.
|
||||
port: The port to release.
|
||||
block_port: Whether to block the port from being allocated again. Use this in
|
||||
situations where the port is being released due some other process is using it.
|
||||
"""
|
||||
if replica_id not in self._allocated_ports:
|
||||
raise ValueError(
|
||||
f"{self._protocol} port not allocated for replica {replica_id} on node {self._node_id}"
|
||||
)
|
||||
expected_port = self._allocated_ports[replica_id]
|
||||
assert expected_port == port, (
|
||||
f"{self._protocol} port mismatch for replica {replica_id} on node {self._node_id}: "
|
||||
f"expected {expected_port}, got {port}"
|
||||
)
|
||||
|
||||
del self._allocated_ports[replica_id]
|
||||
|
||||
if block_port:
|
||||
# Block is a stronger guarantee than quarantine; skip quarantine.
|
||||
self._blocked_ports.add(port)
|
||||
heapq.heappush(self._available_ports, port)
|
||||
logger.info(
|
||||
f"Released and blocked {self._protocol} port {port} for replica "
|
||||
f"{replica_id} on node {self._node_id}"
|
||||
)
|
||||
elif RAY_SERVE_PORT_QUARANTINE_S > 0:
|
||||
deadline = time.monotonic() + RAY_SERVE_PORT_QUARANTINE_S
|
||||
self._quarantined_ports[port] = deadline
|
||||
heapq.heappush(self._quarantine_heap, (deadline, port))
|
||||
logger.info(
|
||||
f"Quarantined {self._protocol} port {port} for "
|
||||
f"{RAY_SERVE_PORT_QUARANTINE_S}s after releasing for replica "
|
||||
f"{replica_id} on node {self._node_id}"
|
||||
)
|
||||
else:
|
||||
heapq.heappush(self._available_ports, port)
|
||||
logger.info(
|
||||
f"Released {self._protocol} port {port} for replica {replica_id} on node {self._node_id}"
|
||||
)
|
||||
|
||||
def prune(self, active_replica_ids: Set[str]):
|
||||
for replica_id in list(self._allocated_ports.keys()):
|
||||
if replica_id not in active_replica_ids:
|
||||
port = self._allocated_ports[replica_id]
|
||||
logger.info(
|
||||
f"Cleaning up {self._protocol} port {port} for stale replica {replica_id} on node {self._node_id}"
|
||||
)
|
||||
self.release(replica_id, port)
|
||||
|
||||
def get_port(self, replica_id: str) -> int:
|
||||
if replica_id not in self._allocated_ports:
|
||||
raise ValueError(
|
||||
f"{self._protocol} port not allocated for replica {replica_id} on node {self._node_id}"
|
||||
)
|
||||
return self._allocated_ports[replica_id]
|
||||
|
||||
def is_port_allocated(self, replica_id: str) -> bool:
|
||||
return replica_id in self._allocated_ports
|
||||
|
||||
|
||||
class NodePortManager:
|
||||
"""
|
||||
This class is responsible for managing replica-specific port allocations on a node,
|
||||
and is only used in direct ingress mode, where each Serve replica is exposed individually
|
||||
via a Kubernetes or GCP or AWS Ingress.
|
||||
|
||||
The primary goal of this class is to assign ports in a consistent and efficient manner,
|
||||
minimizing EndpointSlice fragmentation in Kubernetes. It uses a min-heap strategy to
|
||||
allocate ports incrementally, ensuring that all nodes tend to reuse the same port numbers.
|
||||
|
||||
Background:
|
||||
Kubernetes groups endpoints into EndpointSlices based on the set of ports exposed by each Pod.
|
||||
If Pods expose different port combinations (e.g., due to random port assignment), Kubernetes
|
||||
generates separate EndpointSlices per unique port list. This leads to unnecessary fragmentation
|
||||
and increased resource consumption.
|
||||
|
||||
By allocating ports deterministically, we ensure:
|
||||
- Consistent port usage across all nodes
|
||||
- Fewer unique port lists, reducing the number of EndpointSlices created
|
||||
- Improved performance and resource utilization
|
||||
|
||||
Although Kubernetes does not allow users to explicitly configure the ports included in
|
||||
EndpointSlices, maintaining a uniform port layout across nodes is still beneficial.
|
||||
|
||||
Port lifecycle:
|
||||
- Replicas are expected to release their ports when stopped
|
||||
- If a replica crashes without releasing its port, the controller loop will detect and
|
||||
reclaim leaked ports during reconciliation
|
||||
|
||||
Note:
|
||||
Although this strategy is designed with Kubernetes in mind, it is applied uniformly
|
||||
across all platforms for consistency.
|
||||
"""
|
||||
|
||||
_node_managers: Dict[str, "NodePortManager"] = {}
|
||||
|
||||
@classmethod
|
||||
def get_node_manager(cls, node_id: str) -> "NodePortManager":
|
||||
# this doesn't need to be behind a lock because it will already be called from same thread
|
||||
if node_id not in cls._node_managers:
|
||||
logger.info(f"Creating node manager for node {node_id}")
|
||||
cls._node_managers[node_id] = cls(node_id)
|
||||
return cls._node_managers[node_id]
|
||||
|
||||
@classmethod
|
||||
def prune(cls, node_id_to_alive_replica_ids: Dict[str, Set[str]]):
|
||||
# this doesn't need to be behind a lock because it will already be called from same thread
|
||||
for node_id in list(cls._node_managers):
|
||||
manager = cls._node_managers[node_id]
|
||||
alive = node_id_to_alive_replica_ids.get(node_id, set())
|
||||
# Skip the reclaim scan for a node whose set of alive replicas is unchanged
|
||||
# since the last prune. Comparing the whole alive set means a replica that
|
||||
# left, arrived, or moved nodes is still caught; skipping only defers reclaim
|
||||
# and never reuses a live port (allocate() guards that). The emptied-manager
|
||||
# teardown below still runs every tick.
|
||||
if alive == manager._last_pruned_alive:
|
||||
pass
|
||||
else:
|
||||
# Release ports of replicas no longer alive (quarantines them).
|
||||
manager._prune_replica_ports(alive)
|
||||
manager._last_pruned_alive = set(alive)
|
||||
# Keep the manager until its quarantine drains; dropping it early
|
||||
# would discard the deadline and let the port be reused immediately.
|
||||
if not alive and not manager.has_pending_quarantine():
|
||||
logger.info(f"Removing node manager for node {node_id}")
|
||||
del cls._node_managers[node_id]
|
||||
|
||||
@classmethod
|
||||
def update_ports(cls, ingress_replicas_info: List[Tuple[str, str, int, int]]):
|
||||
"""Update port values for ingress replicas."""
|
||||
for node_id, replica_id, http_port, grpc_port in ingress_replicas_info:
|
||||
if node_id is None:
|
||||
continue
|
||||
node_port_manager = cls.get_node_manager(node_id)
|
||||
if http_port is not None:
|
||||
node_port_manager._http_allocator.update_port_if_missing(
|
||||
replica_id,
|
||||
http_port,
|
||||
)
|
||||
if grpc_port is not None:
|
||||
node_port_manager._grpc_allocator.update_port_if_missing(
|
||||
replica_id,
|
||||
grpc_port,
|
||||
)
|
||||
|
||||
def __init__(self, node_id: str):
|
||||
self._node_id = node_id
|
||||
# Set of alive replicas this manager last pruned against; when unchanged, the
|
||||
# reclaim scan is skipped. None means never pruned, so the first prune runs.
|
||||
self._last_pruned_alive: Optional[Set[str]] = None
|
||||
|
||||
self._http_allocator = PortAllocator(
|
||||
RAY_SERVE_DIRECT_INGRESS_MIN_HTTP_PORT,
|
||||
RAY_SERVE_DIRECT_INGRESS_MAX_HTTP_PORT,
|
||||
protocol=RequestProtocol.HTTP,
|
||||
node_id=node_id,
|
||||
)
|
||||
self._grpc_allocator = PortAllocator(
|
||||
RAY_SERVE_DIRECT_INGRESS_MIN_GRPC_PORT,
|
||||
RAY_SERVE_DIRECT_INGRESS_MAX_GRPC_PORT,
|
||||
protocol=RequestProtocol.GRPC,
|
||||
node_id=node_id,
|
||||
)
|
||||
|
||||
def _prune_replica_ports(self, active_replica_ids: Set[str]):
|
||||
self._http_allocator.prune(active_replica_ids)
|
||||
self._grpc_allocator.prune(active_replica_ids)
|
||||
|
||||
def allocate_port(self, replica_id: str, protocol: RequestProtocol) -> int:
|
||||
if protocol == RequestProtocol.HTTP:
|
||||
return self._http_allocator.allocate(replica_id)
|
||||
elif protocol == RequestProtocol.GRPC:
|
||||
return self._grpc_allocator.allocate(replica_id)
|
||||
else:
|
||||
raise ValueError(f"Unsupported protocol: {protocol}")
|
||||
|
||||
def release_port(
|
||||
self,
|
||||
replica_id: str,
|
||||
port: int,
|
||||
protocol: RequestProtocol,
|
||||
block_port: bool = False,
|
||||
):
|
||||
if protocol == RequestProtocol.HTTP:
|
||||
self._http_allocator.release(replica_id, port, block_port)
|
||||
elif protocol == RequestProtocol.GRPC:
|
||||
self._grpc_allocator.release(replica_id, port, block_port)
|
||||
else:
|
||||
raise ValueError(f"Unsupported protocol: {protocol}")
|
||||
|
||||
def get_port(self, replica_id: str, protocol: RequestProtocol) -> int:
|
||||
if protocol == RequestProtocol.HTTP:
|
||||
return self._http_allocator.get_port(replica_id)
|
||||
elif protocol == RequestProtocol.GRPC:
|
||||
return self._grpc_allocator.get_port(replica_id)
|
||||
else:
|
||||
raise ValueError(f"Unsupported protocol: {protocol}")
|
||||
|
||||
def is_port_allocated(self, replica_id: str, protocol: RequestProtocol) -> bool:
|
||||
if protocol == RequestProtocol.HTTP:
|
||||
return self._http_allocator.is_port_allocated(replica_id)
|
||||
elif protocol == RequestProtocol.GRPC:
|
||||
return self._grpc_allocator.is_port_allocated(replica_id)
|
||||
else:
|
||||
raise ValueError(f"Unsupported protocol: {protocol}")
|
||||
|
||||
def has_pending_quarantine(self) -> bool:
|
||||
"""True if either allocator still holds a quarantined port."""
|
||||
# Evaluate both (no short-circuit) so each allocator drains expired ports.
|
||||
http_pending = self._http_allocator.has_pending_quarantine()
|
||||
grpc_pending = self._grpc_allocator.has_pending_quarantine()
|
||||
return http_pending or grpc_pending
|
||||
Reference in New Issue
Block a user