Full parameter fine-tuning sampler sub-system

.github/workflows/build-and-push.yml

@@ -28,6 +28,13 @@ jobs:
       - name: Checkout
         uses: actions/checkout@v4
 
+      - name: Free Disk Space
+        run: |
+          sudo rm -rf /usr/share/dotnet
+          sudo rm -rf /opt/ghc
+          sudo rm -rf "/usr/local/share/boost"
+          sudo rm -rf "$AGENT_TOOLSDIRECTORY"
+
       - name: Set up Docker Buildx
         uses: docker/setup-buildx-action@v3
 

.github/workflows/build-pr.yml

@@ -24,6 +24,13 @@ jobs:
       - name: Checkout
         uses: actions/checkout@v4
 
+      - name: Free Disk Space
+        run: |
+          sudo rm -rf /usr/share/dotnet
+          sudo rm -rf /opt/ghc
+          sudo rm -rf "/usr/local/share/boost"
+          sudo rm -rf "$AGENT_TOOLSDIRECTORY"
+
       - name: Set up Docker Buildx
         uses: docker/setup-buildx-action@v3
 

.gitignore

@@ -7,6 +7,8 @@ env/
 **/venv/
 
 # Training Logs & Plots
+scratch/
+docs/scratch/
 *.txt
 *.png
 *.log

Makefile

@@ -169,8 +169,8 @@ REMOTE_HOST ?= <PLACE_HOLDER_FOR_REMOTE_HOST_ADDRESS>
 
 # Push local workspace changes to the remote VM
 push-vm:
-	rsync -avz --exclude '.git' --exclude '.venv' --exclude '__pycache__' --exclude '*.pyc' --exclude '.DS_Store' ./ $(REMOTE_HOST):~/open-rl
+	rsync -avz --exclude '.git' --exclude '.venv' --exclude '__pycache__' --exclude '*.pyc' --exclude '.DS_Store' --exclude 'scratch' ./ $(REMOTE_HOST):~/open-rl
 
 # Pull changes from the remote VM back to the local workspace
 pull-vm:
-	rsync -avz --exclude '.git' --exclude '.venv' --exclude '__pycache__' --exclude '*.pyc' --exclude '.DS_Store' $(REMOTE_HOST):~/open-rl/ ./
+	rsync -avz --exclude '.git' --exclude '.venv' --exclude '__pycache__' --exclude '*.pyc' --exclude '.DS_Store' --exclude 'scratch' $(REMOTE_HOST):~/open-rl/ ./

docs/fft/pod_placement_architecture.md

