1. Introduzione
Questo lab descrive in dettaglio come creare, eseguire il provisioning ed eseguire un loop di addestramento di apprendimento per rinforzo (RL) distribuito ad alte prestazioni su GKE Standard con GKE Agent Sandbox (gVisor), utilizzando l'algoritmo Group Relative Policy Optimization (GRPO) con la libreria trl.
L'obiettivo è dimostrare come valutare in modo sicuro il codice non attendibile generato da LLM durante un loop di addestramento RL. A questo scopo, disaccoppiamo il piano di orchestrazione (Ray) dal piano di esecuzione (GKE Agent Sandbox).
La sfida tecnica della valutazione del codice RL
Quando si addestrano gli agenti LLM utilizzando l'apprendimento per rinforzo (ad es. addestrando un modello per scrivere codice valutando il suo output sui test unitari), il loop di addestramento deve eseguire in parallelo migliaia di script Python non attendibili generati da LLM. Ciò introduce sfide critiche:
- Il collo di bottiglia del churn dei pod: i framework di valutazione tradizionali avviano un nuovo container Docker per ogni attività. Eseguire questa operazione in modo dinamico per centinaia di implementazioni parallele durante un loop di addestramento RL causa un carico elevato sul piano di controllo Kubernetes. La latenza rende impossibile l'addestramento RL ad alta frequenza.
- Il rischio per la sicurezza: l'esecuzione di codice arbitrario generato da LLM all'interno dei runtime dei container standard condivide il kernel del sistema operativo host. Una singola vulnerabilità di escape può compromettere i nodi.
- Furto di token IAM: il codice generato da LLM in esecuzione all'interno di un pod Kubernetes può eseguire query sul server di metadati del provider di servizi cloud per rubare i token dell'account di servizio IAM del nodo.
La soluzione: orchestrazione ed esecuzione disaccoppiate
Questa architettura disaccoppia l'orchestrazione dall'esecuzione:
- L'orchestratore (Ray): un cluster Ray distribuito gestisce il loop di addestramento RL e distribuisce la generazione di rollout.
- Il piano di esecuzione (GKE Agent Sandbox): anziché creare dinamicamente i pod Kubernetes, i worker Ray effettuano semplici chiamate HTTP a un router sandbox dedicato. Il router assegna immediatamente al worker un container isolato e pre-riscaldato in esecuzione in gVisor (GKE Sandbox).
- Latenza inferiore al secondo: poiché le sandbox vengono pre-riscaldate in un
SandboxWarmPoolgestito e gestite tramite un gateway HTTP ad alta velocità, la creazione dell'ambiente scende a meno di 200 ms, bypassando completamente il piano di controllo Kubernetes.
Obiettivi del lab
In questo codelab imparerai a:
- Le sfide architettoniche e le soluzioni per la valutazione del codice non attendibile nei loop RL.
- Come creare immagini sandbox personalizzate per rollout efficienti.
- Come configurare e utilizzare GKE Agent Sandbox e SandboxWarmPool.
- Come isolare in modo sicuro le sandbox per impedire il furto di token IAM.
- Come eseguire un job di addestramento RL di base con SweBench e TRL utilizzando Ray per disaccoppiare l'orchestrazione dall'esecuzione.
2. Creazione del cluster e prerequisiti
Prima di procedere, devi disporre di un cluster GKE con un pool di nodi GPU ad alte prestazioni e dell'operatore Ray installato per gestire il carico di lavoro di addestramento.
Prerequisiti
Questo codelab presuppone che i seguenti strumenti siano installati e configurati:
- Google Cloud SDK (
gcloud) - Docker (obbligatorio per la creazione di immagini personalizzate in locale)
kubectl
Variabili di ambiente
Innanzitutto, imposta le variabili di ambiente che verranno utilizzate durante questo codelab. I comandi riportati di seguito utilizzano valori predefiniti ragionevoli, ma puoi modificarli in base alle esigenze per adattarli al tuo ambiente Google Cloud specifico:
export PROJECT_ID=$(gcloud config get-value project)
export REGION="us-west3"
export ZONE="us-west3-a"
export REPO_NAME="rl-sandbox-repo"
Crea un repository Artifact Registry per contenere le immagini container personalizzate:
gcloud artifacts repositories create $REPO_NAME \
--repository-format=docker \
--location=$REGION \
--description="Repository for RL Sandbox images"
Configurazione cluster
Per una procedura completa sul provisioning di un cluster GKE ottimizzato per i carichi di lavoro AI (inclusi i cablaggi di rete GPUDirect RDMA), consulta la documentazione ufficiale: Crea un cluster personalizzato GKE AI Hypercompute
Prerequisito fondamentale: quando crei il cluster o un pool di nodi di esecuzione specifico, assicurati di passare i flag --enable-agent-sandbox e --sandbox type=gvisor per installare le definizioni di risorse personalizzate (CRD) richieste per i pool di riscaldamento sandbox.
Supponendo che il cluster, le GPU e l'operatore Ray siano in esecuzione, tutto ciò che segue descrive come configurare il piano di esecuzione ed eseguire il loop RL.
3. Crea immagini personalizzate
Un aspetto fondamentale dell'esecuzione di RL ad alte prestazioni è l'incorporamento delle dipendenze nelle immagini. Sono necessarie due immagini distinte: una per i worker GPU che eseguono il modello e una per le sandbox isolate che eseguono il codice di valutazione non attendibile.
1. Crea l'immagine del worker GPU
Il worker GPU Ray ha bisogno di librerie per eseguire il modello linguistico e orchestrare il loop di addestramento. Creiamo questa immagine sopra l'immagine vLLM ufficiale in modo che supporti le GPU più recenti e abbia PyTorch/CUDA preinstallato.
Esegui il comando seguente per creare 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
Crea l'immagine ed eseguine il push nel repository 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
Nota: questa guida utilizza i comandi docker locali per creare le immagini. Se preferisci creare le immagini da remoto, puoi utilizzare Cloud Build (ad es. utilizzando gcloud builds submit).
2. Crea l'immagine dell'head CPU
Il nodo head Ray orchestra solo il cluster e non esegue i modelli di addestramento GPU pesanti. Per evitare un collo di bottiglia di pull di immagini di grandi dimensioni (in genere 15 GB o più) sui nodi CPU standard, creiamo un'immagine leggera solo CPU per il nodo head. Questa immagine contiene Ray e le librerie Python richieste, ma esclude le librerie GPU pesanti come CUDA e vLLM.
Esegui il comando seguente per creare 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
Crea l'immagine ed eseguine il push:
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. Crea l'immagine sandbox
La sandbox ha bisogno delle dipendenze specifiche per l'attività che stiamo valutando, in modo che l'installazione del runtime sia istantanea. Per questo codelab, utilizzeremo un problema dal repository django/django in SWE-bench. Pre-cloneremo il repository e pre-creeremo gli ambienti Python in modo che gli script del modello non perdano tempo a scaricarli nel loop RL.
Esegui il comando seguente per creare 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
Crea l'immagine ed eseguine il push:
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. Configura l'orchestrazione e l'esecuzione
Ora eseguiamo il deployment del cluster Ray per l'orchestrazione e delle risorse sandbox per l'esecuzione.
1. Configurazione del cluster Ray
Esegui il deployment di una risorsa personalizzata RayCluster. Tieni presente che le risorse disponibili del cluster (come memoria, CPU o tipo di GPU) potrebbero variare. Modifica di conseguenza le richieste e i limiti di resources.
Esegui il comando seguente per creare raycluster.yaml. Questo utilizza cat << EOF per sostituire automaticamente le variabili di ambiente nel manifest:
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
Applica:
kubectl apply -f raycluster.yaml
Verifica che il cluster sia stato creato e sia in esecuzione (l'operazione potrebbe richiedere alcuni minuti):
kubectl get raycluster
Output previsto:
NAME DESIRED WORKERS AVAILABLE WORKERS CPUS MEMORY GPUS STATUS AGE rl-cluster 1 1 ready 2m
2. Configurazione di SandboxRouter
SandboxRouter funge da gateway HTTP ad alta velocità, gestendo le richieste dei worker Ray e collegandole immediatamente ai pod gVisor disponibili, bypassando il ciclo di vita dei pod del server API Kubernetes più lento.
Esegui il comando seguente per creare 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
Applica:
kubectl apply -f sandbox_router.yaml
Verifica che il deployment sia in esecuzione:
kubectl get deployment sandbox-router-deployment
Output previsto:
NAME READY UP-TO-DATE AVAILABLE AGE sandbox-router-deployment 2/2 2 2 1m
3. Configurazione di SandboxTemplate e WarmPool
GKE Agent Sandbox consente l'assegnazione immediata di container isolati e pre-riscaldati utilizzando Sandbox Router. Definiamo un SandboxTemplate e un SandboxWarmPool per mantenere i pod pronti.
Esegui il comando seguente per creare sandbox_warmpool.yaml con le variabili di 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
Applica:
kubectl apply -f sandbox_warmpool.yaml
Verifica che SandboxWarmPool sia inizializzato:
kubectl get sandboxwarmpool
Output previsto:
NAME READY AGE swe-bench-django-warmpool 10 1m
4. Isolamento di sicurezza
Un NetworkPolicy isola rigorosamente le sandbox, impedendo l'uscita dal server di metadati GCP e quindi il furto di token IAM.
Esegui il comando seguente per creare 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
Applica la policy:
kubectl apply -f network_policy.yaml
Verifica che NetworkPolicy sia stato creato:
kubectl get networkpolicy
Output previsto:
NAME POD-SELECTOR AGE block-metadata-egress sandbox.gke.io/runtime=gvisor 1m
5. Job RL di base con SweBench e TRL
Una volta preparati il cluster e le sandbox, possiamo eseguire un loop di addestramento GRPO. Utilizzeremo la libreria trl per orchestrare l'algoritmo GRPO e le funzioni remote Ray per valutare il codice generato all'interno delle sandbox isolate.
Per velocizzare l'esecuzione di questo codelab, filtreremo un singolo problema di Django. La logica di routing riportata di seguito mostra come selezionare diversi pool di riscaldamento per repository diversi, il che è utile quando si espande al set di dati SWE-bench completo.
Lo script di addestramento
Esegui il comando seguente per creare 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
Invia il job al cluster
Innanzitutto, esegui il port forwarding alla dashboard Ray Head e invia il job di addestramento dalla macchina locale:
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
Monitora l'esecuzione
Puoi monitorare l'avanzamento dell'esecuzione:
- Dashboard Ray: apri
http://localhost:8265nel browser. - Richieste sandbox: guarda GKE richiedere e rilasciare dinamicamente le sandbox in gVisor:
watch -n 1 "kubectl get sandboxclaims,sandboxes,pods"
6. Conclusione
Complimenti! Hai configurato ed eseguito correttamente un loop di addestramento RL distribuito ad alte prestazioni in modo sicuro su GKE Standard utilizzando GKE Agent Sandbox.