http_server.py 27.7 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
# Copyright 2023-2024 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""
The entry point of inference server. (SRT = SGLang Runtime)

Wang Ran (汪然)'s avatar
Wang Ran (汪然) committed
17
This file implements HTTP APIs for the inference engine via fastapi.
18
19
20
21
"""

import asyncio
import dataclasses
22
import json
23
24
25
26
27
28
import logging
import multiprocessing as multiprocessing
import os
import threading
import time
from http import HTTPStatus
Lianmin Zheng's avatar
Lianmin Zheng committed
29
from typing import AsyncIterator, Callable, Dict, Optional
30
31
32
33

# Fix a bug of Python threading
setattr(threading, "_register_atexit", lambda *args, **kwargs: None)

34
35
36
from contextlib import asynccontextmanager

import numpy as np
37
38
39
40
41
42
43
44
45
import orjson
import requests
import uvicorn
import uvloop
from fastapi import FastAPI, File, Form, Request, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import ORJSONResponse, Response, StreamingResponse

from sglang.srt.entrypoints.engine import _launch_subprocesses
YAMY's avatar
YAMY committed
46
from sglang.srt.function_call_parser import FunctionCallParser
47
48
49
50
51
52
53
54
from sglang.srt.managers.io_struct import (
    CloseSessionReqInput,
    ConfigureLoggingReq,
    EmbeddingReqInput,
    GenerateReqInput,
    GetWeightsByNameReqInput,
    InitWeightsUpdateGroupReqInput,
    OpenSessionReqInput,
55
    ParseFunctionCallReq,
56
    ProfileReqInput,
57
58
    ReleaseMemoryOccupationReqInput,
    ResumeMemoryOccupationReqInput,
Xihuai Wang's avatar
Xihuai Wang committed
59
    SeparateReasoningReqInput,
60
    SetInternalStateReq,
61
62
    UpdateWeightFromDiskReqInput,
    UpdateWeightsFromDistributedReqInput,
63
    UpdateWeightsFromTensorReqInput,
64
    VertexGenerateReqInput,
65
)
66
from sglang.srt.managers.tokenizer_manager import TokenizerManager
67
68
69
70
71
72
73
74
75
76
77
78
79
80
from sglang.srt.metrics.func_timer import enable_func_timer
from sglang.srt.openai_api.adapter import (
    v1_batches,
    v1_cancel_batch,
    v1_chat_completions,
    v1_completions,
    v1_delete_file,
    v1_embeddings,
    v1_files_create,
    v1_retrieve_batch,
    v1_retrieve_file,
    v1_retrieve_file_content,
)
from sglang.srt.openai_api.protocol import ModelCard, ModelList
Xihuai Wang's avatar
Xihuai Wang committed
81
from sglang.srt.reasoning_parser import ReasoningParser
82
83
84
85
86
87
88
89
from sglang.srt.server_args import ServerArgs
from sglang.srt.utils import (
    add_api_key_middleware,
    add_prometheus_middleware,
    delete_directory,
    kill_process_tree,
    set_uvicorn_logging_configs,
)
90
from sglang.srt.warmup import execute_warmups
91
92
93
94
95
96
97
98
99
100
from sglang.utils import get_exception_traceback
from sglang.version import __version__

logger = logging.getLogger(__name__)
asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())


# Store global states
@dataclasses.dataclass
class _GlobalState:
101
    tokenizer_manager: TokenizerManager
102
103
104
105
106
107
108
109
110
111
112
    scheduler_info: Dict


_global_state: Optional[_GlobalState] = None


def set_global_state(global_state: _GlobalState):
    global _global_state
    _global_state = global_state


113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
@asynccontextmanager
async def lifespan(fast_api_app: FastAPI):
    server_args: ServerArgs = fast_api_app.server_args
    if server_args.warmups is not None:
        await execute_warmups(
            server_args.warmups.split(","), _global_state.tokenizer_manager
        )
        logger.info("Warmup ended")

    warmup_thread = getattr(fast_api_app, "warmup_thread", None)
    if warmup_thread is not None:
        warmup_thread.start()
    yield


# Fast API
app = FastAPI(lifespan=lifespan)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

HEALTH_CHECK_TIMEOUT = int(os.getenv("SGLANG_HEALTH_CHECK_TIMEOUT", 20))


141
142
143
144
145
146
147
148
149
150
151
152
153
##### Native API endpoints #####


@app.get("/health")
async def health() -> Response:
    """Check the health of the http server."""
    return Response(status_code=200)


@app.get("/health_generate")
async def health_generate(request: Request) -> Response:
    """Check the health of the inference server by generating one token."""

154
155
    sampling_params = {"max_new_tokens": 1, "temperature": 0.0}
    rid = f"HEALTH_CHECK_{time.time()}"
156

157
158
159
    if _global_state.tokenizer_manager.is_image_gen:
        raise NotImplementedError()
    elif _global_state.tokenizer_manager.is_generation:
160
        gri = GenerateReqInput(
161
162
163
164
            rid=rid,
            input_ids=[0],
            sampling_params=sampling_params,
            log_metrics=False,
165
166
167
        )
    else:
        gri = EmbeddingReqInput(
168
            rid=rid, input_ids=[0], sampling_params=sampling_params, log_metrics=False
169
170
        )

171
    async def gen():
172
        async for _ in _global_state.tokenizer_manager.generate_request(gri, request):
173
            break
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195

    tic = time.time()
    task = asyncio.create_task(gen())
    while time.time() < tic + HEALTH_CHECK_TIMEOUT:
        await asyncio.sleep(1)
        if _global_state.tokenizer_manager.last_receive_tstamp > tic:
            task.cancel()
            _global_state.tokenizer_manager.rid_to_state.pop(rid, None)
            return Response(status_code=200)

    task.cancel()
    tic_time = time.strftime("%H:%M:%S", time.localtime(tic))
    last_receive_time = time.strftime(
        "%H:%M:%S", time.localtime(_global_state.tokenizer_manager.last_receive_tstamp)
    )
    logger.error(
        f"Health check failed. Server couldn't get a response from detokenizer for last "
        f"{HEALTH_CHECK_TIMEOUT} seconds. tic start time: {tic_time}. "
        f"last_heartbeat time: {last_receive_time}"
    )
    _global_state.tokenizer_manager.rid_to_state.pop(rid, None)
    return Response(status_code=503)
196
197
198
199
200
201


@app.get("/get_model_info")
async def get_model_info():
    """Get the model information."""
    result = {
202
203
204
        "model_path": _global_state.tokenizer_manager.model_path,
        "tokenizer_path": _global_state.tokenizer_manager.server_args.tokenizer_path,
        "is_generation": _global_state.tokenizer_manager.is_generation,
205
206
207
208
209
210
    }
    return result


@app.get("/get_server_info")
async def get_server_info():
211
    internal_states = await _global_state.tokenizer_manager.get_internal_state()
212
    return {
213
        **dataclasses.asdict(_global_state.tokenizer_manager.server_args),
214
        **_global_state.scheduler_info,
215
        **internal_states,
216
217
218
219
        "version": __version__,
    }


220
221
222
223
224
225
@app.api_route("/set_internal_state", methods=["POST", "PUT"])
async def set_internal_state(obj: SetInternalStateReq, request: Request):
    res = await _global_state.tokenizer_manager.set_internal_state(obj)
    return res


226
227
228
229
230
231
232
233
# fastapi implicitly converts json in the request to obj (dataclass)
@app.api_route("/generate", methods=["POST", "PUT"])
async def generate_request(obj: GenerateReqInput, request: Request):
    """Handle a generate request."""
    if obj.stream:

        async def stream_results() -> AsyncIterator[bytes]:
            try:
234
                async for out in _global_state.tokenizer_manager.generate_request(
235
236
237
238
239
240
241
                    obj, request
                ):
                    yield b"data: " + orjson.dumps(
                        out, option=orjson.OPT_NON_STR_KEYS
                    ) + b"\n\n"
            except ValueError as e:
                out = {"error": {"message": str(e)}}
242
                logger.error(f"Error: {e}")
243
244
245
246
247
248
249
250
                yield b"data: " + orjson.dumps(
                    out, option=orjson.OPT_NON_STR_KEYS
                ) + b"\n\n"
            yield b"data: [DONE]\n\n"

        return StreamingResponse(
            stream_results(),
            media_type="text/event-stream",
251
            background=_global_state.tokenizer_manager.create_abort_task(obj),
252
253
254
        )
    else:
        try:
255
            ret = await _global_state.tokenizer_manager.generate_request(
256
257
258
259
260
261
262
263
                obj, request
            ).__anext__()
            return ret
        except ValueError as e:
            logger.error(f"Error: {e}")
            return _create_error_response(e)


264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
@app.api_route("/generate_from_file", methods=["POST"])
async def generate_from_file_request(file: UploadFile, request: Request):
    """Handle a generate request, this is purely to work with input_embeds."""
    content = await file.read()
    input_embeds = json.loads(content.decode("utf-8"))

    obj = GenerateReqInput(
        input_embeds=input_embeds,
        sampling_params={
            "repetition_penalty": 1.2,
            "temperature": 0.2,
            "max_new_tokens": 512,
        },
    )

    try:
        ret = await _global_state.generate_request(obj, request).__anext__()
        return ret
    except ValueError as e:
        logger.error(f"Error: {e}")
        return _create_error_response(e)


287
288
289
290
@app.api_route("/encode", methods=["POST", "PUT"])
async def encode_request(obj: EmbeddingReqInput, request: Request):
    """Handle an embedding request."""
    try:
291
        ret = await _global_state.tokenizer_manager.generate_request(
292
293
294
295
296
297
298
299
300
301
302
            obj, request
        ).__anext__()
        return ret
    except ValueError as e:
        return _create_error_response(e)


@app.api_route("/classify", methods=["POST", "PUT"])
async def classify_request(obj: EmbeddingReqInput, request: Request):
    """Handle a reward model request. Now the arguments and return values are the same as embedding models."""
    try:
303
        ret = await _global_state.tokenizer_manager.generate_request(
304
305
306
307
308
309
310
            obj, request
        ).__anext__()
        return ret
    except ValueError as e:
        return _create_error_response(e)


311
@app.api_route("/flush_cache", methods=["GET", "POST"])
312
313
async def flush_cache():
    """Flush the radix cache."""
314
    ret = await _global_state.tokenizer_manager.flush_cache()
315
316
317
    return Response(
        content="Cache flushed.\nPlease check backend logs for more details. "
        "(When there are running or waiting requests, the operation will not be performed.)\n",
318
        status_code=200 if ret.success else HTTPStatus.BAD_REQUEST,
319
320
321
322
    )


@app.api_route("/start_profile", methods=["GET", "POST"])
323
async def start_profile_async(obj: Optional[ProfileReqInput] = None):
324
    """Start profiling."""
325
326
327
328
329
330
    if obj is None:
        obj = ProfileReqInput()

    await _global_state.tokenizer_manager.start_profile(
        obj.output_dir, obj.num_steps, obj.activities
    )
331
332
333
334
335
336
337
338
339
    return Response(
        content="Start profiling.\n",
        status_code=200,
    )


@app.api_route("/stop_profile", methods=["GET", "POST"])
async def stop_profile_async():
    """Stop profiling."""
340
    _global_state.tokenizer_manager.stop_profile()
341
342
343
344
345
346
    return Response(
        content="Stop profiling. This will take some time.\n",
        status_code=200,
    )


347
348
349
@app.api_route("/start_expert_distribution_record", methods=["GET", "POST"])
async def start_expert_distribution_record_async():
    """Start recording the expert distribution. Clear the previous record if any."""
350
    await _global_state.tokenizer_manager.start_expert_distribution_record()
351
352
353
354
355
356
357
358
359
    return Response(
        content="Start recording the expert distribution.\n",
        status_code=200,
    )


@app.api_route("/stop_expert_distribution_record", methods=["GET", "POST"])
async def stop_expert_distribution_record_async():
    """Stop recording the expert distribution."""
360
    await _global_state.tokenizer_manager.stop_expert_distribution_record()
361
362
363
364
365
366
367
368
369
    return Response(
        content="Stop recording the expert distribution.\n",
        status_code=200,
    )


@app.api_route("/dump_expert_distribution_record", methods=["GET", "POST"])
async def dump_expert_distribution_record_async():
    """Dump expert distribution record."""
370
    await _global_state.tokenizer_manager.dump_expert_distribution_record()
371
372
373
374
375
376
    return Response(
        content="Dump expert distribution record.\n",
        status_code=200,
    )


377
378
@app.post("/update_weights_from_disk")
async def update_weights_from_disk(obj: UpdateWeightFromDiskReqInput, request: Request):
379
380
381
    """Update the weights from disk inplace without re-launching the server."""
    success, message, num_paused_requests = (
        await _global_state.tokenizer_manager.update_weights_from_disk(obj, request)
382
    )
383
384
385
386
387
    content = {
        "success": success,
        "message": message,
        "num_paused_requests": num_paused_requests,
    }
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
    if success:
        return ORJSONResponse(
            content,
            status_code=HTTPStatus.OK,
        )
    else:
        return ORJSONResponse(
            content,
            status_code=HTTPStatus.BAD_REQUEST,
        )


@app.post("/init_weights_update_group")
async def init_weights_update_group(
    obj: InitWeightsUpdateGroupReqInput, request: Request
):
    """Initialize the parameter update group."""
405
    success, message = await _global_state.tokenizer_manager.init_weights_update_group(
406
407
408
409
410
411
412
413
414
        obj, request
    )
    content = {"success": success, "message": message}
    if success:
        return ORJSONResponse(content, status_code=200)
    else:
        return ORJSONResponse(content, status_code=HTTPStatus.BAD_REQUEST)


415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
@app.post("/update_weights_from_tensor")
async def update_weights_from_tensor(
    obj: UpdateWeightsFromTensorReqInput, request: Request
):
    """Update the weights from tensor inplace without re-launching the server.
    Notes:
    1. Ensure that the model is on the correct device (e.g., GPU) before calling this endpoint. If the model is moved to the CPU unexpectedly, it may cause performance issues or runtime errors.
    2. HTTP will transmit only the metadata of the tensor, while the tensor itself will be directly copied to the model.
    3. Any binary data in the named tensors should be base64 encoded.
    """

    success, message = await _global_state.tokenizer_manager.update_weights_from_tensor(
        obj, request
    )
    content = {"success": success, "message": message}
    return ORJSONResponse(
        content, status_code=200 if success else HTTPStatus.BAD_REQUEST
    )


435
436
437
438
439
@app.post("/update_weights_from_distributed")
async def update_weights_from_distributed(
    obj: UpdateWeightsFromDistributedReqInput, request: Request
):
    """Update model parameter from distributed online."""
440
441
442
443
    success, message = (
        await _global_state.tokenizer_manager.update_weights_from_distributed(
            obj, request
        )
444
445
446
447
448
449
450
451
452
453
454
455
    )
    content = {"success": success, "message": message}
    if success:
        return ORJSONResponse(content, status_code=200)
    else:
        return ORJSONResponse(content, status_code=HTTPStatus.BAD_REQUEST)


@app.api_route("/get_weights_by_name", methods=["GET", "POST"])
async def get_weights_by_name(obj: GetWeightsByNameReqInput, request: Request):
    """Get model parameter by name."""
    try:
456
        ret = await _global_state.tokenizer_manager.get_weights_by_name(obj, request)
457
458
459
460
461
462
463
464
465
466
467
468
        if ret is None:
            return _create_error_response("Get parameter by name failed")
        else:
            return ORJSONResponse(ret, status_code=200)
    except Exception as e:
        return _create_error_response(e)


@app.api_route("/release_memory_occupation", methods=["GET", "POST"])
async def release_memory_occupation(
    obj: ReleaseMemoryOccupationReqInput, request: Request
):
469
    """Release GPU memory occupation temporarily."""
470
    try:
471
        await _global_state.tokenizer_manager.release_memory_occupation(obj, request)
472
473
474
475
476
477
478
479
    except Exception as e:
        return _create_error_response(e)


@app.api_route("/resume_memory_occupation", methods=["GET", "POST"])
async def resume_memory_occupation(
    obj: ResumeMemoryOccupationReqInput, request: Request
):
480
    """Resume GPU memory occupation."""
481
    try:
482
        await _global_state.tokenizer_manager.resume_memory_occupation(obj, request)
483
484
485
486
487
488
489
490
    except Exception as e:
        return _create_error_response(e)


@app.api_route("/open_session", methods=["GET", "POST"])
async def open_session(obj: OpenSessionReqInput, request: Request):
    """Open a session, and return its unique session id."""
    try:
491
        session_id = await _global_state.tokenizer_manager.open_session(obj, request)
492
493
494
495
496
497
498
499
500
501
502
        if session_id is None:
            raise Exception(
                "Failed to open the session. Check if a session with the same id is still open."
            )
        return session_id
    except Exception as e:
        return _create_error_response(e)


@app.api_route("/close_session", methods=["GET", "POST"])
async def close_session(obj: CloseSessionReqInput, request: Request):
503
    """Close the session."""
504
    try:
505
        await _global_state.tokenizer_manager.close_session(obj, request)
506
507
508
509
510
511
512
        return Response(status_code=200)
    except Exception as e:
        return _create_error_response(e)


@app.api_route("/configure_logging", methods=["GET", "POST"])
async def configure_logging(obj: ConfigureLoggingReq, request: Request):
513
    """Configure the request logging options."""
514
    _global_state.tokenizer_manager.configure_logging(obj)
515
516
517
    return Response(status_code=200)


518
519
@app.post("/parse_function_call")
async def parse_function_call_request(obj: ParseFunctionCallReq, request: Request):
YAMY's avatar
YAMY committed
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
    """
    A native API endpoint to parse function calls from a text.
    """
    # 1) Initialize the parser based on the request body
    parser = FunctionCallParser(tools=obj.tools, tool_call_parser=obj.tool_call_parser)

    # 2) Call the non-stream parsing method (non-stream)
    normal_text, calls = parser.parse_non_stream(obj.text)

    # 3) Organize the response content
    response_data = {
        "normal_text": normal_text,
        "calls": [
            call.model_dump() for call in calls
        ],  # Convert pydantic objects to dictionaries
    }

    return ORJSONResponse(content=response_data, status_code=200)


Xihuai Wang's avatar
Xihuai Wang committed
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
@app.post("/separate_reasoning")
async def separate_reasoning_request(obj: SeparateReasoningReqInput, request: Request):
    """
    A native API endpoint to separate reasoning from a text.
    """
    # 1) Initialize the parser based on the request body
    parser = ReasoningParser(model_type=obj.reasoning_parser)

    # 2) Call the non-stream parsing method (non-stream)
    reasoning_text, normal_text = parser.parse_non_stream(obj.text)

    # 3) Organize the response content
    response_data = {
        "reasoning_text": reasoning_text,
        "text": normal_text,
    }

    return ORJSONResponse(content=response_data, status_code=200)


560
561
562
563
564
##### OpenAI-compatible API endpoints #####


@app.post("/v1/completions")
async def openai_v1_completions(raw_request: Request):
565
    return await v1_completions(_global_state.tokenizer_manager, raw_request)
566
567
568
569


@app.post("/v1/chat/completions")
async def openai_v1_chat_completions(raw_request: Request):
570
    return await v1_chat_completions(_global_state.tokenizer_manager, raw_request)
571
572
573
574


@app.post("/v1/embeddings", response_class=ORJSONResponse)
async def openai_v1_embeddings(raw_request: Request):
575
    response = await v1_embeddings(_global_state.tokenizer_manager, raw_request)
576
577
578
579
580
581
    return response


@app.get("/v1/models", response_class=ORJSONResponse)
def available_models():
    """Show available models."""
582
    served_model_names = [_global_state.tokenizer_manager.served_model_name]
583
584
    model_cards = []
    for served_model_name in served_model_names:
585
586
587
588
589
590
591
        model_cards.append(
            ModelCard(
                id=served_model_name,
                root=served_model_name,
                max_model_len=_global_state.tokenizer_manager.model_config.context_len,
            )
        )
592
593
594
595
596
597
    return ModelList(data=model_cards)


@app.post("/v1/files")
async def openai_v1_files(file: UploadFile = File(...), purpose: str = Form("batch")):
    return await v1_files_create(
598
        file, purpose, _global_state.tokenizer_manager.server_args.file_storage_path
599
600
601
602
603
604
605
606
607
608
609
    )


@app.delete("/v1/files/{file_id}")
async def delete_file(file_id: str):
    # https://platform.openai.com/docs/api-reference/files/delete
    return await v1_delete_file(file_id)


@app.post("/v1/batches")
async def openai_v1_batches(raw_request: Request):
610
    return await v1_batches(_global_state.tokenizer_manager, raw_request)
611
612
613
614
615


@app.post("/v1/batches/{batch_id}/cancel")
async def cancel_batches(batch_id: str):
    # https://platform.openai.com/docs/api-reference/batch/cancel
616
    return await v1_cancel_batch(_global_state.tokenizer_manager, batch_id)
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635


@app.get("/v1/batches/{batch_id}")
async def retrieve_batch(batch_id: str):
    return await v1_retrieve_batch(batch_id)


@app.get("/v1/files/{file_id}")
async def retrieve_file(file_id: str):
    # https://platform.openai.com/docs/api-reference/files/retrieve
    return await v1_retrieve_file(file_id)


@app.get("/v1/files/{file_id}/content")
async def retrieve_file_content(file_id: str):
    # https://platform.openai.com/docs/api-reference/files/retrieve-contents
    return await v1_retrieve_file_content(file_id)


636
637
638
639
640
641
642
643
644
645
646
647
## SageMaker API
@app.get("/ping")
async def sagemaker_health() -> Response:
    """Check the health of the http server."""
    return Response(status_code=200)


@app.post("/invocations")
async def sagemaker_chat_completions(raw_request: Request):
    return await v1_chat_completions(_global_state.tokenizer_manager, raw_request)


648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
## Vertex AI API
@app.post(os.environ.get("AIP_PREDICT_ROUTE", "/vertex_generate"))
async def vertex_generate(vertex_req: VertexGenerateReqInput, raw_request: Request):
    if not vertex_req.instances:
        return []
    inputs = {}
    for input_key in ("text", "input_ids", "input_embeds"):
        if vertex_req.instances[0].get(input_key):
            inputs[input_key] = [
                instance.get(input_key) for instance in vertex_req.instances
            ]
            break
    image_data = [
        instance.get("image_data")
        for instance in vertex_req.instances
        if instance.get("image_data") is not None
    ] or None
    req = GenerateReqInput(
        **inputs,
        image_data=image_data,
        **(vertex_req.parameters or {}),
    )
    ret = await generate_request(req, raw_request)
    return ORJSONResponse({"predictions": ret})


674
675
676
677
678
679
680
681
682
def _create_error_response(e):
    return ORJSONResponse(
        {"error": {"message": str(e)}}, status_code=HTTPStatus.BAD_REQUEST
    )


def launch_server(
    server_args: ServerArgs,
    pipe_finish_writer: Optional[multiprocessing.connection.Connection] = None,
683
    launch_callback: Optional[Callable[[], None]] = None,
684
685
686
687
688
689
690
691
):
    """
    Launch SRT (SGLang Runtime) Server.

    The SRT server consists of an HTTP server and an SRT engine.

    - HTTP server: A FastAPI server that routes requests to the engine.
    - The engine consists of three components:
