Module 4 — Custom training and containers

The parquet snapshot from module 3 is ready. This module moves the training out of the notebook and into a CustomTrainingJob — a container Vertex launches on a managed VM, runs to completion, and cleans up. That single move unlocks reproducibility, right-sized hardware, and everything the rest of the course depends on: tuning, registry, endpoints, pipelines.

Prebuilt containers, or your own

Vertex offers two families of training containers.

Prebuilt containers are Google-maintained images with a framework already installed: europe-docker.pkg.dev/vertex-ai/training/xgboost-cpu.1-7, .../tf-cpu.2-15, .../pytorch-gpu.2-3, etc. You point them at a script, they run it. Zero build time, zero image to store, patched by Google.

job = aiplatform.CustomTrainingJob(
    display_name="fraud-xgb-baseline",
    script_path="trainer/task.py",
    container_uri="europe-docker.pkg.dev/vertex-ai/training/xgboost-cpu.1-7:latest",
    requirements=["pandas==2.2.2", "scikit-learn==1.5.1"],
    model_serving_container_image_uri=(
        "europe-docker.pkg.dev/vertex-ai/prediction/xgboost-cpu.1-7:latest"
    ),
    staging_bucket="gs://fraud-detection-dev-vertex-eu",
)

Custom containers are yours: a Dockerfile you build and push to Artifact Registry, containing exactly the libraries and versions the code needs. Choose them when the framework version is not in the prebuilt list, when the training needs a system package (libpq-dev for a Postgres reader, ffmpeg for audio), or when the same image must be used by other jobs — for reproducibility.

The red thread will use a custom container for the tuning of module 5, because the exact XGBoost version and a small feature-engineering library must be pinned and shared across trials.

The training script Vertex expects

Whatever the container, the training script is a plain Python module that reads inputs from environment variables Vertex sets, writes the model artifact to a path Vertex tells it to, and reports metrics to standard output. A minimal script for the fraud model:

# trainer/task.py
import argparse
import os
import xgboost as xgb
import pandas as pd
from sklearn.metrics import average_precision_score

parser = argparse.ArgumentParser()
parser.add_argument("--training_data", required=True)  # gs://.../training/*.parquet
parser.add_argument("--n_estimators", type=int, default=400)
parser.add_argument("--max_depth", type=int, default=6)
args = parser.parse_args()

model_dir = os.environ["AIP_MODEL_DIR"]            # gs:// path Vertex reserves
tb_dir    = os.environ.get("AIP_TENSORBOARD_LOG_DIR")

df = pd.read_parquet(args.training_data)
y  = df.pop("is_fraud")
X  = pd.get_dummies(df, columns=["merchant_category", "device_type"])

clf = xgb.XGBClassifier(
    n_estimators=args.n_estimators,
    max_depth=args.max_depth,
    scale_pos_weight=(y == 0).sum() / (y == 1).sum(),  # imbalance
    tree_method="hist",
    n_jobs=-1,
)
clf.fit(X, y)

ap = average_precision_score(y, clf.predict_proba(X)[:, 1])
print(f"aucpr: {ap:.4f}")                          # picked up by Vizier in module 5

clf.save_model(f"/tmp/model.bst")
import subprocess
subprocess.check_call(["gsutil", "cp", "/tmp/model.bst", f"{model_dir}/model.bst"])

Two conventions to internalise. AIP_MODEL_DIR is where the model file must be written; Vertex reads it back when constructing the Model object for the registry. Print the metric on stdout with a key: value line — that is exactly how Vizier will read tuning trial results in the next module.

Launching the job from the notebook

job = aiplatform.CustomTrainingJob(
    display_name="fraud-xgb-baseline",
    script_path="trainer/task.py",
    container_uri="europe-docker.pkg.dev/vertex-ai/training/xgboost-cpu.1-7:latest",
    requirements=["pandas==2.2.2", "scikit-learn==1.5.1"],
    model_serving_container_image_uri=(
        "europe-docker.pkg.dev/vertex-ai/prediction/xgboost-cpu.1-7:latest"
    ),
    staging_bucket="gs://fraud-detection-dev-vertex-eu",
)

