Function test_run_engine_and_streaming

tests/entrypoints/test_llm.py:227–285

(monkeypatch)

def test_run_engine_and_streaming(monkeypatch):
    llm = _make_llm(_make_engine(vocab_size=3, enable_logprob=True))
    llm._build_sample_logprobs = lambda *_: [{1: None}]
    llm._build_prompt_logprobs = lambda *_: [None]

    top_logprobs = LogprobsLists(logprob_token_ids=[[1]], logprobs=[[-0.1]], sampled_token_ranks=[0])
    prompt_logprobs = LogprobsTensors(
        paddle.to_tensor([[1]], dtype=paddle.int64),
        paddle.to_tensor([[-0.1]], dtype=paddle.float32),
        paddle.to_tensor([1], dtype=paddle.int64),
    )
    result = DummyResult("r1", [1], top_logprobs=top_logprobs, prompt_logprobs=prompt_logprobs, finished=True)
    llm.req_output["r1"] = result

    class DummyTqdm:
        last_instance = None

        def __init__(self, **_):
            self.updated = 0
            self.closed = False
            DummyTqdm.last_instance = self

        def update(self, n):
            self.updated += n

        def close(self):
            self.closed = True
            return None

    llm_module.tqdm = DummyTqdm
    out = llm._run_engine(["r1"], use_tqdm=True, topk_logprobs=-1, num_prompt_logprobs=-1)
    assert out[0] is result

    current = DummyResult(
        "r2",
        [1, 2],
        top_logprobs=top_logprobs,
        prompt_logprobs=None,
        finished=True,
    )
    if "r2" in llm.req_output:
        llm.req_output.pop("r2")

    def fake_sleep(_):
        if "r2" not in llm.req_output:
            llm.req_output["r2"] = current
        return None

    monkeypatch.setattr(llm_module.time, "sleep", fake_sleep)
    it = llm._run_engine_stream(
        ["r2"],
        prompts=["hi"],
        use_tqdm=True,
        topk_logprobs=1,
        chat_template_kwargs={"enable_thinking": True},
    )
    batches = list(it)
    assert batches[0][0].prompt == "hi"