sampling: let FFT samplers time-slice with their trainer

Makefile

@@ -19,7 +19,7 @@ HOST           ?= 127.0.0.1
 PORT           ?= 9003
 # The fully qualified base URL used by local CLI tools and clients
 BASE_URL       ?= http://$(HOST):$(PORT)
-UNIT_TESTS ?= tests.test_session_lifecycle tests.test_proto_codec tests.test_forward_only tests.test_fft_batch_failure tests.test_api_server_paths tests.test_accel_timeslicer tests.test_trainer_optimizer_correctness tests.test_worker_manager tests.test_scheduler_worker_manager tests.test_estimator tests.test_redis_store tests.test_cluster_eval_script tests.test_delta_weight_sync tests.test_diffing_backends tests.test_sampler_weight_rotation tests.test_commands tests.test_weight_sync_config tests.test_exclusive_trainer tests.test_delete_model tests.test_automodel_worker tests.test_worker_thread_device tests.test_trainer_data_parallel tests.test_sampler_batch_limits tests.test_openrl_status tests.test_skill_docs tests.test_trainer_ready tests.test_sampling_store tests.test_lora_snapshots tests.test_pool_accept
+UNIT_TESTS ?= tests.test_session_lifecycle tests.test_proto_codec tests.test_forward_only tests.test_fft_batch_failure tests.test_api_server_paths tests.test_accel_timeslicer tests.test_trainer_optimizer_correctness tests.test_worker_manager tests.test_scheduler_worker_manager tests.test_sampler_launcher tests.test_sampler_gpu_leases tests.test_estimator tests.test_redis_store tests.test_cluster_eval_script tests.test_delta_weight_sync tests.test_diffing_backends tests.test_sampler_weight_rotation tests.test_commands tests.test_weight_sync_config tests.test_exclusive_trainer tests.test_delete_model tests.test_automodel_worker tests.test_worker_thread_device tests.test_trainer_data_parallel tests.test_openrl_status tests.test_skill_docs tests.test_trainer_ready tests.test_sampling_store tests.test_lora_snapshots tests.test_pool_accept tests.test_sampling_dispatcher
 # Only forward BASE_URL to e2e when the user supplied it. The Makefile default
 # is for local CLI usage; e2e should start its own backend by default.
 TRAINING_TEST_BASE_URL ?= $(if $(filter environment command line,$(origin BASE_URL)),$(BASE_URL),)
@@ -114,7 +114,7 @@ test:
 
 .PHONY: test-weight-transfer
 test-weight-transfer:
-	uv run --frozen --extra vllm --extra gpu --extra cluster python -m unittest tests.test_delta_weight_transfer_engine tests.test_delta_weight_sync tests.test_weight_sync_config tests.test_sampler_patch tests.test_weight_transfer_gpu
+	uv run --frozen --extra vllm --extra gpu --extra cluster python -m unittest tests.test_delta_weight_transfer_engine tests.test_delta_weight_sync tests.test_weight_sync_config tests.test_weight_transfer_gpu
 
 lint:
 	uv run --extra dev ruff check .

docs/configuration.md

@@ -48,7 +48,7 @@ REDIS_URL=redis://127.0.0.1:6379/0 VLLM_ARCHITECTURE_OVERRIDE=Gemma4ForCausalLM
 | Env var | Default | What it does |
 | --- | --- | --- |
 | `BASE_MODEL` | unset | Hugging Face model id loaded by the trainer and, when using vLLM, by the sampler. |
-| `SAMPLING_BACKEND` | `torch` in one process, `vllm` when `REDIS_URL` is set | Sampling backend selector. `torch` samples in the training process. `vllm` queues sampling requests for a vLLM sampler worker and needs `REDIS_URL`. |
+| `SAMPLING_BACKEND` | `torch` in one process, `vllm` when `REDIS_URL` is set | Sampling backend selector. `torch` samples in the training process. `vllm` samples through the sampling pools described below and needs `REDIS_URL`. |
 | `REDIS_URL` | unset | Enables distributed mode by switching the request store to Redis. Leave unset for a single-machine run. |
 | `OPEN_RL_FUTURE_TTL_S` | `300` | How long resolved request results stay readable by `retrieve_future` after a worker resolves them. |
 
