diff --git a/apis/modeldeployments/definition.yaml b/apis/modeldeployments/definition.yaml index 382785c8f..8ec7bf50e 100644 --- a/apis/modeldeployments/definition.yaml +++ b/apis/modeldeployments/definition.yaml @@ -26,6 +26,16 @@ spec: spec: type: object required: [replicas, workers, nodeSelector] + # A disaggregated deployment (prefill set) declares everything + # explicitly — counts and routing are not defaulted. workers.count + # and prefill.workers.count carry no schema default, so these has() + # checks are meaningful; the composition function treats an omitted + # count as 1 for unified deployments, where these rules don't apply. + x-kubernetes-validations: + - rule: "!has(self.prefill) || (has(self.workers.count) && has(self.prefill.workers.count))" + message: "a disaggregated deployment (prefill set) must set both workers.count and prefill.workers.count" + - rule: "!has(self.prefill) || has(self.routing)" + message: "a disaggregated deployment (prefill set) must set routing" properties: replicas: type: integer @@ -144,9 +154,7 @@ spec: count: type: integer minimum: 1 - default: 1 - description: >- - Number of workers per replica. Defaults to 1. + description: Number of workers per replica. topology: type: object required: [tensor] @@ -291,6 +299,266 @@ spec: properties: name: type: string + prefill: + type: object + description: >- + Prefill role for disaggregated serving. When set, the deployment + is disaggregated: the top-level workers is the decode role and this + is the prefill role, each self-contained with its own topology, + template, and nodeSelector. Prefill and decode of a replica are + co-located on one InferenceCluster. Unset means unified serving. + required: [workers, nodeSelector] + properties: + workers: + type: object + required: [topology, template] + description: >- + Compute shape of one worker. Modelplane composes one + worker (or workers.count workers) per ModelReplica. + properties: + count: + type: integer + minimum: 1 + description: Number of workers per replica. + topology: + type: object + required: [tensor] + description: >- + Compute topology for one worker. The axes are + independent and compose multiplicatively: GPUs per + node = tensor, nodes per worker = pipeline. + properties: + tensor: + type: integer + minimum: 1 + description: >- + GPUs per node. Required. + pipeline: + type: integer + minimum: 1 + default: 1 + description: >- + Nodes per worker. Defaults to 1 (single-node). + Values greater than 1 enable multi-node serving + via LeaderWorkerSet. + template: + type: object + description: >- + Pod template for inference workers. A curated subset + of PodTemplateSpec. + properties: + metadata: + type: object + description: >- + Metadata applied to inference pods. Useful for + labels and annotations that control cluster-level + features like service mesh injection. + properties: + labels: + type: object + additionalProperties: + type: string + annotations: + type: object + additionalProperties: + type: string + spec: + type: object + required: [containers] + description: >- + Pod spec for inference workers. + properties: + containers: + type: array + minItems: 1 + maxItems: 1 + description: >- + Containers for the inference pod. v0.1 supports a + single container, which must be named "engine" (the + inference engine). Sidecar / multi-container support + is tracked separately. + x-kubernetes-validations: + - rule: "self.exists_one(c, c.name == 'engine')" + message: "the single container must be named 'engine'" + items: + type: object + required: [name, image] + properties: + name: + type: string + minLength: 1 + description: >- + Container name. The container named "engine" + is the inference engine. + image: + type: string + minLength: 1 + description: Container image. + command: + type: array + description: >- + Container entrypoint override. When set on the + engine container of a multi-node deployment, + it bypasses the built-in vLLM/Ray bootstrap and + runs on every gang pod — the command owns + cross-node coordination against the LWS_* + environment (LWS_WORKER_INDEX, + LWS_LEADER_ADDRESS, LWS_GROUP_SIZE). Use for + non-vLLM engines (e.g. SGLang). + items: + type: string + args: + type: array + description: >- + Container args. For the engine container, + these are passed through to the serving + engine. Includes the model identifier + (e.g. --model=...). + items: + type: string + env: + type: array + description: >- + Environment variables. Supports + valueFrom.secretKeyRef for secrets + like HF_TOKEN. + items: + type: object + required: [name] + properties: + name: + type: string + value: + type: string + valueFrom: + type: object + properties: + secretKeyRef: + type: object + required: [name, key] + properties: + name: + type: string + key: + type: string + optional: + type: boolean + configMapKeyRef: + type: object + required: [name, key] + properties: + name: + type: string + key: + type: string + optional: + type: boolean + imagePullSecrets: + type: array + description: >- + Image pull secrets for private registries + (NGC etc.). + items: + type: object + required: [name] + properties: + name: + type: string + nodeSelector: + type: object + description: >- + Node-level matching, a list of device requests mirroring a + DRA ResourceClaim. The scheduler matches each request against a + candidate pool's InferenceClass devices (surfaced on + InferenceCluster status.gpuPools) and pins the replica to a + pool that satisfies every request. claim: DRA requests also + become DeviceRequests in the ResourceClaim the serving pods + bind GPUs through. Required: GPUs bind only via DRA, so a + deployment must declare the devices its model needs. At least + one request must resolve to a claimable (claim: DRA) device; + the serving workload binds its GPUs through the resulting + ResourceClaim. Synthetic devices refine placement but are never + claimed, so a nodeSelector that matches only synthetic devices + leaves the workload nothing to claim - the scheduler treats + such a pool as ineligible and the deployment reports + InsufficientCapacity. + required: [devices] + properties: + devices: + type: array + description: >- + Device requests. A pool matches a request when it has a + device whose count covers the request and whose driver, + attributes, and capacity satisfy every selector. + minItems: 1 + maxItems: 16 + x-kubernetes-list-type: map + x-kubernetes-list-map-keys: [name] + items: + type: object + required: [name, selectors] + properties: + name: + type: string + description: >- + Name of this request. Mirrors a DRA DeviceRequest + name; carried through to the ResourceClaim. + minLength: 1 + maxLength: 63 + count: + type: integer + description: >- + How many matching devices a node must have. For a GPU + request this is the per-node GPU count (matches the + worker topology's GPUs per node). + default: 1 + minimum: 1 + maximum: 64 + selectors: + type: array + description: >- + Selectors a device must satisfy, all ANDed. Each is a + one-of; today only cel is supported. + minItems: 1 + maxItems: 8 + x-kubernetes-list-type: atomic + items: + type: object + # A selector must carry at least one selector kind + # (today only cel). Without this an empty {} selector + # would match every device, and since nodeSelector is + # the only path to a GPU that silently claims an + # arbitrary one. + minProperties: 1 + properties: + cel: + type: string + description: >- + A DRA CEL expression evaluated against one + device. Reads device.driver, + device.attributes[""]. (typed), + and device.capacity[""]. (a + Quantity), with quantity() and semver() helpers, + e.g. + device.capacity["gpu.nvidia.com"].memory.compareTo(quantity("141Gi")) >= 0. + minLength: 1 + maxLength: 10240 + routing: + type: object + description: >- + Endpoint-picker (EPP) for a disaggregated deployment. The + backend builds the EPP Deployment, Service, ConfigMap, and + RBAC from this block; the EPP sequences each request prefill + then decode. template is a curated PodSpec subset, same shape + and owner as the engine: its "epp" container's image and args + override the pinned llm-d EPP default. + properties: + template: + type: object + properties: + spec: + type: object + x-kubernetes-preserve-unknown-fields: true status: type: object properties: diff --git a/apis/modelreplicas/definition.yaml b/apis/modelreplicas/definition.yaml index 22fa6ba28..c230a36b3 100644 --- a/apis/modelreplicas/definition.yaml +++ b/apis/modelreplicas/definition.yaml @@ -118,7 +118,6 @@ spec: count: type: integer minimum: 1 - default: 1 topology: type: object required: [tensor] @@ -214,5 +213,185 @@ spec: properties: name: type: string + prefill: + type: object + description: >- + Prefill role placement for a disaggregated replica. Mirrors the + top-level decode fields (workers, nodePoolName, deviceRequests), + pinned to the prefill pool the scheduler chose. Absent for + unified replicas. When present, all three fields are required: + compose-model-deployment always populates them, and the llm-d + backend reads them unconditionally. + required: [workers, nodePoolName, deviceRequests] + properties: + workers: + type: object + required: [topology, template] + properties: + count: + type: integer + minimum: 1 + topology: + type: object + required: [tensor] + properties: + tensor: + type: integer + minimum: 1 + pipeline: + type: integer + minimum: 1 + default: 1 + template: + type: object + properties: + metadata: + type: object + properties: + labels: + type: object + additionalProperties: + type: string + annotations: + type: object + additionalProperties: + type: string + spec: + type: object + required: [containers] + properties: + containers: + type: array + minItems: 1 + maxItems: 1 + x-kubernetes-validations: + - rule: "self.exists_one(c, c.name == 'engine')" + message: "the single container must be named 'engine'" + items: + type: object + required: [name, image] + properties: + name: + type: string + minLength: 1 + image: + type: string + minLength: 1 + command: + type: array + items: + type: string + args: + type: array + items: + type: string + env: + type: array + items: + type: object + required: [name] + properties: + name: + type: string + value: + type: string + valueFrom: + type: object + properties: + secretKeyRef: + type: object + required: [name, key] + properties: + name: + type: string + key: + type: string + optional: + type: boolean + configMapKeyRef: + type: object + required: [name, key] + properties: + name: + type: string + key: + type: string + optional: + type: boolean + imagePullSecrets: + type: array + items: + type: object + required: [name] + properties: + name: + type: string + nodePoolName: + type: string + description: The prefill pool on this replica's cluster. + deviceRequests: + type: array + description: >- + Resolved DRA device requests for the matched pool. The parent + ModelDeployment's compose function joins the nodeSelector + requests with the matched InferenceClass devices and stamps the + claim: DRA devices here. This function turns each into a + DeviceRequest in a DRA ResourceClaim for the serving pods. At + least one request is always present: the scheduler only pins a + replica to a pool that yields a claimable device, so the + serving workload always has a ResourceClaim to bind through. + minItems: 1 + maxItems: 16 + x-kubernetes-list-type: map + x-kubernetes-list-map-keys: [name] + items: + type: object + required: [name, deviceClassName] + properties: + name: + type: string + description: Request name; becomes the DeviceRequest name. + minLength: 1 + maxLength: 63 + deviceClassName: + type: string + description: >- + Cluster-scoped DRA DeviceClass to claim through, from the + matched InferenceClass device. + minLength: 1 + maxLength: 253 + count: + type: integer + description: How many devices to claim. + default: 1 + minimum: 1 + maximum: 64 + selectors: + type: array + description: >- + DRA CEL selectors copied verbatim from the nodeSelector + request, ANDed in the DeviceRequest. + maxItems: 8 + x-kubernetes-list-type: atomic + items: + type: object + # A selector must carry a selector kind (today only cel). + # An empty {} selector would match any device, silently + # widening the DRA claim the serving pod binds through. + minProperties: 1 + properties: + cel: + type: string + minLength: 1 + maxLength: 10240 + routing: + type: object + properties: + template: + type: object + properties: + spec: + type: object + x-kubernetes-preserve-unknown-fields: true status: type: object diff --git a/design/disaggregation-routing-phase2-decisions.md b/design/disaggregation-routing-phase2-decisions.md new file mode 100644 index 000000000..82a87b700 --- /dev/null +++ b/design/disaggregation-routing-phase2-decisions.md @@ -0,0 +1,598 @@ +# Disaggregation Routing: Phase 2 Decisions + +**Status:** Complete — gates EPP-emission code in compose-model-replica +**Date:** June 2026 +**Branch:** dennis/disagg-impl +**Researched from:** upstream chart source and CI workflows; no blog posts; URLs cited inline + +## Verdict + +All three questions are answered at HIGH confidence from primary source (Helm templates, CI release +workflow, and upstream examples). Compose-model-replica does not use a per-model Helm install; +it emits provider-kubernetes Objects directly. The llm-d `routerlib` chart confirms the exact set of +Objects that must be emitted and their required fields. The Envoy AI Gateway inference-pool addon is +a two-line YAML addition to the base `envoy-gateway-values.yaml` and is required for `InferencePool` +as an HTTPRoute backendRef. The pd-sidecar image is published to +`ghcr.io/llm-d/llm-d-routing-sidecar` at every semver release tag; v0.8.0 +(latest as of June 2026) should be used as the default. No live-cluster unknowns remain that +block writing code. + +--- + +## Q1: EPP deploy split — helm chart vs compose-model-replica + +### Answer + +The `llm-d-router-gateway` Helm chart (at +`config/charts/llm-d-router-gateway/templates/epp.yaml`) emits ALL of the following in one +render: InferencePool, ConfigMap (EndpointPickerConfig), EPP Deployment, EPP Service (port 9002), +ServiceAccount, Role+RoleBinding for pod watch, and (if HA) a leader-election Role+RoleBinding. +There is no separate "install EPP by hand" step. Because Modelplane does not use per-model helm +installs, **compose-model-replica must emit all of these Objects directly** as +provider-kubernetes `Object` resources. + +The EPP container image comes from the user's `spec.routing.template`; the chart's default is +`ghcr.io/llm-d/llm-d-router-endpoint-picker-dev:main` (a rolling dev tag — production deployments +should override with the release image `ghcr.io/llm-d/llm-d-inference-scheduler:v0.8.0`). + +**EPP container required args** (from `_deployment.yaml`): +``` +--pool-name # name of the InferencePool object +--pool-namespace # namespace where InferencePool lives +--pool-group inference.networking.k8s.io +--zap-encoder json +--config-file /config/ # default: default-plugins.yaml +``` + +No `--grpc-port` arg in the llm-d chart; the port is set via container `ports[].containerPort: 9002` +and the `GAIE base.yaml` reference manifest uses `--grpc-port 9002` explicitly. Include it. + +**Required env vars** (from `_deployment.yaml`): +```yaml +env: + - name: NAMESPACE + valueFrom: + fieldRef: + fieldPath: metadata.namespace + - name: POD_NAME + valueFrom: + fieldRef: + fieldPath: metadata.name +``` + +**ConfigMap mount path:** `/config/` — the EPP Deployment mounts a volume named +`plugins-config-volume` at `/config`, backed by a ConfigMap whose key is the filename passed to +`--config-file` (e.g., `pd-epp-config.yaml`). + +**RBAC (namespace-scoped Role, confirmed from `_rbac.yaml` and `rbac.yaml` in +`llm-d-router-gateway`):** + +Primary SA Role (`-sa`): +```yaml +rules: +- apiGroups: [""] + resources: ["pods"] + verbs: ["get", "watch", "list"] +``` + +Non-SA Role (`-non-sa`) — watches inference CRDs: +```yaml +rules: +- apiGroups: ["inference.networking.x-k8s.io"] + resources: ["inferenceobjectives", "inferencemodelrewrites"] + verbs: ["get", "watch", "list"] +- apiGroups: ["llm-d.ai"] + resources: ["inferenceobjectives", "inferencemodelrewrites"] + verbs: ["get", "watch", "list"] +- apiGroups: ["inference.networking.k8s.io"] + resources: ["inferencepools"] + verbs: ["get", "watch", "list"] +``` + +Additionally the GAIE `base.yaml` example adds a ClusterRole for `tokenreviews` and +`subjectaccessreviews` (used by the metrics auth path): +```yaml +rules: +- apiGroups: ["authentication.k8s.io"] + resources: ["tokenreviews"] + verbs: ["create"] +- apiGroups: ["authorization.k8s.io"] + resources: ["subjectaccessreviews"] + verbs: ["create"] +``` + +Both namespace-scoped Roles bind to the same ServiceAccount. The ClusterRole/ClusterRoleBinding +is conditional on `monitoring.prometheus.enabled` in the chart but required by the GAIE +reference manifest unconditionally. Include it unconditionally for correctness. + +All RBAC is namespace-scoped (Role, not ClusterRole) for pod watch, confirming that the EPP can +only watch pods in its own namespace. This is fine because all serving workloads are co-located +in `default`. + +**Confidence: HIGH** — all fields read directly from chart templates. + +**Sources:** +- `config/charts/llm-d-router-gateway/templates/epp.yaml` — + https://github.com/llm-d/llm-d-inference-scheduler/blob/main/config/charts/llm-d-router-gateway/templates/epp.yaml +- `config/charts/routerlib/templates/_deployment.yaml` — + https://github.com/llm-d/llm-d-inference-scheduler/blob/main/config/charts/routerlib/templates/_deployment.yaml +- `config/charts/routerlib/templates/_rbac.yaml` — + https://github.com/llm-d/llm-d-inference-scheduler/blob/main/config/charts/routerlib/templates/_rbac.yaml +- `config/charts/llm-d-router-gateway/templates/rbac.yaml` — + https://github.com/llm-d/llm-d-inference-scheduler/blob/main/config/charts/llm-d-router-gateway/templates/rbac.yaml +- GAIE `base.yaml` reference manifest — + https://github.com/envoyproxy/ai-gateway/blob/main/examples/inference-pool/base.yaml + +--- + +## Q2: Inference-pool gateway addon values + +### Answer + +Two `values` files are required when installing `gateway-helm` (Envoy Gateway) for +InferencePool + HTTPRoute support: + +**File 1 — base: `manifests/envoy-gateway-values.yaml`** +```yaml +config: + envoyGateway: + gateway: + controllerName: gateway.envoyproxy.io/gatewayclass-controller + logging: + level: + default: info + provider: + type: Kubernetes + extensionApis: + enableEnvoyPatchPolicy: true + enableBackend: true # Required: enables Backend API for AI service backends + extensionManager: + hooks: + xdsTranslator: + translation: + listener: { includeAll: true } + route: { includeAll: true } + cluster: { includeAll: true } + secret: { includeAll: true } + post: + - Translation + - Cluster + - Route + service: + fqdn: + hostname: ai-gateway-controller.envoy-ai-gateway-system.svc.cluster.local + port: 1063 +``` + +**File 2 — inference-pool addon: `examples/inference-pool/envoy-gateway-values-addon.yaml`** +```yaml +config: + envoyGateway: + extensionManager: + backendResources: + - group: inference.networking.k8s.io + kind: InferencePool + version: v1 +``` + +The addon is **required** for `HTTPRoute -> InferencePool` backendRef. It registers +`InferencePool` as a recognized backend resource type in the Envoy Gateway extension manager. +Without it, Envoy Gateway does not know to hand off InferencePool resolution to the AI Gateway +controller's ext-proc path, and the HTTPRoute will fail to program. + +The comment in the addon file confirms this is composed with the base: +``` +helm upgrade -i eg oci://docker.io/envoyproxy/gateway-helm \ + -f ../../manifests/envoy-gateway-values.yaml \ + -f envoy-gateway-values-addon.yaml +``` + +ServingStack must pass BOTH files as `-f` overrides to the `gateway-helm` install. These two +files can be merged into one values object in the Crossplane `Release` resource. + +**Confidence: HIGH** — file content read directly from upstream repo; the addon comment +explicitly states its purpose. + +**Sources:** +- `manifests/envoy-gateway-values.yaml` — + https://github.com/envoyproxy/ai-gateway/blob/main/manifests/envoy-gateway-values.yaml +- `examples/inference-pool/envoy-gateway-values-addon.yaml` — + https://github.com/envoyproxy/ai-gateway/blob/main/examples/inference-pool/envoy-gateway-values-addon.yaml +- `examples/inference-pool/httproute.yaml` (confirms HTTPRoute + InferencePool backendRef shape) — + https://github.com/envoyproxy/ai-gateway/blob/main/examples/inference-pool/httproute.yaml + +--- + +## Q3: pd-sidecar image + +### Answer + +The pd-sidecar image is **published** to ghcr.io on every semver release of +`llm-d/llm-d-inference-scheduler`. The CI release workflow (`.github/workflows/ci-release.yaml`) +sets `sidecar-image-name` to `${repo}-disagg-sidecar` (where `repo` is the GitHub repository +name `llm-d-inference-scheduler`), and the build/push action pushes to: + +``` +ghcr.io/llm-d/llm-d-routing-sidecar: +``` + +For stable releases (non-prerelease), the `latest` tag is also pushed. The latest release is +v0.8.0 (published 2026-04-28). + +**Pin for compose-model-replica:** +``` +ghcr.io/llm-d/llm-d-routing-sidecar:v0.8.0 +``` + +Direct package API access was blocked (no `read:packages` token scope), so this image reference +is derived from the CI workflow source rather than live registry inspection. The tag pattern +`ghcr.io/llm-d/-disagg-sidecar:` is unambiguous. Mark as "derived from CI +source; not live-registry confirmed" and add a cluster smoke-test step: `docker pull +ghcr.io/llm-d/llm-d-routing-sidecar:v0.8.0` before wiring into the decode +pod spec. + +The EPP image (for reference) follows the same convention: +``` +ghcr.io/llm-d/llm-d-inference-scheduler:v0.8.0 +``` + +The `values.yaml` default `ghcr.io/llm-d/llm-d-router-endpoint-picker-dev:main` is the +rolling dev image and must NOT be used in production; the release image above should be the +default when `spec.routing.template` is omitted or when providing a fallback. + +**Confidence: HIGH** for image naming convention (derived from CI source); **MEDIUM** for +confirming the image is actually pullable on ghcr.io without a live registry check. + +**Sources:** +- `.github/workflows/ci-release.yaml` — sidecar_name derivation — + https://github.com/llm-d/llm-d-inference-scheduler/blob/main/.github/workflows/ci-release.yaml +- `.github/actions/docker-build-and-push/action.yml` — tag + registry pattern — + https://github.com/llm-d/llm-d-inference-scheduler/blob/main/.github/actions/docker-build-and-push/action.yml +- `Dockerfile.sidecar` (confirms binary is `pd-sidecar`, base image `gcr.io/distroless/static:nonroot`) — + https://github.com/llm-d/llm-d-inference-scheduler/blob/main/Dockerfile.sidecar +- Latest release v0.8.0 — https://github.com/llm-d/llm-d-inference-scheduler/releases/tag/v0.8.0 + +--- + +## What compose-model-replica emits (disagg replica) + +This section is the implementer's reference. Each item is a provider-kubernetes `Object`. +All objects are in the model's namespace (co-located in `default`). + +### 1. ServiceAccount + +```yaml +apiVersion: v1 +kind: ServiceAccount +metadata: + name: -epp + namespace: +``` + +### 2. Role — pod + inference CRD watch (namespace-scoped) + +```yaml +apiVersion: rbac.authorization.k8s.io/v1 +kind: Role +metadata: + name: -epp-sa + namespace: +rules: +- apiGroups: [""] + resources: ["pods"] + verbs: ["get", "watch", "list"] +- apiGroups: ["inference.networking.x-k8s.io"] + resources: ["inferenceobjectives", "inferencemodelrewrites"] + verbs: ["get", "watch", "list"] +- apiGroups: ["llm-d.ai"] + resources: ["inferenceobjectives", "inferencemodelrewrites"] + verbs: ["get", "watch", "list"] +- apiGroups: ["inference.networking.k8s.io"] + resources: ["inferencepools"] + verbs: ["get", "watch", "list"] +``` + +### 3. RoleBinding + +```yaml +apiVersion: rbac.authorization.k8s.io/v1 +kind: RoleBinding +metadata: + name: -epp-sa + namespace: +subjects: +- kind: ServiceAccount + name: -epp + namespace: +roleRef: + apiGroup: rbac.authorization.k8s.io + kind: Role + name: -epp-sa +``` + +### 4. ClusterRole + ClusterRoleBinding — metrics auth reviewer + +```yaml +apiVersion: rbac.authorization.k8s.io/v1 +kind: ClusterRole +metadata: + name: -epp-auth-reviewer +rules: +- apiGroups: ["authentication.k8s.io"] + resources: ["tokenreviews"] + verbs: ["create"] +- apiGroups: ["authorization.k8s.io"] + resources: ["subjectaccessreviews"] + verbs: ["create"] +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: ClusterRoleBinding +metadata: + name: -epp-auth-reviewer +subjects: +- kind: ServiceAccount + name: -epp + namespace: +roleRef: + apiGroup: rbac.authorization.k8s.io + kind: ClusterRole + name: -epp-auth-reviewer +``` + +### 5. ConfigMap — EndpointPickerConfig (disagg profile) + +Key name matches `--config-file` arg (e.g., `pd-epp-config.yaml`). Content is the official +`deploy/config/pd-epp-config.yaml` from upstream, verbatim, for the default disagg profile. +Custom configs can be injected via `spec.routing.template` annotations or a dedicated API field. + +```yaml +apiVersion: v1 +kind: ConfigMap +metadata: + name: -epp + namespace: +data: + pd-epp-config.yaml: | + apiVersion: inference.networking.x-k8s.io/v1alpha1 + kind: EndpointPickerConfig + plugins: + - type: approx-prefix-cache-producer + parameters: + maxPrefixBlocksToMatch: 256 + lruCapacityPerServer: 31250 + - type: prefix-cache-scorer + - type: queue-scorer + - type: prefill-filter + - type: decode-filter + - type: max-score-picker + - type: prefix-based-pd-decider + parameters: + nonCachedTokens: 16 + - type: disagg-profile-handler + parameters: + deciders: + prefill: prefix-based-pd-decider + schedulingProfiles: + - name: prefill + plugins: + - pluginRef: prefill-filter + - pluginRef: max-score-picker + - pluginRef: prefix-cache-scorer + weight: 2 + - pluginRef: queue-scorer + weight: 1 + - name: decode + plugins: + - pluginRef: decode-filter + - pluginRef: max-score-picker + - pluginRef: prefix-cache-scorer + weight: 2 + - pluginRef: queue-scorer + weight: 1 +``` + +### 6. EPP Deployment + +Image source: from `spec.routing.template` (user-supplied PodSpec subset). The EPP container +image comes from the user's `routing.template` containers list; the function injects the required +args and env listed below around it. Default fallback image: +`ghcr.io/llm-d/llm-d-inference-scheduler:v0.8.0`. + +```yaml +apiVersion: apps/v1 +kind: Deployment +metadata: + name: -epp + namespace: +spec: + replicas: 1 + strategy: + type: Recreate # required: single-replica stateful EPP; rolling update not safe + selector: + matchLabels: + app: -epp + template: + metadata: + labels: + app: -epp + spec: + serviceAccountName: -epp + terminationGracePeriodSeconds: 130 + containers: + - name: epp + image: # e.g. ghcr.io/llm-d/llm-d-inference-scheduler:v0.8.0 + args: + - --pool-name + - -pool + - --pool-namespace + - + - --pool-group + - inference.networking.k8s.io + - --zap-encoder + - json + - --config-file + - /config/pd-epp-config.yaml + - --grpc-port + - "9002" + env: + - name: NAMESPACE + valueFrom: + fieldRef: + fieldPath: metadata.namespace + - name: POD_NAME + valueFrom: + fieldRef: + fieldPath: metadata.name + ports: + - name: grpc + containerPort: 9002 + - name: grpc-health + containerPort: 9003 + - name: metrics + containerPort: 9090 + livenessProbe: + grpc: + port: 9003 + service: inference-extension + initialDelaySeconds: 5 + periodSeconds: 10 + readinessProbe: + grpc: + port: 9003 + service: inference-extension + initialDelaySeconds: 5 + periodSeconds: 10 + volumeMounts: + - name: plugins-config-volume + mountPath: /config + volumes: + - name: plugins-config-volume + configMap: + name: -epp +``` + +### 7. EPP Service + +```yaml +apiVersion: v1 +kind: Service +metadata: + name: -epp + namespace: +spec: + selector: + app: -epp + ports: + - name: grpc-ext-proc + protocol: TCP + port: 9002 + targetPort: 9002 + appProtocol: http2 + - name: http-metrics + protocol: TCP + port: 9090 + type: ClusterIP +``` + +### 8. InferencePool + +Selector must match **both** prefill and decode pods (the EPP partitions by `llm-d.ai/role` +internally). Use a shared label present on all model replica pods (e.g., `app: `). + +```yaml +apiVersion: inference.networking.k8s.io/v1 +kind: InferencePool +metadata: + name: -pool + namespace: +spec: + targetPorts: + - number: 8000 # the port on which vLLM (via pd-sidecar on decode pods) listens + selector: + matchLabels: + app: # must match BOTH prefill and decode pods + endpointPickerRef: + name: -epp + port: + number: 9002 + failureMode: FailOpen +``` + +### 9. HTTPRoute (per-model routing rule, one per InferencePool) + +```yaml +apiVersion: gateway.networking.k8s.io/v1 +kind: HTTPRoute +metadata: + name: + namespace: +spec: + parentRefs: + - group: gateway.networking.k8s.io + kind: Gateway + name: + namespace: + rules: + - backendRefs: + - group: inference.networking.k8s.io + kind: InferencePool + name: -pool + namespace: + weight: 1 + matches: + - path: + type: PathPrefix + value: / + timeouts: + request: 60s +``` + +### 10. Decode pod — pd-sidecar container injection + +Every decode `Deployment` (emitted by compose-model-replica) must include the pd-sidecar as an +additional container alongside the vLLM container. The sidecar listens on the external port +(8000) and forwards to vLLM on the inner port (8001). The decode pod's container port visible to +the InferencePool is 8000 (the sidecar). + +**Sidecar image:** `ghcr.io/llm-d/llm-d-routing-sidecar:v0.8.0` +(confirm pullable with `docker pull` before wiring in; see Q3 residual note). + +```yaml +- name: pd-sidecar + image: ghcr.io/llm-d/llm-d-routing-sidecar:v0.8.0 + args: + - --kv-connector=nixlv2 # matches vLLM --kv-transfer-config NixlConnector + - --vllm-port=8001 # vLLM inner port; sidecar listens on 8000 + ports: + - name: http + containerPort: 8000 +``` + +Pod labels required on decode pods: `app: `, `llm-d.ai/role: decode` +Pod labels required on prefill pods: `app: `, `llm-d.ai/role: prefill` + +--- + +## Live validation (GKE, 2026-06-11) — corrections + +Validated on a real GKE cluster. Several CI-derived guesses were wrong and are now fixed: + +1. **Image references (FIXED).** The CI-derived names were wrong and 403'd on ghcr.io. + Verified against the registry, the published public images are + `ghcr.io/llm-d/llm-d-inference-scheduler:v0.8.0` (EPP) and + `ghcr.io/llm-d/llm-d-routing-sidecar:v0.8.0` (pd-sidecar) — both public, no + `imagePullSecret` needed. The earlier `-endpoint-picker` / `-disagg-sidecar` + suffixes do not exist as packages. + +2. **EndpointPickerConfig apiVersion (FIXED).** `llm-d.ai/v1alpha1` is not registered + by the EPP binary (crash-loops on parse). The correct group is + `inference.networking.x-k8s.io/v1alpha1`. With it the EPP runs and reconciles the + InferencePool. + +3. **vLLM NixlConnector (FIXED).** The engines were missing `--kv-transfer-config`, so no + KV handoff could occur. Both roles now pass + `--kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_both"}'` + (NixlConnector does not distinguish kv_role; the routing sidecar drives direction). + +4. **Envoy AI Gateway v1.0.0.** v0.7.0 is pinned; bump when v1.0.0 GAs (~June 30). No code + change unless the addon values schema changes. diff --git a/design/disaggregation-routing-spike-phase2.md b/design/disaggregation-routing-spike-phase2.md new file mode 100644 index 000000000..dedf574d4 --- /dev/null +++ b/design/disaggregation-routing-spike-phase2.md @@ -0,0 +1,352 @@ +# Disaggregation Routing Spike: Phase 2 + +**Status:** Complete +**Date:** June 2026 +**Branch:** docs-preview + +## Verdict + +All three questions are now confirmed at high confidence from primary sources (llm-d +upstream Go source, Envoy AI Gateway docs/source, GAIE manifests). The EPP header +is `x-prefiller-host-port` (confirmed in `pkg/common/routing/common.go`). The EPP +discovers both prefill and decode pods through a **single** `InferencePool` that +selects all pods sharing a common label; within the EPP, `prefill-filter` and +`decode-filter` plugins partition the set using the `llm-d.ai/role` label. Envoy AI +Gateway v0.7.0 (released June 4, 2026, targeting v1.0 GA by June 30, 2026) supports +`InferencePool` as an HTTPRoute backendRef and runs **on top of** core Envoy Gateway +(`gateway-helm`), adding two additional charts (`ai-gateway-crds-helm`, +`ai-gateway-helm`). Phase 2 requires replacing the ServingStack's current standalone +`gateway-helm` install with a three-chart stack. + +--- + +## Question 1: Prefill header and handoff component + +**Confirmed header name: `x-prefiller-host-port`** + +The constant is defined in +`pkg/common/routing/common.go` in +`github.com/llm-d/llm-d-inference-scheduler`: + +```go +// PrefillEndpointHeader is the header name used to indicate Prefill worker +PrefillEndpointHeader = "x-prefiller-host-port" +``` + +The value is in `host:port` format (e.g., `10.0.0.5:8000`). + +The EPP's `disagg-profile-handler` sets this header via the `PreRequest` hook +in `disagg_headers_handler.go`: it picks the prefill pod from the prefill +scheduling profile's result and writes `net.JoinHostPort(addr, port)` into +`request.Headers[routing.PrefillEndpointHeader]`. + +The component that reads the header and performs the handoff is the +**pd-sidecar**, a reverse proxy that runs as a sidecar container on every +**decode** pod. The sidecar was previously in the standalone repo +`github.com/llm-d/llm-d-routing-sidecar`, which is now deprecated and marked +for archival. The code lives in the main inference-scheduler repo under +`cmd/pd-sidecar/` and `pkg/sidecar/`. On receiving a request whose +`x-prefiller-host-port` header is set, the sidecar forwards it to the named +prefill pod to perform remote prefill, receives the KV block IDs back, then +sends the decode request to the local vLLM instance with the KV transfer +parameters. If the header is absent, the decode pod runs both prefill and +decode locally. + +The sidecar is **not** an init container and is not in-engine. It runs as a +sidecar container on the decode pod, listening on a port in front of vLLM +(default 8000), with vLLM listening on an inner port (default 8001). The sidecar +binary is built from `Dockerfile.sidecar` in the inference-scheduler repo. + +**Confidence: HIGH**. Header name confirmed from source; sidecar migration +confirmed from repo README deprecation notice and directory structure. + +Sources: +- `pkg/common/routing/common.go` — https://github.com/llm-d/llm-d-inference-scheduler/blob/main/pkg/common/routing/common.go +- `pkg/epp/framework/plugins/scheduling/profilehandler/disagg/disagg_headers_handler.go` — https://github.com/llm-d/llm-d-inference-scheduler/blob/main/pkg/epp/framework/plugins/scheduling/profilehandler/disagg/disagg_headers_handler.go +- `cmd/pd-sidecar/` — https://github.com/llm-d/llm-d-inference-scheduler/tree/main/cmd/pd-sidecar +- Archived sidecar repo — https://github.com/llm-d/llm-d-routing-sidecar + +--- + +## Question 2: Prefill pod discovery and label convention + +**Discovery mechanism: single InferencePool watches all pods; label filters +partition within the EPP.** + +The EPP is started with `--pool-name` and `--pool-namespace` pointing at a +single `InferencePool`. By default, the EPP reconciler watches the +`InferencePool`'s `spec.selector.matchLabels` to enumerate all matching pods and +builds an endpoint datastore from them. For P/D disaggregation, the +`InferencePool` selector must be broad enough to match **both** prefill and decode +pods (e.g., `app: my-model`). The EPP then uses `prefill-filter` and +`decode-filter` scheduling plugins to split the full pod set into the prefill +and decode pools at scheduling time. + +The native llm-d label convention is: + +| Label key | Values | +|------------------|-----------------------------------------------------| +| `llm-d.ai/role` | `prefill`, `decode`, `encode`, `prefill-decode`, `encode-prefill`, `encode-prefill-decode` | + +The `prefill-filter` plugin accepts pods whose `llm-d.ai/role` is `prefill`, +`encode-prefill`, `prefill-decode`, or `encode-prefill-decode`. The `decode-filter` +accepts `decode`, `prefill-decode`, `encode-prefill-decode`, and the deprecated +`both`. Both are implemented in +`pkg/epp/framework/plugins/scheduling/filter/bylabel/roles.go`. + +The reference deployment manifest (`deploy/components/vllm-prefill/deployment.yaml`) +labels prefill pods with both `llm-d.ai/component: prefill` and +`llm-d.ai/role: prefill`. The InferencePool selector uses a shared label like +`llm-d.ai/inference-serving: "true"` or `app: ` that covers all pods. + +**Implications for Modelplane's `modelplane.ai/pd-role` label:** the llm-d EPP +does not use `modelplane.ai/pd-role`; it only reads `llm-d.ai/role`. If +Modelplane uses the built-in `prefill-filter`/`decode-filter` plugins (which is +the simplest path), pods must carry `llm-d.ai/role: prefill` or +`llm-d.ai/role: decode`. Alternatively, the `label-selector-filter` plugin +(generic Kubernetes selector syntax) can be configured to read any label key, +including `modelplane.ai/pd-role`. The disagreement docs explicitly describe +this as the supported path for external workloads with different labeling +conventions. Using `label-selector-filter` with `modelplane.ai/pd-role` avoids +touching Pod labels but requires a custom `EndpointPickerConfig`. + +**Remaining unknown:** it is not confirmed in source whether the EPP RBAC +(pod watch) must be in the same namespace as the InferencePool. The e2e +manifest uses a namespace-scoped Role rather than ClusterRole, implying both +pools must be co-located. Cross-namespace prefill discovery is unconfirmed. + +**Confidence: HIGH** for single-pool + label-filter mechanism; **MEDIUM** for +cross-namespace assumptions. + +Sources: +- `pkg/epp/framework/plugins/scheduling/filter/bylabel/roles.go` — https://github.com/llm-d/llm-d-inference-scheduler/blob/main/pkg/epp/framework/plugins/scheduling/filter/bylabel/roles.go +- `deploy/components/vllm-prefill/deployment.yaml` — https://github.com/llm-d/llm-d-inference-scheduler/blob/main/deploy/components/vllm-prefill/deployment.yaml +- `test/sidecar/config/nixl/inferencepool.yaml` — https://github.com/llm-d/llm-d-inference-scheduler/blob/main/test/sidecar/config/nixl/inferencepool.yaml +- `docs/disaggregation.md` — https://github.com/llm-d/llm-d-inference-scheduler/blob/main/docs/disaggregation.md +- `deploy/config/sim-pd-epp-config.yaml` — https://github.com/llm-d/llm-d-inference-scheduler/blob/main/deploy/config/sim-pd-epp-config.yaml + +--- + +## Question 3: Envoy AI Gateway InferencePool maturity (mid-2026) + +**Latest release: v0.7.0, released June 4, 2026.** + +The roadmap tracking issue (envoyproxy/ai-gateway#2083) targets v1.0.0 GA by +June 30, 2026, with v1.0.0-rc1 on June 12, 2026. As of today (June 10, 2026), +v0.7.0 is the latest stable release; v1.0.0 has not shipped yet. The v0.6.0 +CRDs were promoted to v1beta1. + +**InferencePool support:** Envoy AI Gateway supports `InferencePool` +(`inference.networking.k8s.io/v1`) as a backendRef in both `HTTPRoute` and +`AIGatewayRoute`. Support was introduced at v0.3.0 (August 2025) and updated to +GAIE v1.0 at v0.4.0 (November 2025). The install requires the GAIE manifests +(CRDs + EPP controller) to be applied separately: + +``` +kubectl apply -f https://github.com/kubernetes-sigs/gateway-api-inference-extension/releases/download/v1.0.1/manifests.yaml +``` + +**InferencePool as HTTPRoute backendRef is functioning in v0.7.0.** The +`examples/inference-pool/httproute.yaml` in the ai-gateway repo demonstrates +this pattern with `group: inference.networking.k8s.io`, `kind: InferencePool` +as a backendRef. However, the Envoy AI Gateway project does not yet designate +any capability as formally "GA/stable" at v0.7.0; v1.0 is the first stability +milestone. Treat InferencePool support as "production-quality beta" for the +Phase 2 pin. + +**Install model: Envoy AI Gateway runs ON TOP of Envoy Gateway.** The install +sequence is four steps: + +1. Gateway API CRDs (already in Modelplane's ServingStack) +2. `oci://docker.io/envoyproxy/gateway-helm` — core Envoy Gateway, with + specific values from `manifests/envoy-gateway-values.yaml` in the + ai-gateway repo that enable `extensionManager` hooks and the `Backend` API. + This replaces the plain `gateway-helm` install in the current ServingStack. +3. `oci://docker.io/envoyproxy/ai-gateway-crds-helm` — AI Gateway CRDs + (namespace: `envoy-ai-gateway-system`) +4. `oci://docker.io/envoyproxy/ai-gateway-helm` — AI Gateway controller + (namespace: `envoy-ai-gateway-system`) + +The `gateway-helm` install in step 2 requires overrides to activate the +`extensionManager` pointing to the AI Gateway controller service: + +```yaml +config.envoyGateway.extensionManager.service.fqdn.hostname: + ai-gateway-controller.envoy-ai-gateway-system.svc.cluster.local +config.envoyGateway.extensionApis.enableBackend: true +``` + +For InferencePool support specifically, an additional values addon +(`examples/inference-pool/envoy-gateway-values-addon.yaml`) may be required +to enable the ext-proc extension needed by the EPP. + +**GatewayClass controllerName:** `gateway.envoyproxy.io/gatewayclass-controller` +(unchanged from core Envoy Gateway). Envoy AI Gateway does not introduce a +separate GatewayClass; it extends the existing one via the extensionManager hook. + +**Confidence: HIGH** for release version, install model, GatewayClass, and +InferencePool backendRef functionality; **MEDIUM** for stability designation +(v1.0 has not shipped as of today). + +Sources: +- Release list — https://aigateway.envoyproxy.io/release-notes/ +- v1.0 GA roadmap — https://github.com/envoyproxy/ai-gateway/issues/2083 +- HTTPRoute + InferencePool example — https://github.com/envoyproxy/ai-gateway/blob/main/examples/inference-pool/httproute.yaml +- Installation guide — https://github.com/envoyproxy/ai-gateway/blob/main/site/docs/getting-started/installation.md +- envoy-gateway-values.yaml — https://github.com/envoyproxy/ai-gateway/blob/main/manifests/envoy-gateway-values.yaml +- InferencePool example README — https://github.com/envoyproxy/ai-gateway/blob/main/examples/inference-pool/README.md +- GAIE v1.0.1 manifests — https://github.com/kubernetes-sigs/gateway-api-inference-extension/releases/download/v1.0.1/manifests.yaml + +--- + +## Backend shape for Phase 2 + +### compose-serving-stack must install + +Replace the existing bare `gateway-helm` install with a three-chart stack: + +| Chart | Registry path | Namespace | Notes | +|---|---|---|---| +| `gateway-helm` | `oci://docker.io/envoyproxy/gateway-helm` | `envoy-gateway-system` | Must apply `envoy-gateway-values.yaml` overrides and the InferencePool addon values | +| `ai-gateway-crds-helm` | `oci://docker.io/envoyproxy/ai-gateway-crds-helm` | `envoy-ai-gateway-system` | Installs `AIGatewayRoute` and related CRDs | +| `ai-gateway-helm` | `oci://docker.io/envoyproxy/ai-gateway-helm` | `envoy-ai-gateway-system` | Runs the AI Gateway controller; depends on gateway-helm being ready | + +Additionally install the GAIE CRDs and EPP controller: + +``` +kubectl apply -f https://github.com/kubernetes-sigs/gateway-api-inference-extension/releases/download/v1.0.1/manifests.yaml +``` + +**Version to pin:** `v0.7.0` for all three ai-gateway charts. Do not use +`v0.0.0-latest`; it is overwritten on the registry and is not reproducible. + +The GatewayClass `controllerName` remains `gateway.envoyproxy.io/gatewayclass-controller`. + +### compose-model-replica llm-d backend must emit (disagg replica) + +**InferencePool (decode pods only, or all pods with role-based filter):** + +```yaml +apiVersion: inference.networking.k8s.io/v1 +kind: InferencePool +metadata: + name: -decode +spec: + selector: + matchLabels: + app: # covers both prefill and decode pods + llm-d.ai/inference-serving: "true" + targetPorts: + - number: 8000 + endpointPickerRef: + name: -epp + kind: Service + port: + number: 9002 +``` + +The InferencePool selector must match **all** pods (prefill + decode). The EPP +partitions them internally using the `EndpointPickerConfig`. + +**EPP EndpointPickerConfig (ConfigMap):** + +```yaml +apiVersion: llm-d.ai/v1alpha1 +kind: EndpointPickerConfig +plugins: +- type: prefill-filter # selects pods with llm-d.ai/role: prefill +- type: decode-filter # selects pods with llm-d.ai/role: decode +- type: prefix-cache-scorer +- type: max-score-picker +- type: prefix-based-pd-decider + parameters: + nonCachedTokens: 16 +- type: disagg-profile-handler + parameters: + deciders: + prefill: prefix-based-pd-decider +schedulingProfiles: +- name: prefill + plugins: + - pluginRef: prefill-filter + - pluginRef: max-score-picker + - pluginRef: prefix-cache-scorer +- name: decode + plugins: + - pluginRef: decode-filter + - pluginRef: max-score-picker + - pluginRef: prefix-cache-scorer +``` + +**Pod labels required on model replica pods:** + +Each pod must carry the shared selector label AND the role label: + +| Pod type | Required labels | +|---|---| +| Prefill pod | `app: `, `llm-d.ai/inference-serving: "true"`, `llm-d.ai/role: prefill` | +| Decode pod | `app: `, `llm-d.ai/inference-serving: "true"`, `llm-d.ai/role: decode` | + +**Relationship to `modelplane.ai/pd-role`:** Modelplane's existing +`modelplane.ai/pd-role` label is used internally by Crossplane compositions for +replica selection. The EPP does not read it. To use the built-in `prefill-filter` +and `decode-filter` plugins, replicas must additionally carry `llm-d.ai/role`. +Alternatively, replace the built-in filters with `label-selector-filter` plugins +configured to read `modelplane.ai/pd-role` — this avoids adding the llm-d label +but requires explicit EPP config for each model. + +**Decode pod sidecar container:** + +The pd-sidecar must be injected as a sidecar on every decode pod. The image is +built from `Dockerfile.sidecar` in the inference-scheduler repo. Key flags: + +``` +--kv-connector=nixlv2 # matches vLLM's --kv-transfer-config NixlConnector +--vllm-port=8001 # vLLM listens here; sidecar listens on 8000 +--inference-pool=/ # for SSRF allowlisting (optional) +``` + +**HTTPRoute → InferencePool:** + +```yaml +apiVersion: gateway.networking.k8s.io/v1 +kind: HTTPRoute +spec: + parentRefs: + - kind: Gateway + name: + rules: + - backendRefs: + - group: inference.networking.k8s.io + kind: InferencePool + name: -decode + namespace: +``` + +--- + +## Remaining unknowns that block writing routing code + +1. **EPP cross-namespace pod watch.** The test manifests use a namespace-scoped + Role for pod watch. If Modelplane places prefill and decode pods in separate + namespaces (one per replica), the EPP cannot see both. Confirm whether a single + InferencePool + EPP can watch pods across namespaces, or whether prefill and + decode must be co-located in the same namespace. + +2. **GAIE EPP deploy model in llm-d helm.** The inference-scheduler's helm chart + (`config/charts/routerlib`) generates the InferencePool template but it is + unclear whether it also deploys the EPP Deployment and Service, or whether + compose-model-replica must emit those directly. Verify what the routerlib chart + installs vs what the compose function must emit. + +3. **Envoy AI Gateway v1.0 stability date.** As of June 10, 2026, v1.0.0-rc1 is + due June 12 and GA on June 30. The ServingStack version pin should be updated + to v1.0.0 when it ships. The rc1 is safe to test against but should not be + used in a production pin. + +4. **InferencePool addon Envoy Gateway values.** The ai-gateway inference-pool + example uses `examples/inference-pool/envoy-gateway-values-addon.yaml` in + addition to the base `manifests/envoy-gateway-values.yaml`. The exact content + of that addon and whether it is required for the HTTPRoute + InferencePool + path (vs only AIGatewayRoute) needs to be verified before finalising the + ServingStack helm values. diff --git a/design/disaggregation-routing-spike.md b/design/disaggregation-routing-spike.md new file mode 100644 index 000000000..66171a0f5 --- /dev/null +++ b/design/disaggregation-routing-spike.md @@ -0,0 +1,150 @@ +# Disaggregation Routing Feasibility Spike + +**Status:** Complete +**Date:** June 2026 +**Author:** Dennis Ramdass +**Branch:** dennis/disagg-impl + +## Verdict + +The working assumption — "emit one InferencePool selecting both role pods, EPP as picker, HTTPRoute→InferencePool" — needs modification on two independent axes. First, core Envoy Gateway (which Modelplane's ServingStack installs) does **not** support `InferencePool` as an HTTPRoute backendRef; that capability lives in a separate project, Envoy AI Gateway, or in Istio, kgateway, and GKE Gateway. Second, the llm-d inference-scheduler's disaggregation architecture uses a **single** InferencePool that contains only the decode pods; the EPP picks the decode target, then injects an `x-prefiller-host-port` header into the forwarded request so the decode-side routing sidecar can pull the KV cache from the chosen prefill pod. The assumption that the InferencePool selects "both role pods" and the EPP pair-picks is incorrect — prefill pods are outside the InferencePool entirely, reachable only as sidecar-forwarded targets. These two findings together mean that the Phase 2 backend must either (a) switch the workload gateway from core Envoy Gateway to a GAIE-conformant implementation or (b) retain the current `HTTPRoute → Service` pattern and build disaggregation coordination at a lower level. + +--- + +## Gateway Support (Q1) + +### InferencePool API version and status + +The Gateway API Inference Extension (GAIE) shipped `InferencePool` v1 GA under the API group `inference.networking.k8s.io` (the `x-k8s.io` pre-GA prefix was dropped at the v1.0.0 release). As of v1.0.1 (late 2025), the resource is considered stable. The companion `InferenceModel` type was renamed `InferenceObjective` at v1. All current documentation, examples, and conformance tests use `inference.networking.k8s.io/v1`. + +Sources: +- [InferencePool API type](https://gateway-api-inference-extension.sigs.k8s.io/api-types/inferencepool/) — GA since v1.0.0, `inference.networking.k8s.io/v1` +- [v1 API Reference](https://gateway-api-inference-extension.sigs.k8s.io/reference/spec/) +- [Introducing Gateway API Inference Extension (Kubernetes blog, June 2025)](https://kubernetes.io/blog/2025/06/05/introducing-gateway-api-inference-extension/) + +### Core Envoy Gateway + +**Core Envoy Gateway does not support InferencePool.** Modelplane's ServingStack installs the core Envoy Gateway chart (`oci://docker.io/envoyproxy/gateway-helm`, version v1.3.0 as of the current design). The v1.3.0 release notes contain no mention of InferencePool, GAIE, or inference extension support. The project's extension-types API lists only `Endpoints` and `DynamicResolver` as backend types; InferencePool is absent. The GAIE implementations list on `gateway-api-inference-extension.sigs.k8s.io` does not include core Envoy Gateway. + +This is already noted in the codebase: `functions/compose-model-replica/function/backends/llmd.py` explicitly documents that "Envoy Gateway's `InferencePool` v1 support is unconfirmed; alternatively switch the workload gateway to Istio/agentgateway." + +Sources: +- [Envoy Gateway v1.3.0 release notes](https://gateway.envoyproxy.io/news/releases/notes/v1.3.0/) — no inference extension mention +- [Envoy Gateway extension types](https://gateway.envoyproxy.io/latest/api/extension_types/) — no InferencePool + +### Envoy AI Gateway (separate project) + +**Envoy AI Gateway** is a distinct project (`envoyproxy/ai-gateway`, `aigateway.envoyproxy.io`) that wraps core Envoy Gateway with AI-specific features. It is **not** what Modelplane's ServingStack installs. Envoy AI Gateway v0.3.0 (August 2025) introduced InferencePool support via integration with GAIE v0.5.1. HTTPRoute backendRefs with `group: inference.networking.k8s.io`, `kind: InferencePool` are supported. Only one InferencePool per HTTPRoute rule is permitted. Envoy AI Gateway labels this as a non-alpha capability, but the project as a whole is still pre-v1.0. + +Sources: +- [Envoy AI Gateway InferencePool support (v0.3 docs)](https://aigateway.envoyproxy.io/docs/0.3/capabilities/inference/inferencepool-support/) +- [HTTPRoute + InferencePool guide](https://aigateway.envoyproxy.io/docs/capabilities/inference/httproute-inferencepool/) +- [EPP blog post (July 2025)](https://aigateway.envoyproxy.io/blog/endpoint-picker-for-inference-routing/) +- [Envoy AI Gateway v0.3.x release notes](https://aigateway.envoyproxy.io/release-notes/v0.3/) + +### Alternatives that do support InferencePool + +All three alternatives below are stable or near-stable with GAIE v1: + +| Gateway | InferencePool support | Notes | +|---|---|---| +| **Istio** | v1.28+ (full v1 support); v1.29 promotes to beta | Full service-mesh feature set; significant operational overhead vs. standalone Envoy Gateway | +| **kgateway** | v2.0.x (stable) | Envoy-based; built specifically for AI workload routing; lighter than Istio | +| **GKE Gateway** | Listed as supported | GCP-managed; only relevant if workload clusters are GKE | +| **Envoy AI Gateway** | v0.3.0 (pre-v1.0 project) | Superset of core Envoy Gateway; replacing the Helm chart is the lowest-friction path | + +Switching Modelplane's workload gateway would involve: +1. Replacing the Envoy Gateway Helm chart in `compose-serving-stack` with the chosen alternative's chart. +2. Updating the GatewayClass controller name. +3. Verifying that existing `HTTPRoute → Service` resources (current native and llm-d paths) continue to work unchanged (all three alternatives are Gateway API conformant). +4. For Istio: adding the Istio control-plane operator; the operational footprint is substantially larger. +5. For Envoy AI Gateway: the switch is minimal — it ships its own operator on top of the core Envoy Gateway data plane, so existing HTTPRoutes continue to work and InferencePool becomes available as an additional backendRef kind. + +Sources: +- [Istio GAIE support blog](https://istio.io/latest/blog/2025/inference-extension-support/) +- [Istio 1.28 GA announcement](https://istio.io/latest/news/releases/1.28.x/announcing-1.28/) +- [kgateway inference extension docs](https://kgateway.dev/docs/envoy/2.0.x/integrations/inference-extension/) +- [GAIE implementations page](https://gateway-api-inference-extension.sigs.k8s.io/implementations/gateways/) + +--- + +## Disaggregation Request Mechanism (Q2) + +### Architecture overview + +The llm-d inference-scheduler architecture document explicitly states: **"Single `InferencePool` and single `EPP` due to Envoy limitations."** The InferencePool contains **only the decode pods** — it selects on `llm-d.ai/role: decode` (plus an app label). Prefill pods are separate Kubernetes workloads that are not members of the InferencePool. + +Source: `github.com/llm-d/llm-d-inference-scheduler/blob/main/docs/architecture.md` (confirmed via fetch; single-pool constraint explicitly stated) + +### Request flow + +1. **Client → Gateway**: The client sends an OpenAI-compatible request to the workload cluster's inference gateway (HTTPRoute → InferencePool). +2. **EPP selects decode pod**: The EPP (running as a GAIE ext-proc sidecar) receives the request via Envoy's External Processing filter. It runs the scheduling pipeline against the decode pods in the InferencePool: + - **Filter pass**: The decode-filter (`NewDecodeRole`) retains only pods with `llm-d.ai/role` values of `decode`, `prefill-decode`, or `encode-prefill-decode`. + - **Score pass**: Scorers evaluate KV cache locality, queue depth, prefix hit probability, and session affinity. + - **Select**: The highest-scored decode pod is chosen. +3. **EPP also selects a prefill pod**: For a disaggregated request (prompt length above the disaggregation threshold), the scheduler runs a second scheduling pass — a prefill filter pass using `NewPrefillRole`, which retains pods labelled `prefill`, `prefill-decode`, or `encode-prefill-decode`. A prefill pod is selected for KV-cache locality. +4. **Header injection**: The EPP injects the chosen prefill target as an `x-prefiller-host-port` header (in `host:port` format) into the request before forwarding it to the selected decode pod. +5. **Decode pod → Prefill pod (sidecar)**: A routing sidecar co-located with the decode pod intercepts the request, reads `x-prefiller-host-port`, and proxies the prompt to the designated prefill pod's vLLM (`kv_producer`) engine. The prefill engine processes the prompt and transfers the resulting KV cache to the decode pod via NixlConnector over the fast interconnect (NVLink/RDMA). +6. **Decode pod generates tokens**: The decode pod's vLLM (`kv_consumer`) engine consumes the transferred KV cache and generates tokens. The streaming response returns through the sidecar → EPP → gateway → client. + +**Note**: The routing sidecar project (`llm-d/llm-d-routing-sidecar`) was archived on 3 February 2026; its code has been folded into `llm-d/llm-d-inference-scheduler`. The disaggregation sidecar is now described as a component of the scheduler repo, deployed alongside decode workers. + +### Required pod labels + +| Label | Prefill pods | Decode pods | "Both" pods | +|---|---|---|---| +| `llm-d.ai/role` | `prefill` | `decode` | `prefill-decode` | +| `app` (example) | `-prefill` | `-decode` | `-worker` | + +The InferencePool selector uses `llm-d.ai/role: decode` (plus an app label). Prefill pods are not selected by the InferencePool but must be reachable by cluster-internal DNS for the sidecar's `x-prefiller-host-port` forwarding. + +Sources: +- [llm-d inference-scheduler architecture](https://github.com/llm-d/llm-d-inference-scheduler/blob/main/docs/architecture.md) — single InferencePool constraint, dual filter pass for P/D +- [filter package (pkg.go.dev)](https://pkg.go.dev/github.com/llm-d/llm-d-inference-scheduler/pkg/plugins/filter) — `NewDecodeRole`, `NewPrefillRole`, `llm-d.ai/role` label constants +- [llm-d routing sidecar (archived)](https://github.com/llm-d/llm-d-routing-sidecar) — `x-prefiller-host-port` header, archived Feb 2026, code moved to scheduler repo +- [Solo.io deep dive](https://www.solo.io/blog/deep-dive-into-llm-d-and-distributed-inference) — `x-prefiller-url` header and decode→prefill forwarding flow +- [Spheron deployment guide](https://www.spheron.network/blog/llm-d-kubernetes-disaggregated-inference-guide/) — InferencePool selector targets decode pods only; prefill pods are outside the pool + +--- + +## Implications for the Phase 2 backend + +### Working assumption: "emit one InferencePool selecting both role pods, EPP as picker, HTTPRoute→InferencePool" + +This assumption must be revised on two counts: + +**1. The InferencePool selects decode pods only, not both roles.** +The assumption that the InferencePool selects "both role pods" is incorrect. The InferencePool is a decode-only pool. Prefill pods live outside it. The EPP runs two internal scheduling passes (one per role) and coordinates via the `x-prefiller-host-port` header, but only decode pods are registered in the pool. The correct mental model is: one InferencePool → decode pods only; the EPP has out-of-band knowledge of prefill pod addresses (e.g. from a Kubernetes-watch of pods with `llm-d.ai/role: prefill` in the same namespace). **This part of the working assumption is wrong but does not require an architectural rethink — the Phase 2 backend should emit one InferencePool scoped to decode pods plus a separate (unlabelled-by-pool) set of prefill pods with a headless Service so the sidecar can address them.** + +**2. Core Envoy Gateway does not support HTTPRoute→InferencePool.** +The gateway half of the working assumption ("HTTPRoute→InferencePool") cannot be implemented with the current workload gateway. The Phase 2 backend must make one of three choices before emitting InferencePool resources: + +| Option | Change required | Complexity | +|---|---|---| +| **A: Switch to Envoy AI Gateway** | Replace the `gateway-helm` chart in `compose-serving-stack` with the Envoy AI Gateway chart; update GatewayClass controller name. Existing `HTTPRoute → Service` routes continue to work. | Low — same data plane, same resource model | +| **B: Switch to kgateway** | Replace chart and GatewayClass. Full GAIE v1 support in 2.0.x. Existing HTTPRoutes continue to work. | Low-medium | +| **C: Switch to Istio** | Add Istio control plane operator; GAIE v1 support in Istio 1.28+. Much larger operational footprint. | High | +| **D: Retain HTTPRoute→Service (no InferencePool)** | Keep the current pattern; implement disaggregation coordination entirely in the routing sidecar injected alongside decode pods, without a GAIE EPP. No gateway change needed. | Low — defers GAIE entirely | + +Option A is the lowest-friction path: Envoy AI Gateway is built on core Envoy Gateway, its Helm chart replaces the existing one, and no existing Gateway API resources need to change. + +> **DECISION (Dennis, 2026-06-10): Option A — Envoy AI Gateway.** ServingStack swaps its +> `envoyproxy/gateway-helm` release for the Envoy AI Gateway chart (same EG data plane) and +> installs the GAIE `InferencePool` CRDs. That gateway swap is well-understood and can land +> independently. Before writing the InferencePool/EPP/sidecar emission, confirm the two +> still-open items against `llm-d/llm-d-inference-scheduler` source: (1) the exact prefiller +> header name (`x-prefiller-host-port` vs `x-prefiller-url`) and how the EPP discovers prefill +> pod addresses outside the pool; (2) Envoy AI Gateway's current InferencePool maturity +> (confirmed at v0.3.x; verify the mid-2026 release/stability). Don't emit routing resources +> on the unconfirmed mechanism. + +Option D is acceptable as a v0.1 disaggregation target if GAIE routing is deferred. The current `llmd.py` backend already uses this pattern (HTTPRoute→Service), and disaggregation can be layered on top via sidecar injection without involving the gateway at all. The gap is that prefix-cache-aware endpoint selection (the primary value of the EPP) is unavailable without a GAIE-conformant gateway. + +### Confidence and evidence gaps + +- The single-InferencePool architecture for disaggregation is well-evidenced from the scheduler's own architecture doc and the filter package source, corroborated by secondary sources. **High confidence.** +- The `x-prefiller-host-port` header and sidecar-forwarding flow are confirmed by the (now-archived) routing sidecar repo and the Solo.io deep dive. The routing sidecar code has moved into the scheduler repo; the header may have been renamed (the archived repo uses `x-prefiller-host-port`; the Solo.io post uses `x-prefiller-url`). The exact current header name should be confirmed against `llm-d/llm-d-inference-scheduler` source before implementing. **Medium confidence on exact header name.** +- Core Envoy Gateway's lack of InferencePool support is confirmed by absence of any mention in release notes and extension types documentation, and is consistent with the existing in-repo comment in `llmd.py`. **High confidence.** +- Envoy AI Gateway's InferencePool support at v0.3.0 is confirmed. Whether the project has reached v1.0 or stable status by mid-2026 is unconfirmed — latest confirmed release is v0.3.x (August 2025 series). **Medium confidence on current release.** +- The exact llm-d mechanism for the EPP to discover prefill pod addresses outside the InferencePool (Kubernetes watch, sidecar config, etc.) is not fully documented in publicly available sources. **Low confidence — check llm-d source directly before implementing C2 (pd-role label / decode Service excludes prefill).** diff --git a/functions/compose-model-deployment/function/fn.py b/functions/compose-model-deployment/function/fn.py index 975cbc81c..940a3f26d 100644 --- a/functions/compose-model-deployment/function/fn.py +++ b/functions/compose-model-deployment/function/fn.py @@ -190,6 +190,11 @@ def compose_replicas(self, matched): # are different Pydantic classes (generated from different XRDs # with the same schema). workers = mrv1alpha1.Workers.model_validate(self.xr.spec.workers.model_dump(exclude_none=True)) + # Materialize the effective worker count onto the replica. workers.count + # no longer carries a schema default (a disaggregated deployment must set + # it explicitly; a unified one may omit it and means 1), so record the + # resolved value here rather than leave the replica to re-default it. + workers.count = int(self.xr.spec.workers.count or 1) for cluster_info in matched: replica_key = name.replica_key(cluster_info) @@ -228,6 +233,29 @@ def compose_replicas(self, matched): ) if self.xr.spec.modelCacheRef: replica.spec.modelCacheRef = mrv1alpha1.ModelCacheRef(name=self.xr.spec.modelCacheRef.name) + if self.xr.spec.prefill and cluster_info.prefill_pool: + prefill_workers = mrv1alpha1.Workers.model_validate( + self.xr.spec.prefill.workers.model_dump(exclude_none=True) + ) + prefill_workers.count = int(self.xr.spec.prefill.workers.count or 1) + prefill_requests = [ + mrv1alpha1.DeviceRequest( + name=r.name, + deviceClassName=r.device_class_name, + count=r.count, + selectors=[mrv1alpha1.Selector(cel=c) for c in r.cel_selectors], + ) + for r in cluster_info.prefill_device_requests + ] + replica.spec.prefill = mrv1alpha1.Prefill( + workers=prefill_workers, + nodePoolName=cluster_info.prefill_pool, + deviceRequests=prefill_requests, + ) + if self.xr.spec.routing: + replica.spec.routing = mrv1alpha1.Routing.model_validate( + self.xr.spec.routing.model_dump(exclude_none=True) + ) resource.update(self.rsp.desired.resources[replica_key], replica) def compose_endpoints(self, matched): diff --git a/functions/compose-model-deployment/function/scheduling.py b/functions/compose-model-deployment/function/scheduling.py index 76a1715fa..8128752f7 100644 --- a/functions/compose-model-deployment/function/scheduling.py +++ b/functions/compose-model-deployment/function/scheduling.py @@ -107,6 +107,11 @@ class Candidate: # non-empty: a pool matches only when at least one claim: DRA device # resolves (see _match_pool), so every scheduled replica has a claim. device_requests: list[DeviceRequest] = field(default_factory=list) + # Prefill role placement, for disaggregated replicas. Empty for unified. + # Asymmetric with the decode fields above on purpose: decode is required and + # owns the ModelEndpoint; prefill is optional and internal. + prefill_pool: str = "" + prefill_device_requests: list[DeviceRequest] = field(default_factory=list) @dataclass @@ -132,6 +137,14 @@ def topology_shape(workers) -> Shape: return Shape(nodes_per_replica=nodes_per_worker * count) +def prefill_shape(deployment) -> Shape | None: + """Nodes-per-replica for the prefill role, or None when unified.""" + prefill = getattr(deployment.spec, "prefill", None) + if not prefill: + return None + return topology_shape(prefill.workers) + + def _cluster_ready(cluster: icv1alpha1.InferenceCluster) -> bool: """Check that the cluster is Ready and has a gateway address. @@ -167,8 +180,17 @@ def compile_requests(deployment: mdv1alpha1.ModelDeployment) -> list[_CompiledRe Raises cel.CELCompileError on a malformed expression; the caller turns that into an InvalidNodeSelector condition. """ + return _compile_node_selector_requests(deployment.spec.nodeSelector) + + +def _compile_node_selector_requests(node_selector) -> list[_CompiledRequest]: + """Compile device requests from an arbitrary nodeSelector object. + + Used to compile the prefill role's nodeSelector without a full + ModelDeployment. Raises cel.CELCompileError on a malformed expression. + """ requests = [] - for req in deployment.spec.nodeSelector.devices: + for req in node_selector.devices: cel_selectors = [s.cel for s in req.selectors if s.cel] requests.append( _CompiledRequest( @@ -382,6 +404,9 @@ def charge(cluster_name: str, pool_name: str, nodes: int) -> None: if ours and (r.spec.clusterName, _replica_index(r)) not in retained_ids: continue charge(r.spec.clusterName, r.spec.nodePoolName or "", topology_shape(r.spec.workers).nodes_per_replica) + pf = getattr(r.spec, "prefill", None) + if pf and pf.workers: + charge(r.spec.clusterName, pf.nodePoolName or "", topology_shape(pf.workers).nodes_per_replica) return _Ledger(free=free) @@ -391,6 +416,7 @@ def _retain( clusters_by_name: dict[str, icv1alpha1.InferenceCluster], all_replicas: list[mrv1alpha1.ModelReplica], requests: list[_CompiledRequest], + prefill_requests: list[_CompiledRequest] | None = None, ) -> list[Candidate]: """Keep existing replicas whose cluster exists and pool still matches. @@ -399,6 +425,11 @@ def _retain( the fill phase) when its cluster is gone, or when its pinned pool no longer satisfies the nodeSelector - the Kubernetes "template changed, roll the replica" behavior. A degraded-but-present cluster is retained. + + For disaggregated deployments (prefill_requests is not None) BOTH the decode + pool and the prefill pool are checked. A replica whose prefill pinned pool no + longer satisfies the prefill nodeSelector is re-placed, mirroring the same + rolling-replacement semantics applied to the decode pool. """ retained: list[Candidate] = [] seen: set[tuple[str, int]] = set() @@ -412,18 +443,25 @@ def _retain( if identity in seen: continue cluster = clusters_by_name[cluster_name] - if not _pinned_pool_still_matches(r, cluster, requests): + if not _pinned_pool_still_matches(r, cluster, requests, prefill_requests=prefill_requests): continue seen.add(identity) - retained.append( - Candidate( - name=cluster_name, - index=identity[1], - gateway_address=_gateway_address(cluster), - pool=r.spec.nodePoolName or "", - device_requests=_retained_requests(r, cluster, requests), - ) + candidate = Candidate( + name=cluster_name, + index=identity[1], + gateway_address=_gateway_address(cluster), + pool=r.spec.nodePoolName or "", + device_requests=_retained_requests(r, cluster, requests), ) + # For disagg replicas, resolve and populate the retained prefill placement. + if prefill_requests is not None: + pf = getattr(r.spec, "prefill", None) + if pf and pf.nodePoolName: + pf_pool = _pool_by_name(cluster, pf.nodePoolName) + if pf_pool is not None: + candidate.prefill_pool = pf.nodePoolName + candidate.prefill_device_requests = _match_pool(pf_pool, prefill_requests) or [] + retained.append(candidate) return retained @@ -431,8 +469,9 @@ def _pinned_pool_still_matches( replica: mrv1alpha1.ModelReplica, cluster: icv1alpha1.InferenceCluster, requests: list[_CompiledRequest], + prefill_requests: list[_CompiledRequest] | None = None, ) -> bool: - """Whether a retained replica's pinned pool still satisfies the requests. + """Whether a retained replica's pinned pool(s) still satisfy the requests. Modelplane follows Kubernetes here. A change to the deployment's nodeSelector is a change to the deployment "template", so - like editing a Deployment's @@ -443,18 +482,33 @@ def _pinned_pool_still_matches( IgnoredDuringExecution: node-label drift does not evict a bound Pod). Returns False (re-place) when: - * the replica carries no pool pin (it needs a real pool pin), or - * the pinned pool no longer exists on the cluster, or - * the pinned pool no longer satisfies the requests. + * the replica carries no decode pool pin (it needs a real pool pin), or + * the pinned decode pool no longer exists on the cluster, or + * the pinned decode pool no longer satisfies the decode requests. + + For disaggregated deployments (prefill_requests is not None), also returns + False when: + * the replica carries no prefill pool pin, or + * the pinned prefill pool no longer exists on the cluster, or + * the pinned prefill pool no longer satisfies the prefill requests. """ pool_name = replica.spec.nodePoolName if not pool_name: return False pool = _pool_by_name(cluster, pool_name) if pool is None: - # Pinned pool is gone from the cluster's published capacity. + # Pinned decode pool is gone from the cluster's published capacity. return False - return _match_pool(pool, requests) is not None + if _match_pool(pool, requests) is None: + return False + if prefill_requests is not None: + pf = getattr(replica.spec, "prefill", None) + if not pf or not pf.nodePoolName: + return False + pf_pool = _pool_by_name(cluster, pf.nodePoolName) + if pf_pool is None or _match_pool(pf_pool, prefill_requests) is None: + return False + return True def _retained_requests(replica, cluster, requests: list[_CompiledRequest]) -> list[DeviceRequest]: @@ -470,29 +524,59 @@ def _retained_requests(replica, cluster, requests: list[_CompiledRequest]) -> li return _match_pool(pool, requests) or [] -def _eligible_pool( +def _eligible_placement( cluster: icv1alpha1.InferenceCluster, - shape: Shape, - requests: list[_CompiledRequest], + roles: list[tuple[Shape, list[_CompiledRequest]]], ledger: _Ledger, -) -> tuple[str, list[DeviceRequest]] | None: - """Pick the first pool on a cluster that can host one more replica. - - A pool is eligible when it satisfies the nodeSelector requests AND has at - least nodes-per-replica free in the ledger (which already accounts for - replicas placed earlier in this pass). Pools are considered in published - order, which is deterministic. Returns (pool_name, resolved_requests) or - None if no pool on the cluster is eligible. +) -> list[tuple[str, list[DeviceRequest]]] | None: + """Choose a pool per role on this cluster that jointly fits the ledger. + + roles: list of (shape, requests) - one entry for unified, two for disagg + (decode then prefill). Picks pool assignments jointly so a shared pool's + capacity is never double-committed. Returns a list of (pool_name, + resolved_requests) aligned with roles, or None if no assignment fits. + + The pool count per cluster is small, so an exhaustive search is fine. Each + assignment is evaluated cumulatively: two roles assigned to the same pool + must together fit within that pool's available nodes. """ - for pool in cluster.status.gpuPools or []: - name = pool.name or "" - if ledger.available(cluster.metadata.name, name) < shape.nodes_per_replica: - continue - resolved = _match_pool(pool, requests) - if resolved is None: - continue - return name, resolved - return None + pools = cluster.status.gpuPools or [] + cname = cluster.metadata.name + + # For each role, collect the pools that match its selector (ignoring ledger + # here; the search step applies the ledger cumulatively). + per_role_candidates: list[list[tuple[str, list[DeviceRequest]]]] = [] + for _shape, requests in roles: + role_cands: list[tuple[str, list[DeviceRequest]]] = [] + for pool in pools: + resolved = _match_pool(pool, requests) + if resolved is not None: + role_cands.append((pool.name or "", resolved)) + if not role_cands: + return None # No pool satisfies this role at all. + per_role_candidates.append(role_cands) + + def search(k: int, charged: dict[str, int]) -> list[tuple[str, list[DeviceRequest]]] | None: + """Recursively assign pools to roles k..len(roles)-1. + + charged tracks extra nodes already committed to each pool name during + this replica's joint assignment (above whatever the ledger shows), so + two roles sharing one pool see their costs sum against its capacity. + """ + if k == len(roles): + return [] + shape, _ = roles[k] + for name, resolved in per_role_candidates[k]: + already_charged = charged.get(name, 0) + if ledger.available(cname, name) - already_charged >= shape.nodes_per_replica: + charged2 = dict(charged) + charged2[name] = already_charged + shape.nodes_per_replica + rest = search(k + 1, charged2) + if rest is not None: + return [(name, resolved), *rest] + return None + + return search(0, {}) def _fill( @@ -502,6 +586,8 @@ def _fill( ledger: _Ledger, requests: list[_CompiledRequest], n: int, + prefill_requests: list[_CompiledRequest] | None = None, + pf_shape: Shape | None = None, ) -> list[Candidate]: """Place n new replicas, spreading across clusters and packing when forced. @@ -515,6 +601,10 @@ def _fill( lowest free index on its chosen cluster, so the next iteration sees the updated load. Stops early (placing fewer than n) when no cluster can host another replica - the caller surfaces that as InsufficientCapacity. + + For disaggregated deployments, prefill_requests and pf_shape are set. The + decode and prefill pools are then chosen jointly (see _pick_cluster) so a + shared pool's capacity is never double-committed. """ # Per-cluster load and used indices seeded from retained replicas, so spread # accounts for what's already there and new indices don't collide. @@ -526,26 +616,40 @@ def _fill( placed: list[Candidate] = [] for _ in range(n): - choice = _pick_cluster(shape, clusters, load, ledger, requests) + choice = _pick_cluster( + shape, + clusters, + load, + ledger, + requests, + prefill_requests=prefill_requests, + pf_shape=pf_shape, + ) if choice is None: break - cluster, pool_name, resolved = choice + cluster, placement = choice name = cluster.metadata.name index = _lowest_free_index(used_indices.setdefault(name, set())) - placed.append( - Candidate( - name=name, - index=index, - gateway_address=cluster.status.gateway.address, - pool=pool_name, - device_requests=resolved, - ) + decode_pool, decode_resolved = placement[0] + candidate = Candidate( + name=name, + index=index, + gateway_address=cluster.status.gateway.address, + pool=decode_pool, + device_requests=decode_resolved, ) + if len(placement) > 1: + candidate.prefill_pool, candidate.prefill_device_requests = placement[1] + placed.append(candidate) load[name] = load.get(name, 0) + 1 used_indices[name].add(index) - ledger.consume(name, pool_name, shape.nodes_per_replica) + # Consume ledger for decode pool. + ledger.consume(name, decode_pool, shape.nodes_per_replica) + # Consume ledger for prefill pool (may be the same pool as decode). + if len(placement) > 1 and pf_shape is not None: + ledger.consume(name, candidate.prefill_pool, pf_shape.nodes_per_replica) return placed @@ -556,29 +660,39 @@ def _pick_cluster( load: dict[str, int], ledger: _Ledger, requests: list[_CompiledRequest], -) -> tuple[icv1alpha1.InferenceCluster, str, list[DeviceRequest]] | None: + prefill_requests: list[_CompiledRequest] | None = None, + pf_shape: Shape | None = None, +) -> tuple[icv1alpha1.InferenceCluster, list[tuple[str, list[DeviceRequest]]]] | None: """Pick the eligible cluster hosting the fewest of this deployment's replicas. - Eligible means Ready, with a nodeSelector-matching pool that has free - capacity in the ledger. The chosen key is (load on the cluster, cluster - name): fewest replicas first for spread, name for a deterministic tiebreak. - load already counts only this deployment's replicas (seeded from retained - plus those placed earlier in the pass). Returns (cluster, pool_name, - resolved_requests) or None when no cluster is eligible. + Eligible means Ready, with pool(s) jointly satisfying the decode (and + optionally prefill) roles with enough free capacity in the ledger. The + chosen key is (load on the cluster, cluster name): fewest replicas first + for spread, name for a deterministic tiebreak. load already counts only + this deployment's replicas (seeded from retained plus those placed earlier + in the pass). Returns (cluster, placement) where placement is a list of + (pool_name, resolved_requests) aligned with roles, or None when no cluster + is eligible. + + For unified deployments (prefill_requests=None), behaviour is identical to + the old single-role _eligible_pool path: one role, one pool. """ - best = None + roles: list[tuple[Shape, list[_CompiledRequest]]] = [(shape, requests)] + if prefill_requests is not None and pf_shape is not None: + roles.append((pf_shape, prefill_requests)) + + best: tuple[icv1alpha1.InferenceCluster, list[tuple[str, list[DeviceRequest]]]] | None = None best_key = None for cluster in clusters: if not _cluster_ready(cluster): continue - eligible = _eligible_pool(cluster, shape, requests, ledger) - if eligible is None: + placement = _eligible_placement(cluster, roles, ledger) + if placement is None: continue - pool_name, resolved = eligible key = (load.get(cluster.metadata.name, 0), cluster.metadata.name) if best_key is None or key < best_key: best_key = key - best = (cluster, pool_name, resolved) + best = (cluster, placement) return best @@ -630,6 +744,12 @@ def schedule( shortfall by spreading new replicas across clusters (packing onto fewer only when capacity forces it). Returns up to deployment.spec.replicas candidates, fewer if not enough capacity exists. + + For disaggregated deployments (spec.prefill is set), each replica requires + BOTH a decode pool and a prefill pool on the same cluster. The pools are + chosen jointly from one capacity ledger so a shared pool is never + double-committed. Unified deployments (spec.prefill absent) are unaffected: + the single-role path is identical to before. """ desired = int(deployment.spec.replicas) shape = topology_shape(deployment.spec.workers) @@ -640,7 +760,14 @@ def schedule( # expression - the caller turns that into a condition. requests = compile_requests(deployment) - retained = _retain(deployment, clusters_by_name, all_replicas, requests) + # Compile prefill nodeSelector requests when the deployment is disaggregated. + pf_shape = prefill_shape(deployment) + pf_requests: list[_CompiledRequest] | None = None + prefill = getattr(deployment.spec, "prefill", None) + if prefill is not None: + pf_requests = _compile_node_selector_requests(prefill.nodeSelector) + + retained = _retain(deployment, clusters_by_name, all_replicas, requests, prefill_requests=pf_requests) if len(retained) > desired: retained = _scale_down(retained, desired) @@ -653,7 +780,16 @@ def schedule( placed: list[Candidate] = [] if len(retained) < desired: - placed = _fill(shape, clusters, retained, ledger, requests, desired - len(retained)) + placed = _fill( + shape, + clusters, + retained, + ledger, + requests, + desired - len(retained), + prefill_requests=pf_requests, + pf_shape=pf_shape, + ) result = retained + placed result.sort(key=lambda c: (c.name, c.index)) diff --git a/functions/compose-model-deployment/tests/test_fn.py b/functions/compose-model-deployment/tests/test_fn.py index a48b3c428..1affdca77 100644 --- a/functions/compose-model-deployment/tests/test_fn.py +++ b/functions/compose-model-deployment/tests/test_fn.py @@ -954,3 +954,133 @@ async def test_compose(self) -> None: json_format.MessageToDict(got), "-want, +got", ) + + async def test_disagg_prefill_and_routing(self) -> None: + """Disaggregated deployment composes spec.prefill and spec.routing onto the ModelReplica.""" + + # A disaggregated deployment: top-level workers is the decode role; + # spec.prefill carries the prefill role with its own workers + nodeSelector. + # spec.routing is present and should be carried through verbatim. + xr_disagg = v1alpha1.ModelDeployment( + metadata=metav1.ObjectMeta(name="my-model", namespace="ml-team"), + spec=v1alpha1.SpecModel( + replicas=1, + nodeSelector=v1alpha1.NodeSelector( + devices=[ + v1alpha1.Device( + name="gpu", + count=1, + selectors=[v1alpha1.Selector(cel='device.driver == "gpu.nvidia.com"')], + ), + ], + ), + workers=v1alpha1.Workers( + topology=v1alpha1.Topology(tensor=1), + template=v1alpha1.Template( + spec=v1alpha1.Spec( + containers=[ + v1alpha1.Container( + name="engine", + image="vllm/vllm-openai:latest", + args=["--model=Qwen/Qwen3-0.6B"], + ), + ], + ), + ), + ), + prefill=v1alpha1.Prefill( + nodeSelector=v1alpha1.NodeSelector( + devices=[ + v1alpha1.Device( + name="gpu", + count=1, + selectors=[v1alpha1.Selector(cel='device.driver == "gpu.nvidia.com"')], + ), + ], + ), + workers=v1alpha1.Workers( + topology=v1alpha1.Topology(tensor=1), + template=v1alpha1.Template( + spec=v1alpha1.Spec( + containers=[ + v1alpha1.Container( + name="engine", + image="vllm/vllm-openai:latest", + args=["--model=Qwen/Qwen3-0.6B", "--prefill"], + ), + ], + ), + ), + ), + ), + routing=v1alpha1.Routing(), + ), + ).model_dump(exclude_none=True, mode="json") + + # A cluster with TWO GPU pools so the scheduler can assign the decode + # role to "default" and the prefill role to "prefill-pool". + cluster_disagg = icv1alpha1.InferenceCluster( + metadata=metav1.ObjectMeta(name="cluster-a"), + spec=icv1alpha1.Spec( + cluster=icv1alpha1.Cluster( + source="Existing", + existing=icv1alpha1.Existing(secretRef=icv1alpha1.SecretRef(name="k")), + ), + ), + status=icv1alpha1.Status( + conditions=[ + icv1alpha1.Condition( + type="Ready", + status="True", + reason="Available", + lastTransitionTime="2025-01-01T00:00:00Z", + ) + ], + gateway=icv1alpha1.Gateway(address="10.0.0.1"), + providerConfigRef=icv1alpha1.ProviderConfigRef(name="cluster-a"), + gpuPools=[ + icv1alpha1.GpuPool( + name="default", + nodes=1, + devices=[ + icv1alpha1.Device( + name="gpu", + claim="DRA", + driver="gpu.nvidia.com", + deviceClassName="gpu.nvidia.com", + count=1, + ) + ], + ), + icv1alpha1.GpuPool( + name="prefill-pool", + nodes=1, + devices=[ + icv1alpha1.Device( + name="gpu", + claim="DRA", + driver="gpu.nvidia.com", + deviceClassName="gpu.nvidia.com", + count=1, + ) + ], + ), + ], + ), + ).model_dump(exclude_none=True, mode="json") + + req = _req(xr_disagg, clusters=[cluster_disagg]) + got = await self.runner.RunFunction(req, None) + + # Extract the desired ModelReplica from the response. + got_dict = json_format.MessageToDict(got) + replica_resource = got_dict["desired"]["resources"]["replica-cluster-a-0"]["resource"] + + # spec.prefill must be stamped with the scheduler's prefill placement. + self.assertIn("prefill", replica_resource["spec"], "spec.prefill should be set for disagg deployment") + prefill = replica_resource["spec"]["prefill"] + self.assertEqual(prefill["nodePoolName"], "prefill-pool", "prefill nodePoolName should be the prefill pool") + self.assertGreater(len(prefill.get("deviceRequests", [])), 0, "prefill deviceRequests should be non-empty") + + # spec.routing must be carried through from the XR. + self.assertIn("routing", replica_resource["spec"], "spec.routing should be set when XR has routing") diff --git a/functions/compose-model-deployment/tests/test_scheduling.py b/functions/compose-model-deployment/tests/test_scheduling.py index 5034bdc31..aea038b7d 100644 --- a/functions/compose-model-deployment/tests/test_scheduling.py +++ b/functions/compose-model-deployment/tests/test_scheduling.py @@ -909,5 +909,315 @@ def test_invalid_cel_raises(self) -> None: scheduling.schedule(deployment, [_cluster("cluster-a", pools=[_pool("frontier")])], []) +class TestCandidate(unittest.TestCase): + """Tests for Candidate dataclass fields.""" + + def test_prefill_placement_fields_default_empty(self): + """A Candidate carries an optional prefill placement, empty by default.""" + c = scheduling.Candidate(name="c1", index=0) + self.assertEqual(c.prefill_pool, "") + self.assertEqual(c.prefill_device_requests, []) + + +def _dra_nic_device(*, link_type: str = "infiniband", count: int = 1) -> dict: + """A DRA NIC device dict for a pool, satisfying the _IB selector.""" + return { + "name": "nic", + "claim": "DRA", + "driver": "nic.nvidia.com", + "deviceClassName": "nic.nvidia.com", + "count": count, + "attributes": {"linkType": {"string": link_type}}, + } + + +def _disagg_deployment( + decode_req: mdv1alpha1.Device, + prefill_req: mdv1alpha1.Device, + replicas: int = 1, + decode_count: int = 1, + prefill_count: int = 1, +) -> mdv1alpha1.ModelDeployment: + """A ModelDeployment with both decode (top-level workers) and prefill roles.""" + d = _deployment(replicas=replicas, count=decode_count, requests=[decode_req]) + d.spec.prefill = mdv1alpha1.Prefill( + workers=mdv1alpha1.Workers( + count=prefill_count, + topology=mdv1alpha1.Topology(tensor=1, pipeline=1), + template=mdv1alpha1.Template( + spec=mdv1alpha1.Spec( + containers=[mdv1alpha1.Container(name="engine", image="vllm/vllm-openai:latest")], + ), + ), + ), + nodeSelector=mdv1alpha1.NodeSelector(devices=[prefill_req]), + ) + return d + + +class TestScheduleDisagg(unittest.TestCase): + """Tests for joint decode+prefill placement in disaggregated deployments. + + A disagg replica = decode placement + prefill placement, BOTH on the same + cluster, picked jointly from one ledger so capacity is never double-committed. + """ + + def test_decode_and_prefill_different_pools_one_cluster(self): + """Decode and prefill placed on different pools of the same cluster. + + Cluster has two pools: 'gpu-big' (satisfies _MEM_141, decode) and + 'gpu-nic' (satisfies _IB, prefill). The scheduler must find the pair and + record both on one Candidate. + """ + decode_req = _request(name="gpu", cel_exprs=[_MEM_141]) + prefill_req = _request(name="nic", cel_exprs=[_IB]) + deployment = _disagg_deployment(decode_req, prefill_req) + + # gpu-big: DRA GPU device satisfying _MEM_141 (141 GiB memory). + # gpu-nic: DRA NIC device satisfying _IB (infiniband link type). + cluster = _cluster( + "cluster-a", + pools=[ + _pool("gpu-big", nodes=2, devices=[_gpu_device(memory="141Gi")]), + _pool("gpu-nic", nodes=2, devices=[_dra_nic_device(link_type="infiniband")]), + ], + ) + + got = scheduling.schedule(deployment, [cluster], all_replicas=[]) + + self.assertEqual(len(got), 1, f"want 1 candidate, got {len(got)}: {got}") + c = got[0] + self.assertEqual(c.pool, "gpu-big", "decode pool should be gpu-big") + self.assertEqual(c.prefill_pool, "gpu-nic", "prefill pool should be gpu-nic") + self.assertTrue(len(c.prefill_device_requests) > 0, "prefill_device_requests must be non-empty") + + def test_shared_pool_capacity_sums(self): + """When decode and prefill share one pool, their node costs sum against that pool. + + One cluster, one pool ('shared') with 1 node. Decode needs 1 node, + prefill needs 1 node → 2 total > 1 available → no Candidate. + With 2 nodes in the pool, one Candidate is produced with + pool == prefill_pool == 'shared'. + """ + decode_req = _request(name="gpu", cel_exprs=[_MEM_141]) + # Prefill selector also uses _MEM_141 so it matches the same GPU pool. + prefill_req = _request(name="gpu", cel_exprs=[_MEM_141]) + deployment = _disagg_deployment(decode_req, prefill_req) + + # 1-node pool: decode(1) + prefill(1) = 2 > 1, should not fit. + cluster_1 = _cluster( + "cluster-a", + pools=[_pool("shared", nodes=1, devices=[_gpu_device(memory="141Gi")])], + ) + got = scheduling.schedule(deployment, [cluster_1], all_replicas=[]) + self.assertEqual(len(got), 0, f"1-node pool should not fit both roles; got {got}") + + # 2-node pool: decode(1) + prefill(1) = 2 == 2, should fit. + cluster_2 = _cluster( + "cluster-a", + pools=[_pool("shared", nodes=2, devices=[_gpu_device(memory="141Gi")])], + ) + got2 = scheduling.schedule(deployment, [cluster_2], all_replicas=[]) + self.assertEqual(len(got2), 1, f"2-node pool should fit both roles; got {got2}") + c = got2[0] + self.assertEqual(c.pool, "shared") + self.assertEqual(c.prefill_pool, "shared") + + def test_no_feasible_pair_rejects_cluster(self): + """A cluster where no pool satisfies the prefill selector yields no Candidates. + + Cluster has only one pool matching the decode selector. The prefill + selector (_IB) has no matching pool, so the whole cluster is ineligible. + """ + decode_req = _request(name="gpu", cel_exprs=[_MEM_141]) + prefill_req = _request(name="nic", cel_exprs=[_IB]) + deployment = _disagg_deployment(decode_req, prefill_req) + + # Only gpu-big: satisfies decode but NOT prefill (_IB requires NIC). + cluster = _cluster( + "cluster-a", + pools=[_pool("gpu-big", nodes=2, devices=[_gpu_device(memory="141Gi")])], + ) + + got = scheduling.schedule(deployment, [cluster], all_replicas=[]) + self.assertEqual(len(got), 0, f"cluster with no prefill pool must yield 0 Candidates; got {got}") + + def test_ledger_charges_existing_replica_prefill_pool(self): + """An existing disagg replica consumes its prefill pool's nodes too. + + One cluster, two pools: + 'gpu-big' (decode, _MEM_141) — 2 nodes (enough for two replicas' decodes) + 'gpu-nic' (prefill, _IB) — 1 node (only enough for ONE prefill) + + An existing disagg ModelReplica already occupies 1 node of each pool. + A deployment asking for replicas=2 retains the existing replica but must + NOT place a second one — the gpu-nic prefill pool is fully consumed. + + Without the fix, _build_ledger only charges gpu-big (the decode pool), + leaving gpu-nic with 1 free node. The fill phase then sees gpu-big free + (2 total - 1 used = 1) and gpu-nic free (1 total - 0 wrongly charged = 1) + and incorrectly places a second replica → 2 Candidates returned. + + With the fix, gpu-nic is also charged (1 total - 1 used = 0 free) and + the fill phase correctly finds no room for a second prefill → 1 Candidate. + """ + decode_req = _request(name="gpu", cel_exprs=[_MEM_141]) + prefill_req = _request(name="nic", cel_exprs=[_IB]) + deployment = _disagg_deployment(decode_req, prefill_req, replicas=2) + + # gpu-big has 2 nodes (decode pool — enough for 2 replicas). + # gpu-nic has 1 node (prefill pool — only enough for 1 replica). + cluster = _cluster( + "cluster-a", + pools=[ + _pool("gpu-big", nodes=2, devices=[_gpu_device(memory="141Gi")]), + _pool("gpu-nic", nodes=1, devices=[_dra_nic_device(link_type="infiniband")]), + ], + ) + + # Build an existing disagg replica that belongs to this deployment and + # occupies 1 node from gpu-big (decode) and 1 node from gpu-nic (prefill). + existing = _replica_with_pool("my-model", "cluster-a", pool="gpu-big", index=0) + existing.spec.prefill = mrv1alpha1.Prefill( + workers=mrv1alpha1.Workers( + count=1, + topology=mrv1alpha1.Topology(tensor=1, pipeline=1), + template=mrv1alpha1.Template( + spec=mrv1alpha1.Spec( + containers=[mrv1alpha1.Container(name="engine", image="vllm/vllm-openai:latest")], + ), + ), + ), + nodePoolName="gpu-nic", + deviceRequests=[ + mrv1alpha1.DeviceRequest( + name="nic", + deviceClassName="nic.nvidia.com", + count=1, + selectors=[mrv1alpha1.Selector(cel=_IB)], + ), + ], + ) + + got = scheduling.schedule(deployment, [cluster], all_replicas=[existing]) + + # The existing replica is retained (index 0). The second cannot be placed + # because gpu-nic (1 node total, 1 consumed) has 0 free nodes for prefill. + # Without the fix this returns 2 Candidates (prefill pool not charged). + self.assertEqual( + len(got), + 1, + f"prefill pool must be charged for existing disagg replicas; got {len(got)} candidates: {got}", + ) + self.assertEqual(got[0].name, "cluster-a") + self.assertEqual(got[0].index, 0) + + def test_retain_drops_when_prefill_pool_stops_matching(self): + """A disagg replica is re-placed when its prefill pool no longer satisfies the selector. + + Two arms: + 1. Prefill selector still matches 'gpu-nic' → replica IS retained (1 Candidate, + pool='gpu-big', prefill_pool='gpu-nic'). + 2. Prefill selector changed to _MEM_200 which 'gpu-nic' (a NIC pool, no GPU memory) + does NOT satisfy → replica is NOT retained; with no eligible pool for the new + prefill selector, 0 Candidates are returned. + + Before the fix, _pinned_pool_still_matches only checks the decode pool, so arm 2 + wrongly retains the replica (returns 1 Candidate) instead of 0. + """ + decode_req = _request(name="gpu", cel_exprs=[_MEM_141]) + + # ARM 1: prefill selector still satisfied by 'gpu-nic' → retained. + prefill_req_ok = _request(name="nic", cel_exprs=[_IB]) + deployment_ok = _disagg_deployment(decode_req, prefill_req_ok) + + # ARM 2: prefill selector now requires _MEM_200 — 'gpu-nic' is a NIC pool with no + # GPU memory attribute, so it fails the new selector → re-placed, no eligible pool. + prefill_req_drifted = _request(name="gpu2", cel_exprs=[_MEM_200]) + deployment_drifted = _disagg_deployment(decode_req, prefill_req_drifted) + + cluster = _cluster( + "cluster-a", + pools=[ + _pool("gpu-big", nodes=2, devices=[_gpu_device(memory="141Gi")]), + _pool("gpu-nic", nodes=2, devices=[_dra_nic_device(link_type="infiniband")]), + ], + ) + + # Build an existing disagg replica pinned to decode='gpu-big', prefill='gpu-nic'. + existing = _replica_with_pool("my-model", "cluster-a", pool="gpu-big", index=0) + existing.spec.prefill = mrv1alpha1.Prefill( + workers=mrv1alpha1.Workers( + count=1, + topology=mrv1alpha1.Topology(tensor=1, pipeline=1), + template=mrv1alpha1.Template( + spec=mrv1alpha1.Spec( + containers=[mrv1alpha1.Container(name="engine", image="vllm/vllm-openai:latest")], + ), + ), + ), + nodePoolName="gpu-nic", + deviceRequests=[ + mrv1alpha1.DeviceRequest( + name="nic", + deviceClassName="nic.nvidia.com", + count=1, + selectors=[mrv1alpha1.Selector(cel=_IB)], + ), + ], + ) + + # ARM 1: prefill selector unchanged → replica must be retained. + got_ok = scheduling.schedule(deployment_ok, [cluster], all_replicas=[existing]) + self.assertEqual( + len(got_ok), + 1, + f"arm1 (prefill still matches): want 1 retained Candidate, got {len(got_ok)}: {got_ok}", + ) + self.assertEqual(got_ok[0].pool, "gpu-big", "arm1: decode pool must be gpu-big") + self.assertEqual(got_ok[0].prefill_pool, "gpu-nic", "arm1: prefill pool must be gpu-nic") + + # ARM 2: prefill selector drifted → replica must NOT be retained; no eligible + # pool for the new prefill selector → 0 Candidates. + got_drifted = scheduling.schedule(deployment_drifted, [cluster], all_replicas=[existing]) + self.assertEqual( + len(got_drifted), + 0, + f"arm2 (prefill pool drifted): want 0 Candidates (re-place fails), got {len(got_drifted)}: {got_drifted}", + ) + + def test_decode_only_cluster_skipped_for_one_with_both(self): + """The scheduler must not greedily commit to the decode-cheapest cluster. + + cluster-a (alphabetically first, so it wins the (load, name) tiebreak) + has a decode pool but NO prefill pool. cluster-b has both. A greedy + per-role placement would pick cluster-a for decode then fail to find a + prefill pool there; the joint pair selection must skip cluster-a entirely + and place the whole replica on cluster-b. + """ + decode_req = _request(name="gpu", cel_exprs=[_MEM_141]) + prefill_req = _request(name="nic", cel_exprs=[_IB]) + deployment = _disagg_deployment(decode_req, prefill_req) + + cluster_a = _cluster( + "cluster-a", + pools=[_pool("gpu-only", nodes=2, devices=[_gpu_device(memory="141Gi")])], + ) + cluster_b = _cluster( + "cluster-b", + pools=[ + _pool("gpu-big", nodes=2, devices=[_gpu_device(memory="141Gi")]), + _pool("gpu-nic", nodes=2, devices=[_dra_nic_device(link_type="infiniband")]), + ], + ) + + got = scheduling.schedule(deployment, [cluster_a, cluster_b], all_replicas=[]) + + self.assertEqual(len(got), 1, f"want 1 candidate on cluster-b; got {got}") + self.assertEqual(got[0].name, "cluster-b", "must skip the decode-only cluster-a") + self.assertEqual(got[0].pool, "gpu-big") + self.assertEqual(got[0].prefill_pool, "gpu-nic") + + if __name__ == "__main__": unittest.main() diff --git a/functions/compose-model-replica/function/backends/base.py b/functions/compose-model-replica/function/backends/base.py index cfde409ad..7e3596595 100644 --- a/functions/compose-model-replica/function/backends/base.py +++ b/functions/compose-model-replica/function/backends/base.py @@ -96,6 +96,18 @@ def apply_cache_args(args: list[str], replica: v1alpha1.ModelReplica, engine) -> # Response resource key for the DRA ResourceClaimTemplate. RESOURCE_CLAIM_KEY = "resource-claim" +# Label that distinguishes prefill vs decode pods in a disaggregated replica. +# The decode Service selector includes {LABEL_PD_ROLE: "decode"} so it never +# accidentally routes traffic to prefill pods. +LABEL_PD_ROLE = "modelplane.ai/pd-role" + +# llm-d-native labels consumed by the GAIE EPP for disaggregated replicas. +# LABEL_LLMD_ROLE carries "decode" or "prefill" so the EPP's per-role filters +# can select the right pod set. LABEL_LLMD_SERVING is the shared selector that +# tells the EPP these pods belong to an llm-d inference-serving workload. +LABEL_LLMD_ROLE = "llm-d.ai/role" +LABEL_LLMD_SERVING = "llm-d.ai/inference-serving" + # DRA API the ResourceClaimTemplate targets. The manifest is a raw dict wrapped # in a provider-kubernetes Object, so no generated model is needed. _DRA_API_VERSION = "resource.k8s.io/v1" @@ -171,10 +183,17 @@ def nodes_per_worker(replica: v1alpha1.ModelReplica) -> int: def needs_cross_pod_coordination(replica: v1alpha1.ModelReplica) -> bool: """True when the replica is more than one self-contained pod. - v0.1: true iff nodes_per_worker > 1. Extension points (no-ops until the - fields exist): a `prefill` block (disaggregated P/D) or multi-node data - parallelism (data > dataLocal) also make this true. + A replica needs cross-pod coordination when: + - nodes_per_worker > 1 (pipeline parallelism spanning multiple nodes), or + - spec.prefill is set (disaggregated P/D: the prefill and decode roles are + separate pods that must coordinate over NIXL regardless of per-role + topology, even when each role is a single node). + + Extension point: multi-node data parallelism (data > dataLocal) will also + make this true when that field lands. """ + if getattr(replica.spec, "prefill", None): + return True return nodes_per_worker(replica) > 1 @@ -190,7 +209,7 @@ def select_backend(replica: v1alpha1.ModelReplica) -> str: def claim_template_name(replica: v1alpha1.ModelReplica) -> str: - """ResourceClaimTemplate name on the remote cluster. + """ResourceClaimTemplate name on the remote cluster (decode / unified). Per-replica, derived from the replica's own name so concurrent replicas of the same deployment on one cluster don't collide. @@ -198,6 +217,20 @@ def claim_template_name(replica: v1alpha1.ModelReplica) -> str: return resource.child_name(replica.metadata.name, _POD_CLAIM_NAME) +def claim_template_name_for(replica: v1alpha1.ModelReplica, role: str) -> str: + """Per-role ResourceClaimTemplate name on the remote cluster. + + For the ``"decode"`` (or unified) role this returns the same value as + ``claim_template_name`` so the unified path is unchanged. For the + ``"prefill"`` role a ``-prefill`` suffix is appended to keep the two + templates distinct on the workload cluster. + """ + base_name = claim_template_name(replica) + if role == "prefill": + return f"{base_name}-prefill" + return base_name + + def engine_resources() -> dict: """Container resources for the engine. @@ -232,6 +265,19 @@ def engine_resources() -> dict: _LABEL_POOL = "modelplane.ai/pool" +def place_pod_on(pod_spec: dict, node_pool: str, claim_tmpl_name: str) -> None: + """Pin a pod spec to a specific node pool and DRA ResourceClaimTemplate. + + Low-level helper used by both ``place_pod`` (unified/decode path) and the + prefill path which supplies a distinct pool name and claim-template name. + Callers are responsible for passing the correct pool and template name for + the role being built. + """ + pod_spec["nodeSelector"] = {_LABEL_POOL: node_pool} + pod_spec["resourceClaims"] = [{"name": _POD_CLAIM_NAME, "resourceClaimTemplateName": claim_tmpl_name}] + pod_spec.setdefault("tolerations", []).append(_GPU_TOLERATION) + + def place_pod(pod_spec: dict, replica: v1alpha1.ModelReplica) -> None: """Constrain a serving pod to the placement the scheduler chose. @@ -251,10 +297,11 @@ def place_pod(pod_spec: dict, replica: v1alpha1.ModelReplica) -> None: carries device requests (the XRD requires them), so every serving pod claims through DRA. A template-backed claim (not a shared ResourceClaim) gives each pod in a gang its own claim. + + Delegates to ``place_pod_on`` with the decode/unified pool and claim-template + name, leaving the unified and decode paths identical to before. """ - pod_spec["nodeSelector"] = {_LABEL_POOL: replica.spec.nodePoolName} - pod_spec["resourceClaims"] = [{"name": _POD_CLAIM_NAME, "resourceClaimTemplateName": claim_template_name(replica)}] - pod_spec.setdefault("tolerations", []).append(_GPU_TOLERATION) + place_pod_on(pod_spec, replica.spec.nodePoolName, claim_template_name(replica)) def resource_claim_template(replica: v1alpha1.ModelReplica, provider_config: str) -> k8sobjv1alpha1.Object: @@ -264,26 +311,93 @@ def resource_claim_template(replica: v1alpha1.ModelReplica, provider_config: str matched InferenceClass claim: DRA devices) becomes one DeviceRequest carrying its DeviceClass, count, and CEL selectors verbatim. Every replica carries at least one device request (the XRD requires them). + + Delegates to ``resource_claim_template_for`` with the ``"decode"`` role and + the replica's own device requests so the unified path is unchanged. + """ + return resource_claim_template_for(replica, provider_config, "decode", replica.spec.deviceRequests) + + +def resource_claim_template_for( + replica: v1alpha1.ModelReplica, + provider_config: str, + role: str, + device_reqs, +) -> k8sobjv1alpha1.Object: + """Compose a DRA ResourceClaimTemplate Object for a specific role. + + ``role`` is either ``"decode"`` (or unified) or ``"prefill"``; it determines + the template name via ``claim_template_name_for`` so the two templates on the + workload cluster have distinct names and don't collide. + + ``device_reqs`` is the list of ``DeviceRequest`` objects for this role's pool + (``replica.spec.deviceRequests`` for decode/unified; + ``replica.spec.prefill.deviceRequests`` for prefill). """ - device_requests = [] - for r in replica.spec.deviceRequests: + tmpl_name = claim_template_name_for(replica, role) + raw_requests = [] + for r in device_reqs or []: exactly: dict = {"deviceClassName": r.deviceClassName, "count": int(r.count or 1)} selectors = [{"cel": {"expression": s.cel}} for s in (r.selectors or []) if s.cel] if selectors: exactly["selectors"] = selectors - device_requests.append({"name": r.name, "exactly": exactly}) + raw_requests.append({"name": r.name, "exactly": exactly}) return wrap_object( provider_config, { "apiVersion": _DRA_API_VERSION, "kind": "ResourceClaimTemplate", - "metadata": {"name": claim_template_name(replica), "namespace": REMOTE_NAMESPACE}, - "spec": {"spec": {"devices": {"requests": device_requests}}}, + "metadata": {"name": tmpl_name, "namespace": REMOTE_NAMESPACE}, + "spec": {"spec": {"devices": {"requests": raw_requests}}}, }, ) +def nixl_side_channel_env() -> dict: + # vLLM NixlConnector needs the pod's own IP as the NIXL side-channel host; it + # can't come from user args, so disaggregated pods get it via the downward API. + return {"name": "VLLM_NIXL_SIDE_CHANNEL_HOST", "valueFrom": {"fieldRef": {"fieldPath": "status.podIP"}}} + + +# Image for the llm-d prefill/decode sidecar that intercepts the serving port on +# disaggregated decode pods. The sidecar reads the x-prefiller-host-port header +# injected by the EPP, forwards the prompt to the prefill instance for KV +# transfer, then proxies to the local vLLM (which has moved to ENGINE_PORT+1). +PD_SIDECAR_IMAGE = "ghcr.io/llm-d/llm-d-routing-sidecar:v0.8.0" + +# Port vLLM moves to on disaggregated decode pods so the pd-sidecar can take the +# external serving port (ENGINE_PORT = 8000). This is internal to the pod; the +# Service always targets ENGINE_PORT (8000) which is the sidecar's listen port. +_DECODE_ENGINE_PORT = ENGINE_PORT + 1 # 8001 + + +def pd_sidecar_container() -> dict: + """Return the pd-sidecar container dict for a disaggregated decode leader pod. + + The sidecar takes the external serving port (ENGINE_PORT = 8000). The local + vLLM engine moves to _DECODE_ENGINE_PORT (8001). The sidecar's readiness + probe mirrors the engine's (GET /health) so the leader pod only becomes ready + once the sidecar is up and accepting traffic. + + --secure-proxy is disabled: the routing sidecar defaults to serving HTTPS on + the serving port, but the whole Modelplane serving path is plain HTTP (the + gateway, the Service, and this readiness probe all speak HTTP), so a TLS + listener here fails the probe and rejects the gateway's HTTP traffic. + """ + return { + "name": "pd-sidecar", + "image": PD_SIDECAR_IMAGE, + "ports": [{"containerPort": ENGINE_PORT}], + "args": ["--secure-proxy=false", "--kv-connector=nixlv2", f"--vllm-port={_DECODE_ENGINE_PORT}"], + "readinessProbe": { + "httpGet": {"path": "/health", "port": ENGINE_PORT}, + "initialDelaySeconds": 30, + "periodSeconds": 10, + }, + } + + class Backend(Protocol): """Builds the cluster-level serving resources for one ModelReplica.""" diff --git a/functions/compose-model-replica/function/backends/llmd.py b/functions/compose-model-replica/function/backends/llmd.py index aed79d7eb..e1d62f672 100644 --- a/functions/compose-model-replica/function/backends/llmd.py +++ b/functions/compose-model-replica/function/backends/llmd.py @@ -1,22 +1,25 @@ """llm-d multi-pod backend: LeaderWorkerSet + Service + HTTPRoute. -Selected only for multi-node replicas (pipeline > 1), so this always renders a -LeaderWorkerSet whose gang size is the per-worker node count. - -Routing is plain Gateway API — `HTTPRoute -> Service`, exactly like native.py — -NOT a GAIE `InferencePool`. The HTTPRoute attaches to the *workload* cluster's -inference gateway (Envoy Gateway, named `inference-gateway`, installed by -ServingStack) and the Service selects the LWS *leader* pods (only the leader -serves the OpenAI API; workers just join the gang). - -Why a Service, not a GAIE `InferencePool`: v0.1 does no KV-/load-aware endpoint -picking, so the `InferencePool` + EPP this path originally emitted aren't needed -yet. Reintroducing them is a *workload-gateway* concern — it needs a -GAIE-conformant workload gateway (Envoy Gateway's `InferencePool` v1 support is -unconfirmed; alternatively switch the workload gateway to Istio/agentgateway). -That is independent of the control-plane gateway (Traefik, named `modelplane`), -which never sees these resources. (Issue #8 — inference-aware routing *across -replicas* on the control plane — is a separate problem at that layer.) +Selected for replicas that need cross-pod coordination: multi-node replicas +(pipeline > 1) and disaggregated replicas (a prefill block, even at pipeline 1). +Renders a LeaderWorkerSet whose gang size is the per-worker node count; a +disaggregated replica additionally emits a separate internal prefill pod set +(see _prefill_objects). + +Routing differs by path. UNIFIED serving uses plain Gateway API — +`HTTPRoute -> Service`, exactly like native.py — where the Service selects the +LWS *leader* pods (only the leader serves the OpenAI API; workers just join the +gang). DISAGGREGATED serving (a `prefill` block) instead fronts the pods with a +GAIE `InferencePool` + endpoint-picker (EPP) so requests are sequenced +prefill->decode: it emits the InferencePool (`_inference_pool_object`), the EPP +stack from `routing.template` (`_epp_objects`), a pd-sidecar on the decode pods, +and points the HTTPRoute at the InferencePool. The InferencePool path needs a +GAIE-conformant workload gateway — Envoy AI Gateway, installed by ServingStack. + +Both attach to the workload cluster's inference gateway, independent of the +control-plane gateway (Traefik, named `modelplane`), which never sees these +resources. (Issue #8 — inference-aware routing *across replicas* on the control +plane — is a separate problem at that layer.) Multi-node bootstrap: the LWS leader and worker run different commands (no `LWS_WORKER_INDEX` branch). The leader starts the Ray head then execs the @@ -70,6 +73,491 @@ def _engine_args(engine, tensor: int, pipeline: int) -> list[str]: return args +# vLLM NixlConnector config enabling disaggregated KV transfer. NixlConnector does +# not distinguish kv_role (the routing sidecar drives the prefill->decode direction +# at request time), so both roles run kv_both; see the vLLM NixlConnector usage +# guide. Without this the engines run as plain servers and no KV handoff occurs. +_NIXL_KV_TRANSFER_CONFIG = '{"kv_connector":"NixlConnector","kv_role":"kv_both"}' + + +def _with_nixl_kv_transfer(args: list[str]) -> list[str]: + """Append the NixlConnector --kv-transfer-config unless the user already set one.""" + if any(a.startswith("--kv-transfer-config") for a in args): + return args + return [*args, f"--kv-transfer-config={_NIXL_KV_TRANSFER_CONFIG}"] + + +def _build_commands( + engine, + tensor: int, + pipeline: int, + replica, + *, + disagg: bool, +) -> tuple[list[str] | None, list[str], list[str], list[str]]: + """Compute (user_command, leader_command, worker_command, args) for the decode LWS. + + user_command is None on the turnkey vLLM path (the caller uses it to decide + whether to emit a separate container args field); leader_command and + worker_command are the final commands for the two pod templates; args is the + vLLM arg list the caller sets on the container. + """ + user_cmd = list(engine.command) if engine.command else None + if user_cmd: + return user_cmd, user_cmd, user_cmd, list(engine.args or []) + args = _engine_args(engine, tensor, pipeline) + args = base.apply_cache_args(args, replica, engine) + if disagg and not any(a.startswith("--port=") for a in args): + args = [*args, f"--port={base._DECODE_ENGINE_PORT}"] + if disagg: + args = _with_nixl_kv_transfer(args) + leader_cmd = ["/bin/sh", "-c", _LEADER_BOOTSTRAP, "vllm", *args] + worker_cmd = ["/bin/sh", "-c", _WORKER_BOOTSTRAP] + return None, leader_cmd, worker_cmd, args + + +def _inference_pool_object( + name: str, + provider_config: str, +) -> k8sobjv1alpha1.Object: + """Build a GAIE InferencePool for a disaggregated decode path. + + The pool selects both decode and prefill pods via the shared llm-d.ai labels + (both roles carry app: + llm-d.ai/inference-serving:"true"). The EPP + partitions them by llm-d.ai/role and is referenced via endpointPickerRef. + The HTTPRoute for the disagg path points at this pool instead of the Service, + so the GAIE EPP can apply KV-aware routing. + """ + manifest = { + "apiVersion": "inference.networking.k8s.io/v1", + "kind": "InferencePool", + "metadata": {"name": f"{name}-pool", "namespace": base.REMOTE_NAMESPACE}, + "spec": { + "selector": { + "matchLabels": { + "app": name, + base.LABEL_LLMD_SERVING: "true", + }, + }, + "targetPorts": [{"number": base.ENGINE_PORT}], + "endpointPickerRef": { + "name": f"{name}-epp", + "port": {"number": 9002}, + # failureMode belongs to endpointPickerRef (EndpointPickerRef) in + # the GAIE v1 schema, not to the pool spec. FailOpen so a transient + # EPP outage doesn't black-hole decode traffic. + "failureMode": "FailOpen", + }, + }, + } + return base.wrap_object(provider_config, manifest) + + +def _prefill_objects( + replica: v1alpha1.ModelReplica, + prefill_spec, + name: str, + provider_config: str, + cache_volumes: list[dict], + cache_volume_mounts: list[dict], +) -> dict[str, k8sobjv1alpha1.Object]: + """Build the prefill pod set + ResourceClaimTemplate for a disaggregated replica. + + Returns the response entries for the internal prefill role: a LeaderWorkerSet + (no Service/HTTPRoute — prefill is not an API entrypoint) and a per-role + ResourceClaimTemplate. Pods carry pd-role:prefill, mount the model cache like + decode, and get the NIXL side-channel env. + """ + prefill_name = f"{name}-prefill" + prefill_claim = base.claim_template_name_for(replica, "prefill") + p_engine = next(c for c in prefill_spec.workers.template.spec.containers if c.name == "engine") + p_tensor = int(prefill_spec.workers.topology.tensor) + p_pipeline = int(prefill_spec.workers.topology.pipeline or 1) + p_size = p_pipeline + + p_user_cmd = list(p_engine.command) if p_engine.command else None + if p_user_cmd: + p_leader_cmd = p_worker_cmd = p_user_cmd + p_args = list(p_engine.args or []) + else: + p_args = _engine_args(p_engine, p_tensor, p_pipeline) + p_args = base.apply_cache_args(p_args, replica, p_engine) + p_args = _with_nixl_kv_transfer(p_args) + p_leader_cmd = ["/bin/sh", "-c", _LEADER_BOOTSTRAP, "vllm", *p_args] + p_worker_cmd = ["/bin/sh", "-c", _WORKER_BOOTSTRAP] + + p_env = [e.model_dump(exclude_none=True) for e in p_engine.env] if p_engine.env else None + # Prefill pods also need the NIXL side-channel env (same reason as decode). + p_env = (p_env or []) + [base.nixl_side_channel_env()] + p_tmpl = prefill_spec.workers.template + p_pull_secrets = ( + [s.model_dump(exclude_none=True) for s in p_tmpl.spec.imagePullSecrets] + if p_tmpl.spec.imagePullSecrets + else None + ) + + def p_container(command: list[str]) -> dict: + c = { + "name": "engine", + "image": p_engine.image, + "resources": base.engine_resources(), + # Mounts the model cache like decode: prefill loads the same + # weights to compute the KV it transfers to decode. + "volumeMounts": [{"name": "dshm", "mountPath": "/dev/shm"}, *cache_volume_mounts], + "command": command, + } + if p_user_cmd and p_args: + c["args"] = p_args + if p_env: + c["env"] = p_env + return c + + def p_pod_spec(c: dict) -> dict: + spec = { + "containers": [c], + "volumes": [{"name": "dshm", "emptyDir": {"medium": "Memory"}}, *cache_volumes], + } + base.place_pod_on(spec, prefill_spec.nodePoolName, prefill_claim) + if p_pull_secrets: + spec["imagePullSecrets"] = p_pull_secrets + return spec + + prefill_labels = { + base.LABEL_SERVING: prefill_name, + base.LABEL_PD_ROLE: "prefill", + base.LABEL_LLMD_ROLE: "prefill", + base.LABEL_LLMD_SERVING: "true", + "app": name, + } + p_leader_pod = { + "metadata": {"labels": {**prefill_labels, _LABEL_ROLE: "leader"}}, + "spec": p_pod_spec(p_container(p_leader_cmd)), + } + p_worker_pod = { + "metadata": {"labels": prefill_labels}, + "spec": p_pod_spec(p_container(p_worker_cmd)), + } + + prefill_lws = { + "apiVersion": "leaderworkerset.x-k8s.io/v1", + "kind": "LeaderWorkerSet", + "metadata": {"name": prefill_name, "namespace": base.REMOTE_NAMESPACE}, + "spec": { + "replicas": int(prefill_spec.workers.count or 1), + "leaderWorkerTemplate": { + "size": p_size, + "leaderTemplate": p_leader_pod, + "workerTemplate": p_worker_pod, + }, + }, + } + + return { + "prefill-serving": base.wrap_object(provider_config, prefill_lws, cel_query=base.AVAILABLE_CEL), + "prefill-resource-claim": base.resource_claim_template_for( + replica, provider_config, "prefill", prefill_spec.deviceRequests + ), + } + + +# The EndpointPickerConfig YAML content for the disaggregated profile. +# This is the authoritative upstream config from deploy/config/pd-epp-config.yaml +# (llm-d/llm-d-inference-scheduler), using approx-prefix-cache-producer, +# prefix-based-pd-decider, disagg-profile-handler, prefill/decode profiles. +_PD_EPP_CONFIG_YAML = """\ +apiVersion: inference.networking.x-k8s.io/v1alpha1 +kind: EndpointPickerConfig +plugins: +- type: approx-prefix-cache-producer + parameters: + maxPrefixBlocksToMatch: 256 + lruCapacityPerServer: 31250 +- type: prefix-cache-scorer +- type: queue-scorer +- type: prefill-filter +- type: decode-filter +- type: max-score-picker +- type: prefix-based-pd-decider + parameters: + nonCachedTokens: 16 +- type: disagg-profile-handler + parameters: + deciders: + prefill: prefix-based-pd-decider +schedulingProfiles: +- name: prefill + plugins: + - pluginRef: prefill-filter + - pluginRef: max-score-picker + - pluginRef: prefix-cache-scorer + weight: 2 + - pluginRef: queue-scorer + weight: 1 +- name: decode + plugins: + - pluginRef: decode-filter + - pluginRef: max-score-picker + - pluginRef: prefix-cache-scorer + weight: 2 + - pluginRef: queue-scorer + weight: 1 +""" + +# Default EPP image when routing.template does not supply one. +_EPP_DEFAULT_IMAGE = "ghcr.io/llm-d/llm-d-inference-scheduler:v0.8.0" + +# Name suffix for the EPP Role/RoleBinding (combines pod-watch + inference CRD watch). +_EPP_ROLE_SUFFIX = "epp-sa" + +# ClusterRole/ClusterRoleBinding name suffix for metrics auth reviewer. +_EPP_AUTH_SUFFIX = "epp-auth-reviewer" + + +def _route_backend_refs(name: str, *, disagg: bool) -> list[dict]: + """Return the HTTPRoute backendRefs for disaggregated (InferencePool) or unified (Service) paths.""" + if disagg: + return [{"group": "inference.networking.k8s.io", "kind": "InferencePool", "name": f"{name}-pool"}] + return [{"name": name, "port": 80}] + + +def _epp_container_from_routing(replica) -> dict | None: + """Return the "epp" container from spec.routing.template, or None. + + routing.template.spec is dict[str, Any]; iterate containers to find name==epp. + None means routing is absent or has no epp container, so the caller falls back + to _EPP_DEFAULT_IMAGE and injects all args itself. + """ + routing = getattr(replica.spec, "routing", None) + if routing is None: + return None + tmpl = getattr(routing, "template", None) + if tmpl is None: + return None + spec = getattr(tmpl, "spec", None) or {} + for c in spec.get("containers", []): + if c.get("name") == "epp": + return c + return None + + +def _epp_objects( + replica, + name: str, + provider_config: str, +) -> dict[str, k8sobjv1alpha1.Object]: + """Build all EPP-related Objects for a disaggregated replica. + + Emits 8 provider-kubernetes Objects (keyed as epp-serviceaccount, epp-role, + epp-rolebinding, epp-clusterrole, epp-clusterrolebinding, epp-config, epp, + epp-service). All namespace-scoped objects go into base.REMOTE_NAMESPACE. + ClusterRole/ClusterRoleBinding are cluster-scoped (no namespace in metadata). + + The EPP container image comes from replica.spec.routing.template (the epp + container in that list). Modelplane injects the required operational args; + any user-supplied args in the template are prepended before the injected set. + """ + ns = base.REMOTE_NAMESPACE + epp_name = f"{name}-epp" + role_name = f"{name}-{_EPP_ROLE_SUFFIX}" + auth_name = f"{name}-{_EPP_AUTH_SUFFIX}" + # The epp container (if any) supplies both the image and any user args. + epp = _epp_container_from_routing(replica) + image = epp.get("image", _EPP_DEFAULT_IMAGE) if epp else _EPP_DEFAULT_IMAGE + user_epp_args = list(epp.get("args") or []) if epp else [] + + # Modelplane-injected args (appended after any user args). + injected_args = [ + f"--pool-name={name}-pool", + f"--pool-namespace={ns}", + "--pool-group=inference.networking.k8s.io", + "--zap-encoder=json", + "--config-file=/config/pd-epp-config.yaml", + "--grpc-port=9002", + ] + epp_args = user_epp_args + injected_args + + sa = { + "apiVersion": "v1", + "kind": "ServiceAccount", + "metadata": {"name": epp_name, "namespace": ns}, + } + + role = { + "apiVersion": "rbac.authorization.k8s.io/v1", + "kind": "Role", + "metadata": {"name": role_name, "namespace": ns}, + "rules": [ + { + "apiGroups": [""], + "resources": ["pods"], + "verbs": ["get", "watch", "list"], + }, + { + "apiGroups": ["inference.networking.x-k8s.io"], + "resources": ["inferenceobjectives", "inferencemodelrewrites"], + "verbs": ["get", "watch", "list"], + }, + { + "apiGroups": ["llm-d.ai"], + "resources": ["inferenceobjectives", "inferencemodelrewrites"], + "verbs": ["get", "watch", "list"], + }, + { + "apiGroups": ["inference.networking.k8s.io"], + "resources": ["inferencepools"], + "verbs": ["get", "watch", "list"], + }, + ], + } + + rolebinding = { + "apiVersion": "rbac.authorization.k8s.io/v1", + "kind": "RoleBinding", + "metadata": {"name": role_name, "namespace": ns}, + "subjects": [ + {"kind": "ServiceAccount", "name": epp_name, "namespace": ns}, + ], + "roleRef": { + "apiGroup": "rbac.authorization.k8s.io", + "kind": "Role", + "name": role_name, + }, + } + + clusterrole = { + "apiVersion": "rbac.authorization.k8s.io/v1", + "kind": "ClusterRole", + "metadata": {"name": auth_name}, + "rules": [ + { + "apiGroups": ["authentication.k8s.io"], + "resources": ["tokenreviews"], + "verbs": ["create"], + }, + { + "apiGroups": ["authorization.k8s.io"], + "resources": ["subjectaccessreviews"], + "verbs": ["create"], + }, + ], + } + + clusterrolebinding = { + "apiVersion": "rbac.authorization.k8s.io/v1", + "kind": "ClusterRoleBinding", + "metadata": {"name": auth_name}, + "subjects": [ + {"kind": "ServiceAccount", "name": epp_name, "namespace": ns}, + ], + "roleRef": { + "apiGroup": "rbac.authorization.k8s.io", + "kind": "ClusterRole", + "name": auth_name, + }, + } + + configmap = { + "apiVersion": "v1", + "kind": "ConfigMap", + "metadata": {"name": epp_name, "namespace": ns}, + "data": {"pd-epp-config.yaml": _PD_EPP_CONFIG_YAML}, + } + + deployment = { + "apiVersion": "apps/v1", + "kind": "Deployment", + "metadata": {"name": epp_name, "namespace": ns}, + "spec": { + "replicas": 1, + "strategy": {"type": "Recreate"}, + "selector": {"matchLabels": {"app": epp_name}}, + "template": { + "metadata": {"labels": {"app": epp_name}}, + "spec": { + "serviceAccountName": epp_name, + "terminationGracePeriodSeconds": 130, + "containers": [ + { + "name": "epp", + "image": image, + "args": epp_args, + "env": [ + { + "name": "NAMESPACE", + "valueFrom": {"fieldRef": {"fieldPath": "metadata.namespace"}}, + }, + { + "name": "POD_NAME", + "valueFrom": {"fieldRef": {"fieldPath": "metadata.name"}}, + }, + ], + "ports": [ + {"name": "grpc", "containerPort": 9002}, + {"name": "grpc-health", "containerPort": 9003}, + {"name": "metrics", "containerPort": 9090}, + ], + "livenessProbe": { + "grpc": {"port": 9003, "service": "inference-extension"}, + "initialDelaySeconds": 5, + "periodSeconds": 10, + }, + "readinessProbe": { + "grpc": {"port": 9003, "service": "inference-extension"}, + "initialDelaySeconds": 5, + "periodSeconds": 10, + }, + "volumeMounts": [ + {"name": "plugins-config-volume", "mountPath": "/config"}, + ], + } + ], + "volumes": [ + { + "name": "plugins-config-volume", + "configMap": {"name": epp_name}, + } + ], + }, + }, + }, + } + + service = { + "apiVersion": "v1", + "kind": "Service", + "metadata": {"name": epp_name, "namespace": ns}, + "spec": { + "selector": {"app": epp_name}, + "type": "ClusterIP", + "ports": [ + { + "name": "grpc-ext-proc", + "protocol": "TCP", + "port": 9002, + "targetPort": 9002, + "appProtocol": "http2", + }, + { + "name": "http-metrics", + "protocol": "TCP", + "port": 9090, + }, + ], + }, + } + + return { + "epp-serviceaccount": base.wrap_object(provider_config, sa), + "epp-role": base.wrap_object(provider_config, role), + "epp-rolebinding": base.wrap_object(provider_config, rolebinding), + "epp-clusterrole": base.wrap_object(provider_config, clusterrole), + "epp-clusterrolebinding": base.wrap_object(provider_config, clusterrolebinding), + "epp-config": base.wrap_object(provider_config, configmap), + "epp": base.wrap_object(provider_config, deployment), + "epp-service": base.wrap_object(provider_config, service), + } + + class LLMDBackend: def build( self, @@ -83,32 +571,54 @@ def build( name = replica.metadata.name tensor = int(replica.spec.workers.topology.tensor) - # nodes_per_worker == pipeline: the LWS gang size (leader + workers). - size = base.nodes_per_worker(replica) pipeline = int(replica.spec.workers.topology.pipeline or 1) cache_volumes, cache_volume_mounts = base.cache_mounts(replica) + pull_secrets = ( + [s.model_dump(exclude_none=True) for s in replica.spec.workers.template.spec.imagePullSecrets] + if replica.spec.workers.template.spec.imagePullSecrets + else None + ) + env = [e.model_dump(exclude_none=True) for e in engine.env] if engine.env else None + + # Disaggregated (prefill set): decode pods carry {pd-role: decode}. + # Unified (no prefill): no pd-role label on decode pods (backward compat). + prefill_spec = getattr(replica.spec, "prefill", None) + disagg = prefill_spec is not None + + # Disaggregated pods need VLLM_NIXL_SIDE_CHANNEL_HOST set to their own + # pod IP so NixlConnector can open the KV side-channel. It cannot come + # from user args (it's pod-IP, not a static value), so the backend + # injects it via the Kubernetes downward API. Only disaggregated replicas + # get it; the unified path is left untouched. + if disagg: + env = (env or []) + [base.nixl_side_channel_env()] + + # Disaggregated decode: the pd-sidecar takes the external serving port + # (ENGINE_PORT = 8000); vLLM moves to _DECODE_ENGINE_PORT (8001). + # Unified / prefill paths stay on ENGINE_PORT. + engine_serving_port = base._DECODE_ENGINE_PORT if disagg else base.ENGINE_PORT + + # Build leader/worker commands (and the args list for the container closure). # A user-supplied command owns cross-node coordination: inject neither the # Ray bootstrap nor vLLM-specific parallelism flags. It runs verbatim on # both templates (e.g. SGLang's symmetric launch against the LWS_* env). - user_command = list(engine.command) if engine.command else None - if user_command: - leader_command = worker_command = user_command - args = list(engine.args or []) - else: - # Args are folded into the leader command (consumed as "$@"); the - # worker only joins the gang. - args = _engine_args(engine, tensor, pipeline) - args = base.apply_cache_args(args, replica, engine) - leader_command = ["/bin/sh", "-c", _LEADER_BOOTSTRAP, "vllm", *args] - worker_command = ["/bin/sh", "-c", _WORKER_BOOTSTRAP] - - pull_secrets = None - tmpl = replica.spec.workers.template - if tmpl.spec.imagePullSecrets: - pull_secrets = [s.model_dump(exclude_none=True) for s in tmpl.spec.imagePullSecrets] - env = [e.model_dump(exclude_none=True) for e in engine.env] if engine.env else None + user_command, leader_command, worker_command, args = _build_commands( + engine, tensor, pipeline, replica, disagg=disagg + ) + + decode_claim = base.claim_template_name(replica) + decode_extra: dict = ( + { + base.LABEL_PD_ROLE: "decode", + base.LABEL_LLMD_ROLE: "decode", + base.LABEL_LLMD_SERVING: "true", + "app": name, + } + if disagg + else {} + ) def container(command: list[str], *, serving: bool) -> dict: c = { @@ -128,9 +638,11 @@ def container(command: list[str], *, serving: bool) -> dict: if env: c["env"] = env if serving: - c["ports"] = [{"containerPort": base.ENGINE_PORT}] + # Disaggregated decode: engine moves to _DECODE_ENGINE_PORT (8001); + # unified/prefill stay on ENGINE_PORT (8000). + c["ports"] = [{"containerPort": engine_serving_port}] c["readinessProbe"] = { - "httpGet": {"path": "/health", "port": base.ENGINE_PORT}, + "httpGet": {"path": "/health", "port": engine_serving_port}, "initialDelaySeconds": 30, "periodSeconds": 10, } @@ -143,19 +655,26 @@ def pod_spec(c: dict) -> dict: } # Both the leader and worker pods pin to the scheduled pool and # claim GPUs via DRA. - base.place_pod(spec, replica) + base.place_pod_on(spec, replica.spec.nodePoolName, decode_claim) if pull_secrets: spec["imagePullSecrets"] = pull_secrets return spec # Only the leader serves the OpenAI API → it carries the role label the # Service selects on, plus the serving port and readiness probe. + leader_spec = pod_spec(container(leader_command, serving=True)) + # Disaggregated decode: append the pd-sidecar to the leader pod. The + # sidecar takes ENGINE_PORT (8000) so the Service targetPort is unchanged; + # the engine has already moved to _DECODE_ENGINE_PORT (8001) above. + # Workers don't serve the API, so they get no sidecar. + if disagg: + leader_spec["containers"].append(base.pd_sidecar_container()) leader_pod = { - "metadata": {"labels": {base.LABEL_SERVING: name, _LABEL_ROLE: "leader"}}, - "spec": pod_spec(container(leader_command, serving=True)), + "metadata": {"labels": {base.LABEL_SERVING: name, _LABEL_ROLE: "leader", **decode_extra}}, + "spec": leader_spec, } worker_pod = { - "metadata": {"labels": {base.LABEL_SERVING: name}}, + "metadata": {"labels": {base.LABEL_SERVING: name, **decode_extra}}, "spec": pod_spec(container(worker_command, serving=False)), } @@ -167,25 +686,33 @@ def pod_spec(c: dict) -> dict: "spec": { "replicas": int(replica.spec.workers.count or 1), "leaderWorkerTemplate": { - "size": size, + "size": base.nodes_per_worker(replica), "leaderTemplate": leader_pod, "workerTemplate": worker_pod, }, }, } + # Service selector: always selects leader pods for this replica. + # For a disagg replica also narrow to pd-role:decode so prefill leader + # pods (which are not behind this Service) are never selected. + svc_selector: dict = {base.LABEL_SERVING: name, _LABEL_ROLE: "leader"} + if disagg: + svc_selector[base.LABEL_PD_ROLE] = "decode" + # Service selects the leader pods of every gang in this replica. service = { "apiVersion": "v1", "kind": "Service", "metadata": {"name": name, "namespace": base.REMOTE_NAMESPACE}, "spec": { - "selector": {base.LABEL_SERVING: name, _LABEL_ROLE: "leader"}, + "selector": svc_selector, "ports": [{"port": 80, "targetPort": base.ENGINE_PORT}], }, } - # HTTPRoute -> Service (plain Gateway API; Traefik- and Envoy-compatible). + # HTTPRoute -> InferencePool (disagg) or Service (unified). + # Plain Gateway API; Traefik- and Envoy-compatible. http_route = { "apiVersion": "gateway.networking.k8s.io/v1", "kind": "HTTPRoute", @@ -211,7 +738,7 @@ def pod_spec(c: dict) -> dict: "urlRewrite": {"path": {"type": "ReplacePrefixMatch", "replacePrefixMatch": "/"}}, } ], - "backendRefs": [{"name": name, "port": 80}], + "backendRefs": _route_backend_refs(name, disagg=disagg), } ], }, @@ -223,4 +750,16 @@ def pod_spec(c: dict) -> dict: "model-route": base.wrap_object(pc, http_route), } out[base.RESOURCE_CLAIM_KEY] = base.resource_claim_template(replica, pc) + + # Disaggregated replica: emit the prefill pod set + its ResourceClaimTemplate. + # No prefill Service or HTTPRoute — prefill is internal-only. + # Also emit the InferencePool that the HTTPRoute now points to; it lets the + # GAIE EPP perform KV-aware routing across the decode and prefill pods. + # The EPP objects (ServiceAccount, Role/RoleBinding, ClusterRole/ClusterRoleBinding, + # ConfigMap, Deployment, Service) are also emitted for the disaggregated path. + if disagg: + out["inference-pool"] = _inference_pool_object(name, pc) + out.update(_prefill_objects(replica, prefill_spec, name, pc, cache_volumes, cache_volume_mounts)) + out.update(_epp_objects(replica, name, pc)) + return out diff --git a/functions/compose-model-replica/tests/test_backends.py b/functions/compose-model-replica/tests/test_backends.py index 6919353bd..807190938 100644 --- a/functions/compose-model-replica/tests/test_backends.py +++ b/functions/compose-model-replica/tests/test_backends.py @@ -372,6 +372,597 @@ def test_pipeline_none_defaults_to_single_pod(self): replica.spec.workers.topology.pipeline = None self.assertFalse(base.needs_cross_pod_coordination(replica)) + def test_prefill_forces_llmd_backend(self): + # A single-pod replica with spec.prefill set must route to llm-d even + # though pipeline=1 (disaggregation requires cross-pod coordination + # between the prefill and decode roles regardless of per-role topology). + replica = _replica(tensor=1, pipeline=1) + replica.spec.prefill = v1alpha1.Prefill( + workers=v1alpha1.Workers( + count=1, + topology=v1alpha1.Topology(tensor=1, pipeline=1), + template=v1alpha1.Template( + spec=v1alpha1.Spec( + containers=[ + v1alpha1.Container( + name="engine", + image="vllm/vllm-openai:latest", + args=["--model=Qwen/Qwen3-0.6B"], + ) + ] + ) + ), + ), + nodePoolName="p", + deviceRequests=[_gpu_request(1)], + ) + self.assertEqual(base.select_backend(replica), base.LLMD) + # A single-pod replica WITHOUT prefill still picks native. + self.assertEqual(base.select_backend(_replica(tensor=1, pipeline=1)), base.NATIVE) + + +class TestDisaggregatedLLMD(unittest.TestCase): + """Tests for the disaggregated prefill/decode path in LLMDBackend.""" + + # EPP image used in all disagg tests — comes from the routing.template epp container. + _EPP_IMAGE = "ghcr.io/llm-d/llm-d-inference-scheduler:v0.8.0" + + def _disagg_replica(self): + """A single-pod disaggregated replica: decode on 'frontier', prefill on 'prefill-pool'. + + Sets spec.routing with a template containing an 'epp' container so the backend + can extract the EPP image from it (as required by P2.6). + """ + replica = _replica(name="dr", tensor=1, pipeline=1, args=["--model=Qwen/Qwen3-0.6B"]) + replica.spec.prefill = v1alpha1.Prefill( + workers=v1alpha1.Workers( + count=1, + topology=v1alpha1.Topology(tensor=1, pipeline=1), + template=v1alpha1.Template( + spec=v1alpha1.Spec( + containers=[ + v1alpha1.Container( + name="engine", + image="vllm/vllm-openai:latest", + args=["--model=Qwen/Qwen3-0.6B"], + ) + ] + ) + ), + ), + nodePoolName="prefill-pool", + deviceRequests=[_gpu_request(1)], + ) + # routing.template.spec is dict[str, Any] — the epp container image lives here. + replica.spec.routing = v1alpha1.Routing( + template=v1alpha1.TemplateModel( + spec={ + "containers": [ + {"name": "epp", "image": self._EPP_IMAGE}, + ] + } + ) + ) + return replica + + def _build(self): + return llmd.LLMDBackend().build(self._disagg_replica(), _CLUSTER) + + # --- key / name presence --- + + def test_both_pod_sets_present(self): + out = self._build() + self.assertIn("model-serving", out) + self.assertIn("prefill-serving", out) + + def test_distinct_metadata_names(self): + out = self._build() + decode_name = out["model-serving"].spec.forProvider.manifest["metadata"]["name"] + prefill_name = out["prefill-serving"].spec.forProvider.manifest["metadata"]["name"] + self.assertNotEqual(decode_name, prefill_name) + self.assertEqual(decode_name, "dr") + self.assertEqual(prefill_name, "dr-prefill") + + def test_two_distinct_resource_claim_templates(self): + out = self._build() + self.assertIn(base.RESOURCE_CLAIM_KEY, out) # "resource-claim" + self.assertIn("prefill-resource-claim", out) + decode_tmpl = out[base.RESOURCE_CLAIM_KEY].spec.forProvider.manifest + prefill_tmpl = out["prefill-resource-claim"].spec.forProvider.manifest + self.assertNotEqual( + decode_tmpl["metadata"]["name"], + prefill_tmpl["metadata"]["name"], + ) + + def test_prefill_pods_mount_model_cache(self): + """Prefill loads the same weights as decode, so it must mount the cache too.""" + replica = self._disagg_replica() + replica.spec.modelCacheRef = v1alpha1.ModelCacheRef(name="qwen") + out = llmd.LLMDBackend().build(replica, _CLUSTER) + leader = out["prefill-serving"].spec.forProvider.manifest["spec"]["leaderWorkerTemplate"]["leaderTemplate"][ + "spec" + ] + self.assertIn("model-cache", {v["name"] for v in leader["volumes"]}) + self.assertIn("model-cache", {m["name"] for m in leader["containers"][0]["volumeMounts"]}) + + # --- pd-role labels on pods --- + + def test_decode_leader_pod_carries_pd_role_decode(self): + out = self._build() + lws = out["model-serving"].spec.forProvider.manifest + leader_labels = lws["spec"]["leaderWorkerTemplate"]["leaderTemplate"]["metadata"]["labels"] + self.assertEqual(leader_labels.get(base.LABEL_PD_ROLE), "decode") + + def test_decode_worker_pod_carries_pd_role_decode(self): + out = self._build() + lws = out["model-serving"].spec.forProvider.manifest + worker_labels = lws["spec"]["leaderWorkerTemplate"]["workerTemplate"]["metadata"]["labels"] + self.assertEqual(worker_labels.get(base.LABEL_PD_ROLE), "decode") + + def test_prefill_leader_pod_carries_pd_role_prefill(self): + out = self._build() + lws = out["prefill-serving"].spec.forProvider.manifest + leader_labels = lws["spec"]["leaderWorkerTemplate"]["leaderTemplate"]["metadata"]["labels"] + self.assertEqual(leader_labels.get(base.LABEL_PD_ROLE), "prefill") + + def test_prefill_worker_pod_carries_pd_role_prefill(self): + out = self._build() + lws = out["prefill-serving"].spec.forProvider.manifest + worker_labels = lws["spec"]["leaderWorkerTemplate"]["workerTemplate"]["metadata"]["labels"] + self.assertEqual(worker_labels.get(base.LABEL_PD_ROLE), "prefill") + + # --- Service selector excludes prefill --- + + def test_decode_service_selector_includes_pd_role_decode(self): + out = self._build() + svc = out["model-service"].spec.forProvider.manifest + selector = svc["spec"]["selector"] + self.assertEqual(selector.get(base.LABEL_PD_ROLE), "decode") + + def test_decode_service_selector_does_not_include_pd_role_prefill(self): + out = self._build() + svc = out["model-service"].spec.forProvider.manifest + selector = svc["spec"]["selector"] + self.assertNotEqual(selector.get(base.LABEL_PD_ROLE), "prefill") + + def test_no_prefill_service_or_route(self): + out = self._build() + self.assertNotIn("prefill-service", out) + self.assertNotIn("prefill-route", out) + + # --- node pool pinning --- + + def test_decode_lws_pinned_to_decode_pool(self): + replica = self._disagg_replica() + out = llmd.LLMDBackend().build(replica, _CLUSTER) + lws = out["model-serving"].spec.forProvider.manifest + leader_ns = lws["spec"]["leaderWorkerTemplate"]["leaderTemplate"]["spec"]["nodeSelector"] + self.assertEqual(leader_ns.get("modelplane.ai/pool"), replica.spec.nodePoolName) + + def test_prefill_lws_pinned_to_prefill_pool(self): + replica = self._disagg_replica() + out = llmd.LLMDBackend().build(replica, _CLUSTER) + lws = out["prefill-serving"].spec.forProvider.manifest + leader_ns = lws["spec"]["leaderWorkerTemplate"]["leaderTemplate"]["spec"]["nodeSelector"] + self.assertEqual(leader_ns.get("modelplane.ai/pool"), replica.spec.prefill.nodePoolName) + + # --- unified replica backward-compat --- + + def test_unified_replica_no_prefill_serving_key(self): + out = llmd.LLMDBackend().build( + _replica(tensor=8, pipeline=2, args=["--model=meta-llama/Llama-3.1-405B"]), + _CLUSTER, + ) + self.assertNotIn("prefill-serving", out) + self.assertNotIn("prefill-resource-claim", out) + + def test_unified_replica_no_pd_role_label_on_pods(self): + # Unified replicas must not gain a pd-role label (no behavioral change). + out = llmd.LLMDBackend().build( + _replica(tensor=8, pipeline=2, args=["--model=meta-llama/Llama-3.1-405B"]), + _CLUSTER, + ) + lws = out["model-serving"].spec.forProvider.manifest + leader_labels = lws["spec"]["leaderWorkerTemplate"]["leaderTemplate"]["metadata"]["labels"] + self.assertNotIn(base.LABEL_PD_ROLE, leader_labels) + + def test_unified_replica_service_selector_unchanged(self): + out = llmd.LLMDBackend().build( + _replica(tensor=8, pipeline=2, args=["--model=meta-llama/Llama-3.1-405B"]), + _CLUSTER, + ) + svc = out["model-service"].spec.forProvider.manifest + selector = svc["spec"]["selector"] + self.assertNotIn(base.LABEL_PD_ROLE, selector) + + # --- readiness policies for prefill pod set --- + + def test_prefill_lws_readiness_derives_from_cel(self): + out = self._build() + self.assertEqual(out["prefill-serving"].spec.readiness.policy, "DeriveFromCelQuery") + self.assertEqual(out["prefill-serving"].spec.readiness.celQuery, base.AVAILABLE_CEL) + + def test_prefill_resource_claim_ready_on_create(self): + out = self._build() + self.assertEqual(out["prefill-resource-claim"].spec.readiness.policy, "SuccessfulCreate") + + # --- NIXL side-channel env injection --- + + def test_decode_and_prefill_get_nixl_side_channel_env(self): + """Both decode and prefill leader containers get VLLM_NIXL_SIDE_CHANNEL_HOST (disagg only).""" + out = self._build() + nixl_env = {"name": "VLLM_NIXL_SIDE_CHANNEL_HOST", "valueFrom": {"fieldRef": {"fieldPath": "status.podIP"}}} + + decode_lws = out["model-serving"].spec.forProvider.manifest + decode_containers = decode_lws["spec"]["leaderWorkerTemplate"]["leaderTemplate"]["spec"]["containers"] + decode_env = decode_containers[0].get("env", []) + self.assertIn(nixl_env, decode_env, "decode leader container missing VLLM_NIXL_SIDE_CHANNEL_HOST") + + prefill_lws = out["prefill-serving"].spec.forProvider.manifest + prefill_containers = prefill_lws["spec"]["leaderWorkerTemplate"]["leaderTemplate"]["spec"]["containers"] + prefill_env = prefill_containers[0].get("env", []) + self.assertIn(nixl_env, prefill_env, "prefill leader container missing VLLM_NIXL_SIDE_CHANNEL_HOST") + + def test_decode_and_prefill_engines_enable_nixl_kv_transfer(self): + """Both disagg engines must run vLLM with the NixlConnector kv-transfer-config. + + Without it the engines are plain servers and no prefill->decode KV handoff + occurs. NixlConnector does not distinguish kv_role, so both run kv_both. + """ + out = self._build() + for key in ("model-serving", "prefill-serving"): + lws = out[key].spec.forProvider.manifest + command = lws["spec"]["leaderWorkerTemplate"]["leaderTemplate"]["spec"]["containers"][0]["command"] + kv = [a for a in command if a.startswith("--kv-transfer-config")] + self.assertEqual(len(kv), 1, f"{key} engine missing --kv-transfer-config") + self.assertIn('"kv_connector":"NixlConnector"', kv[0]) + self.assertIn('"kv_role":"kv_both"', kv[0]) + + def test_unified_replica_has_no_nixl_kv_transfer(self): + """Unified replicas must NOT get --kv-transfer-config (no disaggregation).""" + out = llmd.LLMDBackend().build(_replica(tensor=8, pipeline=2), _CLUSTER) + command = out["model-serving"].spec.forProvider.manifest["spec"]["leaderWorkerTemplate"]["leaderTemplate"][ + "spec" + ]["containers"][0]["command"] + self.assertFalse([a for a in command if a.startswith("--kv-transfer-config")]) + + def test_epp_configmap_uses_gie_apiversion(self): + """The EndpointPickerConfig must use the GIE group the EPP binary registers.""" + config = self._build()["epp-config"].spec.forProvider.manifest["data"]["pd-epp-config.yaml"] + self.assertIn("apiVersion: inference.networking.x-k8s.io/v1alpha1", config) + self.assertNotIn("llm-d.ai/v1alpha1", config) + + def test_default_epp_and_sidecar_images(self): + """Fallback EPP image and the pd-sidecar image are the published llm-d refs.""" + replica = self._disagg_replica() + replica.spec.routing = v1alpha1.Routing(template=v1alpha1.TemplateModel(spec={"containers": []})) + out = llmd.LLMDBackend().build(replica, _CLUSTER) + epp_image = out["epp"].spec.forProvider.manifest["spec"]["template"]["spec"]["containers"][0]["image"] + self.assertEqual(epp_image, "ghcr.io/llm-d/llm-d-inference-scheduler:v0.8.0") + sidecar = self._decode_leader_containers(out)[1] + self.assertEqual(sidecar["image"], "ghcr.io/llm-d/llm-d-routing-sidecar:v0.8.0") + + def test_unified_replica_has_no_nixl_env(self): + """Unified multi-node replicas must NOT get VLLM_NIXL_SIDE_CHANNEL_HOST.""" + out = llmd.LLMDBackend().build( + _replica(tensor=8, pipeline=2, args=["--model=meta-llama/Llama-3.1-405B"]), + _CLUSTER, + ) + lws = out["model-serving"].spec.forProvider.manifest + leader_containers = lws["spec"]["leaderWorkerTemplate"]["leaderTemplate"]["spec"]["containers"] + leader_env = leader_containers[0].get("env", []) + env_names = [e.get("name") for e in leader_env] + self.assertNotIn( + "VLLM_NIXL_SIDE_CHANNEL_HOST", env_names, "unified replica must not have VLLM_NIXL_SIDE_CHANNEL_HOST" + ) + + # --- llm-d EPP selector labels (P2.3) --- + + def _lws_pod_labels(self, out, key, template): + """Helper: extract pod metadata labels from an LWS manifest.""" + lws = out[key].spec.forProvider.manifest + return lws["spec"]["leaderWorkerTemplate"][template]["metadata"]["labels"] + + def test_decode_pods_carry_llmd_selector_labels(self): + """Decode LWS leader and worker pods must carry the llm-d EPP selector labels.""" + out = self._build() + replica = self._disagg_replica() + name = replica.metadata.name # "dr" + + for template in ("leaderTemplate", "workerTemplate"): + with self.subTest(template=template): + labels = self._lws_pod_labels(out, "model-serving", template) + self.assertEqual(labels.get("app"), name, f"{template}: app label must equal replica name") + self.assertEqual( + labels.get(base.LABEL_LLMD_SERVING), + "true", + f"{template}: llm-d.ai/inference-serving must be 'true'", + ) + self.assertEqual( + labels.get(base.LABEL_LLMD_ROLE), + "decode", + f"{template}: llm-d.ai/role must be 'decode'", + ) + # Existing modelplane pd-role label must still be present. + self.assertEqual( + labels.get(base.LABEL_PD_ROLE), + "decode", + f"{template}: modelplane.ai/pd-role must still be 'decode'", + ) + + def test_prefill_pods_carry_llmd_selector_labels(self): + """Prefill LWS leader and worker pods must carry the llm-d EPP selector labels.""" + out = self._build() + replica = self._disagg_replica() + name = replica.metadata.name # "dr" + + for template in ("leaderTemplate", "workerTemplate"): + with self.subTest(template=template): + labels = self._lws_pod_labels(out, "prefill-serving", template) + self.assertEqual( + labels.get("app"), name, f"{template}: app label must equal replica name (not prefill_name)" + ) + self.assertEqual( + labels.get(base.LABEL_LLMD_SERVING), + "true", + f"{template}: llm-d.ai/inference-serving must be 'true'", + ) + self.assertEqual( + labels.get(base.LABEL_LLMD_ROLE), + "prefill", + f"{template}: llm-d.ai/role must be 'prefill'", + ) + # Existing modelplane pd-role label must still be present. + self.assertEqual( + labels.get(base.LABEL_PD_ROLE), + "prefill", + f"{template}: modelplane.ai/pd-role must still be 'prefill'", + ) + + def test_app_label_identical_on_both_roles(self): + """The shared InferencePool selector requires the same 'app' value on decode and prefill pods.""" + out = self._build() + decode_labels = self._lws_pod_labels(out, "model-serving", "leaderTemplate") + prefill_labels = self._lws_pod_labels(out, "prefill-serving", "leaderTemplate") + self.assertEqual( + decode_labels.get("app"), + prefill_labels.get("app"), + "app label must be identical on both roles so one InferencePool selects both", + ) + + def test_unified_replica_no_llmd_selector_labels(self): + """A unified (non-disaggregated) replica must NOT receive llm-d EPP labels.""" + out = llmd.LLMDBackend().build( + _replica(tensor=8, pipeline=2, args=["--model=meta-llama/Llama-3.1-405B"]), + _CLUSTER, + ) + lws = out["model-serving"].spec.forProvider.manifest + leader_labels = lws["spec"]["leaderWorkerTemplate"]["leaderTemplate"]["metadata"]["labels"] + self.assertNotIn(base.LABEL_LLMD_ROLE, leader_labels, "unified replica must not have llm-d.ai/role") + self.assertNotIn( + base.LABEL_LLMD_SERVING, leader_labels, "unified replica must not have llm-d.ai/inference-serving" + ) + self.assertNotIn("app", leader_labels, "unified replica must not have app label") + + # --- pd-sidecar injection (P2.4) --- + + def _decode_leader_containers(self, out): + """Helper: return the containers list from the decode LWS leader pod.""" + lws = out["model-serving"].spec.forProvider.manifest + return lws["spec"]["leaderWorkerTemplate"]["leaderTemplate"]["spec"]["containers"] + + def _prefill_leader_containers(self, out): + """Helper: return the containers list from the prefill LWS leader pod.""" + lws = out["prefill-serving"].spec.forProvider.manifest + return lws["spec"]["leaderWorkerTemplate"]["leaderTemplate"]["spec"]["containers"] + + def test_decode_has_pd_sidecar(self): + """Disaggregated decode leader pod has two containers; one named pd-sidecar with correct image, port, and args.""" + out = self._build() + containers = self._decode_leader_containers(out) + names = [c["name"] for c in containers] + self.assertEqual(len(containers), 2, f"expected 2 containers on decode leader, got {names}") + sidecar = next(c for c in containers if c["name"] == "pd-sidecar") + self.assertEqual(sidecar["image"], base.PD_SIDECAR_IMAGE) + ports = sidecar.get("ports", []) + self.assertEqual(len(ports), 1) + self.assertEqual(ports[0]["containerPort"], 8000) + args = sidecar.get("args", []) + self.assertIn("--secure-proxy=false", args) + self.assertIn("--kv-connector=nixlv2", args) + self.assertIn("--vllm-port=8001", args) + + def test_decode_engine_moves_to_8001(self): + """Disaggregated decode engine container uses port 8001 for both containerPort and readinessProbe; turnkey args include --port=8001.""" + out = self._build() + containers = self._decode_leader_containers(out) + engine = next(c for c in containers if c["name"] == "engine") + # containerPort + ports = engine.get("ports", []) + self.assertEqual(len(ports), 1) + self.assertEqual(ports[0]["containerPort"], 8001) + # readinessProbe + probe = engine.get("readinessProbe", {}) + self.assertEqual(probe["httpGet"]["port"], 8001) + # turnkey vLLM path: --port=8001 must appear in the args/command + cmd = engine.get("command", []) + all_args = cmd + engine.get("args", []) + self.assertTrue( + any("--port=8001" in a for a in all_args), + f"--port=8001 not found in engine command/args: {all_args}", + ) + + def test_decode_service_targets_sidecar_8000(self): + """The decode Service targetPort stays 8000 (the sidecar is the serving entry).""" + out = self._build() + svc = out["model-service"].spec.forProvider.manifest + target = svc["spec"]["ports"][0]["targetPort"] + self.assertEqual(target, 8000) + + def test_unified_replica_has_no_sidecar(self): + """A unified multi-node replica (pipeline=2, no prefill) has a single engine container on 8000.""" + out = llmd.LLMDBackend().build( + _replica(tensor=8, pipeline=2, args=["--model=meta-llama/Llama-3.1-405B"]), + _CLUSTER, + ) + lws = out["model-serving"].spec.forProvider.manifest + containers = lws["spec"]["leaderWorkerTemplate"]["leaderTemplate"]["spec"]["containers"] + names = [c["name"] for c in containers] + self.assertEqual(len(containers), 1, f"unified replica should have 1 container, got {names}") + self.assertNotIn("pd-sidecar", names) + # engine still on 8000 + engine = containers[0] + self.assertEqual(engine["ports"][0]["containerPort"], 8000) + + def test_prefill_has_no_sidecar(self): + """Prefill LWS leader pod has a single engine container, no pd-sidecar.""" + out = self._build() + containers = self._prefill_leader_containers(out) + names = [c["name"] for c in containers] + self.assertEqual(len(containers), 1, f"prefill leader should have 1 container, got {names}") + self.assertNotIn("pd-sidecar", names) + + # --- InferencePool + HTTPRoute wiring (P2.5) --- + + def test_emits_inference_pool(self): + """A disaggregated replica emits an InferencePool manifest at key 'inference-pool'.""" + out = self._build() + self.assertIn("inference-pool", out) + manifest = out["inference-pool"].spec.forProvider.manifest + name = self._disagg_replica().metadata.name # "dr" + self.assertEqual(manifest["kind"], "InferencePool") + self.assertEqual(manifest["apiVersion"], "inference.networking.k8s.io/v1") + self.assertEqual(manifest["metadata"]["name"], f"{name}-pool") + self.assertEqual(manifest["metadata"]["namespace"], base.REMOTE_NAMESPACE) + spec = manifest["spec"] + self.assertEqual( + spec["selector"]["matchLabels"], + {"app": name, "llm-d.ai/inference-serving": "true"}, + ) + self.assertEqual(spec["targetPorts"], [{"number": 8000}]) + epp = spec["endpointPickerRef"] + self.assertEqual(epp["name"], f"{name}-epp") + self.assertEqual(epp["port"]["number"], 9002) + # failureMode is a field of endpointPickerRef in the GAIE v1 schema. + self.assertEqual(epp.get("failureMode"), "FailOpen") + self.assertNotIn("failureMode", spec) + + def test_disagg_httproute_targets_inference_pool(self): + """For a disaggregated replica the HTTPRoute backendRefs points at the InferencePool, not the Service.""" + out = self._build() + name = self._disagg_replica().metadata.name # "dr" + route = out["model-route"].spec.forProvider.manifest + backend_refs = route["spec"]["rules"][0]["backendRefs"] + self.assertEqual(len(backend_refs), 1) + ref = backend_refs[0] + self.assertEqual(ref["group"], "inference.networking.k8s.io") + self.assertEqual(ref["kind"], "InferencePool") + self.assertEqual(ref["name"], f"{name}-pool") + # Port field must NOT be present for an InferencePool backendRef + self.assertNotIn("port", ref) + + def test_unified_httproute_targets_service(self): + """A unified multi-node replica's HTTPRoute backendRef still points at the Service, no InferencePool emitted.""" + out = llmd.LLMDBackend().build( + _replica(tensor=8, pipeline=2, args=["--model=meta-llama/Llama-3.1-405B"]), + _CLUSTER, + ) + self.assertNotIn("inference-pool", out) + route = out["model-route"].spec.forProvider.manifest + backend_refs = route["spec"]["rules"][0]["backendRefs"] + self.assertEqual(len(backend_refs), 1) + ref = backend_refs[0] + self.assertEqual(ref["name"], "r") + self.assertEqual(ref["port"], 80) + self.assertNotIn("group", ref) + self.assertNotIn("kind", ref) + + # --- EPP objects (P2.6) --- + + _EPP_KEYS = ( + "epp-serviceaccount", + "epp-role", + "epp-rolebinding", + "epp-clusterrole", + "epp-clusterrolebinding", + "epp-config", + "epp", + "epp-service", + ) + + def test_disagg_emits_all_epp_keys(self): + """A disaggregated replica emits all 8 EPP-related response keys.""" + out = self._build() + for key in self._EPP_KEYS: + with self.subTest(key=key): + self.assertIn(key, out) + + def test_unified_emits_no_epp_keys(self): + """A unified replica emits none of the epp* keys.""" + out = llmd.LLMDBackend().build( + _replica(tensor=8, pipeline=2, args=["--model=meta-llama/Llama-3.1-405B"]), + _CLUSTER, + ) + for key in self._EPP_KEYS: + with self.subTest(key=key): + self.assertNotIn(key, out) + + def test_epp_deployment_image_from_routing_template(self): + """The EPP Deployment's container image comes from routing.template's epp container.""" + out = self._build() + dep = out["epp"].spec.forProvider.manifest + containers = dep["spec"]["template"]["spec"]["containers"] + epp_container = next(c for c in containers if c["name"] == "epp") + self.assertEqual(epp_container["image"], self._EPP_IMAGE) + + def test_epp_deployment_args_include_pool_name(self): + """The EPP Deployment args include --pool-name=-pool.""" + out = self._build() + dep = out["epp"].spec.forProvider.manifest + containers = dep["spec"]["template"]["spec"]["containers"] + epp_container = next(c for c in containers if c["name"] == "epp") + args = epp_container.get("args", []) + name = self._disagg_replica().metadata.name # "dr" + self.assertIn(f"--pool-name={name}-pool", args) + + def test_epp_deployment_args_include_config_file(self): + """The EPP Deployment args include --config-file=/config/pd-epp-config.yaml.""" + out = self._build() + dep = out["epp"].spec.forProvider.manifest + containers = dep["spec"]["template"]["spec"]["containers"] + epp_container = next(c for c in containers if c["name"] == "epp") + args = epp_container.get("args", []) + self.assertIn("--config-file=/config/pd-epp-config.yaml", args) + + def test_epp_service_has_port_9002(self): + """The EPP Service has a port entry with port 9002.""" + out = self._build() + svc = out["epp-service"].spec.forProvider.manifest + ports = svc["spec"]["ports"] + port_numbers = [p["port"] for p in ports] + self.assertIn(9002, port_numbers) + + def test_epp_configmap_has_pd_epp_config_key(self): + """The EPP ConfigMap data has the key 'pd-epp-config.yaml'.""" + out = self._build() + cm = out["epp-config"].spec.forProvider.manifest + self.assertIn("pd-epp-config.yaml", cm["data"]) + + def test_epp_objects_readiness_successful_create(self): + """All EPP objects except the Deployment use SuccessfulCreate readiness.""" + out = self._build() + for key in ( + "epp-serviceaccount", + "epp-role", + "epp-rolebinding", + "epp-clusterrole", + "epp-clusterrolebinding", + "epp-config", + "epp-service", + ): + with self.subTest(key=key): + self.assertEqual(out[key].spec.readiness.policy, "SuccessfulCreate") + class TestDynamoStub(unittest.TestCase): def test_not_selected_in_v01(self): diff --git a/functions/compose-serving-stack/function/fn.py b/functions/compose-serving-stack/function/fn.py index 688685add..12b86061f 100644 --- a/functions/compose-serving-stack/function/fn.py +++ b/functions/compose-serving-stack/function/fn.py @@ -2,10 +2,10 @@ This function composes the serving substrate (the cluster-side CRDs, controllers, and gateway) that the native and llm-d model-serving backends -depend on: cert-manager, Envoy Gateway, Prometheus, LeaderWorkerSet, and an -inference Gateway. Resources are composed as Helm releases and -provider-kubernetes Objects, all targeting the remote cluster via -ProviderConfigs. +depend on: cert-manager, Envoy Gateway, Envoy AI Gateway, GAIE CRDs, +Prometheus, LeaderWorkerSet, and an inference Gateway. Resources are composed +as Helm releases and provider-kubernetes Objects, all targeting the remote +cluster via ProviderConfigs. Usage resources protect ProviderConfigs from premature deletion during teardown, ensuring Helm releases can uninstall before losing connectivity. @@ -36,6 +36,21 @@ # Identity type for GCP service account credentials. _IDENTITY_TYPE_GCP = "GoogleApplicationCredentials" +# Envoy AI Gateway constants. +_AI_GATEWAY_NAMESPACE = "envoy-ai-gateway-system" +_AI_GATEWAY_REPO = "oci://docker.io/envoyproxy" +_AI_GATEWAY_VERSION = "v0.7.0" +_AI_GATEWAY_CONTROLLER_FQDN = f"ai-gateway-controller.{_AI_GATEWAY_NAMESPACE}.svc.cluster.local" +_AI_GATEWAY_CONTROLLER_PORT = 1063 + +# GAIE (Gateway API Inference Extension) CRD chart constants. +# No remote-manifest pattern exists in this codebase; the CRDs are installed +# via the upstream Helm chart published by the kubernetes-sigs project. +_GAIE_CHART = "inferencepool" +_GAIE_REPO = "oci://ghcr.io/kubernetes-sigs/gateway-api-inference-extension/charts" +_GAIE_VERSION = "v1.0.1" +_GAIE_NAMESPACE = "gateway-api-inference-extension" + # Prometheus constants. _PROMETHEUS_NAMESPACE = "monitoring" _PROMETHEUS_FULLNAME_OVERRIDE = "prometheus" @@ -208,6 +223,8 @@ def compose(self): self.compose_usages() self.compose_cert_manager() self.compose_envoy_gateway() + self.compose_ai_gateway() + self.compose_gaie_crds() self.compose_prometheus() self.compose_leader_worker_set() self.compose_node_feature_discovery() @@ -428,7 +445,12 @@ def compose_cert_manager(self): ) def compose_envoy_gateway(self): - """Compose Envoy Gateway. Gated on ProviderConfigs being observed.""" + """Compose Envoy Gateway. Gated on ProviderConfigs being observed. + + Values merge the base envoy-gateway-values.yaml with the inference-pool + addon (envoy-gateway-values-addon.yaml) so that HTTPRoute -> InferencePool + backendRefs are recognised by the Envoy AI Gateway extension manager. + """ pc_observed = self.provider_configs_observed() if not (pc_observed or "envoy-gateway" in self.req.observed.resources): return @@ -446,13 +468,115 @@ def compose_envoy_gateway(self): values={ "config": { "envoyGateway": { - "extensionApis": {"enableBackend": True}, + "gateway": { + "controllerName": "gateway.envoyproxy.io/gatewayclass-controller", + }, + "logging": {"level": {"default": "info"}}, + "provider": {"type": "Kubernetes"}, + "extensionApis": { + "enableEnvoyPatchPolicy": True, + "enableBackend": True, + }, + "extensionManager": { + "hooks": { + "xdsTranslator": { + "translation": { + "listener": {"includeAll": True}, + "route": {"includeAll": True}, + "cluster": {"includeAll": True}, + "secret": {"includeAll": True}, + }, + "post": [ + "Translation", + "Cluster", + "Route", + ], + }, + }, + "service": { + "fqdn": { + "hostname": _AI_GATEWAY_CONTROLLER_FQDN, + "port": _AI_GATEWAY_CONTROLLER_PORT, + }, + }, + "backendResources": [ + { + "group": "inference.networking.k8s.io", + "kind": "InferencePool", + "version": "v1", + }, + ], + }, }, }, }, ), ) + def compose_ai_gateway(self): + """Compose the Envoy AI Gateway CRDs and controller. + + Installs two Helm releases from oci://docker.io/envoyproxy at v0.7.0: + - ai-gateway-crds: the AI Gateway CRD chart + - ai-gateway: the AI Gateway controller chart + + The AI Gateway controller provides the ext-proc extension manager that + Envoy Gateway delegates InferencePool backend resolution to. Both + releases are gated on the Envoy Gateway ProviderConfigs being observed + (same gate as envoy-gateway). + """ + pc_observed = self.provider_configs_observed() + if not (pc_observed or "ai-gateway-crds" in self.req.observed.resources): + return + + resource.update( + self.rsp.desired.resources["ai-gateway-crds"], + _helm_release( + chart="ai-gateway-crds-helm", + repo=_AI_GATEWAY_REPO, + version=_AI_GATEWAY_VERSION, + namespace=_AI_GATEWAY_NAMESPACE, + provider_config=_pc_name(self.xr), + ), + ) + resource.update( + self.rsp.desired.resources["ai-gateway"], + _helm_release( + chart="ai-gateway-helm", + repo=_AI_GATEWAY_REPO, + version=_AI_GATEWAY_VERSION, + namespace=_AI_GATEWAY_NAMESPACE, + provider_config=_pc_name(self.xr), + ), + ) + + def compose_gaie_crds(self): + """Compose the Gateway API Inference Extension (GAIE) CRDs. + + The codebase has no remote-manifest pattern for provider-kubernetes + Objects; the GAIE CRDs are therefore installed via the upstream Helm + chart published by kubernetes-sigs at + oci://ghcr.io/kubernetes-sigs/gateway-api-inference-extension/charts. + This is equivalent to applying the upstream manifests.yaml from + https://github.com/kubernetes-sigs/gateway-api-inference-extension/releases/download/v1.0.1/manifests.yaml + and is the approach recommended by the project for cluster-lifecycle + tooling that cannot apply raw URLs. + """ + pc_observed = self.provider_configs_observed() + if not (pc_observed or "gaie-crds" in self.req.observed.resources): + return + + resource.update( + self.rsp.desired.resources["gaie-crds"], + _helm_release( + chart=_GAIE_CHART, + repo=_GAIE_REPO, + version=_GAIE_VERSION, + namespace=_GAIE_NAMESPACE, + provider_config=_pc_name(self.xr), + ), + ) + def compose_prometheus(self): """Compose the kube-prometheus-stack. Gated on ProviderConfigs being observed. Provides cluster observability (metrics scraping).""" @@ -645,6 +769,9 @@ def mark_readiness(self): condition_ready = [ "cert-manager", "envoy-gateway", + "ai-gateway-crds", + "ai-gateway", + "gaie-crds", "prometheus", "leader-worker-set", "node-feature-discovery", diff --git a/functions/compose-serving-stack/tests/test_fn.py b/functions/compose-serving-stack/tests/test_fn.py index a37a569eb..f6a0bff5e 100644 --- a/functions/compose-serving-stack/tests/test_fn.py +++ b/functions/compose-serving-stack/tests/test_fn.py @@ -196,7 +196,45 @@ def setUpModule() -> None: "values": { "config": { "envoyGateway": { - "extensionApis": {"enableBackend": True}, + "gateway": { + "controllerName": "gateway.envoyproxy.io/gatewayclass-controller", + }, + "logging": {"level": {"default": "info"}}, + "provider": {"type": "Kubernetes"}, + "extensionApis": { + "enableEnvoyPatchPolicy": True, + "enableBackend": True, + }, + "extensionManager": { + "hooks": { + "xdsTranslator": { + "translation": { + "listener": {"includeAll": True}, + "route": {"includeAll": True}, + "cluster": {"includeAll": True}, + "secret": {"includeAll": True}, + }, + "post": [ + "Translation", + "Cluster", + "Route", + ], + }, + }, + "service": { + "fqdn": { + "hostname": "ai-gateway-controller.envoy-ai-gateway-system.svc.cluster.local", + "port": 1063.0, + }, + }, + "backendResources": [ + { + "group": "inference.networking.k8s.io", + "kind": "InferencePool", + "version": "v1", + }, + ], + }, }, }, }, @@ -208,6 +246,63 @@ def setUpModule() -> None: }, } +_AI_GATEWAY_CRDS = { + "apiVersion": "helm.m.crossplane.io/v1beta1", + "kind": "Release", + "spec": { + "forProvider": { + "chart": { + "name": "ai-gateway-crds-helm", + "repository": "oci://docker.io/envoyproxy", + "version": "v0.7.0", + }, + "namespace": "envoy-ai-gateway-system", + }, + "providerConfigRef": { + "kind": "ProviderConfig", + "name": _PC_NAME, + }, + }, +} + +_AI_GATEWAY = { + "apiVersion": "helm.m.crossplane.io/v1beta1", + "kind": "Release", + "spec": { + "forProvider": { + "chart": { + "name": "ai-gateway-helm", + "repository": "oci://docker.io/envoyproxy", + "version": "v0.7.0", + }, + "namespace": "envoy-ai-gateway-system", + }, + "providerConfigRef": { + "kind": "ProviderConfig", + "name": _PC_NAME, + }, + }, +} + +_GAIE_CRDS = { + "apiVersion": "helm.m.crossplane.io/v1beta1", + "kind": "Release", + "spec": { + "forProvider": { + "chart": { + "name": "inferencepool", + "repository": "oci://ghcr.io/kubernetes-sigs/gateway-api-inference-extension/charts", + "version": "v1.0.1", + }, + "namespace": "gateway-api-inference-extension", + }, + "providerConfigRef": { + "kind": "ProviderConfig", + "name": _PC_NAME, + }, + }, +} + _GATEWAY = { "apiVersion": "kubernetes.m.crossplane.io/v1alpha1", "kind": "Object", @@ -512,12 +607,21 @@ async def test_second_pass(self) -> None: resource=resource.dict_to_struct({"status": {}}), ), resources={ + "ai-gateway": fnv1.Resource( + resource=resource.dict_to_struct(_AI_GATEWAY), + ), + "ai-gateway-crds": fnv1.Resource( + resource=resource.dict_to_struct(_AI_GATEWAY_CRDS), + ), "cert-manager": fnv1.Resource( resource=resource.dict_to_struct(_CERT_MANAGER), ), "envoy-gateway": fnv1.Resource( resource=resource.dict_to_struct(_ENVOY_GATEWAY), ), + "gaie-crds": fnv1.Resource( + resource=resource.dict_to_struct(_GAIE_CRDS), + ), "gateway": fnv1.Resource( resource=resource.dict_to_struct(_GATEWAY), ), @@ -628,6 +732,12 @@ async def test_third_pass(self) -> None: ), ), resources={ + "ai-gateway": fnv1.Resource( + resource=resource.dict_to_struct(_AI_GATEWAY), + ), + "ai-gateway-crds": fnv1.Resource( + resource=resource.dict_to_struct(_AI_GATEWAY_CRDS), + ), "cert-manager": fnv1.Resource( resource=resource.dict_to_struct(_CERT_MANAGER), ready=fnv1.READY_TRUE, @@ -635,6 +745,9 @@ async def test_third_pass(self) -> None: "envoy-gateway": fnv1.Resource( resource=resource.dict_to_struct(_ENVOY_GATEWAY), ), + "gaie-crds": fnv1.Resource( + resource=resource.dict_to_struct(_GAIE_CRDS), + ), "gateway": fnv1.Resource( resource=resource.dict_to_struct(_GATEWAY), ), @@ -691,3 +804,137 @@ async def test_third_pass(self) -> None: json_format.MessageToDict(got), "-want, +got", ) + + +def _second_pass_request() -> fnv1.RunFunctionRequest: + """Build a second-pass request (both ProviderConfigs observed).""" + req = _base_request() + req.observed.resources["provider-config-helm"].CopyFrom( + fnv1.Resource( + resource=resource.dict_to_struct({"apiVersion": "helm.m.crossplane.io/v1beta1", "kind": "ProviderConfig"}), + ), + ) + req.observed.resources["provider-config-kubernetes"].CopyFrom( + fnv1.Resource( + resource=resource.dict_to_struct( + {"apiVersion": "kubernetes.m.crossplane.io/v1alpha1", "kind": "ProviderConfig"} + ), + ), + ) + return req + + +class TestAiGatewayReleases(unittest.IsolatedAsyncioTestCase): + """Tests for P2.2: Envoy AI Gateway + GAIE CRD releases.""" + + @classmethod + def setUpClass(cls) -> None: + cls.runner = fn.FunctionRunner() + + async def _run(self) -> fnv1.RunFunctionResponse: + return await self.runner.RunFunction(_second_pass_request(), None) + + async def test_ai_gateway_crds_release_present(self) -> None: + """ai-gateway-crds Release is composed with the correct chart and version.""" + got = await self._run() + resources = json_format.MessageToDict(got).get("desired", {}).get("resources", {}) + self.assertIn("ai-gateway-crds", resources, "ai-gateway-crds Release missing") + chart = resources["ai-gateway-crds"].get("resource", {}).get("spec", {}).get("forProvider", {}).get("chart", {}) + self.assertEqual(chart.get("name"), "ai-gateway-crds-helm") + self.assertEqual(chart.get("version"), "v0.7.0") + self.assertEqual( + resources["ai-gateway-crds"]["resource"]["spec"]["forProvider"]["namespace"], + "envoy-ai-gateway-system", + ) + + async def test_ai_gateway_release_present(self) -> None: + """ai-gateway Release is composed with the correct chart and version.""" + got = await self._run() + resources = json_format.MessageToDict(got).get("desired", {}).get("resources", {}) + self.assertIn("ai-gateway", resources, "ai-gateway Release missing") + chart = resources["ai-gateway"].get("resource", {}).get("spec", {}).get("forProvider", {}).get("chart", {}) + self.assertEqual(chart.get("name"), "ai-gateway-helm") + self.assertEqual(chart.get("version"), "v0.7.0") + self.assertEqual( + resources["ai-gateway"]["resource"]["spec"]["forProvider"]["namespace"], + "envoy-ai-gateway-system", + ) + + async def test_gaie_crds_release_present(self) -> None: + """gaie-crds Release is composed (Helm-based GAIE CRD install).""" + got = await self._run() + resources = json_format.MessageToDict(got).get("desired", {}).get("resources", {}) + self.assertIn("gaie-crds", resources, "gaie-crds Release missing") + chart = resources["gaie-crds"].get("resource", {}).get("spec", {}).get("forProvider", {}).get("chart", {}) + self.assertEqual(chart.get("name"), "inferencepool") + self.assertEqual(chart.get("version"), "v1.0.1") + + async def test_envoy_gateway_values_enable_backend(self) -> None: + """gateway-helm Release values include extensionApis.enableBackend: true.""" + got = await self._run() + resources = json_format.MessageToDict(got).get("desired", {}).get("resources", {}) + values = ( + resources.get("envoy-gateway", {}) + .get("resource", {}) + .get("spec", {}) + .get("forProvider", {}) + .get("values", {}) + ) + extension_apis = values.get("config", {}).get("envoyGateway", {}).get("extensionApis", {}) + self.assertTrue(extension_apis.get("enableBackend"), "extensionApis.enableBackend must be True") + + async def test_envoy_gateway_values_backend_resources_inference_pool(self) -> None: + """gateway-helm Release values include backendResources entry for InferencePool.""" + got = await self._run() + resources = json_format.MessageToDict(got).get("desired", {}).get("resources", {}) + values = ( + resources.get("envoy-gateway", {}) + .get("resource", {}) + .get("spec", {}) + .get("forProvider", {}) + .get("values", {}) + ) + backend_resources = ( + values.get("config", {}).get("envoyGateway", {}).get("extensionManager", {}).get("backendResources", []) + ) + self.assertTrue( + any( + br.get("group") == "inference.networking.k8s.io" + and br.get("kind") == "InferencePool" + and br.get("version") == "v1" + for br in backend_resources + ), + f"InferencePool backendResources entry not found; got: {backend_resources}", + ) + + async def test_envoy_gateway_values_extension_manager_service(self) -> None: + """gateway-helm Release values wire the AI Gateway controller endpoint.""" + got = await self._run() + resources = json_format.MessageToDict(got).get("desired", {}).get("resources", {}) + values = ( + resources.get("envoy-gateway", {}) + .get("resource", {}) + .get("spec", {}) + .get("forProvider", {}) + .get("values", {}) + ) + fqdn = ( + values.get("config", {}) + .get("envoyGateway", {}) + .get("extensionManager", {}) + .get("service", {}) + .get("fqdn", {}) + ) + self.assertEqual( + fqdn.get("hostname"), + "ai-gateway-controller.envoy-ai-gateway-system.svc.cluster.local", + ) + self.assertEqual(fqdn.get("port"), 1063.0) + + async def test_ai_gateway_releases_gated_on_first_pass(self) -> None: + """AI Gateway releases are absent on first pass (ProviderConfigs not yet observed).""" + got = await self.runner.RunFunction(_base_request(), None) + resources = json_format.MessageToDict(got).get("desired", {}).get("resources", {}) + self.assertNotIn("ai-gateway-crds", resources) + self.assertNotIn("ai-gateway", resources) + self.assertNotIn("gaie-crds", resources) diff --git a/schemas/python/models/ai/modelplane/modeldeployment/v1alpha1.py b/schemas/python/models/ai/modelplane/modeldeployment/v1alpha1.py index ef1747833..108ef1007 100644 --- a/schemas/python/models/ai/modelplane/modeldeployment/v1alpha1.py +++ b/schemas/python/models/ai/modelplane/modeldeployment/v1alpha1.py @@ -170,9 +170,9 @@ class Topology(BaseModel): class Workers(BaseModel): - count: conint(ge=1) | None = 1 + count: conint(ge=1) | None = None """ - Number of workers per replica. Defaults to 1. + Number of workers per replica. """ template: Template """ @@ -184,6 +184,36 @@ class Workers(BaseModel): """ +class Prefill(BaseModel): + nodeSelector: NodeSelector + """ + Node-level matching, a list of device requests mirroring a DRA ResourceClaim. The scheduler matches each request against a candidate pool's InferenceClass devices (surfaced on InferenceCluster status.gpuPools) and pins the replica to a pool that satisfies every request. claim: DRA requests also become DeviceRequests in the ResourceClaim the serving pods bind GPUs through. Required: GPUs bind only via DRA, so a deployment must declare the devices its model needs. At least one request must resolve to a claimable (claim: DRA) device; the serving workload binds its GPUs through the resulting ResourceClaim. Synthetic devices refine placement but are never claimed, so a nodeSelector that matches only synthetic devices leaves the workload nothing to claim - the scheduler treats such a pool as ineligible and the deployment reports InsufficientCapacity. + """ + workers: Workers + """ + Compute shape of one worker. Modelplane composes one worker (or workers.count workers) per ModelReplica. + """ + + +class TemplateModel(BaseModel): + spec: dict[str, Any] | None = None + + +class Routing(BaseModel): + template: TemplateModel | None = None + + +class TemplateModel1(BaseModel): + metadata: Metadata | None = None + """ + Metadata applied to inference pods. Useful for labels and annotations that control cluster-level features like service mesh injection. + """ + spec: Spec | None = None + """ + Pod spec for inference workers. + """ + + class SpecModel(BaseModel): clusterSelector: ClusterSelector | None = None """ @@ -201,10 +231,18 @@ class SpecModel(BaseModel): """ Node-level matching, a list of device requests mirroring a DRA ResourceClaim. The scheduler matches each request against a candidate pool's InferenceClass devices (surfaced on InferenceCluster status.gpuPools) and pins the replica to a pool that satisfies every request. claim: DRA requests also become DeviceRequests in the ResourceClaim the serving pods bind GPUs through. Required: GPUs bind only via DRA, so a deployment must declare the devices its model needs. At least one request must resolve to a claimable (claim: DRA) device; the serving workload binds its GPUs through the resulting ResourceClaim. Synthetic devices refine placement but are never claimed, so a nodeSelector that matches only synthetic devices leaves the workload nothing to claim - the scheduler treats such a pool as ineligible and the deployment reports InsufficientCapacity. """ + prefill: Prefill | None = None + """ + Prefill role for disaggregated serving. When set, the deployment is disaggregated: the top-level workers is the decode role and this is the prefill role, each self-contained with its own topology, template, and nodeSelector. Prefill and decode of a replica are co-located on one InferenceCluster. Unset means unified serving. + """ replicas: conint(ge=1, le=10) """ How many ModelReplicas to fan out to. Each replica is a complete serving instance scheduled to one InferenceCluster. """ + routing: Routing | None = None + """ + Routing layer for this deployment. Carried through unconsumed for now; the disaggregation backend uses it later to build the endpoint-picker (EPP). template is a curated PodSpec subset, same shape and owner as the engine, defaulting to the llm-d EPP image. + """ workers: Workers """ Compute shape of one worker. Modelplane composes one worker (or workers.count workers) per ModelReplica. diff --git a/schemas/python/models/ai/modelplane/modelreplica/v1alpha1.py b/schemas/python/models/ai/modelplane/modelreplica/v1alpha1.py index bb9161b5d..c38c0000c 100644 --- a/schemas/python/models/ai/modelplane/modelreplica/v1alpha1.py +++ b/schemas/python/models/ai/modelplane/modelreplica/v1alpha1.py @@ -3,7 +3,7 @@ from __future__ import annotations -from typing import Literal +from typing import Any, Literal from pydantic import AwareDatetime, BaseModel, Field, conint, constr @@ -124,11 +124,36 @@ class Topology(BaseModel): class Workers(BaseModel): - count: conint(ge=1) | None = 1 + count: conint(ge=1) | None = None template: Template topology: Topology +class Prefill(BaseModel): + deviceRequests: list[DeviceRequest] = Field(..., max_length=16, min_length=1) + """ + Resolved DRA device requests for the matched pool. The parent ModelDeployment's compose function joins the nodeSelector requests with the matched InferenceClass devices and stamps the claim: DRA devices here. This function turns each into a DeviceRequest in a DRA ResourceClaim for the serving pods. At least one request is always present: the scheduler only pins a replica to a pool that yields a claimable device, so the serving workload always has a ResourceClaim to bind through. + """ + nodePoolName: str + """ + The prefill pool on this replica's cluster. + """ + workers: Workers + + +class TemplateModel(BaseModel): + spec: dict[str, Any] | None = None + + +class Routing(BaseModel): + template: TemplateModel | None = None + + +class TemplateModel1(BaseModel): + metadata: Metadata | None = None + spec: Spec | None = None + + class SpecModel(BaseModel): clusterName: constr(min_length=1) """ @@ -150,6 +175,11 @@ class SpecModel(BaseModel): """ Name of the node pool on the pinned InferenceCluster the scheduler selected for this replica. The scheduler pins every replica to a specific matching pool, so this is always set. """ + prefill: Prefill | None = None + """ + Prefill role placement for a disaggregated replica. Mirrors the top-level decode fields (workers, nodePoolName, deviceRequests), pinned to the prefill pool the scheduler chose. Absent for unified replicas. When present, all three fields are required: compose-model-deployment always populates them, and the llm-d backend reads them unconditionally. + """ + routing: Routing | None = None workers: Workers