1. 简介
本实验详细介绍了如何使用 trl 库中的群组相对政策优化 (GRPO) 算法，在 GKE Standard 上通过 GKE Agent Sandbox (gVisor) 构建、配置和执行高性能分布式强化学习 (RL) 训练循环。
目标是演示如何在强化学习训练循环中安全地评估不受信任的 LLM 生成的代码。我们通过将编排平面 (Ray) 与执行平面 (GKE Agent Sandbox) 分离来实现这一目标。
强化学习代码评估的技术挑战
使用强化学习训练 LLM 代理时（例如，通过评估模型在单元测试中的输出来训练模型以编写代码），训练循环必须并行执行数千个不受信任的 LLM 生成的 Python 脚本。这会带来严重挑战：
- Pod 抖动瓶颈：传统评估框架会为每个任务启动一个全新的 Docker 容器。在 RL 训练循环期间，为数百个并行推出动态执行此操作会导致 Kubernetes 控制平面上的负载过重。这种延迟使得高频 RL 训练无法实现。
- 安全风险：在标准容器运行时中运行任意 LLM 生成的代码会共享宿主操作系统内核。一个逃逸漏洞就可能导致节点遭到入侵。
- IAM 令牌盗窃：在 Kubernetes pod 内运行的 LLM 生成的代码可以查询云提供商的元数据服务器，以窃取节点 IAM 服务账号令牌。
解决方案：解耦编排和执行
此架构将编排与执行分离：
- 编排程序 (Ray)：分布式 Ray 集群管理 RL 训练循环并分配推出生成任务。
- 执行平面 (GKE Agent Sandbox)：Ray worker 不会动态创建 Kubernetes pod，而是向专用 Sandbox 路由器发出简单的 HTTP 调用。路由器会立即为工作器分配一个在 gVisor (GKE Sandbox) 下运行的隔离的预热容器。
- 亚秒级延迟时间：由于沙盒在受管理的
SandboxWarmPool中预热，并通过高速 HTTP 网关进行管理，因此环境创建时间缩短至不到 200 毫秒，完全绕过 Kubernetes 控制平面。
实验目标
在此 Codelab 中，您将学习：
- 在强化学习循环中评估不可信代码的架构挑战和解决方案。
- 如何构建自定义沙盒映像以实现高效发布。
- 如何配置和使用 GKE 代理沙盒和 SandboxWarmPool。
- 如何安全地隔离沙盒以防止 IAM 令牌被盗。
- 如何使用 Ray 将编排与执行分离，通过 SweBench 和 TRL 运行基本的 RL 训练作业。
2. 集群创建和前提条件
继续操作之前，您需要一个 GKE 集群，其中包含高性能 GPU 节点池，并且已安装 Ray Operator 来管理训练工作负载。
前提条件
本 Codelab 假定您已安装并配置以下工具：
- Google Cloud SDK (
gcloud) - Docker（在本地构建自定义映像时需要）
kubectl
环境变量
首先，设置将在整个 Codelab 中使用的环境变量。以下命令使用合理的默认值，但您可以根据需要更改这些值，以匹配您的特定 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"
创建 Artifact Registry 制品库以存放自定义容器映像：
gcloud artifacts repositories create $REPO_NAME \
--repository-format=docker \
--location=$REGION \
--description="Repository for RL Sandbox images"
集群配置
如需详细了解如何预配针对 AI 工作负载优化的 GKE 集群（包括 GPUDirect RDMA 网络布线），请参阅官方文档：创建 GKE AI 超算自定义集群
重要前提条件：创建集群或特定执行节点池时，请务必传递 --enable-agent-sandbox 和 --sandbox type=gvisor 标志，以安装沙盒暖池所需的自定义资源定义 (CRD)。
假设您的集群、GPU 和 Ray Operator 正在运行，以下内容详细介绍了如何配置执行平面并运行 RL 循环。
3. 构建自定义映像
运行高性能 RL 的一个关键方面是将依赖项烘焙到映像中。我们需要两个不同的映像：一个用于运行模型的 GPU worker，另一个用于运行不受信任的评估代码的隔离沙盒。
1. 构建 GPU 工作器映像
Ray GPU 工作器需要一些库来运行语言模型并编排训练循环。我们在官方 vLLM 映像的基础上构建此映像，因此它支持最新的 GPU，并且预安装了 PyTorch/CUDA。
运行以下命令以创建 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
构建映像并将其推送到 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
注意：本指南使用本地 docker 命令来构建映像。如果您希望远程构建映像，可以改用 Cloud Build（例如使用 gcloud builds submit）。
2. 构建 CPU 头像
Ray 头节点仅用于编排集群，不运行繁重的 GPU 训练模型。为避免标准 CPU 节点上出现大规模映像拉取瓶颈（通常为 15 GB 以上），我们为头节点构建了一个轻量级、仅限 CPU 的映像。此映像包含 Ray 和必需的 Python 库，但不包含 CUDA 和 vLLM 等占用大量资源的 GPU 库。
运行以下命令以创建 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
构建并推送映像：
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. 构建沙盒映像
沙盒需要我们评估的任务的特定依赖项，以便运行时安装能够立即完成。在此 Codelab 中，我们将使用 SWE-bench 中 django/django 代码库的问题。我们将预先克隆代码库并预先构建 Python 环境，这样模型脚本就不会在 RL 循环中浪费时间下载它们。
运行以下命令以创建 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
构建并推送映像：
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. 配置编排和执行
现在，我们部署用于编排的 Ray 集群和用于执行的沙盒资源。
1. Ray 集群配置
部署 RayCluster 自定义资源。请注意，集群的可用资源（例如内存、CPU 或 GPU 类型）可能有所不同。相应地调整 resources 请求和限制。
运行以下命令以创建 raycluster.yaml。此命令使用 cat << EOF 将环境变量自动替换到清单中：
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
应用该政策：
kubectl apply -f raycluster.yaml
验证集群是否已创建且正在运行（这可能需要几分钟的时间）：
kubectl get raycluster
预期输出：
NAME DESIRED WORKERS AVAILABLE WORKERS CPUS MEMORY GPUS STATUS AGE rl-cluster 1 1 ready 2m
2. SandboxRouter 配置
SandboxRouter 充当高速 HTTP 网关，可处理来自 Ray worker 的请求，并立即将这些请求桥接到可用的 gVisor pod，从而绕过速度较慢的 Kubernetes API 服务器 pod 生命周期。
运行以下命令以创建 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
应用该政策：
kubectl apply -f sandbox_router.yaml
验证部署是否正在运行：
kubectl get deployment sandbox-router-deployment
预期输出：
NAME READY UP-TO-DATE AVAILABLE AGE sandbox-router-deployment 2/2 2 2 1m
3. SandboxTemplate 和 WarmPool 配置
借助 GKE Agent Sandbox，您可以使用 Sandbox 路由器立即分配隔离的预热容器。我们定义了 SandboxTemplate 和 SandboxWarmPool，以确保 Pod 处于就绪状态。
运行以下命令，使用环境变量创建 sandbox_warmpool.yaml：
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
应用该政策：
kubectl apply -f sandbox_warmpool.yaml
验证 SandboxWarmPool 是否已初始化：
kubectl get sandboxwarmpool
预期输出：
NAME READY AGE swe-bench-django-warmpool 10 1m
4. 安全隔离
NetworkPolicy 可严格隔离沙盒，防止出站流量流向 GCP 元数据服务器，从而防止 IAM 令牌被盗。
运行以下命令以创建 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
应用政策：
kubectl apply -f network_policy.yaml
验证 NetworkPolicy 是否已创建：
kubectl get networkpolicy
预期输出：
NAME POD-SELECTOR AGE block-metadata-egress sandbox.gke.io/runtime=gvisor 1m
5. 使用 SweBench 和 TRL 的基本 RL 作业
集群和沙盒准备就绪后，我们就可以运行 GRPO 训练循环了。我们将使用 trl 库来编排 GRPO 算法，并使用 Ray 远程函数来评估隔离沙盒中生成的代码。
为了加快此 Codelab 的执行速度，我们将过滤到单个 Django 问题。以下路由逻辑展示了如何为不同的代码库选择不同的暖池，这在扩展到完整的 SWE-bench 数据集时非常有用。
训练脚本
运行以下命令以创建 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
将作业提交到集群
首先，通过端口转发到 Ray Head 信息中心，然后从本地机器提交训练作业：
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
监控跑步活动
您可以监控跑步进度：
- Ray 信息中心：在浏览器中打开
http://localhost:8265。 - 沙盒声明：观看 GKE 在 gVisor 下动态声明和释放沙盒：
watch -n 1 "kubectl get sandboxclaims,sandboxes,pods"
6. 总结
恭喜！您已成功使用 GKE Agent Sandbox 在 GKE Standard 上安全地配置并执行了高性能分布式 RL 训练循环。