2.6 Prototyping in Workbench and Colab Enterprise with PyTorch, scikit-learn and JAX
Key Takeaways
- 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.
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.
A platform team must give 90 analysts notebook access while centrally enforcing machine type, accelerator limits, subnet, service account, and idle timeout, with no analyst provisioning virtual machines. Which product and mechanism fits?
A team has a 180,000-row tabular dataset with 40 mixed numeric and categorical features and needs a strong baseline classifier within a day. Which prototyping approach is most appropriate?
A researcher needs a custom training loop with fine-grained control over gradient transformations and per-device sharding, targeting Cloud TPUs for a large numerical simulation-style model. Which framework is the strongest fit?
A data scientist has a working notebook that trains a model in roughly six hours on a GPU-attached instance and needs to run it weekly. What is the appropriate next step?