wehub-resource-sync
1200 lines
40 KiB
Python
1200 lines
40 KiB
Python
import asyncio
|
|
from copy import deepcopy
|
|
from typing import Dict, List, Optional, Tuple
|
|
from unittest import mock
|
|
|
|
import pytest
|
|
|
|
from ray.serve._private.common import (
|
|
DeploymentID,
|
|
DeploymentStatus,
|
|
DeploymentStatusTrigger,
|
|
ReplicaID,
|
|
ReplicaState,
|
|
RequestProtocol,
|
|
RunningReplicaInfo,
|
|
)
|
|
from ray.serve._private.controller import ServeController
|
|
from ray.serve._private.node_port_manager import NodePortManager
|
|
from ray.serve.config import HTTPOptions, gRPCOptions
|
|
from ray.serve.schema import (
|
|
DeploymentDetails,
|
|
DeploymentSchema,
|
|
ReplicaDetails,
|
|
Target,
|
|
TargetGroup,
|
|
)
|
|
|
|
|
|
# Simple test KV store implementation
|
|
class FakeKVStore:
|
|
def __init__(self):
|
|
self.data = {}
|
|
|
|
def get(self, key):
|
|
return self.data.get(key)
|
|
|
|
def put(self, key, value):
|
|
self.data[key] = value
|
|
|
|
def delete(self, key):
|
|
if key in self.data:
|
|
del self.data[key]
|
|
|
|
|
|
# Simplified LongPollHost for tests
|
|
class FakeLongPollHost:
|
|
def __init__(self):
|
|
self.notified_changes = {}
|
|
|
|
def notify_changed(self, changes):
|
|
self.notified_changes.update(changes)
|
|
|
|
|
|
# Application State Manager for dependency injection
|
|
class FakeApplicationStateManager:
|
|
def __init__(
|
|
self,
|
|
app_statuses,
|
|
route_prefixes,
|
|
ingress_deployments,
|
|
ingress_request_router_deployments=None,
|
|
):
|
|
self.app_statuses = app_statuses
|
|
self.route_prefixes = route_prefixes
|
|
self.ingress_deployments = ingress_deployments
|
|
self.ingress_request_router_deployments = (
|
|
ingress_request_router_deployments or {}
|
|
)
|
|
|
|
def list_app_statuses(self):
|
|
return self.app_statuses
|
|
|
|
def get_route_prefix(self, app_name):
|
|
return self.route_prefixes.get(app_name, f"/{app_name}")
|
|
|
|
def get_ingress_deployment_name(self, app_name):
|
|
return self.ingress_deployments.get(app_name, f"{app_name}_ingress")
|
|
|
|
def get_ingress_request_router_deployment_name(self, app_name):
|
|
return self.ingress_request_router_deployments.get(app_name)
|
|
|
|
|
|
class FakeProxyState:
|
|
def __init__(self, node_id, node_ip, name):
|
|
self.node_id = node_id
|
|
self.node_ip = node_ip
|
|
self.name = name
|
|
|
|
|
|
class FakeProxyStateManager:
|
|
def __init__(self):
|
|
self.proxy_details = {}
|
|
self._http_options = HTTPOptions()
|
|
self._grpc_options = gRPCOptions(
|
|
grpc_servicer_functions=["f1"],
|
|
)
|
|
self._fallback_proxy_targets = {}
|
|
self._started_fallback_proxy = False
|
|
|
|
def add_proxy_details(self, node_id, node_ip, name):
|
|
|
|
self.proxy_details[node_id] = FakeProxyState(
|
|
node_id=node_id,
|
|
node_ip=node_ip,
|
|
name=name,
|
|
)
|
|
|
|
def add_fallback_proxy_target(
|
|
self,
|
|
node_ip,
|
|
port,
|
|
node_instance_id,
|
|
actor_name,
|
|
protocol,
|
|
):
|
|
self._fallback_proxy_targets[protocol] = Target(
|
|
ip=node_ip,
|
|
port=port,
|
|
instance_id=node_instance_id,
|
|
name=actor_name,
|
|
)
|
|
|
|
def get_proxy_details(self):
|
|
return self.proxy_details
|
|
|
|
def get_targets(self, protocol: RequestProtocol):
|
|
if protocol == RequestProtocol.HTTP:
|
|
port = self._http_options.port
|
|
else:
|
|
port = self._grpc_options.port
|
|
return [
|
|
Target(
|
|
ip=proxy_details.node_ip,
|
|
port=port,
|
|
instance_id="",
|
|
name=proxy_details.name,
|
|
)
|
|
for node_id, proxy_details in self.proxy_details.items()
|
|
]
|
|
|
|
def get_grpc_config(self):
|
|
return self._grpc_options
|
|
|
|
def get_fallback_proxy_targets(self):
|
|
return deepcopy(self._fallback_proxy_targets)
|
|
|
|
def started_fallback_proxy_at_least_once(self):
|
|
return self._started_fallback_proxy
|
|
|
|
|
|
# Deployment State Manager for dependency injection
|
|
class FakeDeploymentStateManager:
|
|
def __init__(
|
|
self,
|
|
running_replica_infos: Dict[DeploymentID, List[RunningReplicaInfo]],
|
|
):
|
|
self.running_replica_infos = running_replica_infos
|
|
|
|
def get_running_replica_infos(self):
|
|
return self.running_replica_infos
|
|
|
|
def get_deployment_ids(self):
|
|
return list(self.running_replica_infos.keys())
|
|
|
|
def get_node_id_to_alive_replica_ids(self):
|
|
node_id_to_alive_replica_ids = {}
|
|
for replica_infos in self.running_replica_infos.values():
|
|
for replica_info in replica_infos:
|
|
if replica_info.node_id is None:
|
|
continue
|
|
node_id_to_alive_replica_ids.setdefault(
|
|
replica_info.node_id, set()
|
|
).add(replica_info.replica_id.unique_id)
|
|
return node_id_to_alive_replica_ids
|
|
|
|
def get_replica_details(self, replica_info: RunningReplicaInfo) -> ReplicaDetails:
|
|
return ReplicaDetails(
|
|
replica_id=replica_info.replica_id.unique_id,
|
|
node_id=replica_info.node_id,
|
|
node_ip=replica_info.node_ip,
|
|
node_instance_id="",
|
|
start_time_s=0,
|
|
state=ReplicaState.RUNNING,
|
|
actor_name=replica_info.replica_id.unique_id,
|
|
)
|
|
|
|
def get_deployment_details(self, id: DeploymentID) -> Optional[DeploymentDetails]:
|
|
if id not in self.running_replica_infos:
|
|
return None
|
|
replica_details = [
|
|
self.get_replica_details(replica_info)
|
|
for replica_info in self.running_replica_infos[id]
|
|
]
|
|
return DeploymentDetails(
|
|
name=id.name,
|
|
status=DeploymentStatus.HEALTHY,
|
|
status_trigger=DeploymentStatusTrigger.UNSPECIFIED,
|
|
message="",
|
|
deployment_config=DeploymentSchema(name=id.name),
|
|
target_num_replicas=1,
|
|
required_resources={},
|
|
replicas=replica_details,
|
|
)
|
|
|
|
def get_ingress_replicas_info(self) -> List[Tuple[str, str, int, int]]:
|
|
return []
|
|
|
|
|
|
class FakeTestingDeploymentStateManager:
|
|
def __init__(self):
|
|
self.dumped_deployment_id = None
|
|
self.stopped_deployment_id = None
|
|
self.replica_states = object()
|
|
|
|
def _dump_replica_states_for_testing(self, deployment_id: DeploymentID):
|
|
self.dumped_deployment_id = deployment_id
|
|
return self.replica_states
|
|
|
|
def _stop_one_running_replica_for_testing(self, deployment_id: DeploymentID):
|
|
self.stopped_deployment_id = deployment_id
|
|
|
|
|
|
# Test Controller that overrides methods and dependencies
|
|
class FakeDirectIngressController(ServeController):
|
|
def __init__(
|
|
self,
|
|
kv_store,
|
|
long_poll_host,
|
|
application_state_manager,
|
|
deployment_state_manager,
|
|
proxy_state_manager,
|
|
):
|
|
# Skip parent __init__ since we'll set dependencies directly
|
|
self.kv_store = kv_store
|
|
self.long_poll_host = long_poll_host
|
|
self.application_state_manager = application_state_manager
|
|
self.deployment_state_manager = deployment_state_manager
|
|
self.proxy_state_manager = proxy_state_manager
|
|
self._direct_ingress_enabled = True
|
|
self._last_ingress_port_tuples = set()
|
|
self._ha_proxy_enabled = False
|
|
self._controller_node_id = "head_node_id"
|
|
|
|
self._shutting_down = False
|
|
self.done_recovering_event = asyncio.Event()
|
|
self.done_recovering_event.set()
|
|
|
|
self.node_update_duration_gauge_s = mock.Mock()
|
|
self._last_broadcasted_fallback_targets = {}
|
|
|
|
def _update_proxy_nodes(self):
|
|
pass
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def setup_env(monkeypatch):
|
|
monkeypatch.setenv("RAY_SERVE_ENABLE_DIRECT_INGRESS", "1")
|
|
monkeypatch.setenv("RAY_SERVE_DIRECT_INGRESS_MIN_HTTP_PORT", "30000")
|
|
monkeypatch.setenv("RAY_SERVE_DIRECT_INGRESS_MAX_HTTP_PORT", "30100")
|
|
monkeypatch.setenv("RAY_SERVE_DIRECT_INGRESS_MIN_GRPC_PORT", "40000")
|
|
monkeypatch.setenv("RAY_SERVE_DIRECT_INGRESS_MAX_GRPC_PORT", "40100")
|
|
|
|
yield
|
|
|
|
|
|
# Fixture to clear NodePortManager between tests
|
|
@pytest.fixture(autouse=True)
|
|
def reset_node_port_manager():
|
|
# Save original managers
|
|
original_managers = NodePortManager._node_managers.copy()
|
|
|
|
# Clear all managers
|
|
NodePortManager._node_managers = {}
|
|
|
|
yield
|
|
|
|
# Restore original managers
|
|
NodePortManager._node_managers = original_managers
|
|
|
|
|
|
# Fixture for a minimal controller setup
|
|
@pytest.fixture
|
|
def direct_ingress_controller():
|
|
# Setup test dependencies
|
|
kv_store = FakeKVStore()
|
|
long_poll_host = FakeLongPollHost()
|
|
app_state_manager = FakeApplicationStateManager({}, {}, {})
|
|
deployment_state_manager = FakeDeploymentStateManager({})
|
|
proxy_state_manager = FakeProxyStateManager()
|
|
# Create controller with test dependencies
|
|
controller = FakeDirectIngressController(
|
|
kv_store=kv_store,
|
|
long_poll_host=long_poll_host,
|
|
application_state_manager=app_state_manager,
|
|
deployment_state_manager=deployment_state_manager,
|
|
proxy_state_manager=proxy_state_manager,
|
|
)
|
|
|
|
yield controller
|
|
|
|
|
|
def test_direct_ingress_is_disabled(
|
|
direct_ingress_controller: FakeDirectIngressController,
|
|
):
|
|
"""Test that get_target_groups returns empty list when direct ingress is disabled."""
|
|
direct_ingress_controller._direct_ingress_enabled = False
|
|
target_groups = direct_ingress_controller.get_target_groups()
|
|
assert target_groups == []
|
|
|
|
# proxy has nodes
|
|
direct_ingress_controller.proxy_state_manager.add_proxy_details(
|
|
"node1", "10.0.0.1", "proxy1"
|
|
)
|
|
direct_ingress_controller.proxy_state_manager.add_proxy_details(
|
|
"node2", "10.0.0.2", "proxy2"
|
|
)
|
|
target_groups = direct_ingress_controller.get_target_groups()
|
|
expected_target_groups = [
|
|
TargetGroup(
|
|
protocol=RequestProtocol.HTTP,
|
|
route_prefix="/",
|
|
targets=[
|
|
Target(ip="10.0.0.1", port=8000, instance_id="", name="proxy1"),
|
|
Target(ip="10.0.0.2", port=8000, instance_id="", name="proxy2"),
|
|
],
|
|
),
|
|
TargetGroup(
|
|
protocol=RequestProtocol.GRPC,
|
|
route_prefix="/",
|
|
targets=[
|
|
Target(ip="10.0.0.1", port=9000, instance_id="", name="proxy1"),
|
|
Target(ip="10.0.0.2", port=9000, instance_id="", name="proxy2"),
|
|
],
|
|
),
|
|
]
|
|
assert target_groups == expected_target_groups
|
|
|
|
|
|
# Test with empty applications
|
|
def test_get_target_groups_empty_when_no_apps(
|
|
direct_ingress_controller: FakeDirectIngressController,
|
|
):
|
|
"""Test that get_target_groups returns empty list when no apps are running."""
|
|
target_groups = direct_ingress_controller.get_target_groups()
|
|
assert target_groups == []
|
|
|
|
# proxy has nodes
|
|
direct_ingress_controller.proxy_state_manager.add_proxy_details(
|
|
"node1", "10.0.0.1", "proxy1"
|
|
)
|
|
direct_ingress_controller.proxy_state_manager.add_proxy_details(
|
|
"node2", "10.0.0.2", "proxy2"
|
|
)
|
|
target_groups = direct_ingress_controller.get_target_groups()
|
|
expected_target_groups = [
|
|
TargetGroup(
|
|
protocol=RequestProtocol.HTTP,
|
|
route_prefix="/",
|
|
targets=[
|
|
Target(ip="10.0.0.1", port=8000, instance_id="", name="proxy1"),
|
|
Target(ip="10.0.0.2", port=8000, instance_id="", name="proxy2"),
|
|
],
|
|
),
|
|
TargetGroup(
|
|
protocol=RequestProtocol.GRPC,
|
|
route_prefix="/",
|
|
targets=[
|
|
Target(ip="10.0.0.1", port=9000, instance_id="", name="proxy1"),
|
|
Target(ip="10.0.0.2", port=9000, instance_id="", name="proxy2"),
|
|
],
|
|
),
|
|
]
|
|
assert target_groups == expected_target_groups
|
|
|
|
|
|
# Test with running applications
|
|
def test_get_target_groups_with_running_apps(
|
|
direct_ingress_controller: FakeDirectIngressController,
|
|
):
|
|
"""Test that get_target_groups properly returns target groups for running apps."""
|
|
# Setup test data
|
|
app_statuses = {
|
|
"app1": {},
|
|
"app2": {},
|
|
}
|
|
route_prefixes = {
|
|
"app1": "/app1",
|
|
"app2": "/app2",
|
|
}
|
|
ingress_deployments = {
|
|
"app1": "app1_ingress",
|
|
"app2": "app2_ingress",
|
|
}
|
|
|
|
deployment_id1 = DeploymentID(name="app1_ingress", app_name="app1")
|
|
deployment_id2 = DeploymentID(name="app2_ingress", app_name="app2")
|
|
|
|
# Create replica info
|
|
replica_id1 = ReplicaID(unique_id="replica1", deployment_id=deployment_id1)
|
|
replica_id2 = ReplicaID(unique_id="replica2", deployment_id=deployment_id2)
|
|
|
|
replica_info1 = RunningReplicaInfo(
|
|
replica_id=replica_id1,
|
|
node_id="node1",
|
|
node_ip="10.0.0.1",
|
|
availability_zone="az1",
|
|
actor_name="replica1",
|
|
max_ongoing_requests=100,
|
|
)
|
|
replica_info2 = RunningReplicaInfo(
|
|
replica_id=replica_id2,
|
|
node_id="node2",
|
|
node_ip="10.0.0.2",
|
|
availability_zone="az2",
|
|
actor_name="replica2",
|
|
max_ongoing_requests=100,
|
|
)
|
|
|
|
running_replica_infos = {
|
|
deployment_id1: [replica_info1],
|
|
deployment_id2: [replica_info2],
|
|
}
|
|
|
|
# Setup test application state manager
|
|
direct_ingress_controller.application_state_manager = FakeApplicationStateManager(
|
|
app_statuses=app_statuses,
|
|
route_prefixes=route_prefixes,
|
|
ingress_deployments=ingress_deployments,
|
|
)
|
|
|
|
# Setup test deployment state manager
|
|
direct_ingress_controller.deployment_state_manager = FakeDeploymentStateManager(
|
|
running_replica_infos=running_replica_infos,
|
|
)
|
|
|
|
# setup proxy state manager
|
|
direct_ingress_controller.proxy_state_manager.add_proxy_details(
|
|
"node1", "10.0.0.1", "proxy1"
|
|
)
|
|
direct_ingress_controller.proxy_state_manager.add_proxy_details(
|
|
"node2", "10.0.0.2", "proxy2"
|
|
)
|
|
|
|
# Allocate ports for replicas using controller's methods
|
|
http_port1 = direct_ingress_controller.allocate_replica_port(
|
|
"node1", replica_id1.unique_id, RequestProtocol.HTTP
|
|
)
|
|
grpc_port1 = direct_ingress_controller.allocate_replica_port(
|
|
"node1", replica_id1.unique_id, RequestProtocol.GRPC
|
|
)
|
|
http_port2 = direct_ingress_controller.allocate_replica_port(
|
|
"node2", replica_id2.unique_id, RequestProtocol.HTTP
|
|
)
|
|
grpc_port2 = direct_ingress_controller.allocate_replica_port(
|
|
"node2", replica_id2.unique_id, RequestProtocol.GRPC
|
|
)
|
|
|
|
# Call get_target_groups
|
|
target_groups = direct_ingress_controller.get_target_groups()
|
|
|
|
# Create expected target groups for direct comparison
|
|
expected_target_groups = [
|
|
TargetGroup(
|
|
protocol=RequestProtocol.HTTP,
|
|
route_prefix="/app1",
|
|
app_name="app1",
|
|
targets=[
|
|
Target(ip="10.0.0.1", port=http_port1, instance_id="", name="replica1"),
|
|
],
|
|
),
|
|
TargetGroup(
|
|
protocol=RequestProtocol.GRPC,
|
|
route_prefix="/app1",
|
|
app_name="app1",
|
|
targets=[
|
|
Target(ip="10.0.0.1", port=grpc_port1, instance_id="", name="replica1"),
|
|
],
|
|
),
|
|
TargetGroup(
|
|
protocol=RequestProtocol.HTTP,
|
|
route_prefix="/app2",
|
|
app_name="app2",
|
|
targets=[
|
|
Target(ip="10.0.0.2", port=http_port2, instance_id="", name="replica2"),
|
|
],
|
|
),
|
|
TargetGroup(
|
|
protocol=RequestProtocol.GRPC,
|
|
route_prefix="/app2",
|
|
app_name="app2",
|
|
targets=[
|
|
Target(ip="10.0.0.2", port=grpc_port2, instance_id="", name="replica2"),
|
|
],
|
|
),
|
|
]
|
|
|
|
# Sort both lists to ensure consistent comparison
|
|
target_groups.sort(key=lambda g: (g.protocol, g.route_prefix))
|
|
expected_target_groups.sort(key=lambda g: (g.protocol, g.route_prefix))
|
|
|
|
# Direct comparison
|
|
assert target_groups == expected_target_groups
|
|
|
|
# now release some ports
|
|
direct_ingress_controller.release_replica_port(
|
|
"node1", replica_id1.unique_id, http_port1, RequestProtocol.HTTP
|
|
)
|
|
direct_ingress_controller.release_replica_port(
|
|
"node2", replica_id2.unique_id, grpc_port2, RequestProtocol.GRPC
|
|
)
|
|
|
|
# verify the ports are released
|
|
assert not direct_ingress_controller._is_port_allocated(
|
|
direct_ingress_controller.deployment_state_manager.get_replica_details(
|
|
replica_info1
|
|
),
|
|
RequestProtocol.HTTP,
|
|
)
|
|
assert not direct_ingress_controller._is_port_allocated(
|
|
direct_ingress_controller.deployment_state_manager.get_replica_details(
|
|
replica_info2
|
|
),
|
|
RequestProtocol.GRPC,
|
|
)
|
|
|
|
# get the target groups again
|
|
target_groups = direct_ingress_controller.get_target_groups()
|
|
expected_target_groups = [
|
|
TargetGroup(
|
|
protocol=RequestProtocol.GRPC,
|
|
route_prefix="/app1",
|
|
app_name="app1",
|
|
targets=[
|
|
Target(ip="10.0.0.1", port=grpc_port1, instance_id="", name="replica1"),
|
|
],
|
|
),
|
|
TargetGroup(
|
|
protocol=RequestProtocol.HTTP,
|
|
route_prefix="/app2",
|
|
app_name="app2",
|
|
targets=[
|
|
Target(ip="10.0.0.2", port=http_port2, instance_id="", name="replica2"),
|
|
],
|
|
),
|
|
]
|
|
|
|
# Sort both lists again for the second comparison
|
|
target_groups.sort(key=lambda g: (g.protocol, g.route_prefix))
|
|
expected_target_groups.sort(key=lambda g: (g.protocol, g.route_prefix))
|
|
|
|
assert target_groups == expected_target_groups
|
|
|
|
|
|
def test_get_target_groups_with_port_not_allocated(
|
|
direct_ingress_controller: FakeDirectIngressController,
|
|
):
|
|
"""Test that get_target_groups correctly filters out targets with unallocated ports."""
|
|
# Setup test data
|
|
app_statuses = {"app1": {}}
|
|
route_prefixes = {"app1": "/app1"}
|
|
ingress_deployments = {"app1": "app1_ingress"}
|
|
|
|
# Create replica info for two replicas
|
|
deployment_id1 = DeploymentID(name="app1_ingress", app_name="app1")
|
|
replica_id1 = ReplicaID(unique_id="replica1", deployment_id=deployment_id1)
|
|
replica_id2 = ReplicaID(unique_id="replica2", deployment_id=deployment_id1)
|
|
|
|
replica_info1 = RunningReplicaInfo(
|
|
replica_id=replica_id1,
|
|
node_id="node1",
|
|
node_ip="10.0.0.1",
|
|
availability_zone="az1",
|
|
actor_name="replica1",
|
|
max_ongoing_requests=100,
|
|
)
|
|
replica_info2 = RunningReplicaInfo(
|
|
replica_id=replica_id2,
|
|
node_id="node2",
|
|
node_ip="10.0.0.2",
|
|
availability_zone="az2",
|
|
actor_name="replica2",
|
|
max_ongoing_requests=100,
|
|
)
|
|
|
|
running_replica_infos = {
|
|
deployment_id1: [replica_info1, replica_info2],
|
|
}
|
|
|
|
# Setup test application state manager
|
|
direct_ingress_controller.application_state_manager = FakeApplicationStateManager(
|
|
app_statuses=app_statuses,
|
|
route_prefixes=route_prefixes,
|
|
ingress_deployments=ingress_deployments,
|
|
)
|
|
|
|
# Setup test deployment state manager
|
|
direct_ingress_controller.deployment_state_manager = FakeDeploymentStateManager(
|
|
running_replica_infos=running_replica_infos,
|
|
)
|
|
|
|
# Only allocate ports for the first replica, leave the second one without ports
|
|
http_port1 = direct_ingress_controller.allocate_replica_port(
|
|
"node1", replica_id1.unique_id, RequestProtocol.HTTP
|
|
)
|
|
grpc_port1 = direct_ingress_controller.allocate_replica_port(
|
|
"node1", replica_id1.unique_id, RequestProtocol.GRPC
|
|
)
|
|
|
|
# Call get_target_groups
|
|
target_groups = direct_ingress_controller.get_target_groups()
|
|
|
|
# Create expected target groups
|
|
expected_target_groups = [
|
|
TargetGroup(
|
|
protocol=RequestProtocol.HTTP,
|
|
route_prefix="/app1",
|
|
app_name="app1",
|
|
targets=[
|
|
Target(ip="10.0.0.1", port=http_port1, instance_id="", name="replica1"),
|
|
],
|
|
),
|
|
TargetGroup(
|
|
protocol=RequestProtocol.GRPC,
|
|
route_prefix="/app1",
|
|
app_name="app1",
|
|
targets=[
|
|
Target(ip="10.0.0.1", port=grpc_port1, instance_id="", name="replica1"),
|
|
],
|
|
),
|
|
]
|
|
|
|
# Sort both lists to ensure consistent comparison
|
|
target_groups.sort(key=lambda g: (g.protocol, g.route_prefix))
|
|
expected_target_groups.sort(key=lambda g: (g.protocol, g.route_prefix))
|
|
|
|
# Direct comparison
|
|
assert target_groups == expected_target_groups
|
|
|
|
|
|
def test_get_target_groups_only_includes_ingress_deployments(
|
|
direct_ingress_controller: FakeDirectIngressController,
|
|
):
|
|
"""Test that get_target_groups only includes ingress deployments and not regular deployments."""
|
|
# Setup test data
|
|
app_statuses = {"app1": {}}
|
|
route_prefixes = {"app1": "/app1"}
|
|
ingress_deployments = {"app1": "app1_ingress"}
|
|
|
|
# Create replica info for ingress deployment
|
|
ingress_deployment_id = DeploymentID(name="app1_ingress", app_name="app1")
|
|
ingress_replica_id = ReplicaID(
|
|
unique_id="ingress_replica", deployment_id=ingress_deployment_id
|
|
)
|
|
ingress_replica_info = RunningReplicaInfo(
|
|
replica_id=ingress_replica_id,
|
|
node_id="node1",
|
|
node_ip="10.0.0.1",
|
|
availability_zone="az1",
|
|
actor_name="ingress_replica",
|
|
max_ongoing_requests=100,
|
|
)
|
|
|
|
# Create replica info for regular non-ingress deployment
|
|
regular_deployment_id = DeploymentID(name="app1_model", app_name="app1")
|
|
regular_replica_id = ReplicaID(
|
|
unique_id="regular_replica", deployment_id=regular_deployment_id
|
|
)
|
|
regular_replica_info = RunningReplicaInfo(
|
|
replica_id=regular_replica_id,
|
|
node_id="node2",
|
|
node_ip="10.0.0.2",
|
|
availability_zone="az2",
|
|
actor_name="regular_replica",
|
|
max_ongoing_requests=100,
|
|
)
|
|
|
|
# Set up running replica infos for both deployments
|
|
running_replica_infos = {
|
|
ingress_deployment_id: [ingress_replica_info],
|
|
regular_deployment_id: [regular_replica_info],
|
|
}
|
|
|
|
# Setup test application state manager
|
|
direct_ingress_controller.application_state_manager = FakeApplicationStateManager(
|
|
app_statuses=app_statuses,
|
|
route_prefixes=route_prefixes,
|
|
ingress_deployments=ingress_deployments,
|
|
)
|
|
|
|
# Setup test deployment state manager
|
|
direct_ingress_controller.deployment_state_manager = FakeDeploymentStateManager(
|
|
running_replica_infos=running_replica_infos,
|
|
)
|
|
|
|
# Allocate ports for both replicas
|
|
ingress_http_port = direct_ingress_controller.allocate_replica_port(
|
|
"node1", ingress_replica_id.unique_id, RequestProtocol.HTTP
|
|
)
|
|
ingress_grpc_port = direct_ingress_controller.allocate_replica_port(
|
|
"node1", ingress_replica_id.unique_id, RequestProtocol.GRPC
|
|
)
|
|
_ = direct_ingress_controller.allocate_replica_port(
|
|
"node2", regular_replica_id.unique_id, RequestProtocol.HTTP
|
|
)
|
|
_ = direct_ingress_controller.allocate_replica_port(
|
|
"node2", regular_replica_id.unique_id, RequestProtocol.GRPC
|
|
)
|
|
|
|
# Call get_target_groups
|
|
target_groups = direct_ingress_controller.get_target_groups()
|
|
|
|
# Create expected target groups - only including the ingress deployment
|
|
expected_target_groups = [
|
|
TargetGroup(
|
|
protocol=RequestProtocol.HTTP,
|
|
route_prefix="/app1",
|
|
app_name="app1",
|
|
targets=[
|
|
Target(
|
|
ip="10.0.0.1",
|
|
port=ingress_http_port,
|
|
instance_id="",
|
|
name="ingress_replica",
|
|
),
|
|
],
|
|
),
|
|
TargetGroup(
|
|
protocol=RequestProtocol.GRPC,
|
|
route_prefix="/app1",
|
|
app_name="app1",
|
|
targets=[
|
|
Target(
|
|
ip="10.0.0.1",
|
|
port=ingress_grpc_port,
|
|
instance_id="",
|
|
name="ingress_replica",
|
|
),
|
|
],
|
|
),
|
|
]
|
|
|
|
# Sort both lists to ensure consistent comparison
|
|
target_groups.sort(key=lambda g: (g.protocol, g.route_prefix))
|
|
expected_target_groups.sort(key=lambda g: (g.protocol, g.route_prefix))
|
|
|
|
# Direct comparison - regular deployment should not be included
|
|
assert target_groups == expected_target_groups
|
|
|
|
# Verify all ports are still allocated even though not all are included in target groups
|
|
assert direct_ingress_controller._is_port_allocated(
|
|
direct_ingress_controller.deployment_state_manager.get_replica_details(
|
|
ingress_replica_info
|
|
),
|
|
RequestProtocol.HTTP,
|
|
)
|
|
assert direct_ingress_controller._is_port_allocated(
|
|
direct_ingress_controller.deployment_state_manager.get_replica_details(
|
|
ingress_replica_info
|
|
),
|
|
RequestProtocol.GRPC,
|
|
)
|
|
assert direct_ingress_controller._is_port_allocated(
|
|
direct_ingress_controller.deployment_state_manager.get_replica_details(
|
|
regular_replica_info
|
|
),
|
|
RequestProtocol.HTTP,
|
|
)
|
|
assert direct_ingress_controller._is_port_allocated(
|
|
direct_ingress_controller.deployment_state_manager.get_replica_details(
|
|
regular_replica_info
|
|
),
|
|
RequestProtocol.GRPC,
|
|
)
|
|
|
|
|
|
def test_get_target_groups_populates_ingress_request_router_targets(
|
|
direct_ingress_controller: FakeDirectIngressController,
|
|
):
|
|
app_name = "app1"
|
|
route_prefix = "/app1"
|
|
ingress_deployment_id = DeploymentID(name="app1_ingress", app_name=app_name)
|
|
router_deployment_id = DeploymentID(name="app1_router", app_name=app_name)
|
|
ingress_replica_id = ReplicaID(
|
|
unique_id="ingress_replica", deployment_id=ingress_deployment_id
|
|
)
|
|
router_replica_id = ReplicaID(
|
|
unique_id="router_replica", deployment_id=router_deployment_id
|
|
)
|
|
ingress_replica_info = RunningReplicaInfo(
|
|
replica_id=ingress_replica_id,
|
|
node_id="node1",
|
|
node_ip="10.0.0.1",
|
|
availability_zone="az1",
|
|
actor_name="ingress_replica",
|
|
max_ongoing_requests=100,
|
|
)
|
|
router_replica_info = RunningReplicaInfo(
|
|
replica_id=router_replica_id,
|
|
node_id="node2",
|
|
node_ip="10.0.0.2",
|
|
availability_zone="az2",
|
|
actor_name="router_replica",
|
|
max_ongoing_requests=100,
|
|
)
|
|
|
|
direct_ingress_controller.application_state_manager = FakeApplicationStateManager(
|
|
app_statuses={app_name: {}},
|
|
route_prefixes={app_name: route_prefix},
|
|
ingress_deployments={app_name: ingress_deployment_id.name},
|
|
ingress_request_router_deployments={app_name: router_deployment_id.name},
|
|
)
|
|
direct_ingress_controller.deployment_state_manager = FakeDeploymentStateManager(
|
|
running_replica_infos={
|
|
ingress_deployment_id: [ingress_replica_info],
|
|
router_deployment_id: [router_replica_info],
|
|
},
|
|
)
|
|
|
|
ingress_http_port = direct_ingress_controller.allocate_replica_port(
|
|
"node1", ingress_replica_id.unique_id, RequestProtocol.HTTP
|
|
)
|
|
router_http_port = direct_ingress_controller.allocate_replica_port(
|
|
"node2", router_replica_id.unique_id, RequestProtocol.HTTP
|
|
)
|
|
|
|
assert direct_ingress_controller.get_target_groups(app_name=app_name) == [
|
|
TargetGroup(
|
|
protocol=RequestProtocol.HTTP,
|
|
route_prefix=route_prefix,
|
|
app_name=app_name,
|
|
targets=[
|
|
Target(
|
|
ip="10.0.0.1",
|
|
port=ingress_http_port,
|
|
instance_id="",
|
|
name="ingress_replica",
|
|
)
|
|
],
|
|
ingress_request_router_targets=[
|
|
Target(
|
|
ip="10.0.0.2",
|
|
port=router_http_port,
|
|
instance_id="",
|
|
name="router_replica",
|
|
)
|
|
],
|
|
)
|
|
]
|
|
|
|
|
|
def test_get_target_groups_app_with_no_running_replicas(
|
|
direct_ingress_controller: FakeDirectIngressController,
|
|
):
|
|
"""Test that get_target_groups correctly handles apps with no running replicas."""
|
|
# Setup test data for two apps - one with replicas, one without
|
|
app_statuses = {
|
|
"app1": {}, # App with replicas
|
|
"app2": {}, # App with no replicas
|
|
}
|
|
route_prefixes = {
|
|
"app1": "/app1",
|
|
"app2": "/app2",
|
|
}
|
|
ingress_deployments = {
|
|
"app1": "app1_ingress",
|
|
"app2": "app2_ingress",
|
|
}
|
|
|
|
# Create replica info for app1's ingress deployment
|
|
deployment_id1 = DeploymentID(name="app1_ingress", app_name="app1")
|
|
replica_id1 = ReplicaID(unique_id="replica1", deployment_id=deployment_id1)
|
|
replica_info1 = RunningReplicaInfo(
|
|
replica_id=replica_id1,
|
|
node_id="node1",
|
|
node_ip="10.0.0.1",
|
|
availability_zone="az1",
|
|
actor_name="replica1",
|
|
max_ongoing_requests=100,
|
|
)
|
|
|
|
# Note: No replicas for app2_ingress
|
|
|
|
# Running replica infos only for app1
|
|
running_replica_infos = {
|
|
deployment_id1: [replica_info1],
|
|
# No entry for app2_ingress
|
|
}
|
|
|
|
# Setup test application state manager
|
|
direct_ingress_controller.application_state_manager = FakeApplicationStateManager(
|
|
app_statuses=app_statuses,
|
|
route_prefixes=route_prefixes,
|
|
ingress_deployments=ingress_deployments,
|
|
)
|
|
|
|
# Setup test deployment state manager
|
|
direct_ingress_controller.deployment_state_manager = FakeDeploymentStateManager(
|
|
running_replica_infos=running_replica_infos,
|
|
)
|
|
|
|
# setup proxy state manager
|
|
direct_ingress_controller.proxy_state_manager.add_proxy_details(
|
|
"node1", "10.0.0.1", "proxy1"
|
|
)
|
|
direct_ingress_controller.proxy_state_manager.add_proxy_details(
|
|
"node2", "10.0.0.2", "proxy2"
|
|
)
|
|
|
|
# Allocate ports for the only existing replica
|
|
http_port = direct_ingress_controller.allocate_replica_port(
|
|
"node1", replica_id1.unique_id, RequestProtocol.HTTP
|
|
)
|
|
grpc_port = direct_ingress_controller.allocate_replica_port(
|
|
"node1", replica_id1.unique_id, RequestProtocol.GRPC
|
|
)
|
|
|
|
# Call get_target_groups
|
|
target_groups = direct_ingress_controller.get_target_groups()
|
|
|
|
# Create expected target groups - only including app1, nothing for app2
|
|
expected_target_groups = [
|
|
TargetGroup(
|
|
protocol=RequestProtocol.HTTP,
|
|
route_prefix="/app1",
|
|
app_name="app1",
|
|
targets=[
|
|
Target(ip="10.0.0.1", port=http_port, instance_id="", name="replica1"),
|
|
],
|
|
),
|
|
TargetGroup(
|
|
protocol=RequestProtocol.GRPC,
|
|
route_prefix="/app1",
|
|
app_name="app1",
|
|
targets=[
|
|
Target(ip="10.0.0.1", port=grpc_port, instance_id="", name="replica1"),
|
|
],
|
|
),
|
|
TargetGroup(
|
|
protocol=RequestProtocol.HTTP,
|
|
route_prefix="/app2",
|
|
app_name="app2",
|
|
targets=[
|
|
Target(ip="10.0.0.1", port=8000, instance_id="", name="proxy1"),
|
|
Target(ip="10.0.0.2", port=8000, instance_id="", name="proxy2"),
|
|
],
|
|
),
|
|
TargetGroup(
|
|
protocol=RequestProtocol.GRPC,
|
|
route_prefix="/app2",
|
|
app_name="app2",
|
|
targets=[
|
|
Target(ip="10.0.0.1", port=9000, instance_id="", name="proxy1"),
|
|
Target(ip="10.0.0.2", port=9000, instance_id="", name="proxy2"),
|
|
],
|
|
),
|
|
]
|
|
|
|
# Sort both lists to ensure consistent comparison
|
|
target_groups.sort(key=lambda g: (g.protocol, g.route_prefix))
|
|
expected_target_groups.sort(key=lambda g: (g.protocol, g.route_prefix))
|
|
|
|
# Direct comparison - app2 should have no target groups since it has no replicas
|
|
assert target_groups == expected_target_groups
|
|
|
|
|
|
def test_control_loop_pruning(
|
|
direct_ingress_controller: FakeDirectIngressController,
|
|
):
|
|
"""Test that the controller loop properly prunes stale node port managers."""
|
|
# Setup replica info for testing
|
|
deployment_id = DeploymentID(name="app1_ingress", app_name="app1")
|
|
replica_id1 = ReplicaID(unique_id="replica1", deployment_id=deployment_id)
|
|
replica_id2 = ReplicaID(unique_id="replica2", deployment_id=deployment_id)
|
|
replica_id3 = ReplicaID(unique_id="replica3", deployment_id=deployment_id)
|
|
|
|
replica_info1 = RunningReplicaInfo(
|
|
replica_id=replica_id1,
|
|
node_id="node1",
|
|
node_ip="10.0.0.1",
|
|
availability_zone="az1",
|
|
actor_name="replica1",
|
|
max_ongoing_requests=100,
|
|
)
|
|
replica_info2 = RunningReplicaInfo(
|
|
replica_id=replica_id2,
|
|
node_id="node1",
|
|
node_ip="10.0.0.1",
|
|
availability_zone="az1",
|
|
actor_name="replica2",
|
|
max_ongoing_requests=100,
|
|
)
|
|
replica_info3 = RunningReplicaInfo(
|
|
replica_id=replica_id3,
|
|
node_id="node2",
|
|
node_ip="10.0.0.2",
|
|
availability_zone="az2",
|
|
actor_name="replica3",
|
|
max_ongoing_requests=100,
|
|
)
|
|
|
|
running_replica_infos = {
|
|
deployment_id: [replica_info1, replica_info2, replica_info3],
|
|
}
|
|
|
|
# Set up controller with the replica info
|
|
direct_ingress_controller.deployment_state_manager = FakeDeploymentStateManager(
|
|
running_replica_infos=running_replica_infos,
|
|
)
|
|
|
|
# Allocate ports for testing
|
|
direct_ingress_controller.allocate_replica_port(
|
|
"node1", replica_id1.unique_id, RequestProtocol.HTTP
|
|
)
|
|
direct_ingress_controller.allocate_replica_port(
|
|
"node1", replica_id2.unique_id, RequestProtocol.HTTP
|
|
)
|
|
direct_ingress_controller.allocate_replica_port(
|
|
"node1", replica_id3.unique_id, RequestProtocol.HTTP
|
|
) # This should be pruned
|
|
direct_ingress_controller.allocate_replica_port(
|
|
"node2", replica_id3.unique_id, RequestProtocol.HTTP
|
|
)
|
|
direct_ingress_controller.allocate_replica_port(
|
|
"node3", "replica4", RequestProtocol.HTTP
|
|
) # Node should be pruned
|
|
|
|
# Verify ports are initially allocated
|
|
assert direct_ingress_controller._is_port_allocated(
|
|
direct_ingress_controller.deployment_state_manager.get_replica_details(
|
|
replica_info1
|
|
),
|
|
RequestProtocol.HTTP,
|
|
)
|
|
assert direct_ingress_controller._is_port_allocated(
|
|
direct_ingress_controller.deployment_state_manager.get_replica_details(
|
|
replica_info2
|
|
),
|
|
RequestProtocol.HTTP,
|
|
)
|
|
assert direct_ingress_controller._is_port_allocated(
|
|
direct_ingress_controller.deployment_state_manager.get_replica_details(
|
|
replica_info3
|
|
),
|
|
RequestProtocol.HTTP,
|
|
)
|
|
|
|
# We need to use NodePortManager directly for this check since we don't have a ReplicaInfo for stale_replica
|
|
node1_manager = NodePortManager.get_node_manager("node1")
|
|
node3_manager = NodePortManager.get_node_manager("node3")
|
|
assert node1_manager.is_port_allocated(replica_id3.unique_id, RequestProtocol.HTTP)
|
|
assert node3_manager.is_port_allocated("replica4", RequestProtocol.HTTP)
|
|
|
|
# Call the control loop step - this should trigger port pruning
|
|
direct_ingress_controller._maybe_update_ingress_ports()
|
|
|
|
# Verify the active replicas still have their ports
|
|
assert direct_ingress_controller._is_port_allocated(
|
|
direct_ingress_controller.deployment_state_manager.get_replica_details(
|
|
replica_info1
|
|
),
|
|
RequestProtocol.HTTP,
|
|
)
|
|
assert direct_ingress_controller._is_port_allocated(
|
|
direct_ingress_controller.deployment_state_manager.get_replica_details(
|
|
replica_info2
|
|
),
|
|
RequestProtocol.HTTP,
|
|
)
|
|
assert direct_ingress_controller._is_port_allocated(
|
|
direct_ingress_controller.deployment_state_manager.get_replica_details(
|
|
replica_info3
|
|
),
|
|
RequestProtocol.HTTP,
|
|
)
|
|
|
|
# Verify stale ports were pruned
|
|
assert not node1_manager.is_port_allocated(
|
|
replica_id3.unique_id, RequestProtocol.HTTP
|
|
)
|
|
assert "node3" not in NodePortManager._node_managers # Entire node should be pruned
|
|
|
|
|
|
def test_control_loop_pruning_does_not_require_private_deployment_states(
|
|
direct_ingress_controller: FakeDirectIngressController,
|
|
):
|
|
deployment_id = DeploymentID(name="app1_ingress", app_name="app1")
|
|
replica_id = ReplicaID(unique_id="replica1", deployment_id=deployment_id)
|
|
replica_info = RunningReplicaInfo(
|
|
replica_id=replica_id,
|
|
node_id="node1",
|
|
node_ip="10.0.0.1",
|
|
availability_zone="az1",
|
|
actor_name="replica1",
|
|
max_ongoing_requests=100,
|
|
)
|
|
|
|
direct_ingress_controller.deployment_state_manager = FakeDeploymentStateManager(
|
|
running_replica_infos={deployment_id: [replica_info]},
|
|
)
|
|
|
|
direct_ingress_controller.allocate_replica_port(
|
|
"node1", replica_id.unique_id, RequestProtocol.HTTP
|
|
)
|
|
direct_ingress_controller._maybe_update_ingress_ports()
|
|
|
|
assert direct_ingress_controller._is_port_allocated(
|
|
direct_ingress_controller.deployment_state_manager.get_replica_details(
|
|
replica_info
|
|
),
|
|
RequestProtocol.HTTP,
|
|
)
|
|
|
|
|
|
def test_dump_replica_states_for_testing_delegates_to_manager(
|
|
direct_ingress_controller: FakeDirectIngressController,
|
|
):
|
|
deployment_id = DeploymentID(name="app1_ingress", app_name="app1")
|
|
deployment_state_manager = FakeTestingDeploymentStateManager()
|
|
direct_ingress_controller.deployment_state_manager = deployment_state_manager
|
|
|
|
assert (
|
|
direct_ingress_controller._dump_replica_states_for_testing(deployment_id)
|
|
is deployment_state_manager.replica_states
|
|
)
|
|
assert deployment_state_manager.dumped_deployment_id == deployment_id
|
|
|
|
|
|
def test_stop_one_running_replica_for_testing_delegates_to_manager(
|
|
direct_ingress_controller: FakeDirectIngressController,
|
|
):
|
|
deployment_id = DeploymentID(name="app1_ingress", app_name="app1")
|
|
deployment_state_manager = FakeTestingDeploymentStateManager()
|
|
direct_ingress_controller.deployment_state_manager = deployment_state_manager
|
|
|
|
direct_ingress_controller._stop_one_running_replica_for_testing(deployment_id)
|
|
|
|
assert deployment_state_manager.stopped_deployment_id == deployment_id
|
|
|
|
|
|
def test_recon_port_gate_feeds_only_new_ingress_tuples(
|
|
direct_ingress_controller: FakeDirectIngressController,
|
|
):
|
|
"""C3: update_ports is fed only the set-diff of
|
|
ingress tuples new since the last sync -- the full set on the first tick (empty
|
|
cache) and an empty list on an unchanged second tick."""
|
|
tuples = [("node1", "r1", 30000, 40000), ("node2", "r2", 30001, 40001)]
|
|
direct_ingress_controller.deployment_state_manager.get_ingress_replicas_info = (
|
|
lambda: list(tuples)
|
|
)
|
|
|
|
with mock.patch.object(NodePortManager, "update_ports") as mock_update:
|
|
# First tick: cache empty -> full emit.
|
|
direct_ingress_controller._maybe_update_ingress_ports()
|
|
assert set(mock_update.call_args[0][0]) == set(tuples)
|
|
|
|
# Second tick: ingress set unchanged -> empty diff.
|
|
mock_update.reset_mock()
|
|
direct_ingress_controller._maybe_update_ingress_ports()
|
|
assert mock_update.call_args[0][0] == []
|
|
|
|
|
|
def test_recon_port_gate_prune_skips_unchanged_node(
|
|
direct_ingress_controller: FakeDirectIngressController,
|
|
):
|
|
"""C4: prune skips a node's O(replicas) reclaim
|
|
scan when that node's alive-replica set is unchanged since the last prune."""
|
|
deployment_id = DeploymentID(name="app1_ingress", app_name="app1")
|
|
replica_id = ReplicaID(unique_id="replica1", deployment_id=deployment_id)
|
|
replica_info = RunningReplicaInfo(
|
|
replica_id=replica_id,
|
|
node_id="node1",
|
|
node_ip="10.0.0.1",
|
|
availability_zone="az1",
|
|
actor_name="replica1",
|
|
max_ongoing_requests=100,
|
|
)
|
|
direct_ingress_controller.deployment_state_manager = FakeDeploymentStateManager(
|
|
running_replica_infos={deployment_id: [replica_info]},
|
|
)
|
|
direct_ingress_controller.allocate_replica_port(
|
|
"node1", replica_id.unique_id, RequestProtocol.HTTP
|
|
)
|
|
node1_manager = NodePortManager.get_node_manager("node1")
|
|
|
|
with mock.patch.object(
|
|
node1_manager,
|
|
"_prune_replica_ports",
|
|
wraps=node1_manager._prune_replica_ports,
|
|
) as spy:
|
|
# First tick: fingerprint unset -> reclaim scan runs.
|
|
direct_ingress_controller._maybe_update_ingress_ports()
|
|
assert spy.call_count == 1
|
|
# Second tick: alive set unchanged -> reclaim scan skipped.
|
|
direct_ingress_controller._maybe_update_ingress_ports()
|
|
assert spy.call_count == 1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main()
|