یادگیری ماشین توزیع‌شده با کارایی بالا روی استاندارد GKE: راهنمای کامل

۱. مقدمه

این آزمایشگاه به تفصیل نحوه ساخت، آماده‌سازی و اجرای یک حلقه آموزشی یادگیری تقویتی توزیع‌شده (RL) با کارایی بالا بر روی GKE Standard با GKE Agent Sandboxes (gVisor) و با استفاده از الگوریتم بهینه‌سازی سیاست نسبی گروهی (GRPO) با کتابخانه trl شرح می‌دهد.

هدف این است که نشان دهیم چگونه می‌توان کد تولید شده توسط LLM غیرقابل اعتماد را در طول یک حلقه آموزش RL به طور ایمن ارزیابی کرد. ما با جدا کردن صفحه ارکستراسیون (Ray) از صفحه اجرا (GKE Agent Sandboxes) به این هدف دست می‌یابیم.

چالش فنی ارزیابی کد RL

هنگام آموزش عامل‌های LLM با استفاده از یادگیری تقویتی (مثلاً آموزش یک مدل برای نوشتن کد با ارزیابی خروجی آن در تست‌های واحد)، حلقه آموزشی باید هزاران اسکریپت پایتون تولید شده توسط LLM غیرقابل اعتماد را به صورت موازی اجرا کند. این امر چالش‌های مهمی را ایجاد می‌کند:

  1. گلوگاه Pod Churn: چارچوب‌های ارزیابی سنتی برای هر وظیفه، یک کانتینر داکر جدید راه‌اندازی می‌کنند. انجام این کار به صورت پویا برای صدها اجرای موازی در طول یک حلقه آموزش یادگیری تقویتی (RL) باعث بار شدید روی صفحه کنترل Kubernetes می‌شود. تأخیر، آموزش یادگیری تقویتی با فرکانس بالا را غیرممکن می‌سازد.
  2. خطر امنیتی: اجرای کد دلخواه تولید شده توسط LLM در داخل زمان‌های اجرای استاندارد کانتینر، هسته سیستم عامل میزبان را به اشتراک می‌گذارد. یک آسیب‌پذیری escape می‌تواند گره‌های شما را به خطر بیندازد.
  3. سرقت توکن IAM: کد تولید شده توسط LLM که درون یک پاد Kubernetes اجرا می‌شود، می‌تواند سرور ابرداده ارائه دهنده ابر را برای سرقت توکن‌های حساب سرویس IAM گره جستجو کند.

راه حل: هماهنگی و اجرای جداگانه

این معماری، هماهنگی (orcistination) را از اجرا (run) جدا می‌کند :

  • هماهنگ‌کننده (Ray): یک کلاستر توزیع‌شده Ray، حلقه آموزش RL را مدیریت کرده و تولید نسخه نهایی را توزیع می‌کند.
  • صفحه اجرا (GKE Agent Sandbox): به جای ایجاد پویای پادهای Kubernetes، کارگران Ray فراخوانی‌های ساده HTTP را به یک روتر Sandbox اختصاصی انجام می‌دهند. روتر فوراً یک کانتینر ایزوله و از پیش گرم شده را که تحت gVisor (GKE Sandbox) اجرا می‌شود، به کارگر اختصاص می‌دهد.
  • تأخیر زیر ثانیه: از آنجا که سندباکس‌ها در یک SandboxWarmPool مدیریت‌شده از قبل گرم می‌شوند و از طریق یک دروازه HTTP پرسرعت مدیریت می‌شوند، ایجاد محیط به زیر ۲۰۰ میلی‌ثانیه کاهش می‌یابد و به طور کامل صفحه کنترل Kubernetes را دور می‌زند.

اهداف آزمایشگاه

در این آزمایشگاه کد، شما یاد خواهید گرفت:

  • چالش‌ها و راه‌حل‌های معماری برای ارزیابی کد غیرقابل اعتماد در حلقه‌های یادگیری تقویتی
  • چگونه تصاویر سندباکس سفارشی برای انتشارهای کارآمد بسازیم.
  • نحوه پیکربندی و استفاده از GKE Agent Sandboxes و SandboxWarmPools.
  • چگونه می‌توان سندباکس‌ها را به طور ایمن ایزوله کرد تا از سرقت توکن‌های IAM جلوگیری شود.
  • چگونه یک کار آموزشی مقدماتی RL را با SweBench و TRL با استفاده از Ray اجرا کنیم تا هماهنگی (orchestration) را از اجرا (execution) جدا کنیم.

۲. ایجاد خوشه و پیش‌نیازها

قبل از ادامه، به یک کلاستر GKE با یک استخر گره GPU با کارایی بالا و Ray Operator نصب شده برای مدیریت حجم کار آموزشی نیاز دارید.

پیش‌نیازها

این آزمایشگاه کد فرض می‌کند که ابزارهای زیر نصب و پیکربندی شده‌اند:

متغیرهای محیطی

