Vertex AI Tutorial: ML on Google Cloud

If you search for a Vertex AI tutorial today, the documentation you land on has a different name at the top. Google's platform for building and running models is now called Gemini Enterprise Agent Platform, and its product page carries the old name in brackets: formerly Vertex AI (Google Cloud: Gemini Enterprise Agent Platform). The machine learning part works as it did. Several pieces keep their Vertex names, and the command line still uses gcloud ai.

This tutorial follows one model through the platform, from data to predictions.

The path of a model

Google's documentation now says inference where older material says prediction. They mean the same thing.

Two ways to train

AutoML asks for labelled data and nothing else. You create a managed dataset, choose the column or label to predict, and the platform searches for a model. Managed datasets are required for AutoML and cover data types such as images and tables (Google Cloud: machine learning on Agent Platform).

Custom training runs your own code. The training service starts machines, runs your program in a container, and shuts the machines down. It supports any framework, with built-in support for PyTorch, TensorFlow, scikit-learn and XGBoost, and you pay for the compute you use (Google Cloud: serverless training overview).

The contract for training code

A training job is your ordinary script with two rules about where things go (Google Cloud: prepare training code).

Reading. Cloud Storage buckets appear to the job as folders under /gcs/, and BigQuery can be queried with its client library.

Writing. The platform sets an environment variable, AIP_MODEL_DIR, to the place where the model should be saved. Your code writes there instead of to a path you invented.

There is a third rule if you want Google's ready-made serving container. For scikit-learn, the saved file must be called model.joblib, exactly (Google Cloud: export model artifacts). And the container must match the scikit-learn version you trained with. The newest prebuilt one listed is for version 1.6 (Google Cloud: prebuilt containers for inference).

A model, trained on a laptop first

The script follows that contract and runs anywhere. It predicts how long a New York taxi trip will take from its distance, the hour and whether it starts or ends at an airport (NYC TLC: trip record data). It trains on 15 January 2024 and tests on the 16th, a later day the model has not seen.

It ran locally. AIP_MODEL_DIR is not set on a laptop, so the model goes to a local folder. Nothing was created on Google Cloud.

import os

import joblib
import pandas as pd
from sklearn.dummy import DummyRegressor
from sklearn.ensemble import HistGradientBoostingRegressor
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_absolute_error

MODEL_DIR = os.environ.get("AIP_MODEL_DIR", "model_output")   # set by the platform; a local folder here
FEATURES = ["trip_distance", "hour", "airport_pickup", "airport_dropoff"]
AIRPORTS = [132, 138]                                          # JFK and LaGuardia zones

trips = pd.read_parquet("../../dedata/yellow_tripdata_2024-01.parquet",
                        columns=["tpep_pickup_datetime", "tpep_dropoff_datetime", "trip_distance",
                                 "PULocationID", "DOLocationID"])
trips["minutes"] = (trips.tpep_dropoff_datetime - trips.tpep_pickup_datetime).dt.total_seconds() / 60
trips = trips[trips.trip_distance.between(0.1, 50) & trips.minutes.between(1, 180)].copy()
trips["hour"] = trips.tpep_pickup_datetime.dt.hour
trips["airport_pickup"] = trips.PULocationID.isin(AIRPORTS).astype(int)
trips["airport_dropoff"] = trips.DOLocationID.isin(AIRPORTS).astype(int)
day = trips.tpep_pickup_datetime.dt.strftime("%Y-%m-%d")
train, test = trips[day == "2024-01-15"], trips[day == "2024-01-16"]
print(f"train: 15 January, {len(train):,} trips   test: 16 January, {len(test):,} trips")

models = {"median of the training day": DummyRegressor(strategy="median"),
          "linear regression": LinearRegression(),
          "gradient boosted trees": HistGradientBoostingRegressor(random_state=0)}
print()
print("model                        error on 16 January, minutes   average prediction minus actual")
for name, model in models.items():
    model.fit(train[FEATURES], train.minutes)
    predicted = model.predict(test[FEATURES])
    print(f"{name:28} {mean_absolute_error(test.minutes, predicted):>28.2f}   "
          f"{(predicted - test.minutes).mean():>+31.2f}")

os.makedirs(MODEL_DIR, exist_ok=True)
artifact = os.path.join(MODEL_DIR, "model.joblib")
joblib.dump(models["gradient boosted trees"], artifact)
print()
print("saved", artifact.replace(os.sep, "/"))

served = joblib.load(artifact)                                  # what a serving container does
instances = pd.DataFrame([[1.0, 8, 0, 0], [3.5, 18, 0, 0], [17.0, 15, 1, 0]], columns=FEATURES)
for row, minutes in zip(instances.itertuples(index=False), served.predict(instances)):
    print(f"  {row.trip_distance:>4} miles at {row.hour:02d}:00, airport pickup {row.airport_pickup}  ->  {minutes:.1f} minutes")