@@ -100,3 +100,25 @@ On Kubernetes the release bundles set these variables. The API server runs with
 each trainer and sampler, and the OpenRL scheduler starts their pods. The API
 server writes each worker's environment into the `Workload`'s pod template, so
 there are no worker deployments to configure by hand.
+
+## Sampling pools
+
+Sampling goes through sampling pools. A pool serves one FFT job, one exclusive
+LoRA job, or every shared LoRA job on one base model. Its samplers are plain
+`vllm serve` replicas, and a dispatcher sends each request to the replica the
+pool's router picks, then commits the result. Pending requests stay in a Redis
+stream until a result is committed, so the Redis in `k8s/deploy/base` keeps an
+append-only file on a persistent volume and never evicts keys.
+
+On Kubernetes each pool is `vllm serve` Workloads on exclusive GPUs plus one CPU
+Deployment, named after the pool, that runs the dispatcher, Envoy and the llm-d
+endpoint picker (configured by the `open-rl-sampler-router` ConfigMap in
+`k8s/deploy/base`). A replica is ready only while vLLM reports the weight version
+the dispatcher expects, and a NetworkPolicy lets only the pool's CPU pod reach it.
+On one machine the local worker manager starts one `vllm serve` on
+`SAMPLER_CUDA_VISIBLE_DEVICES` and a dispatcher that sends to it directly.
+
+| Env var | Default | What it does |
+| --- | --- | --- |
+| `OPEN_RL_DISPATCH_CONCURRENCY` | `256` | Requests one pool's dispatcher keeps in flight across its replicas. |
+| `OPEN_RL_DISPATCHER_IMAGE` | the worker image | Image for each pool's dispatcher container. The release sets it to the API-server image, which carries the dispatcher without the GPU stack and pulls in seconds. |

docs/setup/gke-fft-timeslice.md

@@ -182,6 +182,32 @@ in-memory queue and active/checkpointed state for workloads sharing the physical
 GPU. The worker pod labels provide the workload identity llm-d uses to discover
 the relevant pod and process set.
 
