RL Terdistribusi Performa Tinggi di GKE Standard: Panduan Lengkap

1. Pengantar

Lab ini menjelaskan cara membangun, menyediakan, dan mengeksekusi loop pelatihan Reinforcement Learning (RL) terdistribusi berperforma tinggi di GKE Standard dengan GKE Agent Sandbox (gVisor), menggunakan algoritma Group Relative Policy Optimization (GRPO) dengan library trl.

Tujuannya adalah untuk mendemonstrasikan cara mengevaluasi kode yang dihasilkan LLM yang tidak tepercaya secara aman selama loop pelatihan RL. Kami mencapai hal ini dengan memisahkan bidang orkestrasi (Ray) dari bidang eksekusi (GKE Agent Sandboxes).

Tantangan Teknis Evaluasi Kode RL

Saat melatih agen LLM menggunakan Reinforcement Learning (misalnya, melatih model untuk menulis kode dengan mengevaluasi outputnya pada pengujian unit), loop pelatihan harus menjalankan ribuan skrip Python yang dihasilkan LLM dan tidak tepercaya secara paralel. Hal ini menimbulkan tantangan penting:

  1. Bottleneck Perubahan Pod: Framework evaluasi tradisional menjalankan container Docker baru per tugas. Melakukan hal ini secara dinamis untuk ratusan peluncuran paralel selama loop pelatihan RL menyebabkan beban berat pada bidang kontrol Kubernetes. Latensi membuat pelatihan RL frekuensi tinggi menjadi tidak mungkin.
  2. Risiko Keamanan: Menjalankan kode arbitrer yang dihasilkan LLM di dalam runtime container standar akan berbagi kernel OS host. Satu kerentanan escape dapat membahayakan node Anda.
  3. Pencurian Token IAM: Kode yang dihasilkan LLM yang berjalan di dalam pod Kubernetes dapat mengkueri server metadata penyedia cloud untuk mencuri token akun layanan IAM node.

Solusi: Orkestrasi & Eksekusi yang Tidak Terikat

Arsitektur ini memisahkan orkestrasi dari eksekusi:

  • Orchestrator (Ray): Cluster Ray terdistribusi mengelola loop pelatihan RL dan mendistribusikan pembuatan peluncuran.
  • Execution Plane (GKE Agent Sandbox): Alih-alih membuat pod Kubernetes secara dinamis, pekerja Ray melakukan panggilan HTTP sederhana ke Sandbox Router khusus. Router langsung menetapkan container terisolasi yang sudah dipanaskan sebelumnya yang berjalan di bawah gVisor (GKE Sandbox) ke pekerja.
  • Latensi di Bawah Detik: Karena sandbox dipanaskan terlebih dahulu di SandboxWarmPool terkelola dan dikelola melalui gateway HTTP berkecepatan tinggi, pembuatan lingkungan turun menjadi di bawah 200 md, sehingga sepenuhnya melewati bidang kontrol Kubernetes.

Tujuan Lab

Dalam codelab ini, Anda akan mempelajari:

  • Tantangan dan solusi arsitektur untuk mengevaluasi kode yang tidak tepercaya dalam loop RL.
  • Cara membuat image sandbox kustom untuk peluncuran yang efisien.
  • Cara mengonfigurasi dan menggunakan GKE Agent Sandbox dan SandboxWarmPool.
  • Cara mengisolasi sandbox secara aman untuk mencegah pencurian token IAM.
  • Cara menjalankan tugas pelatihan RL dasar dengan SweBench dan TRL menggunakan Ray untuk memisahkan orkestrasi dari eksekusi.

2. Pembuatan Cluster & Prasyarat

Sebelum melanjutkan, Anda memerlukan cluster GKE dengan kumpulan node GPU berperforma tinggi dan Ray Operator yang diinstal untuk mengelola workload pelatihan.

Prasyarat

Codelab ini mengasumsikan bahwa alat berikut telah diinstal dan dikonfigurasi:

Variabel Lingkungan

Pertama, tetapkan variabel lingkungan yang akan digunakan di seluruh codelab ini. Perintah di bawah menggunakan default yang wajar, tetapi Anda dapat mengubahnya sesuai kebutuhan agar cocok dengan lingkungan Google Cloud spesifik Anda:

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

Buat repositori Artifact Registry untuk menyimpan image container kustom:

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

Konfigurasi Cluster

Untuk mengetahui panduan lengkap tentang penyediaan cluster GKE yang dioptimalkan untuk workload AI (termasuk pengkabelan jaringan GPUDirect RDMA), ikuti dokumentasi resmi: Membuat Cluster Kustom Hypercompute AI GKE

Prasyarat Penting: Saat membuat cluster atau node pool eksekusi tertentu, pastikan Anda meneruskan flag --enable-agent-sandbox dan --sandbox type=gvisor untuk menginstal Definisi Resource Kustom (CRD) yang diperlukan untuk pool hangat Sandbox.

Dengan asumsi cluster, GPU, dan Ray Operator Anda berjalan, semua hal di bawah ini menjelaskan cara mengonfigurasi bidang eksekusi dan menjalankan loop RL.

3. Membangun Image Kustom

Aspek penting dalam menjalankan RL berperforma tinggi adalah memasukkan dependensi ke dalam image Anda. Kita memerlukan dua image yang berbeda: satu untuk pekerja GPU yang menjalankan model, dan satu untuk sandbox terisolasi yang menjalankan kode evaluasi yang tidak tepercaya.

1. Membangun Image Pekerja GPU

