From 1e13f0a063be711ac5e109e36286c5b805397ee0 Mon Sep 17 00:00:00 2001 From: Iman Tabrizian Date: Thu, 2 May 2024 16:14:18 +0000 Subject: [PATCH] Review edits --- .../tests/test_iterative_scheduler.py | 52 ++++--------------- 1 file changed, 11 insertions(+), 41 deletions(-) diff --git a/Conceptual_Guide/Part_7-iterative_scheduling/tests/test_iterative_scheduler.py b/Conceptual_Guide/Part_7-iterative_scheduling/tests/test_iterative_scheduler.py index 0ce74663..7988e787 100644 --- a/Conceptual_Guide/Part_7-iterative_scheduling/tests/test_iterative_scheduler.py +++ b/Conceptual_Guide/Part_7-iterative_scheduling/tests/test_iterative_scheduler.py @@ -1,46 +1,22 @@ -# Copyright 2024, 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. - import time import unittest +import pytest from tritonserver import InferenceRequest, Model, Server class IterativeSchedulerTest(unittest.TestCase): - def setUp(self): - self._server = Server( + @classmethod + def setUpClass(cls): + cls._server = Server( model_repository="../model_repository", log_info=True, log_error=True, strict_model_config=False, ) - self._server.start(wait_until_ready=True) + cls._server.start(wait_until_ready=True) - def _infer_and_verify_request(self, max_tokens, prompt, ignore_eos, model_name): + def _infer_and_verify_response(self, max_tokens, prompt, ignore_eos, model_name): parameters = {"max_tokens": max_tokens, "ignore_eos": ignore_eos} model = Model(self._server, name=model_name) self.assertTrue(model.ready()) @@ -55,23 +31,17 @@ def _infer_and_verify_request(self, max_tokens, prompt, ignore_eos, model_name): self.assertEqual(max_tokens, len(responses)) def test_max_tokens(self): - self._infer_and_verify_request( + self._infer_and_verify_response( 32, "Triton Inference Server is", ignore_eos=True, model_name="iterative-gpt2", ) - self._infer_and_verify_request( + self._infer_and_verify_response( 32, "Triton Inference Server is", ignore_eos=True, model_name="simple-gpt2" ) - def tearDown(self): - # Add a small delay to allow graceful shutdown - # This is related to a known issue with server - # stop. + @classmethod + def tearDownClass(cls): time.sleep(5) - self._server.stop() - - -if __name__ == "__main__": - unittest.main() + cls._server.stop()