Inference Preprocessing and Postprocessing on Vertex AI

Vertex AI custom serving is the pattern where inference preprocessing, prediction, and postprocessing all run inside a containerized HTTP server managed by Vertex AI (opens in a new tab). A custom prediction routine extends a prebuilt serving image with custom handler code, while a fully custom container (opens in a new tab) replaces the entire server. Both approaches let the same image deserialize diverse request payloads, apply business logic, and format prediction responses to prevent training-serving skew.

1.1Custom Prediction Routines Compared With Fully Custom Containers

Custom prediction routines and fully custom containers both run on Vertex AI endpoints, but they divide ownership of the HTTP server differently. A custom prediction routine keeps the prebuilt server and adds a custom predictor handler for preprocessing and postprocessing, while a fully custom container requires the team to build and operate the entire HTTP server.

Custom code in either option sits on the request path between Vertex AI and the model. Traffic flows from the Vertex AI endpoint to the container, the handler transforms the payload, the model produces raw outputs, and the handler transforms those outputs back into the client response.

1.2HTTP Server Contract for Serving Containers

The serving container contract requires the Docker image to run an HTTP server indefinitely when the container starts. The serving container contract applies to any language or framework, as long as the server listens and responds to liveness checks, health checks, and inference requests from Vertex AI. Common ways to meet the serving container contract include:

• Flask for a lightweight custom Python server with explicit preprocessing and postprocessing functions
• TensorFlow Serving for high-performance TensorFlow model serving with a built-in HTTP interface
• TorchServe for PyTorch models that need model versioning and batch handling
• KServe Python Server for a Python predictor interface that already separates model load, preprocess, predict, and postprocess steps

The start command for the server comes from the image Dockerfile ENTRYPOINT, the Dockerfile CMD, or both together. Vertex AI can override that startup behavior without rebuilding the image by setting containerSpec.command to replace ENTRYPOINT and containerSpec.args to replace CMD. However the command is set, the resulting process must run indefinitely so the container stays alive to accept checks and predictions.

1.3Liveness Checks, Health Checks, and Inference Requests

Liveness checks, health checks, and inference requests define the serving lifecycle that the container HTTP server must support. Liveness indicates whether the container process is running, health indicates whether the model and dependencies are ready to serve, and inference carries the actual prediction workload. The interaction follows an ordered sequence:

1. Container start: The startup command launches and the HTTP server begins listening; the process stays in a running state indefinitely.
2. Liveness checks: Vertex AI sends repeated liveness probes to the container, and the server must respond so the platform knows not to restart the container.
3. Health checks: Only after the server is live, Vertex AI sends health probes to verify that model artifacts are loaded and preprocessing dependencies are ready; the container moves from a starting state to a healthy state.
4. Inference requests: Only after the container is healthy, Vertex AI forwards client prediction traffic to the container, where the handler deserializes the body, runs preprocessing, invokes the model, runs postprocessing, and returns a formatted response.

This ordering protects latency and reliability because no prediction traffic flows to a container until its server, model, and custom logic have all reported readiness.

1.4Image Publishing, Permissions, and Model Artifact Access

Image publishing, permissions, and model artifact access control whether Vertex AI can deploy and run the serving container. The serving container image must be published in a registry location that Vertex AI can reach, with permissions that allow the Vertex AI service to pull the image at deploy time. The Model resource references that image URI, so the deployment depends on both the image location and the grant that authorizes access. Model artifacts are separate from the container image and are accessed by the running container during startup and inference. The container loads those artifacts into memory during initialization, then applies them to every inference request without reloading them per request.

1.5Environment Variables Available to Inference Logic

Environment variables available to the inference container configure ports, paths, and custom behavior without changing handler code. The inference container merges configuration from three distinct owners, which lets platform routing, image defaults, and per-model settings compose in one runtime.

• Variables set in the container image: Defaults baked into the Dockerfile define the server behavior, library paths, and preprocessing thresholds that ship with the image.
• Variables set by Vertex AI: Platform-injected variables carry deployment context from Vertex AI to the container, such as where to listen and where model artifacts are staged.
• Variables set in the Model resource: Per-model variables defined on the Model resource customize one deployment of the same image, such as selecting a model variant or tuning postprocessing logic.

Custom predictor handlers read these variables at startup to locate artifacts and at request time to control deserialization and response formatting, so the same image can serve different models with different preprocessing rules.

