From 8d7b4363d1e99cabbfc53c00e038c1fc07224002 Mon Sep 17 00:00:00 2001 From: sufubao Date: Wed, 26 Aug 2026 13:41:05 +0800 Subject: [PATCH 1/2] feat(pd): expose stage and node load metrics --- .../httpserver_for_pd_master/manager.py | 60 ++++++++-- lightllm/server/metrics/manager.py | 5 +- lightllm/server/metrics/metrics.py | 16 ++- .../test_pd_master_metrics.py | 110 ++++++++++++++++++ 4 files changed, 174 insertions(+), 17 deletions(-) create mode 100644 test/test_pd_selector/test_pd_master_metrics.py diff --git a/lightllm/server/httpserver_for_pd_master/manager.py b/lightllm/server/httpserver_for_pd_master/manager.py index aebe039426..c5c3b11aa1 100644 --- a/lightllm/server/httpserver_for_pd_master/manager.py +++ b/lightllm/server/httpserver_for_pd_master/manager.py @@ -41,13 +41,16 @@ def __init__( self.metric_client = MetricClient(get_shm_port_args().metric_port) self.id_gen = ReqIDGenerator() - self.pd_manager = PDManager(args) + self.pd_manager = PDManager(args, self.metric_client) self.req_id_to_out_inf: Dict[int, ReqStatus] = {} self.infos_queues = None # 这个需要延迟初始化，否则使用的loop不对 self.health_timeout = int(os.getenv("HEALTH_TIMEOUT", "200")) self.latest_success_infer_time = time.time() self.running_request_count = 0 + self.pd_stage_waiting_request_counts = {"prefill": 0, "decode": 0} + for stage in self.pd_stage_waiting_request_counts: + self.metric_client.gauge_set("lightllm_pd_master_stage_waiting_requests", 0, labels={"stage": stage}) self.tokenizer = get_tokenizer(args.model_dir, args.tokenizer_mode, trust_remote_code=args.trust_remote_code) @@ -354,6 +357,28 @@ async def raise_if_disconnected() -> None: except asyncio.TimeoutError: continue + def _change_pd_stage_waiting_requests(self, stage: str, delta: int) -> None: + self.pd_stage_waiting_request_counts[stage] += delta + self.metric_client.gauge_set( + "lightllm_pd_master_stage_waiting_requests", + self.pd_stage_waiting_request_counts[stage], + labels={"stage": stage}, + ) + + async def _wait_for_pd_stage( + self, + event: asyncio.Event, + request: Request, + timeout: float, + group_request_id: int, + stage: str, + ) -> None: + self._change_pd_stage_waiting_requests(stage, 1) + try: + await self._wait_for_event_or_disconnect(event, request, timeout, group_request_id, stage) + finally: + self._change_pd_stage_waiting_requests(stage, -1) + async def _log_req_header(self, request: Request, group_request_id: int): x_request_id = request.headers.get("X-Request-Id", "") x_session_id = request.headers.get("X-Session-Id", "") @@ -388,7 +413,7 @@ async def fetch_pd_stream( await p_node.websocket.send_bytes(pickle.dumps((ObjType.REQ, (prompt, sampling_params, multimodal_params)))) try: - await self._wait_for_event_or_disconnect( + await self._wait_for_pd_stage( prefill_prompt_ids_event, request, timeout=60, @@ -409,7 +434,7 @@ async def fetch_pd_stream( ) try: - await self._wait_for_event_or_disconnect( + await self._wait_for_pd_stage( up_status_event, request, timeout=180, @@ -431,13 +456,19 @@ async def fetch_pd_stream( first_token_gen = False needs_prefill_first_token = decode_node_info.ready_kv_len != len(prompt_ids) - 1 - prompt_cache_len_from_prefill = await self._wait_for_prefill_token_if_needed( - req_status=req_status, - request=request, - group_request_id=group_request_id, - needs_prefill_first_token=needs_prefill_first_token, - ready_kv_len=decode_node_info.ready_kv_len, - ) + if needs_prefill_first_token: + self._change_pd_stage_waiting_requests("prefill", 1) + try: + prompt_cache_len_from_prefill = await self._wait_for_prefill_token_if_needed( + req_status=req_status, + request=request, + group_request_id=group_request_id, + needs_prefill_first_token=needs_prefill_first_token, + ready_kv_len=decode_node_info.ready_kv_len, + ) + finally: + if needs_prefill_first_token: + self._change_pd_stage_waiting_requests("prefill", -1) while True: await req_status.wait_to_ready() @@ -749,8 +780,9 @@ async def put_tokens_to_front(self, token_list: List[Tuple[int, str, dict, Finis class PDManager: - def __init__(self, args: StartArgs): + def __init__(self, args: StartArgs, metric_client=None): self.args: StartArgs = args + self.metric_client = metric_client self.prefill_nodes: List[PD_Client_Obj] = [] self.decode_nodes: List[PD_Client_Obj] = [] self.url_to_pd_nodes: Dict[str, PD_Client_Obj] = {} @@ -874,6 +906,12 @@ def update_node_load_info(self, load_info: Optional[dict]): total_token_usage_rate = load_info["total_token_usage_rate"] pd_client = self.url_to_pd_nodes.get(client_ip_port) pd_client.run_status.total_token_usage_rate = total_token_usage_rate + if self.metric_client is not None: + self.metric_client.gauge_set( + "lightllm_pd_node_token_usage_ratio", + total_token_usage_rate, + labels={"role": pd_client.mode, "endpoint": client_ip_port}, + ) except BaseException as e: logger.warning(f"udpate node load info failed, load_info: {load_info} error: {str(e)}") return diff --git a/lightllm/server/metrics/manager.py b/lightllm/server/metrics/manager.py index 22f6426a77..f8f56dc8d3 100644 --- a/lightllm/server/metrics/manager.py +++ b/lightllm/server/metrics/manager.py @@ -55,8 +55,9 @@ def exposed_counter_inc_by(self, name: str, amount: float) -> None: def exposed_histogram_observe(self, name: str, value: float, label: str = None) -> None: return self.monitor.histogram_observe(name, value, label) - def exposed_gauge_set(self, name: str, value: float) -> None: - return self.monitor.gauge_set(name, value) + def exposed_gauge_set(self, name: str, value: float, labels: dict = None) -> None: + local_labels = None if labels is None else {key: labels[key] for key in labels} + return self.monitor.gauge_set(name, value, local_labels) def exposed_generate_latest(self) -> bytes: data = generate_latest(self.monitor.registry) diff --git a/lightllm/server/metrics/metrics.py b/lightllm/server/metrics/metrics.py index 0d42462c3f..b58b587b26 100644 --- a/lightllm/server/metrics/metrics.py +++ b/lightllm/server/metrics/metrics.py @@ -32,6 +32,8 @@ "lightllm_cache_hit_rate": "Prefix cache hit rate of latest completed request", "lightllm_gen_throughput": "Generation throughput of latest completed request (tokens/s)", "lightllm_num_running_reqs": "Number of running requests", + "lightllm_pd_node_token_usage_ratio": "Token capacity usage ratio reported by a PD node", + "lightllm_pd_master_stage_waiting_requests": "Requests waiting for a PD stage to become ready", } @@ -111,6 +113,8 @@ def init_metrics(self, args): self.create_gauge("lightllm_cache_hit_rate") self.create_gauge("lightllm_gen_throughput") self.create_gauge("lightllm_num_running_reqs") + self.create_gauge("lightllm_pd_node_token_usage_ratio", labelnames=["role", "endpoint"]) + self.create_gauge("lightllm_pd_master_stage_waiting_requests", labelnames=["stage"]) def create_histogram(self, name, buckets, labelnames=None): all_labels = ["model_name"] + (labelnames or []) @@ -122,8 +126,9 @@ def create_counter(self, name, labelnames=None): counter = Counter(name, MONITOR_INFO[name], labelnames=all_labels, registry=self.registry) self.monitor_registry[name] = counter - def create_gauge(self, name): - gauge = Gauge(name, MONITOR_INFO[name], labelnames=["model_name"], registry=self.registry) + def create_gauge(self, name, labelnames=None): + all_labels = ["model_name"] + (labelnames or []) + gauge = Gauge(name, MONITOR_INFO[name], labelnames=all_labels, registry=self.registry) self.monitor_registry[name] = gauge def counter_inc(self, name, label=None): @@ -141,8 +146,11 @@ def histogram_observe(self, name, value, label=None): else: self.monitor_registry[name].labels(model_name=self.model_name, method=label).observe(value) - def gauge_set(self, name, value): - self.monitor_registry[name].labels(model_name=self.model_name).set(value) + def gauge_set(self, name, value, labels=None): + metric_labels = {"model_name": self.model_name} + if labels: + metric_labels.update(labels) + self.monitor_registry[name].labels(**metric_labels).set(value) def push_metrices(self): if self.gateway_url is not None: diff --git a/test/test_pd_selector/test_pd_master_metrics.py b/test/test_pd_selector/test_pd_master_metrics.py new file mode 100644 index 0000000000..fa2ca30d1f --- /dev/null +++ b/test/test_pd_selector/test_pd_master_metrics.py @@ -0,0 +1,110 @@ +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, call + +import pytest +from prometheus_client import generate_latest + +from lightllm.server.httpserver_for_pd_master.manager import HttpServerManagerForPDMaster, PDManager +from lightllm.server.metrics.manager import MetricServer +from lightllm.server.metrics.metrics import Monitor +from lightllm.server.pd_io_struct import PD_Client_Obj +from lightllm.utils.error_utils import ServerBusyError + + +class RPyCDictProxyLike: + def __init__(self, values): + self.values = values + + def __iter__(self): + return iter(self.values) + + def __getitem__(self, key): + return self.values[key] + + +def test_pd_master_exports_node_load_with_role_and_endpoint_labels(): + args = SimpleNamespace( + metric_gateway=None, + job_name="lightllm", + grouping_key=None, + enable_monitor_auth=False, + model_name="test-model", + max_req_total_len=1024, + mtp_step=0, + select_p_d_node_strategy="random", + ) + monitor = Monitor(args) + metric_client = MagicMock() + metric_client.gauge_set.side_effect = monitor.gauge_set + manager = PDManager(args, metric_client) + manager.url_to_pd_nodes["10.0.0.1:28761"] = PD_Client_Obj( + node_id=1, + client_ip_port="10.0.0.1:28761", + mode="prefill", + start_args={}, + ) + manager.url_to_pd_nodes["10.0.0.2:28764"] = PD_Client_Obj( + node_id=2, + client_ip_port="10.0.0.2:28764", + mode="decode", + start_args={}, + ) + + manager.update_node_load_info({"client_ip_port": "10.0.0.1:28761", "total_token_usage_rate": 0.25}) + manager.update_node_load_info({"client_ip_port": "10.0.0.2:28764", "total_token_usage_rate": 0.75}) + monitor.gauge_set("lightllm_pd_master_stage_waiting_requests", 2, labels={"stage": "decode"}) + + metrics = generate_latest(monitor.registry).decode() + assert ( + 'lightllm_pd_node_token_usage_ratio{endpoint="10.0.0.1:28761",model_name="test-model",role="prefill"} 0.25' + in metrics + ) + assert ( + 'lightllm_pd_node_token_usage_ratio{endpoint="10.0.0.2:28764",model_name="test-model",role="decode"} 0.75' + in metrics + ) + assert 'lightllm_pd_master_stage_waiting_requests{model_name="test-model",stage="decode"} 2.0' in metrics + + +def test_metric_server_copies_rpyc_label_proxy_before_updating_gauge(): + args = SimpleNamespace( + metric_gateway=None, + job_name="lightllm", + grouping_key=None, + enable_monitor_auth=False, + model_name="test-model", + max_req_total_len=1024, + mtp_step=0, + push_interval=10, + ) + server = MetricServer(args) + + server.exposed_gauge_set( + "lightllm_pd_master_stage_waiting_requests", + 3, + RPyCDictProxyLike({"stage": "decode"}), + ) + + metrics = generate_latest(server.monitor.registry).decode() + assert 'lightllm_pd_master_stage_waiting_requests{model_name="test-model",stage="decode"} 3.0' in metrics + + +@pytest.mark.parametrize("stage", ["prefill", "decode"]) +def test_pd_master_stage_waiting_gauge_is_released_on_failure(stage): + async def run(): + manager = HttpServerManagerForPDMaster.__new__(HttpServerManagerForPDMaster) + manager.metric_client = MagicMock() + manager.pd_stage_waiting_request_counts = {"prefill": 0, "decode": 0} + manager._wait_for_event_or_disconnect = AsyncMock(side_effect=ServerBusyError()) + + with pytest.raises(ServerBusyError): + await manager._wait_for_pd_stage(AsyncMock(), AsyncMock(), 1, 123, stage) + + assert manager.pd_stage_waiting_request_counts[stage] == 0 + assert manager.metric_client.gauge_set.call_args_list == [ + call("lightllm_pd_master_stage_waiting_requests", 1, labels={"stage": stage}), + call("lightllm_pd_master_stage_waiting_requests", 0, labels={"stage": stage}), + ] + + asyncio.run(run()) From 7193c56b283565cc66c7274a94f79ec529067b2d Mon Sep 17 00:00:00 2001 From: sufubao Date: Wed, 26 Aug 2026 13:44:11 +0800 Subject: [PATCH 2/2] chore(pd): remove metrics unit tests --- .../test_pd_master_metrics.py | 110 ------------------ 1 file changed, 110 deletions(-) delete mode 100644 test/test_pd_selector/test_pd_master_metrics.py diff --git a/test/test_pd_selector/test_pd_master_metrics.py b/test/test_pd_selector/test_pd_master_metrics.py deleted file mode 100644 index fa2ca30d1f..0000000000 --- a/test/test_pd_selector/test_pd_master_metrics.py +++ /dev/null @@ -1,110 +0,0 @@ -import asyncio -from types import SimpleNamespace -from unittest.mock import AsyncMock, MagicMock, call - -import pytest -from prometheus_client import generate_latest - -from lightllm.server.httpserver_for_pd_master.manager import HttpServerManagerForPDMaster, PDManager -from lightllm.server.metrics.manager import MetricServer -from lightllm.server.metrics.metrics import Monitor -from lightllm.server.pd_io_struct import PD_Client_Obj -from lightllm.utils.error_utils import ServerBusyError - - -class RPyCDictProxyLike: - def __init__(self, values): - self.values = values - - def __iter__(self): - return iter(self.values) - - def __getitem__(self, key): - return self.values[key] - - -def test_pd_master_exports_node_load_with_role_and_endpoint_labels(): - args = SimpleNamespace( - metric_gateway=None, - job_name="lightllm", - grouping_key=None, - enable_monitor_auth=False, - model_name="test-model", - max_req_total_len=1024, - mtp_step=0, - select_p_d_node_strategy="random", - ) - monitor = Monitor(args) - metric_client = MagicMock() - metric_client.gauge_set.side_effect = monitor.gauge_set - manager = PDManager(args, metric_client) - manager.url_to_pd_nodes["10.0.0.1:28761"] = PD_Client_Obj( - node_id=1, - client_ip_port="10.0.0.1:28761", - mode="prefill", - start_args={}, - ) - manager.url_to_pd_nodes["10.0.0.2:28764"] = PD_Client_Obj( - node_id=2, - client_ip_port="10.0.0.2:28764", - mode="decode", - start_args={}, - ) - - manager.update_node_load_info({"client_ip_port": "10.0.0.1:28761", "total_token_usage_rate": 0.25}) - manager.update_node_load_info({"client_ip_port": "10.0.0.2:28764", "total_token_usage_rate": 0.75}) - monitor.gauge_set("lightllm_pd_master_stage_waiting_requests", 2, labels={"stage": "decode"}) - - metrics = generate_latest(monitor.registry).decode() - assert ( - 'lightllm_pd_node_token_usage_ratio{endpoint="10.0.0.1:28761",model_name="test-model",role="prefill"} 0.25' - in metrics - ) - assert ( - 'lightllm_pd_node_token_usage_ratio{endpoint="10.0.0.2:28764",model_name="test-model",role="decode"} 0.75' - in metrics - ) - assert 'lightllm_pd_master_stage_waiting_requests{model_name="test-model",stage="decode"} 2.0' in metrics - - -def test_metric_server_copies_rpyc_label_proxy_before_updating_gauge(): - args = SimpleNamespace( - metric_gateway=None, - job_name="lightllm", - grouping_key=None, - enable_monitor_auth=False, - model_name="test-model", - max_req_total_len=1024, - mtp_step=0, - push_interval=10, - ) - server = MetricServer(args) - - server.exposed_gauge_set( - "lightllm_pd_master_stage_waiting_requests", - 3, - RPyCDictProxyLike({"stage": "decode"}), - ) - - metrics = generate_latest(server.monitor.registry).decode() - assert 'lightllm_pd_master_stage_waiting_requests{model_name="test-model",stage="decode"} 3.0' in metrics - - -@pytest.mark.parametrize("stage", ["prefill", "decode"]) -def test_pd_master_stage_waiting_gauge_is_released_on_failure(stage): - async def run(): - manager = HttpServerManagerForPDMaster.__new__(HttpServerManagerForPDMaster) - manager.metric_client = MagicMock() - manager.pd_stage_waiting_request_counts = {"prefill": 0, "decode": 0} - manager._wait_for_event_or_disconnect = AsyncMock(side_effect=ServerBusyError()) - - with pytest.raises(ServerBusyError): - await manager._wait_for_pd_stage(AsyncMock(), AsyncMock(), 1, 123, stage) - - assert manager.pd_stage_waiting_request_counts[stage] == 0 - assert manager.metric_client.gauge_set.call_args_list == [ - call("lightllm_pd_master_stage_waiting_requests", 1, labels={"stage": stage}), - call("lightllm_pd_master_stage_waiting_requests", 0, labels={"stage": stage}), - ] - - asyncio.run(run())