From 25ffbaa01b696dd8824875b054c39990b7690583 Mon Sep 17 00:00:00 2001 From: Andrew White Date: Fri, 24 Jul 2026 09:30:15 -0500 Subject: [PATCH] fix: use integer dtype for random_seed in query script The standalone query script passed random_seed as np.single (float), but the Triton input spec and the streaming variant use np.int_. This could cause type mismatches when the model expects an integer seed. Adds a unit test that mocks pytriton and verifies the dtype is np.int_. Signed-off-by: Andrew White --- scripts/deploy/nlp/query.py | 2 +- tests/unit_tests/deploy/test_query_script.py | 54 ++++++++++++++++++++ 2 files changed, 55 insertions(+), 1 deletion(-) create mode 100644 tests/unit_tests/deploy/test_query_script.py diff --git a/scripts/deploy/nlp/query.py b/scripts/deploy/nlp/query.py index e953695ad..070c9f32f 100644 --- a/scripts/deploy/nlp/query.py +++ b/scripts/deploy/nlp/query.py @@ -114,7 +114,7 @@ def query_llm( inputs["temperature"] = np.full(prompts.shape, temperature, dtype=np.single) if random_seed is not None: - inputs["random_seed"] = np.full(prompts.shape, random_seed, dtype=np.single) + inputs["random_seed"] = np.full(prompts.shape, random_seed, dtype=np.int_) if stop_words_list is not None: stop_words_list = np.char.encode(stop_words_list, "utf-8") diff --git a/tests/unit_tests/deploy/test_query_script.py b/tests/unit_tests/deploy/test_query_script.py new file mode 100644 index 000000000..cf37a2167 --- /dev/null +++ b/tests/unit_tests/deploy/test_query_script.py @@ -0,0 +1,54 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest.mock import MagicMock, patch + +import numpy as np +import pytest + + +@pytest.fixture +def query_llm_module(): + """Import query.py with pytriton clients mocked.""" + with patch.dict( + "sys.modules", + { + "pytriton": MagicMock(), + "pytriton.client": MagicMock(), + }, + ): + from scripts.deploy.nlp import query as query_module + + yield query_module + + +class TestQueryLLMInputs: + @patch("scripts.deploy.nlp.query.ModelClient") + def test_random_seed_uses_int_dtype(self, mock_client, query_llm_module): + """random_seed must be passed as an integer array, not float.""" + mock_instance = MagicMock() + mock_client.return_value.__enter__.return_value = mock_instance + mock_instance.infer_batch.return_value = {"outputs": np.array([b"response"])} + mock_instance.model_config.outputs = [MagicMock(dtype=np.bytes_)] + + query_llm_module.query_llm( + url="localhost:8000", + model_name="test", + prompts=["hello"], + random_seed=42, + ) + + call_kwargs = mock_instance.infer_batch.call_args.kwargs + assert call_kwargs["random_seed"].dtype == np.int_ + assert call_kwargs["random_seed"][0] == 42