Skip to content

[LightRAG][Step 4] Refactor Scheduler retrieval into an asynchronous pre-routing stage #166

Description

@rickisba

Parent Epic: #162
Depends on: #163, #164
Related infrastructure contracts: #140, #141, #142

Goal

Move knowledge retrieval out of the synchronous Request.build_request() path and execute it as an asynchronous Scheduler stage before pool-level routing.

The Scheduler must continue to select KDN and Proxy pools only. Proxy remains responsible for concrete Instance selection.

Target Request Flow

Scheduler HTTP handler
    |
    +-- parse request and user prompt
    +-- await KnowledgeRetriever.retrieve(...)
    +-- normalize/filter canonical KIDs
    +-- build Request with RetrievalResult
    +-- Scheduler strategy selects KDN + Proxy
    `-- forward to Proxy

Scope

  • Introduce an async pre-routing retrieval stage in both chat/completions and completions handlers without duplicating business logic.
  • Change Request.build_request() or its replacement seam so it consumes an already-computed retrieval result instead of performing external retrieval I/O.
  • Preserve a compatibility path for callers/tests that still construct requests without a retriever.
  • Add backend selection configuration with embedding as the default.
  • Add deterministic fallback: LightRAG failure -> embedding backend -> empty knowledge list while ordinary inference continues.
  • Normalize candidates to canonical KIDs, remove duplicates, validate current KDN availability, filter deleted/stale entries, and enforce candidate and token budgets.
  • Ensure Scheduler strategies receive retrieval metadata plus the existing knowledge list/length context, but do not receive LightRAG client objects or internal chunk IDs.
  • Add timeout, cancellation, and request-disconnect handling.
  • Keep Scheduler pool selection semantics unchanged except for consuming the normalized retrieval result.
  • Add focused tests for chat/completions, completions, RAG disabled, backend failure, fallback, empty candidates, stale KIDs, token-budget truncation, and request cancellation.

Configuration

Support configuration equivalent to:

  • SCHEDULER_RETRIEVAL_BACKEND=embedding|lightrag|shadow
  • LightRAG timeout and mode
  • maximum candidate count
  • maximum retrieval token budget
  • fallback enable/disable

The shadow behavior itself is completed in Step 5; this step may reserve the configuration value but must not change production results.

Invariants

  • Scheduler selects KDN and Proxy, never Instance.
  • Proxy continues to select Instance from its own InstancePool.
  • Retrieval failure never bypasses required routing validation.
  • LightRAG internal IDs never appear in Request.Service.Knowledge_List.
  • RAG-disabled requests do not call either retrieval backend.
  • No blocking HTTP client or asyncio.run() is allowed inside the serving event loop.

Non-Goals

  • Do not implement shadow metrics beyond minimal scaffolding.
  • Do not enable LightRAG as the default backend.
  • Do not modify Proxy scheduling or injection algorithms.
  • Do not implement Knowledge Demand Profile or joint optimization.

Acceptance Criteria

  • External retrieval I/O no longer occurs inside the synchronous request-model construction path.
  • Existing embedding retrieval remains the default and passes compatibility tests.
  • LightRAG can be selected explicitly and falls back according to configuration.
  • Requests continue without knowledge injection when both backends fail.
  • Chat and completions paths share the same retrieval orchestration logic.
  • Cancellation and timeouts do not leak tasks or HTTP connections.
  • CPU-only tests need no live LightRAG, GPU, Redis, vLLM, or KDN service.
  • Proxy Instance selection behavior is unchanged.

Validation

Run focused Scheduler retrieval tests, existing request-building tests, public import compatibility checks, and relevant routing tests. Include actual results in the PR and use Closes #166.

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or request

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions