conftest.py 15.9 KB
Newer Older
1
import sys
2
3
4
5
6
7
import subprocess
import contextlib
import pytest
import asyncio
import os
import docker
8
9
10
import json
import math
import time
11
import random
12
import re
13
14

from docker.errors import NotFound
15
16
17
from typing import Optional, List, Dict
from syrupy.extensions.json import JSONSnapshotExtension
from aiohttp import ClientConnectorError, ClientOSError, ServerDisconnectedError
18
19

from text_generation import AsyncClient
drbh's avatar
drbh committed
20
21
22
23
24
25
26
from text_generation.types import (
    Response,
    Details,
    InputToken,
    Token,
    BestOfSequence,
    Grammar,
drbh's avatar
drbh committed
27
28
    ChatComplete,
    ChatCompletionChunk,
29
    ChatCompletionComplete,
30
    Completion,
drbh's avatar
drbh committed
31
)
32
33
34
35
36
37

DOCKER_IMAGE = os.getenv("DOCKER_IMAGE", None)
HUGGING_FACE_HUB_TOKEN = os.getenv("HUGGING_FACE_HUB_TOKEN", None)
DOCKER_VOLUME = os.getenv("DOCKER_VOLUME", "/data")


38
class ResponseComparator(JSONSnapshotExtension):
39
    rtol = 0.2
OlivierDehaene's avatar
OlivierDehaene committed
40

41
42
43
44
45
46
47
    def serialize(
        self,
        data,
        *,
        exclude=None,
        matcher=None,
    ):
48
49
50
51
52
53
54
        if (
            isinstance(data, Response)
            or isinstance(data, ChatComplete)
            or isinstance(data, ChatCompletionChunk)
            or isinstance(data, ChatCompletionComplete)
        ):
            data = data.model_dump()
55

56
        if isinstance(data, List):
57
            data = [d.model_dump() for d in data]
58
59
60
61
62
63
64
65
66
67
68
69
70
71

        data = self._filter(
            data=data, depth=0, path=(), exclude=exclude, matcher=matcher
        )
        return json.dumps(data, indent=2, ensure_ascii=False, sort_keys=False) + "\n"

    def matches(
        self,
        *,
        serialized_data,
        snapshot_data,
    ) -> bool:
        def convert_data(data):
            data = json.loads(data)
drbh's avatar
drbh committed
72
73
            if isinstance(data, Dict) and "choices" in data:
                choices = data["choices"]
74
75
76
77
78
                if isinstance(choices, List) and len(choices) >= 1:
                    if "delta" in choices[0]:
                        return ChatCompletionChunk(**data)
                    if "text" in choices[0]:
                        return Completion(**data)
drbh's avatar
drbh committed
79
                return ChatComplete(**data)
80
81
82
83

            if isinstance(data, Dict):
                return Response(**data)
            if isinstance(data, List):
84
85
86
87
88
89
                if (
                    len(data) > 0
                    and "object" in data[0]
                    and data[0]["object"] == "text_completion"
                ):
                    return [Completion(**d) for d in data]
90
91
92
93
94
95
96
                return [Response(**d) for d in data]
            raise NotImplementedError

        def eq_token(token: Token, other: Token) -> bool:
            return (
                token.id == other.id
                and token.text == other.text
97
                and math.isclose(token.logprob, other.logprob, rel_tol=self.rtol)
98
99
100
                and token.special == other.special
            )

101
        def eq_prefill_token(prefill_token: InputToken, other: InputToken) -> bool:
