@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 118 lines 3.82 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Disaggregated multimodal serving: render → generate round-trip.
4
5 Demonstrates the two-phase disaggregated flow:
6 1. /v1/chat/completions/render – preprocesses a multimodal chat request
7 into token IDs and serialized tensor features.
8 2. /inference/v1/generate – runs inference on the preprocessed tokens.
9
10 The render response is passed *directly* to generate with only
11 ``sampling_params`` added, showing that the two endpoints compose with
12 zero client-side transformation.
13
14 Launch the server first:
15
16 vllm serve Qwen/Qwen3-VL-2B-Instruct \
17 --dtype bfloat16 --max-model-len 4096 --enforce-eager \
18 --enable-scale-out
19
20 Then run this script:
21
22 python example_mm_serve.py
23 """
24
25 import io
26
27 import pybase64 as base64
28 import requests
29 from PIL import Image
30 from transformers import AutoTokenizer
31
32 BASE_URL = "http://localhost:8000"
33 MODEL_NAME = "Qwen/Qwen3-VL-2B-Instruct"
34
35
36 def make_data_url(image: Image.Image) -> str:
37 """Encode a PIL image as a base64 data URL."""
38 buf = io.BytesIO()
39 image.save(buf, format="PNG")
40 b64 = base64.b64encode(buf.getvalue()).decode()
41 return f"data:image/png;base64,{b64}"
42
43
44 def main():
45 # -- Step 1: Create a test image (solid red) -------------------------
46 image = Image.new("RGB", (224, 224), color=(255, 0, 0))
47 data_url = make_data_url(image)
48 print("Created 224x224 red test image")
49
50 # -- Step 2: Render (preprocess) -------------------------------------
51 render_payload = {
52 "model": MODEL_NAME,
53 "messages": [
54 {
55 "role": "user",
56 "content": [
57 {"type": "image_url", "image_url": {"url": data_url}},
58 {
59 "type": "text",
60 "text": "What color is this image? Answer in one word.",
61 },
62 ],
63 }
64 ],
65 }
66
67 print("\n--- Render ---")
68 render_resp = requests.post(
69 f"{BASE_URL}/v1/chat/completions/render", json=render_payload
70 )
71 render_resp.raise_for_status()
72 render_data = render_resp.json()
73
74 print(f"Response keys: {list(render_data.keys())}")
75 print(f"Number of token_ids: {len(render_data['token_ids'])}")
76
77 features = render_data.get("features")
78 if features and features.get("kwargs_data"):
79 print(f"kwargs_data modalities: {list(features['kwargs_data'].keys())}")
80 for modality, items in features["kwargs_data"].items():
81 print(
82 f" {modality}: {len(items)} item(s), "
83 f"first item type: {type(items[0])} length: {len(items[0])}"
84 if items
85 else "First item: (empty)"
86 )
87 else:
88 print("WARNING: no kwargs_data in render response")
89
90 # -- Step 3: Generate (inference) ------------------------------------
91 # Pass the render output directly — only add sampling_params.
92 generate_payload = render_data
93 generate_payload["sampling_params"] = {
94 "max_tokens": 20,
95 "temperature": 0.0,
96 }
97
98 print("\n--- Generate ---")
99 gen_resp = requests.post(f"{BASE_URL}/inference/v1/generate", json=generate_payload)
100 gen_resp.raise_for_status()
101 gen_data = gen_resp.json()
102
103 # -- Step 4: Decode & print ------------------------------------------
104 output_ids = gen_data["choices"][0]["token_ids"]
105 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
106 text = tokenizer.decode(output_ids, skip_special_tokens=True)
107
108 print(f"Output token count: {len(output_ids)}")
109 print(f"Generated text: {text!r}")
110
111 if "red" in text.lower():
112 print("\nModel correctly identified the red image.")
113 else:
114 print(f"\nWARNING: Expected 'red' in output, got: {text!r}")
115
116
117 if __name__ == "__main__":
118 main()