From a9db5b1cbcbdabf6ceee5afe0a21b0e34690b68d Mon Sep 17 00:00:00 2001 From: Xueshen Liu Date: Tue, 28 Oct 2025 04:38:30 +0000 Subject: [PATCH] attach kubernetes controller code --- examples/scripts/README.md | 32 +- examples/scripts/run_async_grpo_pipeline.sh | 2 - .../scripts/run_async_grpo_pipeline_n2.sh | 61 +++ kube/.gitignore | 14 + kube/README.md | 76 ++++ kube/cluster_setup/README.md | 64 +++ kube/cluster_setup/gke_setup.py | 397 ++++++++++++++++++ kube/config/example.toml | 35 ++ kube/deployment/kube-rollout-manifest.yaml | 57 +++ kube/deployment/kube-sglang-manifest.yaml | 55 +++ kube/deployment/kube-test-manifest.yaml | 35 ++ kube/docker/launch_rollout.sh | 40 ++ kube/docker/launch_sglang.sh | 31 ++ kube/docker/rollout.dockerfile | 30 ++ kube/docker/test.dockerfile | 23 + kube/docker/test_script.py | 56 +++ kube/rust/Cargo.toml | 26 ++ kube/rust/README.md | 89 ++++ kube/rust/src/main.rs | 77 ++++ kube/rust/src/scaler.rs | 95 +++++ kube/rust/src/server.rs | 116 +++++ kube/rust/src/settings.rs | 137 ++++++ kube/rust/src/watcher.rs | 145 +++++++ src/rlboost/rollout-manager/config.toml | 4 +- 24 files changed, 1688 insertions(+), 9 deletions(-) create mode 100755 examples/scripts/run_async_grpo_pipeline_n2.sh create mode 100644 kube/.gitignore create mode 100644 kube/README.md create mode 100644 kube/cluster_setup/README.md create mode 100644 kube/cluster_setup/gke_setup.py create mode 100644 kube/config/example.toml create mode 100644 kube/deployment/kube-rollout-manifest.yaml create mode 100644 kube/deployment/kube-sglang-manifest.yaml create mode 100644 kube/deployment/kube-test-manifest.yaml create mode 100644 kube/docker/launch_rollout.sh create mode 100644 kube/docker/launch_sglang.sh create mode 100644 kube/docker/rollout.dockerfile create mode 100644 kube/docker/test.dockerfile create mode 100644 kube/docker/test_script.py create mode 100644 kube/rust/Cargo.toml create mode 100644 kube/rust/README.md create mode 100644 kube/rust/src/main.rs create mode 100644 kube/rust/src/scaler.rs create mode 100644 kube/rust/src/server.rs create mode 100644 kube/rust/src/settings.rs create mode 100644 kube/rust/src/watcher.rs diff --git a/examples/scripts/README.md b/examples/scripts/README.md index 33d5550..ec77b6f 100644 --- a/examples/scripts/README.md +++ b/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='' +``` + +3. On the root node, start the job by +```bash +RAY_API_SERVER_ADDRESS='' ray job submit --working-dir . -- bash +``` + + diff --git a/examples/scripts/run_async_grpo_pipeline.sh b/examples/scripts/run_async_grpo_pipeline.sh index 29f3886..e8798e3 100755 --- a/examples/scripts/run_async_grpo_pipeline.sh +++ b/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="" # e.g. '192.168.0.0/16' - echo "Starting GRPO training..." # Run PPO training RAY_DEDUP_LOGS=0 \ diff --git a/examples/scripts/run_async_grpo_pipeline_n2.sh b/examples/scripts/run_async_grpo_pipeline_n2.sh new file mode 100755 index 0000000..58db66f --- /dev/null +++ b/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 diff --git a/kube/.gitignore b/kube/.gitignore new file mode 100644 index 0000000..3acedce --- /dev/null +++ b/kube/.gitignore @@ -0,0 +1,14 @@ +# rust +rust/target/ +rust/Cargo.lock + +# python +*.pyc +*.pyo +*.pyd +*.pyw +*.pyz +*.pywz +*.pyzw +*.pyzwz +**/__pycache__/ \ No newline at end of file diff --git a/kube/README.md b/kube/README.md new file mode 100644 index 0000000..95b1611 --- /dev/null +++ b/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 "" \ + --project="" \ + --repository-format="docker" \ + --location="" \ + --description="Docker repository for PolyRL workloads" + +gcloud auth configure-docker "-docker.pkg.dev" --project="" +``` + +#### 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 "-test:v1" -f "docker/test.dockerfile" . +docker push "-test:v1" + +docker build -t "-rollout:v1" -f "docker/rollout.dockerfile" . +docker push "-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`. \ No newline at end of file diff --git a/kube/cluster_setup/README.md b/kube/cluster_setup/README.md new file mode 100644 index 0000000..07b46a0 --- /dev/null +++ b/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/.toml`) or an absolute path. +- Worker pool names default to `--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. + diff --git a/kube/cluster_setup/gke_setup.py b/kube/cluster_setup/gke_setup.py new file mode 100644 index 0000000..aedfb71 --- /dev/null +++ b/kube/cluster_setup/gke_setup.py @@ -0,0 +1,397 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import argparse +import os +import shlex +import subprocess +import sys +from pathlib import Path + + +def _supports_color() -> bool: + return sys.stdout.isatty() and os.environ.get("NO_COLOR") is None + + +def _style(kind: str, text: str) -> str: + if not _supports_color(): + if kind == "title": + return f"== {text} ==" + if kind == "warn": + return f"[!] {text}" + return text + styles = { + "title": "\033[1;36m", # bold cyan + "info": "\033[1;32m", # bold green + "warn": "\033[1;33m", # bold yellow + "dim": "\033[2m", # dim + "reset": "\033[0m", + } + prefix = styles.get(kind, "") + reset = styles["reset"] + return f"{prefix}{text}{reset}" + + +def _print_section(title: str) -> None: + bar = "=" * max(6, min(78, len(title) + 10)) + print("\n" + _style("title", title)) + print(_style("dim", bar)) + + +def load_toml_bytes(path: Path) -> dict: + import toml + + with path.open("r", encoding="utf-8") as f: + return toml.load(f) + + +def repo_root() -> Path: + return Path(__file__).resolve().parents[1] + + +def resolve_config_path(config: str) -> Path: + candidate = Path(config) + if candidate.exists(): + return candidate + # Treat as a short name under config/.toml + by_name = repo_root() / "config" / f"{config}.toml" + if by_name.exists(): + return by_name + raise FileNotFoundError(f"Config not found: {config}") + + +def get_worker_pools(cfg: dict) -> list: + pools = cfg.get("worker_pools", []) + if isinstance(pools, dict): + return [pools] + return list(pools) + + +def run(cmd: list[str], dry_run: bool) -> None: + cmd_str = " ".join(shlex.quote(c) for c in cmd).replace(" --", " \\\n --") + print(_style("info", "Command:")) + print("$", cmd_str) + if dry_run: + return + subprocess.run(cmd, check=True) + + +def create_cluster(cfg: dict, dry_run: bool, skip_connect: bool) -> None: + g = cfg["general"] + dpool = cfg.get("default_pool", {}) + + project = str(g["project"]) + region = str(g["region"]) + zone = str(g["zone"]) + cluster_name = str(g["cluster_name"]) + machine_type = str(dpool.get("node_type", "e2-medium")) + min_nodes = str(dpool.get("min_nodes", 1)) + max_nodes = str(dpool.get("max_nodes", 4)) + + _print_section(f"Create cluster '{cluster_name}' in {region} ({zone})") + print(_style("dim", f"project={project} machine_type={machine_type} autoscaling=[{min_nodes},{max_nodes}]")) + + cmd = [ + "gcloud", + "container", + "clusters", + "create", + cluster_name, + "--project", + project, + "--region", + region, + "--node-locations", + zone, + "--release-channel", + "regular", + "--machine-type", + machine_type, + "--num-nodes", + "1", + "--enable-autoscaling", + "--min-nodes", + min_nodes, + "--max-nodes", + max_nodes, + "--enable-ip-alias", + "--workload-pool", + f"{project}.svc.id.goog", + ] + + if g.get("vpc_net") and g.get("vpc_subnet"): + vpc_net = str(g["vpc_net"]) + vpc_subnet = str(g["vpc_subnet"]) + cmd.extend([ + "--network", + vpc_net, + "--subnetwork", + vpc_subnet, + ]) + + run(cmd, dry_run) + if skip_connect: + print(_style("warn", "Next step:")) + print("$", f"gcloud container clusters get-credentials {cluster_name} --region {region} --project {project}") + else: + _print_section("Connect kubectl context") + credentials_cmd = [ + "gcloud", + "container", + "clusters", + "get-credentials", + cluster_name, + "--region", + region, + "--project", + project, + ] + run(credentials_cmd, dry_run) + + +def create_rollout_pool(cfg: dict, dry_run: bool, delete_first: bool, pool_name: str | None = None, + min_nodes: int | None = None, max_nodes: int | None = None) -> None: + g = cfg["general"] + rpool = cfg["rollout_pool"] + + project = str(g["project"]) + region = str(g["region"]) + zone = str(g["zone"]) + cluster_name = str(g["cluster_name"]) + + machine_type = str(rpool["node_type"]) + accelerator = str(rpool["accelerator"]) + min_nodes_val = str(min_nodes if min_nodes is not None else rpool.get("min_nodes", 0)) + max_nodes_val = str(max_nodes if max_nodes is not None else rpool.get("max_nodes", 4)) + + provision = str(rpool["provision"]).lower() + if provision != "standard" and provision != "spot": + print(_style("warn", f"Unknown provision type: {provision}. Using 'spot' by default. Specify 'spot' or 'standard'.")) + + pool = pool_name or f"{machine_type}-{provision}-gpu-pool" + + _print_section(f"Create node pool '{pool}'") + print(_style("dim", f"cluster={cluster_name} region={region} machine_type={machine_type} accel=[{accelerator}] {provision} autoscaling=[{min_nodes_val},{max_nodes_val}]")) + + if delete_first: + _print_section(f"Delete existing node pool '{pool}' if present") + del_cmd = [ + "gcloud", + "container", + "node-pools", + "delete", + pool, + "--cluster", + cluster_name, + "--project", + project, + "--region", + region, + ] + run(del_cmd, dry_run) + + create_cmd = [ + "gcloud", + "container", + "node-pools", + "create", + pool, + "--cluster", + cluster_name, + "--project", + project, + "--region", + region, + "--node-locations", + zone, + "--machine-type", + machine_type, + "--accelerator", + accelerator, + f"--{provision}", + "--enable-autoscaling", + "--num-nodes", + min_nodes_val, + "--min-nodes", + min_nodes_val, + "--max-nodes", + max_nodes_val, + "--node-taints", + "nvidia.com/gpu=present:NoSchedule", + "--enable-image-streaming", + ] + run(create_cmd, dry_run) + + +def create_worker_pool_entry(cfg: dict, pool_entry: dict, index: int, dry_run: bool, delete_first: bool, + pool_name_override: str | None = None, + min_nodes_override: int | None = None, + max_nodes_override: int | None = None) -> None: + g = cfg["general"] + + project = str(g["project"]) + region = str(g["region"]) + zone = str(g["zone"]) + cluster_name = str(g["cluster_name"]) + + machine_type = str(pool_entry["machine_type"]) + accelerator = str(pool_entry.get("accelerator", "")) + provision = str(pool_entry.get("provision", "spot")).lower() + if provision not in ("standard", "spot"): + print(_style("warn", f"Unknown provision type: {provision}. Using 'spot'.")) + provision = "spot" + + min_nodes_val = str(min_nodes_override if min_nodes_override is not None else pool_entry.get("min_nodes", 0)) + max_nodes_val = str(max_nodes_override if max_nodes_override is not None else pool_entry.get("max_nodes", 4)) + + base_name = pool_name_override or f"{machine_type}-{provision}-gpu-pool" + pool = f"{base_name}-{index}" + + _print_section(f"Create worker pool '{pool}' (index {index})") + print(_style("dim", f"cluster={cluster_name} region={region} machine_type={machine_type} accel=[{accelerator}] {provision} autoscaling=[{min_nodes_val},{max_nodes_val}]")) + + if delete_first: + _print_section(f"Delete existing node pool '{pool}' if present") + del_cmd = [ + "gcloud", + "container", + "node-pools", + "delete", + pool, + "--cluster", + cluster_name, + "--project", + project, + "--region", + region, + ] + run(del_cmd, dry_run) + + create_cmd = [ + "gcloud", + "container", + "node-pools", + "create", + pool, + "--cluster", + cluster_name, + "--project", + project, + "--region", + region, + "--node-locations", + zone, + "--machine-type", + machine_type, + ] + + if accelerator: + create_cmd.extend(["--accelerator", accelerator]) + + create_cmd.extend([ + f"--{provision}", + "--enable-autoscaling", + "--num-nodes", + min_nodes_val, + "--min-nodes", + min_nodes_val, + "--max-nodes", + max_nodes_val, + "--node-taints", + "nvidia.com/gpu=present:NoSchedule", + "--enable-image-streaming", + ]) + + run(create_cmd, dry_run) + + +def create_worker_pools(cfg: dict, index: int, dry_run: bool, delete_first: bool, + pool_name: str | None, min_nodes: int | None, max_nodes: int | None) -> None: + pools = get_worker_pools(cfg) + if index == -1: + for i, entry in enumerate(pools): + create_worker_pool_entry(cfg, entry, i, dry_run, delete_first, pool_name, min_nodes, max_nodes) + else: + if index < 0 or index >= len(pools): + raise IndexError(f"worker pool index out of range: {index} (have {len(pools)})") + create_worker_pool_entry(cfg, pools[index], index, dry_run, delete_first, pool_name, min_nodes, max_nodes) + + +def main(argv: list[str]) -> int: + parser = argparse.ArgumentParser( + description=( + "Manage GKE cluster and GPU node pools using a TOML config (config/*.toml).\n\n" + "Examples:\n" + " ./cluster_setup/gke_setup.py --config example create-cluster\n" + " ./cluster_setup/gke_setup.py --config example create-worker-pool -1\n" + " ./cluster_setup/gke_setup.py --dry-run --config example create-worker-pool 0 --delete-first\n" + ), + formatter_class=argparse.RawTextHelpFormatter, + ) + parser.add_argument( + "--config", + required=True, + help="Path to TOML or short name (e.g., 'example' => config/example.toml)", + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="Print gcloud commands without executing", + ) + + sub = parser.add_subparsers(dest="command", required=True) + # Create cluster + p_cluster = sub.add_parser("create-cluster", help="Create the GKE cluster defined in [general] and [default_pool]") + p_cluster.add_argument("--skip-connect", action="store_true", help="Skip connection to cluster after creation") + + # Create rollout GPU node pool + p_workers = sub.add_parser( + "create-worker-pool", + help="Create worker pool(s) based on [[worker_pools]] entries in TOML. Pass index n, or -1 for all.", + ) + p_workers.add_argument("n", type=int, help="Worker pool index (0-based). Use -1 to create all") + p_workers.add_argument("--delete-first", action="store_true", help="Delete pool if it exists") + p_workers.add_argument("--pool-name", help="Override base pool name (default from machine type and provisioning)") + p_workers.add_argument("--min-nodes", type=int, help="Override min nodes") + p_workers.add_argument("--max-nodes", type=int, help="Override max nodes") + + args = parser.parse_args(argv) + + cfg_path = resolve_config_path(args.config) + cfg = load_toml_bytes(cfg_path) + + if args.command == "create-cluster": + create_cluster(cfg, args.dry_run, args.skip_connect) + elif args.command == "create-worker-pool": + create_worker_pools( + cfg, + index=args.n, + dry_run=args.dry_run, + delete_first=args.delete_first, + pool_name=args.pool_name, + min_nodes=args.min_nodes, + max_nodes=args.max_nodes, + ) + else: + parser.error(f"Unknown command: {args.command}") + + return 0 + + +if __name__ == "__main__": + try: + sys.exit(main(sys.argv[1:])) + except subprocess.CalledProcessError as e: + print(f"Command failed with exit code {e.returncode}", file=sys.stderr) + sys.exit(e.returncode) + diff --git a/kube/config/example.toml b/kube/config/example.toml new file mode 100644 index 0000000..493209d --- /dev/null +++ b/kube/config/example.toml @@ -0,0 +1,35 @@ +[general] +project = "example-project" +region = "us-central1" +zone = "us-central1-a" +cluster_name = "example-cluster" +vpc_net = "projects/example-project/global/networks/example-mgmt-net" +vpc_subnet = "projects/example-project/regions/us-central1/subnetworks/example-mgmt-sub" +namespace = "default" +logs_dir = "logs/example" +bind_addr = "0.0.0.0:5000" + +[default_pool] +machine_type = "e2-medium" +min_nodes = 1 +max_nodes = 4 + +[[worker_pools]] +machine_type = "n1-standard-1" +provision = "spot" +accelerator = "type=nvidia-tesla-t4,count=1,gpu-driver-version=latest" +min_nodes = 0 +max_nodes = 8 + +[[worker_pools]] +machine_type = "a3-highgpu-2g" +provision = "spot" +accelerator = "type=nvidia-h100-80gb,count=2,gpu-driver-version=latest" +min_nodes = 0 +max_nodes = 8 + +[deployments] +files = ["deployment/kube-test-manifest.yaml"] + +[scaler] +scale_up_step = 1 \ No newline at end of file diff --git a/kube/deployment/kube-rollout-manifest.yaml b/kube/deployment/kube-rollout-manifest.yaml new file mode 100644 index 0000000..158ccaf --- /dev/null +++ b/kube/deployment/kube-rollout-manifest.yaml @@ -0,0 +1,57 @@ +apiVersion: apps/v1 +kind: Deployment +metadata: + name: rollout-app +spec: + replicas: 0 + selector: + matchLabels: + app: rollout + template: + metadata: + labels: + app: rollout + spec: + # 1. TOLERATION: Allows this pod to be scheduled on a tainted GPU node + tolerations: + - key: "nvidia.com/gpu" + operator: "Exists" + effect: "NoSchedule" + + # 2. NODE SELECTOR: Ensures this pod only runs on nodes with a T4 GPU + nodeSelector: + cloud.google.com/gke-accelerator: nvidia-h100-80gb + cloud.google.com/gke-spot: "true" + + # MODIFICATION: Add a memory-backed volume for shared memory (/dev/shm) + volumes: + - name: dshm + emptyDir: + medium: Memory + sizeLimit: 512Gi # A generous size for multi-GPU communication + + # Add these two lines to the Pod specification to use host network + hostNetwork: true + hostUsers: true + + containers: + - name: rl-rollout + image: "-rollout:v1" + command: ["bash", "launch_rollout.sh"] + + # MODIFICATION: Add environment variable to set the host + env: + - name: SGLANG_MODEL_NAME + value: "Qwen/Qwen3-8B" + - name: ROLLOUT_MGR + value: "http://:5000" + + # MODIFICATION: Mount the shared memory volume into the container + volumeMounts: + - mountPath: /dev/shm + name: dshm + + # 3. RESOURCE LIMIT: Requests 2 GPUs for this container + resources: + limits: + nvidia.com/gpu: "2" \ No newline at end of file diff --git a/kube/deployment/kube-sglang-manifest.yaml b/kube/deployment/kube-sglang-manifest.yaml new file mode 100644 index 0000000..ad879a2 --- /dev/null +++ b/kube/deployment/kube-sglang-manifest.yaml @@ -0,0 +1,55 @@ +apiVersion: apps/v1 +kind: Deployment +metadata: + name: sglang-app +spec: + replicas: 0 + selector: + matchLabels: + app: sglang + template: + metadata: + labels: + app: sglang + spec: + # 1. TOLERATION: Allows this pod to be scheduled on a tainted GPU node + tolerations: + - key: "nvidia.com/gpu" + operator: "Exists" + effect: "NoSchedule" + + # 2. NODE SELECTOR: Ensures this pod only runs on nodes with a T4 GPU + nodeSelector: + cloud.google.com/gke-accelerator: nvidia-h100-80gb + cloud.google.com/gke-spot: "true" + + # MODIFICATION: Add a memory-backed volume for shared memory (/dev/shm) + volumes: + - name: dshm + emptyDir: + medium: Memory + sizeLimit: 512Gi # A generous size for multi-GPU communication + + # Add these two lines to the Pod specification to use host network + hostNetwork: true + hostUsers: true + + containers: + - name: rl-rollout + image: "-rollout:v1" + command: ["bash", "launch_sglang.sh"] + + # MODIFICATION: Add environment variable to set the host + env: + - name: SGLANG_MODEL_NAME + value: "Qwen/Qwen3-8B" + + # MODIFICATION: Mount the shared memory volume into the container + volumeMounts: + - mountPath: /dev/shm + name: dshm + + # 3. RESOURCE LIMIT: Requests 2 GPUs for this container + resources: + limits: + nvidia.com/gpu: "2" \ No newline at end of file diff --git a/kube/deployment/kube-test-manifest.yaml b/kube/deployment/kube-test-manifest.yaml new file mode 100644 index 0000000..8ebfbd1 --- /dev/null +++ b/kube/deployment/kube-test-manifest.yaml @@ -0,0 +1,35 @@ +apiVersion: apps/v1 +kind: Deployment +metadata: + name: gpu-test-app +spec: + replicas: 0 + selector: + matchLabels: + app: gpu-test + template: + metadata: + labels: + app: gpu-test + spec: + # 1. TOLERATION: Allows this pod to be scheduled on a tainted GPU node + tolerations: + - key: "nvidia.com/gpu" + operator: "Exists" + effect: "NoSchedule" + + # 2. NODE SELECTOR: Ensures this pod only runs on nodes with a T4 GPU + nodeSelector: + cloud.google.com/gke-accelerator: nvidia-tesla-t4 + cloud.google.com/gke-spot: "true" + + containers: + - name: cuda-pytorch-test + image: "-test:v1" + command: ["llm/bin/python3", "test_script.py"] + + # 3. RESOURCE LIMIT: Requests 1 GPU for this container + resources: + limits: + nvidia.com/gpu: "1" + diff --git a/kube/docker/launch_rollout.sh b/kube/docker/launch_rollout.sh new file mode 100644 index 0000000..ffa4363 --- /dev/null +++ b/kube/docker/launch_rollout.sh @@ -0,0 +1,40 @@ +#!/bin/bash + +# check if rollout manager addr is set, otherwise quit +if [[ -z "${ROLLOUT_MGR}" ]]; then + echo "Error: ROLLOUT_MGR is not defined. Please specify as 'http://:' " + exit 1 +fi + +# Read from environment variables, using a default value if not set. +MODEL_NAME=${SGLANG_MODEL_NAME:-"Qwen/Qwen3-1.7B"} +HOST=${SGLANG_HOST:-$(hostname -i)} +PORT=${SGLANG_PORT:-"40000"} +MEM_FRAC=${SGLANG_MEM_FRAC:-"0.6"} +MAX_NUM_REQ=${SGLANG_MAX_REQ:-"128"} +TP_SIZE=${SGLANG_TP_SIZE:-2} + +# --- Sanity Checks & User Feedback --- +echo "✅ Starting SGLang server with the following configuration:" +echo " - Model: $MODEL_NAME" +echo " - Host: $HOST" +echo " - Port: $PORT" +echo " - TP Size: $TP_SIZE" +echo " - Rollout manager: $ROLLOUT_MGR" + +# --- Launch the Server --- +python -m sglang.launch_server \ + --model-path "$MODEL_NAME" \ + --port "$PORT" \ + --host "$HOST" \ + --grammar-backend outlines \ + --tp-size "$TP_SIZE" \ + --mem-fraction-static "$MEM_FRAC" \ + --max-running-requests "$MAX_NUM_REQ" \ + --enable-memory-saver \ + --enable-mixed-chunk \ + --stream-output \ + --stream-interval 10 \ + --mooncake-handshake-port 20000 \ + --enable-weight-transfer-agent \ + --rollout-manager-address "$ROLLOUT_MGR" \ No newline at end of file diff --git a/kube/docker/launch_sglang.sh b/kube/docker/launch_sglang.sh new file mode 100644 index 0000000..b97e03d --- /dev/null +++ b/kube/docker/launch_sglang.sh @@ -0,0 +1,31 @@ +#!/bin/bash + +# Read from environment variables, using a default value if not set. +# This is more flexible than using positional arguments ($1, $2, etc.). +MODEL_NAME=${SGLANG_MODEL_NAME:-"Qwen/Qwen3-1.7B"} +HOST=${SGLANG_HOST:-$(hostname -i)} +PORT=${SGLANG_PORT:-"40000"} +MEM_FRAC=${SGLANG_MEM_FRAC:-"0.6"} +MAX_NUM_REQ=${SGLANG_MAX_REQ:-"128"} +TP_SIZE=${SGLANG_TP_SIZE:-2} + +# --- Sanity Checks & User Feedback --- +echo "✅ Starting SGLang server with the following configuration:" +echo " - Model: $MODEL_NAME" +echo " - Host: $HOST" +echo " - Port: $PORT" +echo " - TP Size: $TP_SIZE" + +# --- Launch the Server --- +python -m sglang.launch_server \ + --model-path "$MODEL_NAME" \ + --port "$PORT" \ + --host "$HOST" \ + --grammar-backend outlines \ + --tp-size "$TP_SIZE" \ + --mem-fraction-static "$MEM_FRAC" \ + --max-running-requests "$MAX_NUM_REQ" \ + --enable-memory-saver \ + --enable-mixed-chunk \ + --stream-output \ + --stream-interval 10 \ No newline at end of file diff --git a/kube/docker/rollout.dockerfile b/kube/docker/rollout.dockerfile new file mode 100644 index 0000000..1e31d8e --- /dev/null +++ b/kube/docker/rollout.dockerfile @@ -0,0 +1,30 @@ +# Start from the official NVIDIA base image, default WORKDIR is /workspace +FROM nvcr.io/nvidia/pytorch:25.05-py3 + +# Set up the environment to be non-interactive and disable Python's output buffering +ENV DEBIAN_FRONTEND=noninteractive +ENV PYTHONUNBUFFERED=1 + +# Layer 1: System dependencies and cleanup +# This layer is cached unless you need to add more apt packages. +RUN apt-get update && \ + # Fix: Corrected typo from ldconfg to ldconfig + ldconfig && \ + # Best Practice: Clean up apt cache to reduce image size + rm -rf /var/lib/apt/lists/* + +# Layer 3: Clone the repository and install the project-specific dependencies +# This layer will only be re-run if the git repo or project requirements change. +RUN git clone -b weight-transfer https://github.com/Terra-Flux/PolyRL.git && \ + cd PolyRL/src/sglang && \ + pip install -e "python[all]" + +# Layer 2: Install standalone Python packages +# This is cached unless you change these specific dependencies. +RUN pip install flash_attn rpyc mooncake-transfer-engine --no-build-isolation && \ + pip cache purge + +# Layer 4: Copy local files +# This is one of the last steps, so changes to the script don't invalidate previous layers. +COPY launch_sglang.sh . +COPY launch_rollout.sh . diff --git a/kube/docker/test.dockerfile b/kube/docker/test.dockerfile new file mode 100644 index 0000000..a9b9f9c --- /dev/null +++ b/kube/docker/test.dockerfile @@ -0,0 +1,23 @@ +# Start from the official NVIDIA base image +FROM nvcr.io/nvidia/cuda:12.8.0-cudnn-devel-ubuntu24.04 + +# Set up the environment to be non-interactive +ENV DEBIAN_FRONTEND=noninteractive + +# *** THE FIX: Add this environment variable to disable output buffering *** +ENV PYTHONUNBUFFERED=1 + +# Install Python, pip +RUN apt-get update && apt-get install -y \ + python3-pip python3.12-venv \ + && rm -rf /var/lib/apt/lists/* + +# Create an /app directory for our code +WORKDIR /app + +# Install PyTorch for CUDA 12.8 in a single, efficient layer +RUN python3 -m venv llm && \ + llm/bin/pip3 install torch torchvision + +# Copy the Python script into the image +COPY test_script.py . diff --git a/kube/docker/test_script.py b/kube/docker/test_script.py new file mode 100644 index 0000000..65d8743 --- /dev/null +++ b/kube/docker/test_script.py @@ -0,0 +1,56 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import time +import os + +print("Launching CUDA test") + +import torch + +# Check if CUDA is available +if not torch.cuda.is_available(): + print('!!! FATAL: CUDA is not available to PyTorch. Exiting.') + exit(1) + +# Set the device to the first GPU +device = torch.device('cuda:0') +pod_name = os.getenv('HOSTNAME', 'unknown-pod') + +print(f'🚀 CUDA test running on pod: {pod_name}') +print(f'--> Using device: {torch.cuda.get_device_name(0)}') + +while True: + try: + # Create two random matrices on the GPU + a = torch.randn(4096, 4096, device=device) + b = torch.randn(4096, 4096, device=device) + + # Perform matrix multiplication + print(f'[{time.strftime("%Y-%m-%d %H:%M:%S")}] Performing 4096x4096 matrix multiplication...') + + start_time = time.time() + c = torch.matmul(a, b) + + # Synchronize to get accurate timing + torch.cuda.synchronize() + end_time = time.time() + + print(f'[{time.strftime("%Y-%m-%d %H:%M:%S")}] Computation complete in {(end_time - start_time):.4f} seconds.') + + except Exception as e: + print(f'An error occurred: {e}') + + # Wait for 10 seconds before the next computation + print("...sleeping for 10 seconds...") + time.sleep(10) diff --git a/kube/rust/Cargo.toml b/kube/rust/Cargo.toml new file mode 100644 index 0000000..106f3bc --- /dev/null +++ b/kube/rust/Cargo.toml @@ -0,0 +1,26 @@ +[package] +name = "gke-rl" +version = "0.1.0" +edition = "2024" + +[dependencies] +anyhow = "1" +axum = { version = "0.7", features = ["json"] } +serde = { version = "1", features = ["derive"] } +serde_json = "1" +toml = "0.8" +serde_yaml = "0.9" +tokio = { version = "1", features = ["full"] } +tracing = "0.1" +tracing-subscriber = { version = "0.3", features = ["env-filter"] } +thiserror = "1" +futures = "0.3" +tokio-stream = "0.1" +hyper = { version = "1", features = ["full"] } +kube = { version = "0.91", features = ["runtime", "derive", "client", "rustls-tls"] } +k8s-openapi = { version = "0.22", features = ["v1_30"] } +tokio-util = { version = "0.7", features = ["codec"] } +tower = "0.5" +tower-http = { version = "0.5", features = ["trace"] } +chrono = { version = "0.4", features = ["clock"] } +uuid = { version = "1", features = ["v4"] } diff --git a/kube/rust/README.md b/kube/rust/README.md new file mode 100644 index 0000000..057bb07 --- /dev/null +++ b/kube/rust/README.md @@ -0,0 +1,89 @@ +## GKE RL Scaler (Rust) + +### Overview +This service exposes a minimal HTTP API to inspect Kubernetes workloads and scale Deployments. It connects to the current kube context (same as `kubectl`) and writes pod logs locally grouped by the pod's `app` label. + +### Prerequisites +- Rust toolchain installed (`cargo`) +- `kubectl` configured to point at your cluster +- Cluster and pools can be created with the Python helper: `./cluster_setup/gke_setup.py --config example create-cluster` and `create-rollout-pool` +- Update the `ROLLOUT_MGR` value `http://:5000` in `kube/deployment/kube-rollout-manifest.yaml` to your rollout manager address. +- Update the `-test:v1` and `-rollout:v1` in `kube/deployment/kube-test-manifest.yaml` and `kube/deployment/kube-rollout-manifest.yaml` + +### Build +```bash +cd rust +cargo build --release +``` + +### Run +You can run with a short config name (resolved under `config/.toml`) or an absolute path. + +```bash +# From repo root +./rust/target/release/gke-rl --config example + +# Or using cargo while developing +cargo run -- --config example + +# Or pass an explicit path +cargo run -- --config /Users/you/path/to/config/example.toml +``` + +Environment overrides (optional): +- `BIND_ADDR` (default `0.0.0.0:5000`) +- `LOGS_DIR` (default `pod-logs`) +- `NAMESPACE` (default `default`) + +Example with overrides: +```bash +BIND_ADDR=0.0.0.0:5001 LOGS_DIR=logs/example NAMESPACE=default cargo run -- --config example +``` + +If you specify `[deployments]` in the `example.toml`, it will automatically apply manifest before start the API service. + +### API +Base URL: `http://:` (default `http://0.0.0.0:5000`) + +- Health +```bash +curl localhost:5000/health +``` + +- List Deployments (in configured namespace) +```bash +curl localhost:5000/deployments +``` + +- List Pods (in configured namespace) +```bash +curl localhost:5000/pods +``` + +- Scale a Deployment +```bash +curl -X POST localhost:5000/scale \ + -H 'content-type: application/json' \ + -d '{"app":"rollout-app","replicas":2}' +``` +Replace `rollout-app` with the Deployment name you want to scale (e.g., the deployment defined in `deployment/kube-rollout-manifest.yaml`). + +### Logs +The service tails logs from running pods in the configured namespace and stores them under `LOGS_DIR//.log`. + +Notes: +- `` is derived from the pod label `app`. If missing, logs are written under `unknown/`. +- The directory structure is created automatically. + +### Configuration +When `--config ` is provided, the service reads TOML and uses the following fields: +- `[general].namespace` +- `[general].bind_addr` +- `[general].logs_dir` + +Example: `config/example.toml` includes these keys under `[general]`. + +### Troubleshooting +- Ensure `kubectl get pods` works in the same shell; the scaler uses the same kube context. +- If you see RBAC errors, make sure your user/service account has permissions to list pods/deployments and patch deployments in the target namespace. + diff --git a/kube/rust/src/main.rs b/kube/rust/src/main.rs new file mode 100644 index 0000000..71b186e --- /dev/null +++ b/kube/rust/src/main.rs @@ -0,0 +1,77 @@ +// Copyright 2025 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +use anyhow::Result; +use axum::{routing::{get, post}, Router}; +use kube::Client; +use std::{net::SocketAddr, sync::Arc}; +use tower_http::trace::TraceLayer; +use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt}; + +mod server; +mod watcher; +mod settings; +mod scaler; + +#[tokio::main] +async fn main() -> Result<()> { + tracing_subscriber::registry() + .with(tracing_subscriber::EnvFilter::new( + std::env::var("RUST_LOG").unwrap_or_else(|_| "info,tower_http=info".into()), + )) + .with(tracing_subscriber::fmt::layer()) + .init(); + + let cfg = { + // Basic, lightweight CLI arg parse for optional --config + let mut args = std::env::args().skip(1); + let mut config_arg: Option = None; + while let Some(a) = args.next() { + if a == "--config" { + if let Some(v) = args.next() { config_arg = Some(v); } + } + } + settings::AppSettings::from_config_arg(config_arg)? + }; + let client = Client::try_default().await?; + let shared = Arc::new(server::AppState::new(client.clone(), cfg.clone())); + + // spawn watcher task + let state_for_watcher = shared.clone(); + tokio::spawn(async move { + if let Err(err) = watcher::run_pod_monitor(state_for_watcher).await { + tracing::error!(error = ?err, "pod monitor terminated with error"); + } + }); + + let app = Router::new() + .route("/health", get(|| async { "ok" })) + .route("/pods", get(server::list_pods)) + .route("/deployments", get(server::list_deployments)) + .route("/scale", post(server::scale_pool)) + .layer(TraceLayer::new_for_http()) + .with_state(shared); + + // Optionally apply deployment manifests on startup if configured + for path in &cfg.deployment_files { + tracing::info!(file=%path.display(), "applying deployment manifest from config"); + if let Err(e) = scaler::apply_yaml_manifest(client.clone(), &cfg.namespace, path).await { + tracing::error!(file=%path.display(), error=?e, "failed to apply deployment manifest"); + } + } + + let addr: SocketAddr = cfg.bind_addr.parse()?; + tracing::info!(%addr, "starting server"); + axum::serve(tokio::net::TcpListener::bind(addr).await?, app).await?; + Ok(()) +} diff --git a/kube/rust/src/scaler.rs b/kube/rust/src/scaler.rs new file mode 100644 index 0000000..c26cfe3 --- /dev/null +++ b/kube/rust/src/scaler.rs @@ -0,0 +1,95 @@ +// Copyright 2025 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +use anyhow::{anyhow, bail, Context, Result}; +use k8s_openapi::api::apps::v1::Deployment; +use kube::{ + api::{Patch, PatchParams}, + core::{DynamicObject, GroupVersionKind}, + discovery::{Discovery, Scope}, + Api, Client, +}; +use kube::ResourceExt; // for name_any +use serde::de::DeserializeOwned; +use std::path::Path; +use kube::error::Error as KubeError; + +pub async fn scale_deployment(client: Client, namespace: &str, deployment: &str, replicas: i32) -> Result<()> { + if replicas < 0 { bail!("replicas must be >= 0"); } + let api: Api = Api::namespaced(client, namespace); + let patch = serde_json::json!({ + "spec": {"replicas": replicas} + }); + tracing::debug!(deployment=%deployment, namespace=%namespace, replicas=replicas, "sending merge patch to scale deployment"); + match api.patch( + deployment, + &PatchParams::default(), + &Patch::Merge(patch), + ).await { + Ok(_) => { + tracing::info!(deployment=%deployment, namespace=%namespace, replicas=replicas, "deployment scaled successfully"); + Ok(()) + } + Err(KubeError::Api(err)) => { + // Surface the Kubernetes API error details directly + Err(anyhow!("Kubernetes API error ({}): {} [reason={}, code={}]", deployment, err.message, err.reason, err.code)) + } + Err(e) => { + Err(e).with_context(|| format!("scaling deployment {} in {}", deployment, namespace)) + } + } +} + +pub async fn apply_yaml_manifest(client: Client, default_namespace: &str, path: &Path) -> Result<()> { + let yaml = std::fs::read_to_string(path) + .with_context(|| format!("reading manifest at {}", path.display()))?; + let discovery = Discovery::new(client.clone()).run().await?; + let ssapply = PatchParams::apply("gke-rl-scaler").force(); + for doc in yaml_documents(&yaml)? { + let obj: DynamicObject = serde_yaml::from_value(doc)?; + let ns = obj + .metadata + .namespace + .as_deref() + .unwrap_or(default_namespace); + let gvk = if let Some(tm) = &obj.types { + GroupVersionKind::try_from(tm)? + } else { + bail!("cannot apply object without valid TypeMeta: {:?}", obj); + }; + let name = obj.name_any(); + if let Some((ar, caps)) = discovery.resolve_gvk(&gvk) { + let api: Api = if caps.scope == Scope::Namespaced { + Api::namespaced_with(client.clone(), ns, &ar) + } else { + Api::all_with(client.clone(), &ar) + }; + let data: serde_json::Value = serde_json::to_value(&obj)?; + tracing::info!(kind=%gvk.kind, name=%name, namespace=%ns, "applying manifest via SSA"); + api.patch(&name, &ssapply, &Patch::Apply(data)).await?; + } else { + tracing::warn!(?gvk, "cannot apply document for unknown GVK"); + } + } + Ok(()) +} + +fn yaml_documents(yaml: &str) -> Result> { + let mut docs = Vec::new(); + for d in serde_yaml::Deserializer::from_str(yaml) { + let v = T::deserialize(d).context("parsing one YAML document")?; + docs.push(v); + } + Ok(docs) +} + diff --git a/kube/rust/src/server.rs b/kube/rust/src/server.rs new file mode 100644 index 0000000..a69a3b1 --- /dev/null +++ b/kube/rust/src/server.rs @@ -0,0 +1,116 @@ +// Copyright 2025 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +use std::sync::Arc; + +use anyhow::Result; +use axum::{extract::State, http::StatusCode, Json}; +use k8s_openapi::api::{apps::v1::Deployment, core::v1::Pod}; +use kube::{Api, Client}; +use serde::{Deserialize, Serialize}; + +use crate::{scaler::scale_deployment, settings::AppSettings}; + +#[derive(Clone)] +pub struct AppState { + pub client: Client, + pub cfg: AppSettings, +} + +impl AppState { + pub fn new(client: Client, cfg: AppSettings) -> Self { + Self { client, cfg } + } +} + +#[derive(Deserialize)] +pub struct ScaleRequest { pub app: String, pub replicas: i32 } + +#[derive(Serialize)] +pub struct ScaleResponse { pub deployment: String, pub replicas: i32 } + +pub async fn scale_pool(State(state): State>, Json(req): Json) -> Result<(StatusCode, Json), (StatusCode, String)> { + let deploy = req.app.clone(); + tracing::info!(deployment=%deploy, replicas=req.replicas, namespace=%state.cfg.namespace, "received scale request"); + match scale_deployment(state.client.clone(), &state.cfg.namespace, &deploy, req.replicas).await { + Ok(()) => { + tracing::info!(deployment=%deploy, replicas=req.replicas, "scale request succeeded"); + Ok((StatusCode::OK, Json(ScaleResponse { deployment: deploy, replicas: req.replicas }))) + } + Err(e) => { + tracing::error!(deployment=%deploy, replicas=req.replicas, error=?e, "scale request failed"); + Err((StatusCode::INTERNAL_SERVER_ERROR, e.to_string())) + } + } +} + +#[derive(Serialize)] +pub struct PodSummary { + pub name: String, + pub phase: Option, + pub node_name: Option, + pub host_ip: Option, + pub pod_ip: Option, +} + +#[derive(Serialize)] +pub struct PodsResponse { pub items: Vec } + +pub async fn list_pods(State(state): State>) -> Result, (StatusCode, String)> { + tracing::debug!(namespace=%state.cfg.namespace, "list pods request"); + let api: Api = Api::namespaced(state.client.clone(), &state.cfg.namespace); + let pods = api + .list(&Default::default()) + .await + .map_err(|e| { + tracing::error!(error=?e, "failed to list pods"); + (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()) + })?; + let items: Vec = pods + .into_iter() + .map(|p| PodSummary { + name: p.metadata.name.unwrap_or_default(), + phase: p.status.as_ref().and_then(|s| s.phase.clone()), + node_name: p.spec.as_ref().and_then(|s| s.node_name.clone()), + host_ip: p.status.as_ref().and_then(|s| s.host_ip.clone()), + pod_ip: p.status.as_ref().and_then(|s| s.pod_ip.clone()), + }) + .collect::>(); + tracing::info!(count=items.len(), "list pods succeeded"); + Ok(Json(PodsResponse { items })) +} + +#[derive(Serialize)] +pub struct DeploymentSummary { pub name: String, pub replicas: Option, pub available: Option } + +#[derive(Serialize)] +pub struct DeploymentsResponse { pub items: Vec } + +pub async fn list_deployments(State(state): State>) -> Result, (StatusCode, String)> { + tracing::debug!(namespace=%state.cfg.namespace, "list deployments request"); + let api: Api = Api::namespaced(state.client.clone(), &state.cfg.namespace); + let list = api.list(&Default::default()).await.map_err(|e| { + tracing::error!(error=?e, "failed to list deployments"); + (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()) + })?; + let items: Vec = list.into_iter().map(|d| { + DeploymentSummary { + name: d.metadata.name.unwrap_or_default(), + replicas: d.spec.as_ref().and_then(|s| s.replicas), + available: d.status.as_ref().and_then(|s| s.available_replicas), + } + }).collect::>(); + tracing::info!(count=items.len(), "list deployments succeeded"); + Ok(Json(DeploymentsResponse { items })) +} + diff --git a/kube/rust/src/settings.rs b/kube/rust/src/settings.rs new file mode 100644 index 0000000..e3e2634 --- /dev/null +++ b/kube/rust/src/settings.rs @@ -0,0 +1,137 @@ +// Copyright 2025 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +use anyhow::{Context, Result}; +use serde::Deserialize; +use std::path::{Path, PathBuf}; + +#[derive(Clone, Debug)] +pub struct AppSettings { + pub bind_addr: String, + pub logs_dir: PathBuf, + pub namespace: String, + pub scale_up_step: i32, + pub deployment_files: Vec, + pub default_pool: Option, + pub worker_pools: Vec, +} + +impl AppSettings { + pub fn from_env() -> Result { Self::from_config_arg(None) } + + pub fn from_config_arg(config_arg: Option) -> Result { + let file_cfg = if let Some(spec) = config_arg { + let path = resolve_config_path(&spec); + let bytes = std::fs::read_to_string(&path) + .with_context(|| format!("reading config at {}", path.display()))?; + let parsed: FileConfig = toml::from_str(&bytes).context("parsing TOML config")?; + Some(parsed) + } else { + None + }; + + // Defaults + let mut bind_addr = "0.0.0.0:5000".to_string(); + let mut logs_dir = PathBuf::from("pod-logs"); + let mut namespace = "default".to_string(); + let mut scale_up_step: i32 = 1; + let mut deployment_files: Vec = Vec::new(); + let mut default_pool: Option = None; + let mut worker_pools: Vec = Vec::new(); + + if let Some(cfg) = &file_cfg { + if let Some(g) = &cfg.general { + if let Some(v) = &g.bind_addr { bind_addr = v.clone(); } + if let Some(v) = &g.logs_dir { logs_dir = PathBuf::from(v); } + if let Some(v) = &g.namespace { namespace = v.clone(); } + } + if let Some(sc) = &cfg.scaler { + if let Some(v) = sc.scale_up_step { scale_up_step = v; } + } + if let Some(d) = &cfg.deployments { + if let Some(files) = &d.files { + for f in files { + deployment_files.push(resolve_config_relative(f)); + } + } + } + if let Some(dp) = &cfg.default_pool { default_pool = Some(dp.clone()); } + if let Some(wps) = &cfg.worker_pools { worker_pools = wps.clone(); } + } + + // Environment overrides + if let Ok(v) = std::env::var("BIND_ADDR") { bind_addr = v; } + if let Ok(v) = std::env::var("LOGS_DIR") { logs_dir = PathBuf::from(v); } + if let Ok(v) = std::env::var("NAMESPACE") { namespace = v; } + if let Ok(v) = std::env::var("SCALE_UP_STEP") { scale_up_step = v.parse().unwrap_or(scale_up_step); } + + Ok(Self { + bind_addr, + logs_dir, + namespace, + scale_up_step, + deployment_files, + default_pool, + worker_pools, + }) + } +} + +#[derive(Debug, Deserialize)] +struct FileConfig { + #[serde(default)] + pub general: Option, + #[serde(default)] + pub scaler: Option, + #[serde(default)] + pub deployments: Option, + #[serde(default)] + pub default_pool: Option, + #[serde(default)] + pub worker_pools: Option>, +} + +#[derive(Debug, Deserialize)] +struct General { pub project: Option, pub region: Option, pub zone: Option, pub cluster_name: Option, pub vpc_net: Option, pub vpc_subnet: Option, pub namespace: Option, pub logs_dir: Option, pub bind_addr: Option } + +#[derive(Debug, Deserialize)] +struct Scaler { pub scale_up_step: Option } + +#[derive(Debug, Deserialize)] +struct DeploymentsSection { pub files: Option> } + +#[derive(Debug, Deserialize, Clone)] +#[allow(dead_code)] +pub struct DefaultPoolToml { pub machine_type: Option, pub min_nodes: Option, pub max_nodes: Option } + +#[derive(Debug, Deserialize, Clone)] +#[allow(dead_code)] +pub struct WorkerPoolToml { pub machine_type: Option, pub provision: Option, pub accelerator: Option, pub min_nodes: Option, pub max_nodes: Option } + +fn repo_root() -> PathBuf { + // CARGO_MANIFEST_DIR points to rust/; repo root is its parent + Path::new(env!("CARGO_MANIFEST_DIR")).parent().unwrap().to_path_buf() +} + +fn resolve_config_path(spec: &str) -> PathBuf { + let candidate = PathBuf::from(spec); + if candidate.exists() { return candidate; } + let by_name = repo_root().join("config").join(format!("{}.toml", spec)); + by_name +} + +fn resolve_config_relative(rel: &str) -> PathBuf { + let from_root = repo_root().join(rel); + from_root +} + diff --git a/kube/rust/src/watcher.rs b/kube/rust/src/watcher.rs new file mode 100644 index 0000000..079ac86 --- /dev/null +++ b/kube/rust/src/watcher.rs @@ -0,0 +1,145 @@ +// Copyright 2025 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +use std::{fs::{self, File}, io::Write, path::PathBuf, sync::Arc}; + +use anyhow::Result; +use futures::{StreamExt, TryStreamExt, AsyncBufReadExt}; +use k8s_openapi::api::core::v1::Pod; +use kube::{api::{ListParams, LogParams}, Api, Client, runtime::watcher}; + +use crate::server::AppState; + +pub async fn run_pod_monitor(state: Arc) -> Result<()> { + // Recreate top-level logs dir on startup to ensure a clean directory + if state.cfg.logs_dir.exists() { + tracing::warn!(dir=%state.cfg.logs_dir.display(), "removing existing logs directory"); + match fs::remove_dir_all(&state.cfg.logs_dir) { + Ok(_) => {} + Err(e) => tracing::warn!(dir=%state.cfg.logs_dir.display(), error=?e, "failed to remove logs directory, continuing"), + } + } + fs::create_dir_all(&state.cfg.logs_dir)?; + tracing::info!(dir=%state.cfg.logs_dir.display(), namespace=%state.cfg.namespace, "starting pod monitor and ensuring logs dir"); + let namespace = state.cfg.namespace.clone(); + let client = state.client.clone(); + let pod_api: Api = Api::namespaced(client.clone(), &namespace); + + // initial sweep: fetch current pods and ensure log files exist + let pods = pod_api.list(&ListParams::default()).await?; + for p in pods { + if let Some(name) = p.metadata.name { + let dir = pod_logs_dir(&state.cfg.logs_dir, p.metadata.labels.as_ref()); + tracing::debug!(pod=%name, dir=%dir.display(), "ensuring initial log file"); + ensure_log_file(&dir, &name)?; + } + } + + // watcher for pod events + let mut w = watcher(pod_api.clone(), watcher::Config::default()).boxed(); + while let Some(ev) = w.try_next().await? { + use kube::runtime::watcher::Event; + match ev { + Event::Applied(p) => { + if let Some(name) = p.metadata.name.clone() { + if is_running(&p) { + let client = client.clone(); + let namespace = namespace.clone(); + let logs_dir = pod_logs_dir(&state.cfg.logs_dir, p.metadata.labels.as_ref()); + tracing::info!(pod=%name, dir=%logs_dir.display(), "starting log stream for running pod"); + tokio::spawn(async move { + if let Err(e) = stream_pod_logs(client, &namespace, &name, logs_dir.clone()).await { + tracing::warn!(pod=%name, dir=%logs_dir.display(), error=?e, "log stream ended with error"); + } + }); + } + } + } + Event::Deleted(p) => { + if let Some(name) = p.metadata.name { tracing::info!(pod=%name, "pod deleted"); } + } + Event::Restarted(ps) => { + for p in ps { + if let Some(name) = p.metadata.name.clone() { + if is_running(&p) { + let client = client.clone(); + let namespace = namespace.clone(); + let logs_dir = pod_logs_dir(&state.cfg.logs_dir, p.metadata.labels.as_ref()); + tracing::info!(pod=%name, dir=%logs_dir.display(), "restarting log stream after watcher restart"); + tokio::spawn(async move { + if let Err(e) = stream_pod_logs(client, &namespace, &name, logs_dir.clone()).await { + tracing::warn!(pod=%name, dir=%logs_dir.display(), error=?e, "log stream ended with error"); + } + }); + } + } + } + } + } + } + + Ok(()) +} + +fn is_running(p: &Pod) -> bool { + p.status + .as_ref() + .and_then(|s| s.phase.as_deref()) + .map(|ph| ph == "Running") + .unwrap_or(false) +} + +fn ensure_log_file(dir: &PathBuf, pod: &str) -> Result { + fs::create_dir_all(dir)?; + let path = dir.join(format!("{}.log", pod)); + if !path.exists() { File::create(&path)?; } + Ok(path) +} + +async fn stream_pod_logs(client: Client, ns: &str, pod: &str, dir: PathBuf) -> Result<()> { + let api: Api = Api::namespaced(client, ns); + let mut lp = LogParams::default(); + lp.follow = true; + lp.tail_lines = Some(100); + tracing::debug!(pod=%pod, namespace=%ns, "requesting log stream from kube API"); + let reader = api.log_stream(pod, &lp).await?; + let log_path = ensure_log_file(&dir, pod)?; + tracing::debug!(pod=%pod, file=%log_path.display(), "opened log file for append"); + let mut file = File::options().append(true).open(&log_path)?; + let mut lines = reader.lines(); + while let Some(line) = lines.next().await { + match line { + Ok(l) => { + file.write_all(l.as_bytes())?; + file.write_all(b"\n")?; + file.flush()?; + // keep this at debug to avoid noise + tracing::debug!(pod=%pod, bytes=l.len(), "wrote log line"); + } + Err(e) => { + tracing::warn!(pod=%pod, error=?e, "error reading log line"); + break; + } + } + } + Ok(()) +} + +fn pod_logs_dir(base: &PathBuf, labels: Option<&std::collections::BTreeMap>) -> PathBuf { + let app = labels + .and_then(|m| m.get("app").cloned()) + .unwrap_or_else(|| "unknown".to_string()); + base.join(app) +} + + diff --git a/src/rlboost/rollout-manager/config.toml b/src/rlboost/rollout-manager/config.toml index 241477e..e860d84 100644 --- a/src/rlboost/rollout-manager/config.toml +++ b/src/rlboost/rollout-manager/config.toml @@ -7,11 +7,11 @@ mooncake_transfer_device_name = "" # Mooncake transfer protocol (e.g., rdma, tcp) mooncake_transfer_protocol = "tcp" +allowed_sender_ips = "192.168.0.0/16" +weight_sender_rpyc_base_port = 18861 # Weight sender RPyC endpoints # These will be automatically populated by the training script # but can be manually configured here -allowed_sender_ips = "192.168.0.0/16" -weight_sender_rpyc_base_port = 18861 weight_sender_rpyc_endpoints = [] # Number of mooncake groups per sender