102
103
104
105
106
            try:
                return (
                    prefill_token.id == other.id
                    and prefill_token.text == other.text
                    and (
OlivierDehaene's avatar
OlivierDehaene committed
107
108
109
                        math.isclose(
                            prefill_token.logprob, other.logprob, rel_tol=self.rtol
                        )
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
                        if prefill_token.logprob is not None
                        else prefill_token.logprob == other.logprob
                    )
                )
            except TypeError:
                return False

        def eq_best_of(details: BestOfSequence, other: BestOfSequence) -> bool:
            return (
                details.finish_reason == other.finish_reason
                and details.generated_tokens == other.generated_tokens
                and details.seed == other.seed
                and len(details.prefill) == len(other.prefill)
                and all(
                    [
                        eq_prefill_token(d, o)
                        for d, o in zip(details.prefill, other.prefill)
                    ]
                )
                and len(details.tokens) == len(other.tokens)
                and all([eq_token(d, o) for d, o in zip(details.tokens, other.tokens)])
            )

        def eq_details(details: Details, other: Details) -> bool:
            return (
                details.finish_reason == other.finish_reason
                and details.generated_tokens == other.generated_tokens
                and details.seed == other.seed
                and len(details.prefill) == len(other.prefill)
                and all(
                    [
                        eq_prefill_token(d, o)
                        for d, o in zip(details.prefill, other.prefill)
                    ]
                )
                and len(details.tokens) == len(other.tokens)
                and all([eq_token(d, o) for d, o in zip(details.tokens, other.tokens)])
                and (
                    len(details.best_of_sequences)
                    if details.best_of_sequences is not None
                    else 0
                )
                == (
                    len(other.best_of_sequences)
                    if other.best_of_sequences is not None
                    else 0
                )
                and (
                    all(
                        [
                            eq_best_of(d, o)
                            for d, o in zip(
                                details.best_of_sequences, other.best_of_sequences
                            )
                        ]
                    )
                    if details.best_of_sequences is not None
                    else details.best_of_sequences == other.best_of_sequences
                )
            )

171
172
173
        def eq_completion(response: Completion, other: Completion) -> bool:
            return response.choices[0].text == other.choices[0].text

drbh's avatar
drbh committed
174
175
176
177
178
179
180
181
182
183
        def eq_chat_complete(response: ChatComplete, other: ChatComplete) -> bool:
            return (
                response.choices[0].message.content == other.choices[0].message.content
            )

        def eq_chat_complete_chunk(
            response: ChatCompletionChunk, other: ChatCompletionChunk
        ) -> bool:
            return response.choices[0].delta.content == other.choices[0].delta.content

184
185
186
187
188
189
190
191
192
193
194
195
196
        def eq_response(response: Response, other: Response) -> bool:
            return response.generated_text == other.generated_text and eq_details(
                response.details, other.details
            )

        serialized_data = convert_data(serialized_data)
        snapshot_data = convert_data(snapshot_data)

        if not isinstance(serialized_data, List):
            serialized_data = [serialized_data]
        if not isinstance(snapshot_data, List):
            snapshot_data = [snapshot_data]

197
198
199
200
201
        if isinstance(serialized_data[0], Completion):
            return len(snapshot_data) == len(serialized_data) and all(
                [eq_completion(r, o) for r, o in zip(serialized_data, snapshot_data)]
            )

drbh's avatar
drbh committed
202
203
204
205
206
207
208
209
210
211
212
213
214
        if isinstance(serialized_data[0], ChatComplete):
            return len(snapshot_data) == len(serialized_data) and all(
                [eq_chat_complete(r, o) for r, o in zip(serialized_data, snapshot_data)]
            )

        if isinstance(serialized_data[0], ChatCompletionChunk):
            return len(snapshot_data) == len(serialized_data) and all(
                [
                    eq_chat_complete_chunk(r, o)
                    for r, o in zip(serialized_data, snapshot_data)
                ]
            )

215
216
217
218
219
        return len(snapshot_data) == len(serialized_data) and all(
            [eq_response(r, o) for r, o in zip(serialized_data, snapshot_data)]
        )


220
221
222
223
class GenerousResponseComparator(ResponseComparator):
    # Needed for GPTQ with exllama which has serious numerical fluctuations.
    rtol = 0.75

OlivierDehaene's avatar
OlivierDehaene committed
224

225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
class LauncherHandle:
    def __init__(self, port: int):
        self.client = AsyncClient(f"http://localhost:{port}")

    def _inner_health(self):
        raise NotImplementedError

    async def health(self, timeout: int = 60):
        assert timeout > 0
        for _ in range(timeout):
            if not self._inner_health():
                raise RuntimeError("Launcher crashed")

            try:
                await self.client.generate("test")
                return
            except (ClientConnectorError, ClientOSError, ServerDisconnectedError) as e:
                time.sleep(1)
        raise RuntimeError("Health check failed")


class ContainerLauncherHandle(LauncherHandle):
    def __init__(self, docker_client, container_name, port: int):
        super(ContainerLauncherHandle, self).__init__(port)
        self.docker_client = docker_client
        self.container_name = container_name

    def _inner_health(self) -> bool:
        container = self.docker_client.containers.get(self.container_name)
        return container.status in ["running", "created"]


class ProcessLauncherHandle(LauncherHandle):
    def __init__(self, process, port: int):
        super(ProcessLauncherHandle, self).__init__(port)
        self.process = process

    def _inner_health(self) -> bool:
        return self.process.poll() is None


266
@pytest.fixture
267
268
def response_snapshot(snapshot):
    return snapshot.use_extension(ResponseComparator)
269

OlivierDehaene's avatar
OlivierDehaene committed
270

271
272
273
274
@pytest.fixture
def generous_response_snapshot(snapshot):
    return snapshot.use_extension(GenerousResponseComparator)

275
276
277
278
279
280
281
282
283
284
285
286

@pytest.fixture(scope="module")
def event_loop():
    loop = asyncio.get_event_loop()
    yield loop
    loop.close()


@pytest.fixture(scope="module")
def launcher(event_loop):
    @contextlib.contextmanager
    def local_launcher(
287
288
289
290
        model_id: str,
        num_shard: Optional[int] = None,
        quantize: Optional[str] = None,
        trust_remote_code: bool = False,
291
        use_flash_attention: bool = True,
drbh's avatar
drbh committed
292
        disable_grammar_support: bool = False,
OlivierDehaene's avatar
OlivierDehaene committed
293
        dtype: Optional[str] = None,
294
        revision: Optional[str] = None,
295
        max_input_length: Optional[int] = None,
Nicolas Patry's avatar
Nicolas Patry committed
296
        max_batch_prefill_tokens: Optional[int] = None,
297
        max_total_tokens: Optional[int] = None,
298
    ):
299
300
        port = random.randint(8000, 10_000)
        master_port = random.randint(10_000, 20_000)
301

302
303
304
        shard_uds_path = (
            f"/tmp/tgi-tests-{model_id.split('/')[-1]}-{num_shard}-{quantize}-server"
        )
305
306
307
308
309
310
311
312
313
314
315
316
317

        args = [
            "text-generation-launcher",
            "--model-id",
            model_id,
            "--port",
            str(port),
            "--master-port",
            str(master_port),
            "--shard-uds-path",
            shard_uds_path,
        ]

318
319
        env = os.environ

drbh's avatar
drbh committed
320
321
        if disable_grammar_support:
            args.append("--disable-grammar-support")
322
323
        if num_shard is not None:
            args.extend(["--num-shard", str(num_shard)])
324
        if quantize is not None:
325
            args.append("--quantize")
326
            args.append(quantize)
327
328
329
        if dtype is not None:
            args.append("--dtype")
            args.append(dtype)
330
331
332
        if revision is not None:
            args.append("--revision")
            args.append(revision)
333
334
        if trust_remote_code:
            args.append("--trust-remote-code")
335
336
337
        if max_input_length:
            args.append("--max-input-length")
            args.append(str(max_input_length))
Nicolas Patry's avatar
Nicolas Patry committed
338
339
340
        if max_batch_prefill_tokens:
            args.append("--max-batch-prefill-tokens")
            args.append(str(max_batch_prefill_tokens))
341
342
343
        if max_total_tokens:
            args.append("--max-total-tokens")
            args.append(str(max_total_tokens))
344

345
346
        env["LOG_LEVEL"] = "info,text_generation_router=debug"

347
348
349
        if not use_flash_attention:
            env["USE_FLASH_ATTENTION"] = "false"

350
        with subprocess.Popen(
351
            args, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=env
352
        ) as process:
353
            yield ProcessLauncherHandle(process, port)
354
355
356
357
358

            process.terminate()
            process.wait(60)

            launcher_output = process.stdout.read().decode("utf-8")
359
            print(launcher_output, file=sys.stderr)
360
361
362
363

            process.stdout.close()
            process.stderr.close()

364
365
366
        if not use_flash_attention:
            del env["USE_FLASH_ATTENTION"]

