Sunday, December 3, 2023
No menu items!
HomeCloud ComputingBuilding a Machine Learning Platform with Kubeflow and Ray on Google Kubernetes...

Building a Machine Learning Platform with Kubeflow and Ray on Google Kubernetes Engine

Increasingly more enterprises adopt Machine Learning (ML) capabilities to enhance their services, products, and operations. As their ML capabilities mature, they build centralized ML Platforms to serve many teams and users across their organization. Machine learning is inherently an experimental process requiring repeated iterations. An ML Platform standardizes the model development and deployment workflow to offer greater consistency for the repeated process. This facilitates productivity and reduces time from prototype to production.

When first trying ML in the cloud, many practitioners will start with fully managed ML platforms like Google Cloud’s Vertex AI. Fully-managed platforms abstract out many complexities to simplify the end-to-end workflow. However, like with most decisions, there are tradeoffs. Organizations may choose to build their own custom, self-managed ML platform for various reasons such as control and flexibility. Building your own platform gives you more control over your resources. You can implement unique resource utilization constraints, access permissions, and infrastructure strategies that fit your organization’s specific needs. You also get more flexibility over tools and frameworks. Since the system is completely open, you can integrate any ML tools that you already are using. And lastly, these benefits help avoid vendor lock-in because cloud-native platforms are by definition portable across cloud providers.

For self-managed ML Platforms, Open Source Software is an important driver of digital innovation. If you are following the evolution of ML technologies, then you are probably aware of the ever growing ecosystem of Open Source Machine Learning frameworks, platforms, and tools. However, no single Open Source library delivers a complete ML solution, so we must integrate multiple Open Source projects to build an ML platform.

To start building an ML Platform, it should support the basic ML user journey of notebook prototyping to scaled training to online serving. For organizations with multiple teams, it additionally needs to support administrative requirements of multi-user support with identity-based authentication and authorization. Two popular Open Source projects – Kubeflow and Ray – together can support these needs. Kubeflow provides the multi-user environment and interactive notebook management. Ray orchestrates distributed computing workloads across the entire ML lifecycle, including training and serving.

Google Kubernetes Engine (GKE) simplifies deploying Open Source ML software in the cloud with autoscaling and auto-provisioning. GKE reduces the effort to deploy and manage the underlying infrastructure at scale and offers the flexibility to use your ML frameworks of choice. In this article, we will show how Kubeflow and Ray can be  assembled into a seamless experience. We will demonstrate how platform builders can deploy them both to GKE to provide a comprehensive, production-ready ML platform.

Kubeflow and Ray

First, let’s take a closer look at these two Open Source projects. While both Kubeflow and Ray deal with the problem of enabling ML at scale, they focus on very different aspects of the puzzle.

Kubeflow is a Kubernetes-native ML platform aimed at simplifying the build-train-deploy lifecycle of ML models. As such, its focus is on general MLOps. Some of the unique features offered by Kubeflow include:

Built-in integration with Jupyter notebooks for prototyping

Multi-user isolation support

Workflow orchestration with Kubeflow Pipelines

Identity-based authentication and authorization through Istio Integration

Out-of-the-box integration with major cloud providers such as GCP, Azure, and AWS

Ray is a general-purpose distributed computing framework with a rich set of libraries for large scale data processing, model training, reinforcement learning, and model serving. It is popular with customers as a simple API for building and scaling AI and Python workloads. Its focus is on the application itself – allowing users to build distributed computing software with a unified and flexible set of APIs. Some of the advanced libraries offered by Ray include:

RLLib for reinforcement learning

Ray Tune for hyperparameter tuning

Ray Train for distributed deep learning

Ray Serve for scalable model serving

Ray Data for preprocessing

It should be noted that Ray is not a Kubernetes-native project. In order to deploy Ray on Kubernetes, the Open Source community has created KubeRay, which is exactly what it sounds like – a toolkit for deploying Ray in Kubernetes. KubeRay offers a powerful set of tools that include many great features, like custom resource APIs and a scalable operator. You can learn more about it here.

Now that we have examined the differences between Kubeflow and Ray, you might be asking which is the right platform for your organization. Kubeflow’s MLOps capabilities and Ray’s distributed computing libraries are both independently useful with different advantages. What if we can combine the benefits of both systems? Imagine having an environment that:

Supports Ray Train with autoscaling and resource provisioning

Integrated with identity-based authentication and authorization

Supports multi-user isolation and collaboration

Contains an interactive notebook server

Let’s now take a look at how we can put these two platforms together and take advantage of the  useful features offered by each. Specifically, we will deploy Kuberay in a GKE cluster installed with Kubeflow. The system looks something like this:

