Skip to content

fix(pipelines): stop shipping a CUDA image for CPU-only embedding inference - #224

Open
Sharkyii wants to merge 3 commits into
kubeflow:mainfrom
Sharkyii:fix/pipeline-cuda-bloat
Open

fix(pipelines): stop shipping a CUDA image for CPU-only embedding inference#224
Sharkyii wants to merge 3 commits into
kubeflow:mainfrom
Sharkyii:fix/pipeline-cuda-bloat

Conversation

@Sharkyii

Copy link
Copy Markdown
Contributor

Fixes the Docker image bloat issue reported against incremental-pipeline.py and Dockerfile.pipeline (both were still installing torch/sentence-transformers the plain way, which pulls the full CUDA build even though nothing in this cluster has a GPU).

What changed

  • incremental-pipeline.py: chunk_and_embed_incremental no longer loads sentence-transformers locally on a pytorch/pytorch:...-cuda12.1-cudnn8-runtime base image. It now calls the shared in-cluster TEI embeddings service over HTTP, same as kubeflow-pipeline.py / code-pipeline.py / issues-pipeline.py already do. Base image is now plain python:3.11-slim.
  • Dockerfile.pipeline: dropped sentence-transformers/torch and the model pre-download step entirely nothing in pipelines/ needs local model inference anymore.
  • README: updated the line describing Dockerfile.pipeline's old "bake PyTorch into the image" design, which no longer reflects reality.

Tested it directly torch --index-url https://download.pytorch.org/whl/cpu on one requirements line is silently ignored by pip, it still resolves the full ~500MB CUDA build. The flag only works as its own standalone line, and then it applies to every package below it, breaking anything not hosted on PyTorch's index. Migrating off local inference entirely sidesteps the problem rather than working around it.

Testing

  • python incremental-pipeline.py compiles the pipeline successfully (KFP compiler validates the full component graph).
  • ruff check docs-agent-mcp/mcp-server tests docs-agent-mcp/pipelines passes.

Closes #198 / #199 (both were stale against a since-refactored path anyway).

@google-oss-prow

Copy link
Copy Markdown

[APPROVALNOTIFIER] This PR is NOT APPROVED

This pull-request has been approved by:
Once this PR has been reviewed and has the lgtm label, please assign chasecadet for approval. For more information see the Kubernetes Code Review Process.

The full list of commands accepted by this bot can be found here.

Details Needs approval from an approver in each of these files:

Approvers can indicate their approval by writing /approve in a comment
Approvers can cancel approval by writing /approve cancel in a comment

Sharkyii added 3 commits July 19, 2026 12:00
…h locally

Signed-off-by: Sneh Kansagara <snehkansagara@gmail.com>
Signed-off-by: Sneh Kansagara <snehkansagara@gmail.com>
Signed-off-by: Sneh Kansagara <snehkansagara@gmail.com>
@Sharkyii
Sharkyii force-pushed the fix/pipeline-cuda-bloat branch from 7ef39bb to 793c712 Compare July 19, 2026 06:31

@ArshVermaGit ArshVermaGit left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice optimization. Switching the pipeline away from the unnecessary CUDA image reduces image bloat while keeping the embedding flow consistent with the shared service. The documentation and test updates round out the change nicely. LGTM.

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

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

bug(kagent-feast-mcp): requirements.txt installs full CUDA torch (~2.5GB) on a CPU-only base image

2 participants