Encrypting gRPC Communication for Distributed AI Training on Kubernetes

Question

Pulumi · Accepted Answer

To set up encrypted gRPC communication for distributed AI training on Kubernetes, we need to create various Kubernetes resources. We'll be using a CertificateSigningRequest (CSR) to request a certificate for securing the gRPC communication, and a Kubernetes Service to expose our AI training application deployed on Kubernetes pods.

The process will involve the following steps:

Create a CertificateSigningRequest resource to obtain a signed certificate from the Kubernetes cluster Certificate Authority (CA). This certificate will be used to encrypt the gRPC communication.
Once the CSR is approved and the certificate is issued, we need to store the certificate in a Kubernetes Secret.
Configure our AI training application's deployment to mount the Secret containing the TLS certificate, so that the application can use it for encrypted communication.
Create a Kubernetes Service of type ClusterIP that routes traffic to our application pods and uses the TLS certificate for encryption.

Below is the Pulumi program written in Python that implements the above steps:

import pulumi
import pulumi_kubernetes as k8s

# Step 1: Create the CertificateSigningRequest to secure gRPC communication
csr_name = "ai-grpc-csr"
# PEM encoded CSR
pem_encoded_csr = '<<PEM_ENCODED_CSR_CONTENT>>'

certificate_signing_request = k8s.certificates.v1.CertificateSigningRequest(
    csr_name,
    metadata=k8s.meta.v1.ObjectMetaArgs(name=csr_name),
    spec=k8s.certificates.v1.CertificateSigningRequestSpecArgs(
        # This CSR requests a certificate with the common name "ai-grpc-service"
        # and usage for server authentication, which is required for gRPC TLS communication.
        request=pem_encoded_csr,
        signerName="kubernetes.io/kube-apiserver-client",
        usages=["digital signature", "key encipherment", "server auth"],
    )
)

# Step 2: Once the CSR is approved (out of scope for this program), 
# the signed certificate would be issued and should be stored in a Secret.

# Step 3: Configure the AI training application deployment
#   - It is expected that the application code is configured to use the mounted certificate
#     for setting up TLS with gRPC.
#   - This example assumes a Deployment resource named `ai-training-deployment` already exists
#     and we will add the certificate using volume mounts.
#   - Ensure to replace `<<DEPLOYMENT_NAME>>` and `<<NAMESPACE>>` with actual values.

deployment_name = "<<DEPLOYMENT_NAME>>"
namespace = "<<NAMESPACE>>"
app_labels = {"app": "ai-training"}

ai_training_deployment = k8s.apps.v1.Deployment.get(
    "ai-training-deployment",
    pulumi.ResourceOptions(
        id=f"{namespace}/{deployment_name}"
    )
)

modified_deployment = k8s.apps.v1.Deployment(
    "modified-ai-training-deployment",
    metadata=ai_training_deployment.metadata,
    spec=ai_training_deployment.spec.apply(lambda spec: spec._replace(
        template=k8s.core.v1.PodTemplateSpecArgs(
            metadata=ai_training_deployment.spec.template.metadata,
            spec=k8s.core.v1.PodSpecArgs(
                containers=[k8s.core.v1.ContainerArgs(
                    name="ai-training-container",
                    image=spec.template.spec.containers[0].image,
                    ports=[k8s.core.v1.ContainerPortArgs(container_port=80)],
                    volume_mounts=[k8s.core.v1.VolumeMountArgs(
                        name="ai-grpc-certs",
                        mount_path="/etc/grpc/certs",
                        read_only=True,
                    )]
                )],
                volumes=[k8s.core.v1.VolumeArgs(
                    name="ai-grpc-certs",
                    secret=k8s.core.v1.SecretVolumeSourceArgs(
                        secret_name=f"{csr_name}-certs",
                    )
                )]
            )
        )
    )),
    opts=pulumi.ResourceOptions(depends_on=[certificate_signing_request])
)

# Step 4: Expose the AI training application using a Kubernetes Service with TLS support
ai_service = k8s.core.v1.Service(
    "ai-grpc-service",
    metadata=k8s.meta.v1.ObjectMetaArgs(
        name="ai-grpc-service",
        labels=app_labels,
    ),
    spec=k8s.core.v1.ServiceSpecArgs(
        type="ClusterIP",
        ports=[k8s.core.v1.ServicePortArgs(
            port=80,
            target_port="grpc-port",  # The port your application uses for gRPC communication
        )],
        selector=app_labels,
    )
)

pulumi.export("service_name", ai_service.metadata["name"])
pulumi.export("service_ports", ai_service.spec["ports"])

This program requires you to:

Replace <<PEM_ENCODED_CSR_CONTENT>> with the actual PEM encoded CSR content.
Approve the CSR manually or by automation out of the scope of this code.
Once approved, the certificate will be stored in a Kubernetes secret that should be referred to by deployment's volume secret_name.
Update <<DEPLOYMENT_NAME>> and <<NAMESPACE>> with the actual deployment name and Kubernetes namespace where your AI training application runs.
The labels in app_labels have to match the ones used by your AI training application pods for the service to route traffic correctly.
The target_port should match the container port that your application uses to listen for gRPC communication.

Finally, execute the Pulumi program to apply the configuration. The pulumi.export lines will output the name and ports of the created Service that you can use to interact with your AI training application.