We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Update Docker to torch 2.1?
### Feature request

H100s have trouble with gptq quants due to not having latest pytorch, can in the next TGI Docker we update torch to this, or have one special for this for use on h100s? 

### Motivation

Cant get tgi + gptq quant to work on h100s

### Your contribution

Sorry I dont have any contribution ^_^ 

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `integration-tests/conftest.py`
Content:
```
1 import sys
2 import subprocess
3 import contextlib
4 import pytest
5 import asyncio
6 import os
7 import docker
8 import json
9 import math
10 import time
11 import random
12 
13 from docker.errors import NotFound
14 from typing import Optional, List, Dict
15 from syrupy.extensions.json import JSONSnapshotExtension
16 from aiohttp import ClientConnectorError, ClientOSError, ServerDisconnectedError
17 
18 from text_generation import AsyncClient
19 from text_generation.types import Response, Details, InputToken, Token, BestOfSequence
20 
21 DOCKER_IMAGE = os.getenv("DOCKER_IMAGE", None)
22 HUGGING_FACE_HUB_TOKEN = os.getenv("HUGGING_FACE_HUB_TOKEN", None)
23 DOCKER_VOLUME = os.getenv("DOCKER_VOLUME", "/data")
24 
25 
26 class ResponseComparator(JSONSnapshotExtension):
27     def serialize(
28         self,
29         data,
30         *,
31         exclude=None,
32         matcher=None,
33     ):
34         if isinstance(data, List):
35             data = [d.dict() for d in data]
36 
37         data = self._filter(
38             data=data, depth=0, path=(), exclude=exclude, matcher=matcher
39         )
40         return json.dumps(data, indent=2, ensure_ascii=False, sort_keys=False) + "\n"
41 
42     def matches(
43         self,
44         *,
45         serialized_data,
46         snapshot_data,
47     ) -> bool:
48         def convert_data(data):
49             data = json.loads(data)
50 
51             if isinstance(data, Dict):
52                 return Response(**data)
53             if isinstance(data, List):
54                 return [Response(**d) for d in data]
55             raise NotImplementedError
56 
57         def eq_token(token: Token, other: Token) -> bool:
58             return (
59                 token.id == other.id
60                 and token.text == other.text
61                 and math.isclose(token.logprob, other.logprob, rel_tol=0.2)
62                 and token.special == other.special
63             )
64 
65         def eq_prefill_token(prefill_token: InputToken, other: InputToken) -> bool:
66             try:
67                 return (
68                     prefill_token.id == other.id
69                     and prefill_token.text == other.text
70                     and (
71                         math.isclose(prefill_token.logprob, other.logprob, rel_tol=0.2)
72                         if prefill_token.logprob is not None
73                         else prefill_token.logprob == other.logprob
74                     )
75                 )
76             except TypeError:
77                 return False
78 
79         def eq_best_of(details: BestOfSequence, other: BestOfSequence) -> bool:
80             return (
81                 details.finish_reason == other.finish_reason
82                 and details.generated_tokens == other.generated_tokens
83                 and details.seed == other.seed
84                 and len(details.prefill) == len(other.prefill)
85                 and all(
86                     [
87                         eq_prefill_token(d, o)
88                         for d, o in zip(details.prefill, other.prefill)
89                     ]
90                 )
91                 and len(details.tokens) == len(other.tokens)
92                 and all([eq_token(d, o) for d, o in zip(details.tokens, other.tokens)])
93             )
94 
95         def eq_details(details: Details, other: Details) -> bool:
96             return (
97                 details.finish_reason == other.finish_reason
98                 and details.generated_tokens == other.generated_tokens
99                 and details.seed == other.seed
100                 and len(details.prefill) == len(other.prefill)
101                 and all(
102                     [
103                         eq_prefill_token(d, o)
104                         for d, o in zip(details.prefill, other.prefill)
105                     ]
106                 )
107                 and len(details.tokens) == len(other.tokens)
108                 and all([eq_token(d, o) for d, o in zip(details.tokens, other.tokens)])
109                 and (
110                     len(details.best_of_sequences)
111                     if details.best_of_sequences is not None
112                     else 0
113                 )
114                 == (
115                     len(other.best_of_sequences)
116                     if other.best_of_sequences is not None
117                     else 0
118                 )
119                 and (
120                     all(
121                         [
122                             eq_best_of(d, o)
123                             for d, o in zip(
124                                 details.best_of_sequences, other.best_of_sequences
125                             )
126                         ]
127                     )
128                     if details.best_of_sequences is not None
129                     else details.best_of_sequences == other.best_of_sequences
130                 )
131             )
132 
133         def eq_response(response: Response, other: Response) -> bool:
134             return response.generated_text == other.generated_text and eq_details(
135                 response.details, other.details
136             )
137 
138         serialized_data = convert_data(serialized_data)
139         snapshot_data = convert_data(snapshot_data)
140 
141         if not isinstance(serialized_data, List):
142             serialized_data = [serialized_data]
143         if not isinstance(snapshot_data, List):
144             snapshot_data = [snapshot_data]
145 
146         return len(snapshot_data) == len(serialized_data) and all(
147             [eq_response(r, o) for r, o in zip(serialized_data, snapshot_data)]
148         )
149 
150 
151 class LauncherHandle:
152     def __init__(self, port: int):
153         self.client = AsyncClient(f"http://localhost:{port}")
154 
155     def _inner_health(self):
156         raise NotImplementedError
157 
158     async def health(self, timeout: int = 60):
159         assert timeout > 0
160         for _ in range(timeout):
161             if not self._inner_health():
162                 raise RuntimeError("Launcher crashed")
163 
164             try:
165                 await self.client.generate("test")
166                 return
167             except (ClientConnectorError, ClientOSError, ServerDisconnectedError) as e:
168                 time.sleep(1)
169         raise RuntimeError("Health check failed")
170 
171 
172 class ContainerLauncherHandle(LauncherHandle):
173     def __init__(self, docker_client, container_name, port: int):
174         super(ContainerLauncherHandle, self).__init__(port)
175         self.docker_client = docker_client
176         self.container_name = container_name
177 
178     def _inner_health(self) -> bool:
179         container = self.docker_client.containers.get(self.container_name)
180         return container.status in ["running", "created"]
181 
182 
183 class ProcessLauncherHandle(LauncherHandle):
184     def __init__(self, process, port: int):
185         super(ProcessLauncherHandle, self).__init__(port)
186         self.process = process
187 
188     def _inner_health(self) -> bool:
189         return self.process.poll() is None
190 
191 
192 @pytest.fixture
193 def response_snapshot(snapshot):
194     return snapshot.use_extension(ResponseComparator)
195 
196 
197 @pytest.fixture(scope="module")
198 def event_loop():
199     loop = asyncio.get_event_loop()
200     yield loop
201     loop.close()
202 
203 
204 @pytest.fixture(scope="module")
205 def launcher(event_loop):
206     @contextlib.contextmanager
207     def local_launcher(
208         model_id: str,
209         num_shard: Optional[int] = None,
210         quantize: Optional[str] = None,
211         trust_remote_code: bool = False,
212         use_flash_attention: bool = True,
213     ):
214         port = random.randint(8000, 10_000)
215         master_port = random.randint(10_000, 20_000)
216 
217         shard_uds_path = (
218             f"/tmp/tgi-tests-{model_id.split('/')[-1]}-{num_shard}-{quantize}-server"
219         )
220 
221         args = [
222             "text-generation-launcher",
223             "--model-id",
224             model_id,
225             "--port",
226             str(port),
227             "--master-port",
228             str(master_port),
229             "--shard-uds-path",
230             shard_uds_path,
231         ]
232 
233         env = os.environ
234 
235         if num_shard is not None:
236             args.extend(["--num-shard", str(num_shard)])
237         if quantize is not None:
238             args.append("--quantize")
239             args.append(quantize)
240         if trust_remote_code:
241             args.append("--trust-remote-code")
242 
243         env["LOG_LEVEL"] = "info,text_generation_router=debug"
244 
245         if not use_flash_attention:
246             env["USE_FLASH_ATTENTION"] = "false"
247 
248         with subprocess.Popen(
249             args, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=env
250         ) as process:
251             yield ProcessLauncherHandle(process, port)
252 
253             process.terminate()
254             process.wait(60)
255 
256             launcher_output = process.stdout.read().decode("utf-8")
257             print(launcher_output, file=sys.stderr)
258 
259             process.stdout.close()
260             process.stderr.close()
261 
262         if not use_flash_attention:
263             del env["USE_FLASH_ATTENTION"]
264 
265     @contextlib.contextmanager
266     def docker_launcher(
267         model_id: str,
268         num_shard: Optional[int] = None,
269         quantize: Optional[str] = None,
270         trust_remote_code: bool = False,
271         use_flash_attention: bool = True,
272     ):
273         port = random.randint(8000, 10_000)
274 
275         args = ["--model-id", model_id, "--env"]
276 
277         if num_shard is not None:
278             args.extend(["--num-shard", str(num_shard)])
279         if quantize is not None:
280             args.append("--quantize")
281             args.append(quantize)
282         if trust_remote_code:
283             args.append("--trust-remote-code")
284 
285         client = docker.from_env()
286 
287         container_name = f"tgi-tests-{model_id.split('/')[-1]}-{num_shard}-{quantize}"
288 
289         try:
290             container = client.containers.get(container_name)
291             container.stop()
292             container.wait()
293         except NotFound:
294             pass
295 
296         gpu_count = num_shard if num_shard is not None else 1
297 
298         env = {"LOG_LEVEL": "info,text_generation_router=debug"}
299         if not use_flash_attention:
300             env["USE_FLASH_ATTENTION"] = "false"
301 
302         if HUGGING_FACE_HUB_TOKEN is not None:
303             env["HUGGING_FACE_HUB_TOKEN"] = HUGGING_FACE_HUB_TOKEN
304 
305         volumes = []
306         if DOCKER_VOLUME:
307             volumes = [f"{DOCKER_VOLUME}:/data"]
308 
309         container = client.containers.run(
310             DOCKER_IMAGE,
311             command=args,
312             name=container_name,
313             environment=env,
314             auto_remove=False,
315             detach=True,
316             device_requests=[
317                 docker.types.DeviceRequest(count=gpu_count, capabilities=[["gpu"]])
318             ],
319             volumes=volumes,
320             ports={"80/tcp": port},
321         )
322 
323         yield ContainerLauncherHandle(client, container.name, port)
324 
325         if not use_flash_attention:
326             del env["USE_FLASH_ATTENTION"]
327 
328         try:
329             container.stop()
330             container.wait()
331         except NotFound:
332             pass
333 
334         container_output = container.logs().decode("utf-8")
335         print(container_output, file=sys.stderr)
336 
337         container.remove()
338 
339     if DOCKER_IMAGE is not None:
340         return docker_launcher
341     return local_launcher
342 
343 
344 @pytest.fixture(scope="module")
345 def generate_load():
346     async def generate_load_inner(
347         client: AsyncClient, prompt: str, max_new_tokens: int, n: int
348     ) -> List[Response]:
349         futures = [
350             client.generate(
351                 prompt, max_new_tokens=max_new_tokens, decoder_input_details=True
352             )
353             for _ in range(n)
354         ]
355 
356         return await asyncio.gather(*futures)
357 
358     return generate_load_inner
359
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/integration-tests/conftest.py b/integration-tests/conftest.py
--- a/integration-tests/conftest.py
+++ b/integration-tests/conftest.py
@@ -318,6 +318,7 @@
             ],
             volumes=volumes,
             ports={"80/tcp": port},
+            shm_size="1G"
         )
 
         yield ContainerLauncherHandle(client, container.name, port)

{"golden_diff": "diff --git a/integration-tests/conftest.py b/integration-tests/conftest.py\n--- a/integration-tests/conftest.py\n+++ b/integration-tests/conftest.py\n@@ -318,6 +318,7 @@\n             ],\n             volumes=volumes,\n             ports={\"80/tcp\": port},\n+            shm_size=\"1G\"\n         )\n \n         yield ContainerLauncherHandle(client, container.name, port)\n", "issue": "Update Docker to torch 2.1?\n### Feature request\n\nH100s have trouble with gptq quants due to not having latest pytorch, can in the next TGI Docker we update torch to this, or have one special for this for use on h100s? \n\n### Motivation\n\nCant get tgi + gptq quant to work on h100s\n\n### Your contribution\n\nSorry I dont have any contribution ^_^ \n", "before_files": [{"content": "import sys\nimport subprocess\nimport contextlib\nimport pytest\nimport asyncio\nimport os\nimport docker\nimport json\nimport math\nimport time\nimport random\n\nfrom docker.errors import NotFound\nfrom typing import Optional, List, Dict\nfrom syrupy.extensions.json import JSONSnapshotExtension\nfrom aiohttp import ClientConnectorError, ClientOSError, ServerDisconnectedError\n\nfrom text_generation import AsyncClient\nfrom text_generation.types import Response, Details, InputToken, Token, BestOfSequence\n\nDOCKER_IMAGE = os.getenv(\"DOCKER_IMAGE\", None)\nHUGGING_FACE_HUB_TOKEN = os.getenv(\"HUGGING_FACE_HUB_TOKEN\", None)\nDOCKER_VOLUME = os.getenv(\"DOCKER_VOLUME\", \"/data\")\n\n\nclass ResponseComparator(JSONSnapshotExtension):\n    def serialize(\n        self,\n        data,\n        *,\n        exclude=None,\n        matcher=None,\n    ):\n        if isinstance(data, List):\n            data = [d.dict() for d in data]\n\n        data = self._filter(\n            data=data, depth=0, path=(), exclude=exclude, matcher=matcher\n        )\n        return json.dumps(data, indent=2, ensure_ascii=False, sort_keys=False) + \"\\n\"\n\n    def matches(\n        self,\n        *,\n        serialized_data,\n        snapshot_data,\n    ) -> bool:\n        def convert_data(data):\n            data = json.loads(data)\n\n            if isinstance(data, Dict):\n                return Response(**data)\n            if isinstance(data, List):\n                return [Response(**d) for d in data]\n            raise NotImplementedError\n\n        def eq_token(token: Token, other: Token) -> bool:\n            return (\n                token.id == other.id\n                and token.text == other.text\n                and math.isclose(token.logprob, other.logprob, rel_tol=0.2)\n                and token.special == other.special\n            )\n\n        def eq_prefill_token(prefill_token: InputToken, other: InputToken) -> bool:\n            try:\n                return (\n                    prefill_token.id == other.id\n                    and prefill_token.text == other.text\n                    and (\n                        math.isclose(prefill_token.logprob, other.logprob, rel_tol=0.2)\n                        if prefill_token.logprob is not None\n                        else prefill_token.logprob == other.logprob\n                    )\n                )\n            except TypeError:\n                return False\n\n        def eq_best_of(details: BestOfSequence, other: BestOfSequence) -> bool:\n            return (\n                details.finish_reason == other.finish_reason\n                and details.generated_tokens == other.generated_tokens\n                and details.seed == other.seed\n                and len(details.prefill) == len(other.prefill)\n                and all(\n                    [\n                        eq_prefill_token(d, o)\n                        for d, o in zip(details.prefill, other.prefill)\n                    ]\n                )\n                and len(details.tokens) == len(other.tokens)\n                and all([eq_token(d, o) for d, o in zip(details.tokens, other.tokens)])\n            )\n\n        def eq_details(details: Details, other: Details) -> bool:\n            return (\n                details.finish_reason == other.finish_reason\n                and details.generated_tokens == other.generated_tokens\n                and details.seed == other.seed\n                and len(details.prefill) == len(other.prefill)\n                and all(\n                    [\n                        eq_prefill_token(d, o)\n                        for d, o in zip(details.prefill, other.prefill)\n                    ]\n                )\n                and len(details.tokens) == len(other.tokens)\n                and all([eq_token(d, o) for d, o in zip(details.tokens, other.tokens)])\n                and (\n                    len(details.best_of_sequences)\n                    if details.best_of_sequences is not None\n                    else 0\n                )\n                == (\n                    len(other.best_of_sequences)\n                    if other.best_of_sequences is not None\n                    else 0\n                )\n                and (\n                    all(\n                        [\n                            eq_best_of(d, o)\n                            for d, o in zip(\n                                details.best_of_sequences, other.best_of_sequences\n                            )\n                        ]\n                    )\n                    if details.best_of_sequences is not None\n                    else details.best_of_sequences == other.best_of_sequences\n                )\n            )\n\n        def eq_response(response: Response, other: Response) -> bool:\n            return response.generated_text == other.generated_text and eq_details(\n                response.details, other.details\n            )\n\n        serialized_data = convert_data(serialized_data)\n        snapshot_data = convert_data(snapshot_data)\n\n        if not isinstance(serialized_data, List):\n            serialized_data = [serialized_data]\n        if not isinstance(snapshot_data, List):\n            snapshot_data = [snapshot_data]\n\n        return len(snapshot_data) == len(serialized_data) and all(\n            [eq_response(r, o) for r, o in zip(serialized_data, snapshot_data)]\n        )\n\n\nclass LauncherHandle:\n    def __init__(self, port: int):\n        self.client = AsyncClient(f\"http://localhost:{port}\")\n\n    def _inner_health(self):\n        raise NotImplementedError\n\n    async def health(self, timeout: int = 60):\n        assert timeout > 0\n        for _ in range(timeout):\n            if not self._inner_health():\n                raise RuntimeError(\"Launcher crashed\")\n\n            try:\n                await self.client.generate(\"test\")\n                return\n            except (ClientConnectorError, ClientOSError, ServerDisconnectedError) as e:\n                time.sleep(1)\n        raise RuntimeError(\"Health check failed\")\n\n\nclass ContainerLauncherHandle(LauncherHandle):\n    def __init__(self, docker_client, container_name, port: int):\n        super(ContainerLauncherHandle, self).__init__(port)\n        self.docker_client = docker_client\n        self.container_name = container_name\n\n    def _inner_health(self) -> bool:\n        container = self.docker_client.containers.get(self.container_name)\n        return container.status in [\"running\", \"created\"]\n\n\nclass ProcessLauncherHandle(LauncherHandle):\n    def __init__(self, process, port: int):\n        super(ProcessLauncherHandle, self).__init__(port)\n        self.process = process\n\n    def _inner_health(self) -> bool:\n        return self.process.poll() is None\n\n\n@pytest.fixture\ndef response_snapshot(snapshot):\n    return snapshot.use_extension(ResponseComparator)\n\n\n@pytest.fixture(scope=\"module\")\ndef event_loop():\n    loop = asyncio.get_event_loop()\n    yield loop\n    loop.close()\n\n\n@pytest.fixture(scope=\"module\")\ndef launcher(event_loop):\n    @contextlib.contextmanager\n    def local_launcher(\n        model_id: str,\n        num_shard: Optional[int] = None,\n        quantize: Optional[str] = None,\n        trust_remote_code: bool = False,\n        use_flash_attention: bool = True,\n    ):\n        port = random.randint(8000, 10_000)\n        master_port = random.randint(10_000, 20_000)\n\n        shard_uds_path = (\n            f\"/tmp/tgi-tests-{model_id.split('/')[-1]}-{num_shard}-{quantize}-server\"\n        )\n\n        args = [\n            \"text-generation-launcher\",\n            \"--model-id\",\n            model_id,\n            \"--port\",\n            str(port),\n            \"--master-port\",\n            str(master_port),\n            \"--shard-uds-path\",\n            shard_uds_path,\n        ]\n\n        env = os.environ\n\n        if num_shard is not None:\n            args.extend([\"--num-shard\", str(num_shard)])\n        if quantize is not None:\n            args.append(\"--quantize\")\n            args.append(quantize)\n        if trust_remote_code:\n            args.append(\"--trust-remote-code\")\n\n        env[\"LOG_LEVEL\"] = \"info,text_generation_router=debug\"\n\n        if not use_flash_attention:\n            env[\"USE_FLASH_ATTENTION\"] = \"false\"\n\n        with subprocess.Popen(\n            args, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=env\n        ) as process:\n            yield ProcessLauncherHandle(process, port)\n\n            process.terminate()\n            process.wait(60)\n\n            launcher_output = process.stdout.read().decode(\"utf-8\")\n            print(launcher_output, file=sys.stderr)\n\n            process.stdout.close()\n            process.stderr.close()\n\n        if not use_flash_attention:\n            del env[\"USE_FLASH_ATTENTION\"]\n\n    @contextlib.contextmanager\n    def docker_launcher(\n        model_id: str,\n        num_shard: Optional[int] = None,\n        quantize: Optional[str] = None,\n        trust_remote_code: bool = False,\n        use_flash_attention: bool = True,\n    ):\n        port = random.randint(8000, 10_000)\n\n        args = [\"--model-id\", model_id, \"--env\"]\n\n        if num_shard is not None:\n            args.extend([\"--num-shard\", str(num_shard)])\n        if quantize is not None:\n            args.append(\"--quantize\")\n            args.append(quantize)\n        if trust_remote_code:\n            args.append(\"--trust-remote-code\")\n\n        client = docker.from_env()\n\n        container_name = f\"tgi-tests-{model_id.split('/')[-1]}-{num_shard}-{quantize}\"\n\n        try:\n            container = client.containers.get(container_name)\n            container.stop()\n            container.wait()\n        except NotFound:\n            pass\n\n        gpu_count = num_shard if num_shard is not None else 1\n\n        env = {\"LOG_LEVEL\": \"info,text_generation_router=debug\"}\n        if not use_flash_attention:\n            env[\"USE_FLASH_ATTENTION\"] = \"false\"\n\n        if HUGGING_FACE_HUB_TOKEN is not None:\n            env[\"HUGGING_FACE_HUB_TOKEN\"] = HUGGING_FACE_HUB_TOKEN\n\n        volumes = []\n        if DOCKER_VOLUME:\n            volumes = [f\"{DOCKER_VOLUME}:/data\"]\n\n        container = client.containers.run(\n            DOCKER_IMAGE,\n            command=args,\n            name=container_name,\n            environment=env,\n            auto_remove=False,\n            detach=True,\n            device_requests=[\n                docker.types.DeviceRequest(count=gpu_count, capabilities=[[\"gpu\"]])\n            ],\n            volumes=volumes,\n            ports={\"80/tcp\": port},\n        )\n\n        yield ContainerLauncherHandle(client, container.name, port)\n\n        if not use_flash_attention:\n            del env[\"USE_FLASH_ATTENTION\"]\n\n        try:\n            container.stop()\n            container.wait()\n        except NotFound:\n            pass\n\n        container_output = container.logs().decode(\"utf-8\")\n        print(container_output, file=sys.stderr)\n\n        container.remove()\n\n    if DOCKER_IMAGE is not None:\n        return docker_launcher\n    return local_launcher\n\n\n@pytest.fixture(scope=\"module\")\ndef generate_load():\n    async def generate_load_inner(\n        client: AsyncClient, prompt: str, max_new_tokens: int, n: int\n    ) -> List[Response]:\n        futures = [\n            client.generate(\n                prompt, max_new_tokens=max_new_tokens, decoder_input_details=True\n            )\n            for _ in range(n)\n        ]\n\n        return await asyncio.gather(*futures)\n\n    return generate_load_inner\n", "path": "integration-tests/conftest.py"}], "after_files": [{"content": "import sys\nimport subprocess\nimport contextlib\nimport pytest\nimport asyncio\nimport os\nimport docker\nimport json\nimport math\nimport time\nimport random\n\nfrom docker.errors import NotFound\nfrom typing import Optional, List, Dict\nfrom syrupy.extensions.json import JSONSnapshotExtension\nfrom aiohttp import ClientConnectorError, ClientOSError, ServerDisconnectedError\n\nfrom text_generation import AsyncClient\nfrom text_generation.types import Response, Details, InputToken, Token, BestOfSequence\n\nDOCKER_IMAGE = os.getenv(\"DOCKER_IMAGE\", None)\nHUGGING_FACE_HUB_TOKEN = os.getenv(\"HUGGING_FACE_HUB_TOKEN\", None)\nDOCKER_VOLUME = os.getenv(\"DOCKER_VOLUME\", \"/data\")\n\n\nclass ResponseComparator(JSONSnapshotExtension):\n    def serialize(\n        self,\n        data,\n        *,\n        exclude=None,\n        matcher=None,\n    ):\n        if isinstance(data, List):\n            data = [d.dict() for d in data]\n\n        data = self._filter(\n            data=data, depth=0, path=(), exclude=exclude, matcher=matcher\n        )\n        return json.dumps(data, indent=2, ensure_ascii=False, sort_keys=False) + \"\\n\"\n\n    def matches(\n        self,\n        *,\n        serialized_data,\n        snapshot_data,\n    ) -> bool:\n        def convert_data(data):\n            data = json.loads(data)\n\n            if isinstance(data, Dict):\n                return Response(**data)\n            if isinstance(data, List):\n                return [Response(**d) for d in data]\n            raise NotImplementedError\n\n        def eq_token(token: Token, other: Token) -> bool:\n            return (\n                token.id == other.id\n                and token.text == other.text\n                and math.isclose(token.logprob, other.logprob, rel_tol=0.2)\n                and token.special == other.special\n            )\n\n        def eq_prefill_token(prefill_token: InputToken, other: InputToken) -> bool:\n            try:\n                return (\n                    prefill_token.id == other.id\n                    and prefill_token.text == other.text\n                    and (\n                        math.isclose(prefill_token.logprob, other.logprob, rel_tol=0.2)\n                        if prefill_token.logprob is not None\n                        else prefill_token.logprob == other.logprob\n                    )\n                )\n            except TypeError:\n                return False\n\n        def eq_best_of(details: BestOfSequence, other: BestOfSequence) -> bool:\n            return (\n                details.finish_reason == other.finish_reason\n                and details.generated_tokens == other.generated_tokens\n                and details.seed == other.seed\n                and len(details.prefill) == len(other.prefill)\n                and all(\n                    [\n                        eq_prefill_token(d, o)\n                        for d, o in zip(details.prefill, other.prefill)\n                    ]\n                )\n                and len(details.tokens) == len(other.tokens)\n                and all([eq_token(d, o) for d, o in zip(details.tokens, other.tokens)])\n            )\n\n        def eq_details(details: Details, other: Details) -> bool:\n            return (\n                details.finish_reason == other.finish_reason\n                and details.generated_tokens == other.generated_tokens\n                and details.seed == other.seed\n                and len(details.prefill) == len(other.prefill)\n                and all(\n                    [\n                        eq_prefill_token(d, o)\n                        for d, o in zip(details.prefill, other.prefill)\n                    ]\n                )\n                and len(details.tokens) == len(other.tokens)\n                and all([eq_token(d, o) for d, o in zip(details.tokens, other.tokens)])\n                and (\n                    len(details.best_of_sequences)\n                    if details.best_of_sequences is not None\n                    else 0\n                )\n                == (\n                    len(other.best_of_sequences)\n                    if other.best_of_sequences is not None\n                    else 0\n                )\n                and (\n                    all(\n                        [\n                            eq_best_of(d, o)\n                            for d, o in zip(\n                                details.best_of_sequences, other.best_of_sequences\n                            )\n                        ]\n                    )\n                    if details.best_of_sequences is not None\n                    else details.best_of_sequences == other.best_of_sequences\n                )\n            )\n\n        def eq_response(response: Response, other: Response) -> bool:\n            return response.generated_text == other.generated_text and eq_details(\n                response.details, other.details\n            )\n\n        serialized_data = convert_data(serialized_data)\n        snapshot_data = convert_data(snapshot_data)\n\n        if not isinstance(serialized_data, List):\n            serialized_data = [serialized_data]\n        if not isinstance(snapshot_data, List):\n            snapshot_data = [snapshot_data]\n\n        return len(snapshot_data) == len(serialized_data) and all(\n            [eq_response(r, o) for r, o in zip(serialized_data, snapshot_data)]\n        )\n\n\nclass LauncherHandle:\n    def __init__(self, port: int):\n        self.client = AsyncClient(f\"http://localhost:{port}\")\n\n    def _inner_health(self):\n        raise NotImplementedError\n\n    async def health(self, timeout: int = 60):\n        assert timeout > 0\n        for _ in range(timeout):\n            if not self._inner_health():\n                raise RuntimeError(\"Launcher crashed\")\n\n            try:\n                await self.client.generate(\"test\")\n                return\n            except (ClientConnectorError, ClientOSError, ServerDisconnectedError) as e:\n                time.sleep(1)\n        raise RuntimeError(\"Health check failed\")\n\n\nclass ContainerLauncherHandle(LauncherHandle):\n    def __init__(self, docker_client, container_name, port: int):\n        super(ContainerLauncherHandle, self).__init__(port)\n        self.docker_client = docker_client\n        self.container_name = container_name\n\n    def _inner_health(self) -> bool:\n        container = self.docker_client.containers.get(self.container_name)\n        return container.status in [\"running\", \"created\"]\n\n\nclass ProcessLauncherHandle(LauncherHandle):\n    def __init__(self, process, port: int):\n        super(ProcessLauncherHandle, self).__init__(port)\n        self.process = process\n\n    def _inner_health(self) -> bool:\n        return self.process.poll() is None\n\n\n@pytest.fixture\ndef response_snapshot(snapshot):\n    return snapshot.use_extension(ResponseComparator)\n\n\n@pytest.fixture(scope=\"module\")\ndef event_loop():\n    loop = asyncio.get_event_loop()\n    yield loop\n    loop.close()\n\n\n@pytest.fixture(scope=\"module\")\ndef launcher(event_loop):\n    @contextlib.contextmanager\n    def local_launcher(\n        model_id: str,\n        num_shard: Optional[int] = None,\n        quantize: Optional[str] = None,\n        trust_remote_code: bool = False,\n        use_flash_attention: bool = True,\n    ):\n        port = random.randint(8000, 10_000)\n        master_port = random.randint(10_000, 20_000)\n\n        shard_uds_path = (\n            f\"/tmp/tgi-tests-{model_id.split('/')[-1]}-{num_shard}-{quantize}-server\"\n        )\n\n        args = [\n            \"text-generation-launcher\",\n            \"--model-id\",\n            model_id,\n            \"--port\",\n            str(port),\n            \"--master-port\",\n            str(master_port),\n            \"--shard-uds-path\",\n            shard_uds_path,\n        ]\n\n        env = os.environ\n\n        if num_shard is not None:\n            args.extend([\"--num-shard\", str(num_shard)])\n        if quantize is not None:\n            args.append(\"--quantize\")\n            args.append(quantize)\n        if trust_remote_code:\n            args.append(\"--trust-remote-code\")\n\n        env[\"LOG_LEVEL\"] = \"info,text_generation_router=debug\"\n\n        if not use_flash_attention:\n            env[\"USE_FLASH_ATTENTION\"] = \"false\"\n\n        with subprocess.Popen(\n            args, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=env\n        ) as process:\n            yield ProcessLauncherHandle(process, port)\n\n            process.terminate()\n            process.wait(60)\n\n            launcher_output = process.stdout.read().decode(\"utf-8\")\n            print(launcher_output, file=sys.stderr)\n\n            process.stdout.close()\n            process.stderr.close()\n\n        if not use_flash_attention:\n            del env[\"USE_FLASH_ATTENTION\"]\n\n    @contextlib.contextmanager\n    def docker_launcher(\n        model_id: str,\n        num_shard: Optional[int] = None,\n        quantize: Optional[str] = None,\n        trust_remote_code: bool = False,\n        use_flash_attention: bool = True,\n    ):\n        port = random.randint(8000, 10_000)\n\n        args = [\"--model-id\", model_id, \"--env\"]\n\n        if num_shard is not None:\n            args.extend([\"--num-shard\", str(num_shard)])\n        if quantize is not None:\n            args.append(\"--quantize\")\n            args.append(quantize)\n        if trust_remote_code:\n            args.append(\"--trust-remote-code\")\n\n        client = docker.from_env()\n\n        container_name = f\"tgi-tests-{model_id.split('/')[-1]}-{num_shard}-{quantize}\"\n\n        try:\n            container = client.containers.get(container_name)\n            container.stop()\n            container.wait()\n        except NotFound:\n            pass\n\n        gpu_count = num_shard if num_shard is not None else 1\n\n        env = {\"LOG_LEVEL\": \"info,text_generation_router=debug\"}\n        if not use_flash_attention:\n            env[\"USE_FLASH_ATTENTION\"] = \"false\"\n\n        if HUGGING_FACE_HUB_TOKEN is not None:\n            env[\"HUGGING_FACE_HUB_TOKEN\"] = HUGGING_FACE_HUB_TOKEN\n\n        volumes = []\n        if DOCKER_VOLUME:\n            volumes = [f\"{DOCKER_VOLUME}:/data\"]\n\n        container = client.containers.run(\n            DOCKER_IMAGE,\n            command=args,\n            name=container_name,\n            environment=env,\n            auto_remove=False,\n            detach=True,\n            device_requests=[\n                docker.types.DeviceRequest(count=gpu_count, capabilities=[[\"gpu\"]])\n            ],\n            volumes=volumes,\n            ports={\"80/tcp\": port},\n            shm_size=\"1G\"\n        )\n\n        yield ContainerLauncherHandle(client, container.name, port)\n\n        if not use_flash_attention:\n            del env[\"USE_FLASH_ATTENTION\"]\n\n        try:\n            container.stop()\n            container.wait()\n        except NotFound:\n            pass\n\n        container_output = container.logs().decode(\"utf-8\")\n        print(container_output, file=sys.stderr)\n\n        container.remove()\n\n    if DOCKER_IMAGE is not None:\n        return docker_launcher\n    return local_launcher\n\n\n@pytest.fixture(scope=\"module\")\ndef generate_load():\n    async def generate_load_inner(\n        client: AsyncClient, prompt: str, max_new_tokens: int, n: int\n    ) -> List[Response]:\n        futures = [\n            client.generate(\n                prompt, max_new_tokens=max_new_tokens, decoder_input_details=True\n            )\n            for _ in range(n)\n        ]\n\n        return await asyncio.gather(*futures)\n\n    return generate_load_inner\n", "path": "integration-tests/conftest.py"}]}