Add kubernetes controller code

examples/scripts/README.md

@@ -1,10 +1,13 @@
 # To run PolyRL training pipeline
 
-# To launch the training process
+## To launch the training process
 
-1. Launch an RL trainer
+1. Config rollout manager
+The configuration file of rollout manager is `src/rlboost/rollout-manager/config.toml`.
+- `allowed_sender_ips` is for you to specify the range of ips for weight transfer from trainer to rollout instances
+- `num_mooncake_groups_per_sender` is the max number of weight transfer groups per sender ip, each group will shard model weight into `num_mooncake_engines_per_group` and transfer weight in parallel to maximize bandwidth.
 
-Update `SENDER_IPS` in `run_async_grpo_pipeline.sh`. It means the network interfaces that you want to use for weight transfer.
+2. Launch an RL trainer
 
 ```bash
 export GSM8K_DATA_DIR="path/to/gsm8k/data"
@@ -13,7 +16,7 @@ bash examples/scripts/run_async_grpo_pipeline.sh
 
 It will automatically compile and launch the rollout manager and the weight transfer agent. The default port of the rollout manager is `5000`.
 
-2. Launch a rollout instance.
+3. Launch a rollout instance.
 
 Update the address in `launch_sglang.sh`
 ```bash
@@ -26,8 +29,27 @@ On each remote rollout engine, launch
 bash examples/scripts/launch_sglang.sh
 ```
 
-3. Run colocated RL baseline
+4. Run colocated RL baseline
 
 ```bash
 bash examples/scripts/run_sync_grpo_default.sh
 ```
+
+## To run trainer on multiple nodes
+
+1. Start a ray cluster on the root node
+```bash
+ray start --head --dashboard-host=0.0.0.0
+```
+
+2. On other node, join the cluster by 
+```bash
+ray start --address='<follow the prompt on the root node>'
+```
+
+3. On the root node, start the job by
+```bash
+RAY_API_SERVER_ADDRESS='<follow the prompt on the root node>' ray job submit --working-dir . -- bash <your multi-node script>
+```
+
+

examples/scripts/run_async_grpo_pipeline.sh

@@ -9,8 +9,6 @@ ROLLOUT_NAME=${ROLLOUT_NAME:-"sglang-disaggregated"}
 EXP_NAME="qwen3_1.7b_grpo"
 TIMESTAMP=$(date +"%Y-%m-%d_%H-%M-%S")
 
-SENDER_IPS="<SENDER_IPs>" # e.g. '192.168.0.0/16'
-
 echo "Starting GRPO training..."
 # Run PPO training
 RAY_DEDUP_LOGS=0 \

examples/scripts/run_async_grpo_pipeline_n2.sh

@@ -0,0 +1,61 @@
+#!/bin/bash
+
+ulimit -n 65536  # Increase max open files
+
+# Set dataset directory - use environment variable GSM8K_DATA_DIR if set, otherwise use default
+GSM8K_DATA_DIR=${GSM8K_DATA_DIR:-"~/data/gsm8k"}
+# Choose from sglang, sglang-disaggregated
+ROLLOUT_NAME=${ROLLOUT_NAME:-"sglang-disaggregated"}
+EXP_NAME="qwen3_1.7b_grpo"
+TIMESTAMP=$(date +"%Y-%m-%d_%H-%M-%S")
+
+echo "Starting GRPO training..."
+# Run PPO training
+RAY_DEDUP_LOGS=0 \
+PYTHONUNBUFFERED=1 python3 -m verl.trainer.main_stream \
+ algorithm.adv_estimator=grpo \
+ data.train_files=${GSM8K_DATA_DIR}/train.parquet \
+ data.val_files=${GSM8K_DATA_DIR}/test.parquet \
+ data.train_batch_size=128 \
+ data.val_batch_size=32 \
+ data.max_prompt_length=512 \
+ data.max_response_length=14336 \
+ data.filter_overlong_prompts=True \
+ data.truncation='error' \
+ actor_rollout_ref.nccl_timeout=2000 \
+ actor_rollout_ref.rollout.name=${ROLLOUT_NAME} \
+ actor_rollout_ref.model.path=Qwen/Qwen3-1.7B \
+ actor_rollout_ref.actor.optim.lr=1e-6 \
+ actor_rollout_ref.model.use_remove_padding=True \
+ actor_rollout_ref.actor.ppo_mini_batch_size=64 \
+ actor_rollout_ref.actor.use_kl_loss=True \
+ actor_rollout_ref.actor.kl_loss_coef=0.001 \
+ actor_rollout_ref.actor.kl_loss_type=low_var_kl \
+ actor_rollout_ref.actor.entropy_coeff=0 \
+ actor_rollout_ref.actor.use_dynamic_bsz=True \
+ actor_rollout_ref.actor.ppo_max_token_len_per_gpu=16384 \
+ actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=8 \
+ actor_rollout_ref.rollout.min_stream_batch_size=16 \
+ actor_rollout_ref.rollout.calculate_log_probs=True \
+ actor_rollout_ref.model.enable_gradient_checkpointing=True \
+ actor_rollout_ref.actor.fsdp_config.fsdp_size=8 \
+ actor_rollout_ref.actor.fsdp_config.param_offload=False \
+ actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \
+ actor_rollout_ref.rollout.tensor_model_parallel_size=2 \
+ actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \
+ actor_rollout_ref.rollout.n=8 \
+ actor_rollout_ref.rollout.free_cache_engine=True \
+ actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=8 \
+ algorithm.use_kl_in_reward=False \
+ trainer.logger=['console','tensorboard'] \
+ trainer.project_name='verl_grpo_example' \
+ trainer.experiment_name="${EXP_NAME}_${TIMESTAMP}" \
+ trainer.val_before_train=False \
+ trainer.balance_batch=True \
+ trainer.critic_warmup=0 \
+ trainer.n_gpus_per_node=8 \
+ trainer.stream_fit=True \
+ trainer.nnodes=2 \
+ trainer.save_freq=20 \
+ trainer.test_freq=10 \
+ trainer.total_epochs=15 2>&1 | tee verl_demo.log

kube/.gitignore

@@ -0,0 +1,14 @@
+# rust
+rust/target/
+rust/Cargo.lock
+
+# python
+*.pyc
+*.pyo
+*.pyd
+*.pyw
+*.pyz
+*.pywz
+*.pyzw
+*.pyzwz
+**/__pycache__/

kube/README.md

@@ -0,0 +1,76 @@
+# GKE setup guide for PolyRL
+
+## Preliminaries
+
+### VPC setup (Recommend)
+
+TODO: add management network setup instructions
+
+### Build and Push Docker Image
+
+```bash
+cd docker
+```
+
+#### Create a Repository on GCloud to Save Docker Images
+
+```bash
+gcloud artifacts repositories create "<your repo name>" \
+    --project="<your project name>" \
+    --repository-format="docker" \
+    --location="<your project region>" \
+    --description="Docker repository for PolyRL workloads"
+
+gcloud auth configure-docker "<your project region>-docker.pkg.dev" --project="<your project name>"
+```
+
+#### Build and push image
+
+We provide a docker file to test functionality of Kubernetes instance.
+You may build and push the repo following
+```bash
+docker build -t "<your image name>-test:v1" -f "docker/test.dockerfile" .
+docker push "<your image name>-test:v1"
+
+docker build -t "<your image name>-rollout:v1" -f "docker/rollout.dockerfile" .
+docker push "<your image name>-rollout:v1"
+```
+
+## Create a Kubernetes Cluster
+
+Use the Python CLI to create the cluster from a TOML config (e.g., `config/example.toml`):
+```bash
+python cluster_setup/gke_setup.py --config example create-cluster
+```
+It initializes a cluster with the `default_pool` settings and optionally connects `kubectl` automatically (add `--skip-connect` to skip).
+
+## Create Node Pools
+
+Create worker pool(s) based on `[[worker_pools]]` entries in your TOML:
+```bash
+# Create all worker pools defined in the config
+python cluster_setup/gke_setup.py --config example create-worker-pool -1
+
+# Create only the first worker pool (index 0)
+python cluster_setup/gke_setup.py --config example create-worker-pool 0
+```
+Pools are named from machine type and provisioning, with an index suffix (e.g., `a3-highgpu-2g-spot-gpu-pool-0`). You can override the base name with `--pool-name`, and autoscaling bounds with `--min-nodes`, `--max-nodes`.
+
+If you want to update an existing pool, add `--delete-first`:
+```bash
+python cluster_setup/gke_setup.py --config example create-worker-pool 0 --delete-first
+```
+
+For a dry-run to preview commands:
+```bash
+python cluster_setup/gke_setup.py --dry-run --config example create-worker-pool -1
+```
+
+## Run PolyRL with Spot Instances
+
+### Start Trainer of PolyRL 
+Please refer to `examples/scripts/README.md`.
+
+### Start Rollout Deployment
+
+Please refer to `kube/rust/README.md`.

kube/cluster_setup/README.md

@@ -0,0 +1,64 @@
+# GKE setup CLI
+
+This directory provides a single Python CLI to create the GKE cluster and GPU worker pools from TOML configs under `config/`.
+
+Prerequisites
+- Python 3.9+ with the `toml` package: `pip install toml`
+- gcloud SDK authenticated and configured
+
+Config shape (excerpt)
+```toml
+[general]
+project = "example-project"
+region = "us-central1"
+zone = "us-central1-a"
+cluster_name = "example-cluster"
+
+[default_pool]
+machine_type = "e2-medium"
+min_nodes = 1
+max_nodes = 4
+
+[[worker_pools]]
+machine_type = "n1-standard-1"
+provision = "spot" # or "standard"
+accelerator = "type=nvidia-tesla-t4,count=1,gpu-driver-version=latest"
+min_nodes = 0
+max_nodes = 4
+```
+
+Basic usage
+```bash
+# Create cluster from config/example.toml
+python cluster_setup/gke_setup.py --config example create-cluster
+
+# Skip connecting kubectl context after creation
+python cluster_setup/gke_setup.py --config example create-cluster --skip-connect
+
+# Create ALL worker pools defined in [[worker_pools]] (index -1 means all)
+python cluster_setup/gke_setup.py --config example create-worker-pool -1
+
+# Create only the Nth worker pool (0-based index)
+python cluster_setup/gke_setup.py --config example create-worker-pool 0
+python cluster_setup/gke_setup.py --config example create-worker-pool 1
+
+# Delete-before-create when updating an existing pool
+python cluster_setup/gke_setup.py --config example create-worker-pool 0 --delete-first
+
+# Override base name and autoscaling bounds (a "-N" suffix will be added)
+python cluster_setup/gke_setup.py --config example create-worker-pool 0 \
+  --pool-name custom-pool --min-nodes 0 --max-nodes 6
+
+# Dry-run to preview the gcloud commands without execution
+python cluster_setup/gke_setup.py --dry-run --config example create-worker-pool -1
+
+# Use an absolute TOML path
+python cluster_setup/gke_setup.py --config /abs/path/to/file.toml create-cluster
+```
+
+Notes
+- `--config` accepts a short name (resolved to `config/<name>.toml`) or an absolute path.
+- Worker pool names default to `<machine_type>-<provision>-gpu-pool-N` (e.g., `a3-highgpu-2g-spot-gpu-pool-0`).
+- Output is colorized where supported and shows a labeled "Command:" preview with line continuations.
+- Set `NO_COLOR=1` to disable colored output.
+