Scaling prototypes into ML models · Topic 2 of 4
Organizing Training Data and Running Training Jobs
Organize training data in Cloud Storage and BigQuery, choose custom training, Kubeflow on GKE, AutoML or Tabular Workflows, and fix training failures.
- 5 min
- Tabular data in BigQuery, files in Cloud Storage, shards for large deep learning sets, same region as training.
- Prebuilt containers for standard frameworks; custom containers for anything else.
- Kubeflow on GKE only when Kubernetes control or portability is required.
What it is
Production training runs as a managed job, not in a notebook. The engineer organizes the training data where jobs can read it fast, picks the training service that fits the code and the team, and diagnoses jobs that fail or run badly.
How it works
Organizing training data
| Data | Store | Format and layout |
|---|---|---|
| Tabular | BigQuery | Tables or views per training snapshot; read with the BigQuery Storage Read API for speed |
| Images, video | Cloud Storage | Files plus a CSV or JSON Lines manifest of URI and label |
| Speech and audio | Cloud Storage | Audio files plus transcripts or labels in a manifest |
| Text | Cloud Storage or BigQuery | JSON Lines for examples; tables for labeled rows |
| Large deep learning sets | Cloud Storage | Many shards (for example TFRecord), so workers read in parallel |
Keep training, validation and test splits fixed and versioned so runs are comparable. Put data in the same region as the training job to avoid slow, billed cross-region reads.
Ingesting data into training
| Source | Way in |
|---|---|
| BigQuery | Read directly in the training code, or export a snapshot to Cloud Storage files |
| Cloud Storage | Custom training jobs can read buckets as a local file system under /gcs/ (Cloud Storage FUSE) |
| Streaming or operational sources | Land them first (Pub/Sub to BigQuery or Cloud Storage with Dataflow), then train on the landed data |
| Unstructured files from many systems | A Dataflow or pipeline step that copies, converts and labels them into one bucket |
| Managed datasets | AutoML and custom jobs can read a managed dataset with its splits |
Inside the training container the service sets environment variables such as AIP_MODEL_DIR, the Cloud Storage path where the job should write the trained model.
Choosing how to train
| Option | What it is | Choose it when |
|---|---|---|
| Agent Platform custom training with a prebuilt container | Your Python script runs in a Google-provided TensorFlow, PyTorch, scikit-learn or XGBoost container | Standard frameworks; least container work |
| Agent Platform custom training with a custom container | Your own Docker image from Artifact Registry | Special dependencies, other languages or frameworks |
| Kubeflow on GKE | Kubeflow and training operators on your own GKE cluster | Existing Kubernetes platform, portability, full cluster control |
| Agent Platform AutoML | Managed architecture and hyperparameter search | Low code on tabular or image data |
| Tabular Workflows | Managed, configurable pipelines for tabular models (including AutoML end to end, TabNet, Wide and Deep, and forecasting) | AutoML-level automation with more control over steps, search space and hardware |
from google.cloud import aiplatform
PREBUILT_TRAINING_IMAGE = "..." # a prebuilt XGBoost training container URI
aiplatform.init(project="my-project", location="us-central1",
staging_bucket="gs://my-staging")
job = aiplatform.CustomTrainingJob(
display_name="churn-xgb",
script_path="train.py",
container_uri=PREBUILT_TRAINING_IMAGE,
requirements=["pandas"],
)
job.run(replica_count=1, machine_type="n1-standard-8",
args=["--n-estimators", "300"])Fully managed training beats self-managed clusters unless the stem says the company already runs Kubernetes, needs portability, or must control the cluster.
Troubleshooting training failures
| Symptom | Likely cause | Fix |
|---|---|---|
| Out of memory on GPU or VM | Batch or model too large | Smaller batch, bigger machine, mixed precision, or distribute the model |
| Permission denied reading data | Job's service account lacks access | Grant the needed role on the bucket or dataset to that service account |
| Job waits or fails for resources | Accelerator quota or regional capacity | Request quota, try another region or machine type |
| Loss becomes NaN | Learning rate too high, unscaled inputs, bad values | Lower the learning rate, normalize, clip gradients, clean data |
| Training slow, GPU mostly idle | Input pipeline is the bottleneck | Parallel reads, prefetching, sharded files, data in the same region |
| Container exits at start | Missing package, wrong entry point | Read the job logs in Cloud Logging; test the container locally first |
| Job killed part way on Spot VMs | Preemption | Save checkpoints to Cloud Storage and resume from the last one |
| Great training score, poor validation | Overfitting | Regularization, more data, early stopping, simpler model |
Job logs go to Cloud Logging, and resource charts show CPU, GPU and memory use; start every investigation there.
A worked example
A PyTorch image model trained on a notebook's GPU takes 30 hours and sometimes crashes. The engineer moves 5 million images into Cloud Storage as shards in the training region and packages the code for a prebuilt PyTorch container. The custom training job reads the shards through /gcs/ with parallel data loaders.
GPU utilization was low because of slow reads; after sharding it stays high. Checkpoints are written every epoch, so Spot VMs can be used safely.
How the exam asks about this
- "A training job fails with permission denied on a bucket." Grant the job's service account read access to the bucket.
- "GPU utilization is low and training is slow." Fix the input pipeline: shard files, read in parallel and prefetch.
- "The team wants AutoML-style tabular training but must control the hardware and search space." Use Tabular Workflows.
What to remember for the exam
- Tabular data in BigQuery, files in Cloud Storage, shards for large deep learning sets, same region as training.
- Prebuilt containers for standard frameworks; custom containers for anything else.
- Kubeflow on GKE only when Kubernetes control or portability is required.
- Tabular Workflows add control to AutoML-style tabular training.
- Debug with Cloud Logging; checkpoint so preempted or failed jobs can resume.
Training from data in another region or continent slows every epoch and adds transfer cost. Keep data and training jobs co-located.
Common mistakes
Training production models from a notebook. It is not reproducible or scalable; submit a training job.
Adding GPUs when the input pipeline is slow. More accelerators sit idle if data cannot reach them; fix reading first.
Using Spot VMs without checkpoints. A preempted job restarts from the beginning and wastes the savings.
Check yourself
Where does a custom training job save the trained model?
In the Cloud Storage path given by the AIP_MODEL_DIR environment variable that the service sets inside the container.
When is a custom container needed instead of a prebuilt one?
When the code needs dependencies, frameworks or languages that the prebuilt training containers do not provide.
What is the first fix for a slow job with a mostly idle GPU?
Speed up the input pipeline with sharded files, parallel reads and prefetching, with data in the same region as the job.
Topic 2 of 4 in Scaling prototypes into ML models
Test yourself on Scaling prototypes into ML models
The full Google Cloud Professional Machine Learning Engineer question bank has 749 questions, with timed practice exams and an explanation for every answer.
This study guide is independent and is not affiliated with, endorsed by, or sponsored by Google Cloud. Vendor names, certification names and exam codes are trademarks of their respective owners and are used for identification only.