test_openai_server.py 19.7 KB
Newer Older
1
2
3
4
5
"""
python3 -m unittest test_openai_server.TestOpenAIServer.test_batch
python3 -m unittest test_openai_server.TestOpenAIServer.test_completion

"""
6
import json
7
import time
8
import unittest
9
10

import openai
11

yichuan~'s avatar
yichuan~ committed
12
from sglang.srt.hf_transformers_utils import get_tokenizer
13
from sglang.srt.utils import kill_child_process
14
15
from sglang.test.test_utils import (
    DEFAULT_MODEL_NAME_FOR_TEST,
16
17
    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
    DEFAULT_URL_FOR_TEST,
18
19
    popen_launch_server,
)
20
21
22
23
24


class TestOpenAIServer(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
Ying Sheng's avatar
Ying Sheng committed
25
        cls.model = DEFAULT_MODEL_NAME_FOR_TEST
26
        cls.base_url = DEFAULT_URL_FOR_TEST
27
28
        cls.api_key = "sk-123456"
        cls.process = popen_launch_server(
29
30
31
32
            cls.model,
            cls.base_url,
            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
            api_key=cls.api_key,
33
        )
34
        cls.base_url += "/v1"
Ying Sheng's avatar
Ying Sheng committed
35
        cls.tokenizer = get_tokenizer(DEFAULT_MODEL_NAME_FOR_TEST)
36
37
38

    @classmethod
    def tearDownClass(cls):
Lianmin Zheng's avatar
Lianmin Zheng committed
39
        kill_child_process(cls.process.pid, include_self=True)
40

yichuan~'s avatar
yichuan~ committed
41
42
43
    def run_completion(
        self, echo, logprobs, use_list_input, parallel_sample_num, token_input
    ):
44
        client = openai.Client(api_key=self.api_key, base_url=self.base_url)
45
        prompt = "The capital of France is"
yichuan~'s avatar
yichuan~ committed
46
47
48
49
50
51
        if token_input:
            prompt_input = self.tokenizer.encode(prompt)
            num_prompt_tokens = len(prompt_input)
        else:
            prompt_input = prompt
            num_prompt_tokens = len(self.tokenizer.encode(prompt))
52
53

        if use_list_input:
yichuan~'s avatar
yichuan~ committed
54
            prompt_arg = [prompt_input, prompt_input]
55
            num_choices = len(prompt_arg)
yichuan~'s avatar
yichuan~ committed
56
            num_prompt_tokens *= 2
57
        else:
yichuan~'s avatar
yichuan~ committed
58
            prompt_arg = prompt_input
59
60
            num_choices = 1

61
62
        response = client.completions.create(
            model=self.model,
63
            prompt=prompt_arg,
yichuan~'s avatar
yichuan~ committed
64
            temperature=0,
65
66
67
            max_tokens=32,
            echo=echo,
            logprobs=logprobs,
yichuan~'s avatar
yichuan~ committed
68
            n=parallel_sample_num,
69
        )
70

yichuan~'s avatar
yichuan~ committed
71
        assert len(response.choices) == num_choices * parallel_sample_num
72

Cody Yu's avatar
Cody Yu committed
73
        if echo:
74
            text = response.choices[0].text
75
            assert text.startswith(prompt)
yichuan~'s avatar
yichuan~ committed
76

Cody Yu's avatar
Cody Yu committed
77
        if logprobs:
78
79
80
            assert response.choices[0].logprobs
            assert isinstance(response.choices[0].logprobs.tokens[0], str)
            assert isinstance(response.choices[0].logprobs.top_logprobs[1], dict)
81
            ret_num_top_logprobs = len(response.choices[0].logprobs.top_logprobs[1])
82

83
            # FIXME: Sometimes, some top_logprobs are missing in the return value. The reason is that some output id maps to the same output token and duplicate in the map
84
            # assert ret_num_top_logprobs == logprobs, f"{ret_num_top_logprobs} vs {logprobs}"
yichuan~'s avatar
yichuan~ committed
85
            assert ret_num_top_logprobs > 0
86
87

            assert response.choices[0].logprobs.token_logprobs[0]
yichuan~'s avatar
yichuan~ committed
88

89
90
        assert response.id
        assert response.created
yichuan~'s avatar
yichuan~ committed
91
92
93
        assert (
            response.usage.prompt_tokens == num_prompt_tokens
        ), f"{response.usage.prompt_tokens} vs {num_prompt_tokens}"
94
95
96
        assert response.usage.completion_tokens > 0
        assert response.usage.total_tokens > 0

97
98
99
    def run_completion_stream(
        self, echo, logprobs, use_list_input, parallel_sample_num, token_input
    ):
100
        client = openai.Client(api_key=self.api_key, base_url=self.base_url)
101
        prompt = "The capital of France is"
yichuan~'s avatar
yichuan~ committed
102
        if token_input:
103
104
            prompt_input = self.tokenizer.encode(prompt)
            num_prompt_tokens = len(prompt_input)
yichuan~'s avatar
yichuan~ committed
105
        else:
106
107
108
109
110
111
112
113
114
115
116
            prompt_input = prompt
            num_prompt_tokens = len(self.tokenizer.encode(prompt))

        if use_list_input:
            prompt_arg = [prompt_input, prompt_input]
            num_choices = len(prompt_arg)
            num_prompt_tokens *= 2
        else:
            prompt_arg = prompt_input
            num_choices = 1

117
118
        generator = client.completions.create(
            model=self.model,
yichuan~'s avatar
yichuan~ committed
119
120
            prompt=prompt_arg,
            temperature=0,
121
122
123
124
            max_tokens=32,
            echo=echo,
            logprobs=logprobs,
            stream=True,
125
            stream_options={"include_usage": True},
126
            n=parallel_sample_num,
127
128
        )

129
        is_firsts = {}
130
        for response in generator:
131
132
133
134
135
136
            usage = response.usage
            if usage is not None:
                assert usage.prompt_tokens > 0
                assert usage.completion_tokens > 0
                assert usage.total_tokens > 0
                continue
137
138
139
140

            index = response.choices[0].index
            is_first = is_firsts.get(index, True)

141
142
143
            if logprobs:
                assert response.choices[0].logprobs
                assert isinstance(response.choices[0].logprobs.tokens[0], str)
144
                if not (is_first and echo):
145
146
147
148
149
150
                    assert isinstance(
                        response.choices[0].logprobs.top_logprobs[0], dict
                    )
                    ret_num_top_logprobs = len(
                        response.choices[0].logprobs.top_logprobs[0]
                    )
151
                    # FIXME: Sometimes, some top_logprobs are missing in the return value. The reason is that some output id maps to the same output token and duplicate in the map
152
                    # assert ret_num_top_logprobs == logprobs, f"{ret_num_top_logprobs} vs {logprobs}"
yichuan~'s avatar
yichuan~ committed
153
                    assert ret_num_top_logprobs > 0
154

155
            if is_first:
156
                if echo:
yichuan~'s avatar
yichuan~ committed
157
158
                    assert response.choices[0].text.startswith(
                        prompt
159
160
                    ), f"{response.choices[0].text} and all args {echo} {logprobs} {token_input} {is_first}"
                is_firsts[index] = False
161
162
163
            assert response.id
            assert response.created

164
165
166
167
168
        for index in [i for i in range(parallel_sample_num * num_choices)]:
            assert not is_firsts.get(
                index, True
            ), f"index {index} is not found in the response"

yichuan~'s avatar
yichuan~ committed
169
    def run_chat_completion(self, logprobs, parallel_sample_num):
170
        client = openai.Client(api_key=self.api_key, base_url=self.base_url)
171
172
173
174
        response = client.chat.completions.create(
            model=self.model,
            messages=[
                {"role": "system", "content": "You are a helpful AI assistant"},
Ying Sheng's avatar
Ying Sheng committed
175
176
177
178
                {
                    "role": "user",
                    "content": "What is the capital of France? Answer in a few words.",
                },
179
180
181
182
            ],
            temperature=0,
            logprobs=logprobs is not None and logprobs > 0,
            top_logprobs=logprobs,
yichuan~'s avatar
yichuan~ committed
183
            n=parallel_sample_num,
184
        )
Ying Sheng's avatar
Ying Sheng committed
185

186
187
188
189
190
191
192
193
194
195
196
        if logprobs:
            assert isinstance(
                response.choices[0].logprobs.content[0].top_logprobs[0].token, str
            )

            ret_num_top_logprobs = len(
                response.choices[0].logprobs.content[0].top_logprobs
            )
            assert (
                ret_num_top_logprobs == logprobs
            ), f"{ret_num_top_logprobs} vs {logprobs}"
Ying Sheng's avatar
Ying Sheng committed
197

yichuan~'s avatar
yichuan~ committed
198
        assert len(response.choices) == parallel_sample_num
199
200
201
202
203
204
205
206
        assert response.choices[0].message.role == "assistant"
        assert isinstance(response.choices[0].message.content, str)
        assert response.id
        assert response.created
        assert response.usage.prompt_tokens > 0
        assert response.usage.completion_tokens > 0
        assert response.usage.total_tokens > 0

207
    def run_chat_completion_stream(self, logprobs, parallel_sample_num=1):
208
        client = openai.Client(api_key=self.api_key, base_url=self.base_url)
209
210
211
212
213
214
215
216
217
218
        generator = client.chat.completions.create(
            model=self.model,
            messages=[
                {"role": "system", "content": "You are a helpful AI assistant"},
                {"role": "user", "content": "What is the capital of France?"},
            ],
            temperature=0,
            logprobs=logprobs is not None and logprobs > 0,
            top_logprobs=logprobs,
            stream=True,
219
            stream_options={"include_usage": True},
220
            n=parallel_sample_num,
221
222
        )

223
        is_firsts = {}
224
        for response in generator:
225
226
227
228
229
230
231
            usage = response.usage
            if usage is not None:
                assert usage.prompt_tokens > 0
                assert usage.completion_tokens > 0
                assert usage.total_tokens > 0
                continue

232
            index = response.choices[0].index
233
            data = response.choices[0].delta
234

235
236
237
            if is_firsts.get(index, True):
                assert data.role == "assistant"
                is_firsts[index] = False
238
239
240
                continue

            if logprobs:
yichuan~'s avatar
yichuan~ committed
241
242
243
244
245
246
247
248
249
250
251
252
253
                assert response.choices[0].logprobs
                assert isinstance(
                    response.choices[0].logprobs.content[0].top_logprobs[0].token, str
                )
                assert isinstance(
                    response.choices[0].logprobs.content[0].top_logprobs, list
                )
                ret_num_top_logprobs = len(
                    response.choices[0].logprobs.content[0].top_logprobs
                )
                assert (
                    ret_num_top_logprobs == logprobs
                ), f"{ret_num_top_logprobs} vs {logprobs}"
254
255
256
257
258

            assert isinstance(data.content, str)
            assert response.id
            assert response.created

259
260
261
262
263
        for index in [i for i in range(parallel_sample_num)]:
            assert not is_firsts.get(
                index, True
            ), f"index {index} is not found in the response"

264
    def _create_batch(self, mode, client):
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
        if mode == "completion":
            input_file_path = "complete_input.jsonl"
            # write content to input file
            content = [
                {
                    "custom_id": "request-1",
                    "method": "POST",
                    "url": "/v1/completions",
                    "body": {
                        "model": "gpt-3.5-turbo-instruct",
                        "prompt": "List 3 names of famous soccer player: ",
                        "max_tokens": 20,
                    },
                },
                {
                    "custom_id": "request-2",
                    "method": "POST",
                    "url": "/v1/completions",
                    "body": {
                        "model": "gpt-3.5-turbo-instruct",
                        "prompt": "List 6 names of famous basketball player:  ",
                        "max_tokens": 40,
                    },
                },
                {
                    "custom_id": "request-3",
                    "method": "POST",
                    "url": "/v1/completions",
                    "body": {
                        "model": "gpt-3.5-turbo-instruct",
                        "prompt": "List 6 names of famous tenniss player:  ",
                        "max_tokens": 40,
                    },
                },
            ]

        else:
            input_file_path = "chat_input.jsonl"
            content = [
                {
                    "custom_id": "request-1",
                    "method": "POST",
                    "url": "/v1/chat/completions",
                    "body": {
                        "model": "gpt-3.5-turbo-0125",
                        "messages": [
                            {
                                "role": "system",
                                "content": "You are a helpful assistant.",
                            },
                            {
                                "role": "user",
                                "content": "Hello! List 3 NBA players and tell a story",
                            },
                        ],
                        "max_tokens": 30,
                    },
                },
                {
                    "custom_id": "request-2",
                    "method": "POST",
                    "url": "/v1/chat/completions",
                    "body": {
                        "model": "gpt-3.5-turbo-0125",
                        "messages": [
                            {"role": "system", "content": "You are an assistant. "},
                            {
                                "role": "user",
                                "content": "Hello! List three capital and tell a story",
                            },
                        ],
                        "max_tokens": 50,
                    },
                },
            ]
340

341
342
343
        with open(input_file_path, "w") as file:
            for line in content:
                file.write(json.dumps(line) + "\n")
344

345
346
347
348
349
350
351
352
353
354
355
356
        with open(input_file_path, "rb") as file:
            uploaded_file = client.files.create(file=file, purpose="batch")
        if mode == "completion":
            endpoint = "/v1/completions"
        elif mode == "chat":
            endpoint = "/v1/chat/completions"
        completion_window = "24h"
        batch_job = client.batches.create(
            input_file_id=uploaded_file.id,
            endpoint=endpoint,
            completion_window=completion_window,
        )
357

358
        return batch_job, content, uploaded_file
359
360
361

    def run_batch(self, mode):
        client = openai.Client(api_key=self.api_key, base_url=self.base_url)
362
        batch_job, content, uploaded_file = self._create_batch(mode=mode, client=client)
363

364
365
366
367
368
369
        while batch_job.status not in ["completed", "failed", "cancelled"]:
            time.sleep(3)
            print(
                f"Batch job status: {batch_job.status}...trying again in 3 seconds..."
            )
            batch_job = client.batches.retrieve(batch_job.id)
370
371
372
        assert (
            batch_job.status == "completed"
        ), f"Batch job status is not completed: {batch_job.status}"
373
374
375
376
377
378
        assert batch_job.request_counts.completed == len(content)
        assert batch_job.request_counts.failed == 0
        assert batch_job.request_counts.total == len(content)

        result_file_id = batch_job.output_file_id
        file_response = client.files.content(result_file_id)
yichuan~'s avatar
yichuan~ committed
379
380
381
382
383
384
        result_content = file_response.read().decode("utf-8")  # Decode bytes to string
        results = [
            json.loads(line)
            for line in result_content.split("\n")
            if line.strip() != ""
        ]
385
        assert len(results) == len(content)
386
387
388
        for delete_fid in [uploaded_file.id, result_file_id]:
            del_pesponse = client.files.delete(delete_fid)
            assert del_pesponse.deleted
389

390
391
    def run_cancel_batch(self, mode):
        client = openai.Client(api_key=self.api_key, base_url=self.base_url)
392
        batch_job, _, uploaded_file = self._create_batch(mode=mode, client=client)
393
394
395
396
397
398
399
400
401
402
403
404
405
406

        assert batch_job.status not in ["cancelling", "cancelled"]

        batch_job = client.batches.cancel(batch_id=batch_job.id)
        assert batch_job.status == "cancelling"

        while batch_job.status not in ["failed", "cancelled"]:
            batch_job = client.batches.retrieve(batch_job.id)
            print(
                f"Batch job status: {batch_job.status}...trying again in 3 seconds..."
            )
            time.sleep(3)

        assert batch_job.status == "cancelled"
407
408
        del_response = client.files.delete(uploaded_file.id)
        assert del_response.deleted
409

410
411
412
    def test_completion(self):
        for echo in [False, True]:
            for logprobs in [None, 5]:
413
                for use_list_input in [True, False]:
yichuan~'s avatar
yichuan~ committed
414
415
416
417
418
419
420
421
422
                    for parallel_sample_num in [1, 2]:
                        for token_input in [False, True]:
                            self.run_completion(
                                echo,
                                logprobs,
                                use_list_input,
                                parallel_sample_num,
                                token_input,
                            )
423
424

    def test_completion_stream(self):
yichuan~'s avatar
yichuan~ committed
425
        # parallel sampling adn list input are not supported in streaming mode
426
427
        for echo in [False, True]:
            for logprobs in [None, 5]:
428
429
430
431
432
433
434
435
436
437
                for use_list_input in [True, False]:
                    for parallel_sample_num in [1, 2]:
                        for token_input in [False, True]:
                            self.run_completion_stream(
                                echo,
                                logprobs,
                                use_list_input,
                                parallel_sample_num,
                                token_input,
                            )
438

439
440
    def test_chat_completion(self):
        for logprobs in [None, 5]:
yichuan~'s avatar
yichuan~ committed
441
442
            for parallel_sample_num in [1, 2]:
                self.run_chat_completion(logprobs, parallel_sample_num)
443
444
445

    def test_chat_completion_stream(self):
        for logprobs in [None, 5]:
446
447
            for parallel_sample_num in [1, 2]:
                self.run_chat_completion_stream(logprobs, parallel_sample_num)
448

449
450
451
452
    def test_batch(self):
        for mode in ["completion", "chat"]:
            self.run_batch(mode)

453
    def test_cancel_batch(self):
454
455
456
        for mode in ["completion", "chat"]:
            self.run_cancel_batch(mode)

457
    def test_regex(self):
458
        client = openai.Client(api_key=self.api_key, base_url=self.base_url)
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486

        regex = (
            r"""\{\n"""
            + r"""   "name": "[\w]+",\n"""
            + r"""   "population": [\d]+\n"""
            + r"""\}"""
        )

        response = client.chat.completions.create(
            model=self.model,
            messages=[
                {"role": "system", "content": "You are a helpful AI assistant"},
                {"role": "user", "content": "Introduce the capital of France."},
            ],
            temperature=0,
            max_tokens=128,
            extra_body={"regex": regex},
        )
        text = response.choices[0].message.content

        try:
            js_obj = json.loads(text)
        except (TypeError, json.decoder.JSONDecodeError):
            print("JSONDecodeError", text)
            raise
        assert isinstance(js_obj["name"], str)
        assert isinstance(js_obj["population"], int)

487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
    def test_penalty(self):
        client = openai.Client(api_key=self.api_key, base_url=self.base_url)

        response = client.chat.completions.create(
            model=self.model,
            messages=[
                {"role": "system", "content": "You are a helpful AI assistant"},
                {"role": "user", "content": "Introduce the capital of France."},
            ],
            temperature=0,
            max_tokens=32,
            frequency_penalty=1.0,
        )
        text = response.choices[0].message.content
        assert isinstance(text, str)

503
504
505
506
    def test_response_prefill(self):
        client = openai.Client(api_key=self.api_key, base_url=self.base_url)

        response = client.chat.completions.create(
507
            model="meta-llama/Llama-3.1-8B-Instruct",
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
            messages=[
                {"role": "system", "content": "You are a helpful AI assistant"},
                {
                    "role": "user",
                    "content": """
Extract the name, size, price, and color from this product description as a JSON object:

<description>
The SmartHome Mini is a compact smart home assistant available in black or white for only $49.99. At just 5 inches wide, it lets you control lights, thermostats, and other connected devices via voice or app—no matter where you place it in your home. This affordable little hub brings convenient hands-free control to your smart devices.
</description>
""",
                },
                {
                    "role": "assistant",
                    "content": "{\n",
                },
            ],
            temperature=0,
        )

        assert (
            response.choices[0]
            .message.content.strip()
            .startswith('"name": "SmartHome Mini",')
        )

534

535
if __name__ == "__main__":
Lianmin Zheng's avatar
Lianmin Zheng committed
536
    unittest.main()