Worker GPU Ray memerlukan library untuk menjalankan model bahasa dan mengatur loop pelatihan. Kita membangun image ini di atas image vLLM resmi sehingga mendukung GPU terbaru dan telah menginstal PyTorch/CUDA.

Jalankan perintah berikut untuk membuat 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

Bangun dan kirim image ke repositori Artifact Registry Anda:

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

Catatan: Panduan ini menggunakan perintah docker lokal untuk membuat image. Jika lebih suka membangun image dari jarak jauh, Anda dapat menggunakan Cloud Build (misalnya, menggunakan gcloud builds submit).

2. Buat Gambar Kepala CPU

Node head Ray hanya mengorkestrasi cluster dan tidak menjalankan model pelatihan GPU berat. Untuk menghindari hambatan penarikan image yang besar (biasanya 15 GB+) pada node CPU standar, kami membuat image ringan khusus CPU untuk node head. Image ini berisi Ray dan library Python yang diperlukan, tetapi tidak menyertakan library GPU berat seperti CUDA dan vLLM.

Jalankan perintah berikut untuk membuat 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

Build dan kirim image:

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

Sandbox memerlukan dependensi khusus untuk tugas yang kita evaluasi sehingga penginstalan runtime dapat dilakukan secara instan. Untuk codelab ini, kita akan menggunakan masalah dari repositori django/django di SWE-bench. Kita akan melakukan pra-kloning repositori dan pra-build lingkungan python sehingga skrip model kita tidak membuang waktu untuk mendownloadnya dalam loop RL.

Jalankan perintah berikut untuk membuat 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

Build dan kirim image:

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. Mengonfigurasi Orkestrasi dan Eksekusi

Sekarang kita men-deploy cluster Ray untuk orkestrasi dan resource Sandbox untuk eksekusi.

1. Konfigurasi Cluster Ray

Deploy resource kustom RayCluster. Perhatikan bahwa resource yang tersedia di cluster Anda (seperti jenis memori, CPU, atau GPU) mungkin berbeda. Sesuaikan permintaan dan batas resources dengan tepat.

Jalankan perintah berikut untuk membuat raycluster.yaml. Tindakan ini menggunakan cat << EOF untuk otomatis mengganti variabel lingkungan Anda ke dalam manifes:

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

Terapkan:

kubectl apply -f raycluster.yaml

Pastikan cluster dibuat dan berjalan (proses ini mungkin memerlukan waktu beberapa menit):

kubectl get raycluster

Output yang diharapkan:

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

2. Konfigurasi SandboxRouter

SandboxRouter bertindak sebagai gateway HTTP berkecepatan tinggi, yang menangani permintaan dari pekerja Ray dan langsung menghubungkannya ke pod gVisor yang tersedia, sehingga melewati siklus proses pod server Kubernetes API yang lebih lambat.

Jalankan perintah berikut untuk membuat 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

Terapkan:

kubectl apply -f sandbox_router.yaml

Verifikasi bahwa deployment sedang berjalan:

kubectl get deployment sandbox-router-deployment

Output yang diharapkan:

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

3. Konfigurasi SandboxTemplate dan WarmPool

GKE Agent Sandbox memungkinkan penetapan instan container terisolasi yang telah di-warm-up menggunakan Sandbox Router. Kita menentukan SandboxTemplate dan SandboxWarmPool agar pod tetap siap.

Jalankan perintah berikut untuk membuat sandbox_warmpool.yaml dengan variabel lingkungan Anda:

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

Terapkan:

kubectl apply -f sandbox_warmpool.yaml

Verifikasi bahwa SandboxWarmPool diinisialisasi:

kubectl get sandboxwarmpool

Output yang diharapkan:

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

4. Isolasi Keamanan

NetworkPolicy mengisolasi sandbox secara ketat, mencegah keluar ke Server Metadata GCP, sehingga mencegah pencurian token IAM.

Jalankan perintah berikut untuk membuat 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

Terapkan kebijakan:

kubectl apply -f network_policy.yaml

Pastikan NetworkPolicy telah dibuat:

kubectl get networkpolicy

Output yang diharapkan:

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

5. Pekerjaan RL Dasar dengan SweBench dan TRL

Setelah cluster dan sandbox disiapkan, kita dapat menjalankan loop pelatihan GRPO. Kita akan menggunakan library trl untuk mengatur algoritma GRPO, dan fungsi jarak jauh ray untuk mengevaluasi kode yang dihasilkan di dalam sandbox terisolasi.

Untuk membuat eksekusi cepat untuk codelab ini, kita akan memfilter hingga satu masalah Django. Logika perutean di bawah menunjukkan cara memilih warmpool yang berbeda untuk repositori yang berbeda, yang berguna saat memperluas ke set data SWE-bench lengkap.

Skrip Pelatihan

Jalankan perintah berikut untuk membuat 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

Mengirimkan Tugas ke Cluster

Pertama, teruskan port ke dasbor Ray Head dan kirimkan tugas pelatihan dari komputer lokal Anda:

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

Memantau Operasi

Anda dapat memantau progres eksekusi:

  • Dasbor Ray: Buka http://localhost:8265 di browser Anda.
  • Klaim Sandbox: Amati GKE yang mengklaim dan melepaskan sandbox secara dinamis di bawah gVisor:
    watch -n 1 "kubectl get sandboxclaims,sandboxes,pods"
    

6. Kesimpulan

Selamat! Anda telah berhasil mengonfigurasi dan menjalankan loop pelatihan RL terdistribusi berperforma tinggi secara aman di GKE Standard menggunakan GKE Agent Sandbox.