Wire Proxy-maintained per-Instance load counters into least_load

[zhy1658858023]

Background

PR #116 added the first Proxy-side resource-aware Instance selection strategy: least_load.

At the moment, least_load reads InstanceInfo.load.inflight and InstanceInfo.load.qps_1m if they are available, and falls back to round-robin when all load metrics are unknown. This is correct as a strategy entry point, but the next step is to make those load fields reflect real Proxy-side runtime state.

Resource data is currently mainly used for observability. This issue should wire real Proxy-maintained per-Instance load counters into the data plane so least_load can make meaningful decisions.

The current route flow is roughly:

Proxy request handler
  -> recover SchedulerRequest
  -> select_instance(...)
  -> create ProxyTask
  -> queue_mgr.enqueue_prepare(task)
  -> queue_mgr.iter_response(task)

select_instance(...) currently selects from InstancePool.list(include_dead=False) and calls the configured strategy. This issue should keep that strategy interface intact while ensuring the selected Instance has runtime load counters updated around request execution.

Goal

Add Proxy-maintained per-Instance load counters and expose them through InstanceInfo.load so least_load can select the least loaded live Instance using real local runtime state.

The first implementation should focus on:

per-instance inflight
optional per-instance queue depth / active queue counters if safely available
simple debug visibility for strategy inputs and scores

Branch hygiene

Follow doc/codex_workflow.md if present.

Start from latest origin/main and use a fresh branch for this issue. Do not continue from previous Codex branches or include previous unmerged PR changes.

Suggested branch name:

codex/issue-117-proxy-instance-load-counters

Scope

Expected files may include:

proxy/proxy.py
proxy/resource/instance_pool.py
proxy/resource/p_control_plane.py
proxy/queue/manager.py
proxy/queue/instance_queues.py
proxy/strategy/least_load.py
proxy/README.md

Only modify files that are necessary for Proxy-local load counters, least_load integration, and debug documentation.

Required behavior

1. Proxy-maintained per-Instance inflight

When Proxy has selected an Instance and the request is accepted into the Proxy data-plane path, increment that Instance's runtime inflight counter:

inflight[instance_id] += 1

When the request finishes, fails, raises, or the client disconnects, decrement it:

inflight[instance_id] -= 1

Important requirements:

• Use try/finally or an equivalent lifecycle guard so inflight is decremented on success, error, and cancellation.
• Never allow the counter to go below 0.
• The counter must be thread-safe / async-safe enough for concurrent requests.
• The counter should update the same load view that least_load reads, preferably InstanceInfo.load.inflight.
• Do not rely on Instance heartbeat to supply inflight for this first implementation; Proxy should maintain its own local counter.

2. Counter ownership and data model

Prefer adding small methods to InstancePool, for example:

begin_request(instance_id: str) -> bool
end_request(instance_id: str) -> bool
snapshot_instance_loads(...) -> Dict[str, Any]

or equivalent names.

The implementation should avoid direct mutation of InstanceInfo.load from many unrelated call sites. Keep the counter logic centralized.

begin_request(...) should:

• find the Instance by id
• increment load.inflight
• optionally refresh or preserve liveness timestamp based on current design
• return False if the Instance disappeared

end_request(...) should:

• find the Instance by id
• decrement load.inflight with floor 0
• return False if the Instance disappeared

3. Request lifecycle integration

Integrate the counters around both Proxy request paths:

POST /v1/chat/completions
POST /v1/completions

For chat streaming, the selected Instance remains active until the streaming generator is exhausted or cancelled. The decrement must happen when streaming finishes, not immediately after returning StreamingResponse.

For non-streaming completions, decrement after queue_mgr.iter_response(task) completes or errors.

Suggested approach:

• For non-streaming path, wrap the enqueue_prepare + response collection block in try/finally.
• For streaming path, wrap the async generator that drains queue_mgr.iter_response(task) so end_request(instance_id) is called from the generator's finally block.
• If increment succeeds but enqueue_prepare fails, still decrement.
• If increment fails because the selected Instance disappeared, return a safe 503 or reselect only if the existing routing flow already supports that cleanly.

4. Queue depth integration

least_load already works with InstanceInfo.load.inflight / qps_1m.

For this issue, queue depth integration is optional but preferred if it can be done without invasive request-flow changes.

Possible minimal path:

• Keep least_load primarily based on Proxy-maintained load.inflight.
• Add or expose a Proxy-local per-instance load debug endpoint that includes queue depth from QueueManager.queue_depth_snapshot() / PerInstanceQueueMap.snapshot_depths() if available.
• Leave a TODO for using queue depth inside the strategy score if the strategy call site does not naturally receive per-instance queue hints yet.

If queue hints are added to least_load, they must remain optional and must not break the existing strategy interface.

5. Debug visibility

