9.2 Custom Training on Agent Platform: Containers, Jobs & Worker Pools
Key Takeaways
- A serverless training job (CustomJob) runs your code from either a Python training application on a prebuilt container or a custom container image.
- In a distributed job, workerPoolSpecs[0] is the primary replica with replicaCount 1, [1] holds workers, [2] holds parameter servers or Reduction Server, and [3] holds evaluators.
- Agent Platform sets AIP_MODEL_DIR, AIP_CHECKPOINT_DIR, and AIP_TENSORBOARD_LOG_DIR when you set a base output directory, so code can write artifacts without hard-coded paths.
- A persistent resource keeps a training cluster running between jobs, which assures capacity and cuts startup time but bills the entire time it runs.
- Spot VMs cut serverless training cost for fault-tolerant jobs, and Dynamic Workload Scheduler's FLEX_START strategy waits for requested GPUs to become available.
The exam guide lists model training using different SDKs (for example, Agent Platform custom training, Kubeflow on GKE, Agent Platform AutoML, and Tabular Workflows) and organizing training on Google Cloud. This section covers custom training. Section 9.3 covers Kubeflow on GKE, AutoML, and Tabular Workflows.
Three Custom Training Resources
| Resource | What it does | Use when |
|---|---|---|
| CustomJob (serverless training job) | Runs your training code on the machines you specify | Standard training runs, single node or distributed |
| HyperparameterTuningJob | Runs many trials of a CustomJob with different hyperparameters | Searching for the best configuration (Chapter 10) |
| TrainingPipeline | Orchestrates a training job plus steps such as loading a managed dataset and uploading the model to Model Registry | You want dataset handling and model upload built in |
Packaging Your Code
| Option | How | Choose when |
|---|---|---|
| Python training application + prebuilt container | Package code as a Python source distribution, and pick a Google prebuilt training image for TensorFlow, PyTorch, scikit-learn, or XGBoost | Standard framework versions. Fastest path with Google-maintained dependencies |
| Custom container | Build a Docker image with your code and dependencies, and push it to Artifact Registry | Unusual libraries, system packages, non-Python code, exact reproducibility, or a framework version without a prebuilt image |
| Autopackaging | gcloud ai custom-jobs create builds an image from local code, pushes it, and creates the job in one command | Quick iteration from a laptop or notebook |
Prebuilt containers follow a framework support policy with end-of-patch and end-of-availability dates. Pin a supported version and plan upgrades. Images must be referenced with their latest tag, which gets the most recent patch.
Anatomy of a CustomJob
A CustomJob includes:
- One or more worker pools (
workerPoolSpecs), each with a machine type, optional accelerators (type and count), disk, replica count, and either aPythonPackageSpecor aContainerSpec. - Scheduling: timeout, restart behavior, and a strategy such as Spot or flex-start.
- A service account, network (VPC peering or Private Service Connect interface), and environment variables.
- A base output directory, from which Agent Platform sets:
AIP_MODEL_DIR: where to export the final model artifactsAIP_CHECKPOINT_DIR: where to write checkpointsAIP_TENSORBOARD_LOG_DIR: where TensorBoard logs go
Training code must export model artifacts to Cloud Storage. Local disk disappears when the job ends. If you want to upload the model to Model Registry for Agent Platform serving, export it in a format the serving containers support, such as TensorFlow SavedModel, a scikit-learn or XGBoost file, or a PyTorch model archive for a PyTorch serving container.
Worker pool order for distributed training
| Position | Role |
|---|---|
workerPoolSpecs[0] | Primary (chief) replica. Required. replicaCount must be 1 |
workerPoolSpecs[1] | Workers |
workerPoolSpecs[2] | Parameter servers or Reduction Server |
workerPoolSpecs[3] | Evaluators (TensorFlow usually expects at most one) |
Leave a position empty ({}) to skip it. Code finds its role through the CLUSTER_SPEC or TF_CONFIG environment variables. Distributed strategies are covered in Chapter 11.
Identity, Networking, and Security
- Assign a custom service account with least privilege instead of relying on the default service agent's project-wide bucket access.
- Use VPC Network Peering or a Private Service Connect interface when training code must reach private resources such as databases or on-premises systems.
- Use VPC Service Controls and CMEK as required (Chapter 18).
- Store secrets in Secret Manager, not in container images or code.
Capacity and Cost Options
| Option | What it does | Trade-off |
|---|---|---|
| On-demand serverless training | Provision resources per job and release them afterward | Startup time. No guaranteed capacity |
| Spot VMs | Big discounts on spare capacity | Preemption. The job fails with a STOCKOUT error and Compute Engine retries up to six times, so code must checkpoint and resume. Not supported for TPU Pods |
Dynamic Workload Scheduler (scheduling.strategy = FLEX_START) | Waits until requested GPUs (L4, A100, H100, H200, B200) are available, and starts all nodes together | Delayed start. Jobs are limited to 7 days, and all worker pools need the same machine configuration |
| Reservations | Use Compute Engine reservations for guaranteed capacity | Pay for the reservation |
| Persistent resource | A long-running cluster that stays up between jobs | Billed the entire time it runs, even when idle. Best for many short jobs, repeated jobs that benefit from image and data caching, or critical peak periods |
| Managed Training Clusters | Reserved, dedicated high-end accelerator clusters | Highest commitment. For large, long-running training |
Making Jobs Resilient
Long or preemptible jobs fail in predictable ways. Build for it:
- Checkpoint regularly to
AIP_CHECKPOINT_DIR(Cloud Storage) and resume from the latest checkpoint at startup. - Make steps idempotent, so rerunning a partially finished job doesn't duplicate outputs.
- Set a sensible timeout, so hung jobs don't burn accelerator hours.
- Log to Cloud Logging and stream metrics to TensorBoard, so failures can be diagnosed after the job ends (Chapter 10).
- Use the interactive shell on a running job when you need to inspect a live replica.
Worked Scenario
A team trains a PyTorch recommendation model nightly on 8 A100 GPUs. Jobs sometimes wait hours for GPUs, and the model occasionally needs an unusual CUDA extension.
- Build a custom container with the extension and push it to Artifact Registry, since no prebuilt image includes it.
- Enable Dynamic Workload Scheduler (
FLEX_START) so the job starts when all A100s are available together, or buy a reservation if the nightly deadline is strict. - Write checkpoints to
AIP_CHECKPOINT_DIRso a restarted job resumes instead of starting over. - Run the job as an Agent Platform Pipelines step that uploads the model to Model Registry and evaluates it (Chapter 15).
A training job needs a system library and a compiled C++ extension that no Google prebuilt training image includes. How should the team package it?
In a distributed CustomJob, which worker pool configuration is required?
A research group runs about 200 short experiments a day. Each takes 4 minutes of training but waits 6 minutes for VMs to provision. Which option most directly cuts total turnaround time?