[Serve] Gate ingress request router body forwarding behind escape hatch (#63183)

Signed-off-by: Seiji Eicher <seiji@anyscale.com>
Signed-off-by: Kourosh Hakhamaneshi <kourosh@anyscale.com>
Co-authored-by: Kourosh Hakhamaneshi <kourosh@anyscale.com>

Author: eicherseiji

python/ray/serve/_private/constants.py

@@ -799,10 +799,25 @@
 # active. Bodies longer than this are truncated; the Lua forwards what it has
 # with an `X-Body-Truncated: <bytes>/<content-length>` header so the router can
 # do best-effort prefix matching. Memory cost is ~2 * bufsize * maxconn.
+# Only consulted when RAY_SERVE_INGRESS_REQUEST_ROUTER_FORWARD_BODY=1.
 RAY_SERVE_HAPROXY_INGRESS_REQUEST_ROUTER_BUFSIZE = get_env_int(
     "RAY_SERVE_HAPROXY_INGRESS_REQUEST_ROUTER_BUFSIZE", 262144
 )
 
+# Escape hatch: when true, HAProxy forwards the (possibly truncated) request
+# body to /internal/route and the router reads it. Off by default because for
+# large payloads the body buffering / re-emit cost adds noticeable time-to-
+# first-response. Skipping the forward is fine for any policy whose decision
+# does not depend on the request body: round-robin and power-of-two ignore
+# the body entirely, and session-aware policies key on the ``x-session-id``
+# header (forwarded with the request line) rather than the body.
+#
+# Flip this to true if the configured request router needs the body for its
+# decision, e.g. prefix-aware / prefix-cache routing.
+RAY_SERVE_INGRESS_REQUEST_ROUTER_FORWARD_BODY = get_env_bool(
+    "RAY_SERVE_INGRESS_REQUEST_ROUTER_FORWARD_BODY", False
+)
+
 RAY_SERVE_DIRECT_INGRESS_MIN_HTTP_PORT = int(
     os.environ.get("RAY_SERVE_DIRECT_INGRESS_MIN_HTTP_PORT", "30000")
 )

python/ray/serve/_private/haproxy.py

@@ -56,6 +56,7 @@
     RAY_SERVE_HAPROXY_TIMEOUT_CLIENT_S,
     RAY_SERVE_HAPROXY_TIMEOUT_CONNECT_S,
     RAY_SERVE_HAPROXY_TIMEOUT_SERVER_S,
+    RAY_SERVE_INGRESS_REQUEST_ROUTER_FORWARD_BODY,
     SERVE_CONTROLLER_NAME,
     SERVE_LOGGER_NAME,
     SERVE_NAMESPACE,
@@ -804,6 +805,7 @@ def _write_ingress_request_router_lua(
 
         content = _load_lua_template().substitute(
             TIMEOUT_S=RAY_SERVE_HAPROXY_INGRESS_REQUEST_ROUTER_TIMEOUT_S,
+            FORWARD_BODY=str(RAY_SERVE_INGRESS_REQUEST_ROUTER_FORWARD_BODY).lower(),
             ROUTERS=_format_routers_lua(routers),
             REPLICA_TARGETS=_format_replica_targets_lua(targets),
         )
@@ -815,7 +817,7 @@ def _write_ingress_request_router_lua(
             logger.debug(f"Wrote Lua routing script to {lua_path}")
         return lua_path
 
-    def _generate_config_file_internal(self) -> None:
+    def _generate_config_file_internal(self) -> bool:
         """Internal config generation without locking (for use within locked sections)."""
         try:
             env = Environment()
@@ -872,6 +874,9 @@ def _generate_config_file_internal(self) -> None:
                     "ingress_request_router_bufsize": (
                         RAY_SERVE_HAPROXY_INGRESS_REQUEST_ROUTER_BUFSIZE
                     ),
+                    "ingress_request_router_forward_body": (
+                        RAY_SERVE_INGRESS_REQUEST_ROUTER_FORWARD_BODY
+                    ),
                 }
             )
 

python/ray/serve/_private/haproxy_templates.py

@@ -28,6 +28,8 @@
     nbthread {{ config.nbthread }}
     {%- if has_ingress_request_router %}
     lua-load-per-thread {{ ingress_request_router_lua_path }}
+    {%- endif %}
+    {%- if has_ingress_request_router and ingress_request_router_forward_body %}
     tune.bufsize {{ ingress_request_router_bufsize }}
     {%- endif %}
     {%- if config.enable_hap_optimization %}
@@ -97,7 +99,9 @@
     {%- endif %}
     {%- endfor %}
     acl has_ingress_request_router_app var(txn.ingress_request_router_app) -m found
+    {%- if ingress_request_router_forward_body %}
     http-request wait-for-body time {{ ingress_request_router_timeout_s }}s if METH_POST has_ingress_request_router_app
+    {%- endif %}
     http-request lua.route_via_ingress_request_router if METH_POST has_ingress_request_router_app
     # Fail loudly when the Lua dispatch did not pick a replica. Must appear
     # before the use_backend rules below so the request never falls back to
@@ -163,6 +167,10 @@
 {%- if has_ingress_request_router and backend.ingress_request_router_servers %}
 backend {{ backend.name or 'unknown' }}-via-ingress-request-router
     log global
+    # Keep the pinned data-plane path on the same connection policy as the
+    # primary backend. For streamed responses, forcing server-close can leave
+    # HAProxy holding unread server-side FINs under a burst while worker
+    # threads are still routing other requests.
     http-reuse always
     # use-server falls through to LB if the pinned server is DOWN.
     option redispatch

python/ray/serve/_private/ingress_request_router.lua.tmpl

@@ -4,8 +4,10 @@
 --
 -- Bodies exceeding tune.bufsize are truncated; we forward what we have with
 -- X-Body-Truncated since prefix-routing only needs the head of the body.
+-- See RAY_SERVE_INGRESS_REQUEST_ROUTER_FORWARD_BODY in constants.py.
 
 local ROUTER_REQUEST_TIMEOUT_S = ${TIMEOUT_S}
+local FORWARD_BODY = ${FORWARD_BODY}
 
 -- Per-app state. The frontend sets txn.ingress_request_router_app to the
 -- backend name whose path prefix matched, and we look up that app's
@@ -97,17 +99,24 @@ core.register_action("route_via_ingress_request_router", {"http-req"}, function(
         return
     end
 
-    local body = txn.sf:req_body()
-    if not body or body == "" then
-        txn:set_var("txn.ingress_request_router_failed", "empty_body")
-        return
-    end
-
-    local truncated = truncated_full_length(txn, body)
-    if truncated then
-        core.log(core.warning,
-            "ingress_request_router: forwarding truncated body to router ("
-                .. #body .. "/" .. truncated .. " bytes)")
+    -- FORWARD_BODY=false: don't read or forward the body; call the router
+    -- with body="" so a Content-Length: 0 POST still goes through routing
+    -- (any policy that needs the body must opt in via FORWARD_BODY=true).
+    -- Empty body in FORWARD_BODY=true mode is treated the same -- a
+    -- legitimate input the router can accept or reject on its own terms;
+    -- we don't synthesize a sentinel for it here.
+    local body = ""
+    local truncated = nil
+    if FORWARD_BODY then
+        body = txn.sf:req_body() or ""
+        if body ~= "" then
+            truncated = truncated_full_length(txn, body)
+            if truncated then
+                core.log(core.warning,
+                    "ingress_request_router: forwarding truncated body to router ("
+                        .. #body .. "/" .. truncated .. " bytes)")
+            end
+        end
     end
 
     local response = call_router(router, body, truncated)

python/ray/serve/tests/test_haproxy_api.py

@@ -583,6 +583,11 @@ def test_ingress_request_router_does_not_leak_into_other_backends(
 
         assert "backend llm-via-ingress-request-router" in cfg
         assert "backend api-via-ingress-request-router" not in cfg
+        assert "option http-buffer-request" not in cfg
+        direct_backend = cfg.split("backend llm-via-ingress-request-router", 1)[1]
+        direct_backend = direct_backend.split("listen stats", 1)[0]
+        assert "http-reuse always" in direct_backend
+        assert "option http-server-close" not in direct_backend
         # Only router-bearing backends contribute a set-var directive that
         # arms the Lua dispatch; the plain `api` backend must not.
         assert "set-var(txn.ingress_request_router_app) str(llm)" in cfg
@@ -651,6 +656,64 @@ def test_router_failure_503_rule_appears_before_use_backend(haproxy_api_cleanup)
         assert "X-Serve-Reason" in cfg, cfg
 
 
+@pytest.mark.parametrize("forward_body", [True, False])
+def test_ingress_request_router_forward_body_gate_renders(
+    haproxy_api_cleanup, monkeypatch, forward_body
+):
+    """The FORWARD_BODY escape hatch must drive both:
+    - HAProxy ``wait-for-body`` + ``tune.bufsize`` directives (memory cost
+      and the per-request body round-trip), and
+    - the Lua ``FORWARD_BODY`` constant (whether the action reads the body
+      and forwards it to ``/internal/route``).
+
+    Off by default: round-robin ignores the body, so neither cost is paid.
+    """
+    monkeypatch.setattr(
+        "ray.serve._private.haproxy.RAY_SERVE_INGRESS_REQUEST_ROUTER_FORWARD_BODY",
+        forward_body,
+    )
+    with tempfile.TemporaryDirectory() as temp_dir:
+        api = _make_api(
+            temp_dir,
+            {
+                "llm": BackendConfig(
+                    name="llm",
+                    path_prefix="/",
+                    app_name="llm",
+                    servers=[
+                        ServerConfig(
+                            name="r1",
+                            host="10.0.0.1",
+                            port=30001,
+                            replica_id="rid_1",
+                        )
+                    ],
+                    ingress_request_router_servers=[
+                        ServerConfig(name="router", host="10.0.0.10", port=9000)
+                    ],
+                ),
+            },
+        )
+        with mock.patch(
+            "ray.serve._private.constants.RAY_SERVE_HAPROXY_CONFIG_FILE_LOC",
+            api.config_file_path,
+        ):
+            api._generate_config_file_internal()
+        with open(api.config_file_path) as f:
+            cfg = f.read()
+        with open(os.path.join(temp_dir, "ingress_request_router.lua")) as f:
+            lua = f.read()
+
+        if forward_body:
+            assert "tune.bufsize" in cfg, cfg
+            assert "wait-for-body" in cfg, cfg
+            assert "local FORWARD_BODY = true" in lua, lua
+        else:
+            assert "tune.bufsize" not in cfg, cfg
+            assert "wait-for-body" not in cfg, cfg
+            assert "local FORWARD_BODY = false" in lua, lua
+
+
 def _create_replica_server(port: int, replica_id_header: str):
     """Fake data-plane replica that echoes its identity in a response header."""
     app = FastAPI()
@@ -670,7 +733,7 @@ async def root(path: str, req: Request, res: Response):
 
 def _create_router_server(port: int, replica_id_to_return: str):
     """Fake /internal/route. Captures bodies so tests can verify HAProxy
-    actually buffered + forwarded the request."""
+    forwards the buffered request body prefix to the router."""
     app = FastAPI()
     captured = {"bodies": []}
 
@@ -689,14 +752,20 @@ def ready():
         )
 
     server, thread = _serve_fastapi_app(app, port, ready)
+    # Discard the readiness-probe body so callers see only client traffic.
+    captured["bodies"].clear()
     return server, thread, captured
 
 
 @pytest.mark.asyncio
-async def test_ingress_request_router_end_to_end(haproxy_api_cleanup):
+async def test_ingress_request_router_end_to_end(haproxy_api_cleanup, monkeypatch):
     """Run actual HAProxy against a fake router + two replicas; verify a POST
     is pinned to the replica the router selects, while a GET (which doesn't
     trigger the router-routed path) is not."""
+    monkeypatch.setattr(
+        "ray.serve._private.haproxy.RAY_SERVE_INGRESS_REQUEST_ROUTER_FORWARD_BODY",
+        True,
+    )
     with tempfile.TemporaryDirectory() as temp_dir:
         haproxy_port = find_free_port()
         stats_port = find_free_port()
@@ -787,12 +856,9 @@ async def test_ingress_request_router_end_to_end(haproxy_api_cleanup):
             assert resp.status_code == 200, resp.text
             assert resp.headers.get("x-replica-id") == "B"
 
-            # The router actually saw the original body (via wait-for-body +
-            # txn.sf:req_body()). Just check the field made it through.
-            assert any(
-                '"prompt"' in body and "hello" in body
-                for body in router_captured["bodies"]
-            )
+            # Direct streaming keeps a bounded request-body path for
+            # prefix-cache-aware routing.
+            assert router_captured["bodies"] == ['{"prompt": "hello"}']
 
             # Repeat to confirm the pin holds across requests.
             for _ in range(3):
@@ -802,9 +868,10 @@ async def test_ingress_request_router_end_to_end(haproxy_api_cleanup):
                     timeout=5,
                 )
                 assert resp.headers.get("x-replica-id") == "B"
+            assert router_captured["bodies"] == ['{"prompt": "hello"}'] * 4
 
-            # GET is not POST, so wait-for-body / Lua never run; the router
-            # should have seen exactly the four POSTs above and nothing more.
+            # GET is not POST, so Lua routing never runs; the router should
+            # have seen exactly the four POSTs above and nothing more.
             n_router_calls_before_get = len(router_captured["bodies"])
             requests.get(
                 f"http://127.0.0.1:{haproxy_port}/health-passthrough", timeout=5
@@ -936,28 +1003,23 @@ async def test_router_failure_fails_loud_with_reason(haproxy_api_cleanup):
                 timeout=10,
             )
 
-            # Every dispatch failure mode must surface as a 5xx with a
-            # reason label, never as a silent primary-backend fallback.
-            # ``router_non_200``: router answered with status != 200 (forced
-            # by the broken-router stub).
-            # ``empty_body``: HAProxy's wait-for-body completed but the
-            # request had no body to forward to the router.
-            failure_cases = [
-                (dict(json={"prompt": "hi"}), "router_non_200"),
-                (dict(data=""), "empty_body"),
-            ]
-            for kwargs, expected_reason in failure_cases:
+            # Every dispatch failure must surface as 5xx with a reason
+            # label, never as a silent primary-backend fallback. The broken
+            # router returns 500 for both empty and non-empty bodies, so
+            # both shapes surface the same ``router_non_200`` reason; the
+            # body shape is parametrized to pin that empty-body POSTs are
+            # routed through the router and not silently bypassed.
+            for body_kwargs in (dict(json={"prompt": "hi"}), dict(data="")):
                 for _ in range(3):
                     resp = requests.post(
                         f"http://127.0.0.1:{haproxy_port}/predict",
                         timeout=5,
-                        **kwargs,
-                    )
-                    assert resp.status_code == 503, (expected_reason, resp.text)
-                    assert resp.headers.get("X-Serve-Reason") == expected_reason, (
-                        expected_reason,
-                        resp.headers,
+                        **body_kwargs,
                     )
+                    assert resp.status_code == 503, resp.text
+                    assert (
+                        resp.headers.get("X-Serve-Reason") == "router_non_200"
+                    ), resp.headers
 
             stats_csv = requests.get(
                 f"http://127.0.0.1:{stats_port}/stats;csv", timeout=5