We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Unable to stop TGI after serving models
### System Info

I use the official docker image: ghcr.io/huggingface/text-generation-inference:2.0.1

### Information

- [X] Docker
- [ ] The CLI directly

### Tasks

- [X] An officially supported command
- [ ] My own modifications

### Reproduction

I used the following command to serve the model. After TGI finished the model sharding/loading and started serving, I cannot use `Ctrl+C` to terminate the server. 

```bash
model=mistralai/Mixtral-8x7B-Instruct-v0.1
volume=/my_path_for_hf_cache
token="myhftokens"
    
docker run --gpus '"device=4,5"' \
    --shm-size 20g \
    -e HUGGING_FACE_HUB_TOKEN=$token \
    -p 8080:80 \
    -v $volume:/data ghcr.io/huggingface/text-generation-inference:2.0.1 \
    --model-id $model \
    --sharded true \
    --quantize eetq \
    --max-input-length 10240 \
    --max-batch-prefill-tokens 10240 \
    --max-total-tokens 32768 \
    --port 80
```

### Expected behavior

In previous version 1.3.0 and 1.4.0, I can use `Ctrl+C` to terminate the server while it is not the case for 2.0.1. My current solution is to use docker command to kill the container. Not sure if this is a good way?

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `server/text_generation_server/server.py`
Content:
```
1 import asyncio
2 import os
3 import torch
4 import time
5 import signal
6 
7 from grpc import aio
8 from loguru import logger
9 
10 from grpc_reflection.v1alpha import reflection
11 from pathlib import Path
12 from typing import List, Optional
13 
14 from text_generation_server.cache import Cache
15 from text_generation_server.interceptor import ExceptionInterceptor
16 from text_generation_server.models import Model, get_model
17 from text_generation_server.models.pali_gemma import PaliGemmaBatch
18 from text_generation_server.models.vlm_causal_lm import (
19     VlmCausalLMBatch,
20 )
21 from text_generation_server.pb import generate_pb2_grpc, generate_pb2
22 from text_generation_server.tracing import UDSOpenTelemetryAioServerInterceptor
23 from text_generation_server.models.idefics_causal_lm import IdeficsCausalLMBatch
24 
25 
26 class SignalHandler:
27     KEEP_PROCESSING = True
28 
29     def __init__(self):
30         signal.signal(signal.SIGINT, self.exit_gracefully)
31         signal.signal(signal.SIGTERM, self.exit_gracefully)
32 
33     def exit_gracefully(self, signum, frame):
34         print(f"Exiting gracefully: Signal {signum}")
35         self.KEEP_PROCESSING = False
36 
37 
38 signal_handler = SignalHandler()
39 
40 
41 class TextGenerationService(generate_pb2_grpc.TextGenerationServiceServicer):
42     def __init__(
43         self,
44         model: Model,
45         cache: Cache,
46         quantize: Optional[str],
47         server_urls: List[str],
48     ):
49         self.cache = cache
50         self.model = model
51         self.quantize = quantize
52         self.server_urls = server_urls
53         # For some reason, inference_mode does not work well with GLOO which we use on CPU
54         if model.device.type == "cuda":
55             # Force inference mode for the lifetime of TextGenerationService
56             self._inference_mode_raii_guard = torch._C._InferenceMode(True)
57 
58     async def Info(self, request, context):
59         return self.model.info
60 
61     async def Health(self, request, context):
62         if self.model.device.type == "cuda":
63             torch.zeros((2, 2)).cuda()
64         return generate_pb2.HealthResponse()
65 
66     async def ServiceDiscovery(self, request, context):
67         return generate_pb2.ServiceDiscoveryResponse(urls=self.server_urls)
68 
69     async def ClearCache(self, request, context):
70         if request.HasField("id"):
71             self.cache.delete(request.id)
72         else:
73             self.cache.clear()
74         return generate_pb2.ClearCacheResponse()
75 
76     async def FilterBatch(self, request, context):
77         batch = self.cache.pop(request.batch_id)
78         if batch is None:
79             raise ValueError(f"Batch ID {request.batch_id} not found in cache.")
80         filtered_batch = batch.filter(request.request_ids)
81         self.cache.set(filtered_batch)
82 
83         return generate_pb2.FilterBatchResponse(batch=filtered_batch.to_pb())
84 
85     async def Warmup(self, request, context):
86         if self.quantize == "gptq":
87             try:
88                 # When using GPTQ, Exllama kernels need some global kernels
89                 # For which we have the finale shapes only after the model has loaded
90                 # This will allocate those buffers.
91                 from text_generation_server.layers.gptq import (
92                     create_exllama_buffers,
93                     set_device,
94                 )
95 
96                 set_device(self.model.device)
97                 create_exllama_buffers(request.max_prefill_tokens)
98             except ImportError:
99                 pass
100 
101         if self.model.batch_type in {
102             IdeficsCausalLMBatch,
103             VlmCausalLMBatch,
104             PaliGemmaBatch,
105         }:  # Hack, i would rather use kwargs in the `from_pb` call
106             batch = self.model.batch_type.from_pb_processor(
107                 request.batch,
108                 self.model.tokenizer,
109                 self.model.processor,
110                 self.model.model.config,
111                 self.model.dtype,
112                 self.model.device,
113             )
114         else:
115             batch = self.model.batch_type.from_pb(
116                 request.batch, self.model.tokenizer, self.model.dtype, self.model.device
117             )
118         max_supported_total_tokens = self.model.warmup(batch)
119 
120         return generate_pb2.WarmupResponse(
121             max_supported_total_tokens=max_supported_total_tokens
122         )
123 
124     async def Prefill(self, request, context):
125         start = time.time_ns()
126         if self.model.batch_type in {
127             IdeficsCausalLMBatch,
128             VlmCausalLMBatch,
129             PaliGemmaBatch,
130         }:  # Hack, i would rather use kwargs in the `from_pb` call
131             batch = self.model.batch_type.from_pb_processor(
132                 request.batch,
133                 self.model.tokenizer,
134                 self.model.processor,
135                 self.model.model.config,
136                 self.model.dtype,
137                 self.model.device,
138             )
139         else:
140             batch = self.model.batch_type.from_pb(
141                 request.batch, self.model.tokenizer, self.model.dtype, self.model.device
142             )
143 
144         generations, next_batch, timings = self.model.generate_token(batch)
145         self.cache.set(next_batch)
146 
147         return generate_pb2.PrefillResponse(
148             generations=[generation.to_pb() for generation in generations],
149             batch=next_batch.to_pb() if next_batch else None,
150             forward_ns=timings[0],
151             decode_ns=timings[1],
152             total_ns=time.time_ns() - start,
153         )
154 
155     async def Decode(self, request, context):
156         start = time.time_ns()
157         if len(request.batches) == 0:
158             raise ValueError("Must provide at least one batch")
159 
160         batches = []
161         for batch_pb in request.batches:
162             batch = self.cache.pop(batch_pb.id)
163             if batch is None:
164                 raise ValueError(f"Batch ID {batch_pb.id} not found in cache.")
165             batches.append(batch)
166 
167         if len(batches) == 0:
168             raise ValueError("All batches are empty")
169 
170         if len(batches) > 1:
171             start_concat = time.time_ns()
172             batch = self.model.batch_type.concatenate(batches)
173             concat_ns = time.time_ns() - start_concat
174         else:
175             batch = batches[0]
176             concat_ns = None
177 
178         generations, next_batch, timings = self.model.generate_token(batch)
179         self.cache.set(next_batch)
180 
181         return generate_pb2.DecodeResponse(
182             generations=[generation.to_pb() for generation in generations],
183             batch=next_batch.to_pb() if next_batch else None,
184             concat_ns=concat_ns,
185             forward_ns=timings[0],
186             decode_ns=timings[1],
187             total_ns=time.time_ns() - start,
188         )
189 
190 
191 def serve(
192     model_id: str,
193     revision: Optional[str],
194     sharded: bool,
195     quantize: Optional[str],
196     speculate: Optional[int],
197     dtype: Optional[str],
198     trust_remote_code: bool,
199     uds_path: Path,
200 ):
201     async def serve_inner(
202         model_id: str,
203         revision: Optional[str],
204         sharded: bool = False,
205         quantize: Optional[str] = None,
206         speculate: Optional[int] = None,
207         dtype: Optional[str] = None,
208         trust_remote_code: bool = False,
209     ):
210         unix_socket_template = "unix://{}-{}"
211         if sharded:
212             server_urls = [
213                 unix_socket_template.format(uds_path, rank)
214                 for rank in range(int(os.environ["WORLD_SIZE"]))
215             ]
216             local_url = server_urls[int(os.environ["RANK"])]
217         else:
218             local_url = unix_socket_template.format(uds_path, 0)
219             server_urls = [local_url]
220 
221         try:
222             model = get_model(
223                 model_id,
224                 revision,
225                 sharded,
226                 quantize,
227                 speculate,
228                 dtype,
229                 trust_remote_code,
230             )
231         except Exception:
232             logger.exception("Error when initializing model")
233             raise
234 
235         server = aio.server(
236             interceptors=[
237                 ExceptionInterceptor(),
238                 UDSOpenTelemetryAioServerInterceptor(),
239             ]
240         )
241         generate_pb2_grpc.add_TextGenerationServiceServicer_to_server(
242             TextGenerationService(model, Cache(), quantize, server_urls), server
243         )
244         SERVICE_NAMES = (
245             generate_pb2.DESCRIPTOR.services_by_name["TextGenerationService"].full_name,
246             reflection.SERVICE_NAME,
247         )
248         reflection.enable_server_reflection(SERVICE_NAMES, server)
249         server.add_insecure_port(local_url)
250 
251         await server.start()
252 
253         logger.info("Server started at {}".format(local_url))
254 
255         while signal_handler.KEEP_PROCESSING:
256             await asyncio.sleep(0.5)
257 
258     asyncio.run(
259         serve_inner(
260             model_id, revision, sharded, quantize, speculate, dtype, trust_remote_code
261         )
262     )
263
```
--- 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/server/text_generation_server/server.py b/server/text_generation_server/server.py
--- a/server/text_generation_server/server.py
+++ b/server/text_generation_server/server.py
@@ -35,9 +35,6 @@
         self.KEEP_PROCESSING = False
 
 
