tokenizer_manager.py 9.4 KB
Newer Older
Lianmin Zheng's avatar
Lianmin Zheng committed
1
2
3
import asyncio
import concurrent.futures
import dataclasses
4
import multiprocessing as mp
Lianmin Zheng's avatar
Lianmin Zheng committed
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
import os
from typing import List

import numpy as np
import transformers
import uvloop
import zmq
import zmq.asyncio
from sglang.srt.hf_transformers_utils import (
    get_config,
    get_context_length,
    get_processor,
    get_tokenizer,
)
from sglang.srt.managers.io_struct import (
    BatchStrOut,
Cody Yu's avatar
Cody Yu committed
21
    DetokenizeReqInput,
22
    FlushCacheReq,
Lianmin Zheng's avatar
Lianmin Zheng committed
23
24
25
    GenerateReqInput,
    TokenizedGenerateReqInput,
)
shiyi.c_98's avatar
shiyi.c_98 committed
26
from sglang.srt.mm_utils import expand2square, process_anyres_image
Lianmin Zheng's avatar
Lianmin Zheng committed
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
from sglang.srt.sampling_params import SamplingParams
from sglang.srt.server_args import PortArgs, ServerArgs
from sglang.srt.utils import get_exception_traceback, is_multimodal_model, load_image

asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())


@dataclasses.dataclass
class ReqState:
    out_list: List
    finished: bool
    event: asyncio.Event


global global_processor


def init_global_processor(server_args: ServerArgs):
    global global_processor
    transformers.logging.set_verbosity_error()
    global_processor = get_processor(
        server_args.tokenizer_path,
        tokenizer_mode=server_args.tokenizer_mode,
        trust_remote_code=server_args.trust_remote_code,
    )


54
55
56
def get_pixel_values(
    image_data, image_aspect_ratio=None, image_grid_pinpoints=None, processor=None
):
Lianmin Zheng's avatar
Lianmin Zheng committed
57
58
59
60
    try:
        processor = processor or global_processor
        image = load_image(image_data)
        image_hash = hash(image_data)
shiyi.c_98's avatar
shiyi.c_98 committed
61
62
63
64
65
66
67
        if image_aspect_ratio == "pad":
            image = expand2square(
                image, tuple(int(x * 255) for x in processor.image_processor.image_mean)
            )
            pixel_values = processor.image_processor(image)["pixel_values"][0]
        elif image_aspect_ratio == "anyres":
            pixel_values = process_anyres_image(
Ying Sheng's avatar
Ying Sheng committed
68
                image, processor.image_processor, image_grid_pinpoints
shiyi.c_98's avatar
shiyi.c_98 committed
69
70
71
            )
        else:
            pixel_values = processor.image_processor(image)["pixel_values"][0]
Lianmin Zheng's avatar
Lianmin Zheng committed
72
        pixel_values = pixel_values.astype(np.float16)
shiyi.c_98's avatar
shiyi.c_98 committed
73
        return pixel_values, image_hash, image.size
Lianmin Zheng's avatar
Lianmin Zheng committed
74
75
76
77
78
79
80
81
82
83
    except Exception:
        print("Exception in TokenizerManager:\n" + get_exception_traceback())


class TokenizerManager:
    def __init__(
        self,
        server_args: ServerArgs,
        port_args: PortArgs,
    ):
Liangsheng Yin's avatar
Liangsheng Yin committed
84
85
        self.server_args = server_args

Lianmin Zheng's avatar
Lianmin Zheng committed
86
87
88
89
90
91
92
93
94
95
96
        context = zmq.asyncio.Context(2)
        self.recv_from_detokenizer = context.socket(zmq.PULL)
        self.recv_from_detokenizer.bind(f"tcp://127.0.0.1:{port_args.tokenizer_port}")

        self.send_to_router = context.socket(zmq.PUSH)
        self.send_to_router.connect(f"tcp://127.0.0.1:{port_args.router_port}")

        self.model_path = server_args.model_path
        self.hf_config = get_config(
            self.model_path, trust_remote_code=server_args.trust_remote_code
        )