model = job.run(
    args=[
        "--training_data=gs://fraud-detection-dev-vertex-eu/training/exp-042/2026-06-*.parquet",
        "--n_estimators=400",
        "--max_depth=6",
    ],
    replica_count=1,
    machine_type="n1-highmem-8",
    service_account="vertex-training@fraud-detection-dev.iam.gserviceaccount.com",
    labels={"env": "dev", "team": "fraud", "experiment": "xgb-baseline"},
)

job.run() blocks until the job finishes, returning a Model that is already in the Model Registry (module 6). The service account attached is the least-privileged one from module 1 — do not skip that argument.

Machine types and accelerators, without paying for what you do not use

Vertex exposes the Compute Engine catalogue. Two families cover 90% of ML training:

Family	Good for	Rough hourly cost
n1-standard-*	Classical ML, small deep learning	$0.20–$0.80
n1-highmem-*	Larger tabular data, feature-heavy	$0.30–$1.60
n1-highcpu-*	CPU-bound preprocessing	$0.15–$0.60
a2-highgpu-1g	One A100 for medium models	$3–$4
a2-highgpu-8g	Eight A100s for large training	$25–$30

Accelerators are attached separately (accelerator_type= and accelerator_count=). The pragmatic rule: start CPU-only. XGBoost on the fraud dataset trains in about two minutes on n1-highmem-8 for $0.02 per run. A GPU here would idle and cost 20×.

A GPU is warranted for the deep-learning baseline of module 5's hyperparameter search, and even there NVIDIA_TESLA_T4 (about $0.35 per hour) covers the exercise — reserve A100s for real production workloads.

"Just try an A100 to see if it's faster"

Testing on an A100 without a workload that saturates it is expensive theatre. A tabular XGBoost job on an A100 costs the same as on a CPU and finishes in the same time, because the training does not touch the GPU. Confirm the framework and the workload can actually consume the accelerator before you attach one.

Where the trained model actually lands

At the end of a successful job.run():

1. The model file the script wrote to AIP_MODEL_DIR sits in gs://<staging>/aiplatform-custom-training-<timestamp>/model/.
2. A Model resource is created in the registry (module 6), pointing at that GCS path.
3. The model_serving_container_image_uri you passed determines which prebuilt container will run predictions on that model when it is deployed.

That third point matters. Even for a custom-trained model, deployment usually re-uses a prebuilt serving container matched to the training framework. The XGBoost serving image knows how to load model.bst, exposes /predict, and handles request batching — you get all of that without writing a Flask app.

Building a custom container: the four-line Dockerfile

For the tuning job of module 5 the container needs a specific XGBoost version and a small internal wheel. The Dockerfile is minimal:

FROM europe-docker.pkg.dev/vertex-ai/training/base-cpu.py311:latest

RUN pip install --no-cache-dir \
    xgboost==2.1.0 \
    pandas==2.2.2 \
    scikit-learn==1.5.1

COPY trainer /trainer
ENTRYPOINT ["python", "-m", "trainer.task"]

Build and push:

gcloud builds submit \
    --tag europe-docker.pkg.dev/fraud-detection-dev/vertex/fraud-trainer:v3

Then use it in the job with container_uri=".../fraud-trainer:v3". Immutable tags (v3, not latest) mean a tuning trial two months from now still runs against the same code.

In summary

• A CustomTrainingJob moves training out of the notebook into a container Vertex runs, cleans up and bills per second — the foundation for tuning, registry and pipelines.
• Prebuilt containers are enough for standard frameworks; write a custom container when you need a specific version or a system package, and pin its tag.
• The training script reads inputs from CLI flags, writes the artifact to AIP_MODEL_DIR, and prints metrics as key: value for Vizier to pick up.
• Start CPU-only; add a GPU only when the workload actually consumes it, and prefer a small accelerator (T4) to a giant idle one (A100).

Next module: turning one training run into a systematic hyperparameter search with Vizier, parallel trials and early stopping.