2.6 Prototyping in Workbench and Colab Enterprise with PyTorch, scikit-learn and JAX

Google Cloud Professional Machine Learning Engineer Exam

Track your progress as you read

Finished sections get a check mark.

1 Introduction & Exam Overview

0.1 Exam Overview, Blueprint & Certification Logistics 0.2 Study Strategy, Scenario Mechanics & 8-Week Roadmap

2 Section 1: Architecting Low-Code AI Solutions (~13%)

1.1 BigQuery ML: Model Types, SQL Syntax and In-Database Evaluation 1.2 Feature Engineering and Feature Selection in BigQuery ML 1.3 Generating Predictions with BigQuery ML 1.4 AutoML on Agent Platform: Tabular, Image, Text and Video Training 1.5 Fine-Tuning Gemini Models from BigQuery 1.6 Model Garden: Evaluating and Selecting the Right Model 1.7 Industry-Specific AI APIs: Document AI, Vision, Translate and Speech 1.8 Generative Media and Models as a Service: Gemini, Imagen and Veo 1.9 Optimizing Gemini Applications for Cost, Latency and Availability

3 Section 2: Collaborating Within and Across Teams to Manage Data and Models (~16%)

2.1 Organizing and Exploring Tabular, Text, Image and Video Data 2.2 Choosing the Preprocessing Tool: BigQuery, Dataflow, Spark and In-Memory Python 2.3 Creating and Consolidating Features in Agent Platform Feature Store 2.4 Data Privacy, PII Handling and Sensitive Information 2.5 Notebook Environments: Collaboration and Security Best Practices 2.6 Prototyping in Workbench and Colab Enterprise with PyTorch, scikit-learn and JAX 2.7 Prototyping with Model Garden Foundational and Open Models in Notebooks 2.8 Choosing the Development and Experimentation Environment 2.9 Evaluating Generative Solutions: LLM-as-a-Judge and RAG Evaluation 2.10 Experiment Tracking, Artifact Lineage and Model Registry Governance

4 Section 3: Scaling Prototypes into ML Models (~21%)

3.1 Choosing the Model Type: ARIMA, DNN, LLM and Classical Baselines 3.2 Choosing the Product: AutoML, BigQuery ML, Custom Training and Pipelines 3.3 Choosing the Deployment Strategy at Design Time 3.4 Modeling Techniques for Interpretability Requirements 3.5 Organizing Training Data on Cloud Storage and BigQuery 3.6 Ingesting Structured and Unstructured Data into Training Pipelines 3.7 Training SDKs and Services: Custom Training, Kubeflow on GKE, AutoML and Tabular Workflows 3.8 Troubleshooting ML Model Training Failures 3.9 Hyperparameter Tuning, Bayesian Optimization and Vizier 3.10 Model Evaluation Metrics, Transfer Learning and Fine-Tuning 3.11 Evaluating Compute and Accelerator Options: CPU, GPU and TPU 3.12 Distributed Training Strategies, Hardware Acceleration and Performance Optimization

5 Section 4: Serving and Scaling Models (~20%)

4.1 Batch and Online Inference: Agent Platform, Model Garden, Cloud Run and GKE 4.2 Rollout Strategies: Canary, Traffic Splitting, Shadow and Blue-Green 4.3 Tuning Models for Production: Vector Search and Model Compression 4.4 Deploying Models to Public and Private Endpoints 4.5 Choosing Serving Hardware: CPU, GPU, TPU and Edge 4.6 Scaling the Serving Backend for Throughput

6 Section 5: Automating and Orchestrating ML Pipelines (~18%)

5.1 Validating Data and Models Inside Pipelines 5.2 Agent Platform Pipelines: Kubeflow Pipelines (KFP) and TFX Architecture 5.3 Reusable Pipeline Components, Artifacts and ML Metadata Lineage 5.4 Ray on Agent Platform for Distributed Python Workloads 5.5 MLOps CI/CD Automation, Retraining Triggers and Skew Prevention

7 Section 6: Monitoring AI Solutions (~13%)

6.1 Securing AI Systems: Exfiltration, Malicious Prompting and Model Armor 6.2 Responsible AI, Explainable AI, Fairness Metrics and Model Governance 6.3 Model Monitoring: Feature Skew and Drift Detection 6.4 Operational Troubleshooting, Performance Profiling and Cost Optimization 6.5 Monitoring, Testing and Evaluating Generative AI in Production

2.6 Prototyping in Workbench and Colab Enterprise with PyTorch, scikit-learn and JAX

Key Takeaways Explain This Topic AI

Agent Platform Workbench provides a managed JupyterLab instance you control, size, and keep; Colab Enterprise provides a managed, IAM-governed Colab experience with shared runtime templates.

Colab Enterprise runtimes are created from runtime templates that centrally fix machine type, accelerator, network, and idle timeout for a whole team.

scikit-learn is the right prototyping choice for tabular problems at in-memory scale; PyTorch dominates deep learning research and custom architectures; JAX suits high-performance numerical work and TPU-first training.

Framework choice at prototyping time constrains later serving, because prebuilt inference containers exist for common framework and version combinations.