In this system, the Kubernetes cluster is partitioned into logically-isolated workspaces, called “profiles”. Each new user will create their own profile, which is a container for all their resources in this Kubernetes cluster. The user can then provision their own resources within their designated namespace, including Ray Clusters and Jupyter Notebooks. If the user’s resources are provisioned through the Kubeflow dashboard, then Kubeflow will automatically place these resources in their profile namespace.

Under this setup, each Ray cluster is by default protected by role-based access control policies (with Istio) preventing unauthorized access. This allows each user to interact with their own Ray clusters independently of each other, and allows them to share Ray clusters with other team members.

For this setup, I used the following versions:

Google Kubernetes Engine 1.21.12-gke.2200 

Kubeflow 1.5.0

Kuberay 0.3.0

Python 3.7

Ray 1.13.1

The configuration files used for this deployment can be found here.

Deploying Kubeflow and Kuberay

For deploying Kubeflow, we will be using the GCP instructions here. For simplicity purposes, I have used mostly default configuration settings. You can freely experiment with customizations before deploying, for example, you can enable GPU nodes in your cluster by following these instructions.

Deploying the KubeRay operator is pretty straightforward. We will be using the latest released version:

code_block[StructValue([(u’code’, u’export KUBERAY_VERSION=v0.3.0rnkubectl create -k “${KUBERAY_VERSION}”rnkubectl apply -k “${KUBERAY_VERSION}”‘), (u’language’, u”), (u’caption’, <wagtail.wagtailcore.rich_text.RichText object at 0x3e086da8cd90>)])]

This will deploy the KubeRay operator in the “ray-systems” namespace in your cluster.

Creating Your Kubeflow User Profile

Before you can deploy and use resources in Kubeflow, you need to first create your user profile. If you follow the GKE installation instructions, you should be able to navigate to https://[cluster].endpoints.[project] in your browser, where [cluster] is the name of your GKE cluster and [project] is your GCP project name.

This should redirect you to a web page where you can use your GCP credentials to authenticate yourself.

Follow the dialogue, and Kubeflow will create a namespace with you as the administrator. We’ll discuss later in this article how to invite others to your workspace.

Build the Ray Worker Image

Next, let’s build the image we’ll be using for the Ray cluster. Ray is very sensitive when it comes to version compatibility (for example, the head and worker nodes must use the same versions of Ray and Python), so it is highly recommended to prepare and version-control your own worker images. Look for the base image you want from their Docker page here: rayproject/ray – Docker Image

The following is a functioning worker image using Ray 1.13 and Python 3.7:

code_block[StructValue([(u’code’, u’FROM rayproject/ray:1.13.1-py37rnrnRUN pip install numpy tensorflowrnrnCMD [“bin/bash”]’), (u’language’, u”), (u’caption’, <wagtail.wagtailcore.rich_text.RichText object at 0x3e088f6f98d0>)])]

Here is the same Dockerfile for a worker image running on GPUs if you prefer GPUs instead of CPUs:

code_block[StructValue([(u’code’, u’FROM rayproject/ray:1.13.1-py37-gpurnrnRUN pip install numpy tensorflowrnrnCMD [“bin/bash”]’), (u’language’, u”), (u’caption’, <wagtail.wagtailcore.rich_text.RichText object at 0x3e088e75b9d0>)])]

Use Docker to build and push both images to your image repository:

code_block[StructValue([(u’code’, u’$ docker build -t <path-to-your-image> -f Dockerfile .rn$ docker push <path-to-your-image>’), (u’language’, u”), (u’caption’, <wagtail.wagtailcore.rich_text.RichText object at 0x3e088e75b790>)])]

Build the Jupyter Notebook Image

Similarly we need to build the notebook image that we are going to use. Because we are going to use this notebook to interact with the Ray cluster, we need to ensure that it uses the same version of Ray and Python as the Ray workers.

The Kubeflow example Jupyter notebooks can be found at Example Notebook Servers. For this example, I changed the PYTHON_VERSION in components/example-notebook-servers/jupyter/Dockerfile to the following:

code_block[StructValue([(u’code’, u’ARG MINIFORGE_VERSION=4.10.1-4rnARG PIP_VERSION=21.1.2rnARG PYTHON_VERSION=3.7.10′), (u’language’, u”), (u’caption’, <wagtail.wagtailcore.rich_text.RichText object at 0x3e088e75b110>)])]

Use Docker to build and push the notebook image to your image repository, similar to the previous step:

code_block[StructValue([(u’code’, u’$ docker build -t <path-to-your-image> -f Dockerfile .rn$ docker push <path-to-your-image>’), (u’language’, u”), (u’caption’, <wagtail.wagtailcore.rich_text.RichText object at 0x3e086da8c450>)])]

