Persistent Volume Management for Stateful AI Applications

Question

Pulumi · Accepted Answer

For managing persistent volumes in stateful AI applications, Kubernetes is an excellent choice due to its native support for orchestration and management of containerized applications. In Kubernetes, a PersistentVolume (PV) is a piece of storage in the cluster that has been provisioned by an administrator or dynamically provisioned using Storage Classes. It is a resource in the cluster just like a node is a Kubernetes resource. PVs are volume plugins like Volumes, but have a lifecycle independent of any individual Pod that uses the PV.

Here’s why you would want to use them:
- **Persistent storage**: Ensures that your data survives Pod restarts and is not deleted when Pods are deleted.
- **Decoupling storage from the Pod lifecycle**: Allows the storage to persist independently of the lifecycle of the Kubernetes Pods or containers that use it.
- **Flexibility**: Supports a wide variety of storage backends, including network storage options and cloud-provided storage solutions.

When working within a stateful application context, a Kubernetes StatefulSet may be a more appropriate resource than a Deployment or Pod. It manages the deployment and scaling of a set of Pods, and provides guarantees about the ordering and uniqueness of these Pods. In contrast to a Deployment, a StatefulSet maintains a sticky identity for each of their Pods.

Let's write a simple program using Pulumi's Kubernetes SDK to create a persistent volume and a StatefulSet for a stateful AI application. Here's what we'll do:
1. Create a PersistentVolume (PV) that will represent the storage volume that persists data.
2. Create a StorageClass if you want to provision storage dynamically.
3. Deploy a StatefulSet with a volume claim template that ensures each Pod gets one PersistentVolume.

Below is the Pulumi program written in Python that accomplishes these tasks.

```python
import pulumi
import pulumi_kubernetes as k8s

# Create a Kubernetes Persistent Volume which is a cluster-wide resource that you can use to persist data.
persistent_volume = k8s.core.v1.PersistentVolume("persistent-volume",
    # The metadata allows you to name the persistent volume and add labels.
    metadata=k8s.meta.v1.ObjectMetaArgs(
        name="my-persistent-volume",
        labels={"type": "local"},
    ),
    # The spec defines the volume type, capacity, access modes, and various volume attributes.
    spec=k8s.core.v1.PersistentVolumeSpecArgs(
        capacity={"storage": "10Gi"},
        access_modes=["ReadWriteOnce"],
        persistent_volume_reclaim_policy="Retain",
        storage_class_name="manual",
        local=k8s.core.v1.LocalVolumeSourceArgs(
            path="/mnt/data",
        ),
        node_affinity=k8s.core.v1.VolumeNodeAffinityArgs(
            required=k8s.core.v1.NodeAffinityArgs(
                node_selector_terms=[k8s.core.v1.NodeSelectorTermArgs(
                    match_expressions=[k8s.core.v1.NodeSelectorRequirementArgs(
                        key="kubernetes.io/hostname",
                        operator="In",
                        values=["my-node"],
                    )],
                )],
            ),
        ),
    )
)

# Create a Kubernetes StorageClass for dynamic provisioning of persistent volumes.
storage_class = k8s.storage.v1.StorageClass("storage-class",
    metadata=k8s.meta.v1.ObjectMetaArgs(
        name="manual",
    ),
    provisioner="kubernetes.io/no-provisioner",
    volume_binding_mode="WaitForFirstConsumer",
)

# Deploy a StatefulSet which is suitable for running stateful applications that require stable unique network identifiers,
# stable persistent storage, and ordered graceful deployment and scaling.
stateful_set = k8s.apps.v1.StatefulSet("stateful-set",
    metadata=k8s.meta.v1.ObjectMetaArgs(
        name="my-stateful-set",
    ),
    spec=k8s.apps.v1.StatefulSetSpecArgs(
        selector=k8s.meta.v1.LabelSelectorArgs(
            match_labels={"app": "my-app"},
        ),
        # Template for the pods that will be created.
        template=k8s.core.v1.PodTemplateSpecArgs(
            metadata=k8s.meta.v1.ObjectMetaArgs(
                labels={"app": "my-app"},
            ),
            spec=k8s.core.v1.PodSpecArgs(
                containers=[
                    k8s.core.v1.ContainerArgs(
                        name="my-app",
                        image="my-app-image",
                        ports=[k8s.core.v1.ContainerPortArgs(
                            container_port=80,
                        )],
                        volume_mounts=[k8s.core.v1.VolumeMountArgs(
                            name="storage",
                            mount_path="/data",
                        )],
                    ),
                ],
            ),
        ),
        # Define a volume claim template which provides stable storage using Persistent Volumes provisioned by the above storage class.
        volume_claim_templates=[
            k8s.core.v1.PersistentVolumeClaimArgs(
                metadata=k8s.meta.v1.ObjectMetaArgs(
                    name="storage",
                ),
                spec=k8s.core.v1.PersistentVolumeClaimSpecArgs(
                    access_modes=["ReadWriteOnce"],
                    storage_class_name="manual",
                    resources=k8s.core.v1.ResourceRequirementsArgs(
                        requests={"storage": "10Gi"},
                    ),
                ),
            ),
        ],
    )
)

# Exports
pulumi.export('persistent_volume_name', persistent_volume.metadata.apply(lambda metadata: metadata.name))
pulumi.export('storage_class_name', storage_class.metadata.apply(lambda metadata: metadata.name))
pulumi.export('stateful_set_name', stateful_set.metadata.apply(lambda metadata: metadata.name))
```

Here’s what each section of the code does:
- **PersistentVolume**: Defines a persistent volume with a name, label, storage capacity, and access mode. It retains the stored data even if the using pod is deleted.
- **StorageClass**: Used to provision storage dynamically. Here we are using manual provisioning since we define the storage ourselves, but you can use cloud-provided solutions.
- **StatefulSet**: Manages the deployment and scaling of a set of Pods, and provides guarantees about the ordering and uniqueness of these Pods. It also uses a volume claim template to provide each Pod its own persistent volume via the StorageClass.

The exports at the end of the program will output the names of the created resources once the Pulumi program is deployed.

It’s important to adjust the `path` in the PersistentVolume to point to an existing directory on your nodes and also use an actual image instead of `"my-app-image"` when deploying your application. The node affinity is set to a hypothetical node labeled `'my-node'`, but you'll want to change that to match the actual node labels used in your cluster.