Google verwendet KI-Technologie, um Inhalte in Ihre bevorzugte Sprache zu übersetzen. KI-Übersetzungen können Fehler enthalten.

Hochleistungsfähiges verteiltes RL in GKE Standard: Der vollständige Leitfaden

1. Einführung

In diesem Lab wird beschrieben, wie Sie einen leistungsstarken verteilten Trainingszyklus für Reinforcement Learning (RL) in GKE Standard mit GKE Agent Sandboxes (gVisor) erstellen, bereitstellen und ausführen. Dabei wird der Algorithmus Group Relative Policy Optimization (GRPO) mit der trl-Bibliothek verwendet.

Ziel ist es, zu zeigen, wie Sie nicht vertrauenswürdigen, von LLMs generierten Code während eines RL-Trainingszyklus sicher auswerten können. Dies erreichen wir, indem wir die Orchestrierungsebene (Ray) von der Ausführungsebene (GKE Agent Sandboxes) entkoppeln.

Die technische Herausforderung der RL-Codebewertung

Beim Trainieren von LLM-Agents mit Reinforcement Learning (z.B. beim Trainieren eines Modells zum Schreiben von Code durch Bewerten der Ausgabe in Unittests) müssen im Trainingszyklus Tausende von nicht vertrauenswürdigen, von LLM generierten Python-Skripts parallel ausgeführt werden. Das führt zu kritischen Herausforderungen:

Engpass durch häufige Pod-Neustarts:Bei herkömmlichen Bewertungsframeworks wird für jede Aufgabe ein neuer Docker-Container gestartet. Wenn dies dynamisch für Hunderte von parallelen Rollouts während einer RL-Trainingsschleife geschieht, führt dies zu einer erheblichen Belastung der Kubernetes-Steuerungsebene. Die Latenz macht ein hochfrequentes RL-Training unmöglich.
Das Sicherheitsrisiko:Wenn beliebiger, von LLMs generierter Code in Standard-Containerlaufzeiten ausgeführt wird, wird der Host-Betriebssystemkernel gemeinsam genutzt. Eine einzelne Escape-Schwachstelle kann Ihre Knoten gefährden.
IAM-Tokendiebstahl:LLM-generierter Code, der in einem Kubernetes-Pod ausgeführt wird, kann den Metadatenserver des Cloud-Anbieters abfragen, um Knoten-IAM-Dienstkonto-Tokens zu stehlen.

Die Lösung: Entkoppelte Orchestrierung und Ausführung

Bei dieser Architektur wird die Orchestrierung von der Ausführung entkoppelt:

Orchestrator (Ray): Ein verteilter Ray-Cluster verwaltet den RL-Trainingszyklus und verteilt die Rollout-Generierung.
Ausführungsebene (GKE Agent Sandbox): Anstatt Kubernetes-Pods dynamisch zu erstellen, führen Ray-Worker einfache HTTP-Aufrufe an einen dedizierten Sandbox-Router aus. Der Router weist dem Worker sofort einen isolierten, vorab aufgewärmten Container zu, der unter gVisor (GKE Sandbox) ausgeführt wird.
Latenz unter einer Sekunde:Da Sandboxes in einer verwalteten SandboxWarmPool vorab aufgewärmt und über ein schnelles HTTP-Gateway verwaltet werden, dauert das Erstellen von Umgebungen weniger als 200 ms. Die Kubernetes-Steuerungsebene wird dabei vollständig umgangen.

Lernziele des Labs

In diesem Codelab erfahren Sie Folgendes:

Die architektonischen Herausforderungen und Lösungen für die Auswertung nicht vertrauenswürdigen Codes in RL-Schleifen.
Benutzerdefinierte Sandbox-Images für effiziente Rollouts erstellen
Konfigurieren und Verwenden von GKE Agent Sandboxes und SandboxWarmPools.
So isolieren Sie Sandboxes sicher, um den Diebstahl von IAM-Tokens zu verhindern.
So führen Sie einen einfachen RL-Trainingsjob mit SweBench und TRL mit Ray aus, um die Orchestrierung von der Ausführung zu entkoppeln.

2. Clustererstellung und Voraussetzungen

Bevor Sie fortfahren, benötigen Sie einen GKE-Cluster mit einem leistungsstarken GPU-Knotenpool und dem Ray-Operator, der die Trainingsarbeitslast verwaltet.

Vorbereitung

In diesem Codelab wird davon ausgegangen, dass die folgenden Tools installiert und konfiguriert sind:

Google Cloud SDK (gcloud)
Docker (erforderlich, um benutzerdefinierte Images lokal zu erstellen)
kubectl

Umgebungsvariablen

