".github/vscode:/vscode.git/clone" did not exist on "dd4b60f174a13631dc08b5de6a0aba35633bb497"
test_session_control.py 28.7 KB
Newer Older
1
2
3
"""
Usage:
python3 -m unittest test_session_control.TestSessionControl.test_session_control
4
5
python3 -m unittest test_session_control.TestSessionControl.test_session_control_with_branching
python3 -m unittest test_session_control.TestSessionControl.test_session_control_backtrack_with_abort
6
python3 -m unittest test_session_control.TestSessionControlVision.test_session_control
7
8
"""

9
10
import asyncio
import json
11
12
import unittest

13
import aiohttp
14
15
import requests

16
from sglang.srt.utils import kill_process_tree
17
from sglang.srt.utils.hf_transformers_utils import get_tokenizer
18
19
20
21
from sglang.test.test_utils import (
    DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
    DEFAULT_URL_FOR_TEST,
22
    CustomTestCase,
23
24
25
26
    popen_launch_server,
)


27
28
29
30
def remove_prefix(text: str, prefix: str) -> str:
    return text[len(prefix) :] if text.startswith(prefix) else text


31
class TestSessionControl(unittest.TestCase):
32
33
34
35
36
    @classmethod
    def setUpClass(cls):
        cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
        cls.base_url = DEFAULT_URL_FOR_TEST
        cls.process = popen_launch_server(
37
38
39
40
41
42
43
            cls.model,
            cls.base_url,
            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
            other_args=[
                "--attention-backend",
                "flashinfer",
            ],
44
45
46
47
        )

    @classmethod
    def tearDownClass(cls):
48
        kill_process_tree(cls.process.pid)
49

50
    def test_session_control(self, gen_len=12):
51
52
53
        chunks = [
            "Let me tell you something about France.",
            "The capital of France is",
54
            "The population of the city is",
55
56
57
58
            "A brief history about that city is",
        ]
        tokenizer = get_tokenizer(self.model)
        chunks_ids = [tokenizer.encode(x) for x in chunks]
59
60
61
        for i in range(1, len(chunks_ids)):
            if chunks_ids[i][0] == tokenizer.bos_token_id:
                chunks_ids[i] = chunks_ids[i][1:]
62
63

        # 1. using session control
64
        requests.post(self.base_url + "/flush_cache")
65
66
67
68
69
70
        session_id = requests.post(
            self.base_url + "/open_session",
            json={"capacity_of_str_len": 1000},
        ).json()
        rid = None

71
        # open an existing session, should get session_id as None
72
        ret = requests.post(
73
74
            self.base_url + "/open_session",
            json={"capacity_of_str_len": 1000, "session_id": session_id},
75
76
        )
        self.assertNotEqual(ret.status_code, 200)
77

78
79
        first_rid = None
        outputs_from_session = []
80
81
        logprobs_from_session = []
        cur_logprob_start_len = 0
82
        for i, chunk_ids in enumerate(chunks_ids):
83
            max_new_tokens = gen_len if i > 0 else 1  # prefill only for the first chunk
84
85
86
87
            response = requests.post(
                self.base_url + "/generate",
                json={
                    "input_ids": chunk_ids,
88
89
90
91
92
93
                    "session_params": {
                        "id": session_id,
                        "rid": rid,
                        "offset": -1,
                        "replace": True,
                    },
94
95
                    "sampling_params": {
                        "temperature": 0,
96
                        "max_new_tokens": max_new_tokens,
97
98
                        "no_stop_trim": True,
                        "skip_special_tokens": False,
99
                    },
100
101
                    "return_logprob": True,
                    "logprob_start_len": cur_logprob_start_len - 1,
102
103
104
105
106
107
108
                },
            ).json()
            rid = response["meta_info"]["id"]
            if i == 0:
                first_rid = rid
            if i > 0:
                outputs_from_session.append(response["text"])
109
110
111
112
113
114
115
116
117
                logprobs_from_session.extend(
                    [
                        round(sublist[0], 2)
                        for sublist in response["meta_info"]["output_token_logprobs"]
                    ]
                )
            cur_logprob_start_len += len(chunk_ids) + max_new_tokens

        # query with a logprob_start_len longer than the request, should see error