692
        1. TokenizerManager: Tokenizes the requests and sends them to the scheduler.
693
694
695
696
        2. Scheduler (subprocess): Receives requests from the Tokenizer Manager, schedules batches, forwards them, and sends the output tokens to the Detokenizer Manager.
        3. DetokenizerManager (subprocess): Detokenizes the output tokens and sends the result back to the Tokenizer Manager.

    Note:
697
    1. The HTTP server, Engine, and TokenizerManager both run in the main process.
698
    2. Inter-process communication is done through IPC (each process uses a different port) via the ZMQ library.
699
    """
700
    tokenizer_manager, scheduler_info = _launch_subprocesses(server_args=server_args)
701
702
    set_global_state(
        _GlobalState(
703
            tokenizer_manager=tokenizer_manager,
704
705
706
707
708
709
710
711
712
713
714
715
716
            scheduler_info=scheduler_info,
        )
    )

    # Add api key authorization
    if server_args.api_key:
        add_api_key_middleware(app, server_args.api_key)

    # Add prometheus middleware
    if server_args.enable_metrics:
        add_prometheus_middleware(app)
        enable_func_timer()

717
718
719
    # Send a warmup request - we will create the thread launch it
    # in the lifespan after all other warmups have fired.
    warmup_thread = threading.Thread(
720
721
722
723
        target=_wait_and_warmup,
        args=(
            server_args,
            pipe_finish_writer,
724
            _global_state.tokenizer_manager.image_token_id,
725
            launch_callback,
726
727
        ),
    )
728
    app.warmup_thread = warmup_thread
729
730
731
732

    try:
        # Update logging configs
        set_uvicorn_logging_configs()
733
        app.server_args = server_args
734
735
736
737
738
739
740
741
742
743
        # Listen for HTTP requests
        uvicorn.run(
            app,
            host=server_args.host,
            port=server_args.port,
            log_level=server_args.log_level_http or server_args.log_level,
            timeout_keep_alive=5,
            loop="uvloop",
        )
    finally:
744
        warmup_thread.join()
745
746


747
748
749
750
751
752
def _wait_and_warmup(
    server_args: ServerArgs,
    pipe_finish_writer: Optional[multiprocessing.connection.Connection],
    image_token_text: str,
    launch_callback: Optional[Callable[[], None]] = None,
):
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
    headers = {}
    url = server_args.url()
    if server_args.api_key:
        headers["Authorization"] = f"Bearer {server_args.api_key}"

    # Wait until the server is launched
    success = False
    for _ in range(120):
        time.sleep(1)
        try:
            res = requests.get(url + "/get_model_info", timeout=5, headers=headers)
            assert res.status_code == 200, f"{res=}, {res.text=}"
            success = True
            break
        except (AssertionError, requests.exceptions.RequestException):
            last_traceback = get_exception_traceback()
            pass

    if not success:
        if pipe_finish_writer is not None:
            pipe_finish_writer.send(last_traceback)
        logger.error(f"Initialization failed. warmup error: {last_traceback}")
        kill_process_tree(os.getpid())
        return

    model_info = res.json()

    # Send a warmup request
    request_name = "/generate" if model_info["is_generation"] else "/encode"
    max_new_tokens = 8 if model_info["is_generation"] else 1
    json_data = {
        "sampling_params": {
            "temperature": 0,
            "max_new_tokens": max_new_tokens,
        },
    }
    if server_args.skip_tokenizer_init:
fzyzcjy's avatar
fzyzcjy committed
790
        json_data["input_ids"] = [[10, 11, 12] for _ in range(server_args.dp_size)]
fzyzcjy's avatar
fzyzcjy committed
791
792
793
        # TODO Workaround the bug that embedding errors for list of size 1
        if server_args.dp_size == 1:
            json_data["input_ids"] = json_data["input_ids"][0]
794
    else:
fzyzcjy's avatar
fzyzcjy committed
795
        json_data["text"] = ["The capital city of France is"] * server_args.dp_size
fzyzcjy's avatar
fzyzcjy committed
796
797
798
        # TODO Workaround the bug that embedding errors for list of size 1
        if server_args.dp_size == 1:
            json_data["text"] = json_data["text"][0]
799

800
801
802
803
804
805
806
807
    # Debug dumping
    if server_args.debug_tensor_dump_input_file:
        json_data.pop("text", None)
        json_data["input_ids"] = np.load(
            server_args.debug_tensor_dump_input_file
        ).tolist()
        json_data["sampling_params"]["max_new_tokens"] = 0

808
    try:
809
810
811
812
813
814
815
816
817
818
819
        if server_args.disaggregation_mode == "null":
            res = requests.post(
                url + request_name,
                json=json_data,
                headers=headers,
                timeout=600,
            )
            assert res.status_code == 200, f"{res}"
        else:
            # Warmup request currently hangs in disaggregation mode, so we skip it.
            logger.info("Skipping warmup request in disaggregation mode")
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
    except Exception:
        last_traceback = get_exception_traceback()
        if pipe_finish_writer is not None:
            pipe_finish_writer.send(last_traceback)
        logger.error(f"Initialization failed. warmup error: {last_traceback}")
        kill_process_tree(os.getpid())
        return

    # Debug print
    # logger.info(f"{res.json()=}")

    logger.info("The server is fired up and ready to roll!")
    if pipe_finish_writer is not None:
        pipe_finish_writer.send("ready")

    if server_args.delete_ckpt_after_loading:
        delete_directory(server_args.model_path)
837
838
839
840
841
842

    if server_args.debug_tensor_dump_input_file:
        kill_process_tree(os.getpid())

    if launch_callback is not None:
        launch_callback()