From db8921fb5c6b911e52a0fa69d1c600c5a5036e7f Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 20 May 2026 22:59:39 +0000 Subject: [PATCH] fix: forward HTTP headers for generate requests Co-authored-by: Mudit Aggarwal --- qa/L0_http/generate_endpoint_test.py | 30 +++++++-- qa/L0_http/test.sh | 4 +- .../generate_models/mock_llm/1/model.py | 16 +++-- .../mock_llm_ensemble/config.pbtxt | 61 +++++++++++++++++++ src/http_server.cc | 2 + 5 files changed, 102 insertions(+), 11 deletions(-) create mode 100644 qa/python_models/generate_models/mock_llm_ensemble/config.pbtxt diff --git a/qa/L0_http/generate_endpoint_test.py b/qa/L0_http/generate_endpoint_test.py index df9317b4dd..fb73bf1dfb 100755 --- a/qa/L0_http/generate_endpoint_test.py +++ b/qa/L0_http/generate_endpoint_test.py @@ -47,8 +47,10 @@ def setUp(self): def _get_infer_url(self, model_name, route): return f"http://localhost:8000/v2/models/{model_name}/{route}" - def generate_stream(self, model_name, inputs, stream=False): - headers = {"Accept": "text/event-stream"} + def generate_stream(self, model_name, inputs, stream=False, headers=None): + request_headers = {"Accept": "text/event-stream"} + if headers is not None: + request_headers.update(headers) url = self._get_infer_url(model_name, "generate_stream") # stream=True used to indicate response can be iterated over, which # should be the common setting for generate_stream. @@ -57,14 +59,16 @@ def generate_stream(self, model_name, inputs, stream=False): return requests.post( url, data=inputs if isinstance(inputs, str) else json.dumps(inputs), - headers=headers, + headers=request_headers, stream=stream, ) - def generate(self, model_name, inputs): + def generate(self, model_name, inputs, headers=None): url = self._get_infer_url(model_name, "generate") return requests.post( - url, data=inputs if isinstance(inputs, str) else json.dumps(inputs) + url, + data=inputs if isinstance(inputs, str) else json.dumps(inputs), + headers=headers, ) def generate_expect_failure(self, model_name, inputs, msg): @@ -201,6 +205,22 @@ def test_request_id(self): self.assertIn("TEXT", data) self.assertEqual(text, data["TEXT"]) + def test_generate_forwards_http_headers_to_parameters(self): + text = "hello world" + header_value = "forwarded header value" + headers = {"x-generate-header": header_value} + inputs = {"PROMPT": text, "STREAM": False} + + r = self.generate("mock_llm_ensemble", inputs, headers=headers) + r.raise_for_status() + + self.assertIn("Content-Type", r.headers) + self.assertEqual(r.headers["Content-Type"], "application/json") + + data = r.json() + self.assertIn("TEXT", data) + self.assertEqual(header_value, data["TEXT"]) + def test_generate_stream(self): # Setup text-based input text = "hello world" diff --git a/qa/L0_http/test.sh b/qa/L0_http/test.sh index e6b48272d2..0088e4d372 100755 --- a/qa/L0_http/test.sh +++ b/qa/L0_http/test.sh @@ -697,7 +697,7 @@ wait $SERVER_PID # https://github.com/mpetazzoni/sseclient pip install sseclient-py psutil -SERVER_ARGS="--model-repository=`pwd`/../python_models/generate_models --log-verbose=1" +SERVER_ARGS="--model-repository=`pwd`/../python_models/generate_models --log-verbose=1 --http-header-forward-pattern x-generate-header" SERVER_LOG="./inference_server_generate_endpoint_test.log" CLIENT_LOG="./generate_endpoint_test.log" run_server @@ -710,7 +710,7 @@ fi ## Python Unit Tests TEST_RESULT_FILE='test_results.txt' PYTHON_TEST=generate_endpoint_test.py -EXPECTED_NUM_TESTS=18 +EXPECTED_NUM_TESTS=19 set +e python $PYTHON_TEST > $CLIENT_LOG 2>&1 if [ $? -ne 0 ]; then diff --git a/qa/python_models/generate_models/mock_llm/1/model.py b/qa/python_models/generate_models/mock_llm/1/model.py index 1c92aaf1ad..c54da063e3 100644 --- a/qa/python_models/generate_models/mock_llm/1/model.py +++ b/qa/python_models/generate_models/mock_llm/1/model.py @@ -1,4 +1,4 @@ -# Copyright 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # Redistribution and use in source and binary forms, with or without # modification, are permitted provided that the following conditions @@ -49,6 +49,7 @@ def exec(self, requests): for request in requests: params = json.loads(request.parameters()) rep_count = params["REPETITION"] if "REPETITION" in params else 1 + header_text = params.get("x-generate-header") input_np = pb_utils.get_input_tensor_by_name(request, "PROMPT").as_numpy() stream_np = pb_utils.get_input_tensor_by_name(request, "STREAM").as_numpy() @@ -62,9 +63,12 @@ def exec(self, requests): ) ) else: - out_tensor = pb_utils.Tensor( - "TEXT", np.repeat(input_np, rep_count, axis=1) + out_np = ( + np.asarray([[header_text]], dtype=object) + if header_text is not None + else np.repeat(input_np, rep_count, axis=1) ) + out_tensor = pb_utils.Tensor("TEXT", out_np) responses.append(pb_utils.InferenceResponse([out_tensor])) return responses @@ -75,11 +79,15 @@ def exec_decoupled(self, requests): fail_last = params["FAIL_LAST"] if "FAIL_LAST" in params else False delay = params["DELAY"] if "DELAY" in params else None output_0_dim = params["OUTPUT_0_DIM"] if "OUTPUT_0_DIM" in params else False + header_text = params.get("x-generate-header") sender = request.get_response_sender() input_np = pb_utils.get_input_tensor_by_name(request, "PROMPT").as_numpy() stream_np = pb_utils.get_input_tensor_by_name(request, "STREAM").as_numpy() - out_value = np.array([]) if output_0_dim else input_np + if header_text is not None: + out_value = np.asarray([[header_text]], dtype=object) + else: + out_value = np.array([]) if output_0_dim else input_np out_tensor = pb_utils.Tensor("TEXT", out_value) response = pb_utils.InferenceResponse([out_tensor]) # If stream enabled, just send multiple copies of response diff --git a/qa/python_models/generate_models/mock_llm_ensemble/config.pbtxt b/qa/python_models/generate_models/mock_llm_ensemble/config.pbtxt new file mode 100644 index 0000000000..9ef1edcb39 --- /dev/null +++ b/qa/python_models/generate_models/mock_llm_ensemble/config.pbtxt @@ -0,0 +1,61 @@ +# Copyright 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions +# are met: +# * Redistributions of source code must retain the above copyright +# notice, this list of conditions and the following disclaimer. +# * Redistributions in binary form must reproduce the above copyright +# notice, this list of conditions and the following disclaimer in the +# documentation and/or other materials provided with the distribution. +# * Neither the name of NVIDIA CORPORATION nor the names of its +# contributors may be used to endorse or promote products derived +# from this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY +# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR +# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, +# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, +# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR +# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY +# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +platform: "ensemble" +max_batch_size: 0 + +input [ + { + name: "PROMPT" + data_type: TYPE_STRING + dims: [ 1, 1 ] + }, + { + name: "STREAM" + data_type: TYPE_BOOL + dims: [ 1, 1 ] + } +] + +output [ + { + name: "TEXT" + data_type: TYPE_STRING + dims: [ 1, -1 ] + } +] + +ensemble_scheduling { + step [ + { + model_name: "mock_llm" + model_version: -1 + input_map { key: "PROMPT", value: "PROMPT" } + input_map { key: "STREAM", value: "STREAM" } + output_map { key: "TEXT", value: "TEXT" } + } + ] +} diff --git a/src/http_server.cc b/src/http_server.cc index d351956449..72253c8e8f 100644 --- a/src/http_server.cc +++ b/src/http_server.cc @@ -3514,6 +3514,8 @@ HTTPAPIServer::HandleGenerate( input_metadata, generate_request->RequestSchema(), request), error_callback); + RETURN_AND_CALLBACK_IF_ERR(ForwardHeaders(req, irequest), error_callback); + auto request_release_payload = std::make_unique(irequest_shared, nullptr); // [FIXME] decompression..