6.2 Prototyping with PyTorch, scikit-learn, JAX & Model Garden in Notebooks

The exam guide expects you to develop models in Agent Platform Workbench or Colab Enterprise notebooks using common frameworks (for example, PyTorch, sklearn, and JAX) and to use foundational and open-source models in Model Garden to create prototypes in notebook environments. The goal is fast iteration that can later become production training without a rewrite.

Choosing a Framework for the Prototype

Setting Up the Environment

• Workbench instances ship with JupyterLab and preinstalled TensorFlow and PyTorch packages. Add conda environments for other kernels, such as a specific framework version or JAX. Use a custom container when the whole team needs the same image.
• Colab Enterprise runtimes use templates for machine type and GPUs. Install extra packages at the start of the notebook or with a post-startup script, and pin versions for reproducibility.
• Pick GPUs for deep learning prototypes. Tabular scikit-learn work usually needs only CPU and memory.
• Read data from BigQuery (the built-in integration or BigQuery DataFrames) and Cloud Storage, not local uploads, so the same paths work in training jobs.

Prototyping Patterns That Scale

1. Sample first. Develop on a stratified sample (for example, 1% of rows) to iterate in seconds. Validate the pipeline end to end before using full data.
2. Structure code as functions (load_data, preprocess, train, evaluate), not a long run of cells that depend on hidden state.
3. Fix random seeds and log library versions.
4. Establish a baseline (a simple model or a BigQuery ML model) before building complex architectures, so improvements are measurable.
5. Log parameters and metrics to Experiments on Agent Platform from the first run (Section 6.3).
6. Profile before scaling. A slow input pipeline is often the bottleneck, not the model.

From notebook to scale

Notebook anti-patterns to avoid

• Cells that only work when run out of order, which leaves hidden state no one else can reproduce.
• Training on the full dataset inside an interactive session on an oversized GPU VM for hours.
• Copying production data to the runtime's local disk, where it's lost at deletion and bypasses governance.
• Hard-coded paths, project IDs, and model versions scattered through cells instead of one configuration cell.

Using Foundation Models in a Notebook

Most gen AI prototypes start in a notebook by calling Gemini through the Google Gen AI SDK configured for Agent Platform:

from google import genai
client = genai.Client(vertexai=True, project=PROJECT_ID, location="global")
resp = client.models.generate_content(
    model=MODEL_ID,  # keep the model version in configuration
    contents=["Summarize this incident report in 3 bullets:", report_text],
)
print(resp.text)

Prototype checklist for gen AI notebooks:

• Build a small evaluation set (inputs plus expected outputs or rubrics) right away.
• Compare prompts and models side by side, then log results as experiment runs.
• Test structured output with a response schema if downstream code parses JSON.
• Track token counts and latency to estimate production cost early (Chapter 4).
• Never paste production PII into prompts during exploration unless it's allowed and protected.

Using Open Models from Model Garden

Model Garden supports notebook-first workflows for open models:

• Open notebook on a model card: most open foundation and fine-tunable models include a sample notebook for tuning, deployment, and inference in Colab Enterprise or Workbench.
• Model Garden SDK: list deployable models, check deployment options (machine type and accelerator), and deploy from code:

from vertexai import model_garden
model = model_garden.OpenModel("google/gemma3@gemma-3-12b-it")
print(model.list_deploy_options())
endpoint = model.deploy(machine_type="g2-standard-48",
                        accelerator_type="NVIDIA_L4",
                        accelerator_count=4,
                        accept_eula=True)

• Hugging Face models can be deployed through the same flow. Google-verified deployment configurations reduce trial and error.
• Tuning: use the model card's fine-tuning notebook or pipeline for LoRA or full tuning, then deploy the tuned weights.

Cost warning: a deployed endpoint bills for its accelerators while it exists. Undeploy prototype endpoints when you finish, or use MaaS APIs for prototypes that don't need your own weights.

Worked Scenario

A retailer wants to test whether an open vision-language model can tag product photos better than its current rules. In Colab Enterprise, a data scientist:

1. Loads 500 labeled product photos from Cloud Storage.
2. Calls Gemini with a tagging prompt and a response schema to get a strong managed baseline.
3. Deploys an open vision-language model from Model Garden to a temporary GPU endpoint with the SDK, and runs the same 500 photos.
4. Logs accuracy, latency, and estimated cost per 1,000 images for both models as experiment runs.
5. Undeploys the endpoint. The team picks the approach that meets the accuracy target at the lowest total cost, then moves the chosen path into a pipeline.