chore: import upstream snapshot with attribution

python/mlc_llm/serve/entrypoints/debug_entrypoints.py

@@ -0,0 +1,128 @@
+"""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()