From 96d663e9fef19d9cc21d154d938432ea7e054c7c Mon Sep 17 00:00:00 2001 From: Sneh Kansagara Date: Sun, 19 Jul 2026 11:16:06 +0530 Subject: [PATCH 1/4] fix(incremental-pipeline): use TEI embeddings instead of loading torch locally Signed-off-by: Sneh Kansagara --- .../pipelines/incremental-pipeline.py | 58 +++++++++++-------- 1 file changed, 35 insertions(+), 23 deletions(-) diff --git a/docs-agent-mcp/pipelines/incremental-pipeline.py b/docs-agent-mcp/pipelines/incremental-pipeline.py index cc0e97f..c552114 100644 --- a/docs-agent-mcp/pipelines/incremental-pipeline.py +++ b/docs-agent-mcp/pipelines/incremental-pipeline.py @@ -8,7 +8,7 @@ except ImportError: # pragma: no cover - optional at compile time k8s = None -from utils import DOCS_COLLECTION +from utils import DEFAULT_EMBEDDING_BATCH_SIZE, DOCS_COLLECTION @dsl.component( base_image="docker.io/library/python:3.9", @@ -171,12 +171,8 @@ def delete_old_vectors( @dsl.component( - base_image="docker.io/pytorch/pytorch:2.3.0-cuda12.1-cudnn8-runtime", - packages_to_install=[ - "sentence-transformers==3.3.1", - "transformers==4.44.2", - "langchain-text-splitters", - ], + base_image="python:3.11-slim", + packages_to_install=["requests", "langchain-text-splitters"], ) def chunk_and_embed_incremental( github_data: dsl.Input[dsl.Dataset], @@ -184,19 +180,18 @@ def chunk_and_embed_incremental( base_url: str, chunk_size: int, chunk_overlap: int, + embeddings_service_url: str, + embedding_batch_size: int, embedded_data: dsl.Output[dsl.Dataset] ): import json import os import re - import torch - from sentence_transformers import SentenceTransformer + import requests from langchain_text_splitters import RecursiveCharacterTextSplitter - device = 'cuda' if torch.cuda.is_available() else 'cpu' - model = SentenceTransformer('sentence-transformers/all-mpnet-base-v2', device=device) - print(f"Model loaded on {device}") - EMBED_BATCH_SIZE = 32 + print(f"Using embeddings service: {embeddings_service_url}") + embedding_batch_size = max(1, int(embedding_batch_size)) records = [] @@ -259,13 +254,7 @@ def chunk_and_embed_incremental( print(f"File: {file_data['path']} -> {len(chunks)} chunks (avg: {sum(len(c) for c in chunks)/len(chunks):.0f} chars)") - # Create embeddings in batches to avoid per-chunk model overhead. - embeddings = model.encode( - chunks, - batch_size=EMBED_BATCH_SIZE, - show_progress_bar=False, - ) - for chunk_idx, (chunk, embedding) in enumerate(zip(chunks, embeddings)): + for chunk_idx, chunk in enumerate(chunks): records.append({ 'file_unique_id': file_unique_id, 'repo_name': repo_name, @@ -274,10 +263,27 @@ def chunk_and_embed_incremental( 'citation_url': citation_url[:1024], 'chunk_index': chunk_idx, 'content_text': chunk[:2000], - 'embedding': embedding.tolist() }) - print(f"Created {len(records)} total chunks for incremental update") + print(f"Created {len(records)} total chunks for incremental update; requesting embeddings from TEI service...") + + # TEI all-mpnet-base-v2 rejects any input >=384 tokens (~1000 chars). + max_tei_chars = 1000 + for i in range(0, len(records), embedding_batch_size): + batch = records[i:i + embedding_batch_size] + texts = [r["content_text"][:max_tei_chars] for r in batch] + response = requests.post( + embeddings_service_url, + json={"inputs": texts}, + headers={"Content-Type": "application/json"}, + timeout=120, + ) + response.raise_for_status() + vectors = response.json() + for idx, vector in enumerate(vectors): + batch[idx]["embedding"] = vector + + print(f"Embedded {len(records)} chunks") with open(embedded_data.path, 'w', encoding='utf-8') as f: for record in records: @@ -404,6 +410,10 @@ def github_rag_incremental_pipeline( base_url: str = "https://www.kubeflow.org/docs", chunk_size: int = 1200, chunk_overlap: int = 100, + embeddings_service_url: str = ( + "http://embeddings-service-predictor.ml-infra.svc.cluster.local/embed" + ), + embedding_batch_size: int = DEFAULT_EMBEDDING_BATCH_SIZE, milvus_host: str = "milvus-milvus.ml-infra.svc.cluster.local", milvus_port: str = "19530", collection_name: str = DOCS_COLLECTION @@ -441,7 +451,9 @@ def github_rag_incremental_pipeline( repo_name=repo_name, base_url=base_url, chunk_size=chunk_size, - chunk_overlap=chunk_overlap + chunk_overlap=chunk_overlap, + embeddings_service_url=embeddings_service_url, + embedding_batch_size=embedding_batch_size, ) # Step 4: Store new vectors in Milvus (after deletion is complete) From 9a7d472107b6d9d5b9eb747cdd005cda1a8d298b Mon Sep 17 00:00:00 2001 From: Sneh Kansagara Date: Sun, 19 Jul 2026 11:16:15 +0530 Subject: [PATCH 2/4] fix(pipelines): drop torch from Dockerfile.pipeline, not needed anymore Signed-off-by: Sneh Kansagara --- docs-agent-mcp/pipelines/Dockerfile.pipeline | 13 +++---------- 1 file changed, 3 insertions(+), 10 deletions(-) diff --git a/docs-agent-mcp/pipelines/Dockerfile.pipeline b/docs-agent-mcp/pipelines/Dockerfile.pipeline index c23b11a..ba8008e 100644 --- a/docs-agent-mcp/pipelines/Dockerfile.pipeline +++ b/docs-agent-mcp/pipelines/Dockerfile.pipeline @@ -1,18 +1,11 @@ FROM python:3.11-slim -# Install python packages without caching to keep image size small +# Install python packages without caching to keep image size small. +# Embeddings are generated by the in-cluster TEI service now (see utils.py), +# so no ML model needs to be loaded or baked into this image anymore. RUN pip install --no-cache-dir \ kfp>=2.0.0 \ pymilvus>=2.4.0 \ - sentence-transformers \ langchain-text-splitters \ beautifulsoup4 \ requests - -# Pre-download the sentence-transformers model so it's baked into the image. -# This prevents downloading it on every pipeline run, saving time, ephemeral storage, -# and avoiding HuggingFace Hub rate limits. -RUN python -c "from sentence_transformers import SentenceTransformer; SentenceTransformer('sentence-transformers/all-mpnet-base-v2')" - -# Set huggingface cache location explicitly -ENV HF_HOME=/root/.cache/huggingface From 793c7127a07127d2036a39c88c010be71adb82f7 Mon Sep 17 00:00:00 2001 From: Sneh Kansagara Date: Sun, 19 Jul 2026 11:16:18 +0530 Subject: [PATCH 3/4] docs: fix stale Dockerfile.pipeline description Signed-off-by: Sneh Kansagara --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 84214d9..577e992 100644 --- a/README.md +++ b/README.md @@ -739,7 +739,7 @@ We use Terraform for declarative, reproducible cluster infrastructure on OKE. ### Pipeline Optimizations (`docs-agent-mcp/pipelines/`) The ingestion pipeline was rewritten to maximize efficiency and avoid Kubernetes ephemeral storage eviction: * **Feast Removal**: The pipeline now writes embeddings directly to Milvus using `pymilvus`, dramatically lowering complexity. -* **Custom Base Image (`Dockerfile.pipeline`)**: We bake the massive PyTorch library and the Hugging Face `all-mpnet-base-v2` model directly into a custom Docker image. This reduces runtime disk usage from 5.5GB to zero, fixing OKE pod eviction errors, and preventing Hugging Face API rate limits. +* **TEI Embeddings**: Chunking components call the in-cluster TEI embeddings service over HTTP instead of loading a model in-process, so `Dockerfile.pipeline` stays a plain `python:3.11-slim` image with no PyTorch/model weights baked in — avoiding both the OKE pod eviction and Hugging Face rate-limit issues without the multi-GB image. ### GitHub Actions CI/CD (`.github/workflows/`) | Workflow | When it runs | Purpose | From 374803dd4205ad2f2794e5c1d414fdac1e572a97 Mon Sep 17 00:00:00 2001 From: Sneh Kansagara Date: Sun, 23 Aug 2026 09:39:25 +0530 Subject: [PATCH 4/4] fix(pipelines): pin package versions in incremental-pipeline components Match versions already pinned in docs-agent-mcp/pipelines/requirements.txt. Signed-off-by: Sneh Kansagara --- docs-agent-mcp/pipelines/incremental-pipeline.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/docs-agent-mcp/pipelines/incremental-pipeline.py b/docs-agent-mcp/pipelines/incremental-pipeline.py index c552114..56eafde 100644 --- a/docs-agent-mcp/pipelines/incremental-pipeline.py +++ b/docs-agent-mcp/pipelines/incremental-pipeline.py @@ -12,7 +12,7 @@ @dsl.component( base_image="docker.io/library/python:3.9", - packages_to_install=["requests", "beautifulsoup4"] + packages_to_install=["requests==2.34.2", "beautifulsoup4==4.15.0"] ) def download_specific_files( repo_owner: str, @@ -94,7 +94,7 @@ def resolve_github_token(token): @dsl.component( base_image="docker.io/library/python:3.9", - packages_to_install=["pymilvus"] + packages_to_install=["pymilvus==2.6.14"] ) def delete_old_vectors( file_paths: str, # JSON string of file paths list @@ -172,7 +172,7 @@ def delete_old_vectors( @dsl.component( base_image="python:3.11-slim", - packages_to_install=["requests", "langchain-text-splitters"], + packages_to_install=["requests==2.34.2", "langchain-text-splitters==1.1.2"], ) def chunk_and_embed_incremental( github_data: dsl.Input[dsl.Dataset], @@ -292,7 +292,7 @@ def chunk_and_embed_incremental( @dsl.component( base_image="docker.io/library/python:3.9", - packages_to_install=["pymilvus", "numpy"] + packages_to_install=["pymilvus==2.6.14", "numpy==2.2.6"] ) def store_milvus_incremental( embedded_data: dsl.Input[dsl.Dataset],