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.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

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