zhifu-platform/tests/test_embedding_batching.py

import asyncio

import numpy as np
import pytest
import torch
from fastapi import HTTPException
from prometheus_client import CollectorRegistry
from sentence_transformers import SentenceTransformer

from zhifu.modules.knowledge.batching import Batcher
from zhifu.modules.knowledge.config import KnowledgeSettings

pytestmark = pytest.mark.integration


@pytest.fixture(scope="module")
def model():
    settings = KnowledgeSettings()
    torch.set_num_threads(4)
    encoder = SentenceTransformer(
        settings.embedding_model,
        revision=settings.embedding_revision,
        cache_folder=settings.model_cache,
        device=settings.inference_device,
    )
    encoder.max_seq_length = 512
    return encoder


async def test_batch_preserves_request_and_text_order(model):
    registry = CollectorRegistry()
    batcher = Batcher(model, KnowledgeSettings(), registry)
    groups = [["键盘如何清洁"], ["订单发票申请", "包裹配送查询"], ["退换货申请"]]
    requests = [asyncio.create_task(batcher.submit(texts)) for texts in groups]
    await asyncio.sleep(0)
    task = asyncio.create_task(batcher.run())
    try:
        results = await asyncio.gather(*requests)
        expected = await asyncio.to_thread(
            model.encode,
            [text for group in groups for text in group],
            normalize_embeddings=True,
            show_progress_bar=False,
        )
        np.testing.assert_allclose(np.concatenate(results), expected, atol=1e-5)
        assert [len(result) for result in results] == [1, 2, 1]
        assert registry.get_sample_value("zhifu_embedding_batch_texts_count") == 1
        assert registry.get_sample_value("zhifu_embedding_batch_texts_sum") == 4
    finally:
        batcher.close()
        await task


async def test_queue_capacity_cancellation_and_shutdown(model):
    batcher = Batcher(model, KnowledgeSettings(queue_texts=32), CollectorRegistry())
    request = asyncio.create_task(batcher.submit(["物流查询"] * 32))
    await asyncio.sleep(0)
    with pytest.raises(HTTPException) as error:
        await batcher.submit(["发票查询"])
    assert error.value.status_code == 429
    request.cancel()
    with pytest.raises(asyncio.CancelledError):
        await request
    assert batcher.waiting == 0 and not batcher.pending
    waiting = asyncio.create_task(batcher.submit(["退换货查询"]))
    await asyncio.sleep(0)
    batcher.close()
    with pytest.raises(HTTPException) as error:
        await waiting
    assert error.value.status_code == 503
    assert batcher.waiting == 0
    with pytest.raises(HTTPException) as error:
        await batcher.submit(["发票查询"])
    assert error.value.status_code == 503


async def test_expired_queue_request_is_not_computed(model):
    registry = CollectorRegistry()
    settings = KnowledgeSettings(queue_wait_seconds=0.001, batch_wait_seconds=0.01)
    batcher = Batcher(model, settings, registry)
    request = asyncio.create_task(batcher.submit(["查询物流"]))
    await asyncio.sleep(0)
    task = asyncio.create_task(batcher.run())
    try:
        with pytest.raises(HTTPException) as error:
            await request
        assert error.value.status_code == 429
        assert registry.get_sample_value("zhifu_embedding_batch_texts_count") == 0
        assert batcher.waiting == 0
    finally:
        batcher.close()
        await task


async def test_cancelled_active_request_does_not_cancel_batch(model):
    batcher = Batcher(model, KnowledgeSettings(), CollectorRegistry())
    first = asyncio.create_task(batcher.submit(["查询包裹物流"]))
    second = asyncio.create_task(batcher.submit(["申请电子发票"]))
    task = asyncio.create_task(batcher.run())
    try:
        async with asyncio.timeout(10):
            await batcher.computing.wait()
        first.cancel()
        with pytest.raises(asyncio.CancelledError):
            await first
        result = await second
        assert len(result) == 1 and len(result[0]) == 1024
        assert not batcher.active and batcher.waiting == 0
    finally:
        batcher.close()
        await task


async def test_shutdown_finishes_active_computation(model):
    batcher = Batcher(model, KnowledgeSettings(), CollectorRegistry())
    request = asyncio.create_task(batcher.submit(["键盘清洁说明"] * 4))
    task = asyncio.create_task(batcher.run())
    try:
        async with asyncio.timeout(10):
            await batcher.computing.wait()
        batcher.close()
        with pytest.raises(HTTPException) as error:
            await request
        assert error.value.status_code == 503
    finally:
        batcher.close()
        await task
    assert not batcher.active and not batcher.accepting