File size: 3,081 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 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 |
# 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.
from unittest.mock import MagicMock, patch
import pytest
from nemo.deploy.deploy_base import DeployBase
class MockDeployable(DeployBase):
def deploy(self):
pass
def serve(self):
pass
def run(self):
pass
def stop(self):
pass
class MockTritonDeployable:
pass
@pytest.fixture
def mock_model():
return MagicMock()
@pytest.fixture
def deploy_base(mock_model):
return MockDeployable(
triton_model_name="test_model",
model=mock_model,
max_batch_size=128,
http_port=8000,
grpc_port=8001,
)
def test_initialization_with_model(deploy_base, mock_model):
assert deploy_base.triton_model_name == "test_model"
assert deploy_base.model == mock_model
assert deploy_base.max_batch_size == 128
assert deploy_base.http_port == 8000
assert deploy_base.grpc_port == 8001
assert deploy_base.address == "0.0.0.0"
assert deploy_base.allow_grpc is True
assert deploy_base.allow_http is True
assert deploy_base.streaming is False
def test_initialization_with_checkpoint():
with patch('nemo.deploy.deploy_base.ModelPT') as mock_model_pt:
mock_model_pt.restore_from.return_value = MagicMock()
deploy_base = MockDeployable(
triton_model_name="test_model",
checkpoint_path="test.ckpt",
)
assert deploy_base.checkpoint_path == "test.ckpt"
def test_initialization_without_model_or_checkpoint():
with pytest.raises(Exception) as exc_info:
MockDeployable(triton_model_name="test_model")
assert "Either checkpoint_path or model should be provided" in str(exc_info.value)
def test_get_module_and_class():
module, class_name = DeployBase.get_module_and_class("nemo.models.test_model.TestModel")
assert module == "nemo.models.test_model"
assert class_name == "TestModel"
def test_is_model_deployable_valid(deploy_base):
deploy_base.model = MockTritonDeployable()
with patch('nemo.deploy.deploy_base.ITritonDeployable', MockTritonDeployable):
assert deploy_base._is_model_deployable() is True
def test_is_model_deployable_invalid(deploy_base):
deploy_base.model = MagicMock()
with patch('nemo.deploy.deploy_base.ITritonDeployable', MockTritonDeployable):
with pytest.raises(Exception) as exc_info:
deploy_base._is_model_deployable()
assert "This model is not deployable to Triton" in str(exc_info.value)
|