freetoken/tests/scheduler/test_abort_inflight_prefill.py

"""Abort vs overlap-scheduling races.

Under overlap scheduling a batch launches one iteration before _process_last_data drains
it, and an abort message is processed in between. Freeing the request's resources inside
the abort handler while its forward is in flight used to corrupt state: the hybrid
prefill-commit dereferenced the None'd GDN ping-pong slots (TypeError killed the
scheduler); plain radix would silently re-read the freed page-table row.

The scheduler now uses the SGLang-style single-owner design: the abort handler frees
immediately ONLY when the request has no forward in flight (not in ``self._last_data``'s
batch); otherwise it just sets ``req.aborted`` and _process_last_data frees the request
when the batch drains, after copy_done.synchronize(). A ``table_idx != -1`` sentinel on
the prefix-commit remains as defense-in-depth.

Tests drive the real (unbound) Scheduler methods against CPU-built hybrid managers.
"""
from __future__ import annotations

from types import SimpleNamespace

import torch

from freetoken.core import Batch, Req, SamplingParams
from freetoken.kvcache.linear_state_pool import LinearStatePool
from freetoken.message import AbortBackendMsg
from freetoken.models.config import LinearGatedDeltaGroupConfig
from freetoken.scheduler.cache import CacheManager
from freetoken.scheduler.decode import DecodeManager
from freetoken.scheduler.prefill import ChunkedReq, PrefillManager
from freetoken.scheduler.scheduler import Scheduler
from freetoken.scheduler.table import TableManager
from freetoken.scheduler.utils import PendingReq

UID = 2


def _pool(num_slots=16):
    g = LinearGatedDeltaGroupConfig(
        name="linear", layer_ids=(0,), num_key_heads=2, num_value_heads=4,
        key_head_dim=16, value_head_dim=16, conv_kernel_dim=4, output_gate="silu",
    )
    return LinearStatePool(group=g, num_slots=num_slots, dtype=torch.bfloat16,
                           device=torch.device("cpu"), tp_size=1)


def _setup():
    """Hybrid managers + a stub Scheduler `self` for the real unbound methods."""
    pool = _pool()
    pt = torch.zeros(4, 64, dtype=torch.int32)
    cm = CacheManager(64, 1, pt, "hybrid_radix", linear_state_pool=pool)
    tm = TableManager(max_running_reqs=4, page_table=pt)
    dm = DecodeManager(page_size=1)
    pm = PrefillManager(cm, tm, dm)
    sent = []
    stub = SimpleNamespace(
        cache_manager=cm,
        table_manager=tm,
        decode_manager=dm,
        prefill_manager=pm,
        finished_reqs=set(),
        eos_token_ids=set(),
        toolcall_anchor_id=None,
        config=SimpleNamespace(page_size=1),
        status_reporter=SimpleNamespace(report_batch=lambda *_, **__: None),
        send_result=sent.extend,
        _kv_usage_pages=cm.page_usage,
        _mamba_slot_usage=lambda: None,
        _swa_token_usage=lambda: None,
        _gpu_mem_bytes=lambda: 0,
        _match_stop_str=lambda _req: None,
        _pending_abort_acks=set(),
        _last_data=None,
    )
    stub._free_req_resources = lambda req: Scheduler._free_req_resources(stub, req)
    return pool, cm, tm, dm, pm, sent, stub


def _launch_req(pool, cm, tm, prompt, *, cls=Req, track_seqlen=None):
    """A launched (forward in flight) hybrid req: handle locked, pages allocated,
    GDN slots held, cached_len advanced -- the state _process_last_data will drain."""
    mr = cm.match_req(SimpleNamespace(input_ids=prompt, input_len=len(prompt),
                                      mm_embeds=None))
    req = cls(input_ids=prompt, table_idx=tm.allocate(), cached_len=0, output_len=4,
              uid=UID, sampling_params=SamplingParams(max_tokens=4),
              cache_handle=mr.cuda_handle)
    req.linear_slot_idx = pool.alloc(1)[0]
    req.mamba_ping_pong = tuple(pool.alloc(2))
    req.mamba_next_track_idx = 1
    cm.lock(mr.cuda_handle)
    cm.allocate_paged([req])
    req.complete_one()
    req.mamba_last_track_seqlen = track_seqlen
    return req


def _as_last_data(batch):
    return (
        SimpleNamespace(batch=batch),
        (None, torch.tensor([42], dtype=torch.int32),
         SimpleNamespace(synchronize=lambda: None)),
    )


def test_abort_inflight_final_chunk_marks_then_drains():
    """Abort while the final prefill chunk (plain Req, already in running_reqs) is in
    flight: the handler only marks; the same iteration's drain frees exactly once."""
    pool, cm, tm, dm, _pm, sent, stub = _setup()
    req = _launch_req(pool, cm, tm, torch.arange(1, 13, dtype=torch.int32),
                      track_seqlen=8)
    batch = Batch(reqs=[req], phase="prefill")
    dm.filter_reqs(batch.reqs)                  # _forward: joins running_reqs at launch
    stub._last_data = _as_last_data(batch)      # overlap_loop exposes the un-drained batch

    Scheduler._process_one_msg(stub, AbortBackendMsg(uid=UID))
    assert req.aborted and req.table_idx != -1  # marked, NOT freed under the forward
    assert req.mamba_ping_pong is not None
    assert req not in dm.running_reqs
    assert UID in stub._pending_abort_acks
    free_after_mark = pool.num_free_slots

    Scheduler._process_last_data(stub, stub._last_data)
    assert req.table_idx == -1                  # freed at the drain point
    assert pool.num_free_slots > free_after_mark
    assert req in stub.finished_reqs
    assert sent == []                           # no DetokenizeMsg: abort ack stays terminal
    cm.check_integrity()