+An FFT job's pool sampler shares the
+trainer's GPU and takes turns with it through the same coordinator, with no
+proxy in the request path. Its pod runs
+`python -m server.sampler_launcher vllm serve ... --enable-sleep-mode` with
+`OPEN_RL_TIME_SLICED_SAMPLER=true`. The launcher registers the sampler's
+workload, starts vLLM inside one turn, puts it to sleep and releases, so vLLM
+never takes GPU memory while the trainer holds the device, then keeps the
+registration open for as long as vLLM runs. The container counts as started only
+then, and the pool's dispatcher, which reads the same variable from the pod,
+leases the replica for each turn through the slicer on its node
+(`status.hostIP` and `OPEN_RL_ACCEL_TIMESLICER_PORT`), wakes vLLM, loads any
+pending weights as a whole, sends requests to it directly rather than through
+the router, puts it to sleep and releases. A turn ends once the replica is idle,
+or after `OPEN_RL_SAMPLER_TURN_SECONDS` (60 by default) with requests still
+pending, so the trainer gets its turn. Requests still running then get ten
+seconds more. After that they are cancelled, aborted inside vLLM by its sleep,
+and sent again in a later turn. If the dispatcher dies holding a lease,
+the slicer checkpoints the replica before anyone else is granted the device. If
+that checkpoint, or the one at the end of a turn, fails, the slicer tells the
+launcher, which stops vLLM and exits, so the device frees for the trainer and
+kubelet restarts the replica. A
+pod deleted while vLLM is still starting is killed by its launcher a few seconds
+after the SIGTERM vLLM ignores at that point, so the turn is freed at once
+rather than at kubelet's SIGKILL. Exclusive jobs' pool samplers and LoRA pool
+samplers keep a GPU of their own and never talk to the slicer.
+
 ## Requirements
 
 - GKE Standard cluster on **1.35 or newer** ([DRA for GPUs](https://docs.cloud.google.com/kubernetes-engine/docs/how-to/set-up-dra)

k8s/deploy/base/api-server.yaml

@@ -71,6 +71,12 @@ spec:
             configMapKeyRef:
               name: open-rl-images
               key: worker
+        - name: OPEN_RL_DISPATCHER_IMAGE
+          valueFrom:
+            configMapKeyRef:
+              name: open-rl-images
+              key: dispatcher
+              optional: true
         resources:
           limits:
             memory: "2Gi"

k8s/deploy/base/images.yaml

@@ -1,3 +1,5 @@
 images:
   - path: data/worker
     kind: ConfigMap
+  - path: data/dispatcher
+    kind: ConfigMap

k8s/deploy/base/kustomization.yaml

@@ -7,6 +7,7 @@ namespace: openrl-system
 resources:
   - ../../../scheduler/deploy/base
   - rbac.yaml
+  - sampler-router.yaml
   - redis.yaml
   - api-server.yaml
 
@@ -20,6 +21,7 @@ configMapGenerator:
   - name: open-rl-images
     literals:
       - worker=ghcr.io/gke-labs/open-rl/server:latest
+      - dispatcher=ghcr.io/gke-labs/open-rl/api-server:latest
 
 # Dynamic worker images receive the same image transformations as containers.
 configurations:

k8s/deploy/base/rbac.yaml

@@ -1,5 +1,6 @@
 # The API server requests one Workload per worker process; the scheduler owns the
-# pods and claims from there, so this account needs only the Workload verbs.
+# pods and claims from there. It also manages each sampling pool's CPU pod and
+# lists sampler pods, because a sampler is ready only once its pod is Ready.
 apiVersion: rbac.authorization.k8s.io/v1
 kind: Role
 metadata:
@@ -8,6 +9,15 @@ rules:
 - apiGroups: ["openrl.io"]
   resources: ["workloads"]
   verbs: ["create", "get", "list", "watch", "delete"]
+- apiGroups: ["apps"]
+  resources: ["deployments"]
+  verbs: ["create", "get", "delete"]
+- apiGroups: ["networking.k8s.io"]
+  resources: ["networkpolicies"]
+  verbs: ["create", "delete"]
+- apiGroups: [""]
+  resources: ["pods"]
+  verbs: ["list"]
 ---
 apiVersion: rbac.authorization.k8s.io/v1
 kind: RoleBinding
@@ -20,3 +30,30 @@ roleRef:
   kind: Role
   name: open-rl-worker-launcher
   apiGroup: rbac.authorization.k8s.io
+---
+# A pool's dispatcher and EPP find its samplers by label, ready or not.
+apiVersion: v1
+kind: ServiceAccount
+metadata:
+  name: open-rl-sampler-dispatcher
+---
+apiVersion: rbac.authorization.k8s.io/v1
+kind: Role
+metadata:
+  name: open-rl-sampler-dispatcher
+rules:
+- apiGroups: [""]
+  resources: ["pods"]
+  verbs: ["get", "list", "watch"]
+---
+apiVersion: rbac.authorization.k8s.io/v1
+kind: RoleBinding
+metadata:
+  name: open-rl-sampler-dispatcher
+subjects:
+- kind: ServiceAccount
+  name: open-rl-sampler-dispatcher
+roleRef:
+  kind: Role
+  name: open-rl-sampler-dispatcher
+  apiGroup: rbac.authorization.k8s.io

k8s/deploy/base/sampler-router.yaml

@@ -0,0 +1,87 @@
+# The router of every sampling pool's CPU pod. Envoy asks the pod's llm-d EPP
+# (ext_proc on 9002) for a sampler and connects to the one the EPP names; it never
+# routes without the EPP. The response names the sampler that answered, so the
+# dispatcher can check its weight version before committing. The EPP keeps a request
+# on the replicas that hold its prefix (precise matching from each pod's KV events)
+# until their in-flight load would cost more time to first token than a cold replica,
+# llm-d's "sticky until saturated". Prefix weighting alone piles long decodes onto one
+# replica, and the slowest decode sets an RL step's time. No token producer is declared
+# because the dispatcher sends token IDs to vLLM's /inference/v1/generate, and the one
+# the EPP creates by default passes them through unchanged, so a vLLM renderer adds nothing.
+apiVersion: v1
+kind: ConfigMap
+metadata:
+  name: open-rl-sampler-router
+data:
+  envoy.yaml: |
+    static_resources:
+      listeners:
+        - name: inference
+          address: {socket_address: {address: 127.0.0.1, port_value: 8081}}
+          filter_chains:
+            - filters:
+                - name: envoy.filters.network.http_connection_manager
+                  typed_config:
+                    "@type": type.googleapis.com/envoy.extensions.filters.network.http_connection_manager.v3.HttpConnectionManager
+                    stat_prefix: inference
+                    # A long generation sends nothing until it finishes, so no idle timeout may cut it off.
+                    stream_idle_timeout: 0s
+                    route_config:
+                      virtual_hosts:
+                        - name: samplers
+                          domains: ["*"]
+                          routes:
+                            - match: {prefix: "/"}
+                              route: {cluster: samplers, timeout: 0s, idle_timeout: 0s}
+                              response_headers_to_add:
+                                - header: {key: x-openrl-upstream, value: "%UPSTREAM_HOST%"}
+                    http_filters:
+                      - name: envoy.filters.http.ext_proc
+                        typed_config:
+                          "@type": type.googleapis.com/envoy.extensions.filters.http.ext_proc.v3.ExternalProcessor
+                          grpc_service:
+                            envoy_grpc: {cluster_name: epp, authority: localhost:9002}
+                            timeout: 10s
+                          processing_mode:
+                            request_body_mode: FULL_DUPLEX_STREAMED
+                            response_body_mode: FULL_DUPLEX_STREAMED
+                            request_trailer_mode: SEND
+                            response_trailer_mode: SEND
+                          message_timeout: 1000s
+                      - name: envoy.filters.http.router
+                        typed_config: {"@type": type.googleapis.com/envoy.extensions.filters.http.router.v3.Router}
+      clusters:
+        - name: samplers
+          type: ORIGINAL_DST
+          lb_policy: CLUSTER_PROVIDED
+          connect_timeout: 10s
+          original_dst_lb_config: {use_http_header: true, http_header_name: x-gateway-destination-endpoint}
+        - name: epp
+          connect_timeout: 10s
+          transport_socket:
+            name: envoy.transport_sockets.tls
+            typed_config:
+              "@type": type.googleapis.com/envoy.extensions.transport_sockets.tls.v3.UpstreamTlsContext
+              common_tls_context: {validation_context: {}}
+          typed_extension_protocol_options:
+            envoy.extensions.upstreams.http.v3.HttpProtocolOptions:
+              "@type": type.googleapis.com/envoy.extensions.upstreams.http.v3.HttpProtocolOptions
+              explicit_http_config: {http2_protocol_options: {}}
+          load_assignment:
+            cluster_name: epp
+            endpoints: [{lb_endpoints: [{endpoint: {address: {socket_address: {address: 127.0.0.1, port_value: 9002}}}}]}]
+  epp.yaml: |
+    apiVersion: llm-d.ai/v1alpha1
+    kind: EndpointPickerConfig
+    plugins:
+      - type: endpoint-notification-source
+      - type: precise-prefix-cache-producer
+      - type: inflight-load-producer
+        parameters: {prefixMatchInfoProducerName: precise-prefix-cache-producer}
+      - type: prefix-cache-affinity-filter
+        parameters: {prefixMatchInfoProducerName: precise-prefix-cache-producer}
+      - type: token-load-scorer
+    dataLayer:
+      sources:
+        - pluginRef: endpoint-notification-source
+          extractors: [{pluginRef: precise-prefix-cache-producer}]

k8s/deploy/kind-dra/kustomization.yaml

@@ -143,6 +143,9 @@ patches:
                   - name: OPEN_RL_WORKER_IMAGE
                     valueFrom: null
                     value: localhost:5001/open-rl-server:kind-dev
+                  - name: OPEN_RL_DISPATCHER_IMAGE
+                    valueFrom: null
+                    value: localhost:5001/open-rl-api-server:kind-dev
                   # The API server also stamps the pull policy into each Workload.
                   - name: OPEN_RL_WORKER_IMAGE_PULL_POLICY
                     value: Always

scheduler/controller/Makefile

@@ -28,15 +28,15 @@ generate: $(CONTROLLER_GEN)
 # become the schema descriptions, so the prose lives in Go, not in the YAML.
 .PHONY: manifests
 manifests: $(CONTROLLER_GEN)
-	$(CONTROLLER_GEN) crd paths="./api/..." output:crd:dir=$(GEN_DIR)/crd
+	$(CONTROLLER_GEN) crd:generateEmbeddedObjectMeta=true paths="./api/..." output:crd:dir=$(GEN_DIR)/crd
 	cp $(GEN_DIR)/crd/openrl.io_workloads.yaml $(CRD_MANIFEST)
 	cp $(GEN_DIR)/crd/openrl.io_claimledgers.yaml $(CRD_GROUP_MANIFEST)
 
 # Fails when the checked-in CRD no longer matches the types. For CI, and for
 # the case where the generator cannot be run on the machine holding the diff.
 .PHONY: manifests-check
 manifests-check: $(CONTROLLER_GEN)
-	$(CONTROLLER_GEN) crd paths="./api/..." output:crd:dir=$(GEN_DIR)/crd
+	$(CONTROLLER_GEN) crd:generateEmbeddedObjectMeta=true paths="./api/..." output:crd:dir=$(GEN_DIR)/crd
 	diff -u $(CRD_MANIFEST) $(GEN_DIR)/crd/openrl.io_workloads.yaml
 	diff -u $(CRD_GROUP_MANIFEST) $(GEN_DIR)/crd/openrl.io_claimledgers.yaml
 

scheduler/deploy/base/00-workload-crd.yaml

@@ -156,6 +156,23 @@ spec:
                     description: |-
                       Standard object's metadata.
                       More info: https://git.k8s.io/community/contributors/devel/sig-architecture/api-conventions.md#metadata
+                    properties:
+                      annotations:
+                        additionalProperties:
+                          type: string
+                        type: object
+                      finalizers:
+                        items:
+                          type: string
+                        type: array
+                      labels:
+                        additionalProperties:
+                          type: string
+                        type: object
+                      name:
+                        type: string
+                      namespace:
+                        type: string
                     type: object
                   spec:
                     description: |-
@@ -7189,6 +7206,23 @@ spec:
                                         May contain labels and annotations that will be copied into the PVC
                                         when creating it. No other fields are allowed and will be rejected during
                                         validation.
+                                      properties:
+                                        annotations:
+                                          additionalProperties:
+                                            type: string
+                                          type: object
+                                        finalizers:
+                                          items:
+                                            type: string
+                                          type: array
+                                        labels:
+                                          additionalProperties:
+                                            type: string
+                                          type: object
+                                        name:
+                                          type: string
+                                        namespace:
+                                          type: string
                                       type: object
                                     spec:
                                       description: |-