wehub-resource-sync
1037 lines
36 KiB
Python
1037 lines
36 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# Standard
|
|
from importlib import resources
|
|
from urllib.parse import unquote
|
|
import argparse
|
|
import asyncio
|
|
import json
|
|
import os
|
|
import threading
|
|
import time
|
|
|
|
# Third Party
|
|
from fastapi import APIRouter, FastAPI, HTTPException, Request
|
|
from fastapi.staticfiles import StaticFiles
|
|
from starlette.responses import FileResponse, PlainTextResponse
|
|
import httpx
|
|
import uvicorn
|
|
|
|
_PACKAGE = "lmcache.lmcache_frontend"
|
|
|
|
|
|
def _package_resource_path(relative: str) -> str:
|
|
"""Return absolute filesystem path for a file shipped inside the package.
|
|
|
|
Replacement for the deprecated ``pkg_resources.resource_filename``;
|
|
works for regular (non-zipped) installs, which is how LMCache ships.
|
|
"""
|
|
return str(resources.files(_PACKAGE).joinpath(relative))
|
|
|
|
|
|
try:
|
|
# Local
|
|
from .heartbeat import HeartbeatService # import as module
|
|
except ImportError:
|
|
# Third Party
|
|
from heartbeat import HeartbeatService # type: ignore # import as script
|
|
|
|
|
|
# Create router
|
|
router = APIRouter()
|
|
|
|
|
|
class _NodeRegistry:
|
|
"""Encapsulates the mutable proxy/node list used by the frontend.
|
|
|
|
Replacing the list is done in-place via :py:meth:`replace` so that
|
|
aliases (the module-level ``target_nodes`` reference and the list
|
|
handed to ``HeartbeatService``) stay in sync.
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
self._nodes: list[dict] = []
|
|
|
|
@property
|
|
def nodes(self) -> list[dict]:
|
|
"""Return the underlying list (mutated in place)."""
|
|
return self._nodes
|
|
|
|
def replace(self, new_nodes: list[dict]) -> None:
|
|
"""Swap the registry content in place with ``new_nodes``."""
|
|
self._nodes[:] = new_nodes
|
|
|
|
def is_allowed(self, host: str, port: str) -> bool:
|
|
"""Return True if ``host:port`` matches any registered node.
|
|
|
|
Used by the SSRF guard in :func:`proxy_request` so the proxy
|
|
only forwards to pre-registered destinations.
|
|
"""
|
|
return self.resolve(host, port) is not None
|
|
|
|
def resolve(self, host: str, port: str) -> tuple[str, str] | None:
|
|
"""Return the registry-owned ``(host, port)`` matching the input.
|
|
|
|
The returned tuple is taken from the registry itself, not from
|
|
the caller-supplied arguments. Using this value to build the
|
|
outbound URL breaks the SSRF taint flow for static analysers.
|
|
"""
|
|
port = str(port)
|
|
for proxy in self._nodes:
|
|
p_host, p_port = str(proxy.get("host")), str(proxy.get("port"))
|
|
if p_host == host and p_port == port:
|
|
return p_host, p_port
|
|
for child in proxy.get("nodes", []):
|
|
c_host, c_port = str(child.get("host")), str(child.get("port"))
|
|
if c_host == host and c_port == port:
|
|
return c_host, c_port
|
|
return None
|
|
|
|
|
|
_node_registry = _NodeRegistry()
|
|
# ``target_nodes`` is a module-level alias to the list owned by the
|
|
# registry. External readers and in-place mutations (append / element
|
|
# update) keep working unchanged; whole-list replacement MUST go
|
|
# through ``_node_registry.replace`` so all aliases stay in sync.
|
|
target_nodes = _node_registry.nodes
|
|
|
|
# Initialize heartbeat service with app context
|
|
heartbeat_service: HeartbeatService = HeartbeatService()
|
|
|
|
global args
|
|
args = None
|
|
|
|
|
|
async def fetch_child_nodes_from_proxy(proxy_node: dict) -> list[dict]:
|
|
"""Fetch child nodes from a single proxy node.
|
|
|
|
Args:
|
|
proxy_node: Proxy dict with ``name``/``host``/``port`` keys.
|
|
``/api/nodes`` of the proxy is queried.
|
|
|
|
Returns:
|
|
A list of leaf node dicts discovered from the proxy. An
|
|
empty list is returned when the request fails.
|
|
"""
|
|
try:
|
|
url = f"http://{proxy_node['host']}:{proxy_node['port']}/api/nodes"
|
|
async with httpx.AsyncClient(timeout=5.0) as client:
|
|
response = await client.get(url)
|
|
response.raise_for_status()
|
|
return_nodes = []
|
|
added_names = set() # Track node names to avoid duplicates
|
|
nodes = response.json().get("nodes", [])
|
|
|
|
# Process each node: only add leaf nodes (nodes without children)
|
|
for node in nodes:
|
|
if node.get("children"):
|
|
# If node has children, add only the children (leaf nodes)
|
|
for child in node["children"]:
|
|
# Skip if already added
|
|
if child["name"] in added_names:
|
|
continue
|
|
child["proxy_id"] = proxy_node["name"]
|
|
return_nodes.append(child)
|
|
added_names.add(child["name"])
|
|
else:
|
|
# Skip if already added
|
|
if node["name"] in added_names:
|
|
continue
|
|
node["proxy_id"] = proxy_node["name"]
|
|
return_nodes.append(node)
|
|
added_names.add(node["name"])
|
|
return return_nodes
|
|
except Exception as e:
|
|
print(f"Failed to fetch nodes from proxy {proxy_node['name']}: {e}")
|
|
return []
|
|
|
|
|
|
async def fetch_all_child_nodes_concurrently(
|
|
proxy_nodes: list[dict],
|
|
) -> list[dict]:
|
|
"""Fetch child nodes from multiple proxy nodes concurrently.
|
|
|
|
Args:
|
|
proxy_nodes: Proxy dicts to query. Each dict is mutated in
|
|
place with a new ``nodes`` field.
|
|
|
|
Returns:
|
|
The same ``proxy_nodes`` list with ``nodes`` populated.
|
|
"""
|
|
tasks = [fetch_child_nodes_from_proxy(proxy) for proxy in proxy_nodes]
|
|
results = await asyncio.gather(*tasks, return_exceptions=True)
|
|
|
|
# Process results and handle exceptions
|
|
for i, result in enumerate(results):
|
|
if isinstance(result, Exception):
|
|
print(f"Error fetching nodes from proxy {proxy_nodes[i]['name']}: {result}")
|
|
proxy_nodes[i]["nodes"] = []
|
|
else:
|
|
proxy_nodes[i]["nodes"] = result
|
|
|
|
return proxy_nodes
|
|
|
|
|
|
async def _fetch_child_nodes_from_api_nodes(
|
|
host: str, port: str, proxy_name: str
|
|
) -> list:
|
|
"""Try to fetch child nodes via /api/nodes from a multiProcess server.
|
|
|
|
Returns a non-empty list of child node dicts when the target is a
|
|
multiProcess lmcache server that exposes /api/nodes. Returns an
|
|
empty list for inProcess nodes (connection error or no children).
|
|
"""
|
|
try:
|
|
url = "http://%s:%s/api/nodes" % (host, port)
|
|
async with httpx.AsyncClient(timeout=5.0) as client:
|
|
response = await client.get(url)
|
|
response.raise_for_status()
|
|
nodes_data = response.json().get("nodes", [])
|
|
children = []
|
|
for node in nodes_data:
|
|
if node.get("children"):
|
|
for child in node["children"]:
|
|
child["proxy_id"] = proxy_name
|
|
children.append(child)
|
|
else:
|
|
node["proxy_id"] = proxy_name
|
|
children.append(node)
|
|
return children
|
|
except Exception:
|
|
return []
|
|
|
|
|
|
async def fetch_nodes_from_supplier(url: str) -> list[dict]:
|
|
"""Fetch node information from node supplier.
|
|
|
|
For each discovered node, attempt to retrieve its child nodes via
|
|
``/api/nodes`` (multiProcess lmcache server). If that call
|
|
succeeds and returns children, those children are used as the
|
|
``nodes`` list (multiProcess mode). Otherwise the node itself is
|
|
used as the sole child (inProcess / leaf mode).
|
|
"""
|
|
try:
|
|
async with httpx.AsyncClient(timeout=10.0) as client:
|
|
response = await client.get(url)
|
|
response.raise_for_status()
|
|
data = response.json()
|
|
|
|
unique_nodes: dict[str, dict] = {}
|
|
for api_address, info in data.get("processInfos", {}).items():
|
|
for entity in info.get("lmCacheInfoEntities", []):
|
|
addr = entity["apiAddress"]
|
|
if addr.startswith("http://"):
|
|
addr = addr[7:]
|
|
host, port = addr.split(":")
|
|
node_key = "%s:%s" % (host, port)
|
|
if node_key not in unique_nodes:
|
|
unique_nodes[node_key] = {
|
|
"host": host,
|
|
"port": port,
|
|
}
|
|
|
|
# Fetch child nodes concurrently for all discovered nodes.
|
|
proxy_list = [
|
|
{
|
|
"name": "proxy_%s" % node_key.replace(":", "_"),
|
|
"host": info["host"],
|
|
"port": info["port"],
|
|
}
|
|
for node_key, info in unique_nodes.items()
|
|
]
|
|
child_tasks = [
|
|
_fetch_child_nodes_from_api_nodes(p["host"], p["port"], p["name"])
|
|
for p in proxy_list
|
|
]
|
|
child_results = await asyncio.gather(*child_tasks, return_exceptions=True)
|
|
|
|
result = []
|
|
for proxy, children in zip(proxy_list, child_results, strict=False):
|
|
if isinstance(children, Exception) or not children:
|
|
# inProcess mode: node itself is the leaf target.
|
|
leaf = {
|
|
"name": proxy["name"],
|
|
"host": proxy["host"],
|
|
"port": proxy["port"],
|
|
}
|
|
nodes = [leaf]
|
|
else:
|
|
# multiProcess mode: use discovered child nodes.
|
|
assert not isinstance(children, BaseException)
|
|
nodes = children
|
|
result.append(
|
|
{
|
|
"name": proxy["name"],
|
|
"host": proxy["host"],
|
|
"port": proxy["port"],
|
|
"nodes": nodes,
|
|
}
|
|
)
|
|
return result
|
|
except Exception as e:
|
|
print(f"Failed to fetch nodes from supplier: {e}")
|
|
return []
|
|
|
|
|
|
def load_config(config_path: str | None = None) -> None:
|
|
"""Load proxy/node list from a JSON config file.
|
|
|
|
Args:
|
|
config_path: Optional path to the JSON file. When ``None`` the
|
|
packaged ``config.json`` is used.
|
|
"""
|
|
try:
|
|
# Prioritize user-specified configuration file
|
|
if config_path:
|
|
with open(config_path, "r") as f:
|
|
_node_registry.replace(json.load(f))
|
|
print(
|
|
f"Loaded {len(target_nodes)} target nodes from specified path: "
|
|
f"{config_path}"
|
|
)
|
|
else:
|
|
# Use package resource path as default configuration
|
|
default_config_path = _package_resource_path("config.json")
|
|
with open(default_config_path, "r") as f:
|
|
_node_registry.replace(json.load(f))
|
|
print(f"Loaded default configuration with {len(target_nodes)} target nodes")
|
|
except Exception as e:
|
|
print(f"Failed to load configuration file: {e}")
|
|
_node_registry.replace([])
|
|
|
|
|
|
def validate_node(node: dict, is_proxy: bool = False) -> bool:
|
|
"""Validate a single node configuration dict.
|
|
|
|
Args:
|
|
node: Candidate node dict.
|
|
is_proxy: Reserved for future proxy-specific validation.
|
|
|
|
Returns:
|
|
True when ``node`` has the required ``name``/``host``/``port``
|
|
keys and (optionally) a string ``proxy_id``.
|
|
"""
|
|
if not isinstance(node, dict):
|
|
return False
|
|
|
|
required_keys = {"name", "host", "port"}
|
|
if not required_keys.issubset(node.keys()):
|
|
return False
|
|
|
|
if "proxy_id" in node and node["proxy_id"]:
|
|
if not isinstance(node["proxy_id"], str):
|
|
return False
|
|
|
|
return True
|
|
|
|
|
|
def validate_nodes(nodes: list) -> bool:
|
|
"""Validate a list of node dicts; see :func:`validate_node`."""
|
|
if not isinstance(nodes, list):
|
|
return False
|
|
|
|
return all(validate_node(node) for node in nodes)
|
|
|
|
|
|
@router.get("/api/nodes")
|
|
async def get_all_nodes() -> dict:
|
|
"""Get all nodes in tree structure (proxies with their child nodes).
|
|
|
|
Returns:
|
|
``{"nodes": [...]}`` where each element is a proxy dict whose
|
|
``children`` list contains the leaf nodes behind it.
|
|
"""
|
|
all_nodes = []
|
|
for proxy in target_nodes:
|
|
# Create proxy node with children property
|
|
proxy_node = {
|
|
"id": f"proxy_{proxy['name']}",
|
|
"name": proxy["name"],
|
|
"host": proxy["host"],
|
|
"port": proxy["port"],
|
|
"is_proxy": True,
|
|
"children": [],
|
|
}
|
|
|
|
# Add child nodes
|
|
for node in proxy.get("nodes", []):
|
|
proxy_node["children"].append(
|
|
{
|
|
"id": "node_%s" % node["name"],
|
|
"name": node["name"],
|
|
"host": node["host"],
|
|
"port": node["port"],
|
|
"is_proxy": False,
|
|
"proxy_id": proxy["name"],
|
|
}
|
|
)
|
|
|
|
all_nodes.append(proxy_node)
|
|
|
|
return {"nodes": all_nodes}
|
|
|
|
|
|
@router.get("/api/proxies/{proxy_name}/refresh")
|
|
async def refresh_proxy_nodes(proxy_name: str):
|
|
"""Refresh child nodes of a proxy"""
|
|
proxy = next((p for p in target_nodes if p["name"] == proxy_name), None)
|
|
if not proxy:
|
|
raise HTTPException(status_code=404, detail="Proxy not found")
|
|
|
|
try:
|
|
child_nodes = await fetch_child_nodes_from_proxy(proxy)
|
|
proxy["nodes"] = child_nodes
|
|
return {"status": "success", "nodes": child_nodes}
|
|
except HTTPException:
|
|
raise
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=str(e)) from e
|
|
|
|
|
|
@router.get("/api/proxies")
|
|
async def get_proxies():
|
|
"""Get all proxy nodes (without child nodes)"""
|
|
proxies = [
|
|
{"name": proxy["name"], "host": proxy["host"], "port": proxy["port"]}
|
|
for proxy in target_nodes
|
|
]
|
|
return {"proxies": proxies}
|
|
|
|
|
|
@router.get("/api/proxies/{proxy_name}/nodes")
|
|
async def get_proxy_nodes(proxy_name: str):
|
|
"""Get all nodes under specified proxy"""
|
|
proxy = next((p for p in target_nodes if p["name"] == proxy_name), None)
|
|
if not proxy:
|
|
raise HTTPException(status_code=404, detail="Proxy not found")
|
|
return {"nodes": proxy["nodes"]}
|
|
|
|
|
|
# ==== Node Management Endpoints ====
|
|
@router.post("/api/proxies")
|
|
async def add_proxy(request: Request):
|
|
"""Add a new proxy node"""
|
|
try:
|
|
new_proxy = await request.json()
|
|
if not validate_node(new_proxy, is_proxy=True):
|
|
raise HTTPException(status_code=400, detail="Invalid proxy format")
|
|
|
|
# Check for duplicate names
|
|
if any(proxy["name"] == new_proxy["name"] for proxy in target_nodes):
|
|
raise HTTPException(status_code=409, detail="Proxy name already exists")
|
|
|
|
# Ensure nodes field exists
|
|
if "nodes" not in new_proxy:
|
|
new_proxy["nodes"] = []
|
|
|
|
target_nodes.append(new_proxy)
|
|
return {"status": "success", "message": "Proxy added"}
|
|
except HTTPException:
|
|
raise
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=str(e)) from e
|
|
|
|
|
|
@router.post("/api/proxies/{proxy_name}/nodes")
|
|
async def add_node_to_proxy(proxy_name: str, request: Request):
|
|
"""Add child node to proxy"""
|
|
try:
|
|
new_node = await request.json()
|
|
if not validate_node(new_node):
|
|
raise HTTPException(status_code=400, detail="Invalid node format")
|
|
|
|
# Find corresponding proxy
|
|
proxy = next((p for p in target_nodes if p["name"] == proxy_name), None)
|
|
if not proxy:
|
|
raise HTTPException(status_code=404, detail="Proxy not found")
|
|
|
|
# Check for duplicate names
|
|
if any(node["name"] == new_node["name"] for node in proxy["nodes"]):
|
|
raise HTTPException(status_code=409, detail="Node name already exists")
|
|
|
|
proxy["nodes"].append(new_node)
|
|
return {"status": "success", "message": "Node added to proxy"}
|
|
except HTTPException:
|
|
raise
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=str(e)) from e
|
|
|
|
|
|
@router.put("/api/nodes/{node_name}")
|
|
async def update_node(node_name: str, request: Request):
|
|
"""Update a proxy or child node identified by ``node_name``.
|
|
|
|
Searches top-level proxies first, then children of every proxy,
|
|
so both direct proxies and supplier-discovered leaf nodes can be
|
|
updated through the same endpoint (mirrors ``proxy_request_by_name``).
|
|
"""
|
|
try:
|
|
updated_node = await request.json()
|
|
if not validate_node(updated_node):
|
|
raise HTTPException(status_code=400, detail="Invalid node format")
|
|
|
|
for i, node in enumerate(target_nodes):
|
|
if node["name"] == node_name:
|
|
# Preserve existing children when the caller did not
|
|
# provide one; proxies are expected to keep ``nodes``.
|
|
if "nodes" not in updated_node and "nodes" in node:
|
|
updated_node["nodes"] = node["nodes"]
|
|
target_nodes[i] = updated_node
|
|
return {"status": "success", "message": "Node updated"}
|
|
|
|
for proxy in target_nodes:
|
|
children = proxy.get("nodes", [])
|
|
for j, child in enumerate(children):
|
|
if child["name"] == node_name:
|
|
children[j] = updated_node
|
|
return {
|
|
"status": "success",
|
|
"message": "Child node updated",
|
|
}
|
|
|
|
raise HTTPException(status_code=404, detail="Node not found")
|
|
except HTTPException:
|
|
raise
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=str(e)) from e
|
|
|
|
|
|
@router.delete("/api/nodes/{node_name}")
|
|
async def delete_node(node_name: str):
|
|
"""Delete a proxy or child node by ``node_name``.
|
|
|
|
Mirrors :func:`update_node`: matches top-level proxies first,
|
|
falls back to scanning children of each proxy.
|
|
"""
|
|
try:
|
|
original_count = len(target_nodes)
|
|
_node_registry.replace(
|
|
[node for node in target_nodes if node["name"] != node_name]
|
|
)
|
|
|
|
if len(target_nodes) != original_count:
|
|
return {"status": "success", "message": "Node deleted"}
|
|
|
|
for proxy in target_nodes:
|
|
children = proxy.get("nodes", [])
|
|
for idx, child in enumerate(children):
|
|
if child["name"] == node_name:
|
|
del children[idx]
|
|
return {
|
|
"status": "success",
|
|
"message": "Child node deleted",
|
|
}
|
|
|
|
raise HTTPException(status_code=404, detail="Node not found")
|
|
except HTTPException:
|
|
raise
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=str(e)) from e
|
|
|
|
|
|
@router.api_route(
|
|
"/proxy2/{node_name}/{path:path}",
|
|
methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS", "HEAD"],
|
|
)
|
|
async def proxy_request_by_name(request: Request, node_name: str, path: str):
|
|
"""Proxy requests using node name as identifier.
|
|
|
|
Searches top-level target_nodes first, then child nodes of every
|
|
proxy, so both direct nodes and supplier-discovered leaf nodes are
|
|
reachable with a single /proxy2/{name}/{path} call.
|
|
"""
|
|
# 1. top-level proxy nodes
|
|
node = next((n for n in target_nodes if n["name"] == node_name), None)
|
|
|
|
# 2. child nodes of every proxy
|
|
if not node:
|
|
for proxy in target_nodes:
|
|
node = next(
|
|
(n for n in proxy.get("nodes", []) if n["name"] == node_name),
|
|
None,
|
|
)
|
|
if node:
|
|
break
|
|
|
|
if not node:
|
|
raise HTTPException(
|
|
status_code=404, detail=f"Node with name '{node_name}' not found"
|
|
)
|
|
|
|
# Use existing proxy_request logic
|
|
return await proxy_request(
|
|
request, target_host=node["host"], target_port_or_socket=node["port"], path=path
|
|
)
|
|
|
|
|
|
@router.api_route(
|
|
"/proxy/{target_host}/{target_port_or_socket}/{path:path}",
|
|
methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS", "HEAD"],
|
|
)
|
|
async def proxy_request(
|
|
request: Request, target_host: str, target_port_or_socket: str | int, path: str
|
|
):
|
|
"""Proxy requests to the specified target host and port or socket path.
|
|
|
|
For security, non-socket targets must match a host/port already
|
|
registered in :data:`_node_registry`; this prevents the endpoint
|
|
from being used as an open relay (SSRF). Socket paths are
|
|
accepted as-is because they are local UDS endpoints.
|
|
"""
|
|
target_port_or_socket = unquote(str(target_port_or_socket))
|
|
# Check if target_port_or_socket is a socket path (contains '/')
|
|
is_socket_path = "/" in target_port_or_socket
|
|
|
|
if is_socket_path:
|
|
# For socket paths, use UDS transport
|
|
socket_path = target_port_or_socket
|
|
target_url = f"http://localhost/{path}"
|
|
|
|
# Create UDS transport
|
|
transport = httpx.AsyncHTTPTransport(uds=socket_path)
|
|
else:
|
|
port = target_port_or_socket
|
|
# SSRF guard: resolve against the registry and reuse the
|
|
# trusted host/port from there when building the outbound URL.
|
|
# This keeps user-controlled values out of the URL sink.
|
|
resolved = _node_registry.resolve(target_host, port)
|
|
if resolved is None:
|
|
raise HTTPException(
|
|
status_code=403,
|
|
detail=("Target %s:%s is not a registered node" % (target_host, port)),
|
|
)
|
|
safe_host, safe_port = resolved
|
|
target_url = f"http://{safe_host}:{safe_port}/{path}"
|
|
transport = None # Use default transport
|
|
|
|
headers = {}
|
|
for key, value in request.headers.items():
|
|
if key.lower() in [
|
|
"host",
|
|
"content-length",
|
|
"connection",
|
|
"keep-alive",
|
|
"proxy-authenticate",
|
|
"proxy-authorization",
|
|
"te",
|
|
"trailers",
|
|
"transfer-encoding",
|
|
"upgrade",
|
|
]:
|
|
continue
|
|
headers[key] = value
|
|
|
|
body = await request.body()
|
|
|
|
# Create client with appropriate transport
|
|
async with httpx.AsyncClient(transport=transport) as client:
|
|
try:
|
|
response = await client.request(
|
|
method=request.method,
|
|
url=target_url,
|
|
headers=headers,
|
|
content=body,
|
|
params=request.query_params,
|
|
timeout=60.0,
|
|
)
|
|
|
|
response_headers = {}
|
|
for key, value in response.headers.items():
|
|
if key.lower() in [
|
|
"connection",
|
|
"keep-alive",
|
|
"proxy-authenticate",
|
|
"proxy-authorization",
|
|
"te",
|
|
"trailers",
|
|
"transfer-encoding",
|
|
"upgrade",
|
|
]:
|
|
continue
|
|
response_headers[key] = value
|
|
|
|
return PlainTextResponse(
|
|
content=response.content,
|
|
headers=response_headers,
|
|
media_type=response.headers.get("content-type", "text/plain"),
|
|
status_code=response.status_code,
|
|
)
|
|
|
|
except httpx.ConnectError as e:
|
|
if is_socket_path:
|
|
detail = f"Failed to connect to socket: {socket_path}"
|
|
else:
|
|
detail = f"Failed to connect to target service {target_host}:{port}"
|
|
raise HTTPException(status_code=502, detail=detail) from e
|
|
except httpx.TimeoutException as e:
|
|
if is_socket_path:
|
|
detail = f"Connection to socket {socket_path} timed out"
|
|
else:
|
|
detail = f"Connection to target service {target_host}:{port} timed out"
|
|
raise HTTPException(status_code=504, detail=detail) from e
|
|
except httpx.HTTPError as e:
|
|
raise HTTPException(
|
|
status_code=502,
|
|
detail=f"Error communicating with target service: {str(e)}",
|
|
) from e
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=f"Proxy error: {str(e)}") from e
|
|
|
|
|
|
@router.get("/health")
|
|
async def health_check():
|
|
return {"status": "healthy", "service": "lmcache-monitor"}
|
|
|
|
|
|
@router.get("/api/heartbeat/status")
|
|
async def get_heartbeat_status():
|
|
"""Get heartbeat status"""
|
|
return heartbeat_service.status()
|
|
|
|
|
|
@router.post("/api/heartbeat/start")
|
|
async def start_heartbeat_api(request: Request):
|
|
"""Start heartbeat service"""
|
|
try:
|
|
data = await request.json()
|
|
heartbeat_url = data.get("heartbeat_url")
|
|
initial_delay = data.get("initial_delay", 0)
|
|
interval = data.get("interval", 30)
|
|
|
|
if not heartbeat_url:
|
|
raise HTTPException(status_code=400, detail="heartbeat_url is required")
|
|
|
|
heartbeat_service.start(heartbeat_url, initial_delay, interval)
|
|
return {"status": "success", "message": "Heartbeat service started"}
|
|
except HTTPException:
|
|
raise
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=str(e)) from e
|
|
|
|
|
|
@router.post("/api/heartbeat/stop")
|
|
async def stop_heartbeat_api():
|
|
"""Stop heartbeat service"""
|
|
try:
|
|
heartbeat_service.stop()
|
|
return {"status": "success", "message": "Heartbeat service stopped"}
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=str(e)) from e
|
|
|
|
|
|
async def load_nodes_from_supplier(node_supplier_url: str | None = None) -> bool:
|
|
"""Load node information from node supplier.
|
|
|
|
Args:
|
|
node_supplier_url: Discovery service URL, e.g. ``/lmcache_infos``.
|
|
|
|
Returns:
|
|
``True`` when a non-empty node list was retrieved and stored.
|
|
"""
|
|
if not node_supplier_url:
|
|
return False
|
|
|
|
print(f"Fetching nodes from supplier: {node_supplier_url}")
|
|
nodes = await fetch_nodes_from_supplier(node_supplier_url)
|
|
if nodes:
|
|
_node_registry.replace(nodes)
|
|
print(f"Loaded {len(target_nodes)} proxy nodes from supplier")
|
|
|
|
# Child nodes are already populated by fetch_nodes_from_supplier;
|
|
# no secondary /api/nodes round-trip is needed.
|
|
for proxy in target_nodes:
|
|
print(f"Proxy {proxy['name']} loaded {len(proxy['nodes'])} child nodes")
|
|
return True
|
|
else:
|
|
print("Warning: No nodes loaded from supplier")
|
|
return False
|
|
|
|
|
|
async def initialize_nodes(node_supplier_url: str | None = None) -> None:
|
|
"""Initialize node configuration from CLI args or supplier URL.
|
|
|
|
Args:
|
|
node_supplier_url: Optional discovery service URL. When set,
|
|
nodes are loaded via :func:`load_nodes_from_supplier`.
|
|
Otherwise falls back to ``args.nodes`` / ``args.config``.
|
|
"""
|
|
global args
|
|
|
|
if args is None:
|
|
raise ValueError("args is not initialized")
|
|
|
|
if node_supplier_url:
|
|
await load_nodes_from_supplier(node_supplier_url)
|
|
elif args.nodes:
|
|
try:
|
|
nodes = json.loads(args.nodes)
|
|
if validate_nodes(nodes):
|
|
_node_registry.replace(
|
|
[
|
|
{
|
|
"name": "local_proxy",
|
|
"host": args.host,
|
|
"port": args.port,
|
|
"nodes": nodes,
|
|
}
|
|
]
|
|
)
|
|
print(f"Loaded {len(nodes)} target nodes from command line arguments")
|
|
except json.JSONDecodeError:
|
|
print("Failed to parse nodes JSON parameter")
|
|
elif args.config:
|
|
load_config(args.config)
|
|
|
|
|
|
# Minimum seconds between two supplier refreshes triggered by ``/``.
|
|
# Each refresh fans out to every registered proxy's ``/api/nodes``
|
|
# endpoint, so an unthrottled ``/`` would DOS both the supplier and
|
|
# every proxy. 30s matches the default heartbeat interval.
|
|
_SUPPLIER_REFRESH_INTERVAL_SEC = 30.0
|
|
_supplier_last_refresh: float = 0.0
|
|
_supplier_refresh_lock = asyncio.Lock()
|
|
|
|
|
|
async def _maybe_refresh_from_supplier(node_supplier_url: str) -> None:
|
|
"""Refresh the node registry from supplier, at most once per interval.
|
|
|
|
The first caller within the interval performs the refresh; other
|
|
concurrent callers return immediately (stale-on-read). This keeps
|
|
the homepage responsive even under high traffic.
|
|
"""
|
|
global _supplier_last_refresh
|
|
now = time.monotonic()
|
|
if now - _supplier_last_refresh < _SUPPLIER_REFRESH_INTERVAL_SEC:
|
|
return
|
|
if _supplier_refresh_lock.locked():
|
|
return
|
|
async with _supplier_refresh_lock:
|
|
now = time.monotonic()
|
|
if now - _supplier_last_refresh < _SUPPLIER_REFRESH_INTERVAL_SEC:
|
|
return
|
|
await initialize_nodes(node_supplier_url)
|
|
_supplier_last_refresh = time.monotonic()
|
|
|
|
|
|
@router.get("/")
|
|
async def serve_frontend():
|
|
"""Return frontend homepage.
|
|
|
|
When a node supplier URL is configured, trigger a throttled
|
|
background-style refresh so opening the homepage repeatedly does
|
|
not hammer the supplier or every proxy's ``/api/nodes``.
|
|
"""
|
|
if args.node_supplier_url:
|
|
await _maybe_refresh_from_supplier(args.node_supplier_url)
|
|
|
|
try:
|
|
# Use package resource path
|
|
index_path = _package_resource_path("static/index.html")
|
|
return FileResponse(index_path)
|
|
except Exception:
|
|
# Development environment uses local files
|
|
return FileResponse("static/index.html")
|
|
|
|
|
|
# Helper function to fetch metrics from a single node
|
|
async def _fetch_node_metrics(node):
|
|
"""Fetch metrics from a single node"""
|
|
try:
|
|
# Check if port is a socket path
|
|
is_socket_path = "/" in node["port"]
|
|
|
|
if is_socket_path:
|
|
# Use UDS transport for socket paths
|
|
transport = httpx.AsyncHTTPTransport(uds=node["port"])
|
|
# Use localhost as host
|
|
url = "http://localhost/metrics"
|
|
async with httpx.AsyncClient(transport=transport, timeout=5.0) as client:
|
|
response = await client.get(url)
|
|
response.raise_for_status()
|
|
return response.text
|
|
else:
|
|
# Build URL for regular port
|
|
url = f"http://{node['host']}:{node['port']}/metrics"
|
|
async with httpx.AsyncClient(timeout=5.0) as client:
|
|
response = await client.get(url)
|
|
response.raise_for_status()
|
|
return response.text
|
|
except Exception as e:
|
|
return f"# ERROR: Failed to get metrics from {node['name']}: {str(e)}\n"
|
|
|
|
|
|
@router.get("/metrics")
|
|
async def aggregated_metrics():
|
|
"""Aggregate /metrics from all leaf (child) nodes across every proxy.
|
|
|
|
Previously only nodes under the ``local_proxy`` entry were
|
|
aggregated, which silently returned nothing in supplier mode
|
|
(where proxy names follow the ``proxy_<host>_<port>`` pattern).
|
|
Now every proxy's ``nodes`` list is flattened.
|
|
"""
|
|
if not target_nodes:
|
|
return PlainTextResponse("# No nodes configured\n", status_code=404)
|
|
|
|
nodes: list[dict] = []
|
|
seen: set[str] = set()
|
|
for proxy in target_nodes:
|
|
for child in proxy.get("nodes", []) or []:
|
|
# De-duplicate by name so the same leaf reported via
|
|
# multiple proxies is only scraped once.
|
|
name = child.get("name")
|
|
if name and name in seen:
|
|
continue
|
|
if name:
|
|
seen.add(name)
|
|
nodes.append(child)
|
|
|
|
if not nodes:
|
|
return PlainTextResponse(
|
|
"# No nodes available for metrics collection\n", status_code=404
|
|
)
|
|
|
|
metrics_results = await asyncio.gather(
|
|
*[_fetch_node_metrics(node) for node in nodes]
|
|
)
|
|
|
|
# Combine all metrics with node name as comment header
|
|
aggregated = ""
|
|
for i, metrics in enumerate(metrics_results):
|
|
node = nodes[i]
|
|
aggregated += (
|
|
f"# Metrics from node: {node['name']} ({node['host']}:{node['port']})\n"
|
|
)
|
|
aggregated += metrics
|
|
aggregated += "\n\n"
|
|
|
|
return PlainTextResponse(aggregated)
|
|
|
|
|
|
def create_app():
|
|
"""Create and configure FastAPI application"""
|
|
app = FastAPI(
|
|
title="Flexible Proxy Server",
|
|
description="HTTP proxy service supporting specified target hosts and ports",
|
|
)
|
|
app.include_router(router)
|
|
|
|
# Get static file path (prefer package resources)
|
|
try:
|
|
static_path = _package_resource_path("static")
|
|
except Exception:
|
|
static_path = os.path.join(os.path.dirname(__file__), "static")
|
|
|
|
# Mount static file service
|
|
app.mount("/static", StaticFiles(directory=static_path), name="static")
|
|
|
|
return app
|
|
|
|
|
|
def main():
|
|
global args
|
|
parser = argparse.ArgumentParser(description="LMCache Cluster Monitoring Tool")
|
|
parser.add_argument(
|
|
"--port", type=int, default=8000, help="Service port, default 8000"
|
|
)
|
|
parser.add_argument(
|
|
"--host", type=str, default="0.0.0.0", help="Bind host address, default 0.0.0.0"
|
|
)
|
|
parser.add_argument(
|
|
"--config",
|
|
type=str,
|
|
default=None,
|
|
help="Specify configuration file path, default uses internal config.json",
|
|
)
|
|
parser.add_argument(
|
|
"--nodes",
|
|
type=str,
|
|
default=None,
|
|
help="Directly specify target nodes as a JSON string. "
|
|
'Example: \'[{"name":"node1","host":"127.0.0.1","port":8001}]\'',
|
|
)
|
|
parser.add_argument(
|
|
"--heartbeat-url",
|
|
type=str,
|
|
default=None,
|
|
help="Heartbeat service URL, e.g.: http://example.com/heartbeat",
|
|
)
|
|
parser.add_argument(
|
|
"--heartbeat-initial-delay",
|
|
type=int,
|
|
default=0,
|
|
help="Initial delay before starting heartbeat (seconds), default 0",
|
|
)
|
|
parser.add_argument(
|
|
"--heartbeat-interval",
|
|
type=int,
|
|
default=30,
|
|
help="Heartbeat interval (seconds), default 30",
|
|
)
|
|
parser.add_argument(
|
|
"--node-supplier-url",
|
|
type=str,
|
|
default=None,
|
|
help="URL to fetch node information from, e.g.: http://example.com/lmcache_infos",
|
|
)
|
|
parser.add_argument(
|
|
"--log-level",
|
|
type=str,
|
|
default="warning",
|
|
choices=["critical", "error", "warning", "warn", "info", "debug", "trace"],
|
|
help="Uvicorn log level, default: warn",
|
|
)
|
|
parser.add_argument(
|
|
"--no-http",
|
|
action="store_true",
|
|
default=False,
|
|
help="Disable HTTP server startup (heartbeat still runs)",
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
# Initialize node configuration
|
|
asyncio.run(initialize_nodes(args.node_supplier_url))
|
|
|
|
app = create_app()
|
|
print(f"Monitoring service running at http://{args.host}:{args.port}")
|
|
print(f"Node management: http://{args.host}:{args.port}/static/index.html")
|
|
|
|
# Start heartbeat service if URL is configured
|
|
if args.heartbeat_url:
|
|
# Set application configuration for heartbeat service
|
|
heartbeat_service.set_app_config(args.host, args.port, target_nodes)
|
|
|
|
print("Starting heartbeat service...")
|
|
print(f"Heartbeat URL: {args.heartbeat_url}")
|
|
print(f"Initial delay: {args.heartbeat_initial_delay}s")
|
|
print(f"Interval: {args.heartbeat_interval}s")
|
|
print(f"API Address: http://{args.host}:{args.port}")
|
|
print(f"Target nodes count: {len(target_nodes)}")
|
|
|
|
heartbeat_service.start(
|
|
args.heartbeat_url, args.heartbeat_initial_delay, args.heartbeat_interval
|
|
)
|
|
else:
|
|
print("Heartbeat URL not configured, heartbeat disabled")
|
|
|
|
if args.no_http:
|
|
print("HTTP server disabled (--no-http), running heartbeat only")
|
|
try:
|
|
stop_event = threading.Event()
|
|
stop_event.wait()
|
|
finally:
|
|
print("Shutting down application...")
|
|
heartbeat_service.stop()
|
|
return
|
|
|
|
try:
|
|
uvicorn.run(app, host=args.host, port=args.port, log_level=args.log_level)
|
|
finally:
|
|
# Stop heartbeat service when app closes
|
|
print("Shutting down application...")
|
|
heartbeat_service.stop()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|