diff --git a/qa/L0_lifecycle/lifecycle_test.py b/qa/L0_lifecycle/lifecycle_test.py index 142fb0494a..36b415e393 100755 --- a/qa/L0_lifecycle/lifecycle_test.py +++ b/qa/L0_lifecycle/lifecycle_test.py @@ -559,7 +559,12 @@ def test_dynamic_model_load_unload(self): # load. Make sure that it has a status and is ready. try: shutil.copytree(libtorch_name, "models/" + libtorch_name) - time.sleep(5) # wait for model to load + tu.wait_for_model_state( + httpclient.InferenceServerClient("localhost:8000", verbose=True), + libtorch_name, + True, + "1", + ) for triton_client in ( httpclient.InferenceServerClient("localhost:8000", verbose=True), grpcclient.InferenceServerClient("localhost:8001", verbose=True), @@ -635,7 +640,12 @@ def test_dynamic_model_load_unload(self): # time to unload. Make sure that it is no longer available. try: shutil.rmtree("models/" + libtorch_name) - time.sleep(5) # wait for model to unload + tu.wait_for_model_state( + httpclient.InferenceServerClient("localhost:8000", verbose=True), + libtorch_name, + False, + "1", + ) for triton_client in ( httpclient.InferenceServerClient("localhost:8000", verbose=True), grpcclient.InferenceServerClient("localhost:8001", verbose=True), @@ -675,7 +685,12 @@ def test_dynamic_model_load_unload(self): # Add back the same model. The status/stats should be reset. try: shutil.copytree(libtorch_name, "models/" + libtorch_name) - time.sleep(5) # wait for model to load + tu.wait_for_model_state( + httpclient.InferenceServerClient("localhost:8000", verbose=True), + libtorch_name, + True, + "1", + ) for triton_client in ( httpclient.InferenceServerClient("localhost:8000", verbose=True), grpcclient.InferenceServerClient("localhost:8001", verbose=True), @@ -718,7 +733,12 @@ def test_dynamic_model_load_unload(self): # time to unload. Make sure that it is unavailable. try: shutil.rmtree("models/" + onnx_name) - time.sleep(5) # wait for model to unload + tu.wait_for_model_state( + httpclient.InferenceServerClient("localhost:8000", verbose=True), + onnx_name, + False, + "1", + ) for triton_client in ( httpclient.InferenceServerClient("localhost:8000", verbose=True), grpcclient.InferenceServerClient("localhost:8001", verbose=True), @@ -929,7 +949,12 @@ def test_dynamic_version_load_unload(self): # unload. Make sure that it is unavailable. try: shutil.rmtree("models/" + libtorch_name + "/1") - time.sleep(5) # wait for version to unload + tu.wait_for_model_state( + httpclient.InferenceServerClient("localhost:8000", verbose=True), + libtorch_name, + False, + "1", + ) for triton_client in ( httpclient.InferenceServerClient("localhost:8000", verbose=True), grpcclient.InferenceServerClient("localhost:8001", verbose=True), @@ -969,7 +994,12 @@ def test_dynamic_version_load_unload(self): shutil.copytree( "models/" + libtorch_name + "/2", "models/" + libtorch_name + "/7" ) - time.sleep(5) # wait for version to load + tu.wait_for_model_state( + httpclient.InferenceServerClient("localhost:8000", verbose=True), + libtorch_name, + True, + "7", + ) for triton_client in ( httpclient.InferenceServerClient("localhost:8000", verbose=True), grpcclient.InferenceServerClient("localhost:8001", verbose=True), @@ -1090,7 +1120,13 @@ def test_dynamic_model_modify(self): "models/" + model_name + "/config.pbtxt", ) - time.sleep(5) # wait for models to reload + for model_name in models: + tu.wait_for_model_state( + httpclient.InferenceServerClient("localhost:8000", verbose=True), + model_name, + True, + str(version), + ) for model_name in models: for model_name, model_shape in zip(models_base, models_shape): try: @@ -1119,7 +1155,13 @@ def test_dynamic_model_modify(self): "config.pbtxt." + base_name, "models/" + model_name + "/config.pbtxt" ) - time.sleep(5) # wait for models to reload + for model_name in models: + tu.wait_for_model_state( + httpclient.InferenceServerClient("localhost:8000", verbose=True), + model_name, + True, + "3", + ) for model_name in models: try: for triton_client in ( @@ -1217,7 +1259,13 @@ def test_dynamic_file_delete(self): for model_name in models: os.remove("models/" + model_name + "/config.pbtxt") - time.sleep(5) # wait for models to reload + for model_name in models: + tu.wait_for_model_state( + httpclient.InferenceServerClient("localhost:8000", verbose=True), + model_name, + True, + "3", + ) for model_name in models: try: for triton_client in ( @@ -1286,7 +1334,12 @@ def test_multiple_model_repository_polling(self): # Add the libtorch to the second model repository, should cause # it to be unloaded due to duplication shutil.copytree(libtorch_name, "models_0/" + libtorch_name) - time.sleep(5) # wait for models to reload + tu.wait_for_model_state( + httpclient.InferenceServerClient("localhost:8000", verbose=True), + libtorch_name, + False, + "1", + ) try: for triton_client in ( httpclient.InferenceServerClient("localhost:8000", verbose=True), @@ -1306,7 +1359,12 @@ def test_multiple_model_repository_polling(self): # properly. In the second model repository libtorch should # have versions 1 and 3. shutil.rmtree("models/" + libtorch_name) - time.sleep(5) # wait for model to unload + tu.wait_for_model_state( + httpclient.InferenceServerClient("localhost:8000", verbose=True), + libtorch_name, + True, + "1", + ) self._infer_success_models( ["libtorch", "openvino", "onnx"], (1, 3), model_shape ) @@ -3441,7 +3499,13 @@ def test_add_custom_config(self): "models/" + model_name + "/configs/custom.pbtxt", ) - time.sleep(5) # wait for models to reload + for model_name in models: + tu.wait_for_model_state( + httpclient.InferenceServerClient("localhost:8000", verbose=True), + model_name, + True, + "2", + ) for model_name in models: try: for triton_client in ( @@ -3483,7 +3547,13 @@ def test_delete_custom_config(self): for model_name in models: os.remove("models/" + model_name + "/configs/custom.pbtxt") - time.sleep(5) # wait for models to reload + for model_name in models: + tu.wait_for_model_state( + httpclient.InferenceServerClient("localhost:8000", verbose=True), + model_name, + True, + "1", + ) for model_name in models: try: for triton_client in ( diff --git a/qa/common/test_util.py b/qa/common/test_util.py index 6f918ded6c..444414ee52 100755 --- a/qa/common/test_util.py +++ b/qa/common/test_util.py @@ -28,6 +28,7 @@ import json import os +import time import unittest import ml_dtypes @@ -79,6 +80,48 @@ def wait_for_stable_rss(server, rss_tolerance_bytes=0.1 * MIB, stable_threshold= return +def wait_for_model_state( + triton_client, + model_name, + expected_ready, + model_version="", + timeout_sec=30, + interval_sec=1, +): + """Poll ``is_model_ready`` until it matches *expected_ready* or timeout. + + Replaces blind ``time.sleep`` calls that wait for model load/unload + with a deterministic polling loop, reducing flaky failures under CI + resource contention. + + Args: + triton_client: A Triton ``InferenceServerClient`` (HTTP or gRPC). + model_name: Name of the model to check. + expected_ready: ``True`` to wait for the model to become ready, + ``False`` to wait for it to become *not* ready. + model_version: Version string (empty string for default/latest). + timeout_sec: Maximum seconds to wait before raising. + interval_sec: Seconds between consecutive polls. + + Raises: + AssertionError: If the model does not reach the expected state + within *timeout_sec*. + """ + deadline = time.time() + timeout_sec + while time.time() < deadline: + if triton_client.is_model_ready(model_name, model_version) == expected_ready: + return + time.sleep(interval_sec) + actual = triton_client.is_model_ready(model_name, model_version) + if actual != expected_ready: + raise AssertionError( + "Model '{}' version '{}' did not reach expected ready={} " + "within {}s (actual ready={})".format( + model_name, model_version, expected_ready, timeout_sec, actual + ) + ) + + def shape_element_count(shape): cnt = 0 for d in shape: