conftest.py 16.5 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
import json
import math
10
11
import shutil
import tempfile
12
import time
13
import random
14
15

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

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

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")


39
class ResponseComparator(JSONSnapshotExtension):
40
    rtol = 0.2
41
    ignore_logprob = False
OlivierDehaene's avatar
OlivierDehaene committed
42

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

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

        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
74
75
            if isinstance(data, Dict) and "choices" in data:
                choices = data["choices"]
76
77
78
79
80
                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
81
                return ChatComplete(**data)
82
83
84
85

            if isinstance(data, Dict):
                return Response(**data)
            if isinstance(data, List):
86
87
88
89
90
91
                if (
                    len(data) > 0
                    and "object" in data[0]
                    and data[0]["object"] == "text_completion"
                ):
                    return [Completion(**d) for d in data]
92
93
94
95
96
97
98
                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
99
100
101
102
                and (
                    self.ignore_logprob
                    or math.isclose(token.logprob, other.logprob, rel_tol=self.rtol)
                )
103
104
105
                and token.special == other.special
            )

106
        def eq_prefill_token(prefill_token: InputToken, other: InputToken) -> bool:
107
108
109
110
111
            try:
                return (
                    prefill_token.id == other.id
                    and prefill_token.text == other.text
                    and (
112
113
114
115
116
                        self.ignore_logprob
                        or math.isclose(
                            prefill_token.logprob,
                            other.logprob,
                            rel_tol=self.rtol,
OlivierDehaene's avatar
OlivierDehaene committed
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
171
172
173
174
175
176
177
178
                        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
                )
            )

179
180
181
        def eq_completion(response: Completion, other: Completion) -> bool:
            return response.choices[0].text == other.choices[0].text

drbh's avatar
drbh committed
182
183
184
185
186
187
188
189
190
191
        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

192
193
194
195
196
197
198
199
200
201
202
203
204
        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]

205
206
207
208
209
        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
210
211
212
213
214
215
216
217
218
219
220
221
222
        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)
                ]
            )

223
224
225
226
227
        return len(snapshot_data) == len(serialized_data) and all(
            [eq_response(r, o) for r, o in zip(serialized_data, snapshot_data)]
        )


228
229
230
231
class GenerousResponseComparator(ResponseComparator):
    # Needed for GPTQ with exllama which has serious numerical fluctuations.
    rtol = 0.75

OlivierDehaene's avatar
OlivierDehaene committed
232

233
234
235
236
class IgnoreLogProbResponseComparator(ResponseComparator):
    ignore_logprob = True


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
266
267
268
269
270
271
272
273
274
275
276
277
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


278
@pytest.fixture
279
280
def response_snapshot(snapshot):
    return snapshot.use_extension(ResponseComparator)
281

OlivierDehaene's avatar
OlivierDehaene committed
282

283
284
285
286
@pytest.fixture
def generous_response_snapshot(snapshot):
    return snapshot.use_extension(GenerousResponseComparator)

287

288
289
290
291
292
@pytest.fixture
def ignore_logprob_response_snapshot(snapshot):
    return snapshot.use_extension(IgnoreLogProbResponseComparator)


293
294
295
296
297
298
299
300
301
302
303
@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(
304
305
306
307
        model_id: str,
        num_shard: Optional[int] = None,
        quantize: Optional[str] = None,
        trust_remote_code: bool = False,
308
        use_flash_attention: bool = True,
drbh's avatar
drbh committed
309
        disable_grammar_support: bool = False,
OlivierDehaene's avatar
OlivierDehaene committed
310
        dtype: Optional[str] = None,
311
        revision: Optional[str] = None,
312
        max_input_length: Optional[int] = None,
Nicolas Patry's avatar
Nicolas Patry committed
313
        max_batch_prefill_tokens: Optional[int] = None,
314
        max_total_tokens: Optional[int] = None,
315
    ):
316
317
        port = random.randint(8000, 10_000)
        master_port = random.randint(10_000, 20_000)
318

319
320
321
        shard_uds_path = (
            f"/tmp/tgi-tests-{model_id.split('/')[-1]}-{num_shard}-{quantize}-server"
        )
322
323
324
325
326
327
328
329
330
331
332
333
334

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

335
336
        env = os.environ

drbh's avatar
drbh committed
337
338
        if disable_grammar_support:
            args.append("--disable-grammar-support")
339
340
        if num_shard is not None:
            args.extend(["--num-shard", str(num_shard)])
341
        if quantize is not None:
342
            args.append("--quantize")
343
            args.append(quantize)
344
345
346
        if dtype is not None:
            args.append("--dtype")
            args.append(dtype)
347
348
349
        if revision is not None:
            args.append("--revision")
            args.append(revision)
350
351
        if trust_remote_code:
            args.append("--trust-remote-code")
352
353
354
        if max_input_length:
            args.append("--max-input-length")
            args.append(str(max_input_length))
Nicolas Patry's avatar
Nicolas Patry committed
355
356
357
        if max_batch_prefill_tokens:
            args.append("--max-batch-prefill-tokens")
            args.append(str(max_batch_prefill_tokens))
358
359
360
        if max_total_tokens:
            args.append("--max-total-tokens")
            args.append(str(max_total_tokens))
361

362
363
        env["LOG_LEVEL"] = "info,text_generation_router=debug"

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

367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
        with tempfile.TemporaryFile("w+") as tmp:
            # We'll output stdout/stderr to a temporary file. Using a pipe
            # cause the process to block until stdout is read.
            with subprocess.Popen(
                args,
                stdout=tmp,
                stderr=subprocess.STDOUT,
                env=env,
            ) as process:
                yield ProcessLauncherHandle(process, port)

                process.terminate()
                process.wait(60)

                tmp.seek(0)
                shutil.copyfileobj(tmp, sys.stderr)
383

384
385
386
        if not use_flash_attention:
            del env["USE_FLASH_ATTENTION"]

387
388
    @contextlib.contextmanager
    def docker_launcher(
389
390
391
392
        model_id: str,
        num_shard: Optional[int] = None,
        quantize: Optional[str] = None,
        trust_remote_code: bool = False,
393
        use_flash_attention: bool = True,
drbh's avatar
drbh committed
394
        disable_grammar_support: bool = False,
OlivierDehaene's avatar
OlivierDehaene committed
395
        dtype: Optional[str] = None,
396
        revision: Optional[str] = None,
397
        max_input_length: Optional[int] = None,
Nicolas Patry's avatar
Nicolas Patry committed
398
        max_batch_prefill_tokens: Optional[int] = None,
399
        max_total_tokens: Optional[int] = None,
400
    ):
401
        port = random.randint(8000, 10_000)
402
403
404

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

drbh's avatar
drbh committed
405
406
        if disable_grammar_support:
            args.append("--disable-grammar-support")
407
408
        if num_shard is not None:
            args.extend(["--num-shard", str(num_shard)])
409
        if quantize is not None:
410
            args.append("--quantize")
411
            args.append(quantize)
412
413
414
        if dtype is not None:
            args.append("--dtype")
            args.append(dtype)
415
416
417
        if revision is not None:
            args.append("--revision")
            args.append(revision)
418
419
        if trust_remote_code:
            args.append("--trust-remote-code")
420
421
422
        if max_input_length:
            args.append("--max-input-length")
            args.append(str(max_input_length))
Nicolas Patry's avatar
Nicolas Patry committed
423
424
425
        if max_batch_prefill_tokens:
            args.append("--max-batch-prefill-tokens")
            args.append(str(max_batch_prefill_tokens))
426
427
428
        if max_total_tokens:
            args.append("--max-total-tokens")
            args.append(str(max_total_tokens))
429
430
431
432
433
434
435
436
437
438
439
440
441
442

        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

443
444
445
        env = {
            "LOG_LEVEL": "info,text_generation_router=debug",
        }
446
447
448
        if not use_flash_attention:
            env["USE_FLASH_ATTENTION"] = "false"

449
450
451
452
453
454
455
456
457
458
459
460
        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,
461
            auto_remove=False,
462
463
464
465
466
467
            detach=True,
            device_requests=[
                docker.types.DeviceRequest(count=gpu_count, capabilities=[["gpu"]])
            ],
            volumes=volumes,
            ports={"80/tcp": port},
OlivierDehaene's avatar
OlivierDehaene committed
468
            shm_size="1G",
469
470
        )

471
        yield ContainerLauncherHandle(client, container.name, port)
472

473
474
475
        if not use_flash_attention:
            del env["USE_FLASH_ATTENTION"]

476
477
478
479
480
        try:
            container.stop()
            container.wait()
        except NotFound:
            pass
481
482

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

485
486
        container.remove()

487
488
489
490
491
492
493
494
    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
495
496
497
498
499
500
501
        client: AsyncClient,
        prompt: str,
        max_new_tokens: int,
        n: int,
        seed: Optional[int] = None,
        grammar: Optional[Grammar] = None,
        stop_sequences: Optional[List[str]] = None,
502
503
    ) -> List[Response]:
        futures = [
504
            client.generate(
drbh's avatar
drbh committed
505
506
507
508
509
510
                prompt,
                max_new_tokens=max_new_tokens,
                decoder_input_details=True,
                seed=seed,
                grammar=grammar,
                stop_sequences=stop_sequences,
511
512
            )
            for _ in range(n)
513
514
        ]

515
        return await asyncio.gather(*futures)
516
517

    return generate_load_inner