In-graph and exported artifact transformations embed preprocessing and postprocessing directly into the served model artifact so Vertex AI (opens in a new tab) receives raw inference payloads and returns finished predictions. The exported artifact flows from the training process to Cloud Storage (opens in a new tab) to a Vertex AI Model resource to an endpoint, and the serving container loads the artifact and executes the embedded logic on each request. This design eliminates training-serving skew because the same transformation code runs in training and serving, and it reduces latency by avoiding a separate preprocessing service hop.

2.1TensorFlow SavedModel Artifacts with Server-Side Batching

A TensorFlow SavedModel artifact packages the computation graph, including any in-graph preprocessing and postprocessing ops, as saved_model.pb for serving by a TensorFlow prebuilt container. The TensorFlow SavedModel artifact lives in the Cloud Storage directory that the Model resource points to, and TensorFlow Serving loads that graph and applies the embedded transformations before and after the core model computation. To enable server-side request batching, include config/batching_parameters_config in the same Cloud Storage directory as saved_model.pb and configure it according to TensorFlow Serving batching configuration. Batching lets the server combine individual inference requests from clients into efficient graph executions, then split the results back to each request destination.

2.2PyTorch Model Archives with Handler-Defined Transformations

A PyTorch model served by the prebuilt container must be packaged as a model archive file named model.mar that contains model weights, supporting artifacts, and a handler script. The PyTorch handler script, either default or custom, defines how the model processes inference requests, with data flowing from the client payload to the handler preprocessing function to the PyTorch model forward pass to the handler postprocessing function and back to the client response. The archive must be created with model-name model so the prebuilt TorchServe image finds model.mar at load time. For advanced workflows that need custom libraries, multi-model logic, or non-standard raw formats, use custom containers instead of relying on the prebuilt PyTorch serving image.

2.3XGBoost and scikit-learn Artifacts with Embedded Pipelines

XGBoost and scikit-learn artifacts embed transformations by exporting a single file whose filename must exactly match what the prebuilt inference container expects. An XGBoost model trained in a prebuilt container is exported either with xgboost.Booster save_model as model.bst or with the joblib library as model.joblib, while a scikit-learn model is exported with joblib as model.joblib. A scikit-learn Pipeline object is the in-artifact mechanism for skew prevention because the pipeline chains transformers and the final estimator into one object, so raw features flow from the request through each transformer step in order to the estimator without separate serving code. Export follows the same ordered lifecycle for both frameworks:

1. Train in the managed environment: Train the booster, estimator, or pipeline using training data.
2. Save locally: Write model.bst or model.joblib to the local filesystem, which does not persist beyond the training worker.
3. Upload to Cloud Storage: Copy the file to the directory indicated by the AIP_MODEL_DIR environment variable using the Cloud Storage client, so the artifact flows from the ephemeral trainer to durable storage for deployment.

To serve inferences, create a Model with the prebuilt container for inference matching the XGBoost or scikit-learn version used for training, because the serving container depends on that version to deserialize the artifact.

2.4Custom Containers for Advanced Preprocessing Logic

A custom container is a Docker image that runs an HTTP server and is used when preprocessing or postprocessing cannot be fully embedded in saved_model.pb, model.mar, model.bst, or model.joblib. The custom container must listen and respond to liveness checks, health checks, and inference requests after Vertex AI starts it, and the server can be implemented with Flask, TensorFlow Serving, TorchServe, or KServe Python Server. The container loads model artifacts from Cloud Storage at startup, then on each request data flows from the endpoint to the custom HTTP handler to the model and back through postprocessing to the response. Choose a custom container when raw payload parsing, external feature lookups, or complex output formatting must live in serving code outside the exported graph or pipeline.

2.5Exported Artifact Requirements Comparison

Exported artifact requirements define the filename contract, the transformation carrier, and the serving dependency for each framework. The choice of artifact determines how transformations travel with the model and what the serving container expects at load time.

The serving container depends on an exact filename match and, for XGBoost and scikit-learn, on a container version match with training, so a mismatched name or version prevents the Model from loading.

Architecting low-latency feature enrichment and stream transformations means combining a real-time event stream with precomputed features for inference without adding lookup delay. In this architecture, Cloud Pub/Sub (Pub/Sub) (opens in a new tab) carries the live transaction stream, Dataflow runs an Apache Beam (Beam) streaming pipeline, Vertex AI Feature Store (Feature Store) (opens in a new tab) supplies low-latency feature values, and Vertex AI endpoints return predictions. The pipeline enriches each event in flight and then scores it, so preprocessing, inference, and post-prediction handling stay in one scalable streaming path.