def test_abort_inflight_intermediate_chunk_marks_then_drains():
    """Abort mid-chunked-prefill: prefill_manager.abort_req pops the pending continuation
    (no next chunk launches) and returns the in-flight ChunkedReq, which is marked and
    freed when its batch drains."""
    pool, cm, tm, _dm, pm, sent, stub = _setup()
    prompt = torch.arange(1, 13, dtype=torch.int32)
    chunk = _launch_req(pool, cm, tm, prompt[:8], cls=ChunkedReq)
    pending = PendingReq(uid=UID, input_ids=prompt,
                         sampling_params=SamplingParams(max_tokens=4))
    pending.chunked_req = chunk
    pm.pending_list = [pending]
    batch = Batch(reqs=[chunk], phase="prefill")
    stub._last_data = _as_last_data(batch)

    Scheduler._process_one_msg(stub, AbortBackendMsg(uid=UID))
    assert pm.pending_list == []                # continuation gone: no next chunk
    assert chunk.aborted and chunk.table_idx != -1

    Scheduler._process_last_data(stub, stub._last_data)
    assert chunk.table_idx == -1
    assert sent == []                           # chunks never reply
    cm.check_integrity()


def test_abort_starved_decode_req_frees_immediately():
    """A request with no forward in flight (e.g. a decode req starved behind a long
    chunked prefill) is freed by the abort handler right away -- deferring would leak
    until its next batch, which strict prefill-priority puts arbitrarily far away."""
    pool, cm, tm, dm, _pm, _sent, stub = _setup()
    req = _launch_req(pool, cm, tm, torch.arange(1, 13, dtype=torch.int32))
    dm.filter_reqs([req])
    # the un-drained batch belongs to some other request's prefill
    stub._last_data = (SimpleNamespace(batch=SimpleNamespace(reqs=[])), None)
    base_free = pool.num_free_slots

    Scheduler._process_one_msg(stub, AbortBackendMsg(uid=UID))
    assert not req.aborted
    assert req.table_idx == -1                  # freed immediately, no drain needed
    assert pool.num_free_slots > base_free
    assert req not in dm.running_reqs
    cm.check_integrity()


def test_prefix_commit_sentinel_guard():
    """Defense-in-depth: even if some future path frees a req early (bypassing the
    aborted mark), the finished=False prefix-commit must skip a freed req instead of
    dereferencing its None'd GDN slots (the original crash, cache.py _cache_req_hybrid)."""
    pool, cm, tm, dm, _pm, sent, stub = _setup()
    req = _launch_req(pool, cm, tm, torch.arange(1, 13, dtype=torch.int32),
                      track_seqlen=8)
    batch = Batch(reqs=[req], phase="prefill")
    dm.filter_reqs(batch.reqs)

    aborted = dm.abort_req(UID)
    assert aborted is req
    Scheduler._free_req_resources(stub, aborted)   # freed WITHOUT the aborted mark
    assert req.table_idx == -1 and req.mamba_ping_pong is None
    free_after_abort = pool.num_free_slots

    Scheduler._process_last_data(stub, _as_last_data(batch))  # pre-guard: TypeError

    assert pool.num_free_slots == free_after_abort  # nothing double-freed
    cm.check_integrity()
    assert [m.uid for m in sent] == [UID]  # un-marked path still publishes the token


if __name__ == "__main__":
    for name, fn in list(globals().items()):
        if name.startswith("test_") and callable(fn):
            fn()
            print(f"{name}: PASS")


def test_post_terminal_overlap_step_is_dropped():
    """Overlap scheduling launches one more decode step for a request that already
    terminated (filter_reqs keeps it while output budget remains). The extra drain
    must not append its token, emit a second DetokenizeMsg, or free twice."""
    from freetoken.message import DetokenizeMsg

    pool, cm, tm, dm, _pm, sent, stub = _setup()
    stub.eos_token_ids = {42}  # the drained token (42) finishes the request by EOS
    req = _launch_req(pool, cm, tm, torch.arange(1, 13, dtype=torch.int32),
                      track_seqlen=8)
    dm.filter_reqs([req])

    Scheduler._process_last_data(stub, _as_last_data(Batch(reqs=[req], phase="prefill")))
    assert req in stub.finished_reqs and req.table_idx == -1
    terminal = [m for m in sent if isinstance(m, DetokenizeMsg)]
    assert len(terminal) == 1 and terminal[0].finished

    # The overlap extra step: the same req sits in the next batch's drain.
    output_len_before = req.output_len
    Scheduler._process_last_data(stub, _as_last_data(Batch(reqs=[req], phase="decode")))
    assert [m for m in sent if isinstance(m, DetokenizeMsg)] == terminal  # no 2nd msg
    assert req.output_len == output_len_before                           # no append
    cm.check_integrity()