ابتدا، متغیرهای محیطی که در سراسر این آزمایشگاه کد استفاده خواهند شد را تنظیم کنید. دستورات زیر از مقادیر پیش‌فرض معقول استفاده می‌کنند، اما می‌توانید در صورت نیاز آنها را تغییر دهید تا با محیط خاص 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 بهینه‌شده برای بارهای کاری هوش مصنوعی (شامل سیم‌کشی شبکه‌ی GPUDirect RDMA)، مستندات رسمی را دنبال کنید: ایجاد یک کلاستر سفارشی GKE AI Hypercompute

پیش‌نیاز بسیار مهم: هنگام ایجاد کلاستر یا یک مجموعه گره اجرایی خاص، مطمئن شوید که پرچم‌های --enable-agent-sandbox و --sandbox type=gvisor را برای نصب تعاریف منابع سفارشی (CRD) مورد نیاز برای مجموعه‌های گرم Sandbox ارسال می‌کنید.

با فرض اینکه کلاستر، پردازنده‌های گرافیکی و اپراتور Ray شما در حال اجرا هستند، موارد زیر جزئیات نحوه پیکربندی صفحه اجرا و اجرای حلقه RL را شرح می‌دهد.

۳. ساخت تصاویر سفارشی

یکی از جنبه‌های حیاتی اجرای یادگیری تقویتی با عملکرد بالا، گنجاندن وابستگی‌ها در تصاویر شماست. ما به دو تصویر مجزا نیاز داریم: یکی برای کارگران GPU که مدل را اجرا می‌کنند و دیگری برای جعبه‌های شنی ایزوله که کد ارزیابی غیرقابل اعتماد را اجرا می‌کنند.

۱. ساخت تصویر کارگر GPU

کارگر 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

ایمیج را بسازید و به مخزن رجیستری مصنوعات خود منتقل کنید:

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 ).

۲. ساخت ایمیج هد پردازنده (CPU Head)

گره هد ری فقط خوشه را هماهنگ می‌کند و مدل‌های آموزشی سنگین GPU را اجرا نمی‌کند. برای جلوگیری از گلوگاه عظیم بارگذاری تصویر (معمولاً ۱۵ گیگابایت یا بیشتر) در گره‌های CPU استاندارد شما، ما یک تصویر سبک و فقط CPU-محور برای گره هد می‌سازیم. این تصویر شامل ری و کتابخانه‌های پایتون مورد نیاز است، اما کتابخانه‌های سنگین 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

۳. ساخت تصویر جعبه شنی

جعبه شنی به وابستگی‌های خاص برای وظیفه‌ای که ارزیابی می‌کنیم نیاز دارد تا نصب در زمان اجرا آنی باشد. برای این آزمایشگاه کد، ما از یک مسئله از مخزن django/django در SWE-bench استفاده خواهیم کرد. ما مخزن را از قبل کلون می‌کنیم و محیط‌های پایتون را از قبل می‌سازیم تا اسکریپت‌های مدل ما برای دانلود آنها در حلقه 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

۴. پیکربندی ارکستراسیون و اجرا

اکنون خوشه Ray را برای هماهنگ‌سازی و منابع Sandbox را برای اجرا مستقر می‌کنیم.

۱. پیکربندی خوشه ری

یک منبع سفارشی 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

۲. پیکربندی روتر سندباکس

SandboxRouter به عنوان یک دروازه HTTP پرسرعت عمل می‌کند، درخواست‌های Ray workerها را دریافت کرده و آنها را فوراً به غلاف‌های 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

۳. پیکربندی SandboxTemplate و WarmPool

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

۴. ایزوله‌سازی امنیتی

یک NetworkPolicy به شدت sandboxها را ایزوله می‌کند و از خروج به سرور GCP Metadata جلوگیری می‌کند و در نتیجه از سرقت توکن 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

۵. کار مقدماتی RL با SweBench و TRL

پس از آماده شدن کلاستر و جعبه‌های شنی، می‌توانیم یک حلقه آموزشی GRPO را اجرا کنیم. ما از کتابخانه trl برای هماهنگ‌سازی الگوریتم GRPO و از توابع ray remote برای ارزیابی کد تولید شده درون جعبه‌های شنی ایزوله استفاده خواهیم کرد.

برای سرعت بخشیدن به اجرا در این آزمایشگاه کد، ما کدها را به یک مسئله‌ی جنگو فیلتر می‌کنیم. منطق مسیریابی زیر نشان می‌دهد که چگونه می‌توانید warmpool های مختلف را برای مخازن مختلف انتخاب کنید، که هنگام گسترش به مجموعه داده‌های کامل 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

نظارت بر اجرا

می‌توانید پیشرفت دویدن خود را زیر نظر داشته باشید:

  • داشبورد ری: http://localhost:8265 را در مرورگر خود باز کنید.
  • ادعاهای مربوط به سندباکس: مشاهده کنید که GKE به صورت پویا سندباکس‌ها را تحت gVisor ادعا و منتشر می‌کند: 
    watch -n 1 "kubectl get sandboxclaims,sandboxes,pods"

۶. نتیجه‌گیری

تبریک! شما با موفقیت یک حلقه آموزشی RL توزیع‌شده با کارایی بالا را به صورت ایمن روی GKE Standard با استفاده از GKE Agent Sandboxes پیکربندی و اجرا کردید.