File size: 2,618 Bytes
b386992 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 |
# Copyright (c) 2025, 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.
import numpy as np
from pytriton.decorators import batch
from pytriton.model_config import Tensor
from nemo.deploy import DeployPyTriton, ITritonDeployable
from nemo.deploy.nlp import NemoQueryLLM
from nemo.deploy.utils import cast_output, str_ndarray2list
class MockModel(ITritonDeployable):
@property
def get_triton_input(self):
inputs = (
Tensor(name="prompts", shape=(-1,), dtype=bytes),
Tensor(name="max_output_len", shape=(-1,), dtype=np.int_, optional=True),
Tensor(name="output_context_logits", shape=(-1,), dtype=np.bool_, optional=False),
Tensor(name="output_generation_logits", shape=(-1,), dtype=np.bool_, optional=False),
)
return inputs
@property
def get_triton_output(self):
outputs = (Tensor(name="outputs", shape=(-1,), dtype=bytes),)
return outputs
@batch
def triton_infer_fn(self, **inputs: np.ndarray):
infer_input = {"input_texts": str_ndarray2list(inputs.pop("prompts"))}
if "max_output_len" in inputs:
infer_input["max_output_len"] = inputs.pop("max_output_len")[0][0]
output_dict = dict()
output_dict["outputs"] = cast_output("I am good, how about you?", np.bytes_)
return output_dict
def test_nemo_deploy_query():
model_name = "mock_model"
model = MockModel()
nm = DeployPyTriton(
model=model,
triton_model_name=model_name,
max_batch_size=32,
http_port=9002,
grpc_port=8001,
address="0.0.0.0",
allow_grpc=True,
allow_http=True,
streaming=False,
)
nm.deploy()
nm.run()
nq = NemoQueryLLM(url="localhost:9002", model_name=model_name)
output_deployed = nq.query_llm(
prompts=["Hey, how is it going?"],
max_output_len=20,
)
nm.stop()
assert output_deployed is not None, "Output cannot be none."
assert output_deployed == "I am good, how about you?", "Output cannot be none."
|