shiyi.c_98's avatar
shiyi.c_98 committed
97

Lianmin Zheng's avatar
Lianmin Zheng committed
98
99
100
101
102
103
104
105
106
107
108
        self.context_len = get_context_length(self.hf_config)

        if is_multimodal_model(self.model_path):
            self.processor = get_processor(
                server_args.tokenizer_path,
                tokenizer_mode=server_args.tokenizer_mode,
                trust_remote_code=server_args.trust_remote_code,
            )
            self.tokenizer = self.processor.tokenizer
            os.environ["TOKENIZERS_PARALLELISM"] = "false"
            self.executor = concurrent.futures.ProcessPoolExecutor(
109
110
111
                initializer=init_global_processor,
                mp_context=mp.get_context("fork"),
                initargs=(server_args,),
Lianmin Zheng's avatar
Lianmin Zheng committed
112
113
114
115
116
117
118
119
120
121
122
123
            )
        else:
            self.tokenizer = get_tokenizer(
                server_args.tokenizer_path,
                tokenizer_mode=server_args.tokenizer_mode,
                trust_remote_code=server_args.trust_remote_code,
            )

        self.to_create_loop = True
        self.rid_to_state = {}  # Dict[str -> ReqState]

    async def get_pixel_values(self, image_data):
Ying Sheng's avatar
Ying Sheng committed
124
        aspect_ratio = getattr(self.hf_config, "image_aspect_ratio", None)
125
126
127
        grid_pinpoints = (
            self.hf_config.image_grid_pinpoints if aspect_ratio == "anyres" else None
        )
Lianmin Zheng's avatar
Lianmin Zheng committed
128
129
130
        if self.executor is not None:
            loop = asyncio.get_event_loop()
            return await loop.run_in_executor(
131
132
133
134
135
                self.executor,
                get_pixel_values,
                image_data,
                aspect_ratio,
                grid_pinpoints,
Lianmin Zheng's avatar
Lianmin Zheng committed
136
137
            )
        else:
138
139
140
            return get_pixel_values(
                image_data, aspect_ratio, grid_pinpoints, self.processor
            )
Lianmin Zheng's avatar
Lianmin Zheng committed
141
142
143
144
145
146
147
148
149
150
151
152
153
154

    async def generate_request(self, obj: GenerateReqInput):
        if self.to_create_loop:
            await self.create_handle_loop()

        is_single = isinstance(obj.text, str)

        if is_single:
            rid = obj.rid
            input_ids = self.tokenizer.encode(obj.text)
            sampling_params = SamplingParams(**obj.sampling_params)
            if sampling_params.max_new_tokens != 0:
                sampling_params.normalize(self.tokenizer)
                sampling_params.verify()
155
156
157
158
159
160

            if isinstance(obj.image_data, list) and len(obj.image_data) > 0:
                pixel_values, image_hash, image_size = await self.get_pixel_values(
                    obj.image_data[0]
                )
            elif isinstance(obj.image_data, str):
shiyi.c_98's avatar
shiyi.c_98 committed
161
162
163
                pixel_values, image_hash, image_size = await self.get_pixel_values(
                    obj.image_data
                )
164
165
            else:
                pixel_values, image_hash, image_size = None, None, None
Lianmin Zheng's avatar
Lianmin Zheng committed
166
167
            tokenized_obj = TokenizedGenerateReqInput(
                rid=rid,
Liangsheng Yin's avatar
Liangsheng Yin committed
168
                input_text=obj.text,
Lianmin Zheng's avatar
Lianmin Zheng committed
169
170
171
                input_ids=input_ids,
                pixel_values=pixel_values,
                image_hash=image_hash,
shiyi.c_98's avatar
shiyi.c_98 committed
172
                image_size=image_size,
Lianmin Zheng's avatar
Lianmin Zheng committed
173
                sampling_params=sampling_params,
174
175
                return_logprob=obj.return_logprob,
                logprob_start_len=obj.logprob_start_len,
Lianmin Zheng's avatar
Lianmin Zheng committed
176
177
178
179
180
                stream=obj.stream,
            )
            self.send_to_router.send_pyobj(tokenized_obj)

            event = asyncio.Event()
