wehub-resource-sync
135 lines
4.6 KiB
Python
135 lines
4.6 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# Standard
|
|
from typing import TYPE_CHECKING
|
|
import asyncio
|
|
import os
|
|
import threading
|
|
|
|
# Third Party
|
|
from fastapi import FastAPI
|
|
import uvicorn
|
|
|
|
# First Party
|
|
from lmcache.logging import init_logger
|
|
|
|
# Local
|
|
from .api_registry import APIRegistry
|
|
|
|
if TYPE_CHECKING:
|
|
# First Party
|
|
from lmcache.v1.manager import LMCacheManager
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
def _build_app() -> FastAPI:
|
|
"""Create a fresh FastAPI app with all internal API routes registered."""
|
|
new_app = FastAPI()
|
|
APIRegistry(new_app).register_all_apis()
|
|
return new_app
|
|
|
|
|
|
# Module-level app kept for backward compatibility with existing tests that
|
|
# import `app` directly. Production code no longer relies on this shared
|
|
# singleton; each InternalAPIServer owns its own FastAPI app so that
|
|
# multiple instances in the same process (e.g. scheduler + worker in TP=1
|
|
# non-MP mode) do not overwrite each other's app.state.lmcache_adapter.
|
|
app = _build_app()
|
|
|
|
|
|
class InternalAPIServer:
|
|
def __init__(self, lmcache_manager: "LMCacheManager"):
|
|
lmcache_engine = lmcache_manager.lmcache_engine
|
|
|
|
# Check if lmcache_engine is None and handle accordingly
|
|
if lmcache_engine is None:
|
|
# Use manager's config directly when engine is not available
|
|
config = lmcache_manager.config
|
|
port_offset = 0 # Default for scheduler mode
|
|
else:
|
|
config = lmcache_engine.config
|
|
# 0 for scheduler, 1 for worker 0, 2 for worker 1, ...
|
|
port_offset = 1 + lmcache_engine.metadata.worker_id
|
|
|
|
self.port = config.internal_api_server_port_start + port_offset
|
|
self.socket_path_prefix = config.internal_api_server_socket_path_prefix
|
|
self.socket_path = (
|
|
f"{self.socket_path_prefix}_{self.port}"
|
|
if self.socket_path_prefix
|
|
else None
|
|
)
|
|
include_index_list = config.internal_api_server_include_index_list
|
|
|
|
self.enable = True
|
|
if not config.internal_api_server_enabled or (
|
|
include_index_list and port_offset not in include_index_list
|
|
):
|
|
logger.info(
|
|
"Internal API server disabled. internal_api_server_enabled=%s, "
|
|
"port_offset=%s, port=%s, socket_path=%s, include_index_list=%s",
|
|
config.internal_api_server_enabled,
|
|
port_offset,
|
|
self.port,
|
|
self.socket_path,
|
|
include_index_list,
|
|
)
|
|
self.enable = False
|
|
return
|
|
|
|
self.app = _build_app()
|
|
|
|
uvicorn_config = {
|
|
"app": self.app,
|
|
"host": config.internal_api_server_host,
|
|
"loop": "uvloop",
|
|
"http": "httptools",
|
|
"access_log": config.get_extra_config_value(
|
|
"internal_api_server_access_log", True
|
|
),
|
|
"log_level": config.get_extra_config_value(
|
|
"internal_api_server_log_level", "warning"
|
|
),
|
|
}
|
|
|
|
if self.socket_path:
|
|
self.server_log_info = f"socket {self.socket_path}"
|
|
logger.info("Init internal API server on %s", self.server_log_info)
|
|
uvicorn_config["uds"] = self.socket_path
|
|
# Ensure socket directory exists
|
|
os.makedirs(os.path.dirname(self.socket_path), exist_ok=True)
|
|
# Remove existing socket file if exists
|
|
if os.path.exists(self.socket_path):
|
|
os.unlink(self.socket_path)
|
|
else:
|
|
self.server_log_info = f"port {self.port}"
|
|
logger.info("Init internal API server on %s", self.server_log_info)
|
|
uvicorn_config["port"] = self.port
|
|
|
|
self.server = uvicorn.Server(uvicorn.Config(**uvicorn_config))
|
|
self.app.state.lmcache_adapter = lmcache_manager
|
|
|
|
async def run(self):
|
|
logger.info("Running LMCache internal API server on %s", self.server_log_info)
|
|
if self.server:
|
|
await self.server.serve()
|
|
|
|
def start(self):
|
|
if not self.enable:
|
|
return
|
|
logger.info("Starting LMCache internal API server on %s", self.server_log_info)
|
|
threading.Thread(
|
|
target=asyncio.run,
|
|
args=(self.run(),),
|
|
daemon=True,
|
|
name="api-server-thread",
|
|
).start()
|
|
|
|
def stop(self):
|
|
if not self.enable:
|
|
return
|
|
logger.info("Stopping LMCache internal API server")
|
|
if self.server:
|
|
self.server.should_exit = True
|
|
if self.socket_path and os.path.exists(self.socket_path):
|
|
os.unlink(self.socket_path)
|