sgl-project--sglang/test/manual/chunked_prefill/test_scripted_abort.py

import unittest

from sglang.srt.environ import envs
from sglang.test.scripted_runtime.context import ScriptedContext
from sglang.test.scripted_runtime.req_handle import ScriptedReqHandle
from sglang.test.scripted_runtime.test_case import ScriptedTestCase
from sglang.test.scripted_runtime_chunked_helpers import (
    DEFAULT_CHUNK_SIZE,
    DEFAULT_MAX_STEPS,
    SMALL_KV_POOL_BALLAST_MAX_NEW_TOKENS,
    SMALL_KV_POOL_BALLAST_PROMPT_LEN,
    SMALL_KV_POOL_MAX_TOTAL_TOKENS,
    VERY_LONG_PROMPT_LEN,
    base_engine_kwargs,
    run_until,
    run_until_all_finished,
    run_until_finished,
)


def _drain_until_released(t: ScriptedContext, *handles: ScriptedReqHandle):
    for _ in range(12):
        if all(
            h.kv_pages == 0
            and h.lock_refs == 0
            and (h.req is None or h.req.req_pool_idx is None)
            for h in handles
        ):
            return
        yield


class TestAbortBasic(ScriptedTestCase):
    ENGINE_KWARGS = base_engine_kwargs(chunked_prefill_size=DEFAULT_CHUNK_SIZE)

    def test_abort_waiting_chunked_resume(self):
        self.server.execute_script(self._script_abort_waiting_chunked_resume)

    @staticmethod
    def _script_abort_waiting_chunked_resume(t: ScriptedContext):
        r = t.start_req(
            prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=100
        )
        yield from run_until(r, lambda h: h.is_chunking)

        pages_before = r.kv_pages
        assert pages_before > 0, "chunked req should own KV pages mid-chunk"

        t.abort(r)
        yield from _drain_until_released(t, r)

        assert r.status in (
            "finished",
            "unknown",
        ), f"after abort r should be finished/unknown, got {r.status}"
        assert (
            r.kv_pages == 0
        ), f"abort must release KV; r.kv_pages={r.kv_pages} after abort"
        assert (
            r.req is None or r.req.req_pool_idx is None
        ), f"abort must release row; r.req={r.req} after abort"
        assert (
            r.lock_refs == 0
        ), f"abort must release lock_refs; r.lock_refs={r.lock_refs}"

    def test_abort_at_chunk_0(self):
        self.server.execute_script(self._script_abort_at_chunk_0)

    @staticmethod
    def _script_abort_at_chunk_0(t: ScriptedContext):
        r = t.start_req(
            prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=110
        )
        yield
        yield from run_until(r, lambda h: h.is_chunking)

        t.abort(r)
        yield from _drain_until_released(t, r)

        assert r.kv_pages == 0
        assert r.req is None or r.req.req_pool_idx is None

    def test_abort_at_chunk_mid(self):
        self.server.execute_script(self._script_abort_at_chunk_mid)

    @staticmethod
    def _script_abort_at_chunk_mid(t: ScriptedContext):
        r = t.start_req(
            prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=120
        )
        yield from run_until(r, lambda h: h.chunks_done >= 2 and h.is_chunking)

        t.abort(r)
        yield from _drain_until_released(t, r)

        assert r.kv_pages == 0

    def test_abort_one_does_not_disturb_other(self):
        self.server.execute_script(self._script_abort_one_does_not_disturb_other)

    @staticmethod
    def _script_abort_one_does_not_disturb_other(t: ScriptedContext):
        r1 = t.start_req(
            prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=130
        )
        r2 = t.start_req(
            prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=131
        )
        yield from run_until(r1, lambda h: h.is_chunking)

        t.abort(r1)
        yield from _drain_until_released(t, r1)

        assert r1.kv_pages == 0
        yield from run_until_finished(r2)
        assert r2.finished, "r2 should still complete after r1 is aborted"

    def test_abort_with_zero_yield(self):
        self.server.execute_script(self._script_abort_with_zero_yield)

    @staticmethod
    def _script_abort_with_zero_yield(t: ScriptedContext):
        r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2)
        t.abort(r)
        yield from _drain_until_released(t, r)
        assert r.kv_pages == 0
        assert r.req is None or r.req.req_pool_idx is None
        assert r.lock_refs == 0

    def test_abort_at_admission_step(self):
        self.server.execute_script(self._script_abort_at_admission_step)

    @staticmethod
    def _script_abort_at_admission_step(t: ScriptedContext):
        r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2)
        yield
        t.abort(r)
        yield from _drain_until_released(t, r)
        assert r.kv_pages == 0
        assert r.req is None or r.req.req_pool_idx is None

    def test_abort_then_start_same_step_new_rid(self):
        self.server.execute_script(self._script_abort_then_start_same_step_new_rid)

    @staticmethod
    def _script_abort_then_start_same_step_new_rid(t: ScriptedContext):
        r1 = t.start_req(
            prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=140
        )
        yield from run_until(r1, lambda h: h.is_chunking)
        t.abort(r1)
        r2 = t.start_req(
            prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=141
        )
        yield from run_until_finished(r2)
        assert r2.finished
        assert r1.kv_pages == 0

    def test_abort_then_start_same_step_same_rid(self):
        self.server.execute_script(self._script_abort_then_start_same_step_same_rid)

    @staticmethod
    def _script_abort_then_start_same_step_same_rid(t: ScriptedContext):
        r1 = t.start_req(
            prompt_len=VERY_LONG_PROMPT_LEN,
            max_new_tokens=2,
            rid="abort-reuse",
            prompt_token=150,
        )
        yield from run_until(r1, lambda h: h.is_chunking)
        t.abort(r1)
        yield
        r2 = t.start_req(prompt_len=16, max_new_tokens=2, rid="abort-reuse")
        yield from run_until_finished(r2)
        assert r2.finished

    def test_abort_five_chunked_in_a_row(self):
        self.server.execute_script(self._script_abort_five_chunked_in_a_row)

    @staticmethod
    def _script_abort_five_chunked_in_a_row(t: ScriptedContext):
        reqs = [
            t.start_req(
                prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=160 + i
            )
            for i in range(5)
        ]
        yield from run_until(reqs[0], lambda h: h.is_chunking)
        for r in reqs:
            t.abort(r)
        yield from _drain_until_released(t, *reqs)
        for r in reqs:
            assert r.kv_pages == 0
            assert r.req is None or r.req.req_pool_idx is None

    def test_abort_unknown_rid_noop(self):
        self.server.execute_script(self._script_abort_unknown_rid_noop)

    @staticmethod
    def _script_abort_unknown_rid_noop(t: ScriptedContext):
        bogus = ScriptedReqHandle(rid="never-submitted-rid", context=t)
        t.abort(bogus, await_arrival=False)
        yield
        r = t.start_req(prompt_len=16, max_new_tokens=2)
        yield from run_until_finished(r)
        assert r.finished
        for _ in range(12):
            if r.kv_pages == 0 and r.lock_refs == 0:
                break
            yield
        assert r.kv_pages == 0
        assert r.lock_refs == 0

    def test_abort_after_finish_noop(self):
        self.server.execute_script(self._script_abort_after_finish_noop)

    @staticmethod
    def _script_abort_after_finish_noop(t: ScriptedContext):
        r = t.start_req(prompt_len=16, max_new_tokens=2)
        yield from run_until_finished(r)
        assert r.finished

        for _ in range(12):
            if t.is_fully_idle:
                break
            yield
        assert r.kv_pages == 0
        assert r.lock_refs == 0
        kv_pool_free_before = t.engine_stats()["kv_pool_free"]

        t.abort(r, await_arrival=False)
        yield
        assert r.kv_pages == 0
        assert r.lock_refs == 0
        kv_pool_free_after = t.engine_stats()["kv_pool_free"]
        assert kv_pool_free_after == kv_pool_free_before, (
            f"abort-after-finish must not move KV pool; "
            f"before={kv_pool_free_before} after={kv_pool_free_after}"
        )

    def test_abort_chunk_last(self):
        self.server.execute_script(self._script_abort_chunk_last)

    @staticmethod
    def _script_abort_chunk_last(t: ScriptedContext):
        r = t.start_req(
            prompt_len=2 * DEFAULT_CHUNK_SIZE, max_new_tokens=4, prompt_token=170
        )
        yield from run_until(r, lambda h: h.chunks_done >= 1 and h.is_chunking)
        t.abort(r)
        yield from _drain_until_released(t, r)
        assert r.kv_pages == 0
        assert r.req is None or r.req.inflight_middle_chunks == 0

    def test_abort_penultimate_chunk(self):
        self.server.execute_script(self._script_abort_penultimate_chunk)

    @staticmethod
    def _script_abort_penultimate_chunk(t: ScriptedContext):
        r = t.start_req(
            prompt_len=4 * DEFAULT_CHUNK_SIZE, max_new_tokens=2, prompt_token=180
        )
        yield from run_until(r, lambda h: h.chunks_done >= 2 and h.is_chunking)
        t.abort(r)
        yield from _drain_until_released(t, r)
        assert r.kv_pages == 0
        assert r.req is None or r.req.req_pool_idx is None
        assert r.lock_refs == 0

    def test_double_abort_idempotent(self):
        self.server.execute_script(self._script_double_abort_idempotent)

    @staticmethod
    def _script_double_abort_idempotent(t: ScriptedContext):
        r = t.start_req(
            prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=190
        )
        yield from run_until(r, lambda h: h.is_chunking)
        t.abort(r)
        t.abort(r)
        yield from _drain_until_released(t, r)
        assert r.kv_pages == 0
        assert r.lock_refs == 0

    def test_abort_during_decode(self):
        self.server.execute_script(self._script_abort_during_decode)

    @staticmethod
    def _script_abort_during_decode(t: ScriptedContext):
        r = t.start_req(prompt_len=16, max_new_tokens=64)
        yield from run_until(r, lambda h: h.status == "running")
        assert r.kv_pages > 0, "decode req must own KV before abort"
        t.abort(r)
        yield from _drain_until_released(t, r)
        assert r.kv_pages == 0
        assert r.lock_refs == 0

    def test_abort_one_of_three_others_finish(self):
        self.server.execute_script(self._script_abort_one_of_three_others_finish)

    @staticmethod
    def _script_abort_one_of_three_others_finish(t: ScriptedContext):
        r1 = t.start_req(prompt_len=16, max_new_tokens=4)
        r2 = t.start_req(prompt_len=16, max_new_tokens=4)
        r3 = t.start_req(prompt_len=16, max_new_tokens=4)
        yield from run_until(r2, lambda h: h.status == "running")
        t.abort(r2)
        yield from run_until_all_finished([r1, r3])
        assert r2.kv_pages == 0

    def test_abort_in_separate_yields(self):
        self.server.execute_script(self._script_abort_in_separate_yields)

    @staticmethod
    def _script_abort_in_separate_yields(t: ScriptedContext):
        reqs = [
            t.start_req(
                prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=200 + i
            )
            for i in range(3)
        ]
        yield from run_until(reqs[0], lambda h: h.is_chunking)
        for r in reqs:
            t.abort(r)
            yield
        yield from _drain_until_released(t, *reqs)
        for r in reqs:
            assert r.kv_pages == 0

    def test_abort_at_chunk_boundary_race(self):
        self.server.execute_script(self._script_abort_at_chunk_boundary_race)

    @staticmethod
    def _script_abort_at_chunk_boundary_race(t: ScriptedContext):
        r = t.start_req(
            prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=210
        )
        yield from run_until(r, lambda h: h.is_chunking)
        yield from run_until(r, lambda h: h.chunks_done >= 1 and h.is_chunking)

        t.abort(r)
        yield from _drain_until_released(t, r)

        assert r.kv_pages == 0

        chunks_after_abort = r.chunks_done
        yield
        yield
        assert not r.is_chunking, "aborted req must not resume chunking"
        assert r.chunks_done == chunks_after_abort, (
            f"aborted req revived and ran another chunk; "
            f"chunks_done went {chunks_after_abort} -> {r.chunks_done}"
        )
        assert r.req is None or r.req.req_pool_idx is None

    def test_abort_mid_chunk_no_extra_radix_node(self):
        self.server.execute_script(self._script_abort_mid_chunk_no_extra_radix_node)

    @staticmethod
    def _script_abort_mid_chunk_no_extra_radix_node(
        t: ScriptedContext,
    ):
        r = t.start_req(
            prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=220
        )
        yield from run_until(r, lambda h: h.is_chunking)

        t.abort(r)
        yield from _drain_until_released(t, r)

        assert r.kv_pages == 0
        assert r.req is None or r.req.req_pool_idx is None
        chunks_after_release = r.chunks_done
        for _ in range(4):
            yield
        assert r.chunks_done == chunks_after_release, (
            f"aborted mid-chunk req revived after release; chunks_done went "
            f"{chunks_after_release} -> {r.chunks_done}"
        )

        lock_refs_after = t.get_all_node_lock_refs()
        assert all(ref == 0 for ref in lock_refs_after.values()), (
            f"abort mid-chunk left a locked radix node behind; "
            f"node lock_refs={lock_refs_after!r}"
        )

    def test_abort_then_resubmit_same_rid_same_step(self):
        self.server.execute_script(self._script_abort_then_resubmit_same_rid_same_step)

    @staticmethod
    def _script_abort_then_resubmit_same_rid_same_step(t: ScriptedContext):
        r1 = t.start_req(
            prompt_len=VERY_LONG_PROMPT_LEN,
            max_new_tokens=2,
            rid="abort-resubmit-same-step",
            prompt_token=230,
        )
        yield from run_until(r1, lambda h: h.is_chunking)
        t.abort(r1)
        r2 = t.start_req(
            prompt_len=16,
            max_new_tokens=2,
            rid="abort-resubmit-same-step",
        )
        yield
        yield from run_until_finished(r2)
        assert r2.finished, "resubmit under same rid must complete independently"
        assert r1.kv_pages == 0, "aborted r1 must release KV before resubmit"
        assert r1.req is None or r1.req.req_pool_idx is None
        assert r1.lock_refs == 0

    def test_abort_during_gap_inflight_middle_chunks_positive(self):
        self.server.execute_script(
            self._script_abort_during_gap_inflight_middle_chunks_positive
        )

    @staticmethod
    def _script_abort_during_gap_inflight_middle_chunks_positive(t: ScriptedContext):
        r = t.start_req(
            prompt_len=2 * DEFAULT_CHUNK_SIZE, max_new_tokens=2, prompt_token=240
        )
        yield from run_until(
            r,
            lambda h: h.is_chunking and h.chunks_done >= 1,
        )
        assert r.req.inflight_middle_chunks > 0

        t.abort(r)
        yield from _drain_until_released(t, r)

        assert r.kv_pages == 0
        assert r.req is None or r.req.req_pool_idx is None

        assert not r.is_chunking, "aborted gap req must not re-enter chunking"
        yield
        assert not r.is_chunking, "aborted gap req must stay out of chunking"

        if r.req is not None:
            assert (
                r.req.inflight_middle_chunks == 0
            ), f"inflight_middle_chunks not cleared; got {r.req.inflight_middle_chunks}"

    def test_abort_when_chunked_only_then_idle(self):
        self.server.execute_script(self._script_abort_when_chunked_only_then_idle)

    @staticmethod
    def _script_abort_when_chunked_only_then_idle(t: ScriptedContext):
        r = t.start_req(
            prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=250
        )
        yield from run_until(r, lambda h: h.is_chunking)
        assert (1 if t.scheduler.chunked_req is not None else 0) == 1

        t.abort(r)
        yield from _drain_until_released(t, r)

        for _ in range(12):
            if t.scheduler.chunked_req is None and t.is_idle:
                break
            yield

        assert r.kv_pages == 0
        assert (1 if t.scheduler.chunked_req is not None else 0) == 0
        assert t.is_idle, "engine must be idle after the only chunked req is aborted"

    def test_chunked_req_then_abort_then_new_short_in_one_yield(self):
        self.server.execute_script(
            self._script_chunked_req_then_abort_then_new_short_in_one_yield
        )

    @staticmethod
    def _script_chunked_req_then_abort_then_new_short_in_one_yield(t: ScriptedContext):
        r1 = t.start_req(
            prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=260
        )
        yield from run_until(r1, lambda h: h.is_chunking)
        assert (
            t.scheduler.chunked_req.rid if t.scheduler.chunked_req is not None else None
        ) == r1.rid, (
            f"r1 should hold the chunked slot before abort; got "
            f"{(t.scheduler.chunked_req.rid if t.scheduler.chunked_req is not None else None)!r}"
        )

        t.abort(r1)
        r2 = t.start_req(prompt_len=16, max_new_tokens=2)
        yield from _drain_until_released(t, r1)

        cur = (
            t.scheduler.chunked_req.rid if t.scheduler.chunked_req is not None else None
        )
        assert cur != r1.rid, f"chunked slot still points to aborted r1; got {cur!r}"
        assert r1.kv_pages == 0
        yield from run_until_finished(r2)
        assert r2.finished, "fresh r2 must admit and complete after combo step"

    def test_force_retract_then_abort_same_yield(self):
        self.server.execute_script(self._script_force_retract_then_abort_same_yield)

    @staticmethod
    def _script_force_retract_then_abort_same_yield(t: ScriptedContext):
        r1 = t.start_req(
            prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=270
        )
        yield from run_until(r1, lambda h: h.is_chunking)
        assert r1.kv_pages > 0

        t.pause_generation(mode="retract")
        t.abort(r1)
        yield from _drain_until_released(t, r1)

        assert r1.kv_pages == 0, (
            f"force_retract + abort same yield must release KV; got " f"{r1.kv_pages}"
        )
        assert r1.req is None or r1.req.req_pool_idx is None, (
            f"force_retract + abort same yield must release row; got " f"{r1.req}"
        )
        assert r1.lock_refs == 0, (
            f"force_retract + abort same yield must release lock_refs; "
            f"got {r1.lock_refs}"
        )
        yield
        t.continue_generation()

    def test_abort_chunked_with_baton_handoff(self):
        self.server.execute_script(self._script_abort_chunked_with_baton_handoff)

    @staticmethod
    def _script_abort_chunked_with_baton_handoff(t: ScriptedContext):
        r1 = t.start_req(
            prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=280
        )
        r2 = t.start_req(
            prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2, prompt_token=281
        )
        yield from run_until(r1, lambda h: h.is_chunking)
        assert (1 if t.scheduler.chunked_req is not None else 0) == 1

        t.abort(r1)
        yield from _drain_until_released(t, r1)

        yield from run_until(r2, lambda h: h.is_chunking)
        assert r1.kv_pages == 0
        assert r1.req is None or r1.req.req_pool_idx is None
        assert r1.lock_refs == 0
        yield from run_until_finished(r2)
        assert r2.finished, "baton handoff must let r2 complete"
        assert r2.lock_refs == 0


