From 899136b857d510f0e19b0e21ea32b49f8aa117ed Mon Sep 17 00:00:00 2001 From: youkaichao Date: Sat, 11 Jan 2025 09:07:24 +0800 Subject: [PATCH] [ci] fix broken distributed-tests-4-gpus (#11937) Signed-off-by: youkaichao --- tests/spec_decode/e2e/test_integration_dist_tp4.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/spec_decode/e2e/test_integration_dist_tp4.py b/tests/spec_decode/e2e/test_integration_dist_tp4.py index 555aef99218c3..2cb10de1c6f55 100644 --- a/tests/spec_decode/e2e/test_integration_dist_tp4.py +++ b/tests/spec_decode/e2e/test_integration_dist_tp4.py @@ -108,7 +108,8 @@ def test_skip_speculation(common_llm_kwargs, per_test_common_llm_kwargs, TODO: fix it to pass without raising Error. (#5814) """ - with pytest.raises(openai.APIConnectionError): + with pytest.raises( + (openai.APIConnectionError, openai.InternalServerError)): run_equality_correctness_test_tp(MAIN_MODEL, common_llm_kwargs, per_test_common_llm_kwargs,