chore: import upstream snapshot with attribution

python/ray/serve/tests/test_callback.py

@@ -0,0 +1,272 @@
+import importlib
+import logging
+import os
+import sys
+from typing import Any, Dict, Generator
+
+import httpx
+import pytest
+import starlette
+from starlette.middleware import Middleware
+
+import ray
+from ray import serve
+from ray._common.test_utils import wait_for_condition
+from ray.exceptions import RayActorError
+from ray.serve._private.test_utils import get_application_url
+from ray.serve._private.utils import call_function_from_import_path
+from ray.serve.config import HTTPOptions, gRPCOptions
+from ray.serve.context import _get_global_client
+from ray.serve.schema import (
+    LoggingConfig,
+    ProxyStatus,
+    ServeInstanceDetails,
+)
+
+
+# ==== Callbacks used in this test ====
+class ASGIMiddleware:
+    def __init__(self, app):
+        self.app = app
+
+    async def __call__(self, scope, receive, send):
+        scope.get("headers").append((b"custom_header_key", "custom_header_value"))
+        await self.app(scope, receive, send)
+
+
+def add_middleware():
+    return [Middleware(ASGIMiddleware)]
+
+
+class MyServeFormatter(logging.Formatter):
+    def format(self, record: logging.LogRecord):
+        log_msg = super().format(record)
+        return "MyCustom message: hello " + log_msg
+
+
+def add_logger():
+    ray_logger = logging.getLogger("ray.serve")
+    handler = logging.StreamHandler()
+    handler.setFormatter(MyServeFormatter())
+    ray_logger.addHandler(handler)
+
+
+def raise_error_callback():
+    raise RuntimeError("this is from raise_error_callback")
+
+
+def return_bad_objects():
+    return [1, 2, 3]
+
+
+ADD_MIDDLEWARE_IMPORT_PATH = "ray.serve.tests.test_callback.add_middleware"
+ADD_LOGGER_IMPORT_PATH = "ray.serve.tests.test_callback.add_logger"
+RAISE_ERROR_IMPORT_PATH = "ray.serve.tests.test_callback.raise_error_callback"
+RETURN_BAD_OBJECTS_IMPORT_PATH = "ray.serve.tests.test_callback.return_bad_objects"
+NOT_CALLABLE_OBJECT = 1
+# ==== end ====
+
+
+@pytest.fixture()
+def ray_instance(
+    request: pytest.FixtureRequest,
+) -> Generator[Dict[str, Any], None, None]:
+    """Starts and stops a Ray instance for this test.
+
+    Args:
+        request: request.param should contain a dictionary of env vars and
+            their values. The Ray instance will be started with these env vars.
+
+    Yields:
+        Dict[str, Any]: The dict returned by ``ray.init`` for the started cluster.
+    """
+
+    original_env_vars = os.environ.copy()
+
+    try:
+        requested_env_vars = request.param
+    except AttributeError:
+        requested_env_vars = {}
+
+    os.environ.update(requested_env_vars)
+    importlib.reload(ray.serve._private.constants)
+    importlib.reload(ray.serve._private.controller)
+    importlib.reload(ray.serve._private.proxy)
+
+    yield ray.init()
+
+    serve.shutdown()
+    # wait_for_processes=True blocks until the raylet/GCS/etc. subprocesses
+    # have fully exited. Without it, this teardown races the next test's
+    # ray.init() and the new raylet can fail to register the driver.
+    ray.shutdown(wait_for_processes=True)
+
+    os.environ.clear()
+    os.environ.update(original_env_vars)
+
+
+def test_call_function_from_import_path():
+    """Basic test for call_function_from_import_path"""
+
+    # basic
+    assert [1, 2, 3] == call_function_from_import_path(RETURN_BAD_OBJECTS_IMPORT_PATH)
+
+    # rasie exception when callback function raise exception
+    with pytest.raises(RuntimeError, match="this is from raise_error_callback"):
+        call_function_from_import_path(RAISE_ERROR_IMPORT_PATH)
+
+    # raise exception when providing invalid import path
+    with pytest.raises(ValueError, match="cannot be imported"):
+        call_function_from_import_path("not_exist")
+
+    # raise exception when providing non callable object
+    with pytest.raises(TypeError, match="is not callable"):
+        call_function_from_import_path(
+            "ray.serve.tests.test_callback.NOT_CALLABLE_OBJECT"
+        )
+
+
+@pytest.mark.parametrize(
+    "ray_instance",
+    [
+        {
+            "RAY_SERVE_CONTROLLER_CALLBACK_IMPORT_PATH": ADD_LOGGER_IMPORT_PATH,
+            "RAY_SERVE_HTTP_PROXY_CALLBACK_IMPORT_PATH": ADD_MIDDLEWARE_IMPORT_PATH,
+        },
+    ],
+    indirect=True,
+)
+def test_callback(ray_instance, capsys):
+    """Test callback function works in http proxy and controller"""
+    serve.start(
+        http_options=HTTPOptions(
+            host="0.0.0.0",
+            request_timeout_s=500,
+        ),
+    )
+
+    @serve.deployment
+    class Model:
+        def __call__(self, request: starlette.requests.Request):
+            headers = request.scope.get("headers")
+            for k, v in headers:
+                if k == b"custom_header_key":
+                    return v
+            return "Not found custom headers"
+
+    serve.run(Model.bind())
+    url = get_application_url()
+    resp = httpx.get(url)
+
+    assert resp.text == "custom_header_value"
+    captured = capsys.readouterr()
+    assert "MyCustom message: hello" in captured.err
+
+
+@pytest.mark.parametrize(
+    "ray_instance",
+    [
+        {
+            "RAY_SERVE_CONTROLLER_CALLBACK_IMPORT_PATH": "not_exist",
+            "RAY_SERVE_HTTP_PROXY_CALLBACK_IMPORT_PATH": RAISE_ERROR_IMPORT_PATH,
+        },
+    ],
+    indirect=True,
+)
+def test_callback_fail(ray_instance):
+    """Test actor call call_function_from_import_path rasing exception.
+
+    Actor will fail to be started and further call will raise RayActorError.
+    """
+
+    actor_def = ray.serve._private.proxy.ProxyActor
+    handle = actor_def.remote(
+        http_options=HTTPOptions(host="http_proxy", root_path="/", port=123),
+        grpc_options=gRPCOptions(),
+        node_ip_address="127.0.0.1",
+        node_id="123",
+        logging_config=LoggingConfig(),
+        long_poll_client="fake_client",
+    )
+    with pytest.raises(RayActorError, match="this is from raise_error_callback"):
+        ray.get(handle.ready.remote())
+
+    serve_controller = ray.serve._private.controller.ServeController
+    actor_def = ray.actor._make_actor(serve_controller, {})
+    handle = actor_def.remote(
+        http_options=HTTPOptions(),
+        grpc_options=gRPCOptions(),
+        global_logging_config=LoggingConfig(),
+    )
+    with pytest.raises(RayActorError, match="cannot be imported"):
+        ray.get(handle.check_alive.remote())
+
+
+@pytest.mark.parametrize(
+    "ray_instance",
+    [
+        {
+            "RAY_SERVE_HTTP_PROXY_CALLBACK_IMPORT_PATH": RETURN_BAD_OBJECTS_IMPORT_PATH,
+        },
+    ],
+    indirect=True,
+)
+def test_http_proxy_return_aribitary_objects(ray_instance):
+    """Test invalid callback path in http proxy"""
+
+    actor_def = ray.serve._private.proxy.ProxyActor
+    handle = actor_def.remote(
+        http_options=HTTPOptions(host="http_proxy", root_path="/", port=123),
+        grpc_options=gRPCOptions(),
+        node_ip_address="127.0.0.1",
+        node_id="123",
+        logging_config=LoggingConfig(),
+        long_poll_client="fake_client",
+    )
+    with pytest.raises(
+        RayActorError, match="must return a list of Starlette middlewares"
+    ):
+        ray.get(handle.ready.remote())
+
+
+@pytest.mark.parametrize(
+    "ray_instance",
+    [
+        {
+            "RAY_SERVE_HTTP_PROXY_CALLBACK_IMPORT_PATH": RAISE_ERROR_IMPORT_PATH,
+        },
+    ],
+    indirect=True,
+)
+def test_http_proxy_callback_failures(ray_instance, capsys):
+    """Test http proxy keeps restarting when callback function fails"""
+
+    try:
+        serve.start()
+    except RayActorError:
+        # serve.start will fail because the http proxy is not started successfully
+        # and client use proxy handle to check the proxy readiness, so it will raise
+        # RayActorError.
+        pass
+
+    client = _get_global_client()
+
+    def check_http_proxy_keep_restarting():
+        # The proxy will be under "STARTING" status and keep restarting.
+        prev_actor_id = None
+        while True:
+            serve_details = ServeInstanceDetails(**client.get_serve_details())
+            for _, proxy_info in serve_details.proxies.items():
+                if proxy_info.status != ProxyStatus.STARTING:
+                    return False
+                if prev_actor_id is None:
+                    prev_actor_id = proxy_info.actor_id
+                    break
+                elif prev_actor_id != proxy_info.actor_id:
+                    return True
+
+    wait_for_condition(check_http_proxy_keep_restarting)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", "-s", __file__]))