118
        ret = requests.post(
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
            self.base_url + "/generate",
            json={
                "input_ids": chunk_ids,
                "session_params": {
                    "id": session_id,
                    "rid": rid,
                    "offset": -1,
                    "replace": True,
                },
                "sampling_params": {
                    "temperature": 0,
                    "max_new_tokens": max_new_tokens,
                    "no_stop_trim": True,
                    "skip_special_tokens": False,
                },
                "return_logprob": True,
                "logprob_start_len": cur_logprob_start_len + len(chunk_ids),
            },
137
138
        )
        self.assertNotEqual(ret.status_code, 200)
139
140

        # backtrack to the first request and regenerate
141
        cur_logprob_start_len = 0
142
143
144
145
        response = requests.post(
            self.base_url + "/generate",
            json={
                "input_ids": chunks_ids[-1],
146
147
148
149
150
151
                "session_params": {
                    "id": session_id,
                    "rid": first_rid,
                    "offset": -1,
                    "replace": True,
                },
152
153
                "sampling_params": {
                    "temperature": 0,
154
                    "max_new_tokens": gen_len,
155
156
                    "no_stop_trim": True,
                    "skip_special_tokens": False,
157
                },
158
159
                "return_logprob": True,
                "logprob_start_len": cur_logprob_start_len,
160
161
162
            },
        ).json()
        outputs_from_session.append(response["text"])
163
164
165
166
167
168
        logprobs_from_session.extend(
            [
                round(sublist[0], 2)
                for sublist in response["meta_info"]["output_token_logprobs"]
            ]
        )
169

170
        # query with a non-existing rid (the last one should be disappeared because of backtrack), should see abort
171
        ret = requests.post(
172
173
174
            self.base_url + "/generate",
            json={
                "input_ids": chunks_ids[-1],
175
176
177
178
179
180
                "session_params": {
                    "id": session_id,
                    "rid": rid,
                    "offset": -1,
                    "replace": True,
                },
181
182
                "sampling_params": {
                    "temperature": 0,
183
                    "max_new_tokens": gen_len,
184
185
                    "no_stop_trim": True,
                    "skip_special_tokens": False,
186
                },
187
                "return_logprob": True,
188
            },
189
190
        )
        self.assertNotEqual(ret.status_code, 200)
191
192
193
194
195

        ret = requests.post(
            self.base_url + "/close_session",
            json={"session_id": session_id},
        )
196
        self.assertEqual(ret.status_code, 200)
197
198

        # send a request to a closed session, should see abort
199
        ret = requests.post(
200
201
202
            self.base_url + "/generate",
            json={
                "input_ids": chunks_ids[-1],
203
204
205
206
207
208
                "session_params": {
                    "id": session_id,
                    "rid": first_rid,
                    "offset": -1,
                    "replace": True,
                },
209
210
                "sampling_params": {
                    "temperature": 0,
211
                    "max_new_tokens": gen_len,
212
213
                    "no_stop_trim": True,
                    "skip_special_tokens": False,
214
                },
215
                "return_logprob": True,
216
            },
217
218
        )
        self.assertNotEqual(ret.status_code, 200)
219
220

        # 2. not use session control
221
222
        requests.post(self.base_url + "/flush_cache")

223
224
225
        input_ids_first_req = None
        input_ids = []
        outputs_normal = []
226
        logprobs_normal = []
227
228
229
230
231
232
233
234
235
        for i, chunk_ids in enumerate(chunks_ids):
            input_ids += chunk_ids
            response = requests.post(
                self.base_url + "/generate",
                json={
                    "input_ids": input_ids,
                    "sampling_params": {
                        "temperature": 0,
                        "max_new_tokens": (
236
                            gen_len if i > 0 else 1
237
                        ),  # prefill only for the first chunk
238
239
                        "no_stop_trim": True,
                        "skip_special_tokens": False,
240
                    },
241
                    "return_logprob": True,
242
243
244
                },
            ).json()
            if i > 0:
245
246
247
                output_ids = tokenizer.encode(response["text"])
                if output_ids[0] == tokenizer.bos_token_id:
                    output_ids = output_ids[1:]
248
                input_ids += output_ids[:-1]
249
                outputs_normal.append(response["text"])
250
251
252
253
254
255
                logprobs_normal.extend(
                    [
                        round(sublist[0], 2)
                        for sublist in response["meta_info"]["output_token_logprobs"]
                    ]
                )
