1. परिचय
इस कोडलैब में, GKE Standard पर GKE एजेंट सैंडबॉक्स (gVisor) का इस्तेमाल करके, ग्रुप रिलेटिव पॉलिसी ऑप्टिमाइज़ेशन (जीआरपीओ) एल्गोरिदम के साथ trl लाइब्रेरी का इस्तेमाल करके, हाई-परफ़ॉर्मेंस डिस्ट्रिब्यूटेड रीइन्फ़ोर्समेंट लर्निंग (आरएल) ट्रेनिंग लूप बनाने, प्रोविज़निंग करने, और एक्ज़ीक्यूट करने का तरीका बताया गया है.
इसका मकसद, आरएल ट्रेनिंग लूप के दौरान, एलएलएम से जनरेट किए गए, ऐसे कोड का सुरक्षित तरीके से आकलन करना है जिस पर भरोसा नहीं किया जा सकता. हम ऑर्केस्ट्रेशन प्लेन (Ray) को एक्ज़ीक्यूशन प्लेन (GKE एजेंट सैंडबॉक्स) से अलग करके, यह काम करते हैं.
आरएल कोड के आकलन से जुड़ी तकनीकी चुनौती
एलएलएम एजेंट को रीइन्फ़ोर्समेंट लर्निंग का इस्तेमाल करके ट्रेन करते समय (उदाहरण के लिए, यूनिट टेस्ट पर उसके आउटपुट का आकलन करके, कोड लिखने के लिए मॉडल को ट्रेन करना), ट्रेनिंग लूप को एलएलएम से जनरेट की गई, ऐसी हज़ारों Python स्क्रिप्ट को एक साथ एक्ज़ीक्यूट करना होता है जिन पर भरोसा नहीं किया जा सकता. इससे ये अहम चुनौतियां सामने आती हैं:
- पॉड चर्न बॉटलनेक: पारंपरिक आकलन फ़्रेमवर्क, हर टास्क के लिए नया Docker कंटेनर स्पिन अप करते हैं. आरएल ट्रेनिंग लूप के दौरान, सैकड़ों पैरलल रोलआउट के लिए ऐसा डाइनैमिक तरीके से करने पर, Kubernetes कंट्रोल प्लेन पर बहुत ज़्यादा लोड पड़ता है. इस वजह से, हाई-फ़्रीक्वेंसी आरएल ट्रेनिंग करना मुमकिन नहीं हो पाता.
- सुरक्षा से जुड़ा जोखिम: एलएलएम से जनरेट किए गए किसी भी कोड को स्टैंडर्ड कंटेनर रनटाइम में चलाने पर, होस्ट ओएस कर्नल शेयर होता है. एक भी एस्केप वलनरेबिलिटी से आपके नोड से समझौता किया जा सकता है.
- आईएएम टोकन की चोरी: Kubernetes पॉड में चल रहा एलएलएम से जनरेट किया गया कोड, नोड आईएएम सेवा खाते के टोकन चुराने के लिए, क्लाउड सेवा देने वाली कंपनी के मेटाडेटा सर्वर से क्वेरी कर सकता है.
समाधान: ऑर्केस्ट्रेशन और एक्ज़ीक्यूशन को अलग करना
इस आर्किटेक्चर में, ऑर्केस्ट्रेशन को एक्ज़ीक्यूशन से अलग किया जाता है:
- ऑर्केस्ट्रेटर (Ray): डिस्ट्रिब्यूटेड Ray क्लस्टर, आरएल ट्रेनिंग लूप को मैनेज करता है और रोलआउट जनरेशन को डिस्ट्रिब्यूट करता है.
- एक्ज़ीक्यूशन प्लेन (GKE एजेंट सैंडबॉक्स): Kubernetes पॉड को डाइनैमिक तरीके से बनाने के बजाय, Ray वर्कर, खास सैंडबॉक्स राऊटर को एचटीटीपी कॉल करते हैं. राऊटर, वर्कर को तुरंत gVisor (GKE सैंडबॉक्स) में चल रहा, अलग किया गया, पहले से वॉर्म किया गया कंटेनर असाइन करता है.
- एक सेकंड से भी कम समय में होने वाली प्रोसेस: सैंडबॉक्स को मैनेज किए गए
SandboxWarmPoolमें पहले से वॉर्म किया जाता है. साथ ही, इन्हें हाई-स्पीड एचटीटीपी गेटवे के ज़रिए मैनेज किया जाता है. इसलिए, एनवायरमेंट बनाने में 200 मिसेकंड से भी कम समय लगता है. इससे Kubernetes कंट्रोल प्लेन पूरी तरह से बाईपास हो जाता है.
लैब के मकसद
इस कोडलैब में, आपको इनके बारे में जानकारी मिलेगी:
- आरएल लूप में, ऐसे कोड का आकलन करने से जुड़ी आर्किटेक्चरल चुनौतियां और उनके समाधान जिन पर भरोसा नहीं किया जा सकता.
- इमेज को बेहतर तरीके से रोल आउट करने के लिए, कस्टम सैंडबॉक्स इमेज बनाने का तरीका.
- GKE एजेंट सैंडबॉक्स और SandboxWarmPools को कॉन्फ़िगर और इस्तेमाल करने का तरीका.
- आईएएम टोकन की चोरी को रोकने के लिए, सैंडबॉक्स को सुरक्षित तरीके से अलग करने का तरीका.
- ऑर्केस्ट्रेशन को एक्ज़ीक्यूशन से अलग करने के लिए, Ray का इस्तेमाल करके, SweBench और TRL के साथ बुनियादी आरएल ट्रेनिंग जॉब चलाने का तरीका.
2. क्लस्टर बनाना और ज़रूरी शर्तें
आगे बढ़ने से पहले, आपके पास GKE क्लस्टर होना चाहिए. इसमें हाई-परफ़ॉर्मेंस जीपीयू नोड पूल और Ray ऑपरेटर इंस्टॉल होना चाहिए, ताकि ट्रेनिंग वर्कलोड को मैनेज किया जा सके.
ज़रूरी शर्तें
इस कोडलैब में, यह माना जाता है कि ये टूल इंस्टॉल और कॉन्फ़िगर किए गए हैं:
- 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"
कस्टम कंटेनर इमेज सेव करने के लिए, Artifact Registry में डेटाबेस बनाएं:
gcloud artifacts repositories create $REPO_NAME \
--repository-format=docker \
--location=$REGION \
--description="Repository for RL Sandbox images"
क्लस्टर कॉन्फ़िगरेशन
एआई वर्कलोड के लिए ऑप्टिमाइज़ किया गया GKE क्लस्टर प्रोविज़निंग करने के बारे में पूरी जानकारी पाने के लिए, आधिकारिक दस्तावेज़ पढ़ें: GKE एआई हाइपरकंप्यूट कस्टम क्लस्टर बनाना. इसमें GPUDirect RDMA नेटवर्क वायरिंग के बारे में भी बताया गया है.
अहम ज़रूरी शर्त: क्लस्टर या किसी खास एक्ज़ीक्यूशन नोड पूल को बनाते समय, पक्का करें कि आपने --enable-agent-sandbox और --sandbox type=gvisor फ़्लैग पास किए हों. इससे, सैंडबॉक्स वॉर्म पूल के लिए ज़रूरी कस्टम रिसॉर्स डेफ़िनिशन (सीआरडी) इंस्टॉल हो जाएंगे.
यह मानकर कि आपका क्लस्टर, जीपीयू, और Ray ऑपरेटर चल रहे हैं, यहां दी गई जानकारी में बताया गया है कि एक्ज़ीक्यूशन प्लेन को कैसे कॉन्फ़िगर किया जाए और आरएल लूप को कैसे चलाया जाए.
3. कस्टम इमेज बनाना
हाई-परफ़ॉर्मेंस आरएल चलाने के लिए, अपनी इमेज में डिपेंडेंसी बेक करना ज़रूरी है. हमें दो अलग-अलग इमेज की ज़रूरत है: एक, मॉडल चलाने वाले जीपीयू वर्कर के लिए और दूसरी, आकलन के ऐसे कोड को चलाने के लिए जिसे अलग किए गए सैंडबॉक्स में चलाया जाता है और जिस पर भरोसा नहीं किया जा सकता.
1. जीपीयू वर्कर इमेज बनाना
Ray जीपीयू वर्कर को, लैंग्वेज मॉडल चलाने और ट्रेनिंग लूप को ऑर्केस्ट्रेट करने के लिए लाइब्रेरी की ज़रूरत होती है. हम इस इमेज को, 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. सैंडबॉक्स इमेज बनाना
सैंडबॉक्स को, आकलन किए जा रहे टास्क के लिए खास डिपेंडेंसी की ज़रूरत होती है, ताकि रनटाइम इंस्टॉलेशन तुरंत हो सके. इस कोडलैब के लिए, हम SWE-bench में django/django डेटाबेस की किसी समस्या का इस्तेमाल करेंगे. हम डेटाबेस को पहले से क्लोन करेंगे और 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. Ray क्लस्टर कॉन्फ़िगर करना
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, हाई-स्पीड एचटीटीपी गेटवे के तौर पर काम करता है. यह Ray वर्कर से अनुरोध लेता है और उन्हें तुरंत उपलब्ध gVisor पॉड से जोड़ता है. इससे, Kubernetes API सर्वर पॉड लाइफ़साइकल को बाईपास किया जा सकता है, जो धीमा होता है.
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 एजेंट सैंडबॉक्स की मदद से, सैंडबॉक्स राऊटर का इस्तेमाल करके, अलग किए गए, पहले से वॉर्म किए गए कंटेनर को तुरंत असाइन किया जा सकता है. पॉड को तैयार रखने के लिए, हम 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, सैंडबॉक्स को पूरी तरह से अलग करता है. इससे, जीसीपी मेटाडेटा सर्वर पर इग्रेस नहीं हो पाता. इस वजह से, आईएएम टोकन की चोरी नहीं हो पाती.
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 के साथ बुनियादी आरएल जॉब
क्लस्टर और सैंडबॉक्स तैयार हो जाने के बाद, हम जीआरपीओ ट्रेनिंग लूप चला सकते हैं. हम जीआरपीओ एल्गोरिदम को ऑर्केस्ट्रेट करने के लिए, trl लाइब्रेरी का इस्तेमाल करेंगे. साथ ही, जनरेट किए गए कोड का आकलन करने के लिए, 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 हेड डैशबोर्ड पर पोर्ट-फ़ॉरवर्ड करें और अपनी लोकल मशीन से ट्रेनिंग जॉब सबमिट करें:
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 एजेंट सैंडबॉक्स का इस्तेमाल करके, GKE Standard पर हाई-परफ़ॉर्मेंस डिस्ट्रिब्यूटेड आरएल ट्रेनिंग लूप को सुरक्षित तरीके से कॉन्फ़िगर और एक्ज़ीक्यूट कर लिया है.