RAG pipeline context injection failure after MLflow model update

Problem: RAG pipeline context injection failure after MLflow model update

During a blue‑green deployment, traffic was shifted from the green version of the rag-service endpoint to the newly promoted blue version. After the switch, the Retrieval‑Augmented Generation (RAG) pipeline began sending prompts to the LLM that contained an empty or malformed context variable. Typical symptoms observed in production logs include:

  • KeyError: 'context' from the LangChain prompt renderer.
  • ValueError: retrieved_documents is empty raised by the Retriever.
  • HTTP 500 responses from the MLflow serving endpoint: {"error":"No context provided"}.

The issue manifested intermittently—requests routed to the green deployment succeeded, while those hitting the blue deployment failed.

Root Cause Analysis

The failure originates from a schema drift between the two model versions and the way MLflow propagates environment variables and request payload schemas during a blue‑green rollout.

  1. Model signature mismatch: The green model was exported with a custom inference signature that included a context field (see MLflow Model Serving documentation). The blue model, built from a newer code branch, omitted this field in its signature.json. When the serving container loads the blue artifact, the inference engine expects only the question key.
  2. Traffic‑shifting semantics: According to the MLflow Model Registry API guide, a traffic split is performed by updating the MLFLOW_MODEL_PATH environment variable for the serving pod. During the promotion, the variable was overwritten before the new artifact was fully staged, causing the container to load a stale model version that lacked the context field.
  3. Vector store deserialization: The blue version introduced a change in the vector store path (e.g., /opt/mlflow/artifacts/v2/index.faiss/opt/mlflow/artifacts/v3/index.faiss). Because the path was hard‑coded in the model’s init() method, the blue container continued to reference the old index, returning an empty list of documents. This matches the community observation in the GitHub issue “RAG context variable empty after model version bump” (LangChain #3421).

Combined, these factors cause the retrieval step to return [], leading the prompt template to render Context: [] or raise a KeyError when the variable is missing.

Investigation and Debugging Steps

Below is a reproducible debugging workflow that was used to isolate the problem.

1. Verify the active model version and endpoint schema

curl -s -X GET http://mlflow-server/api/2.0/mlflow/registered-models/get/latest-versions?name=rag-service | jq .

Expected output (blue version):

{
  "model_version": "3",
  "current_stage": "Production",
  "status": "READY",
  "signature": {
    "inputs": [{"name":"question","type":"string"}],
    "outputs": [{"name":"answer","type":"string"}]
  }
}

If the signature.inputs array does not contain "context", the model version is the culprit.

2. Inspect the container environment

kubectl exec -it $(kubectl get pod -l app=rag-service -o jsonpath="{.items[0].metadata.name}") -- env | grep MLFLOW_MODEL_PATH

Check that the path points to the correct artifact directory (e.g., /opt/mlflow/artifacts/v3/).

3. Capture the retrieval payload

curl -s -X POST http://mlflow-server/invocations \
  -H "Content-Type: application/json" \
  -d '{"question":"What is the capital of France?","context":[]}' \
  -o response.json
cat response.json

Log snippet from the blue pod when the payload is empty:

2026-08-14 10:22:31,112 INFO  mlflow.models: Retrieval returned 0 documents for request id=abc123
2026-08-14 10:22:31,115 ERROR langchain.prompts: Failed to render prompt: missing variable 'context'

4. Compare vector store index files

ls -l /opt/mlflow/artifacts/v2/
ls -l /opt/mlflow/artifacts/v3/

Missing or zero‑byte index.faiss in the new version confirms the deserialization issue.

5. Review model artifact metadata

cat /opt/mlflow/artifacts/v3/MLmodel

Look for the signature block; absence of the context field is a red flag.

Resolution

The fix required three coordinated changes:

1. Align model signatures

Re‑export the blue model with the same inference signature as the green version. In the training script:

# Before (blue version)
mlflow.pyfunc.log_model(
    "rag-model",
    python_model=MyRagModel(),
    artifacts={"vector_store": "/opt/mlflow/artifacts/v3/index.faiss"}
)

# After – include custom input schema
import json
signature = {
    "inputs": [
        {"name": "question", "type": "string"},
        {"name": "context", "type": "list[string]"}
    ],
    "outputs": [{"name": "answer", "type": "string"}]
}
mlflow.pyfunc.log_model(
    "rag-model",
    python_model=MyRagModel(),
    artifacts={"vector_store": "/opt/mlflow/artifacts/v3/index.faiss"},
    signature=signature
)

2. Preserve the vector store path during promotion

Update the model’s init() method to read the index location from an environment variable rather than a hard‑coded path.

# Before
self.index_path = "/opt/mlflow/artifacts/v2/index.faiss"

# After
self.index_path = os.getenv("VECTOR_STORE_PATH", "/opt/mlflow/artifacts/default/index.faiss")

Set the variable in the deployment manifest:

env:
  - name: VECTOR_STORE_PATH
    value: "/opt/mlflow/artifacts/v3/index.faiss"

3. Adjust the blue‑green traffic shift sequence

Use the MLflow Model Registry API to mark the new version Staging, verify the artifact layout, then promote to Production only after the MLFLOW_MODEL_PATH variable is correctly set.

# Promote after validation
curl -X POST http://mlflow-server/api/2.0/mlflow/model-versions/transition-stage \
  -d '{"name":"rag-service","version":"3","stage":"Production","archive_existing_versions":true}'

Validation

After applying the fixes, run the following checks:

  1. Signature verification – the /invocations endpoint should now accept a payload containing both question and context keys without error.
  2. Vector store retrieval – a test request should return a non‑empty context list.
  3. End‑to‑end prompt rendering – the LLM response must include the retrieved context.

Sample successful request:

curl -s -X POST http://mlflow-server/invocations \
  -H "Content-Type: application/json" \
  -d '{"question":"What is the capital of France?","context":["Paris is the capital of France."]}'

Expected fragment in the LLM prompt:

Context: Paris is the capital of France.
Question: What is the capital of France?
Answer:

Metrics from Prometheus should show rag_retrieval_success_total increasing and rag_prompt_render_error_total dropping to zero.

Prevention and Best Practices

  • Version‑locked signatures: Store the expected input schema in a separate JSON file and enforce it during CI/CD with a lint step.
  • Environment variable immutability: Use Kubernetes ConfigMap immutable entries for MLFLOW_MODEL_PATH and VECTOR_STORE_PATH to avoid accidental overwrites during rollout.
  • Canary validation: Before full traffic shift, route a small percentage of requests to the new version and assert that retrieved_documents count > 0.
  • Artifact integrity checks: Add a post‑deployment hook that runs faiss_index_reader --verify /opt/mlflow/artifacts/${VERSION}/index.faiss and fails the deployment if the index is empty.
  • Monitoring: Alert on spikes in KeyError: 'context' or ValueError: retrieved_documents is empty within the first 5 minutes after a model promotion.

Related Topic Hub: Model Serving Troubleshooting Hub

FAQ

  1. Why does the context variable disappear only after the blue deployment?
    Because the blue model was exported without the custom context input in its signature, causing MLflow to drop the field during request deserialization.
  2. Can I keep the old vector store index and still use the new model?
    Yes, if the new model’s init() method reads the index path from an environment variable or configuration file, you can point both versions to the same index.
  3. How do I detect schema drift before a rollout?
    Add a CI step that compares the signature.json of the candidate artifact against the currently deployed version and fails if they differ in required fields.
  4. What should I monitor to catch this issue early?
    Track the Prometheus counters rag_prompt_render_error_total and rag_retrieval_success_total. An alert on a sudden rise of render errors within 2 minutes of a traffic shift is a strong indicator.
  5. Is the problem related to LangChain’s Retriever implementation?
    Only indirectly. LangChain correctly propagates the retrieved documents, but it relies on the serving endpoint to return a populated context field. A mismatched model signature or missing vector store index prevents LangChain from receiving any documents.