256
257
258
259
260
261
262
263
264
265
            if i == 0:
                input_ids_first_req = input_ids.copy()

        input_ids_first_req += chunks_ids[-1]
        response = requests.post(
            self.base_url + "/generate",
            json={
                "input_ids": input_ids_first_req,
                "sampling_params": {
                    "temperature": 0,
266
                    "max_new_tokens": gen_len,
267
268
269
                    "no_stop_trim": True,
                    "skip_special_tokens": False,
                },
270
                "return_logprob": True,
271
272
273
            },
        ).json()
        outputs_normal.append(response["text"])
274
275
276
277
278
279
        logprobs_normal.extend(
            [
                round(sublist[0], 2)
                for sublist in response["meta_info"]["output_token_logprobs"]
            ]
        )
280
281
282
283
284

        print("outputs from chunked queries with session control:")
        print(outputs_from_session)
        print("outputs from normal queries:")
        print(outputs_normal)
285
        self.assertEqual(outputs_from_session, outputs_normal)
286
287
288
289
290
291
292
293
        print("logprobs from chunked queries with session control:")
        print(logprobs_from_session)
        print("logprobs from normal queries:")
        print(logprobs_normal)
        assert len(logprobs_from_session) == len(
            logprobs_normal
        ), "logprobs must have equal length"
        for a, b in zip(logprobs_from_session, logprobs_normal):
294
            assert abs(a - b) <= 0.15, f"logprobs {a} and {b} differ by more than 0.15"
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
    async def async_generate(self, payload):
        url = self.base_url + "/generate"
        async with aiohttp.ClientSession() as session:
            async with session.post(url=url, json=payload) as response:
                assert response.status == 200
                async for chunk_bytes in response.content:
                    chunk_bytes = chunk_bytes.strip()
                    if not chunk_bytes:
                        continue
                    chunk = remove_prefix(chunk_bytes.decode("utf-8"), "data: ")
                    if chunk == "[DONE]":
                        yield "", None, ""
                    else:
                        data = json.loads(chunk)
                        finish_reason = (
                            data["meta_info"]["finish_reason"]["type"]
                            if data["meta_info"]["finish_reason"]
                            else ""
                        )
                        yield data["text"], data["meta_info"]["id"], finish_reason

    async def run_session_control_backtrack_with_abort(self, replace):
        chunks = [
            "Let me tell you something about France.",
            "The capital of France is",
        ]
        tokenizer = get_tokenizer(self.model)
        chunks_ids = [tokenizer.encode(x) for x in chunks]
        for i in range(1, len(chunks_ids)):
            if chunks_ids[i][0] == tokenizer.bos_token_id:
                chunks_ids[i] = chunks_ids[i][1:]

        # 1. using session control
329
        requests.post(self.base_url + "/flush_cache")
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
        session_id = requests.post(
            self.base_url + "/open_session",
            json={"capacity_of_str_len": 1000},
        ).json()
        rid = None

        payload = {
            "input_ids": chunks_ids[0],
            "session_params": {
                "id": session_id,
                "rid": rid,
                "offset": -1,
                "replace": True,
            },
            "sampling_params": {
                "temperature": 0,
                "max_new_tokens": 100,
                "no_stop_trim": True,
                "skip_special_tokens": False,
                "ignore_eos": True,
            },
            "stream": True,
        }
        gen_so_far = ""
        finish_reason = ""
        second_output = ""
        async for chunk, rid, finish_reason_chunk in self.async_generate(payload):
            gen_so_far += chunk
            if finish_reason == "":
                finish_reason = finish_reason_chunk
            if len(gen_so_far) > 50 and second_output == "":
                payload2 = {
                    "input_ids": chunks_ids[1],
                    "session_params": {
                        "id": session_id,
                        "rid": rid,
                        "offset": 50,
                        "replace": replace,
                    },
                    "sampling_params": {
                        "temperature": 0,
                        "max_new_tokens": 32,
                        "no_stop_trim": True,
                        "skip_special_tokens": False,
                    },
                    "stream": False,
                    "stream_output": True,
                }
                response = requests.post(
                    url=self.base_url + "/generate", json=payload2
                ).json()
                second_output = response["text"]
        if replace:
            assert finish_reason == "abort"
        print("first request output:")
        print(gen_so_far)
        print("second request output:")
        print(second_output)

        # close the session
        ret = requests.post(
            self.base_url + "/close_session",
            json={"session_id": session_id},
        )
        assert ret.status_code == 200

        if not replace:
            assert response["meta_info"]["finish_reason"]["type"] == "abort"
        else:
            # 2. not using session control