-signal_handler = SignalHandler()
-
-
 class TextGenerationService(generate_pb2_grpc.TextGenerationServiceServicer):
     def __init__(
         self,
@@ -251,7 +248,7 @@
         await server.start()
 
         logger.info("Server started at {}".format(local_url))
-
+        signal_handler = SignalHandler()
         while signal_handler.KEEP_PROCESSING:
             await asyncio.sleep(0.5)

{"golden_diff": "diff --git a/server/text_generation_server/server.py b/server/text_generation_server/server.py\n--- a/server/text_generation_server/server.py\n+++ b/server/text_generation_server/server.py\n@@ -35,9 +35,6 @@\n         self.KEEP_PROCESSING = False\n \n \n-signal_handler = SignalHandler()\n-\n-\n class TextGenerationService(generate_pb2_grpc.TextGenerationServiceServicer):\n     def __init__(\n         self,\n@@ -251,7 +248,7 @@\n         await server.start()\n \n         logger.info(\"Server started at {}\".format(local_url))\n-\n+        signal_handler = SignalHandler()\n         while signal_handler.KEEP_PROCESSING:\n             await asyncio.sleep(0.5)\n", "issue": "Unable to stop TGI after serving models\n### System Info\n\nI use the official docker image: ghcr.io/huggingface/text-generation-inference:2.0.1\n\n### Information\n\n- [X] Docker\n- [ ] The CLI directly\n\n### Tasks\n\n- [X] An officially supported command\n- [ ] My own modifications\n\n### Reproduction\n\nI used the following command to serve the model. After TGI finished the model sharding/loading and started serving, I cannot use `Ctrl+C` to terminate the server. \r\n\r\n```bash\r\nmodel=mistralai/Mixtral-8x7B-Instruct-v0.1\r\nvolume=/my_path_for_hf_cache\r\ntoken=\"myhftokens\"\r\n    \r\ndocker run --gpus '\"device=4,5\"' \\\r\n    --shm-size 20g \\\r\n    -e HUGGING_FACE_HUB_TOKEN=$token \\\r\n    -p 8080:80 \\\r\n    -v $volume:/data ghcr.io/huggingface/text-generation-inference:2.0.1 \\\r\n    --model-id $model \\\r\n    --sharded true \\\r\n    --quantize eetq \\\r\n    --max-input-length 10240 \\\r\n    --max-batch-prefill-tokens 10240 \\\r\n    --max-total-tokens 32768 \\\r\n    --port 80\r\n```\n\n### Expected behavior\n\nIn previous version 1.3.0 and 1.4.0, I can use `Ctrl+C` to terminate the server while it is not the case for 2.0.1. My current solution is to use docker command to kill the container. Not sure if this is a good way?\n", "before_files": [{"content": "import asyncio\nimport os\nimport torch\nimport time\nimport signal\n\nfrom grpc import aio\nfrom loguru import logger\n\nfrom grpc_reflection.v1alpha import reflection\nfrom pathlib import Path\nfrom typing import List, Optional\n\nfrom text_generation_server.cache import Cache\nfrom text_generation_server.interceptor import ExceptionInterceptor\nfrom text_generation_server.models import Model, get_model\nfrom text_generation_server.models.pali_gemma import PaliGemmaBatch\nfrom text_generation_server.models.vlm_causal_lm import (\n    VlmCausalLMBatch,\n)\nfrom text_generation_server.pb import generate_pb2_grpc, generate_pb2\nfrom text_generation_server.tracing import UDSOpenTelemetryAioServerInterceptor\nfrom text_generation_server.models.idefics_causal_lm import IdeficsCausalLMBatch\n\n\nclass SignalHandler:\n    KEEP_PROCESSING = True\n\n    def __init__(self):\n        signal.signal(signal.SIGINT, self.exit_gracefully)\n        signal.signal(signal.SIGTERM, self.exit_gracefully)\n\n    def exit_gracefully(self, signum, frame):\n        print(f\"Exiting gracefully: Signal {signum}\")\n        self.KEEP_PROCESSING = False\n\n\nsignal_handler = SignalHandler()\n\n\nclass TextGenerationService(generate_pb2_grpc.TextGenerationServiceServicer):\n    def __init__(\n        self,\n        model: Model,\n        cache: Cache,\n        quantize: Optional[str],\n        server_urls: List[str],\n    ):\n        self.cache = cache\n        self.model = model\n        self.quantize = quantize\n        self.server_urls = server_urls\n        # For some reason, inference_mode does not work well with GLOO which we use on CPU\n        if model.device.type == \"cuda\":\n            # Force inference mode for the lifetime of TextGenerationService\n            self._inference_mode_raii_guard = torch._C._InferenceMode(True)\n\n    async def Info(self, request, context):\n        return self.model.info\n\n    async def Health(self, request, context):\n        if self.model.device.type == \"cuda\":\n            torch.zeros((2, 2)).cuda()\n        return generate_pb2.HealthResponse()\n\n    async def ServiceDiscovery(self, request, context):\n        return generate_pb2.ServiceDiscoveryResponse(urls=self.server_urls)\n\n    async def ClearCache(self, request, context):\n        if request.HasField(\"id\"):\n            self.cache.delete(request.id)\n        else:\n            self.cache.clear()\n        return generate_pb2.ClearCacheResponse()\n\n    async def FilterBatch(self, request, context):\n        batch = self.cache.pop(request.batch_id)\n        if batch is None:\n            raise ValueError(f\"Batch ID {request.batch_id} not found in cache.\")\n        filtered_batch = batch.filter(request.request_ids)\n        self.cache.set(filtered_batch)\n\n        return generate_pb2.FilterBatchResponse(batch=filtered_batch.to_pb())\n\n    async def Warmup(self, request, context):\n        if self.quantize == \"gptq\":\n            try:\n                # When using GPTQ, Exllama kernels need some global kernels\n                # For which we have the finale shapes only after the model has loaded\n                # This will allocate those buffers.\n                from text_generation_server.layers.gptq import (\n                    create_exllama_buffers,\n                    set_device,\n                )\n\n                set_device(self.model.device)\n                create_exllama_buffers(request.max_prefill_tokens)\n            except ImportError:\n                pass\n\n        if self.model.batch_type in {\n            IdeficsCausalLMBatch,\n            VlmCausalLMBatch,\n            PaliGemmaBatch,\n        }:  # Hack, i would rather use kwargs in the `from_pb` call\n            batch = self.model.batch_type.from_pb_processor(\n                request.batch,\n                self.model.tokenizer,\n                self.model.processor,\n                self.model.model.config,\n                self.model.dtype,\n                self.model.device,\n            )\n        else:\n            batch = self.model.batch_type.from_pb(\n                request.batch, self.model.tokenizer, self.model.dtype, self.model.device\n            )\n        max_supported_total_tokens = self.model.warmup(batch)\n\n        return generate_pb2.WarmupResponse(\n            max_supported_total_tokens=max_supported_total_tokens\n        )\n\n    async def Prefill(self, request, context):\n        start = time.time_ns()\n        if self.model.batch_type in {\n            IdeficsCausalLMBatch,\n            VlmCausalLMBatch,\n            PaliGemmaBatch,\n        }:  # Hack, i would rather use kwargs in the `from_pb` call\n            batch = self.model.batch_type.from_pb_processor(\n                request.batch,\n                self.model.tokenizer,\n                self.model.processor,\n                self.model.model.config,\n                self.model.dtype,\n                self.model.device,\n            )\n        else:\n            batch = self.model.batch_type.from_pb(\n                request.batch, self.model.tokenizer, self.model.dtype, self.model.device\n            )\n\n        generations, next_batch, timings = self.model.generate_token(batch)\n        self.cache.set(next_batch)\n\n        return generate_pb2.PrefillResponse(\n            generations=[generation.to_pb() for generation in generations],\n            batch=next_batch.to_pb() if next_batch else None,\n            forward_ns=timings[0],\n            decode_ns=timings[1],\n            total_ns=time.time_ns() - start,\n        )\n\n    async def Decode(self, request, context):\n        start = time.time_ns()\n        if len(request.batches) == 0:\n            raise ValueError(\"Must provide at least one batch\")\n\n        batches = []\n        for batch_pb in request.batches:\n            batch = self.cache.pop(batch_pb.id)\n            if batch is None:\n                raise ValueError(f\"Batch ID {batch_pb.id} not found in cache.\")\n            batches.append(batch)\n\n        if len(batches) == 0:\n            raise ValueError(\"All batches are empty\")\n\n        if len(batches) > 1:\n            start_concat = time.time_ns()\n            batch = self.model.batch_type.concatenate(batches)\n            concat_ns = time.time_ns() - start_concat\n        else:\n            batch = batches[0]\n            concat_ns = None\n\n        generations, next_batch, timings = self.model.generate_token(batch)\n        self.cache.set(next_batch)\n\n        return generate_pb2.DecodeResponse(\n            generations=[generation.to_pb() for generation in generations],\n            batch=next_batch.to_pb() if next_batch else None,\n            concat_ns=concat_ns,\n            forward_ns=timings[0],\n            decode_ns=timings[1],\n            total_ns=time.time_ns() - start,\n        )\n\n\ndef serve(\n    model_id: str,\n    revision: Optional[str],\n    sharded: bool,\n    quantize: Optional[str],\n    speculate: Optional[int],\n    dtype: Optional[str],\n    trust_remote_code: bool,\n    uds_path: Path,\n):\n    async def serve_inner(\n        model_id: str,\n        revision: Optional[str],\n        sharded: bool = False,\n        quantize: Optional[str] = None,\n        speculate: Optional[int] = None,\n        dtype: Optional[str] = None,\n        trust_remote_code: bool = False,\n    ):\n        unix_socket_template = \"unix://{}-{}\"\n        if sharded:\n            server_urls = [\n                unix_socket_template.format(uds_path, rank)\n                for rank in range(int(os.environ[\"WORLD_SIZE\"]))\n            ]\n            local_url = server_urls[int(os.environ[\"RANK\"])]\n        else:\n            local_url = unix_socket_template.format(uds_path, 0)\n            server_urls = [local_url]\n\n        try:\n            model = get_model(\n                model_id,\n                revision,\n                sharded,\n                quantize,\n                speculate,\n                dtype,\n                trust_remote_code,\n            )\n        except Exception:\n            logger.exception(\"Error when initializing model\")\n            raise\n\n        server = aio.server(\n            interceptors=[\n                ExceptionInterceptor(),\n                UDSOpenTelemetryAioServerInterceptor(),\n            ]\n        )\n        generate_pb2_grpc.add_TextGenerationServiceServicer_to_server(\n            TextGenerationService(model, Cache(), quantize, server_urls), server\n        )\n        SERVICE_NAMES = (\n            generate_pb2.DESCRIPTOR.services_by_name[\"TextGenerationService\"].full_name,\n            reflection.SERVICE_NAME,\n        )\n        reflection.enable_server_reflection(SERVICE_NAMES, server)\n        server.add_insecure_port(local_url)\n\n        await server.start()\n\n        logger.info(\"Server started at {}\".format(local_url))\n\n        while signal_handler.KEEP_PROCESSING:\n            await asyncio.sleep(0.5)\n\n    asyncio.run(\n        serve_inner(\n            model_id, revision, sharded, quantize, speculate, dtype, trust_remote_code\n        )\n    )\n", "path": "server/text_generation_server/server.py"}], "after_files": [{"content": "import asyncio\nimport os\nimport torch\nimport time\nimport signal\n\nfrom grpc import aio\nfrom loguru import logger\n\nfrom grpc_reflection.v1alpha import reflection\nfrom pathlib import Path\nfrom typing import List, Optional\n\nfrom text_generation_server.cache import Cache\nfrom text_generation_server.interceptor import ExceptionInterceptor\nfrom text_generation_server.models import Model, get_model\nfrom text_generation_server.models.pali_gemma import PaliGemmaBatch\nfrom text_generation_server.models.vlm_causal_lm import (\n    VlmCausalLMBatch,\n)\nfrom text_generation_server.pb import generate_pb2_grpc, generate_pb2\nfrom text_generation_server.tracing import UDSOpenTelemetryAioServerInterceptor\nfrom text_generation_server.models.idefics_causal_lm import IdeficsCausalLMBatch\n\n\nclass SignalHandler:\n    KEEP_PROCESSING = True\n\n    def __init__(self):\n        signal.signal(signal.SIGINT, self.exit_gracefully)\n        signal.signal(signal.SIGTERM, self.exit_gracefully)\n\n    def exit_gracefully(self, signum, frame):\n        print(f\"Exiting gracefully: Signal {signum}\")\n        self.KEEP_PROCESSING = False\n\n\nclass TextGenerationService(generate_pb2_grpc.TextGenerationServiceServicer):\n    def __init__(\n        self,\n        model: Model,\n        cache: Cache,\n        quantize: Optional[str],\n        server_urls: List[str],\n    ):\n        self.cache = cache\n        self.model = model\n        self.quantize = quantize\n        self.server_urls = server_urls\n        # For some reason, inference_mode does not work well with GLOO which we use on CPU\n        if model.device.type == \"cuda\":\n            # Force inference mode for the lifetime of TextGenerationService\n            self._inference_mode_raii_guard = torch._C._InferenceMode(True)\n\n    async def Info(self, request, context):\n        return self.model.info\n\n    async def Health(self, request, context):\n        if self.model.device.type == \"cuda\":\n            torch.zeros((2, 2)).cuda()\n        return generate_pb2.HealthResponse()\n\n    async def ServiceDiscovery(self, request, context):\n        return generate_pb2.ServiceDiscoveryResponse(urls=self.server_urls)\n\n    async def ClearCache(self, request, context):\n        if request.HasField(\"id\"):\n            self.cache.delete(request.id)\n        else:\n            self.cache.clear()\n        return generate_pb2.ClearCacheResponse()\n\n    async def FilterBatch(self, request, context):\n        batch = self.cache.pop(request.batch_id)\n        if batch is None:\n            raise ValueError(f\"Batch ID {request.batch_id} not found in cache.\")\n        filtered_batch = batch.filter(request.request_ids)\n        self.cache.set(filtered_batch)\n\n        return generate_pb2.FilterBatchResponse(batch=filtered_batch.to_pb())\n\n    async def Warmup(self, request, context):\n        if self.quantize == \"gptq\":\n            try:\n                # When using GPTQ, Exllama kernels need some global kernels\n                # For which we have the finale shapes only after the model has loaded\n                # This will allocate those buffers.\n                from text_generation_server.layers.gptq import (\n                    create_exllama_buffers,\n                    set_device,\n                )\n\n                set_device(self.model.device)\n                create_exllama_buffers(request.max_prefill_tokens)\n            except ImportError:\n                pass\n\n        if self.model.batch_type in {\n            IdeficsCausalLMBatch,\n            VlmCausalLMBatch,\n            PaliGemmaBatch,\n        }:  # Hack, i would rather use kwargs in the `from_pb` call\n            batch = self.model.batch_type.from_pb_processor(\n                request.batch,\n                self.model.tokenizer,\n                self.model.processor,\n                self.model.model.config,\n                self.model.dtype,\n                self.model.device,\n            )\n        else:\n            batch = self.model.batch_type.from_pb(\n                request.batch, self.model.tokenizer, self.model.dtype, self.model.device\n            )\n        max_supported_total_tokens = self.model.warmup(batch)\n\n        return generate_pb2.WarmupResponse(\n            max_supported_total_tokens=max_supported_total_tokens\n        )\n\n    async def Prefill(self, request, context):\n        start = time.time_ns()\n        if self.model.batch_type in {\n            IdeficsCausalLMBatch,\n            VlmCausalLMBatch,\n            PaliGemmaBatch,\n        }:  # Hack, i would rather use kwargs in the `from_pb` call\n            batch = self.model.batch_type.from_pb_processor(\n                request.batch,\n                self.model.tokenizer,\n                self.model.processor,\n                self.model.model.config,\n                self.model.dtype,\n                self.model.device,\n            )\n        else:\n            batch = self.model.batch_type.from_pb(\n                request.batch, self.model.tokenizer, self.model.dtype, self.model.device\n            )\n\n        generations, next_batch, timings = self.model.generate_token(batch)\n        self.cache.set(next_batch)\n\n        return generate_pb2.PrefillResponse(\n            generations=[generation.to_pb() for generation in generations],\n            batch=next_batch.to_pb() if next_batch else None,\n            forward_ns=timings[0],\n            decode_ns=timings[1],\n            total_ns=time.time_ns() - start,\n        )\n\n    async def Decode(self, request, context):\n        start = time.time_ns()\n        if len(request.batches) == 0:\n            raise ValueError(\"Must provide at least one batch\")\n\n        batches = []\n        for batch_pb in request.batches:\n            batch = self.cache.pop(batch_pb.id)\n            if batch is None:\n                raise ValueError(f\"Batch ID {batch_pb.id} not found in cache.\")\n            batches.append(batch)\n\n        if len(batches) == 0:\n            raise ValueError(\"All batches are empty\")\n\n        if len(batches) > 1:\n            start_concat = time.time_ns()\n            batch = self.model.batch_type.concatenate(batches)\n            concat_ns = time.time_ns() - start_concat\n        else:\n            batch = batches[0]\n            concat_ns = None\n\n        generations, next_batch, timings = self.model.generate_token(batch)\n        self.cache.set(next_batch)\n\n        return generate_pb2.DecodeResponse(\n            generations=[generation.to_pb() for generation in generations],\n            batch=next_batch.to_pb() if next_batch else None,\n            concat_ns=concat_ns,\n            forward_ns=timings[0],\n            decode_ns=timings[1],\n            total_ns=time.time_ns() - start,\n        )\n\n\ndef serve(\n    model_id: str,\n    revision: Optional[str],\n    sharded: bool,\n    quantize: Optional[str],\n    speculate: Optional[int],\n    dtype: Optional[str],\n    trust_remote_code: bool,\n    uds_path: Path,\n):\n    async def serve_inner(\n        model_id: str,\n        revision: Optional[str],\n        sharded: bool = False,\n        quantize: Optional[str] = None,\n        speculate: Optional[int] = None,\n        dtype: Optional[str] = None,\n        trust_remote_code: bool = False,\n    ):\n        unix_socket_template = \"unix://{}-{}\"\n        if sharded:\n            server_urls = [\n                unix_socket_template.format(uds_path, rank)\n                for rank in range(int(os.environ[\"WORLD_SIZE\"]))\n            ]\n            local_url = server_urls[int(os.environ[\"RANK\"])]\n        else:\n            local_url = unix_socket_template.format(uds_path, 0)\n            server_urls = [local_url]\n\n        try:\n            model = get_model(\n                model_id,\n                revision,\n                sharded,\n                quantize,\n                speculate,\n                dtype,\n                trust_remote_code,\n            )\n        except Exception:\n            logger.exception(\"Error when initializing model\")\n            raise\n\n        server = aio.server(\n            interceptors=[\n                ExceptionInterceptor(),\n                UDSOpenTelemetryAioServerInterceptor(),\n            ]\n        )\n        generate_pb2_grpc.add_TextGenerationServiceServicer_to_server(\n            TextGenerationService(model, Cache(), quantize, server_urls), server\n        )\n        SERVICE_NAMES = (\n            generate_pb2.DESCRIPTOR.services_by_name[\"TextGenerationService\"].full_name,\n            reflection.SERVICE_NAME,\n        )\n        reflection.enable_server_reflection(SERVICE_NAMES, server)\n        server.add_insecure_port(local_url)\n\n        await server.start()\n\n        logger.info(\"Server started at {}\".format(local_url))\n        signal_handler = SignalHandler()\n        while signal_handler.KEEP_PROCESSING:\n            await asyncio.sleep(0.5)\n\n    asyncio.run(\n        serve_inner(\n            model_id, revision, sharded, quantize, speculate, dtype, trust_remote_code\n        )\n    )\n", "path": "server/text_generation_server/server.py"}]}