wehub-resource-sync
201 lines
7.2 KiB
Python
201 lines
7.2 KiB
Python
"""The MLC LLM server launched in a subprocess."""
|
|
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Literal, Optional, Union
|
|
|
|
import psutil
|
|
import requests
|
|
from tvm.runtime import Device
|
|
|
|
from mlc_llm.serve.config import EngineConfig
|
|
from mlc_llm.serve.engine_base import _check_engine_config
|
|
|
|
|
|
class PopenServer:
|
|
"""The wrapper of MLC LLM server, which runs the server in
|
|
a background subprocess.
|
|
|
|
This server can be used for debugging purposes.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
model: str,
|
|
device: Union[str, Device] = "auto",
|
|
*,
|
|
model_lib: Optional[str] = None,
|
|
mode: Literal["local", "interactive", "server"] = "local",
|
|
engine_config: Optional[EngineConfig] = None,
|
|
enable_debug: bool = True,
|
|
enable_tracing: bool = False,
|
|
host: str = "127.0.0.1",
|
|
port: int = 8082,
|
|
) -> None:
|
|
"""Please check out `python/mlc_llm/cli/serve.py` for the server arguments."""
|
|
# - Check the fields fields of `engine_config`.
|
|
if engine_config is None:
|
|
engine_config = EngineConfig()
|
|
_check_engine_config(model, model_lib, mode, engine_config)
|
|
|
|
self.model = model
|
|
self.model_lib = model_lib
|
|
self.device = device
|
|
self.mode = mode
|
|
self.enable_debug = enable_debug
|
|
self.engine_config = engine_config
|
|
self.enable_tracing = enable_tracing
|
|
self.enable_debug = enable_debug
|
|
self.host = host
|
|
self.port = port
|
|
self._proc: Optional[subprocess.Popen] = None
|
|
|
|
self.base_url = ""
|
|
self.openai_v1_base_url = ""
|
|
|
|
def start(self, extra_env=None) -> None:
|
|
"""Launch the server in a popen subprocess.
|
|
Wait until the server becomes ready before return.
|
|
"""
|
|
extra_env = extra_env or {}
|
|
cmd = [sys.executable]
|
|
cmd += ["-m", "mlc_llm", "serve", self.model]
|
|
if self.model_lib is not None:
|
|
cmd += ["--model-lib", self.model_lib]
|
|
cmd += ["--device", self.device]
|
|
|
|
if self.enable_debug:
|
|
cmd += ["--enable-debug"]
|
|
|
|
if self.mode is not None:
|
|
cmd += ["--mode", self.mode]
|
|
|
|
if len(self.engine_config.additional_models) > 0:
|
|
args_additional_model = []
|
|
for additional_model in self.engine_config.additional_models:
|
|
if isinstance(additional_model, str):
|
|
args_additional_model.append(additional_model)
|
|
else:
|
|
args_additional_model.append(additional_model[0] + "," + additional_model[1])
|
|
cmd += ["--additional-models", *args_additional_model]
|
|
cmd += ["--speculative-mode", self.engine_config.speculative_mode]
|
|
cmd += ["--prefix-cache-mode", self.engine_config.prefix_cache_mode]
|
|
|
|
args_overrides = []
|
|
if self.engine_config.max_num_sequence is not None:
|
|
args_overrides.append(f"max_num_sequence={self.engine_config.max_num_sequence}")
|
|
if self.engine_config.max_total_sequence_length is not None:
|
|
args_overrides.append(
|
|
f"max_total_seq_length={self.engine_config.max_total_sequence_length}"
|
|
)
|
|
if self.engine_config.prefill_chunk_size is not None:
|
|
args_overrides.append(f"prefill_chunk_size={self.engine_config.prefill_chunk_size}")
|
|
if self.engine_config.max_history_size is not None:
|
|
args_overrides.append(f"max_history_size={self.engine_config.max_history_size}")
|
|
if self.engine_config.gpu_memory_utilization is not None:
|
|
args_overrides.append(
|
|
f"gpu_memory_utilization={self.engine_config.gpu_memory_utilization}"
|
|
)
|
|
if self.engine_config.spec_draft_length is not None:
|
|
args_overrides.append(f"spec_draft_length={self.engine_config.spec_draft_length}")
|
|
if self.engine_config.prefix_cache_max_num_recycling_seqs is not None:
|
|
args_overrides.append(
|
|
"prefix_cache_max_num_recycling_seqs="
|
|
+ str(self.engine_config.prefix_cache_max_num_recycling_seqs)
|
|
)
|
|
if len(args_overrides) > 0:
|
|
cmd += ["--overrides", ";".join(args_overrides)]
|
|
|
|
if self.enable_tracing:
|
|
cmd += ["--enable-tracing"]
|
|
if self.enable_debug:
|
|
cmd += ["--enable-debug"]
|
|
|
|
cmd += ["--host", self.host]
|
|
cmd += ["--port", str(self.port)]
|
|
process_path = str(Path(__file__).resolve().parents[4])
|
|
final_env = os.environ.copy()
|
|
for key, value in extra_env.items():
|
|
final_env[key] = value
|
|
self._proc = subprocess.Popen(cmd, cwd=process_path, env=final_env)
|
|
# NOTE: DO NOT USE `stdout=subprocess.PIPE, stderr=subprocess.PIPE`
|
|
# in subprocess.Popen here. PIPE has a fixed-size buffer with may block
|
|
# and hang forever.
|
|
|
|
# Try to query the server until it is ready.
|
|
self.base_url = f"http://{self.host}:{str(self.port)}"
|
|
self.openai_v1_base_url = f"http://{self.host}:{str(self.port)}/v1"
|
|
openai_v1_models_url = f"{self.base_url}/v1/models"
|
|
|
|
query_result = None
|
|
timeout = 120
|
|
attempts = 0.0
|
|
while query_result is None and attempts < timeout:
|
|
try:
|
|
query_result = requests.get(openai_v1_models_url, timeout=60)
|
|
if query_result.status_code != 200:
|
|
query_result = None
|
|
attempts += 0.1
|
|
time.sleep(0.1)
|
|
except Exception:
|
|
attempts += 0.1
|
|
time.sleep(0.1)
|
|
|
|
# Check if the subprocess terminates unexpectedly or
|
|
# the queries reach the timeout.
|
|
process_return_code = self._proc.poll()
|
|
if process_return_code is not None:
|
|
raise RuntimeError(
|
|
"The server fails to launch. "
|
|
f'Please check if "{self.model}" is a valid model compiled by MLC LLM.'
|
|
)
|
|
if attempts == timeout:
|
|
self.terminate()
|
|
raise RuntimeError(f"The server fails to launch in {timeout} seconds.")
|
|
|
|
def terminate(self) -> None:
|
|
"""Terminate the server subprocess."""
|
|
if self._proc is None:
|
|
return
|
|
|
|
# Kill all the child processes.
|
|
def kill_child_processes():
|
|
try:
|
|
parent = psutil.Process(self._proc.pid)
|
|
children = parent.children(recursive=True)
|
|
except psutil.NoSuchProcess:
|
|
return
|
|
|
|
for process in children:
|
|
try:
|
|
process.kill()
|
|
except psutil.NoSuchProcess:
|
|
pass
|
|
|
|
kill_child_processes()
|
|
|
|
# Kill the process.
|
|
try:
|
|
self._proc.kill()
|
|
except OSError:
|
|
pass
|
|
|
|
# Join the process to avoid zombies.
|
|
try:
|
|
self._proc.wait(timeout=10.0)
|
|
except subprocess.TimeoutExpired:
|
|
pass
|
|
self._proc = None
|
|
|
|
def __enter__(self):
|
|
"""Start the server."""
|
|
self.start()
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
"""Terminate the server."""
|
|
self.terminate()
|