feat: build the vLLM engine settings in one place

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: |-

src/server/api_server.py

@@ -575,7 +575,7 @@ async def launch_worker_and_enqueue(command: Command) -> str:
     await asyncio.to_thread(worker_manager.ensure, command.model_id, "trainer")
   except Exception as exc:
     traceback.print_exc()
-    await store.set_future(request_id, {"type": "RequestFailedResponse", "error_message": str(exc)})
+    await store.set_result(request_id, {"type": "RequestFailedResponse", "error_message": str(exc)})
     return request_id
   return await enqueue(command)
 
@@ -985,7 +985,7 @@ async def retrieve_future(req: RetrieveFutureRequest, accept: str = Header(defau
   pending, failed, and every other result stay JSON.
   """
   request_id = req.request_id
-  result = await store.get_future(request_id, timeout=60.0)
+  result = await store.get_result(request_id, timeout=60.0)
   if result is None:
     return JSONResponse(status_code=400, content={"type": "RequestFailedResponse", "error_message": "Future not found"})
   if isinstance(result, dict) and result.get("type") == "RequestFailedResponse":

src/server/sampling_store.py

@@ -7,7 +7,7 @@
 import redis.asyncio as redis
 
 RESULT_TTL_SECONDS = 300
-# An accepted request may never be committed, so its idempotency key expires instead of staying until Redis is full.
+# An accepted request may never get a result, so its idempotency key expires instead of staying until Redis is full.
 PENDING_TTL_SECONDS = 3 * 24 * 3600
 
 
@@ -49,8 +49,8 @@ async def pending(self, pool: str, limit: int | None = 1000) -> list[dict[str, A
     entries = await self.redis.xrange(stream_key(pool), count=limit)
     return [{**json.loads(fields["record"]), "entry_id": entry_id} for entry_id, fields in entries]
 
-  async def commit(self, record: dict[str, Any], result: dict[str, Any]) -> bool:
-    """Save the first final result and drop the pending entry. False if it was already committed."""
+  async def set_result(self, record: dict[str, Any], result: dict[str, Any]) -> bool:
+    """Save the first final result and drop the pending entry. False if it already has one."""
     stream, future, entry_id = stream_key(record["pool"]), future_key(record["request_id"]), record["entry_id"]
 
     async def save(pipe) -> bool:
@@ -64,7 +64,7 @@ async def save(pipe) -> bool:
         pipe.expire(idempotency_key(record["pool"], key), RESULT_TTL_SECONDS)
       return True
 
-    # Every commit writes the result key, so watching it lets only one through.
+    # Every set_result writes the result key, so watching it lets only one through.
     return await self.redis.transaction(save, future, value_from_callable=True)
 
   async def pinned_revisions(self, pool: str) -> set[str]:

src/server/store.py

@@ -42,12 +42,12 @@ async def get_sampling_requests_for_model(self, model_id: str) -> list[dict[str,
     pass
 
   @abstractmethod
-  async def set_future(self, req_id: str, result: dict[str, Any]) -> None:
+  async def set_result(self, req_id: str, result: dict[str, Any]) -> None:
     """Resolve a future by its request ID."""
     pass
 
   @abstractmethod
-  async def get_future(self, req_id: str, timeout: float) -> dict[str, Any] | None:
+  async def get_result(self, req_id: str, timeout: float) -> dict[str, Any] | None:
     """Block until the future resolves or the timeout is reached."""
     pass
 
@@ -129,14 +129,14 @@ async def get_sampling_requests_for_model(self, model_id: str) -> list[dict[str,
       batch.append(queue.get_nowait())
     return batch
 
-  async def set_future(self, req_id: str, result: dict[str, Any]) -> None:
+  async def set_result(self, req_id: str, result: dict[str, Any]) -> None:
     if result.get("status") == "pending":
       return
     async with self.futures_cv:
       self.futures_store[req_id] = result
       self.futures_cv.notify_all()
 
-  async def get_future(self, req_id: str, timeout: float) -> dict[str, Any] | None:
+  async def get_result(self, req_id: str, timeout: float) -> dict[str, Any] | None:
     async with self.futures_cv:
       try:
         await asyncio.wait_for(self.futures_cv.wait_for(lambda: req_id in self.futures_store), timeout=timeout)
@@ -259,7 +259,7 @@ async def get_sampling_requests_for_model(self, model_id: str) -> list[dict[str,
 
     return batch
 
-  async def set_future(self, req_id: str, result: dict[str, Any]) -> None:
+  async def set_result(self, req_id: str, result: dict[str, Any]) -> None:
     if result.get("status") == "pending":
       return
     key = f"open_rl:future:{req_id}"
@@ -270,7 +270,7 @@ async def set_future(self, req_id: str, result: dict[str, Any]) -> None:
       pipeline.expire(key, 300)
       await pipeline.execute()
 
-  async def get_future(self, req_id: str, timeout: float) -> dict[str, Any] | None:
+  async def get_result(self, req_id: str, timeout: float) -> dict[str, Any] | None:
     key = f"open_rl:future:{req_id}"
     deadline = time.monotonic() + timeout
     while True:

src/server/training_requests_processor.py

@@ -201,7 +201,7 @@ async def fetch_requests(self) -> list[dict[str, Any]]:
   async def publish_result(self, request_id: str | None, result: dict[str, Any]) -> None:
     # Only rank 0 writes to the store; the other ranks computed the same result.
     if request_id is not None and is_primary():
-      await self.store.set_future(request_id, result)
+      await self.store.set_result(request_id, result)
 
   async def run_once(self) -> None:
     batch = await self.next_batch()