High-Performance Distributed RL on GKE Standard: The Complete Guide

1. מבוא

במעבדה הזו נסביר איך ליצור, להקצות ולהריץ לולאת אימון מבוזרת של למידת חיזוק (RL) בביצועים גבוהים ב-GKE Standard עם ארגזי חול לסוכנים ב-GKE‏ (gVisor), באמצעות אלגוריתם Group Relative Policy Optimization (GRPO) עם ספריית trl.

המטרה היא להדגים איך להעריך באופן מאובטח קוד לא מהימן שנוצר על ידי מודל שפה גדול (LLM) במהלך לולאת אימון של RL. אנחנו עושים את זה על ידי הפרדה בין מישור התזמור (Ray) לבין מישור הביצוע (ארגזי חול לסוכנים ב-GKE).

האתגר הטכני בהערכת קוד RL

כשמאמנים סוכני LLM באמצעות למידת חיזוק (לדוגמה, אימון מודל לכתיבת קוד על ידי הערכת הפלט שלו בבדיקות יחידה), לולאת האימון צריכה להריץ אלפי סקריפטים של Python שנוצרו על ידי LLM במקביל. הדבר יוצר אתגרים קריטיים:

  1. צוואר הבקבוק של Pod Churn: מסגרות הערכה מסורתיות התחלת הרצה קונטיינר Docker חדש לכל משימה. ביצוע הפעולה הזו באופן דינמי למאות השקות מקבילות במהלך לולאת אימון של RL גורם לעומס חמור במישור הבקרה של Kubernetes. זמן האחזור לא מאפשר אימון RL בתדירות גבוהה.
  2. סיכון האבטחה: הפעלת קוד שרירותי שנוצר על ידי LLM בתוך סביבות זמן ריצה רגילות של קונטיינרים משתפת את ליבת מערכת ההפעלה של המארח. פגיעות אחת מסוג escape עלולה לסכן את הצמתים שלכם.
  3. גניבת אסימוני IAM: קוד שנוצר על ידי LLM שפועל בתוך pod של Kubernetes יכול להריץ שאילתה על שרת המטא-נתונים של ספק הענן כדי לגנוב אסימונים של חשבון שירות IAM של הצומת.

הפתרון: תזמור וביצוע מופרדים

בארכיטקטורה הזו התזמור מופרד מהביצוע:

  • הכלי לניהול תזמור (Ray): אשכול Ray מבוזר מנהל את לולאת האימון של RL ומפיץ את יצירת הפריסה.
  • מישור ההפעלה (GKE Agent Sandbox): במקום ליצור באופן דינמי פודים של Kubernetes, עובדי Ray מבצעים קריאות HTTP פשוטות לנתב Sandbox ייעודי. הנתב מקצה לעובד באופן מיידי קונטיינר מבודד שחומם מראש ופועל במסגרת GKE Sandbox.
  • זמן אחזור של פחות משנייה: ארגזי חול מחוממים מראש ב-SandboxWarmPool מנוהל ומנוהלים באמצעות שער HTTP מהיר, ולכן יצירת סביבה מתבצעת תוך פחות מ-200 אלפיות השנייה, בלי לעבור דרך מישור הבקרה של Kubernetes.

מטרות ה-Lab

ב-codelab הזה תלמדו:

  • האתגרים האדריכליים והפתרונות להערכת קוד לא מהימן בלולאות RL.
  • איך ליצור תמונות מותאמות אישית של ארגז חול לפריסות יעילות.
  • איך מגדירים ומשתמשים בארגזי חול לסוכנים ב-GKE וב-SandboxWarmPools.
  • איך לבודד ארגזי חול בצורה מאובטחת כדי למנוע גניבה של טוקנים של IAM.
  • איך מריצים משימת אימון בסיסית של RL עם SweBench ו-TRL באמצעות Ray כדי להפריד בין תזמור לבין ביצוע.

2. יצירת אשכולות ותנאים מוקדמים

לפני שממשיכים, צריך לוודא שיש לכם אשכול GKE עם מאגר צמתים של GPU בעל ביצועים גבוהים, ושה-Ray Operator מותקן כדי לנהל את עומס העבודה של האימון.

דרישות מוקדמות

ב-codelab הזה אנחנו יוצאים מנקודת הנחה שהכלים הבאים מותקנים ומוגדרים:

משתני סביבה

קודם מגדירים את משתני הסביבה שישמשו לאורך כל ה-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 בהתאמה אישית.

דרישה מוקדמת חשובה: כשיוצרים את האשכול או מאגר ספציפי של צמתים להרצה, צריך להעביר את הדגלים --enable-agent-sandbox ו---sandbox type=gvisor כדי להתקין את ההגדרות הנדרשות של משאבים בהתאמה אישית (CRD) למאגרי החמים של ארגז החול.

בהנחה שהאשכול, המעבדים הגרפיים ו-Ray Operator פועלים, בהמשך מוסבר איך להגדיר את מישור הביצוע ולהריץ את לולאת ה-RL.

3. יצירת תמונות בהתאמה אישית

היבט חשוב בהפעלת RL עם ביצועים גבוהים הוא הטמעת תלות בתמונות. אנחנו צריכים שתי תמונות שונות: אחת בשביל עובדי ה-GPU שמריצים את המודל, ואחת בשביל ארגזי החול המבודדים שמריצים את קוד ההערכה הלא מהימן.

1. בניית תמונת GPU Worker

תהליך העבודה של Ray GPU צריך ספריות כדי להריץ את מודל השפה ולתזמן את לולאת האימון. אנחנו יוצרים את התמונה הזו על בסיס תמונת 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. יצירת תמונה של ראש המעבד

צומת הראש של Ray רק מתזמן את האשכול ולא מריץ את מודלי האימון הכבדים של ה-GPU. כדי למנוע צוואר בקבוק של משיכת תמונות גדולות (בדרך כלל 15GB ומעלה) בצמתי CPU רגילים, אנחנו יוצרים תמונה קלה שמתאימה רק ל-CPU עבור צומת הראש. התמונה הזו מכילה את 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. יצירת תמונת ארגז החול

ארגז החול צריך את התלות הספציפית למשימה שאנחנו בודקים, כדי שההתקנה בזמן הריצה תהיה מיידית. ב-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. שימו לב שהמשאבים הזמינים של האשכול (כמו זיכרון, מעבד או סוג 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

2. הגדרת SandboxRouter

ה-SandboxRouter פועל כשער HTTP מהיר, שמקבל בקשות מעובדי 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 כדי שה-pods יהיו מוכנים.

מריצים את הפקודה הבאה כדי ליצור את 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

בודקים שנוצר 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. הלוגיקה של הניתוב שמוצגת בהמשך מראה איך בוחרים מאגרי זיכרון שונים לריפוזיטורים שונים. זה שימושי כשמרחיבים את השימוש לכל מערך הנתונים של 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 בדפדפן.
  • Sandbox Claims: בסרטון הבא אפשר לראות איך GKE תובע ומשחרר באופן דינמי ארגזי חול ב-gVisor:
    watch -n 1 "kubectl get sandboxclaims,sandboxes,pods"
    

6. סיכום

מעולה! הגדרתם והפעלתם בהצלחה לולאת אימון מבוזרת של RL בביצועים גבוהים ובאופן מאובטח ב-GKE Standard באמצעות ארגזי חול לסוכנים ב-GKE.