From ac039bd04ed31f7dc09f69d89fb36c69276ad551 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Wed, 1 Apr 2026 04:26:11 -0700 Subject: [PATCH] Use CustomTestCase for TestSessionControl to enable CI retry (#21830) --- test/registered/sessions/test_session_control.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/test/registered/sessions/test_session_control.py b/test/registered/sessions/test_session_control.py index 48dd5f434..0e6b2a7f2 100644 --- a/test/registered/sessions/test_session_control.py +++ b/test/registered/sessions/test_session_control.py @@ -31,7 +31,7 @@ def remove_prefix(text: str, prefix: str) -> str: return text[len(prefix) :] if text.startswith(prefix) else text -class TestSessionControl(unittest.TestCase): +class TestSessionControl(CustomTestCase): @classmethod def setUpClass(cls): cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST @@ -43,6 +43,8 @@ class TestSessionControl(unittest.TestCase): other_args=[ "--attention-backend", "triton", + "--disable-cuda-graph", + "--disable-piecewise-cuda-graph", ], )