Legen Sie zuerst die Umgebungsvariablen fest, die in diesem Codelab verwendet werden. In den folgenden Befehlen werden sinnvolle Standardwerte verwendet, die Sie aber nach Bedarf an Ihre spezifische Google Cloud-Umgebung anpassen können:

export PROJECT_ID=$(gcloud config get-value project)
export REGION="us-west3"
export ZONE="us-west3-a"
export REPO_NAME="rl-sandbox-repo"

Erstellen Sie ein Artifact Registry-Repository für die benutzerdefinierten Container-Images:

gcloud artifacts repositories create $REPO_NAME \
    --repository-format=docker \
    --location=$REGION \
    --description="Repository for RL Sandbox images"

Clusterkonfiguration

Eine vollständige Anleitung zum Bereitstellen eines für KI-Arbeitslasten optimierten GKE-Clusters (einschließlich GPUDirect RDMA-Netzwerkverkabelung) finden Sie in der offiziellen Dokumentation: Benutzerdefinierten GKE AI Hypercompute-Cluster erstellen

Wichtige Voraussetzung:Wenn Sie Ihren Cluster oder einen bestimmten Ausführungsknotenpool erstellen, müssen Sie die Flags --enable-agent-sandbox und --sandbox type=gvisor übergeben, um die erforderlichen benutzerdefinierten Ressourcendefinitionen (Custom Resource Definitions, CRDs) für die Sandbox-Warmpools zu installieren.

Wenn Ihr Cluster, Ihre GPUs und Ihr Ray-Operator ausgeführt werden, wird unten beschrieben, wie Sie die Ausführungsebene konfigurieren und den RL-Loop ausführen.

3. Benutzerdefinierte Images erstellen

Ein wichtiger Aspekt beim Ausführen von leistungsstarkem RL ist das Einbetten von Abhängigkeiten in Ihre Images. Wir benötigen zwei verschiedene Images: eines für die GPU-Worker, auf denen das Modell ausgeführt wird, und eines für die isolierten Sandboxes, in denen der nicht vertrauenswürdige Auswertungscode ausgeführt wird.

1. GPU-Worker-Image erstellen

Der Ray-GPU-Worker benötigt Bibliotheken, um das Language Model auszuführen und den Trainingszyklus zu orchestrieren. Wir erstellen dieses Image auf dem offiziellen vLLM-Image, sodass es die neuesten GPUs unterstützt und PyTorch/CUDA vorinstalliert ist.

Führen Sie den folgenden Befehl aus, um Dockerfile.gpu_worker zu erstellen:

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

Erstellen Sie das Image und übertragen Sie es per Push in Ihr Artifact Registry-Repository:

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

Hinweis:In dieser Anleitung werden lokale docker-Befehle zum Erstellen von Images verwendet. Wenn Sie die Images lieber remote erstellen möchten, können Sie stattdessen Cloud Build verwenden, z.B. mit gcloud builds submit.

2. CPU-Head-Image erstellen

Der Ray-Head-Knoten orchestriert nur den Cluster und führt nicht die rechenintensiven GPU-Trainingsmodelle aus. Um einen massiven Engpass beim Abrufen von Images (in der Regel über 15 GB) auf Ihren Standard-CPU-Knoten zu vermeiden, erstellen wir ein schlankes, reines CPU-Image für den Head-Knoten. Dieses Image enthält Ray und die erforderlichen Python-Bibliotheken, aber keine GPU-Bibliotheken wie CUDA und vLLM.

Führen Sie den folgenden Befehl aus, um Dockerfile.head zu erstellen:

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

Erstellen Sie das Image und übertragen Sie es per Push:

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. Sandbox-Image erstellen

Für die Sandbox sind die spezifischen Abhängigkeiten für die Aufgabe erforderlich, die wir bewerten, damit die Installation zur Laufzeit sofort erfolgt. In diesem Codelab verwenden wir ein Problem aus dem django/django-Repository in SWE-bench. Wir klonen das Repository vorab und erstellen die Python-Umgebungen vorab, damit unsere Modellskripts im RL-Loop keine Zeit mit dem Herunterladen verschwenden.

Führen Sie den folgenden Befehl aus, um Dockerfile.sandbox zu erstellen:

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

Erstellen Sie das Image und übertragen Sie es per Push:

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. Orchestrierung und Ausführung konfigurieren

Jetzt stellen wir den Ray-Cluster für die Orchestrierung und die Sandbox-Ressourcen für die Ausführung bereit.

1. Ray-Clusterkonfiguration

