228 lines
7.6 KiB
Python
228 lines
7.6 KiB
Python
import pytest
|
|
|
|
from ray.serve._private.common import DeploymentID, EndpointInfo
|
|
from ray.serve._private.proxy_router import (
|
|
NO_REPLICAS_MESSAGE,
|
|
NO_ROUTES_MESSAGE,
|
|
ProxyRouter,
|
|
)
|
|
from ray.serve._private.test_utils import MockDeploymentHandle
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_router():
|
|
def mock_get_handle(endpoint, info):
|
|
return MockDeploymentHandle(endpoint.name, endpoint.app_name)
|
|
|
|
yield ProxyRouter(mock_get_handle)
|
|
|
|
|
|
def test_no_match(mock_router: ProxyRouter):
|
|
router = mock_router
|
|
router.update_routes(
|
|
{
|
|
DeploymentID(name="endpoint", app_name="default"): EndpointInfo(
|
|
route="/hello"
|
|
),
|
|
DeploymentID(name="endpoint2", app_name="app2"): EndpointInfo(
|
|
route="/hello2"
|
|
),
|
|
},
|
|
)
|
|
assert router.match_route("/nonexistent") is None
|
|
assert router.get_handle_for_endpoint("/nonexistent") is None
|
|
|
|
|
|
def test_default_route(mock_router: ProxyRouter):
|
|
router = mock_router
|
|
router.update_routes(
|
|
{
|
|
DeploymentID(name="endpoint", app_name="default"): EndpointInfo(
|
|
route="/endpoint"
|
|
),
|
|
DeploymentID(name="endpoint2", app_name="app2"): EndpointInfo(
|
|
route="/endpoint2"
|
|
),
|
|
},
|
|
)
|
|
|
|
# Route based matching
|
|
assert router.match_route("/nonexistent") is None
|
|
|
|
route, handle, app_is_cross_language = router.match_route("/endpoint")
|
|
assert route == "/endpoint"
|
|
assert handle == ("endpoint", "default")
|
|
assert not app_is_cross_language
|
|
|
|
# Endpoint based matching
|
|
assert router.get_handle_for_endpoint("/nonexistent") is None
|
|
|
|
route, handle, app_is_cross_language = router.get_handle_for_endpoint("default")
|
|
assert route == "/endpoint"
|
|
assert handle == ("endpoint", "default")
|
|
assert not app_is_cross_language
|
|
|
|
|
|
def test_trailing_slash(mock_router: ProxyRouter):
|
|
router = mock_router
|
|
router.update_routes(
|
|
{
|
|
DeploymentID(name="endpoint", app_name="default"): EndpointInfo(
|
|
route="/test"
|
|
)
|
|
},
|
|
)
|
|
|
|
route, handle, _ = router.match_route("/test/")
|
|
assert route == "/test" and handle == ("endpoint", "default")
|
|
|
|
router.update_routes(
|
|
{
|
|
DeploymentID(name="endpoint", app_name="default"): EndpointInfo(
|
|
route="/test/"
|
|
)
|
|
},
|
|
)
|
|
|
|
assert router.match_route("/test") is None
|
|
|
|
|
|
def test_prefix_match(mock_router):
|
|
router = mock_router
|
|
router.update_routes(
|
|
{
|
|
DeploymentID(name="endpoint1", app_name="default"): EndpointInfo(
|
|
route="/test/test2"
|
|
),
|
|
DeploymentID(name="endpoint2", app_name="default"): EndpointInfo(
|
|
route="/test"
|
|
),
|
|
DeploymentID(name="endpoint3", app_name="default"): EndpointInfo(route="/"),
|
|
},
|
|
)
|
|
|
|
route, handle, _ = router.match_route("/test/test2/subpath")
|
|
assert route == "/test/test2" and handle == ("endpoint1", "default")
|
|
route, handle, _ = router.match_route("/test/test2/")
|
|
assert route == "/test/test2" and handle == ("endpoint1", "default")
|
|
route, handle, _ = router.match_route("/test/test2")
|
|
assert route == "/test/test2" and handle == ("endpoint1", "default")
|
|
|
|
route, handle, _ = router.match_route("/test/subpath")
|
|
assert route == "/test" and handle == ("endpoint2", "default")
|
|
route, handle, _ = router.match_route("/test/")
|
|
assert route == "/test" and handle == ("endpoint2", "default")
|
|
route, handle, _ = router.match_route("/test")
|
|
assert route == "/test" and handle == ("endpoint2", "default")
|
|
|
|
route, handle, _ = router.match_route("/test2")
|
|
assert route == "/" and handle == ("endpoint3", "default")
|
|
route, handle, _ = router.match_route("/")
|
|
assert route == "/" and handle == ("endpoint3", "default")
|
|
|
|
|
|
def test_update_routes(mock_router):
|
|
router = mock_router
|
|
router.update_routes(
|
|
{DeploymentID("endpoint", "app1"): EndpointInfo(route="/endpoint")},
|
|
)
|
|
|
|
route, handle, app_is_cross_language = router.match_route("/endpoint")
|
|
assert route == "/endpoint"
|
|
assert handle == ("endpoint", "app1")
|
|
assert not app_is_cross_language
|
|
|
|
route, handle, app_is_cross_language = router.get_handle_for_endpoint("app1")
|
|
assert route == "/endpoint"
|
|
assert handle == ("endpoint", "app1")
|
|
assert not app_is_cross_language
|
|
|
|
router.update_routes(
|
|
{
|
|
DeploymentID(name="endpoint2", app_name="app2"): EndpointInfo(
|
|
route="/endpoint2",
|
|
app_is_cross_language=True,
|
|
),
|
|
DeploymentID(name="endpoint3", app_name="app3"): EndpointInfo(
|
|
route="/endpoint3",
|
|
app_is_cross_language=True,
|
|
),
|
|
},
|
|
)
|
|
|
|
assert router.match_route("/endpoint") is None
|
|
assert router.match_route("app1") is None
|
|
|
|
route, handle, app_is_cross_language = router.match_route("/endpoint2")
|
|
assert route == "/endpoint2"
|
|
assert handle == ("endpoint2", "app2")
|
|
assert app_is_cross_language
|
|
|
|
route, handle, app_is_cross_language = router.get_handle_for_endpoint("app2")
|
|
assert route == "/endpoint2"
|
|
assert handle == ("endpoint2", "app2")
|
|
assert app_is_cross_language
|
|
|
|
|
|
class TestReadyForTraffic:
|
|
@pytest.mark.parametrize("is_head", [False, True])
|
|
def test_route_table_not_populated(self, mock_router, is_head: bool):
|
|
"""Proxy router should NOT be ready for traffic if:
|
|
- it has not received route table from controller
|
|
"""
|
|
|
|
ready_for_traffic, msg = mock_router.ready_for_traffic(is_head=is_head)
|
|
assert not ready_for_traffic
|
|
assert msg == NO_ROUTES_MESSAGE
|
|
|
|
def test_head_route_table_populated_no_replicas(self, mock_router):
|
|
"""Proxy router should be ready for traffic if:
|
|
- it has received route table from controller
|
|
- it hasn't received any replicas yet
|
|
- it lives on head node
|
|
"""
|
|
|
|
d_id = DeploymentID(name="A", app_name="B")
|
|
mock_router.update_routes({d_id: EndpointInfo(route="/")})
|
|
mock_router.handles[d_id].set_running_replicas_populated(False)
|
|
|
|
ready_for_traffic, msg = mock_router.ready_for_traffic(is_head=True)
|
|
assert ready_for_traffic
|
|
assert not msg
|
|
|
|
def test_worker_route_table_populated_no_replicas(self, mock_router):
|
|
"""Proxy router should NOT be ready for traffic if:
|
|
- it has received route table from controller
|
|
- it hasn't received any replicas yet
|
|
- it lives on a worker node
|
|
"""
|
|
|
|
d_id = DeploymentID(name="A", app_name="B")
|
|
mock_router.update_routes({d_id: EndpointInfo(route="/")})
|
|
mock_router.handles[d_id].set_running_replicas_populated(False)
|
|
|
|
ready_for_traffic, msg = mock_router.ready_for_traffic(is_head=False)
|
|
assert not ready_for_traffic
|
|
assert msg == NO_REPLICAS_MESSAGE
|
|
|
|
@pytest.mark.parametrize("is_head", [False, True])
|
|
def test_route_table_populated_with_replicas(self, mock_router, is_head: bool):
|
|
"""Proxy router should be ready for traffic if:
|
|
- it has received route table from controller
|
|
- it has received replicas from controller
|
|
"""
|
|
|
|
d_id = DeploymentID(name="A", app_name="B")
|
|
mock_router.update_routes({d_id: EndpointInfo(route="/")})
|
|
mock_router.handles[d_id].set_running_replicas_populated(True)
|
|
|
|
ready_for_traffic, msg = mock_router.ready_for_traffic(is_head=is_head)
|
|
assert ready_for_traffic
|
|
assert not msg
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", "-s", __file__]))
|