367
368
    @contextlib.contextmanager
    def docker_launcher(
369
370
371
372
        model_id: str,
        num_shard: Optional[int] = None,
        quantize: Optional[str] = None,
        trust_remote_code: bool = False,
373
        use_flash_attention: bool = True,
drbh's avatar
drbh committed
374
        disable_grammar_support: bool = False,
OlivierDehaene's avatar
OlivierDehaene committed
375
        dtype: Optional[str] = None,
376
        revision: Optional[str] = None,
377
        max_input_length: Optional[int] = None,
Nicolas Patry's avatar
Nicolas Patry committed
378
        max_batch_prefill_tokens: Optional[int] = None,
379
        max_total_tokens: Optional[int] = None,
380
    ):
381
        port = random.randint(8000, 10_000)
382
383
384

        args = ["--model-id", model_id, "--env"]

drbh's avatar
drbh committed
385
386
        if disable_grammar_support:
            args.append("--disable-grammar-support")
387
388
        if num_shard is not None:
            args.extend(["--num-shard", str(num_shard)])
389
        if quantize is not None:
390
            args.append("--quantize")
391
            args.append(quantize)
392
393
394
        if dtype is not None:
            args.append("--dtype")
            args.append(dtype)
395
396
397
        if revision is not None:
            args.append("--revision")
            args.append(revision)
398
399
        if trust_remote_code:
            args.append("--trust-remote-code")
400
401
402
        if max_input_length:
            args.append("--max-input-length")
            args.append(str(max_input_length))
Nicolas Patry's avatar
Nicolas Patry committed
403
404
405
        if max_batch_prefill_tokens:
            args.append("--max-batch-prefill-tokens")
            args.append(str(max_batch_prefill_tokens))
406
407
408
        if max_total_tokens:
            args.append("--max-total-tokens")
            args.append(str(max_total_tokens))
409
410
411
412
413
414
415
416
417
418
419
420
421
422

        client = docker.from_env()

        container_name = f"tgi-tests-{model_id.split('/')[-1]}-{num_shard}-{quantize}"

        try:
            container = client.containers.get(container_name)
            container.stop()
            container.wait()
        except NotFound:
            pass

        gpu_count = num_shard if num_shard is not None else 1

423
424
425
        env = {
            "LOG_LEVEL": "info,text_generation_router=debug",
        }
426
427
428
        if not use_flash_attention:
            env["USE_FLASH_ATTENTION"] = "false"

429
430
431
432
433
434
435
436
437
438
439
440
        if HUGGING_FACE_HUB_TOKEN is not None:
            env["HUGGING_FACE_HUB_TOKEN"] = HUGGING_FACE_HUB_TOKEN

        volumes = []
        if DOCKER_VOLUME:
            volumes = [f"{DOCKER_VOLUME}:/data"]

        container = client.containers.run(
            DOCKER_IMAGE,
            command=args,
            name=container_name,
            environment=env,
441
            auto_remove=False,
442
443
444
445
446
447
            detach=True,
            device_requests=[
                docker.types.DeviceRequest(count=gpu_count, capabilities=[["gpu"]])
            ],
            volumes=volumes,
            ports={"80/tcp": port},
OlivierDehaene's avatar
OlivierDehaene committed
448
            shm_size="1G",
449
450
        )

451
        yield ContainerLauncherHandle(client, container.name, port)
452

453
454
455
        if not use_flash_attention:
            del env["USE_FLASH_ATTENTION"]

456
457
458
459
460
        try:
            container.stop()
            container.wait()
        except NotFound:
            pass
461
462

        container_output = container.logs().decode("utf-8")
463
        print(container_output, file=sys.stderr)
464

465
466
        container.remove()

467
468
469
470
471
472
473
474
    if DOCKER_IMAGE is not None:
        return docker_launcher
    return local_launcher


@pytest.fixture(scope="module")
def generate_load():
    async def generate_load_inner(
drbh's avatar
drbh committed
475
476
477
478
479
480
481
        client: AsyncClient,
        prompt: str,
        max_new_tokens: int,
        n: int,
        seed: Optional[int] = None,
        grammar: Optional[Grammar] = None,
        stop_sequences: Optional[List[str]] = None,
482
483
    ) -> List[Response]:
        futures = [
484
            client.generate(
drbh's avatar
drbh committed
485
486
487
488
489
490
                prompt,
                max_new_tokens=max_new_tokens,
                decoder_input_details=True,
                seed=seed,
                grammar=grammar,
                stop_sequences=stop_sequences,
491
492
            )
            for _ in range(n)
493
494
        ]

495
        return await asyncio.gather(*futures)
496
497

    return generate_load_inner