Wysokowydajne rozproszone uczenie ze wzmocnieniem w GKE Standard: pełny przewodnik

1. Wprowadzenie

W tym laboratorium dowiesz się, jak tworzyć, udostępniać i wykonywać wydajną rozproszoną pętlę trenowania uczenia ze wzmocnieniem (RL) w GKE Standard z użyciem piaskownic agentów GKE (gVisor), korzystając z algorytmu Group Relative Policy Optimization (GRPO) z biblioteką trl.

Celem jest pokazanie, jak bezpiecznie oceniać niezaufany kod wygenerowany przez LLM w pętli trenowania RL. Osiągamy to przez oddzielenie platformy orkiestracji (Ray) od platformy wykonywania (piaskownice agentów GKE).

Problem techniczny związany z oceną kodu RL

Podczas trenowania agentów LLM za pomocą uczenia ze wzmocnieniem (np. trenowania modelu do pisania kodu przez ocenę jego danych wyjściowych na podstawie testów jednostkowych) pętla trenowania musi równolegle wykonywać tysiące niezaufanych skryptów w języku Python wygenerowanych przez LLM. Stwarza to poważne problemy:

  1. Wąskie gardło w przypadku podów: tradycyjne platformy oceny uruchamiają nowy kontener Docker dla każdego zadania. Dynamiczne wykonywanie tej czynności w przypadku setek równoległych wdrożeń podczas pętli trenowania RL powoduje duże obciążenie platformy sterującej Kubernetes. Opóźnienie uniemożliwia trenowanie RL z wysoką częstotliwością.
  2. Ryzyko związane z bezpieczeństwem: uruchamianie dowolnego kodu wygenerowanego przez LLM w standardowych środowiskach wykonawczych kontenerów powoduje współdzielenie jądra systemu operacyjnego hosta. Pojedyncza luka w zabezpieczeniach może narazić węzły na ryzyko.
  3. Kradzież tokenów uprawnień: kod wygenerowany przez LLM działający w podzie Kubernetes może wysyłać zapytania do serwera metadanych dostawcy usług w chmurze, aby wykraść tokeny kont usługi uprawnień węzła.

Rozwiązanie: rozdzielenie orkiestracji i wykonywania

Ta architektura oddziela orkiestrację od wykonywania:

  • Orchestrator (Ray): rozproszony klaster Ray zarządza pętlą trenowania RL i rozprowadza generowanie wdrożeń.
  • Płaszczyzna wykonawcza (piaskownica agentów GKE): zamiast dynamicznie tworzyć pody Kubernetes, instancje robocze Ray wykonują proste wywołania HTTP do dedykowanego routera piaskownicy. Router natychmiast przypisuje instancji roboczej izolowany, wstępnie rozgrzany kontener działający w gVisor (GKE Sandbox).
  • Opóźnienie poniżej sekundy: ponieważ piaskownice są wstępnie rozgrzewane w zarządzanym SandboxWarmPool i zarządzane za pomocą szybkiej bramy HTTP, tworzenie środowiska trwa mniej niż 200 ms, co całkowicie omija platformę sterującą Kubernetes.

Cele modułu

Z tego przewodnika dowiesz się:

  • Wyzwania architektoniczne i rozwiązania związane z ocenianiem niezaufanego kodu w pętlach RL.
  • Jak tworzyć niestandardowe obrazy piaskownicy na potrzeby sprawnego wdrażania.
  • Jak skonfigurować i używać piaskownic agentów GKE i puli SandboxWarmPools.
  • Jak bezpiecznie odizolować piaskownice, aby zapobiec kradzieży tokenów IAM.
  • Jak uruchomić podstawowe zadanie trenowania RL za pomocą SweBench i TRL z użyciem Ray, aby oddzielić orkiestrację od wykonania.

2. Tworzenie klastra i wymagania wstępne

Zanim przejdziesz dalej, musisz mieć klaster GKE z pulą węzłów GPU o wysokiej wydajności i zainstalowanym operatorem Ray do zarządzania zbiorem zadań trenowania.

