test_api_server.py 2.9 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
import subprocess
import sys
import time
from multiprocessing import Pool
from pathlib import Path

import pytest
import requests


11
def _query_server(prompt: str, max_tokens: int = 5) -> dict:
12
13
14
    response = requests.post("http://localhost:8000/generate",
                             json={
                                 "prompt": prompt,
15
                                 "max_tokens": max_tokens,
16
17
18
19
20
21
22
                                 "temperature": 0,
                                 "ignore_eos": True
                             })
    response.raise_for_status()
    return response.json()


23
24
25
26
def _query_server_long(prompt: str) -> dict:
    return _query_server(prompt, max_tokens=500)


27
28
29
30
31
@pytest.fixture
def api_server():
    script_path = Path(__file__).parent.joinpath(
        "api_server_async_engine.py").absolute()
    uvicorn_process = subprocess.Popen([
Simon Mo's avatar
Simon Mo committed
32
33
34
35
36
37
38
        sys.executable,
        "-u",
        str(script_path),
        "--model",
        "facebook/opt-125m",
        "--host",
        "127.0.0.1",
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
    ])
    yield
    uvicorn_process.terminate()


def test_api_server(api_server):
    """
    Run the API server and test it.

    We run both the server and requests in separate processes.

    We test that the server can handle incoming requests, including
    multiple requests at the same time, and that it can handle requests
    being cancelled without crashing.
    """
    with Pool(32) as pool:
        # Wait until the server is ready
56
        prompts = ["warm up"] * 1
57
58
59
        result = None
        while not result:
            try:
60
61
                for r in pool.map(_query_server, prompts):
                    result = r
62
                    break
63
            except requests.exceptions.ConnectionError:
64
65
66
67
68
69
70
71
72
73
74
75
                time.sleep(1)

        # Actual tests start here
        # Try with 1 prompt
        for result in pool.map(_query_server, prompts):
            assert result

        num_aborted_requests = requests.get(
            "http://localhost:8000/stats").json()["num_aborted_requests"]
        assert num_aborted_requests == 0

        # Try with 100 prompts
76
        prompts = ["test prompt"] * 100
77
78
79
        for result in pool.map(_query_server, prompts):
            assert result

80
    with Pool(32) as pool:
81
        # Cancel requests
82
        prompts = ["canceled requests"] * 100
83
84
        pool.map_async(_query_server_long, prompts)
        time.sleep(0.01)
85
86
87
88
        pool.terminate()
        pool.join()

        # check cancellation stats
Simon Mo's avatar
Simon Mo committed
89
90
91
        # give it some times to update the stats
        time.sleep(1)

92
93
94
95
96
97
98
        num_aborted_requests = requests.get(
            "http://localhost:8000/stats").json()["num_aborted_requests"]
        assert num_aborted_requests > 0

    # check that server still runs after cancellations
    with Pool(32) as pool:
        # Try with 100 prompts
99
        prompts = ["test prompt after canceled"] * 100
100
101
        for result in pool.map(_query_server, prompts):
            assert result