Guide complet sur le RL distribué à hautes performances sur GKE Standard

1. Introduction

Cet atelier explique en détail comment créer, provisionner et exécuter une boucle d'entraînement par apprentissage par renforcement (RL) distribuée et performante sur GKE Standard avec GKE Agent Sandboxes (gVisor), à l'aide de l'algorithme Group Relative Policy Optimization (GRPO) avec la bibliothèque trl.

L'objectif est de montrer comment évaluer de manière sécurisée le code non fiable généré par LLM lors d'une boucle d'entraînement RL. Pour ce faire, nous découplons le plan d'orchestration (Ray) du plan d'exécution (GKE Agent Sandboxes).

Le défi technique de l'évaluation du code RL

Lorsque vous entraînez des agents LLM à l'aide de l'apprentissage par renforcement (par exemple, en entraînant un modèle à écrire du code en évaluant sa sortie sur des tests unitaires), la boucle d'entraînement doit exécuter des milliers de scripts Python non fiables générés par LLM en parallèle. Cela pose des problèmes critiques :

  1. Goulot d'étranglement du churn de pods : les frameworks d'évaluation traditionnels lancent un nouveau conteneur Docker pour chaque tâche. Effectuer cette opération de manière dynamique pour des centaines de déploiements parallèles au cours d'une boucle d'entraînement par renforcement entraîne une charge importante sur le plan de contrôle Kubernetes. La latence rend impossible l'entraînement RL à haute fréquence.
  2. Risque de sécurité : l'exécution de code arbitraire généré par LLM dans des environnements d'exécution de conteneurs standards partage le noyau de l'OS hôte. Une seule faille d'échappement peut compromettre vos nœuds.
  3. Vol de jetons IAM : le code généré par LLM s'exécutant dans un pod Kubernetes peut interroger le serveur de métadonnées du fournisseur de cloud pour voler les jetons de compte de service IAM du nœud.

La solution : orchestration et exécution dissociées

Cette architecture dissocie l'orchestration de l'exécution :

  • Orchestrateur (Ray) : un cluster Ray distribué gère la boucle d'entraînement RL et distribue la génération de déploiements.
  • Plan d'exécution (GKE Agent Sandbox) : au lieu de créer des pods Kubernetes de manière dynamique, les nœuds de calcul Ray effectuent de simples appels HTTP vers un routeur Sandbox dédié. Le routeur attribue instantanément au worker un conteneur isolé et préchauffé exécuté sous gVisor (GKE Sandbox).
  • Latence inférieure à la seconde : les bacs à sable étant préchauffés dans un SandboxWarmPool géré et gérés via une passerelle HTTP à haut débit, la création d'environnements prend moins de 200 ms, en contournant complètement le plan de contrôle Kubernetes.

Objectifs de l'atelier

Cet atelier de programmation traite des points suivants :

  • Les défis et les solutions architecturales pour évaluer le code non approuvé dans les boucles RL.
  • Découvrez comment créer des images de bac à sable personnalisées pour des déploiements efficaces.
  • Comment configurer et utiliser les GKE Agent Sandboxes et les SandboxWarmPools.
  • Comment isoler les bacs à sable de manière sécurisée pour éviter le vol de jetons IAM.
  • Découvrez comment exécuter un job d'entraînement RL de base avec SweBench et TRL à l'aide de Ray pour dissocier l'orchestration de l'exécution.

2. Création et prérequis des clusters

Avant de continuer, vous avez besoin d'un cluster GKE avec un pool de nœuds GPU hautes performances et de l'opérateur Ray installé pour gérer la charge de travail d'entraînement.

Prérequis

Pour suivre cet atelier de programmation, vous devez avoir installé et configuré les outils suivants :

Variables d'environnement

Commencez par définir les variables d'environnement qui seront utilisées tout au long de cet atelier de programmation. Les commandes ci-dessous utilisent des valeurs par défaut raisonnables, mais vous pouvez les modifier si nécessaire pour les adapter à votre environnement Google Cloud spécifique :

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

Créez un dépôt Artifact Registry pour stocker les images de conteneurs personnalisés :

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

Configuration du cluster

Pour obtenir un guide complet sur le provisionnement d'un cluster GKE optimisé pour les charges de travail d'IA (y compris le câblage réseau GPUDirect RDMA), consultez la documentation officielle : Créer un cluster GKE personnalisé AI Hypercompute.

Prérequis essentiel : lorsque vous créez votre cluster ou un pool de nœuds d'exécution spécifique, veillez à transmettre les indicateurs --enable-agent-sandbox et --sandbox type=gvisor pour installer les définitions de ressources personnalisées (CRD) requises pour les pools de préchauffage du bac à sable.

En supposant que votre cluster, vos GPU et votre opérateur Ray soient en cours d'exécution, tout ce qui suit explique comment configurer le plan d'exécution et exécuter la boucle RL.

3. Créer des images personnalisées

Un aspect essentiel de l'exécution de l'apprentissage par renforcement hautes performances consiste à intégrer les dépendances dans vos images. Nous avons besoin de deux images distinctes : une pour les nœuds de calcul GPU exécutant le modèle et une pour les bacs à sable isolés exécutant le code d'évaluation non fiable.

1. Créer l'image de nœud de calcul GPU

Le nœud de calcul GPU Ray a besoin de bibliothèques pour exécuter le modèle de langage et orchestrer la boucle d'entraînement. Nous créons cette image sur la base de l'image vLLM officielle. Elle est donc compatible avec les derniers GPU et inclut PyTorch/CUDA préinstallés.