181
            state = ReqState([], False, event)
Lianmin Zheng's avatar
Lianmin Zheng committed
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
            self.rid_to_state[rid] = state

            while True:
                await event.wait()
                yield state.out_list[-1]
                state.out_list = []
                if state.finished:
                    del self.rid_to_state[rid]
                    break
                event.clear()
        else:
            assert obj.stream is False
            bs = len(obj.text)
            for i in range(bs):
                rid = obj.rid[i]
                input_ids = self.tokenizer.encode(obj.text[i])
                sampling_params = SamplingParams(**obj.sampling_params[i])
                if sampling_params.max_new_tokens != 0:
                    sampling_params.normalize(self.tokenizer)
                    sampling_params.verify()
                if obj.image_data[i] is None:
shiyi.c_98's avatar
shiyi.c_98 committed
203
                    pixel_values, image_hash, image_size = None, None, None
Lianmin Zheng's avatar
Lianmin Zheng committed
204
                else:
shiyi.c_98's avatar
shiyi.c_98 committed
205
                    pixel_values, image_hash, image_size = await self.get_pixel_values(
Lianmin Zheng's avatar
Lianmin Zheng committed
206
207
208
209
                        obj.image_data[i]
                    )
                tokenized_obj = TokenizedGenerateReqInput(
                    rid=rid,
210
                    input_text=obj.text[i],
Lianmin Zheng's avatar
Lianmin Zheng committed
211
212
213
                    input_ids=input_ids,
                    pixel_values=pixel_values,
                    image_hash=image_hash,
shiyi.c_98's avatar
shiyi.c_98 committed
214
                    image_size=image_size,
Lianmin Zheng's avatar
Lianmin Zheng committed
215
                    sampling_params=sampling_params,
216
217
                    return_logprob=obj.return_logprob[i],
                    logprob_start_len=obj.logprob_start_len[i],
Lianmin Zheng's avatar
Lianmin Zheng committed
218
219
220
221
222
                    stream=obj.stream,
                )
                self.send_to_router.send_pyobj(tokenized_obj)

                event = asyncio.Event()
223
                state = ReqState([], False, event)
Lianmin Zheng's avatar
Lianmin Zheng committed
224
225
226
227
228
229
230
231
232
233
234
235
236
                self.rid_to_state[rid] = state

            output_list = []
            for i in range(bs):
                rid = obj.rid[i]
                state = self.rid_to_state[rid]
                await state.event.wait()
                output_list.append(state.out_list[-1])
                assert state.finished
                del self.rid_to_state[rid]

            yield output_list

Cody Yu's avatar
Cody Yu committed
237
238
239
240
    async def detokenize(self, obj: DetokenizeReqInput):
        token_texts = self.tokenizer.convert_ids_to_tokens(obj.input_ids)
        return [t.decode() if isinstance(t, bytes) else t for t in token_texts]

Liangsheng Yin's avatar
Liangsheng Yin committed
241
242
243
244
    async def flush_cache(self):
        flush_cache_req = FlushCacheReq()
        self.send_to_router.send_pyobj(flush_cache_req)

Lianmin Zheng's avatar
Lianmin Zheng committed
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
    async def create_handle_loop(self):
        self.to_create_loop = False
        loop = asyncio.get_event_loop()
        loop.create_task(self.handle_loop())

    async def handle_loop(self):
        while True:
            recv_obj = await self.recv_from_detokenizer.recv_pyobj()

            if isinstance(recv_obj, BatchStrOut):
                for i, rid in enumerate(recv_obj.rids):
                    recv_obj.meta_info[i]["id"] = rid
                    out_dict = {
                        "text": recv_obj.output_str[i],
                        "meta_info": recv_obj.meta_info[i],
                    }
                    state = self.rid_to_state[rid]
                    state.out_list.append(out_dict)
                    state.finished = recv_obj.finished[i]
                    state.event.set()
            else:
                raise ValueError(f"Invalid object: {recv_obj}")