train: 15 January, 75,169 trips   test: 16 January, 90,660 trips

model                        error on 16 January, minutes   average prediction minus actual
median of the training day                           8.76                             -5.22
linear regression                                    5.58                             -2.37
gradient boosted trees                               4.72                             -2.68

saved model_output/model.joblib
   1.0 miles at 08:00, airport pickup 0  ->  7.6 minutes
   3.5 miles at 18:00, airport pickup 0  ->  19.6 minutes
  17.0 miles at 15:00, airport pickup 1  ->  44.1 minutes

• Start with a baseline. Always guessing the median trip is wrong by 8.76 minutes on average. That is the number a model has to beat.
• The models earn their place. A linear model cuts the error to 5.58 minutes, and boosted trees to 4.72.
• Every model guesses too low. The last column is negative in all three rows. The likely reason is the calendar: the 15th was a public holiday in the United States, Martin Luther King Jr. Day, and the 16th was a working Tuesday. The model learned a quiet city and met a busier one.
• Serving is the last four lines. Load the file, call predict. A serving container does the same behind a web address.

The third point is the lesson a random split would have hidden. Shuffle both days together and the error looks better, because the test set resembles the training set. Test on later data and you see what will happen in use. Once a model is deployed, Model Monitoring watches for exactly this: inputs that have moved away from what the model was trained on (Google Cloud: MLOps on Agent Platform).

One practical note. The run used scikit-learn 1.8.0. To use the prebuilt serving container, pin the version the container supports when you train.

The same script as a training job

Because the script already follows the contract, moving it to the platform is one command. This form builds a container from your local code on top of a prebuilt training image and starts the job (Google Cloud: create a serverless training job):

gcloud ai custom-jobs create \
  --region=us-central1 \
  --display-name=trip-duration-train \
  --worker-pool-spec=machine-type=n1-standard-4,replica-count=1,executor-image-uri=EXECUTOR_IMAGE_URI,local-package-path=.,script=train.py

When it finishes, the model file sits in Cloud Storage. Registering it in Model Registry gives it a name and a version, and later retraining adds versions under the same name (Google Cloud: Model Registry).

Online or batch

There are two ways to get predictions (Google Cloud: getting inferences):

The cost difference is large and easy to overlook. You pay for each model deployed to an endpoint whether or not anyone calls it, until you undeploy it (Google Cloud: Agent Platform pricing). A forgotten endpoint is the classic surprise on a learner's bill. See cost controls and budgets.

Most analytical uses are batch. If predictions are needed once a day and land in a table, there is no reason to keep a server waiting.

When you do not need the platform

If the data is in BigQuery and the model is a standard one, BigQuery ML trains and predicts in SQL, with no training code or containers. See BigQuery ML. Reach for a training job when you need a library, a model type or a level of control that SQL does not offer.

An order to learn it in

1. Pick one tabular problem and compute a baseline.
2. Train locally, and test on later data.
3. Make the script follow the contract: read from a path you pass in, write to AIP_MODEL_DIR.
4. Run it as a training job.
5. Run a batch inference job and read the output.
6. Deploy to an endpoint once, send a request, and undeploy.
7. Then automate the steps. See Vertex AI Pipelines for MLOps.

Common mistakes

• Leaving a model deployed.
• No baseline, so nobody knows whether 4.72 is good.
• A random split on data that has an order in time.
• Training with one library version and serving with another.
• Saving the model under a different filename and wondering why the container cannot find it.
• Starting on the platform before the model works locally. Every mistake then costs a job start.

Machine learning, model evaluation and deployment are covered in the Data Science course.

Sources

• Gemini Enterprise Agent Platform (formerly Vertex AI) (Google Cloud) — Google Cloud
• Introduction to machine learning on Gemini Enterprise Agent Platform (Google Cloud Documentation) — Google Cloud
• Serverless training overview (Google Cloud Documentation) — Google Cloud
• Prepare training code (Google Cloud Documentation) — Google Cloud
• Export model artifacts for inference and explanation (Google Cloud Documentation) — Google Cloud
• Prebuilt containers for inference and explanation (Google Cloud Documentation) — Google Cloud
• Create a serverless training job (Google Cloud Documentation) — Google Cloud
• Introduction to Model Registry (Google Cloud Documentation) — Google Cloud
• Overview of getting inferences (Google Cloud Documentation) — Google Cloud
• MLOps on Gemini Enterprise Agent Platform (Google Cloud Documentation) — Google Cloud
• Gemini Enterprise Agent Platform pricing (Google Cloud) — Google Cloud
• TLC Trip Record Data (NYC Taxi and Limousine Commission) — NYC Taxi and Limousine Commission