3.1Vertex AI Feature Store Online Serving

Vertex AI Feature Store is the central system for creating, maintaining, sharing, and serving machine learning features for both training and inference. Feature data lives in BigQuery as a table or view, which retains all data including historical feature values for offline work such as training and batch analysis. To serve online, you must first create an online store instance and then define feature views that map entity keys, such as user_id, to their precomputed features. Because the same BigQuery-backed definitions feed both offline training and the online store, the serving path reuses the same feature logic and reduces training-serving skew.

Online serving is optimized for workloads that need low latency during streaming inference. As of Apache Beam SDK version 2.55.0, online feature serving is supported through Bigtable online serving and the Vertex AI Feature Store legacy method, with Bigtable online serving used for the enrichment pattern.

3.2Streaming Ingest With Cloud Pub/Sub

Cloud Pub/Sub is the ingestion layer that delivers the unbounded transaction stream into the Beam pipeline running on Dataflow. Each Pub/Sub message carries only the live event fields, for example product_id, user_id, and sale_price in an ecommerce recommendation case. The full model input requires additional context such as age, gender, state, and country, which are not in the event and must be joined at inference time. This split keeps the event small and fast while the pipeline fetches the remaining features by key during enrichment.

3.3Stream Enrichment With the Apache Beam Enrichment Transform

The Apache Beam enrichment transform is a turnkey element-wise transform that enriches each streaming element by key-value lookup against Vertex AI Feature Store. Data flows from the Pub/Sub source to the enrichment step, the enrichment step calls Feature Store with the entity key from the event, and the returned feature values flow forward joined to the original event. The enriched record then becomes the complete inference request, so preprocessing happens inline without a separate batch join. The enrichment transform provides three peer capabilities for production streaming:

• Built-in Vertex AI handler: the handler interacts directly with Vertex AI Feature Store to fetch precomputed feature values for each key.
• Client-side throttling: the transform limits its own request rate to manage rate limiting against the online store under burst load.
• Optional Redis cache: a configured Redis cache serves repeated key lookups locally to improve efficiency and lower online-store load.

3.4Streaming Inference and Post-Prediction Processing

Streaming inference is performed by the Beam RunInference transform using the VertexAIModelHandlerJSON interface to call a deployed Vertex AI model. The VertexAIModelHandlerJSON interface holds the connection from the pipeline to the Vertex AI endpoint, formats each enriched record as JSON, and routes the prediction request to the deployed model. The model returns a prediction, such as a product recommendation, and that prediction flows to downstream post-prediction steps such as writing to a sink, publishing to another Pub/Sub topic, or triggering an action. The end-to-end streaming flow runs in strict order:

1. Publish events: producers publish transaction messages with product_id, user_id, and sale_price to Pub/Sub.
2. Read stream: the Dataflow pipeline with Apache Beam reads the unbounded Pub/Sub stream element by element.
3. Enrich elements: the enrichment transform looks up user_id in Vertex AI Feature Store and joins age, gender, state, and country onto each event.
4. Run inference: the RunInference transform sends the enriched JSON record through VertexAIModelHandlerJSON to the Vertex AI endpoint for online prediction.
5. Handle predictions: the pipeline forwards each prediction for post-prediction processing, such as storage, alerting, or downstream serving.

This order matters because inference depends on enrichment, and enrichment depends on both the live key from Pub/Sub and the pre-provisioned online store and feature views.

Key Facts

• Traffic flows from the Vertex AI endpoint to the serving container, where the custom handler deserializes the payload, runs preprocessing, invokes the model, and formats the response.
• Vertex AI forwards inference traffic only after a container passes liveness checks and health checks, and containerSpec.command replaces Dockerfile ENTRYPOINT while containerSpec.args replaces Dockerfile CMD.
• TensorFlow serving uses saved_model.pb with optional config/batching_parameters_config, PyTorch uses model.mar with model-name model, XGBoost uses model.bst or model.joblib, and scikit-learn uses model.joblib.
• Streaming inference on Dataflow reads Cloud Pub/Sub events, enriches each event with a Vertex AI Feature Store lookup by entity key, scores the enriched record with RunInference through VertexAIModelHandlerJSON, and forwards predictions for post-processing.