@@ -0,0 +1,132 @@
+# Open-RL Kubernetes Pod Placement Architecture Walkthrough
+
+This document provides an end-to-end architectural walkthrough of how Open-RL schedules, places, and virtualizes multi-tenant Supervised Fine-Tuning (SFT) and Reinforcement Learning (RL) workloads across a distributed Kubernetes infrastructure.
+
+---
+
+## High-Level Topology
+
+Open-RL decouples policy gradient computation (Trainers) from rollout generation (Samplers). In a multi-tenant environment (such as concurrent `job-a` and `job-b` experiments), Gateway orchestrates worker pods across dedicated physical GPU machines while ensuring strict boundary isolation.
+
+```mermaid
+graph TD
+    classDef gw fill:#1e293b,stroke:#3b82f6,stroke-width:2px,color:#fff;
+    classDef n1 fill:#0f172a,stroke:#10b981,stroke-width:2px,color:#fff;
+    classDef n2 fill:#0f172a,stroke:#8b5cf6,stroke-width:2px,color:#fff;
+    classDef podA fill:#064e3b,stroke:#34d399,stroke-width:1px,color:#fff;
+    classDef podB fill:#4c1d95,stroke:#a78bfa,stroke-width:1px,color:#fff;
+
+    Client["RL Client SDK (e.g. tiny_rl.py)"] -->|POST /api/v1/create_model| GW["Open-RL Gateway Service"]:::gw
+    Client -->|POST /api/v1/create_sampling_client| GW
+
+    subgraph Cluster ["GKE Regional Standard Cluster (open-rl-dra)"]
+        subgraph Node1 ["Physical Machine 1: dcbk (g2-standard-12)<br/>DRA Group: trainers"]
+            SA1["Snapshot Agent DaemonSet (tcp://:9753)"]:::n1
+            TrA["open-rl-trainer-job-a<br/>RAM: 16GiB / Claim: trainer-gpu-1"]:::podA
+            TrB["open-rl-trainer-job-b<br/>RAM: 16GiB / Claim: trainer-gpu-1"]:::podB
+            TrA <-->|CRIU Time-Slice| SA1
+            TrB <-->|CRIU Time-Slice| SA1
+        end
+
+        subgraph Node2 ["Physical Machine 2: hzp3 (g2-standard-12)<br/>DRA Group: samplers"]
+            SA2["Snapshot Agent DaemonSet (tcp://:9753)"]:::n2
+            SmA["open-rl-sampler-job-a<br/>RAM: 16GiB / Claim: sampler-gpu-1"]:::podA
+            SmB["open-rl-sampler-job-b<br/>RAM: 16GiB / Claim: sampler-gpu-1"]:::podB
+            SmA <-->|vLLM Sleep VRAM Yield| SA2
+            SmB <-->|vLLM Sleep VRAM Yield| SA2
+        end
+
+        NFS[("Managed GKE Filestore NFS (/mnt/shared)")]
+        TrA -->|Write Checkpoints| NFS
+        TrB -->|Write Checkpoints| NFS
+        NFS -->|In-Place Safetensor Reload| SmA
+        NFS -->|In-Place Safetensor Reload| SmB
+    end
+```
+
+---
+
+## 1. Decoupled Dynamic Pod Rendering
+
+When a client initiates a training loop, Gateway intercepts the session requests inside `k8s_worker_manager.py`. Rather than using static Kubernetes Deployments, Gateway dynamically deepcopies role-specific Pod YAML templates mounted from ConfigMaps:
+
+- **Trainer Pod Template**: Loaded from ConfigMap `open-rl-config` defined in `05-worker-pod-template.yaml`.
+- **Sampler Pod Template**: Loaded from ConfigMap `open-rl-sampler-worker-pod-template` defined in `09-sampler-pod-template.yaml`.
+
+Gateway injects unique runtime identifiers (`model_id`, `job_id`, and `OPEN_RL_WORKER_IMAGE` overrides) before submitting imperative `create_namespaced_pod` API calls.
+
+---
+
+## 2. Dynamic Resource Allocation (DRA) Claim Sharing
+
+Standard Kubernetes device plugins (`resources.limits: nvidia.com/gpu: 1`) enforce exclusive physical GPU locks: once Pod A lands on a node, `kube-scheduler` rejects Pod B until Pod A terminates.
+
+Open-RL bypasses this limitation using **Kubernetes Dynamic Resource Allocation (DRA)** exact allocation claims:
+
+```yaml
+# Inside 05-worker-pod-template.yaml (Trainer Spec)
+spec:
+  resourceClaims:
+  - name: trainer-gpu
+    resourceClaimName: open-rl-trainer-gpu-1
+```
+
+```yaml
+# Inside 09-sampler-pod-template.yaml (Sampler Spec)
+spec:
+  resourceClaims:
+  - name: sampler-gpu
+    resourceClaimName: open-rl-sampler-gpu-1
+```
+
+### How Claim Co-Scheduling Works:
+1. **First Tenant (`job-a`)**: When `open-rl-trainer-job-a` spawns, it binds singleton claim `open-rl-trainer-gpu-1` to Physical Machine 1 (`dcbk`).
+2. **Second Tenant (`job-b`)**: When `open-rl-trainer-job-b` spawns seconds later, `kube-scheduler` inspects its `resourceClaimName`. Because `open-rl-trainer-gpu-1` is already allocated on `dcbk`, **Kubernetes co-schedules Job B directly onto `dcbk` alongside Job A**!
+
+---
+
+## 3. Strict Role Segregation via `nodeSelector`
+
+> [!WARNING]
+> Co-locating PyTorch AdamW optimizers and vLLM KV caches on the same physical GPU causes immediate CUDA out-of-memory crashes (`CUDA error: out of memory`).
+
+To prevent cross-role contamination, nodes in the `gpu-dra` node pool are tagged with explicit role labels:
+- Machine 1 (`dcbk`): `group.timeslice.io/trainers="true"`
+- Machine 2 (`hzp3`): `group.timeslice.io/samplers="true"`
+
+Pod specs enforce strict landing boundaries:
+- **Trainers**: Enforce `nodeSelector: { group.timeslice.io/trainers: "true" }` and set `OPEN_RL_TIME_SLICE_GROUP=trainers`.
+- **Samplers**: Enforce `nodeSelector: { group.timeslice.io/samplers: "true" }` and set `OPEN_RL_TIMESLICE_GROUP=samplers`.
+
+---
+
+## 4. Host RAM Oversubscription Tuning
+
+A standard GKE `g2-standard-12` virtual machine provides **48 GiB** of system CPU RAM. 
+
+When scheduling multiple tenant pods onto a single machine, `kube-scheduler` calculates memory feasibility based on `resources.requests.memory`.
+
+| Component / Tenant Pod | Requested CPU Memory | Cumulative Allocated RAM | Node Feasibility on `g2-standard-12` (48 GiB Total) |
+| :--- | :---: | :---: | :---: |
+| **System Overhead** *(DaemonSets, CSI, Calico)* | ~4 GiB | 4 GiB | Schedulable (44 GiB Remaining) |
+| **Tenant 1 Trainer** (`job-a`) | 16 GiB | 20 GiB | Schedulable (28 GiB Remaining) |
+| **Tenant 2 Trainer** (`job-b`) | 16 GiB | 36 GiB | **Schedulable (12 GiB Remaining)** $\checkmark$ |
+
+> [!TIP]
+> Prior to tuning, templates requested `24Gi` and `32Gi` of CPU memory. Under those defaults, Tenant 1 allocated $24 + 4 = 28\text{ GiB}$, leaving only $20\text{ GiB}$ remaining. When Tenant 2 requested `24Gi`, Kubernetes rejected the pod with `FailedScheduling: Insufficient memory`. Lowering requests to `16Gi` unlocked true multi-tenant concurrency.
+
+---
+
+## 5. Node-Local Time-Slicing Virtualization
+
+Once co-scheduled onto the same physical GPU, workloads are virtualized in-flight by the node-local DaemonSet defined in `07-snapshot-agent-daemonset.yaml`.
+
+### A. Trainer Virtualization (CRIU Process Swapping)
+On the Trainer Node (`dcbk`), the Snapshot Agent intercepts PyTorch CUDA allocations over `tcp://status.hostIP:9753`. When Job A finishes its microbatch gradient calculation:
+1. Snapshot Agent freezes Job A's Linux process via CRIU (`checkpointed pid 34715 in 1.99s`).
+2. Snapshot Agent restores Job B's memory state into VRAM (`restored pid 34716 in 0.39s`).
+
+### B. Sampler Virtualization (vLLM Cooperative Sleep)
+On the Sampler Node (`hzp3`), vLLM inference engines time-slice cooperatively inside `vllm_sampler.py`:
+1. **Sleep Preemption**: Upon completing a sampling batch, vLLM invokes `await engine.sleep(level=2)`, instantly discarding physical GPU memory pages (`freed 19.45 GiB`).
+2. **NFS Weight Synchronization**: When Trainer A writes new SFT weights to `/mnt/shared`, Sampler A detects the modification, wakes up physical VRAM (`0.04s`), and reloads the checkpoint safetensors in-place directly from NFS page cache (`took 1.19 seconds`)!

