diff --git a/proxy/README.md b/proxy/README.md index 179440d..447e5b6 100644 --- a/proxy/README.md +++ b/proxy/README.md @@ -151,7 +151,7 @@ POST /v1/instance/unregister GET /v1/instance/list?include_dead=true ``` -Instances register static information such as `instance_id`, `host`, `port`, `endpoints`, `tags`, `weight`, and `meta`. Heartbeats refresh `last_seen_at` and can optionally report lightweight load fields such as `inflight`, `qps_1m`, and `gpu_util`. +Instances register static information such as `instance_id`, `host`, `port`, `endpoints`, `tags`, `weight`, and `meta`. Heartbeats refresh `last_seen_at` and can optionally report lightweight non-lifecycle load fields such as `qps_1m` and `gpu_util`; Proxy request lifecycle hooks maintain `inflight`. ### Instance resource snapshots @@ -160,8 +160,11 @@ After PR #87, demo Instances can report host resource snapshots to the Proxy con ```text POST /v1/instance/resource_snapshot GET /debug/instance_resources +GET /debug/instance_loads ``` +Use `/debug/instance_loads` to inspect Proxy-maintained per-Instance `inflight` counters, `qps_1m` when present, and queue-depth hints when the queue snapshot provider is available. + The reporting path is: ```text diff --git a/proxy/proxy.py b/proxy/proxy.py index f19e369..02451c1 100644 --- a/proxy/proxy.py +++ b/proxy/proxy.py @@ -393,7 +393,12 @@ def _sse_meta_event(task: ProxyTask) -> bytes: ).encode("utf-8") -async def _wrap_chat_stream_with_meta(task: ProxyTask, queue_mgr: QueueManager) -> AsyncGenerator[bytes, None]: +async def _wrap_chat_stream_with_meta( + task: ProxyTask, + queue_mgr: QueueManager, + instance_pool: InstancePool, + instance_id: str, +) -> AsyncGenerator[bytes, None]: """ Forward downstream chat SSE, but delay [DONE] and insert one cacheroute_meta event first. """ @@ -423,11 +428,14 @@ async def _wrap_chat_stream_with_meta(task: ProxyTask, queue_mgr: QueueManager) task.trace["stream_exception_ms"] = int(time.time() * 1000) logger.exception("[Proxy] stream wrapper failed rid=%s", task.request_id) finally: - if pending: - yield pending - pending = b"" - yield _sse_meta_event(task) - yield b"data: [DONE]\n\n" + try: + if pending: + yield pending + pending = b"" + yield _sse_meta_event(task) + yield b"data: [DONE]\n\n" + finally: + instance_pool.end_request(instance_id) def select_instance(app: FastAPI, req_obj: SchedulerRequest): @@ -593,8 +601,15 @@ async def proxy_chat_completions(request: FastAPIRequest): # ==================================== # Send into the queue: enqueue -> manager -> forward # ==================================== + instance_pool: InstancePool = proxy.state.instance_pool # type: ignore + request_counted = instance_pool.begin_request(chosen.instance_id) + if not request_counted: + return JSONResponse( + status_code=503, + content={"error": "no_instance", "detail": "selected instance is no longer registered"}, + ) try: - # 1) Wrap the task (note: chosen comes from RR and has instance_id/host/port fields) + # 1) Wrap the task (note: chosen comes from strategy and has instance_id/host/port fields) task = ProxyTask( request_id=getattr(req_obj, "Request_ID", None), req_obj=req_obj, @@ -611,10 +626,12 @@ async def proxy_chat_completions(request: FastAPIRequest): await queue_mgr.enqueue_prepare(task) - stream_gen = _wrap_chat_stream_with_meta(task, queue_mgr) + stream_gen = _wrap_chat_stream_with_meta(task, queue_mgr, instance_pool, chosen.instance_id) return StreamingResponse(stream_gen, media_type="text/event-stream") except Exception as e: + if request_counted: + instance_pool.end_request(chosen.instance_id) logger.exception("[Proxy] failed to call Worker(chat)") return JSONResponse( status_code=502, @@ -762,6 +779,13 @@ async def proxy_completions(request: FastAPIRequest): # ================================== # Send into the queue: enqueue -> drain -> forward # ================================== + instance_pool: InstancePool = proxy.state.instance_pool # type: ignore + request_counted = instance_pool.begin_request(chosen.instance_id) + if not request_counted: + return JSONResponse( + status_code=503, + content={"error": "no_instance", "detail": "selected instance is no longer registered"}, + ) try: task = ProxyTask( request_id=getattr(req_obj, "Request_ID", None), @@ -780,9 +804,14 @@ async def proxy_completions(request: FastAPIRequest): await queue_mgr.enqueue_prepare(task) content_bytes = b"" - async for chunk in queue_mgr.iter_response(task): - if chunk: - content_bytes += chunk + try: + async for chunk in queue_mgr.iter_response(task): + if chunk: + content_bytes += chunk + finally: + if request_counted: + instance_pool.end_request(chosen.instance_id) + request_counted = False # completions is non-streaming: the worker should return a one-shot JSON response if not content_bytes: @@ -806,6 +835,8 @@ async def proxy_completions(request: FastAPIRequest): ) except Exception as e: + if request_counted: + instance_pool.end_request(chosen.instance_id) logger.exception("[Proxy] failed to call Worker(completions)") return JSONResponse( status_code=502, diff --git a/proxy/resource/instance_pool.py b/proxy/resource/instance_pool.py index 25b8921..38fea8b 100644 --- a/proxy/resource/instance_pool.py +++ b/proxy/resource/instance_pool.py @@ -90,6 +90,8 @@ def upsert( it.weight = float(weight) if meta: it.meta.update(meta) + if it.load.inflight is None: + it.load.inflight = 0 it.last_seen_at = now return it @@ -101,6 +103,7 @@ def upsert( tags=tags or [], weight=float(weight), meta=meta or {}, + load=InstanceLoad(inflight=0), registered_at=now, last_seen_at=now, ) @@ -120,9 +123,8 @@ def heartbeat( if not it: return False it.last_seen_at = now - # Only update when provided to avoid overwriting with 0/None - if inflight is not None: - it.load.inflight = int(inflight) + # Inflight is Proxy-maintained via begin_request/end_request. + # Keep heartbeat load updates for non-lifecycle signals only. if qps_1m is not None: it.load.qps_1m = float(qps_1m) if gpu_util is not None: @@ -145,6 +147,80 @@ def report_resource_snapshot( return True + def begin_request(self, instance_id: str) -> bool: + """Increment the Proxy-maintained inflight counter for an Instance.""" + with self._lock: + it = self._items.get(instance_id) + if not it: + return False + current = int(it.load.inflight or 0) + it.load.inflight = current + 1 + return True + + def end_request(self, instance_id: str) -> bool: + """Decrement the Proxy-maintained inflight counter without going below zero.""" + with self._lock: + it = self._items.get(instance_id) + if not it: + return False + current = int(it.load.inflight or 0) + it.load.inflight = max(0, current - 1) + return True + + def snapshot_instance_loads( + self, + queue_depths: Optional[Dict[str, Any]] = None, + include_dead: bool = True, + ) -> Dict[str, Any]: + """Return per-Instance local load counters and optional queue-depth hints.""" + now = int(time.time()) + per_instance_queue = {} + if queue_depths: + raw = queue_depths.get("per_instance", {}) + if isinstance(raw, dict): + per_instance_queue = raw + + with self._lock: + items = list(self._items.values()) + + instances: List[Dict[str, Any]] = [] + for it in items: + is_alive = (now - int(it.last_seen_at)) <= self._ttl_s + if not include_dead and not is_alive: + continue + queue_item = per_instance_queue.get(it.instance_id, {}) + inflight = it.load.inflight + qps_1m = it.load.qps_1m + instances.append({ + "instance_id": it.instance_id, + "is_alive": is_alive, + "inflight": inflight, + "qps_1m": qps_1m, + "prepare_queue_depth": queue_item.get("prepare_queue_depth"), + "ready_queue_depth": queue_item.get("ready_queue_depth"), + "active_prepare": queue_item.get("active_prepare"), + "active_ready": queue_item.get("active_ready"), + "least_load_score": { + "primary": "inflight" if inflight is not None else None, + "primary_value": inflight, + "primary_source": "proxy_lifecycle_counter" if inflight is not None else "unavailable", + "secondary": "qps_1m" if qps_1m is not None else None, + "secondary_value": qps_1m, + "secondary_source": "instance_heartbeat" if qps_1m is not None else "unavailable", + }, + }) + + return { + "ttl_s": self._ttl_s, + "generated_at": time.time(), + "metric_source": { + "inflight": "proxy_lifecycle_counter", + "qps_1m": "instance_heartbeat", + "queue_depth": "proxy_queue_manager" if queue_depths is not None else "unavailable", + }, + "instances": instances, + } + def build_pool_resource_snapshot( self, proxy_id: str, @@ -216,7 +292,7 @@ def build_pool_resource_snapshot( metric_source = { "instances": "proxy_instance_pool_ttl", "resource": "instance_resource_snapshot" if reporting else "unavailable", - "inflight_total": "instance_heartbeat" if inflight_values else "unavailable", + "inflight_total": "proxy_lifecycle_counter" if inflight_values else "unavailable", "qps_1m_total": "instance_heartbeat" if qps_values else "unavailable", "load_ratio": "derived_from_inflight_capacity" if load_ratio is not None else "unavailable", "capacity": "proxy_config", diff --git a/proxy/resource/p_control_plane.py b/proxy/resource/p_control_plane.py index 2610914..d2c99ba 100644 --- a/proxy/resource/p_control_plane.py +++ b/proxy/resource/p_control_plane.py @@ -339,6 +339,22 @@ async def debug_instance_resources(include_dead: bool = True) -> Dict[str, Any]: }) return {"instances": resources} + +@_control_plane.get("/debug/instance_loads") +async def debug_instance_loads(include_dead: bool = True) -> Dict[str, Any]: + pool = get_pool() + queue_depths = None + if _queue_snapshot_provider is not None: + try: + queue_depths = _queue_snapshot_provider() or {} + except Exception: + logger.warning( + "[ProxyCP] queue snapshot provider failed; instance load queue metrics unavailable", + exc_info=True, + ) + snapshot = pool.snapshot_instance_loads(queue_depths=queue_depths, include_dead=include_dead) + return {"ok": True, **snapshot} + @_control_plane.post("/v1/topology/report") async def report_topology(req: TopologyReportReq) -> Dict[str, Any]: merged = 0