1. Introdução
Este laboratório detalha como criar, provisionar e executar um loop de treinamento de aprendizado por reforço (RL, na sigla em inglês) distribuído de alta performance no GKE Standard com sandboxes de agente do GKE (gVisor), usando o algoritmo de otimização de política relativa de grupo (GRPO) com a biblioteca trl.
O objetivo é demonstrar como avaliar com segurança o código não confiável gerado por LLM durante um loop de treinamento de RL. Para isso, separamos o plano de orquestração (Ray) do plano de execução (sandboxes de agente do GKE).
O desafio técnico da avaliação de código de RL
Ao treinar agentes de LLM usando o aprendizado por reforço (por exemplo, treinar um modelo para escrever código avaliando a saída dele em testes de unidade), o loop de treinamento precisa executar milhares de scripts Python não confiáveis gerados por LLM em paralelo. Isso apresenta desafios críticos:
- O gargalo de rotatividade de pods:as estruturas de avaliação tradicionais criam um contêiner Docker novo por tarefa. Fazer isso dinamicamente para centenas de implementações paralelas durante um loop de treinamento de RL causa uma carga severa no plano de controle do Kubernetes. A latência impossibilita o treinamento de RL de alta frequência.
- O risco de segurança:a execução de código arbitrário gerado por LLM em ambientes de execução de contêiner padrão compartilha o kernel do SO do host. Uma única vulnerabilidade de escape pode comprometer seus nós.
- Roubo de token do IAM:o código gerado por LLM em execução em um pod do Kubernetes pode consultar o servidor de metadados do provedor de nuvem para roubar tokens de conta de serviço do IAM do nó.
A solução: orquestração e execução separadas
Essa arquitetura separa a orquestração da execução:
- O orquestrador (Ray) : um cluster Ray distribuído gerencia o loop de treinamento de RL e distribui a geração de implementação.
- O plano de execução (sandbox de agente do GKE) : em vez de criar pods do Kubernetes dinamicamente, os workers do Ray fazem chamadas HTTP simples para um roteador de sandbox dedicado. O roteador atribui instantaneamente ao worker um contêiner isolado e pré-aquecido em execução no gVisor (GKE Sandbox).
- Latência abaixo de um segundo: como os sandboxes são pré-aquecidos em um
SandboxWarmPoolgerenciado e gerenciados por um gateway HTTP de alta velocidade, a criação do ambiente cai para menos de 200 ms, ignorando completamente o plano de controle do Kubernetes.
Objetivos do laboratório
Neste codelab, você vai aprender o seguinte:
- Os desafios e soluções arquitetônicos para avaliar código não confiável em loops de RL.
- Como criar imagens de sandbox personalizadas para implementações eficientes.
- Como configurar e usar sandboxes de agente do GKE e SandboxWarmPools.
- Como isolar sandboxes com segurança para evitar o roubo de token do IAM.
- Como executar um job de treinamento de RL básico com SweBench e TRL usando o Ray para separar a orquestração da execução.
2. Criação de cluster e pré-requisitos
Antes de continuar, você precisa de um cluster do GKE com um pool de nós de GPU de alta performance e o operador Ray instalado para gerenciar a carga de trabalho de treinamento.
Pré-requisitos
Este codelab pressupõe que as seguintes ferramentas estejam instaladas e configuradas:
- SDK do Google Cloud (
gcloud) - Docker (necessário para criar imagens personalizadas localmente)
kubectl
Variáveis de ambiente
Primeiro, defina as variáveis de ambiente que serão usadas neste codelab. Os comandos abaixo usam padrões razoáveis, mas você pode alterá-los conforme necessário para corresponder ao seu ambiente específico do Google Cloud:
export PROJECT_ID=$(gcloud config get-value project)
export REGION="us-west3"
export ZONE="us-west3-a"
export REPO_NAME="rl-sandbox-repo"
Crie um repositório do Artifact Registry para armazenar as imagens de contêiner personalizadas:
gcloud artifacts repositories create $REPO_NAME \
--repository-format=docker \
--location=$REGION \
--description="Repository for RL Sandbox images"
Configuração do cluster
Para um tutorial completo sobre como provisionar um cluster do GKE otimizado para cargas de trabalho de IA (incluindo a fiação de rede GPUDirect RDMA), siga a documentação oficial: Criar um cluster personalizado de hipercomputação de IA do GKE
Pré-requisito essencial:ao criar o cluster ou um pool de nós de execução específico, transmita as flags --enable-agent-sandbox e --sandbox type=gvisor para instalar as definições de recursos personalizados (CRDs, na sigla em inglês) necessárias para os pools de aquecimento de sandbox.
Supondo que o cluster, as GPUs e o operador Ray estejam em execução, tudo abaixo detalha como configurar o plano de execução e executar o loop de RL.
3. Criar imagens personalizadas
Um aspecto crucial da execução de RL de alta performance é incorporar dependências às imagens. Precisamos de duas imagens distintas: uma para os workers de GPU que executam o modelo e outra para os sandboxes isolados que executam o código de avaliação não confiável.
1. Criar a imagem do worker de GPU
O worker de GPU do Ray precisa de bibliotecas para executar o modelo de linguagem e orquestrar o loop de treinamento. Criamos essa imagem com base na imagem oficial do vLLM para que ela seja compatível com as GPUs mais recentes e tenha o PyTorch/CUDA pré-instalado.
Execute o seguinte comando para criar Dockerfile.gpu_worker:
cat << 'EOF' > Dockerfile.gpu_worker
# ==============================================================================
# Base Image: Use the official vLLM production image.
# This image comes pre-baked with PyTorch 2.11, CUDA 13.0, and vLLM.
# It supports sm_100 Blackwell GPUs natively!
# ==============================================================================
FROM vllm/vllm-openai:latest
USER root
# Install system dependencies
RUN apt-get update && apt-get install -y --no-install-recommends \
build-essential \
numactl \
libnuma-dev \
wget \
ca-certificates \
&& apt-get clean \
&& rm -rf /var/lib/apt/lists/*
# Install Ray, TRL, and Sandbox tools
# TRL does not require compiling flash_attn from source.
RUN pip install --no-cache-dir \
"ray[default]==2.55.1" \
"numpy<2.0" \
gymnasium>=0.28.1 \
k8s-agent-sandbox>=0.4.6 \
trl transformers packaging ninja cachetools accelerate datasets peft
EOF
Crie e envie a imagem para o repositório do Artifact Registry:
export WORKER_REPO="${REGION}-docker.pkg.dev/${PROJECT_ID}/${REPO_NAME}/ray-gpu-worker:v1"
docker build -f Dockerfile.gpu_worker -t $WORKER_REPO .
docker push $WORKER_REPO
Observação:este guia usa comandos docker locais para criar imagens. Se preferir criar as imagens remotamente, use o Cloud Build (por exemplo, usando gcloud builds submit).
2. Criar a imagem principal da CPU
O nó principal do Ray apenas orquestra o cluster e não executa os modelos de treinamento de GPU pesados. Para evitar um gargalo de extração de imagem enorme (normalmente 15 GB ou mais) nos nós de CPU padrão, criamos uma imagem leve e somente de CPU para o nó principal. Essa imagem contém o Ray e as bibliotecas Python necessárias, mas exclui bibliotecas de GPU pesadas, como CUDA e vLLM.
Execute o seguinte comando para criar Dockerfile.head:
cat << 'EOF' > Dockerfile.head
# ==============================================================================
# Base Image: Use the official Python slim image for the exact patch version.
# This aligns the Python version (3.12.13) with the GPU worker node.
# ==============================================================================
FROM python:3.12.13-slim
USER root
# Install system dependencies
RUN apt-get update && apt-get install -y --no-install-recommends \
build-essential \
wget \
ca-certificates \
&& apt-get clean \
&& rm -rf /var/lib/apt/lists/*
# Install Ray, TRL, and Sandbox tools (CPU versions where applicable)
# We install torch CPU first to avoid pulling the 2GB+ CUDA torch package.
RUN pip install --no-cache-dir torch --index-url https://download.pytorch.org/whl/cpu && \
pip install --no-cache-dir \
"ray[default]==2.55.1" \
"numpy<2.0" \
gymnasium>=0.28.1 \
k8s-agent-sandbox>=0.4.6 \
trl transformers packaging ninja cachetools accelerate datasets peft
# Create a 'ray' user to run the container securely and match Ray conventions
RUN useradd -ms /bin/bash ray
USER ray
WORKDIR /home/ray
EOF
Crie e envie a imagem:
export HEAD_REPO="${REGION}-docker.pkg.dev/${PROJECT_ID}/${REPO_NAME}/ray-head:v1"
docker build -f Dockerfile.head -t $HEAD_REPO .
docker push $HEAD_REPO
3. Criar a imagem de sandbox
O sandbox precisa das dependências específicas da tarefa que estamos avaliando para que a instalação do ambiente de execução seja instantânea. Para este codelab, vamos usar um problema do repositório django/django no SWE-bench. Vamos pré-clonar o repositório e pré-criar os ambientes Python para que nossos scripts de modelo não percam tempo fazendo o download deles no loop de RL.
Execute o seguinte comando para criar Dockerfile.sandbox:
cat << 'EOF' > Dockerfile.sandbox
# Use a stable Debian-based Miniconda image
FROM condaforge/miniforge3:latest
# 1. Install essential system libraries (including sqlite3 for Django tests)
RUN apt-get update && apt-get install -y \
git \
build-essential \
libsqlite3-dev \
&& rm -rf /var/lib/apt/lists/*
# 2. Set up the /workspace directory and grant ownership to the pre-existing non-root 'ubuntu' user (UID 1000)
RUN mkdir -p /workspace \
&& chown -R 1000:1000 /workspace
# 3. Switch to the non-root user
USER ubuntu
WORKDIR /workspace
# 4. Pre-configure Git globally so the agent can run git commands
RUN git config --global user.email "agent@gke-sandbox.local" \
&& git config --global user.name "Agent"
# 5. Pre-clone the repository as the non-root user
RUN git clone https://github.com/django/django.git .
# 6. Pre-build Conda environments and pre-cache common dependencies
# We do NOT run "pip install -e ." here to avoid Python version conflicts with the main branch.
# Instead, we pre-install the heavy dependencies so that runtime installation is instantaneous.
RUN conda create -y -n django-py39 python=3.9 \
&& conda run -n django-py39 pip install --no-cache-dir asgiref sqlparse tzdata pytest pytest-django
RUN conda create -y -n django-py310 python=3.10 \
&& conda run -n django-py310 pip install --no-cache-dir asgiref sqlparse tzdata pytest pytest-django
# --- Add Agent Server ---
# We use a multi-stage build to copy the agent server from the official python-runtime-sandbox image
COPY --from=registry.k8s.io/agent-sandbox/python-runtime-sandbox:v0.1.0 /app /opt/sandbox-agent
USER root
RUN chown -R 1000:1000 /opt/sandbox-agent \
&& /opt/conda/bin/pip install --no-cache-dir -r /opt/sandbox-agent/requirements.txt \
&& sed -i 's|"/app"|"/workspace"|g' /opt/sandbox-agent/main.py
USER ubuntu
# ------------------------
# Prepend the django-py39 conda environment bin to PATH for commands executed inside the container
ENV PATH=/home/ubuntu/.conda/envs/django-py39/bin:$PATH
# Keep the container alive and run the agent server using the system Python
CMD ["/opt/conda/bin/python3", "-m", "uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8888", "--log-level", "trace", "--app-dir", "/opt/sandbox-agent"]
EOF
Crie e envie a imagem:
export SANDBOX_REPO="${REGION}-docker.pkg.dev/${PROJECT_ID}/${REPO_NAME}/django-sandbox:v1"
docker build -f Dockerfile.sandbox -t $SANDBOX_REPO .
docker push $SANDBOX_REPO
4. Configurar orquestração e execução
Agora vamos implantar o cluster Ray para orquestração e os recursos de sandbox para execução.
1. Configuração do cluster Ray
Implante um recurso personalizado do RayCluster. Os recursos disponíveis do cluster (como memória, CPU ou tipo de GPU) podem ser diferentes. Ajuste as solicitações e os limites de resources de acordo.
Execute o seguinte comando para criar raycluster.yaml. Ele usa cat << EOF para substituir automaticamente as variáveis de ambiente no manifesto:
cat << EOF > raycluster.yaml
apiVersion: ray.io/v1
kind: RayCluster
metadata:
name: grpo-cluster
namespace: default
spec:
rayVersion: "2.55.1"
headGroupSpec:
rayStartParams:
dashboard-host: "0.0.0.0"
template:
spec:
containers:
- name: ray-head
image: ${REGION}-docker.pkg.dev/${PROJECT_ID}/${REPO_NAME}/ray-head:v1
ports:
- containerPort: 6379
name: gcs-server
- containerPort: 8265
name: dashboard
- containerPort: 10001
name: client
resources:
limits:
cpu: "2"
memory: "8Gi"
requests:
cpu: "2"
memory: "8Gi"
workerGroupSpecs:
- groupName: gpu-group
replicas: 1
minReplicas: 1
maxReplicas: 1
rayStartParams: {}
template:
spec:
containers:
- name: ray-worker
image: ${REGION}-docker.pkg.dev/${PROJECT_ID}/${REPO_NAME}/ray-gpu-worker:v1
resources:
limits:
cpu: "12"
memory: "120Gi"
nvidia.com/gpu: "1"
requests:
cpu: "12"
memory: "120Gi"
nvidia.com/gpu: "1"
EOF
Aplique-o:
kubectl apply -f raycluster.yaml
Verifique se o cluster foi criado e está em execução (isso pode levar alguns minutos):
kubectl get raycluster
Saída esperada:
NAME DESIRED WORKERS AVAILABLE WORKERS CPUS MEMORY GPUS STATUS AGE rl-cluster 1 1 ready 2m
2. Configuração do SandboxRouter
O SandboxRouter atua como um gateway HTTP de alta velocidade, recebendo solicitações de workers do Ray e conectando-as instantaneamente aos pods gVisor disponíveis, ignorando o ciclo de vida do pod do servidor da API Kubernetes mais lento.
Execute o seguinte comando para criar sandbox_router.yaml:
cat << 'EOF' > sandbox_router.yaml
apiVersion: rbac.authorization.k8s.io/v1
kind: Role
metadata:
namespace: default
name: sandbox-claim-manager
rules:
- apiGroups: ["extensions.agents.x-k8s.io"]
resources: ["sandboxclaims"]
verbs: ["get", "list", "watch", "create", "update", "patch", "delete"]
- apiGroups: ["agents.x-k8s.io"]
resources: ["sandboxes"]
verbs: ["get", "list", "watch", "create", "update", "patch", "delete"]
---
apiVersion: rbac.authorization.k8s.io/v1
kind: RoleBinding
metadata:
name: sandbox-claim-manager-binding
namespace: default
subjects:
- kind: ServiceAccount
name: default
namespace: default
roleRef:
kind: Role
name: sandbox-claim-manager
apiGroup: rbac.authorization.k8s.io
---
apiVersion: v1
kind: Service
metadata:
name: sandbox-router
namespace: default
spec:
type: ClusterIP
selector:
app: sandbox-router
ports:
- name: http
protocol: TCP
port: 8080
targetPort: 8080
---
apiVersion: apps/v1
kind: Deployment
metadata:
name: sandbox-router-deployment
namespace: default
spec:
replicas: 2
selector:
matchLabels:
app: sandbox-router
template:
metadata:
labels:
app: sandbox-router
spec:
containers:
- name: router
image: us-central1-docker.pkg.dev/k8s-staging-images/agent-sandbox/sandbox-router:latest-main
ports:
- containerPort: 8080
env:
- name: ALLOW_UNAUTHENTICATED_ROUTER
value: "true"
EOF
Aplique-o:
kubectl apply -f sandbox_router.yaml
Verifique se a implantação está em execução:
kubectl get deployment sandbox-router-deployment
Saída esperada:
NAME READY UP-TO-DATE AVAILABLE AGE sandbox-router-deployment 2/2 2 2 1m
3. Configuração do SandboxTemplate e do WarmPool
O sandbox de agente do GKE permite a atribuição instantânea de contêineres isolados e pré-aquecidos usando o roteador de sandbox. Definimos um SandboxTemplate e um SandboxWarmPool para manter os pods prontos.
Execute o seguinte comando para criar sandbox_warmpool.yaml com as variáveis de ambiente:
cat << EOF > sandbox_warmpool.yaml
apiVersion: extensions.agents.x-k8s.io/v1alpha1
kind: SandboxTemplate
metadata:
name: swe-bench-django
namespace: default
spec:
podTemplate:
spec:
runtimeClassName: gvisor
securityContext:
runAsNonRoot: true
runAsUser: 1000
nodeSelector:
sandbox.gke.io/runtime: gvisor
tolerations:
- key: sandbox.gke.io/runtime
operator: Equal
value: gvisor
effect: NoSchedule
containers:
- name: sandbox
image: ${REGION}-docker.pkg.dev/${PROJECT_ID}/${REPO_NAME}/django-sandbox:v1
securityContext:
allowPrivilegeEscalation: false
capabilities:
drop:
- ALL
resources:
requests:
cpu: "2"
memory: "4Gi"
limits:
cpu: "2"
memory: "4Gi"
---
apiVersion: extensions.agents.x-k8s.io/v1alpha1
kind: SandboxWarmPool
metadata:
name: swe-bench-django-warmpool
namespace: default
spec:
replicas: 10
sandboxTemplateRef:
name: swe-bench-django
EOF
Aplique-o:
kubectl apply -f sandbox_warmpool.yaml
Verifique se o SandboxWarmPool foi inicializado:
kubectl get sandboxwarmpool
Saída esperada:
NAME READY AGE swe-bench-django-warmpool 10 1m
4. Isolamento de segurança
Uma NetworkPolicy isola estritamente os sandboxes, impedindo a saída para o servidor de metadados do GCP e, portanto, evitando o roubo de token do IAM.
Execute o seguinte comando para criar network_policy.yaml:
cat << 'EOF' > network_policy.yaml
apiVersion: networking.k8s.io/v1
kind: NetworkPolicy
metadata:
name: block-metadata-egress
namespace: default
spec:
podSelector:
matchLabels:
sandbox.gke.io/runtime: gvisor
policyTypes:
- Egress
egress:
- to:
- ipBlock:
cidr: 0.0.0.0/0
except:
- 169.254.169.254/32
EOF
Aplique a política:
kubectl apply -f network_policy.yaml
Verifique se a NetworkPolicy foi criada:
kubectl get networkpolicy
Saída esperada:
NAME POD-SELECTOR AGE block-metadata-egress sandbox.gke.io/runtime=gvisor 1m
5. Job de RL básico com SweBench e TRL
Depois que o cluster e os sandboxes estiverem preparados, poderemos executar um loop de treinamento de GRPO. Vamos usar a biblioteca trl para orquestrar o algoritmo GRPO e as funções remotas do Ray para avaliar o código gerado nos sandboxes isolados.
Para acelerar a execução deste codelab, vamos filtrar um único problema do Django. A lógica de roteamento abaixo mostra como selecionar diferentes warmpools para diferentes repositórios, o que é útil ao expandir para o conjunto de dados completo do SWE-bench.
O script de treinamento
Execute o seguinte comando para criar train_trl.py:
cat << 'EOF' > train_trl.py
import ray
from k8s_agent_sandbox import SandboxClient
from k8s_agent_sandbox.models import SandboxDirectConnectionConfig
from trl import GRPOConfig, GRPOTrainer
from transformers import AutoModelForCausalLM, AutoTokenizer
from datasets import load_dataset
import urllib.request
import re
ray.init(ignore_reinit_error=True)
# 1. Define the Ray remote evaluation function
@ray.remote
def evaluate_rollout(code, prompt_data):
client = SandboxClient(connection_config=SandboxDirectConnectionConfig(api_url="http://sandbox-router.default.svc.cluster.local:8080"))
# Claim a pre-warmed sandbox instantly based on the repo
repo = prompt_data.get("repo")
# In a full system, you'd route to different warmpools based on repo
# Here we default to django for our single task
sandbox = client.create_sandbox(
template="swe-bench-django",
warmpool="swe-bench-django-warmpool",
sandbox_ready_timeout=600
)
try:
# Check if the code is correctly formatted
bash_match = re.search(r"```bash\n(.*?)\n```", code, re.DOTALL)
if not bash_match:
return 0.0
script = bash_match.group(1)
# In a real environment, we would apply the base commit and install here
# For simplicity, we just execute the script
import shlex
script_cmd = f"bash -c {shlex.quote(script)}"
result = sandbox.commands.run(script_cmd, timeout=60)
# Calculate continuous reward based on test passage ratio
if result.exit_code == 0:
return 1.0
# Very simple heuristic reward
return 0.1
finally:
# Clean up and release the sandbox back to the pool
client.delete_sandbox(sandbox.claim_name)
# 2. Define the Reward Function for TRL
def sandbox_reward_func(prompts, completions, **kwargs):
# Dispatch evaluation to Ray cluster
futures = [
evaluate_rollout.remote(completion, {
"repo": kwargs.get('repo', [])[i] if 'repo' in kwargs else None,
"base_commit": kwargs.get('base_commit', [])[i] if 'base_commit' in kwargs else None
}) for i, completion in enumerate(completions)
]
# Block and wait for all sandbox evaluations to complete
rewards = ray.get(futures)
return rewards
# 3. Setup GRPO Trainer
@ray.remote(num_gpus=1, num_cpus=8)
def train():
# Load dataset
dataset = load_dataset("princeton-nlp/SWE-bench_Lite", split="test")
# Filter to our selected target issue
dataset = dataset.filter(lambda x: x["instance_id"] == "django__django-15388")
def format_dataset(example):
files = re.findall(r'^\+\+\+ b/(.+)$', example["patch"], re.MULTILINE)
target_file = files[0] if files else ""
file_content = ""
if target_file:
try:
github_repo = example["repo"]
url = f"https://raw.githubusercontent.com/{github_repo}/{example['base_commit']}/{target_file}"
with urllib.request.urlopen(url) as response:
file_content = response.read().decode('utf-8')
except Exception as e:
pass
prompt = f"""You are an expert software engineer.
You are given a GitHub issue and the content of the file that contains the bug.
Write an executable bash script that will modify the target file to fix the bug (e.g. using cat << 'EOF' > {target_file} or inline python edits).
Wrap your bash script in ```bash ... ``` tags. Do not output raw python code directly.
Target File: {target_file}
Original File Content:
```python
{file_content}
```
Issue:
{example['problem_statement']}
"""
return {
"prompt": prompt,
"repo": example["repo"],
"instance_id": example["instance_id"],
"base_commit": example["base_commit"],
}
dataset = dataset.map(format_dataset)
model_name = "Qwen/Qwen2.5-Coder-1.5B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)
training_args = GRPOConfig(
output_dir="outputs",
learning_rate=5e-6,
max_steps=50,
per_device_train_batch_size=1,
gradient_accumulation_steps=4,
num_generations=4,
)
trainer = GRPOTrainer(
model=model_name,
processing_class=tokenizer,
reward_funcs=[sandbox_reward_func],
args=training_args,
train_dataset=dataset,
)
print("Starting GRPO training with GKE Agent Sandboxes...")
trainer.train()
def main():
print("Submitting training job to GPU worker...")
ray.get(train.remote())
if __name__ == "__main__":
main()
EOF
Enviar o job para o cluster
Primeiro, encaminhe a porta para o painel do Ray Head e envie o job de treinamento da máquina local:
kubectl port-forward service/grpo-cluster-head-svc 8265:8265 &
ray job submit \
--address http://localhost:8265 \
--runtime-env-json '{"working_dir": "."}' \
-- python train_trl.py
Monitorar a execução
Você pode monitorar o progresso da execução:
- Painel do Ray:abra
http://localhost:8265no navegador. - Reivindicações de sandbox:observe o GKE reivindicar e liberar sandboxes dinamicamente no gVisor:
watch -n 1 "kubectl get sandboxclaims,sandboxes,pods"
6. Conclusão
Parabéns! Você configurou e executou com sucesso um loop de treinamento de RL distribuído de alta performance com segurança no GKE Standard usando sandboxes de agente do GKE.