1. บทนำ
Lab นี้อธิบายรายละเอียดเกี่ยวกับวิธีสร้าง จัดเตรียม และเรียกใช้ลูปการฝึกการเรียนรู้แบบเสริมกำลัง (RL) แบบกระจายที่มีประสิทธิภาพสูงบน GKE Standard ด้วย GKE Agent Sandbox (gVisor) โดยใช้อัลกอริทึม Group Relative Policy Optimization (GRPO) กับไลบรารี trl
เป้าหมายคือการแสดงวิธีประเมินโค้ดที่ LLM สร้างขึ้นซึ่งไม่น่าเชื่อถืออย่างปลอดภัยในระหว่างลูปการฝึก RL เราบรรลุเป้าหมายนี้ได้โดยการแยกระนาบการจัดการเป็นกลุ่ม (Ray) ออกจากระนาบการดำเนินการ (GKE Agent Sandbox)
ความท้าทายทางเทคนิคของการประเมินโค้ด RL
เมื่อฝึก Agent ของ LLM โดยใช้การเรียนรู้แบบเสริมกำลัง (เช่น การฝึกโมเดลให้เขียนโค้ดโดยการประเมินเอาต์พุตในการทดสอบหน่วย) ลูปการฝึกต้องเรียกใช้สคริปต์ Python ที่ LLM สร้างขึ้นซึ่งไม่น่าเชื่อถือหลายพันรายการแบบขนาน ซึ่งก่อให้เกิดความท้าทายที่สำคัญดังนี้
- คอขวดของการหมุนเวียนพ็อด: เฟรมเวิร์กการประเมินแบบเดิมจะสร้างคอนเทนเนอร์ Docker ใหม่สำหรับแต่ละงาน การดำเนินการนี้แบบไดนามิกสำหรับการเปิดตัวแบบขนานหลายร้อยรายการในระหว่างลูปการฝึก RL จะทำให้เกิดการโหลดอย่างรุนแรงในระนาบการควบคุม Kubernetes เวลาในการตอบสนองทำให้การฝึก RL ที่มีความถี่สูงเป็นไปไม่ได้
- ความเสี่ยงด้านความปลอดภัย: การเรียกใช้โค้ดที่ LLM สร้างขึ้นโดยพลการภายในรันไทม์ของคอนเทนเนอร์มาตรฐานจะแชร์เคอร์เนลของระบบปฏิบัติการโฮสต์ ช่องโหว่ในการหลบหนีเพียงครั้งเดียวอาจทำให้โหนดของคุณตกอยู่ในความเสี่ยง
- การขโมยโทเค็น IAM: โค้ดที่ LLM สร้างขึ้นซึ่งทำงานภายในพ็อด Kubernetes สามารถค้นหาเซิร์ฟเวอร์ข้อมูลเมตาของผู้ให้บริการระบบคลาวด์เพื่อขโมยโทเค็นบัญชีบริการ IAM ของโหนด
โซลูชัน: การจัดการเป็นกลุ่มและการดำเนินการที่แยกออกจากกัน
สถาปัตยกรรมนี้แยกการจัดระเบียบจากการดำเนินการ ดังนี้
- ตัวจัดระเบียบ (Ray): คลัสเตอร์ Ray แบบกระจายจะจัดการลูปการฝึก RL และกระจายการสร้างการเปิดตัว
- ระนาบการดำเนินการ (GKE Agent Sandbox): แทนที่จะสร้างพ็อด Kubernetes แบบไดนามิก Worker ของ Ray จะเรียกใช้ HTTP อย่างง่ายไปยัง เราเตอร์แซนด์บ็อกซ์ เฉพาะ เราเตอร์จะกำหนดคอนเทนเนอร์ที่แยกต่างหากซึ่งอุ่นเครื่องไว้ล่วงหน้าซึ่งทำงานภายใต้ gVisor (GKE Sandbox) ให้กับ Worker ทันที
- เวลาในการตอบสนองต่ำกว่า 1 วินาที: เนื่องจากแซนด์บ็อกซ์ได้รับการอุ่นเครื่องไว้ล่วงหน้าใน
SandboxWarmPoolที่มีการจัดการและจัดการผ่านเกตเวย์ HTTP ความเร็วสูง การสร้างสภาพแวดล้อมจึงลดลงเหลือต่ำกว่า 200 มิลลิวินาที ซึ่งข้ามระนาบการควบคุม Kubernetes ไปโดยสิ้นเชิง
วัตถุประสงค์ของ Lab
ใน Codelab นี้ คุณจะได้เรียนรู้สิ่งต่อไปนี้
- ความท้าทายและโซลูชันด้านสถาปัตยกรรมสำหรับการประเมินโค้ดที่ไม่น่าเชื่อถือในลูป RL
- วิธีสร้างอิมเมจแซนด์บ็อกซ์ที่กำหนดเองสำหรับการเปิดตัวที่มีประสิทธิภาพ
- วิธีกำหนดค่าและใช้ GKE Agent Sandbox และ SandboxWarmPool
- วิธีแยกแซนด์บ็อกซ์อย่างปลอดภัยเพื่อป้องกันการขโมยโทเค็น IAM
- วิธีเรียกใช้งานการฝึก RL ขั้นพื้นฐานด้วย SweBench และ TRL โดยใช้ Ray เพื่อแยกการจัดการเป็นกลุ่มจากการดำเนินการ
2. การสร้างคลัสเตอร์และข้อกำหนดเบื้องต้น
ก่อนดำเนินการต่อ คุณต้องมีคลัสเตอร์ GKE ที่มี Node Pool 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"
การกำหนดค่าคลัสเตอร์
หากต้องการดูคำแนะนำแบบละเอียดเกี่ยวกับการจัดเตรียมคลัสเตอร์ GKE ที่ปรับให้เหมาะกับภาระงาน AI (รวมถึงการเชื่อมต่อเครือข่าย GPUDirect RDMA) ให้ทำตามเอกสารอย่างเป็นทางการ: สร้างคลัสเตอร์ที่กำหนดเองสำหรับ GKE AI Hypercompute
ข้อกำหนดเบื้องต้นที่สำคัญ: เมื่อสร้างคลัสเตอร์หรือ Node Pool การดำเนินการที่เฉพาะเจาะจง ให้ส่งแฟล็ก --enable-agent-sandbox และ --sandbox type=gvisor เพื่อติดตั้งคำจำกัดความทรัพยากรที่กำหนดเอง (CRD) ที่จำเป็นสำหรับ Warm Pool ของแซนด์บ็อกซ์
สมมติว่าคลัสเตอร์ GPU และ Ray Operator ทำงานอยู่ ทุกอย่างด้านล่างนี้จะอธิบายรายละเอียดเกี่ยวกับวิธีกำหนดค่าระนาบการดำเนินการและเรียกใช้ลูป RL
3. สร้างอิมเมจที่กำหนดเอง
สิ่งสำคัญอย่างหนึ่งในการเรียกใช้ RL ประสิทธิภาพสูงคือการรวมทรัพยากร Dependency ไว้ในอิมเมจ เราต้องมีอิมเมจที่แตกต่างกัน 2 รายการ ได้แก่ อิมเมจสำหรับ Worker ของ GPU ที่เรียกใช้โมเดล และอิมเมจสำหรับแซนด์บ็อกซ์ที่แยกต่างหากซึ่งเรียกใช้โค้ดการประเมินที่ไม่น่าเชื่อถือ
1. สร้างอิมเมจ Worker ของ GPU
Worker ของ GPU ของ Ray ต้องมีไลบรารีเพื่อเรียกใช้โมเดลภาษาและจัดระเบียบลูปการฝึก เราสร้างอิมเมจนี้โดยอิงตามอิมเมจ 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. สร้างอิมเมจ Head ของ CPU
โหนด Head ของ Ray จะจัดระเบียบคลัสเตอร์เท่านั้นและไม่ได้เรียกใช้โมเดลการฝึก GPU ที่มีขนาดใหญ่ เพื่อหลีกเลี่ยงคอขวดขนาดใหญ่ในการดึงอิมเมจ (โดยปกติคือ 15 GB ขึ้นไป) ในโหนด CPU มาตรฐาน เราจะสร้างอิมเมจที่มีน้ำหนักเบาซึ่งใช้ CPU เท่านั้นสำหรับโหนด Head อิมเมจนี้มี Ray และไลบรารี Python ที่จำเป็น แต่ไม่รวมไลบรารี GPU ขนาดใหญ่ เช่น 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. สร้างอิมเมจแซนด์บ็อกซ์
แซนด์บ็อกซ์ต้องมีทรัพยากร Dependency ที่เฉพาะเจาะจงสำหรับงานที่เรากำลังประเมินเพื่อให้การติดตั้งรันไทม์เป็นไปอย่างรวดเร็ว สำหรับ Codelab นี้ เราจะใช้ปัญหาจากที่เก็บ django/django ใน SWE-bench เราจะโคลนที่เก็บไว้ล่วงหน้าและสร้างสภาพแวดล้อม 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 เพื่อแทนที่ตัวแปรสภาพแวดล้อมลงในไฟล์ 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
ใช้ไฟล์ Manifest
kubectl apply -f raycluster.yaml
ยืนยันว่าคลัสเตอร์ถูกสร้างขึ้นและทำงานอยู่ (ขั้นตอนนี้อาจใช้เวลา 2-3 นาที)
kubectl get raycluster
เอาต์พุตที่คาดไว้
NAME DESIRED WORKERS AVAILABLE WORKERS CPUS MEMORY GPUS STATUS AGE rl-cluster 1 1 ready 2m
2. การกำหนดค่า SandboxRouter
SandboxRouter ทำหน้าที่เป็นเกตเวย์ HTTP ความเร็วสูง โดยรับคำขอจาก Worker ของ 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
ใช้ไฟล์ Manifest
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 อนุญาตให้กำหนดคอนเทนเนอร์ที่แยกต่างหากซึ่งอุ่นเครื่องไว้ล่วงหน้าได้ทันทีโดยใช้เราเตอร์แซนด์บ็อกซ์ เรากำหนด 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
ใช้ไฟล์ Manifest
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. งาน RL ขั้นพื้นฐานด้วย SweBench และ TRL
เมื่อคลัสเตอร์และแซนด์บ็อกซ์พร้อมแล้ว เราจะเรียกใช้ลูปการฝึก GRPO เราจะใช้ไลบรารี trl เพื่อจัดระเบียบอัลกอริทึม GRPO และฟังก์ชันระยะไกลของ Ray เพื่อประเมินโค้ดที่สร้างขึ้นภายในแซนด์บ็อกซ์ที่แยกต่างหาก
เพื่อให้การดำเนินการเป็นไปอย่างรวดเร็วสำหรับ Codelab นี้ เราจะกรองให้เหลือเพียงปัญหาเดียวของ Django ตรรกะการกำหนดเส้นทางด้านล่างแสดงวิธีเลือก Warm Pool ที่แตกต่างกันสำหรับที่เก็บที่แตกต่างกัน ซึ่งจะเป็นประโยชน์เมื่อขยายไปยังชุดข้อมูล 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
ส่งงานไปยังคลัสเตอร์
ขั้นแรก ให้ส่งต่อพอร์ตไปยังแดชบอร์ด Head ของ Ray แล้วส่งงานการฝึกจากเครื่องในเครื่อง
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. บทสรุป
ยินดีด้วย คุณกำหนดค่าและเรียกใช้ลูปการฝึก RL แบบกระจายที่มีประสิทธิภาพสูงบน GKE Standard โดยใช้ GKE Agent Sandbox ได้อย่างปลอดภัย