wehub-resource-sync
129 lines
4.7 KiB
Python
129 lines
4.7 KiB
Python
"""MLC LLM server debug entrypoints"""
|
|
|
|
import json
|
|
from http import HTTPStatus
|
|
|
|
import fastapi
|
|
|
|
from mlc_llm.protocol import error_protocol
|
|
from mlc_llm.serve.server import ServerContext
|
|
|
|
app = fastapi.APIRouter()
|
|
|
|
################ /debug/dump_event_trace ################
|
|
|
|
|
|
@app.post("/debug/dump_event_trace")
|
|
async def debug_dump_event_trace(request: fastapi.Request):
|
|
"""Return the recorded events in Chrome Trace Event Format in JSON string.
|
|
The input request payload should have only one field, specifying the
|
|
model to query. For example: `{"model": "Llama-2-7b-chat-hf-q0f16"}`.
|
|
"""
|
|
# Get the raw request body as bytes
|
|
request_raw_data = await request.body()
|
|
request_json_str = request_raw_data.decode("utf-8")
|
|
try:
|
|
# Parse the JSON string
|
|
request_dict = json.loads(request_json_str)
|
|
except json.JSONDecodeError:
|
|
return error_protocol.create_error_response(
|
|
HTTPStatus.BAD_REQUEST, message=f"Invalid request {request_json_str}"
|
|
)
|
|
if "model" not in request_dict:
|
|
return error_protocol.create_error_response(
|
|
HTTPStatus.BAD_REQUEST, message=f"Invalid request {request_json_str}"
|
|
)
|
|
|
|
# Check the requested model.
|
|
model = request_dict["model"]
|
|
|
|
server_context: ServerContext = ServerContext.current()
|
|
async_engine = server_context.get_engine(model)
|
|
|
|
if async_engine is None:
|
|
return error_protocol.create_error_response(
|
|
HTTPStatus.BAD_REQUEST,
|
|
message=f'The requested model "{model}" is not served.',
|
|
)
|
|
if async_engine.state.trace_recorder is None:
|
|
return error_protocol.create_error_response(
|
|
HTTPStatus.BAD_REQUEST,
|
|
message=f'The requested model "{model}" does not enable tracing',
|
|
)
|
|
|
|
return json.loads(async_engine.state.trace_recorder.dump_json())
|
|
|
|
|
|
################ /debug/cuda_profiler_start/end ################
|
|
|
|
|
|
@app.post("/debug/cuda_profiler_start")
|
|
async def debug_cuda_profiler_start(_request: fastapi.Request):
|
|
"""Start the cuda profiler for the engine. Only for debug purpose."""
|
|
server_context: ServerContext = ServerContext.current()
|
|
# Since the CUDA profiler is process-wise, call the function for one model is sufficient.
|
|
for model in server_context.get_model_list():
|
|
async_engine = server_context.get_engine(model)
|
|
async_engine._debug_call_func_on_all_worker("mlc.debug_cuda_profiler_start")
|
|
break
|
|
|
|
|
|
@app.post("/debug/cuda_profiler_stop")
|
|
async def debug_cuda_profiler_stop(_request: fastapi.Request):
|
|
"""Stop the cuda profiler for the engine. Only for debug purpose."""
|
|
server_context: ServerContext = ServerContext.current()
|
|
# Since the CUDA profiler is process-wise, call the function for one model is sufficient.
|
|
for model in server_context.get_model_list():
|
|
async_engine = server_context.get_engine(model)
|
|
async_engine._debug_call_func_on_all_worker("mlc.debug_cuda_profiler_stop")
|
|
break
|
|
|
|
|
|
@app.post("/debug/dump_engine_metrics")
|
|
async def debug_dump_engine_metrics(request: fastapi.Request):
|
|
"""Dump the engine metrics for the engine. Only for debug purpose."""
|
|
# Get the raw request body as bytes
|
|
request_raw_data = await request.body()
|
|
request_json_str = request_raw_data.decode("utf-8")
|
|
try:
|
|
# Parse the JSON string
|
|
request_dict = json.loads(request_json_str)
|
|
except json.JSONDecodeError:
|
|
return error_protocol.create_error_response(
|
|
HTTPStatus.BAD_REQUEST, message=f"Invalid request {request_json_str}"
|
|
)
|
|
|
|
# Check the requested model.
|
|
model = request_dict.get("model", None)
|
|
|
|
server_context: ServerContext = ServerContext.current()
|
|
async_engine = server_context.get_engine(model)
|
|
res = await async_engine.metrics()
|
|
return res
|
|
|
|
|
|
@app.post("/debug/reset_engine")
|
|
async def debug_reset_engine_stats(request: fastapi.Request):
|
|
"""Reset the engine, clean up all running data and metrics."""
|
|
# Get the raw request body as bytes
|
|
request_raw_data = await request.body()
|
|
request_json_str = request_raw_data.decode("utf-8")
|
|
try:
|
|
# Parse the JSON string
|
|
request_dict = json.loads(request_json_str)
|
|
except json.JSONDecodeError:
|
|
return error_protocol.create_error_response(
|
|
HTTPStatus.BAD_REQUEST, message=f"Invalid request {request_json_str}"
|
|
)
|
|
if "model" not in request_dict:
|
|
return error_protocol.create_error_response(
|
|
HTTPStatus.BAD_REQUEST, message=f"Invalid request {request_json_str}"
|
|
)
|
|
|
|
# Check the requested model.
|
|
model = request_dict["model"]
|
|
|
|
server_context: ServerContext = ServerContext.current()
|
|
async_engine = server_context.get_engine(model)
|
|
async_engine.reset()
|