Exécutez la commande suivante pour créer 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

Créez l'image et transférez-la vers votre dépôt 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

Remarque : Ce guide utilise des commandes docker locales pour créer des images. Si vous préférez créer les images à distance, vous pouvez utiliser Cloud Build (par exemple, avec gcloud builds submit).

2. Créer l'image de l'en-tête du processeur

Le nœud principal Ray orchestre uniquement le cluster et n'exécute pas les modèles d'entraînement GPU intensifs. Pour éviter un goulot d'étranglement massif lors de l'extraction d'images (généralement 15 Go ou plus) sur vos nœuds de processeur standards, nous créons une image légère, réservée au processeur, pour le nœud principal. Cette image contient Ray et les bibliothèques Python requises, mais exclut les bibliothèques GPU lourdes telles que CUDA et vLLM.

Exécutez la commande suivante pour créer 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

Créez et transférez l'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. Créer l'image du bac à sable

Le bac à sable a besoin des dépendances spécifiques à la tâche que nous évaluons afin que l'installation au moment de l'exécution soit instantanée. Pour cet atelier de programmation, nous allons utiliser un problème du dépôt django/django dans SWE-bench. Nous allons précloner le dépôt et précompiler les environnements Python afin que nos scripts de modèle ne perdent pas de temps à les télécharger dans la boucle RL.

Exécutez la commande suivante pour créer 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

Créez et transférez l'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. Configurer l'orchestration et l'exécution

Nous allons maintenant déployer le cluster Ray pour l'orchestration et les ressources du bac à sable pour l'exécution.

1. Configuration du cluster Ray

Déployez une ressource personnalisée RayCluster. Notez que les ressources disponibles de votre cluster (comme la mémoire ou le type de processeur ou de GPU) peuvent varier. Ajustez les demandes et les limites resources en conséquence.

Exécutez la commande suivante pour créer raycluster.yaml. Cette commande utilise cat << EOF pour remplacer automatiquement vos variables d'environnement dans le fichier manifeste :

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

Appliquez-le :

kubectl apply -f raycluster.yaml

Vérifiez que le cluster est créé et en cours d'exécution (cela peut prendre quelques minutes) :

kubectl get raycluster

Résultat attendu :

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

2. Configuration de SandboxRouter

SandboxRouter sert de passerelle HTTP à haut débit. Il traite les requêtes des nœuds de calcul Ray et les transmet instantanément aux pods gVisor disponibles, en contournant le cycle de vie plus lent des pods du serveur d'API Kubernetes.

Exécutez la commande suivante pour créer 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

Appliquez-le :

kubectl apply -f sandbox_router.yaml

Vérifiez que le déploiement est en cours d'exécution :

kubectl get deployment sandbox-router-deployment

Résultat attendu :

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

3. Configuration de SandboxTemplate et WarmPool

GKE Agent Sandbox permet d'attribuer instantanément des conteneurs isolés préchauffés à l'aide du routeur Sandbox. Nous définissons un SandboxTemplate et un SandboxWarmPool pour que les pods restent prêts.

Exécutez la commande suivante pour créer sandbox_warmpool.yaml avec vos variables d'environnement :

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

Appliquez-le :

kubectl apply -f sandbox_warmpool.yaml

Vérifiez que SandboxWarmPool est initialisé :

kubectl get sandboxwarmpool

Résultat attendu :

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

4. Isolement de sécurité

Une NetworkPolicy isole strictement les bacs à sable, empêchant toute sortie vers le serveur de métadonnées GCP et, par conséquent, tout vol de jetons IAM.

Exécutez la commande suivante pour créer 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

Appliquez la règle :

kubectl apply -f network_policy.yaml

Vérifiez que la règle NetworkPolicy a été créée :

kubectl get networkpolicy

Résultat attendu :

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

5. Job RL de base avec SweBench et TRL

Une fois le cluster et les bacs à sable préparés, nous pouvons exécuter une boucle d'entraînement GRPO. Nous utiliserons la bibliothèque trl pour orchestrer l'algorithme GRPO et les fonctions à distance Ray pour évaluer le code généré dans les bacs à sable isolés.

Pour que l'exécution soit rapide dans cet atelier de programmation, nous allons filtrer pour n'afficher qu'un seul problème Django. La logique de routage ci-dessous montre comment sélectionner différentes warmpools pour différents dépôts, ce qui est utile lors de l'extension à l'ensemble de données SWE-bench complet.

Script d'entraînement

Exécutez la commande suivante pour créer 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

Envoyer le job au cluster

Commencez par transférer le port vers le tableau de bord Ray Head et envoyez la tâche d'entraînement depuis votre ordinateur local :

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

Surveiller l'exécution

Vous pouvez suivre la progression de votre exécution :

  • Tableau de bord Ray : ouvrez http://localhost:8265 dans votre navigateur.
  • Revendications de bac à sable : regardez GKE revendiquer et libérer dynamiquement des bacs à sable sous gVisor :
    watch -n 1 "kubectl get sandboxclaims,sandboxes,pods"
    

6. Conclusion

Félicitations ! Vous avez réussi à configurer et à exécuter une boucle d'entraînement RL distribuée à hautes performances de manière sécurisée sur GKE Standard à l'aide des bacs à sable de l'agent GKE.