server.py 5.76 KB
Newer Older
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
1
import asyncio
Olivier Dehaene's avatar
Olivier Dehaene committed
2
import os
3
import torch
Olivier Dehaene's avatar
Olivier Dehaene committed
4

Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
5
from grpc import aio
6
from loguru import logger
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
7
8
9

from grpc_reflection.v1alpha import reflection
from pathlib import Path
10
from typing import List, Optional
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
11

12
13
14
15
16
from text_generation_server.cache import Cache
from text_generation_server.interceptor import ExceptionInterceptor
from text_generation_server.models import Model, get_model
from text_generation_server.pb import generate_pb2_grpc, generate_pb2
from text_generation_server.tracing import UDSOpenTelemetryAioServerInterceptor
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
17
18


Olivier Dehaene's avatar
Olivier Dehaene committed
19
class TextGenerationService(generate_pb2_grpc.TextGenerationServiceServicer):
20
    def __init__(self, model: Model, cache: Cache, server_urls: List[str]):
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
21
22
23
        self.cache = cache
        self.model = model
        self.server_urls = server_urls
24
25
26
27
        # For some reason, inference_mode does not work well with GLOO which we use on CPU
        if model.device.type == "cuda":
            # Force inference mode for the lifetime of TextGenerationService
            self._inference_mode_raii_guard = torch._C._InferenceMode(True)
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
28

29
30
31
    async def Info(self, request, context):
        return self.model.info

32
33
34
35
36
    async def Health(self, request, context):
        if self.model.device.type == "cuda":
            torch.zeros((2, 2)).cuda()
        return generate_pb2.HealthResponse()

Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
37
38
39
40
    async def ServiceDiscovery(self, request, context):
        return generate_pb2.ServiceDiscoveryResponse(urls=self.server_urls)

    async def ClearCache(self, request, context):
41
42
43
44
        if request.HasField("id"):
            self.cache.delete(request.id)
        else:
            self.cache.clear()
Olivier Dehaene's avatar
Olivier Dehaene committed
45
        return generate_pb2.ClearCacheResponse()
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
46

47
48
49
50
    async def FilterBatch(self, request, context):
        batch = self.cache.pop(request.batch_id)
        if batch is None:
            raise ValueError(f"Batch ID {request.batch_id} not found in cache.")
51
        filtered_batch = batch.filter(request.request_ids)
52
53
54
55
        self.cache.set(filtered_batch)

        return generate_pb2.FilterBatchResponse(batch=filtered_batch.to_pb())

56
57
58
59
    async def Warmup(self, request, context):
        batch = self.model.batch_type.from_pb(
            request.batch, self.model.tokenizer, self.model.dtype, self.model.device
        )
60
        max_supported_total_tokens = self.model.warmup(batch)
61

62
63
64
        return generate_pb2.WarmupResponse(
            max_supported_total_tokens=max_supported_total_tokens
        )
65

66
    async def Prefill(self, request, context):
67
        batch = self.model.batch_type.from_pb(
68
            request.batch, self.model.tokenizer, self.model.dtype, self.model.device
69
        )
Olivier Dehaene's avatar
Olivier Dehaene committed
70

71
        generations, next_batch = self.model.generate_token(batch)
Olivier Dehaene's avatar
Olivier Dehaene committed
72
73
        self.cache.set(next_batch)

74
75
        return generate_pb2.PrefillResponse(
            generations=[generation.to_pb() for generation in generations],
Olivier Dehaene's avatar
Olivier Dehaene committed
76
            batch=next_batch.to_pb() if next_batch else None,
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
77
78
        )

79
    async def Decode(self, request, context):
Olivier Dehaene's avatar
Olivier Dehaene committed
80
81
82
83
84
85
86
87
        if len(request.batches) == 0:
            raise ValueError("Must provide at least one batch")

        batches = []
        for batch_pb in request.batches:
            batch = self.cache.pop(batch_pb.id)
            if batch is None:
                raise ValueError(f"Batch ID {batch_pb.id} not found in cache.")
88
            batches.append(batch)
89
90
91

        if len(batches) == 0:
            raise ValueError("All batches are empty")
Olivier Dehaene's avatar
Olivier Dehaene committed
92
93

        if len(batches) > 1:
94
            batch = self.model.batch_type.concatenate(batches)
Olivier Dehaene's avatar
Olivier Dehaene committed
95
96
97
        else:
            batch = batches[0]

98
        generations, next_batch = self.model.generate_token(batch)
Olivier Dehaene's avatar
Olivier Dehaene committed
99
100
        self.cache.set(next_batch)

101
102
        return generate_pb2.DecodeResponse(
            generations=[generation.to_pb() for generation in generations],
Olivier Dehaene's avatar
Olivier Dehaene committed
103
104
105
            batch=next_batch.to_pb() if next_batch else None,
        )

Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
106

Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
107
def serve(
108
    model_id: str,
109
    revision: Optional[str],
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
110
    sharded: bool,
111
    quantize: Optional[str],
112
    dtype: Optional[str],
113
    trust_remote_code: bool,
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
114
115
    uds_path: Path,
):
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
116
    async def serve_inner(
117
        model_id: str,
118
        revision: Optional[str],
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
119
        sharded: bool = False,
120
        quantize: Optional[str] = None,
121
        dtype: Optional[str] = None,
122
        trust_remote_code: bool = False,
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
123
    ):
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
124
        unix_socket_template = "unix://{}-{}"
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
125
126
        if sharded:
            server_urls = [
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
127
                unix_socket_template.format(uds_path, rank)
128
                for rank in range(int(os.environ["WORLD_SIZE"]))
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
129
            ]
130
            local_url = server_urls[int(os.environ["RANK"])]
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
131
        else:
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
132
            local_url = unix_socket_template.format(uds_path, 0)
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
133
134
            server_urls = [local_url]

135
        try:
136
137
138
            model = get_model(
                model_id, revision, sharded, quantize, dtype, trust_remote_code
            )
139
140
141
        except Exception:
            logger.exception("Error when initializing model")
            raise
142

143
144
145
146
147
148
        server = aio.server(
            interceptors=[
                ExceptionInterceptor(),
                UDSOpenTelemetryAioServerInterceptor(),
            ]
        )
Olivier Dehaene's avatar
Olivier Dehaene committed
149
150
        generate_pb2_grpc.add_TextGenerationServiceServicer_to_server(
            TextGenerationService(model, Cache(), server_urls), server
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
151
152
        )
        SERVICE_NAMES = (
Olivier Dehaene's avatar
Olivier Dehaene committed
153
            generate_pb2.DESCRIPTOR.services_by_name["TextGenerationService"].full_name,
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
154
155
156
157
            reflection.SERVICE_NAME,
        )
        reflection.enable_server_reflection(SERVICE_NAMES, server)
        server.add_insecure_port(local_url)
158

Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
159
        await server.start()
160

161
        logger.info("Server started at {}".format(local_url))
162

Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
163
164
165
        try:
            await server.wait_for_termination()
        except KeyboardInterrupt:
166
            logger.info("Signal received. Shutting down")
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
167
            await server.stop(0)
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
168

169
170
171
    asyncio.run(
        serve_inner(model_id, revision, sharded, quantize, dtype, trust_remote_code)
    )