Wymagania wstępne

W tym laboratorium zakłada się, że te narzędzia są zainstalowane i skonfigurowane:

Zmienne środowiskowe

Najpierw ustaw zmienne środowiskowe, których będziesz używać podczas naszych ćwiczeń z programowania. Poniższe polecenia używają rozsądnych wartości domyślnych, ale możesz je w razie potrzeby zmienić, aby dopasować je do konkretnego środowiska 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"

Utwórz repozytorium Artifact Registry, w którym będą przechowywane niestandardowe obrazy kontenerów:

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

Konfiguracja klastra

Pełne instrukcje dotyczące udostępniania klastra GKE zoptymalizowanego pod kątem zbiorów zadań AI (w tym okablowania sieci GPUDirect RDMA) znajdziesz w oficjalnej dokumentacji: Tworzenie niestandardowego klastra GKE AI Hypercompute

Ważny warunek wstępny: podczas tworzenia klastra lub konkretnej puli węzłów wykonawczych przekaż flagi --enable-agent-sandbox i --sandbox type=gvisor, aby zainstalować wymagane definicje zasobów niestandardowych (CRD) dla pul wstępnego rozruchu piaskownicy.

Zakładając, że klaster, procesory GPU i operator Ray działają, poniżej znajdziesz szczegółowe informacje o tym, jak skonfigurować płaszczyznę wykonawczą i uruchomić pętlę RL.

3. Tworzenie obrazów niestandardowych

Kluczowym aspektem działania RL o wysokiej wydajności jest uwzględnianie zależności w obrazach. Potrzebujemy 2 różnych obrazów: jednego dla procesów roboczych GPU, które uruchamiają model, i jednego dla odizolowanych piaskownic, które uruchamiają niezaufany kod oceny.

1. Tworzenie obrazu instancji roboczej GPU

Proces roboczy GPU Ray potrzebuje bibliotek do uruchamiania modelu językowego i koordynowania pętli trenowania. Ten obraz jest oparty na oficjalnym obrazie vLLM, więc obsługuje najnowsze procesory graficzne i ma wstępnie zainstalowane PyTorch/CUDA.

Aby utworzyć Dockerfile.gpu_worker, uruchom to polecenie:

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

Skompiluj obraz i przenieś go do repozytorium 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

Uwaga: w tym przewodniku do tworzenia obrazów używane są lokalne polecenia docker. Jeśli wolisz tworzyć obrazy zdalnie, możesz zamiast tego użyć Cloud Build (np. za pomocą gcloud builds submit).

2. Tworzenie obrazu procesora

Węzeł główny Ray tylko zarządza klastrem i nie uruchamia modeli trenowania GPU o dużym obciążeniu. Aby uniknąć wąskiego gardła związanego z pobieraniem dużych obrazów (zwykle ponad 15 GB) na standardowych węzłach procesora, tworzymy uproszczony obraz tylko na procesor dla węzła głównego. Ten obraz zawiera Ray i wymagane biblioteki Pythona, ale nie zawiera bibliotek GPU, takich jak CUDA i vLLM.

Aby utworzyć Dockerfile.head, uruchom to polecenie:

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

Utwórz i wypchnij obraz:

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. Tworzenie obrazu piaskownicy

Środowisko testowe potrzebuje konkretnych zależności dla zadania, które oceniamy, aby instalacja w czasie działania była natychmiastowa. W tym laboratorium użyjemy problemu z repozytorium django/django w SWE-bench. Wcześniej sklonujemy repozytorium i utworzymy środowiska Pythona, aby skrypty modelu nie traciły czasu na ich pobieranie w pętli RL.

Aby utworzyć Dockerfile.sandbox, uruchom to polecenie:

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

Utwórz i wypchnij obraz:

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. Konfigurowanie orkiestracji i wykonywania

Teraz wdrażamy klaster Ray do orkiestracji i zasoby piaskownicy do wykonywania.