Prototypes should graduate into a training script or pipeline component rather than being re-executed by hand.

Last updated: September 2026

2.6 Prototyping in Workbench and Colab Enterprise with PyTorch, scikit-learn and JAX

Blueprint reference: Section 2.2, "Developing models in Agent Platform Workbench or Colab Enterprise notebooks using common frameworks (e.g., PyTorch, sklearn, and JAX)."

Two managed notebook products, three frameworks, one recurring question: which one, and why.

Workbench Versus Colab Enterprise

Dimension Agent Platform Workbench Colab Enterprise

Mental model A managed VM running JupyterLab that you own A managed Colab experience with runtimes created on demand

Provisioning You create an instance with a chosen machine type and accelerator A runtime template defines the shape; runtimes are created from it

Persistence Persistent disk survives stop/start; environment is yours Runtimes are ephemeral; notebooks persist as resources

Customization Custom container images, arbitrary apt/pip installs, long-lived environments Standard image with pip installs per session

Governance Per-instance IAM, service account, network config Central runtime templates enforce machine type, network, and idle timeout for everyone

Sharing Share the instance or the repository Notebooks are Google Cloud resources shared with IAM, familiar Colab UX

Best for Deep customization, long-lived environments, heavy local dependencies Teams that want a consistent, governed, low-friction notebook with no VM to manage

The exam discriminator is governance versus customization . If a scenario emphasizes that an administrator must enforce a consistent machine shape, network configuration, and idle timeout across a large team of analysts who should not be provisioning VMs, that is Colab Enterprise with runtime templates. If it emphasizes a bespoke environment — an unusual CUDA build, a compiled dependency, a container image the team maintains — that is Workbench.

Both run as a service account, both should sit on private IP, and both are subject to the security and cost practices covered in the previous section.

Runtime Templates

A runtime template is the reusable specification a Colab Enterprise runtime is created from: machine type, accelerator, disk, network and subnet, service account, and idle shutdown. Administrators publish templates; users create runtimes from them.

This solves a real governance problem. Without templates, every analyst chooses their own machine and network settings and the organization discovers a fleet of over-provisioned, publicly-addressed instances. With templates, the shape is set once and inherited.

Choosing the Framework

Framework Sweet spot Strengths Watch out for

scikit-learn Tabular problems that fit in memory Fast iteration, excellent classical algorithms, pipelines and cross-validation built in Single-machine; no GPU; does not scale past host memory

XGBoost / LightGBM Tabular problems where gradient boosting wins Usually the strongest tabular baseline; handles missing values natively Still largely single-node unless deliberately distributed

PyTorch Deep learning, custom architectures, transfer learning Dominant research ecosystem, dynamic graphs, huge pretrained model availability Distributed training needs explicit setup

TensorFlow / Keras Production deep learning with a mature serving story TFX integration, SavedModel format, TF Serving, tf.data More ceremony for research-style iteration

JAX High-performance numerical computing, TPU-first training Function transformations ( jit , grad , vmap , pmap ), excellent TPU performance Smaller ecosystem; more of the training loop is yours to write

The most common framework question on the exam is not "which is best" but "which is appropriate given the data." A 200,000-row tabular dataset does not need PyTorch; scikit-learn or gradient boosting will train in seconds and beat a hastily-built neural network. Conversely, fine-tuning a vision transformer is not a scikit-learn task.

JAX and TPUs deserve one specific note: JAX's transformation model maps cleanly onto TPU execution, and pmap /sharding APIs express data and model parallelism directly. When a scenario pairs "TPU" with "custom research training loop," JAX is a strong signal.

Prototyping Habits That Survive Promotion

A notebook that will become production code should be written with that in mind:

Parameterize instead of hard-coding. Dates, paths, and hyperparameters as variables at the top, so the notebook can later be executed with different inputs.

Read from the same source production will use. Prototyping on a hand-exported CSV guarantees a skew surprise later.

Log to Experiments. Even during prototyping, record parameters and metrics to Experiments on Agent Platform so the run is comparable rather than lost in scrollback.

Keep transformation logic in functions. A function can be imported into a training script; a cell cannot.

Fix the seed and record it. Reproducibility starts in the notebook.

The graduation path is: notebook → training script in a container → custom training job → pipeline component. The exam expects you to know that long or repeated training belongs in a job, not an interactive kernel, both because kernels die and because the accelerator bills for the whole session rather than the training duration.

Model Garden Inside a Notebook

Both surfaces can pull models from Model Garden directly, which is covered in the next section. The relevant point here is that a notebook is the intended place to compare candidate models before committing to one — loading two or three, running them over a representative sample, and recording the scores as an experiment.

Exam Traps

Choosing a deep learning framework for a small tabular dataset. Gradient boosting is usually the better and faster answer.

Workbench when the requirement is central governance. Runtime templates are a Colab Enterprise capability.

Assuming Colab Enterprise runtimes persist. They are ephemeral; the notebook resource persists.

Training for hours in a kernel. Submit a training job.

Choosing a framework without checking serving support. Prebuilt inference containers cover specific framework and version pairs.