docs/setup/gke-fft-timeslice.md

@@ -17,17 +17,10 @@ three ideas that build on each other:
 There are three separate responsibilities in this PR.
 
 First, DRA is used only for GPU allocation and placement. The deployment creates
-one `ResourceClaim` named `open-rl-trainer-gpu-1`. Trainer worker pods reference
-that same claim. Kubernetes allocates one matching NVIDIA GPU to
-the claim and schedules those pods onto a node where that device is available.
-DRA does not serialize CUDA execution, perform checkpoint/restore, or decide
-which process runs next.
+two static `ResourceClaims` named `open-rl-trainer-gpu-1` and `open-rl-sampler-gpu-1`. Trainer worker pods reference `open-rl-trainer-gpu-1`, while Sampler worker pods reference `open-rl-sampler-gpu-1`. Kubernetes allocates one matching NVIDIA GPU to each claim and schedules those pods onto separate physical nodes where those devices are available.
 
 Second, the Kubernetes worker manager is the deployment launcher. It runs inside
-the gateway process today. When the gateway receives `create_model` or
-`create_model_from_state` in FFT mode, it
-creates a pod for that `model_id` from the trainer worker pod template, stamps
-the pod name, labels, job-id env var, and `--model-id`, then enqueues the request
+the gateway process today. When the gateway receives `create_model` in FFT mode, it creates a trainer pod for that `model_id`. When it receives `create_sampling_client`, it creates a dedicated vLLM sampler pod for that `model_id` from the sampler worker pod template. It enqueues the request
 on the model-specific Redis queue. It is idempotent: if the trainer worker pod
 for a model is already running, it reuses it.
 