class TestAbortPP(ScriptedTestCase):
    ENGINE_KWARGS = base_engine_kwargs(
        chunked_prefill_size=DEFAULT_CHUNK_SIZE,
        pp_size=2,
    )

    def test_abort_at_last_chunk_in_flight_pp(self):
        self.server.execute_script(self._script_abort_at_last_chunk_in_flight_pp)

    @staticmethod
    def _script_abort_at_last_chunk_in_flight_pp(t: ScriptedContext):
        r = t.start_req(
            prompt_len=2 * DEFAULT_CHUNK_SIZE, max_new_tokens=4, prompt_token=290
        )
        yield from run_until(
            r,
            lambda h: h.chunks_done >= 1 and h.is_chunking,
        )

        t.abort(r)
        yield from _drain_until_released(t, r)

        assert r.kv_pages == 0
        assert r.req is None or r.req.req_pool_idx is None
        assert r.lock_refs == 0
        assert r.finished


class TestAbortSmallPool(ScriptedTestCase):
    ENGINE_KWARGS = base_engine_kwargs(
        chunked_prefill_size=DEFAULT_CHUNK_SIZE,
        max_total_tokens=SMALL_KV_POOL_MAX_TOTAL_TOKENS,
    )

    def test_waiting_timeout_sweep_aborts_pressured_waiting_req(self):
        self.server.execute_script(
            self._script_waiting_timeout_sweep_aborts_pressured_waiting_req
        )

    @staticmethod
    def _script_waiting_timeout_sweep_aborts_pressured_waiting_req(t: ScriptedContext):

        b1 = t.start_req(
            prompt_len=SMALL_KV_POOL_BALLAST_PROMPT_LEN,
            max_new_tokens=SMALL_KV_POOL_BALLAST_MAX_NEW_TOKENS,
            ignore_eos=True,
            prompt_token=300,
        )
        b2 = t.start_req(
            prompt_len=SMALL_KV_POOL_BALLAST_PROMPT_LEN,
            max_new_tokens=SMALL_KV_POOL_BALLAST_MAX_NEW_TOKENS,
            ignore_eos=True,
            prompt_token=301,
        )
        yield from run_until(b1, lambda h: h.status == "running")
        yield from run_until(b2, lambda h: h.status == "running")

        r = t.start_req(prompt_len=16, max_new_tokens=2)

        def waiting_rids():
            return {req.rid for req in t.scheduler.waiting_queue}

        yield from run_until(r, lambda h: r.rid in waiting_rids())
        assert r.kv_pages == 0, "pressured waiting req must not own KV before admission"

        with envs.SGLANG_REQ_WAITING_TIMEOUT.override(1e-6):
            for _ in range(DEFAULT_MAX_STEPS):
                if r.rid not in waiting_rids():
                    break
                yield
            else:
                raise AssertionError(
                    f"waiting-timeout sweep never removed the req from "
                    f"waiting_queue after {DEFAULT_MAX_STEPS} steps; "
                    f"waiting_rids={waiting_rids()!r}"
                )

        assert r.rid not in waiting_rids(), (
            f"the loop's waiting-timeout sweep must drop the timed-out waiting "
            f"req from waiting_queue; got {waiting_rids()!r}"
        )
        assert r.status in ("finished", "unknown"), (
            f"swept-out req must be aborted (gone from every live scheduler "
            f"structure); got status={r.status!r}"
        )
        assert r.kv_pages == 0, "timeout-abort of an unadmitted req owns no KV"

        t.abort(b1)
        t.abort(b2)
        yield from _drain_until_released(t, b1, b2)


if __name__ == "__main__":
    unittest.main()