chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,200 @@
|
||||
"""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()
|
||||
Reference in New Issue
Block a user