16.2 Ray on Agent Platform for Distributed ML Workloads
Key Takeaways
- Ray on Agent Platform runs open-source Ray code on managed clusters with minimal changes and integrates with BigQuery and Agent Platform online inference.
- Ray clusters on Agent Platform stay available until deleted, which helps repeated or many short Ray jobs but bills for the whole time the cluster exists.
- Ray worker pools autoscale between min_replica_count and max_replica_count, and Google suggests new Ray users start with manual scaling.
- The Ray head node's logical CPU count is set to 0 so workloads don't run on the head node.
- Google recommends Agent Platform online inference over Ray Serve for most serving, because it's fully managed and adds Model Monitoring, traffic splitting, Model Registry, and Pipelines integration.
The exam guide lists Ray on Gemini Enterprise Agent Platform (formerly Ray on Vertex AI) as a way to build and orchestrate ML workloads. Ray is an open-source framework for scaling Python and AI applications across many machines. It includes libraries for data processing (Ray Data), distributed training (Ray Train), hyperparameter search (Ray Tune), and serving (Ray Serve).
When Ray on Agent Platform Fits
| Situation | Why Ray |
|---|---|
| The team already uses Ray code or libraries | Run the same open-source Ray code on Google Cloud with minimal changes |
| Python-native distributed processing and training in one framework | Ray Data preprocessing feeds Ray Train, instead of stitching Spark and separate training jobs together |
| Many parallel Python tasks (simulations, batch feature computation, embedding generation) | Ray tasks and actors scale arbitrary Python functions |
| Hyperparameter search inside the same cluster | Ray Tune runs many trials on shared cluster resources |
| Spark on Ray, or mixed CPU and GPU workloads | Flexible heterogeneous worker pools |
If the team doesn't use Ray and just needs managed training, Agent Platform custom training is simpler (Chapter 9).
Workflow
- Set up: install the Agent Platform SDK with Ray support (
pip install "google-cloud-aiplatform[ray]"), and optionally configure private networking. - Create a cluster: in the console or with
vertex_ray.create_ray_cluster(), specifying the head node type and worker node types (machine type, accelerators, counts). This needs the Agent Platform Administrator role. - Develop: connect through the Ray Client (
ray.init("vertex_ray://<cluster>")) from Colab Enterprise (a console button opens a connected notebook) or any Python environment, or submit scripts with the Ray Jobs API. The client's Ray version must match the cluster's. This needs the Agent Platform User role. - Use BigQuery: read and write with
vertex_ray.data.read_bigquery()inside Ray tasks. The maximum query response size is 10 GB. - Deploy models: export the model from the Ray checkpoint, upload to Model Registry, deploy to an endpoint, or run batch inference.
- Monitor: logs in Cloud Logging, metrics in Cloud Monitoring, and the Ray Dashboard.
- Delete the cluster when finished to stop billing.
Cluster behavior
- Ray clusters remain available until deleted. Unlike custom jobs, which release resources when they finish, they give capacity assurance and reuse of cached data and images for repeated jobs, or for many jobs shorter than their startup time.
- Following open-source Ray best practice, the head node's logical CPU count is 0, so no workloads run on it.
- Networking can be public or private (VPC peering or a Private Service Connect interface, which Google recommends for private networks).
Scaling Ray Clusters
| Mode | How | Guidance |
|---|---|---|
| Autoscaling | Set min_replica_count and max_replica_count on every worker pool through AutoscalingSpec | Lets the cluster add workers for resource-hungry tasks and cut cost when idle. Adds node launch overhead. Custom scale-up and scale-down speeds aren't supported |
| Manual scaling | Change worker counts yourself | Google suggests starting here if you're new to Ray. Manual scaling has a limitation with VPC peering, so use a PSC interface for private networks |
autoscaling_spec = AutoscalingSpec(min_replica_count=1, max_replica_count=3)
head = Resources(machine_type="n1-standard-16", node_count=1)
workers = [Resources(machine_type="n1-standard-16",
accelerator_type="NVIDIA_TESLA_T4", accelerator_count=1,
autoscaling_spec=autoscaling_spec)]
cluster = vertex_ray.create_ray_cluster(head_node_type=head, worker_node_types=workers, ...)
Serving Ray-Trained Models
Google's comparison recommends Agent Platform online inference for most cases:
| Aspect | Agent Platform online inference (recommended) | Ray Serve |
|---|---|---|
| Management | Fully managed, less operational work | More setup and management |
| Features | Model Monitoring, A/B testing and traffic splitting, Model Registry and Pipelines integration, explainability | Advanced model composition, ensembles, custom inference logic in the Ray ecosystem |
| Formats | Prebuilt containers (TensorFlow, PyTorch, scikit-learn, XGBoost) or custom containers | TensorFlow, PyTorch, scikit-learn |
Choose Ray Serve only when complex model composition inside Ray outweighs the managed features.
Ray vs. Other Options
| Need | Best fit |
|---|---|
| Existing Ray code, Python-native data + training + tuning | Ray on Agent Platform |
| Existing Spark jobs | Managed Service for Apache Spark |
| Single training job with standard distributed strategies | Agent Platform custom training |
| Repeatable multi-step ML workflow with lineage | Agent Platform Pipelines (it can include steps that use Ray) |
| Cross-system data workflow scheduling | Managed Airflow (CreateRayClusterOperator is available) |
Cost and Operations Checklist
| Concern | Practice |
|---|---|
| Idle cost | Delete clusters after work finishes, or use autoscaling with a low minimum for worker pools |
| Right-sizing | Separate CPU and GPU worker pools so preprocessing doesn't hold GPUs |
| Reproducibility | Pin the Ray version in the cluster and client, and record the cluster configuration with experiment runs |
| Security | Use private networking (PSC interface) and least-privilege service accounts for data access |
| Orchestration | Create and delete clusters as pipeline or DAG steps (for example, CreateRayClusterOperator and DeleteRayClusterOperator in Airflow) |
Worked Scenario
A research team has Ray code that generates embeddings for 400 million documents with a GPU model, then runs Ray Tune over a ranking model. Data lives in BigQuery.
- Create a Ray on Agent Platform cluster with a CPU head node and autoscaling GPU worker pools.
- Read documents with
vertex_ray.data.read_bigquery, respecting the 10 GB query-response limit by reading tables or partitioned queries. - Run embedding generation with Ray Data, then Ray Tune for the ranking model.
- Export the best model, upload to Model Registry, and deploy to an Agent Platform endpoint.
- Delete the cluster when the work completes.
A company's data scientists already use Ray Data and Ray Tune on premises and want to move those workloads to Google Cloud with minimal code changes and BigQuery integration. Which option fits best?
A team finishes a Ray on Agent Platform experiment on Friday and leaves the cluster in place over the weekend. What is the main consequence?
A team trained a PyTorch model on Ray on Agent Platform and needs online inference with traffic splitting for canary releases and drift monitoring. What does Google recommend?