File size: 4,679 Bytes
35eafc3
 
 
 
 
ede9e6a
 
 
 
35eafc3
 
 
ede9e6a
35eafc3
ede9e6a
 
35eafc3
 
ede9e6a
35eafc3
 
196f56a
35eafc3
 
 
 
 
ede9e6a
35eafc3
ede9e6a
 
 
35eafc3
 
196f56a
35eafc3
 
 
 
 
 
 
ede9e6a
 
 
35eafc3
 
 
 
196f56a
35eafc3
 
 
196f56a
35eafc3
 
 
 
 
 
 
ede9e6a
 
 
35eafc3
 
 
 
 
ede9e6a
35eafc3
ede9e6a
 
35eafc3
 
 
 
 
ede9e6a
35eafc3
 
196f56a
ede9e6a
35eafc3
 
 
 
 
 
 
 
ede9e6a
35eafc3
ede9e6a
 
 
35eafc3
ede9e6a
35eafc3
 
 
 
 
 
ede9e6a
35eafc3
ede9e6a
 
35eafc3
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
import asyncio

import anyio
import pytest

from faster_whisper_server.config import Config, WhisperConfig
from tests.conftest import DEFAULT_WHISPER_MODEL, AclientFactory

MODEL = DEFAULT_WHISPER_MODEL  # just to make the test more readable


@pytest.mark.asyncio
async def test_model_unloaded_after_ttl(aclient_factory: AclientFactory) -> None:
    ttl = 5
    config = Config(whisper=WhisperConfig(model=MODEL, ttl=ttl), enable_ui=False)
    async with aclient_factory(config) as aclient:
        res = (await aclient.get("/api/ps")).json()
        assert len(res["models"]) == 0
        await aclient.post(f"/api/ps/{MODEL}")
        res = (await aclient.get("/api/ps")).json()
        assert len(res["models"]) == 1
        await asyncio.sleep(ttl + 1)  # wait for the model to be unloaded
        res = (await aclient.get("/api/ps")).json()
        assert len(res["models"]) == 0


@pytest.mark.asyncio
async def test_ttl_resets_after_usage(aclient_factory: AclientFactory) -> None:
    ttl = 5
    config = Config(whisper=WhisperConfig(model=MODEL, ttl=ttl), enable_ui=False)
    async with aclient_factory(config) as aclient:
        await aclient.post(f"/api/ps/{MODEL}")
        res = (await aclient.get("/api/ps")).json()
        assert len(res["models"]) == 1
        await asyncio.sleep(ttl - 2)  # sleep for less than the ttl. The model should not be unloaded
        res = (await aclient.get("/api/ps")).json()
        assert len(res["models"]) == 1

        async with await anyio.open_file("audio.wav", "rb") as f:
            data = await f.read()
        res = (
            await aclient.post(
                "/v1/audio/transcriptions",
                files={"file": ("audio.wav", data, "audio/wav")},
                data={"model": MODEL},
            )
        ).json()
        res = (await aclient.get("/api/ps")).json()
        assert len(res["models"]) == 1
        await asyncio.sleep(ttl - 2)  # sleep for less than the ttl. The model should not be unloaded
        res = (await aclient.get("/api/ps")).json()
        assert len(res["models"]) == 1

        await asyncio.sleep(3)  # sleep for a bit more. The model should be unloaded
        res = (await aclient.get("/api/ps")).json()
        assert len(res["models"]) == 0

        # test the model can be used again after being unloaded
        # this just ensures the model can be loaded again after being unloaded
        res = (
            await aclient.post(
                "/v1/audio/transcriptions",
                files={"file": ("audio.wav", data, "audio/wav")},
                data={"model": MODEL},
            )
        ).json()


@pytest.mark.asyncio
async def test_model_cant_be_unloaded_when_used(aclient_factory: AclientFactory) -> None:
    ttl = 0
    config = Config(whisper=WhisperConfig(model=MODEL, ttl=ttl), enable_ui=False)
    async with aclient_factory(config) as aclient:
        async with await anyio.open_file("audio.wav", "rb") as f:
            data = await f.read()

        task = asyncio.create_task(
            aclient.post(
                "/v1/audio/transcriptions", files={"file": ("audio.wav", data, "audio/wav")}, data={"model": MODEL}
            )
        )
        await asyncio.sleep(0.1)  # wait for the server to start processing the request
        res = await aclient.delete(f"/api/ps/{MODEL}")
        assert res.status_code == 409

        await task
        res = (await aclient.get("/api/ps")).json()
        assert len(res["models"]) == 0


@pytest.mark.asyncio
async def test_model_cant_be_loaded_twice(aclient_factory: AclientFactory) -> None:
    ttl = -1
    config = Config(whisper=WhisperConfig(model=MODEL, ttl=ttl), enable_ui=False)
    async with aclient_factory(config) as aclient:
        res = await aclient.post(f"/api/ps/{MODEL}")
        assert res.status_code == 201
        res = await aclient.post(f"/api/ps/{MODEL}")
        assert res.status_code == 409
        res = (await aclient.get("/api/ps")).json()
        assert len(res["models"]) == 1


@pytest.mark.asyncio
async def test_model_is_unloaded_after_request_when_ttl_is_zero(aclient_factory: AclientFactory) -> None:
    ttl = 0
    config = Config(whisper=WhisperConfig(model=MODEL, ttl=ttl), enable_ui=False)
    async with aclient_factory(config) as aclient:
        async with await anyio.open_file("audio.wav", "rb") as f:
            data = await f.read()
        res = await aclient.post(
            "/v1/audio/transcriptions",
            files={"file": ("audio.wav", data, "audio/wav")},
            data={"model": "Systran/faster-whisper-tiny.en"},
        )
        res = (await aclient.get("/api/ps")).json()
        assert len(res["models"]) == 0