400
            requests.post(self.base_url + "/flush_cache")
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
            output_ids = tokenizer.encode(gen_so_far)
            if output_ids[0] == tokenizer.bos_token_id:
                output_ids = output_ids[1:]
            input_ids = chunks_ids[0] + output_ids
            input_ids = input_ids[:50] + chunks_ids[1]
            payload = {
                "input_ids": input_ids,
                "sampling_params": {
                    "temperature": 0,
                    "max_new_tokens": 32,
                    "no_stop_trim": True,
                    "skip_special_tokens": False,
                },
                "stream": False,
                "stream_output": True,
            }
            response = requests.post(
                url=self.base_url + "/generate", json=payload
            ).json()
            output_no_session = response["text"]
            print("second request output without session:")
            print(output_no_session)
423
424
425
            assert (
                second_output == output_no_session
            ), f"second_output: {second_output}, output_no_session: {output_no_session}"
426

427
    @unittest.skip("broken")
428
429
430
431
432
433
434
435
436
437
438
    def test_session_control_backtrack_with_abort(self):
        asyncio.run(self.run_session_control_backtrack_with_abort(replace=True))
        asyncio.run(self.run_session_control_backtrack_with_abort(replace=False))

    def run_session_control_with_branching(
        self, root_prompt, chunks_per_step, gen_len=16
    ):
        for x in chunks_per_step:
            assert len(x) == len(chunks_per_step[0])

        # 1. using session control
439
        requests.post(self.base_url + "/flush_cache")
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
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
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
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
534
535
536
537
538
539
540
541
542
543
        session_id = requests.post(
            self.base_url + "/open_session",
            json={"capacity_of_str_len": 1000},
        ).json()

        outputs_from_session = []
        # send the root prompt
        response = requests.post(
            self.base_url + "/generate",
            json={
                "text": root_prompt,
                "session_params": {
                    "id": session_id,
                    "rid": None,
                    "offset": 0,
                    "replace": False,
                },
                "sampling_params": {
                    "temperature": 0,
                    "max_new_tokens": gen_len,
                    "no_stop_trim": True,
                    "skip_special_tokens": False,
                },
            },
        ).json()
        rid_per_branch = [response["meta_info"]["id"]] * len(chunks_per_step[0])
        outputs_from_session.append(response["text"])

        # send the prompts in branches
        for chunks_for_branches in chunks_per_step:
            for j, chunk in enumerate(chunks_for_branches):
                response = requests.post(
                    self.base_url + "/generate",
                    json={
                        "text": chunk,
                        "session_params": {
                            "id": session_id,
                            "rid": rid_per_branch[j],
                            "offset": 0,
                            "replace": False,
                        },
                        "sampling_params": {
                            "temperature": 0,
                            "max_new_tokens": gen_len,
                            "no_stop_trim": True,
                            "skip_special_tokens": False,
                        },
                    },
                ).json()
                rid = response["meta_info"]["id"]
                rid_per_branch[j] = rid
                outputs_from_session.append(response["text"])

        # close the session
        ret = requests.post(
            self.base_url + "/close_session",
            json={"session_id": session_id},
        )
        assert ret.status_code == 200

        # 2. not use session control
        requests.post(self.base_url + "/flush_cache")

        outputs_normal = []
        input_texts = [root_prompt] * len(chunks_per_step[0])
        # send the root prompt
        response = requests.post(
            self.base_url + "/generate",
            json={
                "text": root_prompt,
                "sampling_params": {
                    "temperature": 0,
                    "max_new_tokens": gen_len,
                    "no_stop_trim": True,
                    "skip_special_tokens": False,
                },
            },
        ).json()
        outputs_normal.append(response["text"])
        input_texts = [x + response["text"] for x in input_texts]

        # send the prompts in branches
        for chunks_for_branches in chunks_per_step:
            for j, chunk in enumerate(chunks_for_branches):
                input_texts[j] += chunk
                response = requests.post(
                    self.base_url + "/generate",
                    json={
                        "text": input_texts[j],
                        "sampling_params": {
                            "temperature": 0,
                            "max_new_tokens": gen_len,
                            "no_stop_trim": True,
                            "skip_special_tokens": False,
                        },
                    },
                ).json()
                outputs_normal.append(response["text"])
                input_texts[j] += response["text"]

        print("====== outputs from chunked queries with session control: =======")
        print(outputs_from_session)
        print("====== outputs from normal queries: =======")
        print(outputs_normal)
