1. Введение
В этой лабораторной работе подробно описано, как создать, настроить и запустить высокопроизводительный распределенный цикл обучения с подкреплением (RL) на GKE Standard с использованием песочниц агентов GKE (gVisor) и алгоритма групповой относительной оптимизации политики (GRPO) с библиотекой trl .
Цель состоит в том, чтобы продемонстрировать, как безопасно оценивать ненадежный код, сгенерированный LLM, во время цикла обучения RL. Мы достигаем этого, разделяя плоскость оркестрации (Ray) от плоскости выполнения (песочницы агентов GKE).
Технические сложности оценки кода с подкреплением
При обучении агентов LLM с использованием обучения с подкреплением (например, обучении модели написанию кода путем оценки ее выходных данных в модульных тестах) цикл обучения должен параллельно выполнять тысячи ненадежных, сгенерированных LLM скриптов на Python. Это создает серьезные проблемы:
- Проблема "узкого места" при создании подов: традиционные фреймворки оценки запускают новый контейнер Docker для каждой задачи. Динамическое выполнение этого процесса для сотен параллельных развертываний во время цикла обучения с подкреплением создает серьезную нагрузку на плоскость управления Kubernetes. Задержка делает высокочастотное обучение с подкреплением невозможным.
- Риск безопасности: выполнение произвольного кода, сгенерированного LLM, внутри стандартных сред выполнения контейнеров использует ядро хостовой операционной системы. Одна-единственная уязвимость может поставить под угрозу безопасность ваших узлов.
- Кража токенов IAM: сгенерированный LLM-код, работающий внутри пода Kubernetes, может запрашивать у облачного провайдера метаданные для кражи токенов учетных записей служб IAM узлов.
Решение: Разделение оркестровки и выполнения
Эта архитектура разделяет процессы оркестровки и выполнения :
- Оркестратор (Ray): Распределенный кластер Ray управляет циклом обучения с подкреплением и распределяет генерацию развертываний.
- Плоскость выполнения (песочница агента GKE): Вместо динамического создания подов Kubernetes, рабочие процессы Ray выполняют простые HTTP-запросы к выделенному маршрутизатору песочницы . Маршрутизатор мгновенно назначает рабочему процессу изолированный, предварительно подготовленный контейнер, работающий под управлением gVisor (песочница GKE) .
- Задержка менее секунды: поскольку песочницы предварительно прогреваются в управляемом
SandboxWarmPoolи управляются через высокоскоростной HTTP-шлюз, время создания среды сокращается до менее чем 200 мс , полностью минуя плоскость управления Kubernetes.
Цели лабораторной работы
В этом практическом занятии вы узнаете:
- Архитектурные проблемы и решения для оценки ненадежного кода в циклах обучения с подкреплением.
- Как создавать пользовательские образы для песочницы для эффективного развертывания.
- Как настроить и использовать песочницы агентов GKE и пулы теплых песочниц.
- Как безопасно изолировать песочницы, чтобы предотвратить кражу токенов IAM.
- Как запустить базовое обучение с подкреплением с использованием SweBench и TRL, применяя Ray для разделения процессов оркестровки и выполнения.
2. Создание кластера и необходимые условия
Прежде чем продолжить, вам потребуется кластер GKE с высокопроизводительным пулом узлов GPU и установленный Ray Operator для управления нагрузкой обучения.
Предварительные требования
В этом практическом занятии предполагается, что следующие инструменты установлены и настроены:
- Google Cloud SDK (
gcloud) - Docker (необходим для локальной сборки пользовательских образов)
-
kubectl
Переменные окружающей среды
Сначала задайте переменные среды, которые будут использоваться на протяжении всего этого практического занятия. Приведенные ниже команды используют разумные значения по умолчанию, но вы можете изменить их по мере необходимости, чтобы они соответствовали вашей конкретной среде 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"
Создайте репозиторий в реестре артефактов для хранения пользовательских образов контейнеров:
gcloud artifacts repositories create $REPO_NAME \
--repository-format=docker \
--location=$REGION \
--description="Repository for RL Sandbox images"
Конфигурация кластера
Для получения полной пошаговой инструкции по настройке кластера GKE, оптимизированного для рабочих нагрузок ИИ (включая сетевую инфраструктуру GPUDirect RDMA), следуйте официальной документации: Создание пользовательского кластера GKE AI Hypercompute.
Важное предварительное условие: при создании кластера или пула конкретных узлов выполнения убедитесь, что вы передали флаги --enable-agent-sandbox и --sandbox type=gvisor для установки необходимых определений пользовательских ресурсов (CRD) для пулов "теплых" песочниц.
Предполагая, что ваш кластер, графические процессоры и оператор лучей запущены, ниже приведено подробное описание настройки плоскости выполнения и запуска цикла RL.
3. Создание пользовательских образов
Ключевым аспектом высокопроизводительного обучения с подкреплением является закладка зависимостей в ваши образы. Нам нужны два разных образа: один для рабочих процессов на графических процессорах, выполняющих модель, и один для изолированных песочниц, выполняющих ненадежный код оценки.
1. Создайте образ GPU Worker.
Для работы Ray GPU-процесса требуются библиотеки для запуска языковой модели и организации цикла обучения. Мы создали этот образ на основе официального образа vLLM, поэтому он поддерживает новейшие графические процессоры и имеет предустановленные 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. Создайте образ головки ЦП.
Головной узел Ray только управляет кластером и не запускает ресурсоемкие модели обучения на графическом процессоре. Чтобы избежать чрезмерного объема загружаемых изображений (обычно более 15 ГБ) на стандартных узлах с центральным процессором, мы создаем облегченный образ только для центрального узла. Этот образ содержит Ray и необходимые библиотеки Python, но исключает ресурсоемкие библиотеки для графических процессоров, такие как CUDA и vLLM.
Выполните следующую команду, чтобы создать 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. Создайте образ песочницы.
Для выполнения задачи, которую мы оцениваем, в песочнице необходимы специфические зависимости, чтобы установка во время выполнения происходила мгновенно. Для этого практического занятия мы будем использовать задачу из репозитория django/django в SWE-bench. Мы предварительно клонируем репозиторий и предварительно соберем среды Python, чтобы наши скрипты модели не тратили время на их загрузку в цикле обучения с подкреплением.
Выполните следующую команду, чтобы создать 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. Конфигурация кластера лучей
Разверните пользовательский ресурс RayCluster. Обратите внимание, что доступные ресурсы вашего кластера (например, память, ЦП или тип графического процессора) могут отличаться. Соответственно скорректируйте запросы и ограничения 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 и мгновенно перенаправляя их к доступным подам gVisor, минуя более медленный жизненный цикл подов API-сервера Kubernetes.
Выполните следующую команду для создания файла 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. Настройка шаблона песочницы и теплого бассейна
GKE Agent Sandbox позволяет мгновенно назначать изолированные, предварительно подготовленные контейнеры с помощью Sandbox Router. Мы определяем SandboxTemplate и SandboxWarmPool для поддержания готовности подов.
Выполните следующую команду, чтобы создать файл 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
Убедитесь, что политика сети создана:
kubectl get networkpolicy
Ожидаемый результат:
NAME POD-SELECTOR AGE block-metadata-egress sandbox.gke.io/runtime=gvisor 1m
5. Базовая задача обучения с подкреплением с использованием SweBench и TRL.
После подготовки кластера и песочниц мы можем запустить цикл обучения GRPO. Мы будем использовать библиотеку trl для управления алгоритмом GRPO, а также функции удаленного доступа Ray для оценки сгенерированного кода внутри изолированных песочниц.
Чтобы ускорить выполнение этого практического задания, мы ограничимся одной задачей 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 Standard, используя песочницы агентов GKE.