GPU-Based Machine Learning Workloads on Kubernetes

Question

Pulumi · Accepted Answer

To set up GPU-based machine learning workloads on a Kubernetes cluster, one approach is to ensure that your Kubernetes nodes are equipped with the necessary GPU resources and that your workloads are able to request and utilize these resources.

Here is how you can approach this task using Pulumi with Kubernetes and a cloud provider such as AWS, Azure, or GCP:

1. **Create a Kubernetes Cluster**: First, you need to provision a Kubernetes cluster on your chosen cloud provider with nodes that have GPUs attached. In cloud providers like AWS, this typically means selecting specific EC2 instance types such as the `p2` or `p3` families.

2. **Install GPU Drivers**: Once the cluster is set up, you need to ensure that the GPU drivers are installed on the nodes. This can be done by deploying a DaemonSet that runs on every node and installs the required drivers.

3. **Configure Kubernetes to Expose GPU Resources**: You'll need to set up Kubernetes so that it's aware of the GPU resources. Kubernetes does this through the use of Device Plugins.

4. **Create Workloads that Request GPU Resources**: Finally, when creating your workload definitions (e.g., Deployments, Jobs, or Pods), you must specify the GPU resource requirements so that Kubernetes schedules these workloads onto nodes with available GPU resources.

To illustrate this, we'll write a simple Pulumi program that sets up an AWS EKS cluster with GPU-enabled nodes:

```python
import pulumi
import pulumi_aws as aws
import pulumi_eks as eks
import pulumi_kubernetes as k8s

# Step 1: Provision an EKS Cluster with GPU-enabled node group.
# Create an EKS cluster with the desired configuration.
cluster = eks.Cluster('gpu-cluster')

# Create a node group with GPU instances, using an instance type that has GPUs.
gpu_node_group = eks.NodeGroup('gpu-node-group', 
    cluster=cluster.core,
    instance_type='p2.xlarge',
    desired_capacity=2,
    # Make sure to include labels that indicate the presence of GPUs
    labels={'accelerator': 'nvidia-tesla-k80'}
)

# Step 2: Configure the Kubernetes Provider to interact with the created cluster.
k8s_provider = k8s.Provider('k8s-provider', kubeconfig=cluster.kubeconfig)

# Step 3: Apply the NVIDIA device plugin as a DaemonSet to the cluster.
nvidia_daemonset = k8s.apps.v1.DaemonSet('nvidia-daemonset',
    metadata={'name': 'nvidia-device-plugin-daemonset'},
    spec={
        'selector': {'matchLabels': {'name': 'nvidia-device-plugin-ds'}},
        'template': {
            'metadata': {'labels': {'name': 'nvidia-device-plugin-ds'}},
            'spec': {
                'containers': [{
                    'name': 'nvidia-device-plugin-ctr',
                    'image': 'nvidia/k8s-device-plugin:1.0.0-beta',
                    'security_context': {'privileged': True},
                    'volume_mounts': [{'name': 'dev', 'mount_path': '/dev'}]
                }],
                'volumes': [{'name': 'dev', 'host_path': {'path': '/dev'}}]
            },
        }
    },
    opts=pulumi.ResourceOptions(provider=k8s_provider)
)

# Remember to replace 'YOUR_NAMESPACE' with the actual namespace and configure the rest of the Pod specification to match your application requirements.
# Step 4: Deploy a sample workload that requests GPU resources.
gpu_pod = k8s.core.v1.Pod('gpu-pod',
    metadata={
        'name': 'gpu-pod',
        'namespace': 'YOUR_NAMESPACE',
    },
    spec={
        'containers': [{
            'name': 'cuda-container',
            'image': 'nvidia/cuda:10.0-runtime',
            'resources': {
                'limits': {
                    'nvidia.com/gpu': 1  # This line requests one GPU.
                }
            }
        }]
    },
    opts=pulumi.ResourceOptions(provider=k8s_provider)
)

# Export the cluster name and kubeconfig
pulumi.export('cluster_name', cluster.eks_cluster.name)
pulumi.export('kubeconfig', cluster.kubeconfig)
```

This program does the following:

- Provisions an AWS EKS cluster with GPU-enabled nodes by specifying an instance type that includes GPUs.
- Utilizes the Pulumi EKS package to simplify the cluster creation process.
- Creates a Kubernetes Provider (`k8s.Provider`) that uses the generated `kubeconfig` to interact with the cluster.
- Deploys the NVIDIA device plugin as a `DaemonSet` using Pulumi's Kubernetes SDK. The device plugin is responsible for advertising the NVIDIA GPU devices to the Kubernetes scheduler.
- Defines a sample workload that requests a single GPU resource. Note that the GPU resource is requested under `resources.limits` in the Pod definition.
- Exports useful outputs such as the cluster name and `kubeconfig` which can be used to access the cluster with `kubectl` or other Kubernetes tools.

Please make sure to replace `'YOUR_NAMESPACE'` with the actual namespace where you want your pods to be scheduled.

This is a starting point, and you may need to tailor the program to fit the specific requirements of your machine learning workload, such as installing additional ML libraries or frameworks.