544
545
546
        assert (
            outputs_from_session == outputs_normal
        ), f"outputs_from_session: {outputs_from_session}, outputs_normal: {outputs_normal}"
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569

    def test_session_control_with_branching(self):
        root_prompt = "First, let me explain in one sentence about AI"
        chunks_per_step = [
            [
                "Then, briefly, the positive side of AI is",
                "But, briefly, AI could be harmful to human",
            ],
            ["For example", "For example"],
        ]
        self.run_session_control_with_branching(
            root_prompt=root_prompt, chunks_per_step=chunks_per_step, gen_len=8
        )

        root_prompt = "I have three apples."
        chunks_per_step = [
            ["I then give one apple to my friend", "My friend give me another apple."],
            ["I still have", "I now have"],
        ]
        self.run_session_control_with_branching(
            root_prompt=root_prompt, chunks_per_step=chunks_per_step, gen_len=8
        )

570

571
@unittest.skip("broken")
572
class TestSessionControlVision(CustomTestCase):
573
574
575
576
577
578
579
580
581
582
583
584
585
    @classmethod
    def setUpClass(cls):
        cls.model = "lmms-lab/llava-onevision-qwen2-7b-ov"
        cls.base_url = DEFAULT_URL_FOR_TEST
        cls.process = popen_launch_server(
            cls.model,
            cls.base_url,
            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
            # other_args={"--disable-radix"},
        )

    @classmethod
    def tearDownClass(cls):
586
        kill_process_tree(cls.process.pid)
587
588
589
590
591

    def test_session_control(self):
        text_chunks = [
            "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n",
            "<|im_start|>user\n<image>\nDescribe this image in a very short sentence.<|im_end|>\n<|im_start|>assistant\n",
592
593
594
            "<|im_start|>user\n<image>\nIs this image same with one of the previous images?<|im_end|>\n<|im_start|>assistant\n",
            "<|im_start|>user\n<image>\nIs this image same with one of the previous images?<|im_end|>\n<|im_start|>assistant\n",
            "<|im_start|>user\nDescribe this image in a very short sentence.<|im_end|>\nassistant:",
595
596
597
598
        ]
        image_chunks = [
            "https://raw.githubusercontent.com/sgl-project/sglang/main/test/lang/example_image.png",
            "https://raw.githubusercontent.com/sgl-project/sglang/main/test/lang/example_image.png",
599
            "https://raw.githubusercontent.com/sgl-project/sglang/main/assets/logo.png",
600
        ]
601

602
603
        self.assertEqual(
            len(text_chunks), len(image_chunks) + 2
604
        )  # the first and the last prompt does not contain images
605
606
        tokenizer = get_tokenizer(self.model)
        text_input_ids = [tokenizer.encode(x) for x in text_chunks]
607
608
609
610
        for i in range(1, len(text_input_ids)):
            if text_input_ids[i][0] == tokenizer.bos_token_id:
                text_input_ids[i] = text_input_ids[i][1:]
        gen_len = 32
611
612

        # 1. using session control
613
        requests.post(self.base_url + "/flush_cache")
614
615
616
617
618
619
        session_id = requests.post(
            self.base_url + "/open_session",
            json={"capacity_of_str_len": 1000},
        ).json()
        rid = None

620
        # open an existing session, should get session_id as None
621
        ret = requests.post(
622
623
            self.base_url + "/open_session",
            json={"capacity_of_str_len": 1000, "session_id": session_id},
624
625
        )
        self.assertNotEqual(ret.status_code, 200)
626

627
628
        first_rid = None
        outputs_from_session = []
629
        for i in range(len(text_input_ids[:-1])):
