Introducción a la implementación de IA multimodal en K8s

El despliegue de modelos de lenguaje de gran tamaño (LLM) con capacidades de visión, como Gemma-3-12b-it, demanda una infraestructura robusta que gestione eficientemente el consumo de VRAM y procese solicitudes concurrentes. Kubernetes (K8s) se presenta como la solución ideal para estas aplicaciones, permitiendo la orquestación de servicios, el autoescalado y la alta disponibilidad de los modelos de inteligencia artificial.

Gemma-3 Pixel Studio no solo requiere potencia de cálculo bruta, sino una gestión inteligente de los recursos. En este artículo, exploraremos cómo configurar un entorno de múltiples instancias para garantizar un servicio estable y escalable.

Requisitos del sistema y preparación del entorno

Antes de proceder con el despliegue, es fundamental validar que el clúster cumple con las especificaciones técnicas necesarias para soportar cargas de trabajo de GPU.

Hardware recomendado por Pod

• Unidad de Procesamiento Gráfico: NVIDIA GPU con un mínimo de 24 GB de memoria de video (ej. RTX 4090, A10 o A100).
• Memoria RAM: 32 GB o superior.
• Almacenamiento: 60 GB de espacio libre para la caché del modelo.

Dependencias de Software

• Clúster de Kubernetes versión 1.24 o superior.
• NVIDIA GPU Operator instalado y funcionando correctamente.
• Runtime de contenedores compatible con NVIDIA (containerd o Docker).
• Configuración de almacenamiento persistente (NFS, Longhorn o soluciones Cloud).

Configuración inicial en Kubernetes

Comenzamos definiendo un espacio de nombres dedicado y cuotas de recursos para evitar que la aplicación consuma la totalidad de los recursos del clúster.

# namespace-setup.yaml
apiVersion: v1
kind: Namespace
metadata:
  name: pixel-ai-space
---
apiVersion: v1
kind: ResourceQuota
metadata:
  name: ai-resource-limits
  namespace: pixel-ai-space
spec:
  hard:
    requests.cpu: "12"
    requests.memory: 48Gi
    requests.nvidia.com/gpu: "4"
    limits.cpu: "24"
    limits.memory: 96Gi
    limits.nvidia.com/gpu: "4"

Contenerización del servicio

Para empaquetar Gemma-3 Pixel Studio, utilizaremos una imagen base optimizada para Python y PyTorch. Es recomendable separar los archivos del modelo del código de la aplicación mediante volúmenes persistentes.

Dockerfile de la aplicación

FROM python:3.10-slim-bullseye

RUN apt-get update && apt-get install -y \
    ffmpeg \
    libsm6 \
    libxext6 \
    && rm -rf /var/lib/apt/lists/*

WORKDIR /opt/pixel-studio

COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

COPY . .

RUN useradd -u 1001 studio-worker
USER studio-worker

EXPOSE 8501

CMD ["streamlit", "run", "main.py", "--server.port=8501", "--server.address=0.0.0.0"]

Dependencias (requirements.txt)

torch==2.1.2
transformers==4.36.0
accelerate==0.25.0
streamlit==1.29.0
einops==0.7.0

Definición de recursos de Kubernetes

La arquitectura se basa en un Deployment que gestiona los Pods y un Service para el balanceo de carga.

Configuración de parámetros (ConfigMap)

apiVersion: v1
kind: ConfigMap
metadata:
  name: studio-runtime-cfg
  namespace: pixel-ai-space
data:
  MODEL_ID: "google/gemma-3-12b-it"
  GPU_LAYERS: "max"
  INFERENCE_PRECISION: "bfloat16"

Despliegue de la aplicación (Deployment)

apiVersion: apps/v1
kind: Deployment
metadata:
  name: pixel-studio-node
  namespace: pixel-ai-space
spec:
  replicas: 3
  selector:
    matchLabels:
      app: pixel-studio
  template:
    metadata:
      labels:
        app: pixel-studio
    spec:
      containers:
      - name: inference-engine
        image: custom-registry/pixel-studio:v1.0
        resources:
          limits:
            nvidia.com/gpu: 1
            memory: "30Gi"
            cpu: "6"
        volumeMounts:
        - name: model-weights
          mountPath: /models
          readOnly: true
        envFrom:
        - configMapRef:
            name: studio-runtime-cfg
      volumes:
      - name: model-weights
        persistentVolumeClaim:
          claimName: gemma-weights-pvc
      nodeSelector:
        hardware-type: gpu-node

Escalamiento y Optimización de Sesiones

Para aplicaciones de IA que mantienen estados de inferencia en memoria, es vital configurar correctamente el autoescalado y la afinidad de sesión.

Autoescalado Horizontal (HPA)

Configuramos el sistema para que escale cuando el uso de la CPU o el procesamiento de inferencias supere un umbral determinado.

apiVersion: autoscaling/v2
kind: HorizontalPodAutoscaler
metadata:
  name: pixel-studio-hpa
  namespace: pixel-ai-space
spec:
  scaleTargetRef:
    apiVersion: apps/v1
    kind: Deployment
    name: pixel-studio-node
  minReplicas: 2
  maxReplicas: 8
  metrics:
  - type: Resource
    resource:
      name: cpu
      target:
        type: Utilization
        averageUtilization: 75

Persistencia de Sesión en el Balanceador

Dado que Streamlit utiliza WebSockets, debemos asegurar que el cliente se mantenga conectado al mismo Pod mediante sessionAffinity.

apiVersion: v1
kind: Service
metadata:
  name: pixel-studio-svc
  namespace: pixel-ai-space
spec:
  selector:
    app: pixel-studio
  ports:
  - protocol: TCP
    port: 80
    targetPort: 8501
  sessionAffinity: ClientIP
  sessionAffinityConfig:
    clientIP:
      timeoutSeconds: 1800

Estrategias de Monitoreo y Verificación

Una vez desplegado el servicio, es necesario implementar sondas de salud (probes) para asegurar que el modelo está cargado antes de recibir tráfico.

# Añadir a la sección del contenedor en el Deployment
readinessProbe:
  httpGet:
    path: /healthz
    port: 8501
  initialDelaySeconds: 45
  periodSeconds: 10
livenessProbe:
  tcpSocket:
    port: 8501
  initialDelaySeconds: 90
  periodSeconds: 20

Para diagnosticar problemas de aisgnación de memoria GPU dentro del clúster, se puede ejecutar el siguiente comando directamente en el Pod:

kubectl exec -it <pod-name> -n pixel-ai-space -- nvidia-smi

Consideraciones de Producción

• Cuantización: Si la VRAM es limitada, considere cargar el modelo en formato 4-bit o 8-bit utilizando la biblioteca bitsandbytes.
• Aislamiento de nodos: Use Taints y Tolerations para asegurar que solo las cargas de trabajo de IA se ejecuten en nodos con GPU.
• Gestión de Modelos: Utilice un sistema de almacenamiento compartido como CephFS para que todas las instancias accedan a los mismos pesos del modelo sin duplicar el almacenamiento.