@@ -98,14 +91,13 @@ scheduler.
 
 ## 1. DRA pins the GPU allocation
 
-`k8s/deploy/distributed-fft-timeslice/06-gpu-resourceclaim.yaml` creates a
-single namespace-scoped `ResourceClaim`:
+`k8s/deploy/distributed-fft-timeslice/06-gpu-resourceclaim.yaml` and `08-sampler-resourceclaim.yaml` create dedicated namespace-scoped `ResourceClaims` for Trainers and Samplers:
 
 ```yaml
 apiVersion: resource.k8s.io/v1
 kind: ResourceClaim
 metadata:
-  name: open-rl-trainer-gpu-1
+  name: open-rl-trainer-gpu-1 # (and open-rl-sampler-gpu-1)
 spec:
   devices:
     requests:
@@ -114,22 +106,15 @@ spec:
         deviceClassName: gpu.nvidia.com
 ```
 
-Trainer worker pods reference that same claim:
+Trainer worker pods reference `open-rl-trainer-gpu-1`, while Sampler worker pods reference `open-rl-sampler-gpu-1`:
 
 ```yaml
-resources:
-  claims:
-  - name: trainer-gpu
 resourceClaims:
-- name: trainer-gpu
-  resourceClaimName: open-rl-trainer-gpu-1
+- name: trainer-gpu # (or sampler-gpu)
+  resourceClaimName: open-rl-trainer-gpu-1 # (or open-rl-sampler-gpu-1)
 ```
 
-Because this is a shared `ResourceClaim`, Kubernetes allocates a single matching
-device to the claim and schedules all referencing pods where that allocated
-device is accessible. Do not use a `ResourceClaimTemplate` for this PR's pinning
-behavior: templates generate per-pod claims, which is the pattern for separate
-devices.
+Because these are shared `ResourceClaims`, Kubernetes allocates a single matching device to each claim and schedules referencing pods onto the dedicated nodes where those claims reside (`group.timeslice.io/trainers` vs `samplers`).
 
 DRA is the allocation and placement layer. It does not serialize CUDA execution
 by itself. This PR is intentionally an oversubscription model: multiple trainer
@@ -230,8 +215,8 @@ gcloud container node-pools create gpu-dra \
   --cluster "${CLUSTER}" --zone "${ZONE}" \
   --machine-type g2-standard-24 \
   --accelerator "type=nvidia-l4,count=1,gpu-driver-version=disabled" \
-  --node-labels="group.timeslice.io/trainers=true,gke-no-default-nvidia-gpu-device-plugin=true,nvidia.com/gpu.present=true" \
-  --num-nodes 1
+  --node-labels="group.timeslice.io/trainers=true,group.timeslice.io/samplers=true,gke-no-default-nvidia-gpu-device-plugin=true,nvidia.com/gpu.present=true" \
+  --num-nodes 2
 ```
 
 Install the GPU driver manually. Use the `latest` installer so the
@@ -287,17 +272,15 @@ The deployment assumes one base model per rollout: set `BASE_MODEL` in
 `kustomization.yaml`, and the gateway uses that value for `get_info` and
 `create_model` requests that do not explicitly pass a base model.
 
-There is no static trainer worker. Every `create_model` call makes the gateway
-create a trainer worker pod named `open-rl-trainer-<model-id>`, labeled:
+There are no static worker deployments. Every `create_model` call makes the gateway create a trainer pod named `open-rl-trainer-<model-id>`, and every `create_sampling_client` call makes the gateway create a dedicated vLLM sampler pod named `open-rl-sampler-<model-id>`. Both are labeled:
 
 ```yaml
-snapshot-agent: "true"          # OpenRL/future coordinator marker
-timeslice.io/group: trainers    # snapshot-agent group
+snapshot-agent: "true"          # OpenRL coordinator marker
+timeslice.io/group: trainers    # (or samplers for vLLM rollout workers)
 timeslice.io/job-id: <model-id> # per-worker identity
 ```
 
-The gateway's `open-rl-sa` service account has a Role allowing pod CRUD in the
-workload namespace (`03-rbac.yaml`).
+The gateway's `open-rl-sa` service account has a Role allowing pod CRUD in the workload namespace (`03-rbac.yaml`). When weight updates occur during FFT training, Trainers write checkpoints to NFS `/mnt/shared`, and Samplers dynamically reload those checkpoint safetensors in-place in ~1.1 seconds while yielding GPU VRAM via cooperative sleep.
 
 ## Setup 3: Run training on the cluster