1. Konfiguracja klastra Ray

Wdróż niestandardowy zasób RayCluster. Pamiętaj, że dostępne zasoby klastra (np. pamięć, CPU lub typ GPU) mogą się różnić. Dostosuj odpowiednio żądania i limity resources.

Aby utworzyć raycluster.yaml, uruchom to polecenie: Używa to funkcji cat << EOF do automatycznego zastępowania zmiennych środowiskowych w pliku manifestu:

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

Zastosuj go:

kubectl apply -f raycluster.yaml

Sprawdź, czy klaster został utworzony i działa (może to potrwać kilka minut):

kubectl get raycluster

Oczekiwane dane wyjściowe:

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

2. Konfiguracja SandboxRouter

SandboxRouter działa jako szybka brama HTTP, która obsługuje żądania od instancji roboczych Ray i natychmiast przekazuje je do dostępnych podów gVisor, omijając wolniejszy cykl życia podów serwera API Kubernetes.

Aby utworzyć sandbox_router.yaml, uruchom to polecenie:

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

Aby ją zastosować:

kubectl apply -f sandbox_router.yaml

Sprawdź, czy wdrożenie jest uruchomione:

kubectl get deployment sandbox-router-deployment

Oczekiwane dane wyjściowe:

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

3. Konfiguracja szablonu piaskownicy i puli wstępnej

Piaskownica agentów GKE umożliwia natychmiastowe przypisywanie izolowanych, wstępnie rozgrzanych kontenerów za pomocą routera piaskownicy. Definiujemy SandboxTemplate i SandboxWarmPool, aby utrzymać gotowość zasobników.

Aby utworzyć sandbox_warmpool.yaml ze zmiennymi środowiskowymi, uruchom to polecenie:

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

Aby ją zastosować:

kubectl apply -f sandbox_warmpool.yaml

Sprawdź, czy pula SandboxWarmPool została zainicjowana:

kubectl get sandboxwarmpool

Oczekiwane dane wyjściowe:

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

4. Izolacja zabezpieczeń

Zasady sieciowe ściśle izolują piaskownice, uniemożliwiając ruch wychodzący do serwera metadanych GCP, a tym samym zapobiegając kradzieży tokenów IAM.

Aby utworzyć network_policy.yaml, uruchom to polecenie:

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

Zastosuj zasadę:

kubectl apply -f network_policy.yaml

Sprawdź, czy zasada NetworkPolicy została utworzona:

kubectl get networkpolicy

Oczekiwane dane wyjściowe:

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

5. Podstawowe zadanie RL z SweBench i TRL

Gdy klaster i piaskownice będą gotowe, możemy uruchomić pętlę trenowania GRPO. Do koordynowania algorytmu GRPO będziemy używać biblioteki trl, a do oceny wygenerowanego kodu w izolowanych piaskownicach – funkcji zdalnych Ray.

Aby przyspieszyć wykonanie tego laboratorium, ograniczymy się do jednego problemu z Django. Poniższa logika routingu pokazuje, jak wybrać różne pule wstępne dla różnych repozytoriów. Jest to przydatne podczas rozszerzania na pełny zbiór danych SWE-bench.

Skrypt trenowania

Aby utworzyć train_trl.py, uruchom to polecenie:

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

Przesyłanie zadania do klastra

Najpierw przekieruj port do panelu Ray Head i prześlij zadanie trenowania z komputera lokalnego:

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

Monitorowanie uruchomienia

Możesz monitorować postępy biegu:

  • Panel Ray: otwórz http://localhost:8265 w przeglądarce.
  • Sandbox Claims: zobacz, jak GKE dynamicznie przejmuje i zwalnia piaskownice w gVisor: 
    watch -n 1 "kubectl get sandboxclaims,sandboxes,pods"

6. Podsumowanie

Gratulacje! Udało Ci się skonfigurować i bezpiecznie wykonać w GKE Standard pętlę trenowania rozproszonego RL o wysokiej wydajności przy użyciu piaskownic agentów GKE.