Add lightweight debug visibility so runtime tests can verify the counters and future strategy decisions.

Suggested endpoint:

GET /debug/instance_loads

Possible response:

{
  "ok": true,
  "strategy": "least_load",
  "instances": [
    {
      "instance_id": "instance_127.0.0.1:9001",
      "is_alive": true,
      "inflight": 2,
      "qps_1m": null,
      "prepare_queue_depth": 1,
      "ready_queue_depth": 0,
      "active_prepare": 1,
      "active_ready": 0,
      "score": 2,
      "score_source": "proxy_runtime_inflight"
    }
  ]
}

The exact shape can differ, but it should clearly expose:

• per-instance inflight
• qps_1m if present
• queue depths if available
• liveness state
• enough score/provenance information to understand why least_load would choose an Instance

6. qps_1m handling

This issue does not need to implement a full rolling QPS counter.

Acceptable options:

• leave qps_1m unchanged / optional
• document that qps remains unavailable until a later rolling counter issue
• if implementing qps, keep it minimal and clearly separated from inflight

Do not block this issue on QPS.

7. Preserve least_load semantics

least_load must continue to obey the existing semantics from PR #116:

• unknown load is not treated as real zero
• if all load metrics are unknown, fallback to round-robin
• ties are broken by round-robin among tied candidates
• empty input raises RuntimeError("no instances")

Once Proxy-maintained inflight is wired, selected Instances should normally have known inflight values.

Non-goals

Do not implement yet:

• Scheduler resource-aware routing
• Scheduler Proxy selection based on pool_resource
• KV-aware Instance routing
• KDN-aware routing
• KVCache locality scoring
• changes to IWS decision logic
• changes to request forwarding semantics beyond lifecycle-safe load counter hooks
• Resource Agent sampling changes
• pool_resource schema changes
• full rolling QPS calculation unless it is very small and isolated

This issue is only about making Proxy-side Instance load counters real and observable for least_load.

Suggested implementation direction

A practical first implementation:

1. Add thread-safe per-instance inflight update methods in InstancePool.
2. Initialize InstanceInfo.load.inflight to 0 only when the Proxy has begun maintaining that counter for the Instance.
3. In proxy_chat_completions, call begin_request(chosen.instance_id) before enqueueing the ProxyTask.
4. In the streaming wrapper, call end_request(chosen.instance_id) in finally after the stream ends or is cancelled.
5. In proxy_completions, call begin_request(...) before enqueueing and end_request(...) in finally after response collection finishes.
6. Add /debug/instance_loads in Proxy control plane.
7. Optionally include queue depth data from QueueManager.queue_depth_snapshot() in debug output.
8. Update docs with a short note that least_load now uses Proxy-maintained runtime inflight.

Validation plan

Compile check

PYTHONDONTWRITEBYTECODE=1 python3 -m py_compile \
  proxy/resource/instance_pool.py \
  proxy/resource/p_control_plane.py \
  proxy/proxy.py \
  proxy/queue/*.py \
  proxy/strategy/least_load.py \
  proxy/strategy/factory.py \
  test/demo_proxy.py

Unit-style smoke test

Use fake InstanceInfo objects or a real InstancePool instance to verify:

• begin_request(instance_id) increments inflight
• end_request(instance_id) decrements inflight
• repeated end_request does not make inflight negative
• unknown instance returns false / no crash
• least_load chooses lower inflight
• equal inflight ties are round-robin

Runtime smoke test

Start Scheduler, Proxy, and two Instances.

Start Proxy with:

python3 test/demo_proxy.py --strategy least_load --injection-strategy iws

Inspect before load:

curl -sS http://127.0.0.1:8002/debug/instance_loads | python3 -m json.tool

Send several concurrent requests. During active requests, verify:

inflight increases on selected Instances
least_load tends to avoid Instances with higher inflight
inflight returns to 0 after requests complete

Also verify existing round-robin still starts:

python3 test/demo_proxy.py --strategy round_robin

Failure/cancellation sanity check

Trigger or simulate a failed request path and verify:

inflight does not leak
/debug/instance_loads eventually returns inflight=0

Acceptance criteria

• Proxy maintains real per-Instance inflight counters locally.
• least_load sees known InstanceInfo.load.inflight values during request execution.
• Inflight increments when a selected Instance receives work and decrements when work completes, errors, or is cancelled.
• Inflight never becomes negative.
• Chat streaming decrements only when the stream finishes or is cancelled.
• Non-streaming completions decrement after response collection finishes or errors.
• A debug endpoint exposes per-instance load state and optional queue depths.
• round_robin remains default and unchanged.
• least_load continues to fallback to round-robin when all loads are unknown.
• No Scheduler routing, KDN, KVCache, IWS decision logic, Resource Agent, or pool_resource schema changes are introduced.