native_api.ipynb 11 KB
Newer Older
Chayenne's avatar
Chayenne committed
1
2
3
4
5
6
{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
Chayenne's avatar
Chayenne committed
7
    "# Native APIs\n",
Chayenne's avatar
Chayenne committed
8
    "\n",
Chayenne's avatar
Chayenne committed
9
    "Apart from the OpenAI compatible APIs, the SGLang Runtime also provides its native server APIs. We introduce these following APIs:\n",
Chayenne's avatar
Chayenne committed
10
    "\n",
Chayenne's avatar
Chayenne committed
11
    "- `/generate` (text generation model)\n",
Chayenne's avatar
Chayenne committed
12
13
14
15
16
17
    "- `/get_server_args`\n",
    "- `/get_model_info`\n",
    "- `/health`\n",
    "- `/health_generate`\n",
    "- `/flush_cache`\n",
    "- `/get_memory_pool_size`\n",
18
    "- `/get_max_total_num_tokens`\n",
Chayenne's avatar
Chayenne committed
19
    "- `/update_weights`\n",
Chayenne's avatar
Chayenne committed
20
    "- `/encode`(embedding model)\n",
21
    "- `/classify`(reward model)\n",
Chayenne's avatar
Chayenne committed
22
23
24
25
26
27
28
29
30
31
32
33
34
35
    "\n",
    "We mainly use `requests` to test these APIs in the following examples. You can also use `curl`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Launch A Server"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
36
   "metadata": {},
Chayenne's avatar
Chayenne committed
37
38
39
40
41
42
43
44
45
   "outputs": [],
   "source": [
    "from sglang.utils import (\n",
    "    execute_shell_command,\n",
    "    wait_for_server,\n",
    "    terminate_process,\n",
    "    print_highlight,\n",
    ")\n",
    "\n",
Chayenne's avatar
Chayenne committed
46
47
    "import requests\n",
    "\n",
Chayenne's avatar
Chayenne committed
48
    "server_process = execute_shell_command(\n",
Chayenne's avatar
Chayenne committed
49
    "    \"\"\"\n",
Chayenne's avatar
Chayenne committed
50
51
52
53
54
55
56
57
58
59
60
    "python3 -m sglang.launch_server --model-path meta-llama/Llama-3.2-1B-Instruct --port=30010\n",
    "\"\"\"\n",
    ")\n",
    "\n",
    "wait_for_server(\"http://localhost:30010\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
Chayenne's avatar
Chayenne committed
61
    "## Generate (text generation model)\n",
62
    "Generate completions. This is similar to the `/v1/completions` in OpenAI API. Detailed parameters can be found in the [sampling parameters](../references/sampling_params.md)."
Chayenne's avatar
Chayenne committed
63
64
65
66
67
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
68
   "metadata": {},
Chayenne's avatar
Chayenne committed
69
70
71
   "outputs": [],
   "source": [
    "url = \"http://localhost:30010/generate\"\n",
Chayenne's avatar
Chayenne committed
72
    "data = {\"text\": \"What is the capital of France?\"}\n",
Chayenne's avatar
Chayenne committed
73
74
    "\n",
    "response = requests.post(url, json=data)\n",
Lianmin Zheng's avatar
Lianmin Zheng committed
75
    "print_highlight(response.json())"
Chayenne's avatar
Chayenne committed
76
77
78
79
80
81
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
Chayenne's avatar
Chayenne committed
82
    "## Get Server Args\n",
Lianmin Zheng's avatar
Lianmin Zheng committed
83
    "Get the arguments of a server."
Chayenne's avatar
Chayenne committed
84
85
86
87
88
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
89
   "metadata": {},
Chayenne's avatar
Chayenne committed
90
91
92
93
94
95
96
97
98
99
100
101
102
103
   "outputs": [],
   "source": [
    "url = \"http://localhost:30010/get_server_args\"\n",
    "\n",
    "response = requests.get(url)\n",
    "print_highlight(response.json())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get Model Info\n",
    "\n",
Lianmin Zheng's avatar
Lianmin Zheng committed
104
    "Get the information of the model.\n",
Chayenne's avatar
Chayenne committed
105
106
107
108
109
110
111
112
    "\n",
    "- `model_path`: The path/name of the model.\n",
    "- `is_generation`: Whether the model is used as generation model or embedding model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
113
   "metadata": {},
Chayenne's avatar
Chayenne committed
114
115
116
117
118
119
120
121
   "outputs": [],
   "source": [
    "url = \"http://localhost:30010/get_model_info\"\n",
    "\n",
    "response = requests.get(url)\n",
    "response_json = response.json()\n",
    "print_highlight(response_json)\n",
    "assert response_json[\"model_path\"] == \"meta-llama/Llama-3.2-1B-Instruct\"\n",
Lianmin Zheng's avatar
Lianmin Zheng committed
122
    "assert response_json[\"is_generation\"] is True\n",
Chayenne's avatar
Chayenne committed
123
124
125
126
127
128
129
    "assert response_json.keys() == {\"model_path\", \"is_generation\"}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
Lianmin Zheng's avatar
Lianmin Zheng committed
130
    "## Health Check\n",
Chayenne's avatar
Chayenne committed
131
132
133
134
135
136
137
    "- `/health`: Check the health of the server.\n",
    "- `/health_generate`: Check the health of the server by generating one token."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
138
   "metadata": {},
Chayenne's avatar
Chayenne committed
139
140
141
142
143
144
145
146
147
148
149
   "outputs": [],
   "source": [
    "url = \"http://localhost:30010/health_generate\"\n",
    "\n",
    "response = requests.get(url)\n",
    "print_highlight(response.text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
150
   "metadata": {},
Chayenne's avatar
Chayenne committed
151
152
153
154
155
156
157
158
159
160
161
162
163
164
   "outputs": [],
   "source": [
    "url = \"http://localhost:30010/health\"\n",
    "\n",
    "response = requests.get(url)\n",
    "print_highlight(response.text)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Flush Cache\n",
    "\n",
Lianmin Zheng's avatar
Lianmin Zheng committed
165
    "Flush the radix cache. It will be automatically triggered when the model weights are updated by the `/update_weights` API."
Chayenne's avatar
Chayenne committed
166
167
168
169
170
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
171
   "metadata": {},
Chayenne's avatar
Chayenne committed
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
   "outputs": [],
   "source": [
    "# flush cache\n",
    "\n",
    "url = \"http://localhost:30010/flush_cache\"\n",
    "\n",
    "response = requests.post(url)\n",
    "print_highlight(response.text)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get Memory Pool Size\n",
    "\n",
    "Get the memory pool size in number of tokens.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
194
   "metadata": {},
Chayenne's avatar
Chayenne committed
195
196
197
198
199
200
201
202
203
204
   "outputs": [],
   "source": [
    "# get_memory_pool_size\n",
    "\n",
    "url = \"http://localhost:30010/get_memory_pool_size\"\n",
    "\n",
    "response = requests.get(url)\n",
    "print_highlight(response.text)"
   ]
  },
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get Maximum Total Number of Tokens\n",
    "\n",
    "Exposes the maximum number of tokens SGLang can handle based on the current configuration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get_max_total_num_tokens\n",
    "\n",
    "url = \"http://localhost:30010/get_max_total_num_tokens\"\n",
    "\n",
    "response = requests.get(url)\n",
    "print_highlight(response.text)"
   ]
  },
Chayenne's avatar
Chayenne committed
228
229
230
231
232
233
234
235
236
237
238
239
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Update Weights\n",
    "\n",
    "Update model weights without restarting the server. Use for continuous evaluation during training. Only applicable for models with the same architecture and parameter size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
240
   "metadata": {},
Chayenne's avatar
Chayenne committed
241
242
243
244
245
246
247
248
249
   "outputs": [],
   "source": [
    "# successful update with same architecture and size\n",
    "\n",
    "url = \"http://localhost:30010/update_weights\"\n",
    "data = {\"model_path\": \"meta-llama/Llama-3.2-1B\"}\n",
    "\n",
    "response = requests.post(url, json=data)\n",
    "print_highlight(response.text)\n",
250
    "assert response.json()[\"success\"] is True\n",
Chayenne's avatar
Chayenne committed
251
252
253
254
255
256
257
    "assert response.json()[\"message\"] == \"Succeeded to update model weights.\"\n",
    "assert response.json().keys() == {\"success\", \"message\"}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
258
   "metadata": {},
Chayenne's avatar
Chayenne committed
259
260
261
262
263
264
265
266
267
268
   "outputs": [],
   "source": [
    "# failed update with different parameter size\n",
    "\n",
    "url = \"http://localhost:30010/update_weights\"\n",
    "data = {\"model_path\": \"meta-llama/Llama-3.2-3B\"}\n",
    "\n",
    "response = requests.post(url, json=data)\n",
    "response_json = response.json()\n",
    "print_highlight(response_json)\n",
269
    "assert response_json[\"success\"] is False\n",
Chayenne's avatar
Chayenne committed
270
271
272
273
274
275
276
    "assert response_json[\"message\"] == (\n",
    "    \"Failed to update weights: The size of tensor a (2048) must match \"\n",
    "    \"the size of tensor b (3072) at non-singleton dimension 1.\\n\"\n",
    "    \"Rolling back to original weights.\"\n",
    ")"
   ]
  },
Chayenne's avatar
Chayenne committed
277
278
279
280
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
Chayenne's avatar
Chayenne committed
281
    "## Encode (embedding model)\n",
Chayenne's avatar
Chayenne committed
282
    "\n",
Chayenne's avatar
Chayenne committed
283
284
    "Encode text into embeddings. Note that this API is only available for [embedding models](openai_api_embeddings.html#openai-apis-embedding) and will raise an error for generation models.\n",
    "Therefore, we launch a new server to server an embedding model."
Chayenne's avatar
Chayenne committed
285
286
287
288
289
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
290
   "metadata": {},
Chayenne's avatar
Chayenne committed
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
   "outputs": [],
   "source": [
    "terminate_process(server_process)\n",
    "\n",
    "embedding_process = execute_shell_command(\n",
    "    \"\"\"\n",
    "python -m sglang.launch_server --model-path Alibaba-NLP/gte-Qwen2-7B-instruct \\\n",
    "    --port 30020 --host 0.0.0.0 --is-embedding\n",
    "\"\"\"\n",
    ")\n",
    "\n",
    "wait_for_server(\"http://localhost:30020\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
308
   "metadata": {},
Chayenne's avatar
Chayenne committed
309
310
311
312
313
314
315
316
317
318
319
320
   "outputs": [],
   "source": [
    "# successful encode for embedding model\n",
    "\n",
    "url = \"http://localhost:30020/encode\"\n",
    "data = {\"model\": \"Alibaba-NLP/gte-Qwen2-7B-instruct\", \"text\": \"Once upon a time\"}\n",
    "\n",
    "response = requests.post(url, json=data)\n",
    "response_json = response.json()\n",
    "print_highlight(f\"Text embedding (first 10): {response_json['embedding'][:10]}\")"
   ]
  },
Chayenne's avatar
Chayenne committed
321
322
323
324
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
325
    "## Classify (reward model)\n",
Chayenne's avatar
Chayenne committed
326
    "\n",
327
    "SGLang Runtime also supports reward models. Here we use a reward model to classify the quality of pairwise generations."
Chayenne's avatar
Chayenne committed
328
329
330
331
332
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
333
   "metadata": {},
Chayenne's avatar
Chayenne committed
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
   "outputs": [],
   "source": [
    "terminate_process(embedding_process)\n",
    "\n",
    "# Note that SGLang now treats embedding models and reward models as the same type of models.\n",
    "# This will be updated in the future.\n",
    "\n",
    "reward_process = execute_shell_command(\n",
    "    \"\"\"\n",
    "python -m sglang.launch_server --model-path Skywork/Skywork-Reward-Llama-3.1-8B-v0.2 --port 30030 --host 0.0.0.0 --is-embedding\n",
    "\"\"\"\n",
    ")\n",
    "\n",
    "wait_for_server(\"http://localhost:30030\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
353
   "metadata": {},
Chayenne's avatar
Chayenne committed
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer\n",
    "\n",
    "PROMPT = (\n",
    "    \"What is the range of the numeric output of a sigmoid node in a neural network?\"\n",
    ")\n",
    "\n",
    "RESPONSE1 = \"The output of a sigmoid node is bounded between -1 and 1.\"\n",
    "RESPONSE2 = \"The output of a sigmoid node is bounded between 0 and 1.\"\n",
    "\n",
    "CONVS = [\n",
    "    [{\"role\": \"user\", \"content\": PROMPT}, {\"role\": \"assistant\", \"content\": RESPONSE1}],\n",
    "    [{\"role\": \"user\", \"content\": PROMPT}, {\"role\": \"assistant\", \"content\": RESPONSE2}],\n",
    "]\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"Skywork/Skywork-Reward-Llama-3.1-8B-v0.2\")\n",
    "prompts = tokenizer.apply_chat_template(CONVS, tokenize=False)\n",
    "\n",
373
    "url = \"http://localhost:30030/classify\"\n",
Chayenne's avatar
Chayenne committed
374
    "data = {\"model\": \"Skywork/Skywork-Reward-Llama-3.1-8B-v0.2\", \"text\": prompts}\n",
Chayenne's avatar
Chayenne committed
375
376
377
378
379
380
    "\n",
    "responses = requests.post(url, json=data).json()\n",
    "for response in responses:\n",
    "    print_highlight(f\"reward: {response['embedding'][0]}\")"
   ]
  },
Chayenne's avatar
Chayenne committed
381
382
  {
   "cell_type": "code",
383
384
   "execution_count": null,
   "metadata": {},
Chayenne's avatar
Chayenne committed
385
386
   "outputs": [],
   "source": [
Chayenne's avatar
Chayenne committed
387
    "terminate_process(reward_process)"
Chayenne's avatar
Chayenne committed
388
389
390
391
392
393
394
395
396
397
398
399
400
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
401
   "pygments_lexer": "ipython3"
Chayenne's avatar
Chayenne committed
402
403
404
405
406
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}