Deploy a Ray Cluster

Now we are ready to configure and deploy our Ray cluster.

1. Copy the following sample yaml file from GitHub:

code_block[StructValue([(u’code’, u’curl -o ray-cluster.serve.yaml’), (u’language’, u”), (u’caption’, <wagtail.wagtailcore.rich_text.RichText object at 0x3e086db809d0>)])]

2. Edit the settings in the file:

a. For the user namespace, change the value to match with your Kubeflow profile name:

code_block[StructValue([(u’code’, u’namespace: %your_name%’), (u’language’, u”), (u’caption’, <wagtail.wagtailcore.rich_text.RichText object at 0x3e086e369850>)])]

b. For the Ray head and worker settings, change the value to point to the image you have built previously:

code_block[StructValue([(u’code’, u’image: %your_image%’), (u’language’, u”), (u’caption’, <wagtail.wagtailcore.rich_text.RichText object at 0x3e086da8c550>)])]

c. Edit resource requests and limits, as required. For example, you can change the CPU or GPU requirements for worker nodes here:

code_block[StructValue([(u’code’, u’resources: rn limits:rn cpu: 1rn requests:rn cpu: 200m’), (u’language’, u”), (u’caption’, <wagtail.wagtailcore.rich_text.RichText object at 0x3e086da8c0d0>)])]

3. Deploy the cluster:

code_block[StructValue([(u’code’, u’kubectl apply -f raycluster.serve.yaml’), (u’language’, u”), (u’caption’, <wagtail.wagtailcore.rich_text.RichText object at 0x3e088d5f6a10>)])]

4. Your cluster should be ready to go momentarily. If you have enabled node auto-provisioning on your GKE cluster, you should be able to see the cluster dynamically scale up and down according to usage. You can check the status of your cluster by doing:

code_block[StructValue([(u’code’, u’$ kubectl get pods -n <user name>rnNAME READY STATUS RESTARTS AGErnexample-cluster-head-8cbwb 1/1 Running 0 12srnexample-cluster-worker-large-group-75lsr 1/1 Running 0 12srnexample-cluster-worker-large-group-jqvtp 1/1 Running 0 11srnexample-cluster-worker-large-group-t7t4n 1/1 Running 0 12s’), (u’language’, u”), (u’caption’, <wagtail.wagtailcore.rich_text.RichText object at 0x3e088d5f6490>)])]

You can also verify that the service endpoints are created:

code_block[StructValue([(u’code’, u’$ kubectl get services -n <user name>rnNAME TYPE CLUSTER-IP EXTERNAL-IP PORT(S) AGErnexample-cluster-head-svc ClusterIP <none> 8265/TCP,10001/TCP,8000/TCP,6379/TCP 18s’), (u’language’, u”), (u’caption’, <wagtail.wagtailcore.rich_text.RichText object at 0x3e088d5f6590>)])]

Remember this service name – we will come back to it later.

Now our ML Platform is all set up and we are ready to start Training a model.

Training a ML Model

We are going to use a Notebook to orchestrate our model training. We can access Ray from a Jupyter notebook session.

1. In the Kubeflow dashboard, navigate to the “Notebooks” tab.

2. Click on “New Notebook”.

3. In the “Image” section, click on “Custom Image”, and input the path to the Jupyter notebook image that you have built here.

4. Configure resource requirements for the notebook as needed. The default notebook uses half a CPU and 1G of memory. Note that these resources are only for the notebook session, and not for the Training resources. Later, we use Ray to orchestrated resources at scale on GKE.

5. Click on “LAUNCH”.

6. When the notebook finishes deploying, click on “Connect” to start a new notebook session.

7. Inside the notebook, open a terminal by clicking on File -> New -> Terminal. 

8. Install Ray 1.13 in the terminal:

code_block[StructValue([(u’code’, u’pip install ray==1.13′), (u’language’, u”), (u’caption’, <wagtail.wagtailcore.rich_text.RichText object at 0x3e086da8c290>)])]

9. Now you are ready to run an actual Ray application, using this notebook and the Ray cluster you just deployed in the previous section. I have made a .ipynb file using the canonical Ray trainer example here.

10. Run through the cells in the notebook. The magic line that connects to the Ray cluster is:

