|
| 1 | +package e2e |
| 2 | + |
| 3 | +import ( |
| 4 | + "bytes" |
| 5 | + "encoding/json" |
| 6 | + "fmt" |
| 7 | + "io" |
| 8 | + "net/http" |
| 9 | + "time" |
| 10 | + |
| 11 | + "github.com/google/uuid" |
| 12 | + "github.com/onsi/ginkgo/v2" |
| 13 | + "github.com/onsi/gomega" |
| 14 | + |
| 15 | + testutils "github.com/llm-d/llm-d-router/test/utils" |
| 16 | +) |
| 17 | + |
| 18 | +const ( |
| 19 | + generatePath = "/inference/v1/generate" |
| 20 | + |
| 21 | + requestTimeout = 60 * time.Second |
| 22 | +) |
| 23 | + |
| 24 | +// imageSpec describes one multimodal image entry as it appears in encode |
| 25 | +// and prefill request bodies. Hash, Offset and Length come from the |
| 26 | +// (mocked) render stage; tests pick fixed values so the bodies are |
| 27 | +// deterministic. |
| 28 | +type imageSpec struct { |
| 29 | + Hash string |
| 30 | + Offset int |
| 31 | + Length int |
| 32 | +} |
| 33 | + |
| 34 | +// singleImage is the canonical one-image spec the basic encode/prefill |
| 35 | +// tests share. Offset is always 1 in encode requests (right after BOS), |
| 36 | +// length matches the placeholder span. |
| 37 | +var singleImage = imageSpec{Hash: "e2e-image-hash", Offset: 1, Length: 3} |
| 38 | + |
| 39 | +// twoImages is the canonical two-image spec for fan-out coverage. |
| 40 | +// Distinct lengths so the two placeholder spans aren't accidentally identical. |
| 41 | +var twoImages = []imageSpec{ |
| 42 | + {Hash: "e2e-image-hash-0", Offset: 1, Length: 3}, |
| 43 | + {Hash: "e2e-image-hash-1", Offset: 4, Length: 5}, |
| 44 | +} |
| 45 | + |
| 46 | +var _ = ginkgo.Describe("Direct gateway /inference/v1/generate encode against encode-only", ginkgo.Label(extendedTestLabel), func() { |
| 47 | + // Uses single-profile-handler (generateEncodeConfig) so the EPP routes |
| 48 | + // directly to encode pods without requiring a decode stage first. |
| 49 | + ginkgo.It("returns ec_transfer_params for encode bodies", func() { |
| 50 | + nsName := getNamespace() |
| 51 | + infPoolObjects = createInferencePool(1, true) |
| 52 | + |
| 53 | + encodeReplicas := 1 |
| 54 | + modelServers := createModelServersEncodeOnly(encodeReplicas) |
| 55 | + epp := createEndPointPicker(generateEncodeConfig) |
| 56 | + ginkgo.DeferCleanup(func() { |
| 57 | + testutils.DeleteObjects(testConfig, epp, nsName) |
| 58 | + testutils.DeleteObjects(testConfig, modelServers, nsName) |
| 59 | + }) |
| 60 | + |
| 61 | + encodePods := getPodNames(encodeSelector) |
| 62 | + gomega.Expect(encodePods).Should(gomega.HaveLen(encodeReplicas)) |
| 63 | + |
| 64 | + ginkgo.By("Encode_Generate: single-image encode body returns ec_transfer_params") |
| 65 | + { |
| 66 | + resp, raw := doGenerate(encodeBody(singleImage)) |
| 67 | + parsed := expectGenerateOK(resp, raw) |
| 68 | + expectECTransferParams(parsed, raw) |
| 69 | + } |
| 70 | + |
| 71 | + ginkgo.By("TwoImages_Encode: per-image fan-out, every encode body returns ec_transfer_params") |
| 72 | + for _, img := range twoImages { |
| 73 | + resp, raw := doGenerate(encodeBody(img)) |
| 74 | + parsed := expectGenerateOK(resp, raw) |
| 75 | + expectECTransferParams(parsed, raw) |
| 76 | + } |
| 77 | + }) |
| 78 | +}) |
| 79 | + |
| 80 | +var _ = ginkgo.Describe("Direct gateway /inference/v1/generate prefill against prefill-only", ginkgo.Label(extendedTestLabel), func() { |
| 81 | + // Uses single-profile-handler (generatePrefillConfig) so the EPP routes |
| 82 | + // directly to prefill pods without requiring a decode stage first. |
| 83 | + ginkgo.It("returns kv_transfer_params for prefill bodies", func() { |
| 84 | + nsName := getNamespace() |
| 85 | + infPoolObjects = createInferencePool(1, true) |
| 86 | + |
| 87 | + prefillReplicas := 1 |
| 88 | + modelServers := createModelServersPrefillOnly(prefillReplicas) |
| 89 | + epp := createEndPointPicker(generatePrefillConfig) |
| 90 | + ginkgo.DeferCleanup(func() { |
| 91 | + testutils.DeleteObjects(testConfig, epp, nsName) |
| 92 | + testutils.DeleteObjects(testConfig, modelServers, nsName) |
| 93 | + }) |
| 94 | + |
| 95 | + prefillPods := getPodNames(prefillSelector) |
| 96 | + gomega.Expect(prefillPods).Should(gomega.HaveLen(prefillReplicas)) |
| 97 | + |
| 98 | + ginkgo.By("TwoImages_Prefill: combined two-image prefill body returns kv_transfer_params") |
| 99 | + { |
| 100 | + tokenIDs := []int{1, 32000, 32000, 32000, 32000, 32000, 32000, 32000, 32000, 2345, 6789} |
| 101 | + resp, raw := doGenerate(prefillBody(tokenIDs, twoImages)) |
| 102 | + parsed := expectGenerateOK(resp, raw) |
| 103 | + expectKVTransferParams(parsed, raw) |
| 104 | + } |
| 105 | + }) |
| 106 | +}) |
| 107 | + |
| 108 | +// imageFeatures builds the features map that encode and prefill share: |
| 109 | +// mm_hashes, mm_placeholders, kwargs_data, all keyed by modality |
| 110 | +// ("image"). kwargs_data is a placeholder "AA==" per entry; the |
| 111 | +// simulator does not validate it. |
| 112 | +func imageFeatures(images []imageSpec) map[string]any { |
| 113 | + hashes := make([]string, len(images)) |
| 114 | + placeholders := make([]map[string]any, len(images)) |
| 115 | + kwargs := make([]string, len(images)) |
| 116 | + for i, img := range images { |
| 117 | + hashes[i] = img.Hash |
| 118 | + placeholders[i] = map[string]any{"offset": img.Offset, "length": img.Length} |
| 119 | + kwargs[i] = "AA==" |
| 120 | + } |
| 121 | + return map[string]any{ |
| 122 | + "mm_hashes": map[string]any{"image": hashes}, |
| 123 | + "mm_placeholders": map[string]any{"image": placeholders}, |
| 124 | + "kwargs_data": map[string]any{"image": kwargs}, |
| 125 | + } |
| 126 | +} |
| 127 | + |
| 128 | +// encodeBody builds a single-image encode request. |
| 129 | +// token_ids = [BOS, placeholder*length]; offset is always 1 since each |
| 130 | +// encode carries exactly one image. |
| 131 | +func encodeBody(image imageSpec) []byte { |
| 132 | + tokenIDs := make([]int, 1+image.Length) |
| 133 | + tokenIDs[0] = 1 |
| 134 | + for i := 1; i < len(tokenIDs); i++ { |
| 135 | + tokenIDs[i] = 32000 |
| 136 | + } |
| 137 | + body := map[string]any{ |
| 138 | + "model": simModelName, |
| 139 | + "token_ids": tokenIDs, |
| 140 | + "features": imageFeatures([]imageSpec{{Hash: image.Hash, Offset: 1, Length: image.Length}}), |
| 141 | + "sampling_params": map[string]any{"max_tokens": 1}, |
| 142 | + } |
| 143 | + return mustMarshal(body) |
| 144 | +} |
| 145 | + |
| 146 | +// ecTransferEntries builds the per-image entries of |
| 147 | +// prefill.ec_transfer_params.image, one map per image keyed by mm_hash. |
| 148 | +// Values are dummy NIXL transfer params; the simulator does not |
| 149 | +// validate. |
| 150 | +func ecTransferEntries(images []imageSpec) []map[string]any { |
| 151 | + entries := make([]map[string]any, len(images)) |
| 152 | + for i, img := range images { |
| 153 | + entries[i] = map[string]any{ |
| 154 | + img.Hash: map[string]any{ |
| 155 | + "peer_host": "10.0.0.1", |
| 156 | + "peer_port": 5501 + i, |
| 157 | + "size_bytes": 0, |
| 158 | + "nixl_agent_metadata_b64": "", |
| 159 | + }, |
| 160 | + } |
| 161 | + } |
| 162 | + return entries |
| 163 | +} |
| 164 | + |
| 165 | +// prefillBody builds a prefill request covering every image in one body. |
| 166 | +func prefillBody(tokenIDs []int, images []imageSpec) []byte { |
| 167 | + body := map[string]any{ |
| 168 | + "request_id": "e2e-prefill-" + uuid.NewString(), |
| 169 | + "model": simModelName, |
| 170 | + "token_ids": tokenIDs, |
| 171 | + "features": imageFeatures(images), |
| 172 | + "ec_transfer_params": map[string]any{"image": ecTransferEntries(images)}, |
| 173 | + "sampling_params": map[string]any{ |
| 174 | + "max_tokens": 1, |
| 175 | + "extra_args": map[string]any{ |
| 176 | + "kv_transfer_params": map[string]any{"do_remote_decode": true}, |
| 177 | + }, |
| 178 | + }, |
| 179 | + } |
| 180 | + return mustMarshal(body) |
| 181 | +} |
| 182 | + |
| 183 | +// doRequest POSTs body to <gateway><path>. Always sets Content-Type and a |
| 184 | +// fresh X-Request-ID. Returns the live *http.Response (its body is already |
| 185 | +// drained) and the raw response body. |
| 186 | +func doRequest(path string, body []byte) (*http.Response, []byte) { |
| 187 | + req, err := http.NewRequest(http.MethodPost, fmt.Sprintf("http://localhost:%d%s", getPort(), path), bytes.NewReader(body)) |
| 188 | + gomega.Expect(err).NotTo(gomega.HaveOccurred(), "build POST request") |
| 189 | + req.Header.Set("Content-Type", "application/json") |
| 190 | + req.Header.Set("X-Request-ID", uuid.NewString()) |
| 191 | + |
| 192 | + client := &http.Client{Timeout: requestTimeout} |
| 193 | + resp, err := client.Do(req) |
| 194 | + gomega.Expect(err).NotTo(gomega.HaveOccurred(), "POST %s", req.URL) |
| 195 | + defer func() { |
| 196 | + gomega.Expect(resp.Body.Close()).To(gomega.Succeed()) |
| 197 | + }() |
| 198 | + raw, err := io.ReadAll(resp.Body) |
| 199 | + gomega.Expect(err).NotTo(gomega.HaveOccurred(), "read response body") |
| 200 | + return resp, raw |
| 201 | +} |
| 202 | + |
| 203 | +// doGenerate is a thin wrapper over doRequest targeting /inference/v1/generate. |
| 204 | +func doGenerate(body []byte) (*http.Response, []byte) { |
| 205 | + return doRequest(generatePath, body) |
| 206 | +} |
| 207 | + |
| 208 | +// expectGenerateOK asserts a 2xx status and that the response body parses |
| 209 | +// as a JSON object, returning the parsed map for further phase-specific |
| 210 | +// assertions. |
| 211 | +func expectGenerateOK(resp *http.Response, raw []byte) map[string]any { |
| 212 | + gomega.Expect(resp.StatusCode).To(gomega.Equal(http.StatusOK), |
| 213 | + "non-200 from gateway: status=%d body=%s", resp.StatusCode, string(raw)) |
| 214 | + var parsed map[string]any |
| 215 | + gomega.Expect(json.Unmarshal(raw, &parsed)).To(gomega.Succeed(), |
| 216 | + "response is not valid JSON: %s", string(raw)) |
| 217 | + return parsed |
| 218 | +} |
| 219 | + |
| 220 | +// expectECTransferParams asserts the encode response carries a non-empty |
| 221 | +// ec_transfer_params object keyed by mm_hash. |
| 222 | +func expectECTransferParams(parsed map[string]any, raw []byte) { |
| 223 | + ec, ok := parsed["ec_transfer_params"].(map[string]any) |
| 224 | + gomega.Expect(ok).To(gomega.BeTrue(), |
| 225 | + "missing or malformed ec_transfer_params in encode response: %s", string(raw)) |
| 226 | + gomega.Expect(ec).NotTo(gomega.BeEmpty(), |
| 227 | + "ec_transfer_params is empty: %s", string(raw)) |
| 228 | +} |
| 229 | + |
| 230 | +// expectKVTransferParams asserts the prefill response carries a non-empty |
| 231 | +// kv_transfer_params object (handoff metadata for the decode worker). |
| 232 | +func expectKVTransferParams(parsed map[string]any, raw []byte) { |
| 233 | + kv, ok := parsed["kv_transfer_params"].(map[string]any) |
| 234 | + gomega.Expect(ok).To(gomega.BeTrue(), |
| 235 | + "missing or malformed kv_transfer_params in prefill response: %s", string(raw)) |
| 236 | + gomega.Expect(kv).NotTo(gomega.BeEmpty(), |
| 237 | + "kv_transfer_params is empty: %s", string(raw)) |
| 238 | +} |
| 239 | + |
| 240 | +var _ = ginkgo.Describe("P/D gateway /inference/v1/generate disaggregates via sidecar", ginkgo.Label(sharedStorageTestLabel, disaggTestLabel), func() { |
| 241 | + // Regression test for https://github.com/llm-d/llm-d-router/issues/1461: |
| 242 | + // the pd-sidecar previously had no route for /inference/v1/generate, so |
| 243 | + // token-in P/D requests silently fell through to decode-only. This test |
| 244 | + // verifies that the prefill pod receives the generate request when a full |
| 245 | + // P/D setup is used, proving the sidecar routes it through |
| 246 | + // disaggregatedPrefillHandler rather than the decoder catch-all. |
| 247 | + ginkgo.It("routes token-in generate to the prefill pod", func() { |
| 248 | + nsName := getNamespace() |
| 249 | + infPoolObjects = createInferencePool(1, true) |
| 250 | + |
| 251 | + prefillReplicas := 1 |
| 252 | + decodeReplicas := 1 |
| 253 | + modelServers := createModelServersPDSharedStorage(decodeReplicas) |
| 254 | + epp := createEndPointPicker(pdConfig) |
| 255 | + ginkgo.DeferCleanup(func() { |
| 256 | + testutils.DeleteObjects(testConfig, epp, nsName) |
| 257 | + testutils.DeleteObjects(testConfig, modelServers, nsName) |
| 258 | + }) |
| 259 | + |
| 260 | + prefillPods, decodePods := getModelServerPods(podSelector, prefillSelector, decodeSelector) |
| 261 | + gomega.Expect(prefillPods).Should(gomega.HaveLen(prefillReplicas)) |
| 262 | + gomega.Expect(decodePods).Should(gomega.HaveLen(decodeReplicas)) |
| 263 | + |
| 264 | + prefillCountBefore := getPodRequestCount(nsName, prefillPods[0]) |
| 265 | + ginkgo.By(fmt.Sprintf("prefill request count before test: %d", prefillCountBefore)) |
| 266 | + |
| 267 | + ginkgo.By("sending /inference/v1/generate through P/D EPP") |
| 268 | + resp, raw := doGenerate(simpleTokenGenerateBody()) |
| 269 | + gomega.Expect(resp.StatusCode).To(gomega.Equal(http.StatusOK), |
| 270 | + "non-200 from gateway: status=%d body=%s", resp.StatusCode, string(raw)) |
| 271 | + |
| 272 | + prefillCountAfter := getPodRequestCount(nsName, prefillPods[0]) |
| 273 | + ginkgo.By(fmt.Sprintf("prefill request count after test: %d", prefillCountAfter)) |
| 274 | + |
| 275 | + gomega.Expect(prefillCountAfter).To(gomega.BeNumerically(">", prefillCountBefore), |
| 276 | + "prefill pod should have received the generate request; sidecar must route "+ |
| 277 | + "/inference/v1/generate through disaggregatedPrefillHandler, not the decoder catch-all") |
| 278 | + }) |
| 279 | +}) |
| 280 | + |
| 281 | +// simpleTokenGenerateBody builds a minimal /inference/v1/generate body with |
| 282 | +// enough token IDs to exceed the prefix-based-pd-decider nonCachedTokens |
| 283 | +// threshold (16) and guarantee P/D routing on a cold cache. |
| 284 | +func simpleTokenGenerateBody() []byte { |
| 285 | + tokenIDs := make([]int, 20) |
| 286 | + for i := range tokenIDs { |
| 287 | + tokenIDs[i] = 1000 + i |
| 288 | + } |
| 289 | + return mustMarshal(map[string]any{ |
| 290 | + "model": simModelName, |
| 291 | + "token_ids": tokenIDs, |
| 292 | + "sampling_params": map[string]any{ |
| 293 | + "max_tokens": 1, |
| 294 | + }, |
| 295 | + }) |
| 296 | +} |
| 297 | + |
| 298 | +func mustMarshal(v any) []byte { |
| 299 | + b, err := json.Marshal(v) |
| 300 | + gomega.Expect(err).NotTo(gomega.HaveOccurred(), "json.Marshal failed: %v", v) |
| 301 | + return b |
| 302 | +} |
0 commit comments