1. Giới thiệu
Lớp học lập trình này trình bày chi tiết cách xây dựng, cung cấp và thực thi một vòng lặp huấn luyện Học tăng cường (RL) phân tán hiệu suất cao trên GKE Standard bằng Hộp cát tác nhân GKE (gVisor), sử dụng thuật toán Tối ưu hoá chính sách tương đối theo nhóm (GRPO) với thư viện trl.
Mục tiêu là minh hoạ cách đánh giá an toàn mã không đáng tin cậy do LLM tạo trong vòng lặp huấn luyện RL. Chúng tôi đạt được mục tiêu này bằng cách tách riêng mặt phẳng điều phối (Ray) khỏi mặt phẳng thực thi (Hộp cát tác nhân GKE).
Thử thách kỹ thuật khi đánh giá mã RL
Khi huấn luyện các tác nhân LLM bằng Học tăng cường (ví dụ: huấn luyện một mô hình viết mã bằng cách đánh giá đầu ra của mô hình đó trên các bài kiểm thử đơn vị), vòng lặp huấn luyện phải thực thi hàng nghìn tập lệnh Python không đáng tin cậy do LLM tạo song song. Điều này đặt ra những thử thách quan trọng:
- Nút thắt cổ chai của Pod Churn: Các khung đánh giá truyền thống khởi động một vùng chứa Docker mới cho mỗi tác vụ. Việc thực hiện thao tác này một cách linh động cho hàng trăm lần triển khai song song trong vòng lặp huấn luyện RL gây ra tải trọng lớn cho mặt phẳng điều khiển Kubernetes. Độ trễ khiến việc huấn luyện RL tần suất cao trở nên không thể thực hiện được.
- Rủi ro bảo mật: Việc chạy mã tuỳ ý do LLM tạo bên trong thời gian chạy vùng chứa tiêu chuẩn sẽ chia sẻ nhân hệ điều hành máy chủ. Một lỗ hổng thoát duy nhất có thể làm tổn hại đến các nút của bạn.
- Đánh cắp mã thông báo IAM: Mã do LLM tạo chạy bên trong một pod Kubernetes có thể truy vấn máy chủ siêu dữ liệu của nhà cung cấp dịch vụ đám mây để đánh cắp mã thông báo tài khoản dịch vụ IAM của nút.
Giải pháp: Tách riêng việc điều phối và thực thi
Kiến trúc này tách riêng việc điều phối khỏi việc thực thi:
- Trình điều phối (Ray): Một cụm Ray phân tán quản lý vòng lặp huấn luyện RL và phân phối quá trình tạo bản triển khai.
- Mặt phẳng thực thi (Hộp cát tác nhân GKE): Thay vì tạo các pod Kubernetes một cách linh động, các worker Ray sẽ thực hiện các lệnh gọi HTTP đơn giản đến một Bộ định tuyến hộp cát chuyên dụng. Bộ định tuyến sẽ ngay lập tức chỉ định cho worker một vùng chứa được cách ly và khởi động trước chạy trong GKE Sandbox (Hộp cát GKE).
- Độ trễ dưới một giây: Vì các hộp cát được khởi động trước trong
SandboxWarmPoolđược quản lý và được quản lý thông qua một cổng HTTP tốc độ cao, nên quá trình tạo môi trường sẽ giảm xuống dưới 200 mili giây, hoàn toàn bỏ qua mặt phẳng điều khiển Kubernetes.
Mục tiêu của lớp học lập trình
Trong lớp học lập trình này, bạn sẽ tìm hiểu:
- Các thử thách và giải pháp về kiến trúc để đánh giá mã không đáng tin cậy trong các vòng lặp RL.
- Cách tạo hình ảnh hộp cát tuỳ chỉnh để triển khai hiệu quả.
- Cách định cấu hình và sử dụng Hộp cát tác nhân GKE và SandboxWarmPools.
- Cách cách ly an toàn các hộp cát để ngăn chặn việc đánh cắp mã thông báo IAM.
- Cách chạy một công việc huấn luyện RL cơ bản với SweBench và TRL bằng Ray để tách riêng việc điều phối khỏi việc thực thi.
2. Tạo cụm và điều kiện tiên quyết
Trước khi tiếp tục, bạn cần có một cụm GKE có bộ nút GPU hiệu suất cao và Trình điều khiển Ray được cài đặt để quản lý khối lượng công việc huấn luyện.
Điều kiện tiên quyết
Lớp học lập trình này giả định rằng các công cụ sau đã được cài đặt và định cấu hình:
- SDK Google Cloud (
gcloud) - Docker (bắt buộc để tạo hình ảnh tuỳ chỉnh cục bộ)
kubectl
Biến môi trường
Trước tiên, hãy thiết lập các biến môi trường sẽ được sử dụng trong suốt lớp học lập trình này. Các lệnh bên dưới sử dụng các giá trị mặc định hợp lý, nhưng bạn có thể thay đổi các giá trị này theo nhu cầu để phù hợp với môi trường Google Cloud cụ thể của mình:
export PROJECT_ID=$(gcloud config get-value project)
export REGION="us-west3"
export ZONE="us-west3-a"
export REPO_NAME="rl-sandbox-repo"
Tạo kho lưu trữ Artifact Registry để lưu giữ các hình ảnh vùng chứa tuỳ chỉnh:
gcloud artifacts repositories create $REPO_NAME \
--repository-format=docker \
--location=$REGION \
--description="Repository for RL Sandbox images"
Cấu hình cụm
Để xem hướng dẫn đầy đủ về cách cung cấp một cụm GKE được tối ưu hoá cho khối lượng công việc AI (bao gồm cả hệ thống dây mạng GPUDirect RDMA), hãy xem tài liệu chính thức: Tạo cụm tuỳ chỉnh GKE AI Hypercompute
Điều kiện tiên quyết quan trọng: Khi tạo cụm hoặc bộ nút thực thi cụ thể, hãy đảm bảo bạn truyền các cờ --enable-agent-sandbox và --sandbox type=gvisor để cài đặt các Định nghĩa tài nguyên tuỳ chỉnh (CRD) cần thiết cho các nhóm khởi động trước hộp cát.
Giả sử cụm, GPU và Trình điều khiển Ray của bạn đang chạy, mọi thông tin bên dưới đều trình bày chi tiết cách định cấu hình mặt phẳng thực thi và chạy vòng lặp RL.
3. Tạo hình ảnh tuỳ chỉnh
Một khía cạnh quan trọng của việc chạy RL hiệu suất cao là đưa các phần phụ thuộc vào hình ảnh của bạn. Chúng ta cần 2 hình ảnh riêng biệt: một cho các worker GPU chạy mô hình và một cho các hộp cát được cách ly chạy mã đánh giá không đáng tin cậy.
1. Tạo hình ảnh worker GPU
Worker GPU Ray cần các thư viện để chạy mô hình ngôn ngữ và điều phối vòng lặp huấn luyện. Chúng ta tạo hình ảnh này dựa trên hình ảnh vLLM chính thức để hình ảnh này hỗ trợ các GPU mới nhất và đã cài đặt sẵn PyTorch/CUDA.
Chạy lệnh sau để tạo 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
Tạo và đẩy hình ảnh vào kho lưu trữ 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
Lưu ý: Hướng dẫn này sử dụng các lệnh docker cục bộ để tạo hình ảnh. Nếu muốn tạo hình ảnh từ xa, bạn có thể sử dụng Cloud Build (ví dụ: sử dụng gcloud builds submit).
2. Tạo hình ảnh đầu CPU
Nút đầu Ray chỉ điều phối cụm và không chạy các mô hình huấn luyện GPU nặng. Để tránh nút thắt cổ chai kéo hình ảnh lớn (thường là 15 GB trở lên) trên các nút CPU tiêu chuẩn, chúng ta sẽ tạo một hình ảnh chỉ dành cho CPU và có dung lượng nhẹ cho nút đầu. Hình ảnh này chứa Ray và các thư viện Python cần thiết, nhưng không bao gồm các thư viện GPU nặng như CUDA và vLLM.
Chạy lệnh sau để tạo 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
Tạo và đẩy hình ảnh:
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. Tạo hình ảnh hộp cát
Hộp cát cần các phần phụ thuộc cụ thể cho tác vụ mà chúng ta đang đánh giá để quá trình cài đặt thời gian chạy diễn ra ngay lập tức. Trong lớp học lập trình này, chúng ta sẽ sử dụng một vấn đề từ kho lưu trữ django/django trong SWE-bench. Chúng ta sẽ sao chép trước kho lưu trữ và tạo trước các môi trường python để tập lệnh mô hình không lãng phí thời gian tải các môi trường này xuống trong vòng lặp RL.
Chạy lệnh sau để tạo 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
Tạo và đẩy hình ảnh:
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. Định cấu hình việc điều phối và thực thi
Bây giờ, chúng ta sẽ triển khai cụm Ray để điều phối và các tài nguyên Hộp cát để thực thi.
1. Cấu hình cụm Ray
Triển khai tài nguyên tuỳ chỉnh RayCluster. Xin lưu ý rằng các tài nguyên hiện có của cụm (chẳng hạn như bộ nhớ, CPU hoặc loại GPU) có thể khác nhau. Điều chỉnh các yêu cầu và giới hạn resources cho phù hợp.
Chạy lệnh sau để tạo raycluster.yaml. Lệnh này sử dụng cat << EOF để tự động thay thế các biến môi trường của bạn vào tệp kê khai:
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
Áp dụng tệp kê khai này:
kubectl apply -f raycluster.yaml
Xác minh rằng cụm đã được tạo và đang chạy (quá trình này có thể mất vài phút):
kubectl get raycluster
Kết quả đầu ra dự kiến:
NAME DESIRED WORKERS AVAILABLE WORKERS CPUS MEMORY GPUS STATUS AGE rl-cluster 1 1 ready 2m
2. Cấu hình SandboxRouter
SandboxRouter hoạt động như một cổng HTTP tốc độ cao, xử lý các yêu cầu từ các worker Ray và kết nối các yêu cầu này ngay lập tức với các pod gVisor hiện có, bỏ qua vòng đời pod máy chủ API Kubernetes chậm hơn.
Chạy lệnh sau để tạo 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
Áp dụng tệp kê khai này:
kubectl apply -f sandbox_router.yaml
Xác minh rằng quá trình triển khai đang chạy:
kubectl get deployment sandbox-router-deployment
Kết quả đầu ra dự kiến:
NAME READY UP-TO-DATE AVAILABLE AGE sandbox-router-deployment 2/2 2 2 1m
3. Cấu hình SandboxTemplate và WarmPool
Hộp cát tác nhân GKE cho phép chỉ định ngay lập tức các vùng chứa được cách ly và khởi động trước bằng Bộ định tuyến hộp cát. Chúng ta xác định SandboxTemplate và SandboxWarmPool để giữ cho các pod luôn sẵn sàng.
Chạy lệnh sau để tạo sandbox_warmpool.yaml bằng các biến môi trường của bạn:
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
Áp dụng tệp kê khai này:
kubectl apply -f sandbox_warmpool.yaml
Xác minh rằng SandboxWarmPool đã được khởi chạy:
kubectl get sandboxwarmpool
Kết quả đầu ra dự kiến:
NAME READY AGE swe-bench-django-warmpool 10 1m
4. Công tác chia cách bảo mật
NetworkPolicy cách ly nghiêm ngặt các hộp cát, ngăn chặn việc gửi dữ liệu ra Máy chủ siêu dữ liệu GCP, do đó ngăn chặn việc đánh cắp mã thông báo IAM.
Chạy lệnh sau để tạo 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
Áp dụng chính sách này:
kubectl apply -f network_policy.yaml
Xác minh rằng NetworkPolicy đã được tạo:
kubectl get networkpolicy
Kết quả đầu ra dự kiến:
NAME POD-SELECTOR AGE block-metadata-egress sandbox.gke.io/runtime=gvisor 1m
5. Công việc RL cơ bản với SweBench và TRL
Sau khi cụm và các hộp cát được chuẩn bị, chúng ta có thể chạy vòng lặp huấn luyện GRPO. Chúng ta sẽ sử dụng thư viện trl để điều phối thuật toán GRPO và các hàm từ xa của Ray để đánh giá mã được tạo bên trong các hộp cát được cách ly.
Để thực thi nhanh trong lớp học lập trình này, chúng ta sẽ lọc xuống một vấn đề Django duy nhất. Logic định tuyến bên dưới cho thấy cách bạn chọn các nhóm khởi động trước khác nhau cho các kho lưu trữ khác nhau. Điều này hữu ích khi mở rộng sang tập dữ liệu SWE-bench đầy đủ.
Tập lệnh huấn luyện
Chạy lệnh sau để tạo 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
Gửi công việc đến cụm
Trước tiên, hãy chuyển tiếp cổng đến trang tổng quan Ray Head và gửi công việc huấn luyện từ máy cục bộ của bạn:
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
Giám sát quá trình chạy
Bạn có thể giám sát tiến trình chạy:
- Trang tổng quan Ray: Mở
http://localhost:8265trong trình duyệt. - Yêu cầu hộp cát: Theo dõi cách GKE yêu cầu và phát hành các hộp cát một cách linh động trong gVisor:
watch -n 1 "kubectl get sandboxclaims,sandboxes,pods"
6. Kết luận
Xin chúc mừng! Bạn đã định cấu hình và thực thi thành công một vòng lặp huấn luyện RL phân tán hiệu suất cao một cách an toàn trên GKE Standard bằng Hộp cát tác nhân GKE.