Generative AI Google Brain

open-models-use-cases-cloud-run-ollama-gemma2-rag-qa

4 min read

Cloud Run GPU Inference: Gemma 2 RAG Q&A with Ollama and LangChain

Source notebook

Repo path: open-models/use-cases/cloud_run_ollama_gemma2_rag_qa.ipynb · Open on GitHub · advanced

Deploys Gemma 2 on Cloud Run GPU with Ollama and builds a LangChain RAG Q&A chain.

Summary

This notebook teaches how to package Ollama with the Gemma 2 9B model, build the container with Cloud Build, and deploy it to Cloud Run with an NVIDIA L4 GPU. It then invokes the deployed Ollama API with curl, requests, and LangChain ChatOllama. The end-to-end workflow builds a RAG Q&A chain over Cloud Run documentation using Vertex AI embeddings, a SKLearn vector store, query rewriting, retrieval, and answer generation.

Key code patterns

Initialize Vertex AI

PROJECT_ID = os.environ.get("GOOGLE_CLOUD_PROJECT")
LOCATION = os.environ.get("GOOGLE_CLOUD_REGION", "us-central1")
vertexai.init(project=PROJECT_ID, location=LOCATION)

Sets project and region context before using Vertex AI embeddings.

Build Ollama Image

MODEL_NAME = "gemma2:9b"
ENV OLLAMA_HOST 0.0.0.0:8080
ENV OLLAMA_MODELS /models
ENV OLLAMA_KEEP_ALIVE -1
RUN ollama serve & sleep 5 && ollama pull $MODEL

Stores the Gemma 2 model in the container image and keeps weights loaded for Cloud Run GPU inference.

Deploy GPU Cloud Run Service

gcloud beta run deploy $SERVICE_NAME \
  --image $CONTAINER_URI \
  --concurrency 4 \
  --cpu 8 --gpu 1 \
  --gpu-type nvidia-l4 \
  --memory 32Gi \
  --no-allow-unauthenticated

Deploys a private Cloud Run service sized for GPU-backed Ollama inference.

Use ChatOllama Against Cloud Run

llm = ChatOllama(
    model=MODEL_NAME,
    base_url=SERVICE_URL,
    num_predict=300,
    headers={"Authorization": f"Bearer {ID_TOKEN}"},
)

Lets LangChain call the Cloud Run-hosted Ollama endpoint with an identity token.

Create Retriever

embeddings = VertexAIEmbeddings(
    project=PROJECT_ID,
    model_name="text-embedding-005",
    credentials=credentials,
)
docs = WebBaseLoader(url).load()
documents = CharacterTextSplitter(chunk_size=800, chunk_overlap=100).split_documents(docs)
retriever = SKLearnVectorStore.from_documents(documents, embeddings).as_retriever()

Loads Cloud Run documentation, chunks it, embeds it, and exposes semantic retrieval.

Compose RAG Chain

rag_chain = (
    {"context": extract_query | retriever | format_docs,
     "messages": RunnablePassthrough()}
    | answer_generation_template
    | llm
    | StrOutputParser()
)

Combines query rewriting, retrieval, context formatting, prompting, model invocation, and parsing.

Refresh Identity Tokens

class GoogleCloudAuth(AuthBase):
    def __call__(self, r):
        r.headers["Authorization"] = f"Bearer {self.get_token()}"
        return r
    def get_token(self):
        if time.time() >= self.expiry_time:
            self.refresh_token()
        return self.token

Handles hourly Cloud Run identity token expiration for long-running LangChain usage.

Models & APIs used

  • Models: gemma2:9b, text-embedding-005
  • APIs / services: Vertex AI, Cloud Run, Cloud Build, Artifact Registry, Cloud Storage
  • SDKs / libraries: google-cloud-aiplatform, vertexai, langchain-community, langchainhub, langchain_google_vertexai, langchain, google.auth, requests

When to use this

Use this pattern when serving an open Ollama model on Cloud Run GPU and grounding answers with LangChain RAG over web documentation.

Gotchas & caveats

  • Cloud Run GPU support is a guarded feature and the project must be enabled through g.co/cloudrun/gpu.
  • The notebook requires IAM roles for Artifact Registry, Cloud Build, Cloud Run, service account use, Service Usage, and Storage.
  • Vertex AI API must be enabled before using Vertex AI SDK and embeddings.
  • The runtime must be restarted after installing notebook packages.
  • Cloud Run service is deployed with —no-allow-unauthenticated, so requests require an identity token.
  • Identity tokens expire hourly by default and need refresh handling for uninterrupted use.
  • Initial deployment can be slower because the container image is pulled for the first time.
  • Concurrency should be tuned for the latency versus throughput tradeoff.

Best practices

  • Store Gemma 2 9B and similarly sized model weights directly in the container image for startup time and scalability.
  • Consider storage requirements before placing larger model weights in the image.
  • Use e2-highcpu-32 for Cloud Build to speed up parallel downloads.
  • Set OLLAMA_KEEP_ALIVE=-1 to avoid unloading model weights from GPU memory.
  • Use a dummy startup request to load the model into GPU memory.
  • Use retrieved context and instruct the model to say it does not know when the answer is unavailable.
  • Compare direct LLM answers with RAG answers to validate grounding quality.
  • Delete the Cloud Run service after the tutorial to clean up resources.

Graph View

Backlinks