From ce647cbcfdd06c192a71437b0d503b9c6c52ed38 Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Sun, 31 May 2026 09:51:39 +0800 Subject: [PATCH] Add a SimplePhaseChecker for execution-phase assertions (#26803) --- python/sglang/srt/environ.py | 1 + python/sglang/srt/utils/phase_checker.py | 109 +++++ test/registered/utils/test_phase_checker.py | 483 ++++++++++++++++++++ 3 files changed, 593 insertions(+) create mode 100644 python/sglang/srt/utils/phase_checker.py create mode 100644 test/registered/utils/test_phase_checker.py diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 8ba2f9079..9ca64d7fc 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -238,6 +238,7 @@ class Envs: SGLANG_FORCE_SHUTDOWN = EnvBool(False) SGLANG_DEBUG_MEMORY_POOL = EnvBool(False) SGLANG_DEBUG_REVERT_PR = EnvInt(0) + SGLANG_PHASE_CHECKER_DEBUG = EnvBool(False) SGLANG_TEST_REQUEST_TIME_STATS = EnvBool(False) SGLANG_DISABLE_TP_MEMORY_INBALANCE_CHECK = EnvBool(False) SGLANG_SIMULATE_ACC_LEN = EnvFloat(-1) diff --git a/python/sglang/srt/utils/phase_checker.py b/python/sglang/srt/utils/phase_checker.py new file mode 100644 index 000000000..84c1a9891 --- /dev/null +++ b/python/sglang/srt/utils/phase_checker.py @@ -0,0 +1,109 @@ +from __future__ import annotations + +from enum import IntEnum + +import torch +import triton +import triton.language as tl + +from sglang.srt.environ import envs + + +def _phase_repr(phase: int | IntEnum) -> str: + if isinstance(phase, IntEnum): + return f"{phase.name}({int(phase)})" + return str(int(phase)) + + +def _host_debug(msg: str) -> None: + if envs.SGLANG_PHASE_CHECKER_DEBUG.get(): + print(msg, flush=True) + + +# debug=True so tl.device_assert below actually raises. Without it the assert +# is stripped at compile time and only tl.device_print fires (the assert is +# gated on the TRITON_DEBUG env var by default — see tl.device_assert docstring). +@triton.jit(debug=True) +def _phase_check_kernel( + phase_ptr, + enable_assert_ptr, + EXPECT_PHASE: tl.constexpr, + NEXT_PHASE: tl.constexpr, + CALLER_TAG: tl.constexpr, +): + cur = tl.load(phase_ptr) + enable_assert = tl.load(enable_assert_ptr) + if enable_assert != 0: + if cur != EXPECT_PHASE: + # constexpr values get baked into the prefix string at compile time; + # only `cur` is runtime. + tl.device_print( + f"[SimplePhaseChecker FAIL] caller_tag={CALLER_TAG} " + f"expect={EXPECT_PHASE} next={NEXT_PHASE} actual=", + cur, + ) + tl.device_assert(cur == EXPECT_PHASE, "SimplePhaseChecker: phase mismatch") + tl.store(phase_ptr, NEXT_PHASE) + + +class SimplePhaseChecker: + """GPU-side state machine for any int-keyed phase sequence.""" + + def __init__(self, *, initial_phase: int | IntEnum, device: torch.device) -> None: + self._initial_phase = int(initial_phase) + self._phase = torch.tensor( + self._initial_phase, dtype=torch.int32, device=device + ) + self._enable_assert_device = torch.zeros(1, dtype=torch.int32, device=device) + self._caller_tag_registry: dict[str, int] = {} + _host_debug( + f"[SimplePhaseChecker.__init__] device={device} " + f"initial_phase={_phase_repr(initial_phase)} " + f"enable_assert=OFF (call enable_assert() after init is done)" + ) + + def enable_assert(self) -> None: + """Reset phase to initial_phase, then enable the device-side assert.""" + self._reset_to_idle() + self._enable_assert_device.fill_(1) + _host_debug(f"[SimplePhaseChecker.enable_assert] assert ENABLED") + + def update( + self, + *, + expect_phase: int | IntEnum, + next_phase: int | IntEnum, + caller_name: str = "", + ) -> None: + caller_tag = self._resolve_caller_tag(caller_name) + _host_debug( + f"[SimplePhaseChecker.update] caller={caller_name!r} " + f"caller_tag={caller_tag} " + f"expect={_phase_repr(expect_phase)} " + f"next={_phase_repr(next_phase)} " + f"capturing={torch.cuda.is_current_stream_capturing()}" + ) + _phase_check_kernel[(1,)]( + self._phase, + self._enable_assert_device, + EXPECT_PHASE=int(expect_phase), + NEXT_PHASE=int(next_phase), + CALLER_TAG=caller_tag, + ) + + def _reset_to_idle(self) -> None: + self._phase.fill_(self._initial_phase) + _host_debug( + f"[SimplePhaseChecker._reset_to_idle] phase reset to " + f"{self._initial_phase}" + ) + + def _resolve_caller_tag(self, caller_name: str) -> int: + registry = self._caller_tag_registry + if caller_name not in registry: + registry[caller_name] = len(registry) + 1 + _host_debug( + f"[SimplePhaseChecker] registered caller_tag " + f"{registry[caller_name]} <- {caller_name!r}" + ) + return registry[caller_name] diff --git a/test/registered/utils/test_phase_checker.py b/test/registered/utils/test_phase_checker.py new file mode 100644 index 000000000..03de94d33 --- /dev/null +++ b/test/registered/utils/test_phase_checker.py @@ -0,0 +1,483 @@ +"""Unit tests for srt/utils/phase_checker.py — SimplePhaseChecker.""" + +from __future__ import annotations + +import subprocess +import sys +import textwrap +import unittest +from enum import IntEnum + +import torch + +from sglang.srt.utils.phase_checker import SimplePhaseChecker +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=120, stage="base-b", runner_config="1-gpu-small") + + +class _Phase(IntEnum): + IDLE = 0 + A = 1 + B = 2 + C = 3 + + +def _phase_value(checker: SimplePhaseChecker) -> int: + """Read the device-side phase tensor as a host int.""" + return int(checker._phase.item()) + + +def _assert_flag(checker: SimplePhaseChecker) -> int: + """Read the device-side assert-enable flag as a host int.""" + return int(checker._enable_assert_device.item()) + + +class TestConstruction(CustomTestCase): + """Constructor sets phase and leaves assert OFF so init-time work is tolerated.""" + + @classmethod + def setUpClass(cls) -> None: + cls.device = torch.device("cuda:0") + + def test_init_stores_initial_phase_int(self) -> None: + checker = SimplePhaseChecker(initial_phase=7, device=self.device) + self.assertEqual(_phase_value(checker), 7) + + def test_init_stores_initial_phase_intenum(self) -> None: + checker = SimplePhaseChecker(initial_phase=_Phase.B, device=self.device) + self.assertEqual(_phase_value(checker), int(_Phase.B)) + + def test_init_assert_flag_is_off(self) -> None: + """assert flag must default to OFF so warmup / cuda-graph-capture don't trip.""" + checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device) + self.assertEqual(_assert_flag(checker), 0) + + def test_init_caller_registry_is_empty(self) -> None: + checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device) + self.assertEqual(checker._caller_tag_registry, {}) + + def test_init_phase_tensor_is_on_requested_device(self) -> None: + checker = SimplePhaseChecker(initial_phase=0, device=self.device) + self.assertEqual(checker._phase.device, self.device) + self.assertEqual(checker._enable_assert_device.device, self.device) + self.assertEqual(checker._phase.dtype, torch.int32) + self.assertEqual(checker._enable_assert_device.dtype, torch.int32) + + +class TestUpdateAssertDisabled(CustomTestCase): + """With the assert flag OFF, update always advances the phase and never raises.""" + + @classmethod + def setUpClass(cls) -> None: + cls.device = torch.device("cuda:0") + + def test_update_advances_phase_on_match(self) -> None: + checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device) + checker.update(expect_phase=_Phase.IDLE, next_phase=_Phase.A, caller_name="t") + torch.cuda.synchronize() + self.assertEqual(_phase_value(checker), int(_Phase.A)) + + def test_update_advances_phase_on_mismatch(self) -> None: + """assert OFF tolerates mismatches — store still happens unconditionally.""" + checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device) + checker.update(expect_phase=_Phase.C, next_phase=_Phase.B, caller_name="t") + torch.cuda.synchronize() + self.assertEqual(_phase_value(checker), int(_Phase.B)) + + def test_init_time_lifecycle_violations_tolerated(self) -> None: + """Documented use case: pre-enable_assert work may freely violate the lifecycle.""" + checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device) + # Bogus sequence — would all mismatch if assert were on. + checker.update(expect_phase=_Phase.B, next_phase=_Phase.C, caller_name="warmup") + checker.update( + expect_phase=_Phase.IDLE, next_phase=_Phase.A, caller_name="warmup" + ) + checker.update( + expect_phase=_Phase.B, next_phase=_Phase.IDLE, caller_name="warmup" + ) + torch.cuda.synchronize() # no raise + self.assertEqual(_phase_value(checker), int(_Phase.IDLE)) + + +class TestUpdateAssertEnabled(CustomTestCase): + """With assert ON, matching updates pass and walk a multi-state lifecycle.""" + + @classmethod + def setUpClass(cls) -> None: + cls.device = torch.device("cuda:0") + + def test_update_advances_phase_on_match(self) -> None: + checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device) + checker.enable_assert() + checker.update(expect_phase=_Phase.IDLE, next_phase=_Phase.A, caller_name="t") + torch.cuda.synchronize() + self.assertEqual(_phase_value(checker), int(_Phase.A)) + + def test_full_4_state_lifecycle_round_trip(self) -> None: + """IDLE -> A -> B -> C -> IDLE under assert ON, no synchronization mid-cycle.""" + checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device) + checker.enable_assert() + for _ in range(3): # multiple cycles to exercise stable behavior + checker.update( + expect_phase=_Phase.IDLE, next_phase=_Phase.A, caller_name="p1" + ) + checker.update(expect_phase=_Phase.A, next_phase=_Phase.B, caller_name="p2") + checker.update(expect_phase=_Phase.B, next_phase=_Phase.C, caller_name="p3") + checker.update( + expect_phase=_Phase.C, next_phase=_Phase.IDLE, caller_name="p4" + ) + torch.cuda.synchronize() + self.assertEqual(_phase_value(checker), int(_Phase.IDLE)) + + def test_update_mismatch_after_enable_raises_in_subprocess(self) -> None: + """A mismatched update with assert ON must fire device_assert at the next sync. + + Run in a subprocess because device-side asserts poison the CUDA context. + """ + script = textwrap.dedent(""" + import sys + + import torch + + from sglang.srt.utils.phase_checker import SimplePhaseChecker + + device = torch.device("cuda:0") + checker = SimplePhaseChecker(initial_phase=0, device=device) + checker.enable_assert() + # phase=0 but we claim expect=99 — kernel must fire device_assert. + checker.update(expect_phase=99, next_phase=1, caller_name="bad") + try: + torch.cuda.synchronize() + except RuntimeError as e: + msg = str(e).lower() + if "device-side assert" in msg or "phase mismatch" in msg: + sys.exit(0) + print(f"Unexpected RuntimeError: {e}", file=sys.stderr) + sys.exit(2) + print("expected RuntimeError but none was raised", file=sys.stderr) + sys.exit(1) + """) + result = subprocess.run( + [sys.executable, "-c", script], + capture_output=True, + text=True, + timeout=180, + ) + # Accept both exit-0 (RuntimeError caught) and SIGABRT (-6, the kernel's + # device_assert killed the process directly before sync could raise). + # Either way, the kernel-side check fired — which is what we're verifying. + # The presence of the SimplePhaseChecker FAIL line in stdout confirms it. + self.assertIn( + "SimplePhaseChecker FAIL", + result.stdout, + f"stdout: {result.stdout}\nstderr: {result.stderr}", + ) + self.assertIn( + result.returncode, + (0, -6), + f"unexpected returncode {result.returncode}; " + f"stdout: {result.stdout}\nstderr: {result.stderr}", + ) + + +class TestEnableAssert(CustomTestCase): + """enable_assert flips the device flag and resets phase to the initial value.""" + + @classmethod + def setUpClass(cls) -> None: + cls.device = torch.device("cuda:0") + + def test_enable_assert_sets_flag_to_one(self) -> None: + checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device) + self.assertEqual(_assert_flag(checker), 0) + checker.enable_assert() + self.assertEqual(_assert_flag(checker), 1) + + def test_enable_assert_resets_phase_to_initial(self) -> None: + """Phase advanced during init must be wiped by enable_assert. + + Otherwise post-init lifecycle would start from whatever captured kernels + left in the phase tensor — the explicit purpose of enable_assert. + """ + checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device) + # Mess with the phase as init-time work might. + checker.update( + expect_phase=_Phase.IDLE, next_phase=_Phase.C, caller_name="warmup" + ) + torch.cuda.synchronize() + self.assertEqual(_phase_value(checker), int(_Phase.C)) + + checker.enable_assert() + self.assertEqual(_phase_value(checker), int(_Phase.IDLE)) + + def test_enable_assert_with_nonzero_initial_phase(self) -> None: + """Reset target tracks the original initial_phase, not 0.""" + checker = SimplePhaseChecker(initial_phase=42, device=self.device) + checker.update(expect_phase=42, next_phase=7, caller_name="t") + torch.cuda.synchronize() + self.assertEqual(_phase_value(checker), 7) + + checker.enable_assert() + self.assertEqual(_phase_value(checker), 42) + + def test_enable_assert_is_idempotent(self) -> None: + checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device) + checker.enable_assert() + checker.enable_assert() + self.assertEqual(_assert_flag(checker), 1) + self.assertEqual(_phase_value(checker), int(_Phase.IDLE)) + + +class TestResetToIdle(CustomTestCase): + """_reset_to_idle only touches the phase tensor; the assert flag is untouched.""" + + @classmethod + def setUpClass(cls) -> None: + cls.device = torch.device("cuda:0") + + def test_reset_after_update_restores_initial_phase(self) -> None: + checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device) + checker.update(expect_phase=_Phase.IDLE, next_phase=_Phase.B, caller_name="t") + torch.cuda.synchronize() + self.assertEqual(_phase_value(checker), int(_Phase.B)) + + checker._reset_to_idle() + self.assertEqual(_phase_value(checker), int(_Phase.IDLE)) + + def test_reset_does_not_toggle_assert_flag_when_off(self) -> None: + checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device) + self.assertEqual(_assert_flag(checker), 0) + checker._reset_to_idle() + self.assertEqual(_assert_flag(checker), 0) + + def test_reset_does_not_toggle_assert_flag_when_on(self) -> None: + checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device) + checker.enable_assert() + self.assertEqual(_assert_flag(checker), 1) + checker._reset_to_idle() + self.assertEqual(_assert_flag(checker), 1) + + def test_reset_with_nonzero_initial_phase(self) -> None: + checker = SimplePhaseChecker(initial_phase=5, device=self.device) + checker.update(expect_phase=5, next_phase=9, caller_name="t") + torch.cuda.synchronize() + checker._reset_to_idle() + self.assertEqual(_phase_value(checker), 5) + + +class TestCallerTagRegistry(CustomTestCase): + """Caller tags are interned 1-indexed integers, assigned in first-seen order.""" + + @classmethod + def setUpClass(cls) -> None: + cls.device = torch.device("cuda:0") + + def test_first_caller_gets_tag_one(self) -> None: + checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device) + self.assertEqual(checker._resolve_caller_tag("first"), 1) + + def test_distinct_callers_get_distinct_increasing_tags(self) -> None: + checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device) + self.assertEqual(checker._resolve_caller_tag("a"), 1) + self.assertEqual(checker._resolve_caller_tag("b"), 2) + self.assertEqual(checker._resolve_caller_tag("c"), 3) + + def test_repeated_caller_returns_same_tag(self) -> None: + checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device) + first = checker._resolve_caller_tag("repeat") + # Mix other callers in between. + checker._resolve_caller_tag("other") + second = checker._resolve_caller_tag("repeat") + self.assertEqual(first, second) + + def test_empty_caller_name_is_interned_like_any_other(self) -> None: + """default caller_name="" must be a legal key, not crash on use.""" + checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device) + checker.enable_assert() + checker.update(expect_phase=_Phase.IDLE, next_phase=_Phase.A) # caller_name="" + torch.cuda.synchronize() + self.assertIn("", checker._caller_tag_registry) + self.assertEqual(_phase_value(checker), int(_Phase.A)) + + def test_update_populates_registry(self) -> None: + checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device) + checker.update( + expect_phase=_Phase.IDLE, next_phase=_Phase.A, caller_name="alpha" + ) + checker.update( + expect_phase=_Phase.A, next_phase=_Phase.IDLE, caller_name="beta" + ) + torch.cuda.synchronize() + self.assertEqual(checker._caller_tag_registry, {"alpha": 1, "beta": 2}) + + +class TestMultipleInstances(CustomTestCase): + """Independent checkers must not share phase, assert flag, or caller registry.""" + + @classmethod + def setUpClass(cls) -> None: + cls.device = torch.device("cuda:0") + + def test_phase_tensors_are_independent(self) -> None: + a = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device) + b = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device) + a.update(expect_phase=_Phase.IDLE, next_phase=_Phase.B, caller_name="a") + torch.cuda.synchronize() + self.assertEqual(_phase_value(a), int(_Phase.B)) + self.assertEqual(_phase_value(b), int(_Phase.IDLE)) + + def test_assert_flags_are_independent(self) -> None: + a = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device) + b = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device) + a.enable_assert() + self.assertEqual(_assert_flag(a), 1) + self.assertEqual(_assert_flag(b), 0) + + def test_caller_registries_are_independent(self) -> None: + a = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device) + b = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device) + a._resolve_caller_tag("shared_name") + # Same string, different checker — should still get tag=1, not 2. + self.assertEqual(b._resolve_caller_tag("shared_name"), 1) + + +class TestCudaGraphCapture(CustomTestCase): + """The kernel is launched unconditionally so it is capture-safe; the device flag + decides at replay time whether the assert fires. + """ + + @classmethod + def setUpClass(cls) -> None: + cls.device = torch.device("cuda:0") + + def _capture_one_update( + self, + checker: SimplePhaseChecker, + *, + expect_phase: int, + next_phase: int, + caller_name: str, + ) -> torch.cuda.CUDAGraph: + """Warm up + capture a single update() call into a fresh CUDAGraph. + + Both warmup and capture launches actually execute on the GPU and advance + the phase tensor. Caller must keep assert OFF here — otherwise the second + launch (capture) would see phase=next_phase and mismatch the expected. + """ + assert _assert_flag(checker) == 0, "capture helper requires assert OFF" + stream = torch.cuda.Stream(self.device) + # Warmup launch on a side stream — required before torch.cuda.graph. + stream.wait_stream(torch.cuda.current_stream(self.device)) + with torch.cuda.stream(stream): + checker.update( + expect_phase=expect_phase, + next_phase=next_phase, + caller_name=caller_name, + ) + torch.cuda.current_stream(self.device).wait_stream(stream) + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + checker.update( + expect_phase=expect_phase, + next_phase=next_phase, + caller_name=caller_name, + ) + return graph + + def test_captured_update_advances_phase_on_replay(self) -> None: + """Replay reissues the captured kernel; phase tensor is updated each replay.""" + checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device) + # Capture under assert OFF (the documented init-time pattern). + graph = self._capture_one_update( + checker, + expect_phase=int(_Phase.IDLE), + next_phase=int(_Phase.B), + caller_name="captured", + ) + + # Enable assert (resets phase -> IDLE) and replay — captured expect=IDLE matches. + checker.enable_assert() + graph.replay() + torch.cuda.synchronize() + self.assertEqual(_phase_value(checker), int(_Phase.B)) + + # Reset + replay again — same result, no raise. + checker._reset_to_idle() + graph.replay() + torch.cuda.synchronize() + self.assertEqual(_phase_value(checker), int(_Phase.B)) + + def test_assert_flag_toggle_visible_to_replayed_graph(self) -> None: + """Key contract: capture with assert OFF, enable assert later, replay sees the new flag. + + We don't try to make the replay raise (that would poison the context); we + verify the *positive* side — capture under assert OFF (with deliberate + phase mismatches tolerated) and confirm the same captured kernel still + works correctly when assert is later turned on. + """ + checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device) + graph = self._capture_one_update( + checker, + expect_phase=int(_Phase.IDLE), + next_phase=int(_Phase.A), + caller_name="captured", + ) + self.assertEqual(_assert_flag(checker), 0) + + # Replay with assert OFF tolerates a deliberately diverged phase. + checker._phase.fill_(999) + graph.replay() + torch.cuda.synchronize() # no raise — flag is OFF + self.assertEqual(_phase_value(checker), int(_Phase.A)) + + # Now turn on asserts (also resets phase -> IDLE) and replay. + checker.enable_assert() + self.assertEqual(_assert_flag(checker), 1) + self.assertEqual(_phase_value(checker), int(_Phase.IDLE)) + + graph.replay() + torch.cuda.synchronize() # no raise — phase matched expect + self.assertEqual(_phase_value(checker), int(_Phase.A)) + + +class TestPhaseReprNoCrash(CustomTestCase): + """Debug logging path (host_debug) tolerates both int and IntEnum phase args.""" + + @classmethod + def setUpClass(cls) -> None: + cls.device = torch.device("cuda:0") + + def test_update_with_int_phases_does_not_crash(self) -> None: + checker = SimplePhaseChecker(initial_phase=0, device=self.device) + checker.enable_assert() + checker.update(expect_phase=0, next_phase=1, caller_name="ints") + torch.cuda.synchronize() + self.assertEqual(_phase_value(checker), 1) + + def test_update_with_intenum_phases_does_not_crash(self) -> None: + checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device) + checker.enable_assert() + checker.update( + expect_phase=_Phase.IDLE, next_phase=_Phase.A, caller_name="enums" + ) + torch.cuda.synchronize() + self.assertEqual(_phase_value(checker), int(_Phase.A)) + + def test_update_mixing_int_and_intenum_phases(self) -> None: + """expect_phase is an IntEnum, next_phase is a plain int (and vice versa).""" + checker = SimplePhaseChecker(initial_phase=_Phase.IDLE, device=self.device) + checker.enable_assert() + checker.update(expect_phase=_Phase.IDLE, next_phase=5, caller_name="mix1") + torch.cuda.synchronize() + self.assertEqual(_phase_value(checker), 5) + checker.update(expect_phase=5, next_phase=_Phase.IDLE, caller_name="mix2") + torch.cuda.synchronize() + self.assertEqual(_phase_value(checker), int(_Phase.IDLE)) + + +if __name__ == "__main__": + unittest.main()