code_block[StructValue([(u’code’, u’ray.init(“ray://example-cluster-head-svc:10001″)’), (u’language’, u”), (u’caption’, <wagtail.wagtailcore.rich_text.RichText object at 0x3e085eda2b50>)])]

This should match with the service endpoint that you created earlier. If you have several different Ray clusters, you can simply change the endpoint here to connect to a different one.

11. The next few lines will start a Ray Trainer process on the cluster:

code_block[StructValue([(u’code’, u’trainer = Trainer(backend=”tensorflow”, num_workers=4)rntrainer.start()rnresults =’), (u’language’, u”), (u’caption’, <wagtail.wagtailcore.rich_text.RichText object at 0x3e085eda2050>)])]

Note here that we specify 4 workers, which matches with our Ray cluster’s number of replicas. If we change this number, the Ray cluster will automatically scale up or down according to resource demands.

Serving a ML Model

In this section we will look at how we can serve the machine learning model that we have just trained in the last section.

1. Using the same notebook, wait for the training steps to complete. You should see some output logs with metrics for the model that we have trained.

2 Run the next cell:

code_block[StructValue([(u’code’, u’serve.start(detached=True, http_options={“host”: “”})rnTFMnistModel.deploy(TRAINED_MODEL_PATH)’), (u’language’, u”), (u’caption’, <wagtail.wagtailcore.rich_text.RichText object at 0x3e085eda2950>)])]

This will start serving the model that we have just trained, using the same service endpoint we created before.

3. To verify that the inference endpoint is now working, we can create a new notebook. You can use this one here.

4. Note that we are calling the same inference endpoint as before, but using a different port:

code_block[StructValue([(u’code’, u’resp = requests.get(rn “http://example-cluster-head-svc:8000/mnist”,rn json={“array”: np.random.randn(28 * 28).tolist()})’), (u’language’, u”), (u’caption’, <wagtail.wagtailcore.rich_text.RichText object at 0x3e088c4398d0>)])]

5. You should see the inference results displayed in your notebook session.

Sharing the Ray Cluster with Others

Now that you have a functional workspace with an interactive notebook and a Ray cluster, let’s invite others to collaborate.

1. On Cloud Console, grant the user minimal cluster access here.

2. In the left-hand panel of the Kubeflow dashboard, select “Manage Contributors”.

3. In the “Contributors to your namespace” section, enter the email address of the user to whom you are granting access. Press enter.

4. That user can now select your namespace and access your notebooks, including your Ray cluster.

Using Ray Dashboard

Finally, you can also bring up the Ray Dashboard using Istio virtual services. Using these steps, you can bring up a dashboard UI inside the Kubeflow central dashboard console:

1. Create an Istio Virtual Service config file:

code_block[StructValue([(u’code’, u”apiVersion: VirtualServicernmetadata:rn name: example-cluster-virtual-servicern Namespace: kubeflowrnspec:rn gateways:rn – kubeflow-gatewayrn hosts:rn – ‘*’rn http:rn – match:rn – uri:rn prefix: /example-cluster/rn rewrite:rn uri: /rn route:rn – destination:rn host: example-cluster-head-svc.$(USER_NAMESPACE).svc.localrn port:rn number: 8265″), (u’language’, u”), (u’caption’, <wagtail.wagtailcore.rich_text.RichText object at 0x3e088c439910>)])]

Replace $(USER_NAMESPACE) with the namespace of your user profile. Save this to a local file.

 2. Deploy the virtual service:

code_block[StructValue([(u’code’, u’kubectl apply -f virtual_service.yaml’), (u’language’, u”), (u’caption’, <wagtail.wagtailcore.rich_text.RichText object at 0x3e088c439bd0>)])]

3. In your browser window, navigate to https://<host>/_/example-cluster/. The Ray dashboard should be displayed in the window:


Let’s take a minute to recap what we have done. In this article, we have demonstrated how to deploy two popular ML frameworks, Kubeflow and Ray, in the same GCP Kubernetes cluster. The setup also takes advantage of GCP features like IAP (Identity-Aware Proxy) for user authentication, which protects your applications while simplifying the experience for cloud admins. The end result is a well-integrated and production-ready system that pulls in useful features offered by each system:

Orchestrating distributed computing workloads using Ray APIs;

Multi-user isolation using Kubeflow;

Interactive notebook environment using Kubeflow notebooks;

Cluster autoscaling and auto-provisioning using Google Kubernetes Engine

We’ve only scratched the surface of the possibilities, and you can expand from here:

Integrations with other MLOps offerings, such as Vertex Model monitoring;

Faster and safer image storage and management, through the Artifact Repository;

High throughput storage for unstructured data using GCSFuse;

Improve network throughput for collective communication with NCCL Fast Socket.

We look forward to the growth of your ML Platform and how your team innovates with Machine Learning. Look out for future articles on how to enable additional ML Platform features.

Related Article

Enabling real-time AI with Streaming Ingestion in Vertex AI

Many machine learning (ML) use cases, like fraud detection, ad targeting, and recommendation engines, require near real-time predictions….

Read Article

Cloud BlogRead More



Please enter your comment!
Please enter your name here

Most Popular

Recent Comments