630
631
632
633
634
635
            response = requests.post(
                self.base_url + "/generate",
                json={
                    "input_ids": text_input_ids[i],
                    "image_data": image_chunks[i - 1] if i > 0 else None,
                    "modalities": ["multi-images"],
636
637
638
639
640
641
                    "session_params": {
                        "id": session_id,
                        "rid": rid,
                        "offset": 0,
                        "replace": True,
                    },
642
643
644
                    "sampling_params": {
                        "temperature": 0,
                        "max_new_tokens": (
645
                            gen_len if i > 0 else 0
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
                        ),  # prefill only for the first chunk
                        "no_stop_trim": True,
                        "skip_special_tokens": False,
                    },
                },
            ).json()
            rid = response["meta_info"]["id"]
            if i == 0:
                first_rid = rid
            if i > 0:
                outputs_from_session.append(response["text"])

        # backtrack to the first request and regenerate
        response = requests.post(
            self.base_url + "/generate",
            json={
                "input_ids": text_input_ids[-1],
663
664
665
666
667
668
                "session_params": {
                    "id": session_id,
                    "rid": first_rid,
                    "offset": 0,
                    "replace": True,
                },
669
670
                "sampling_params": {
                    "temperature": 0,
671
                    "max_new_tokens": gen_len,
672
673
674
675
676
677
678
                    "no_stop_trim": True,
                    "skip_special_tokens": False,
                },
            },
        ).json()
        outputs_from_session.append(response["text"])

679
        # query with a non-existing rid (the last one should be disappeared because of backtrack), should see abort
680
        ret = requests.post(
681
682
683
            self.base_url + "/generate",
            json={
                "input_ids": text_input_ids[-1],
684
685
686
687
688
689
                "session_params": {
                    "id": session_id,
                    "rid": rid,
                    "offset": 0,
                    "replace": True,
                },
690
691
                "sampling_params": {
                    "temperature": 0,
692
                    "max_new_tokens": gen_len,
693
694
695
696
                    "no_stop_trim": True,
                    "skip_special_tokens": False,
                },
            },
697
698
        )
        self.assertNotEqual(ret.status_code, 200)
699
700
701
702
703

        ret = requests.post(
            self.base_url + "/close_session",
            json={"session_id": session_id},
        )
704
        self.assertEqual(ret.status_code, 200)
705
706

        # send a request to a closed session, should see abort
707
        ret = requests.post(
708
709
710
            self.base_url + "/generate",
            json={
                "input_ids": text_input_ids[-1],
711
712
713
714
715
716
                "session_params": {
                    "id": session_id,
                    "rid": first_rid,
                    "offset": 0,
                    "replace": True,
                },
717
718
                "sampling_params": {
                    "temperature": 0,
719
                    "max_new_tokens": gen_len,
720
721
722
723
                    "no_stop_trim": True,
                    "skip_special_tokens": False,
                },
            },
724
725
        )
        self.assertNotEqual(ret.status_code, 200)
726
727

        # 2. not use session control
728
729
        requests.post(self.base_url + "/flush_cache")

730
731
732
        input_ids_first_req = None
        input_ids = []
        outputs_normal = []
733
        for i in range(len(text_input_ids[:-1])):
734
735
736
737
738
739
740
741
742
743
744
            input_ids += text_input_ids[i]
            image_data = image_chunks[:i] if i > 0 else None
            response = requests.post(
                self.base_url + "/generate",
                json={
                    "input_ids": input_ids,
                    "image_data": image_data,
                    "modalities": ["multi-images"],
                    "sampling_params": {
                        "temperature": 0,
                        "max_new_tokens": (
745
                            gen_len if i > 0 else 0
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
                        ),  # prefill only for the first chunk
                        "no_stop_trim": True,
                        "skip_special_tokens": False,
                    },
                },
            ).json()
            if i > 0:
                output_ids = tokenizer.encode(response["text"])
                if output_ids[0] == tokenizer.bos_token_id:
                    output_ids = output_ids[1:]
                input_ids += output_ids
                outputs_normal.append(response["text"])
            if i == 0:
                input_ids_first_req = input_ids.copy()

        input_ids_first_req += text_input_ids[-1]
        response = requests.post(
            self.base_url + "/generate",
            json={
                "input_ids": input_ids_first_req,
                "sampling_params": {
                    "temperature": 0,
768
                    "max_new_tokens": gen_len,
769
770
                    "no_stop_trim": True,
                    "skip_special_tokens": False,
771
772
773
774
775
776
777
778
779
                },
            },
        ).json()
        outputs_normal.append(response["text"])

        print("outputs from chunked queries with session control:")
        print(outputs_from_session)
        print("outputs from normal queries:")
        print(outputs_normal)
780
781
782
        assert (
            outputs_from_session == outputs_normal
        ), f"outputs_from_session: {outputs_from_session}, outputs_normal: {outputs_normal}"
783
784
785
786


if __name__ == "__main__":
    unittest.main()