Stellen Sie eine benutzerdefinierte RayCluster-Ressource bereit. Die verfügbaren Ressourcen Ihres Clusters (z. B. Arbeitsspeicher, CPU oder GPU-Typ) können sich unterscheiden. Passen Sie die resources-Anfragen und ‑Limits entsprechend an.

Führen Sie den folgenden Befehl aus, um raycluster.yaml zu erstellen. Dabei wird cat << EOF verwendet, um Ihre Umgebungsvariablen automatisch in das Manifest einzufügen:

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

So wenden Sie es an:

kubectl apply -f raycluster.yaml

Prüfen Sie, ob der Cluster erstellt wurde und ausgeführt wird. Dies kann einige Minuten dauern:

kubectl get raycluster

Erwartete Ausgabe:

NAME       DESIRED WORKERS   AVAILABLE WORKERS   CPUS   MEMORY   GPUS   STATUS   AGE
rl-cluster   1                 1                                            ready    2m

2. SandboxRouter-Konfiguration

Der SandboxRouter fungiert als schnelles HTTP-Gateway, das Anfragen von Ray-Workern entgegennimmt und sie sofort an verfügbare gVisor-Pods weiterleitet. So wird der langsamere Kubernetes API-Server-Pod-Lebenszyklus umgangen.

Führen Sie den folgenden Befehl aus, um sandbox_router.yaml zu erstellen:

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

So wenden Sie es an:

kubectl apply -f sandbox_router.yaml

Prüfen Sie, ob das Deployment ausgeführt wird:

kubectl get deployment sandbox-router-deployment

Erwartete Ausgabe:

NAME                        READY   UP-TO-DATE   AVAILABLE   AGE
sandbox-router-deployment   2/2     2            2           1m

3. SandboxTemplate- und WarmPool-Konfiguration

Mit GKE Agent Sandbox können isolierte, vorab aufgewärmte Container mithilfe des Sandbox-Routers sofort zugewiesen werden. Wir definieren einen SandboxTemplate und einen SandboxWarmPool, damit die Pods bereit bleiben.

Führen Sie den folgenden Befehl aus, um sandbox_warmpool.yaml mit Ihren Umgebungsvariablen zu erstellen:

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

So wenden Sie es an:

kubectl apply -f sandbox_warmpool.yaml

Prüfen Sie, ob der SandboxWarmPool initialisiert wurde:

kubectl get sandboxwarmpool

Erwartete Ausgabe:

NAME                        READY   AGE
swe-bench-django-warmpool   10      1m

4. Sicherheitsisolation

Eine NetworkPolicy isoliert Sandboxes strikt und verhindert den Egress zum GCP-Metadatenserver, wodurch der Diebstahl von IAM-Tokens verhindert wird.

Führen Sie den folgenden Befehl aus, um network_policy.yaml zu erstellen:

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

Wenden Sie die Richtlinie an:

kubectl apply -f network_policy.yaml

Prüfen Sie, ob die NetworkPolicy erstellt wurde:

kubectl get networkpolicy

Erwartete Ausgabe:

NAME                 POD-SELECTOR     AGE
block-metadata-egress             sandbox.gke.io/runtime=gvisor     1m

5. Einfacher RL-Job mit SweBench und TRL

Sobald der Cluster und die Sandboxes vorbereitet sind, können wir einen GRPO-Trainingslauf ausführen. Wir verwenden die trl-Bibliothek, um den GRPO-Algorithmus zu orchestrieren, und Ray-Remote-Funktionen, um den generierten Code in den isolierten Sandboxes auszuwerten.

Damit die Ausführung in diesem Codelab schnell geht, filtern wir nach einem einzelnen Django-Problem. Die folgende Routing-Logik zeigt, wie Sie verschiedene Warmpools für verschiedene Repositories auswählen. Das ist nützlich, wenn Sie den gesamten SWE-bench-Datensatz verwenden.

Das Trainingsskript

Führen Sie den folgenden Befehl aus, um train_trl.py zu erstellen:

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

Job an den Cluster senden

Leiten Sie zuerst den Port zum Ray Head-Dashboard weiter und senden Sie den Trainingsjob von Ihrem lokalen Computer aus:

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

Lauf überwachen

Sie können den Fortschritt Ihres Laufs so verfolgen:

Ray-Dashboard:Öffnen Sie http://localhost:8265 in Ihrem Browser.
Sandbox-Ansprüche:Beobachten Sie, wie GKE Sandboxes unter gVisor dynamisch beansprucht und freigibt:
```
watch -n 1 "kubectl get sandboxclaims,sandboxes,pods"
```

6. Fazit

Glückwunsch! Sie haben einen leistungsstarken verteilten RL-Trainingszyklus mit GKE Agent Sandboxes in GKE Standard sicher konfiguriert und ausgeführt.