test_update_weights_from_disk.py 3.35 KB
Newer Older
1
2
3
4
5
import json
import unittest

import requests

6
from sglang.srt.utils import kill_process_tree
7
from sglang.test.test_utils import (
Lianmin Zheng's avatar
Lianmin Zheng committed
8
    DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
9
10
    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
    DEFAULT_URL_FOR_TEST,
11
12
13
14
    popen_launch_server,
)


15
class TestUpdateWeights(unittest.TestCase):
16
17
    @classmethod
    def setUpClass(cls):
Lianmin Zheng's avatar
Lianmin Zheng committed
18
        cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
19
20
21
22
        cls.base_url = DEFAULT_URL_FOR_TEST
        cls.process = popen_launch_server(
            cls.model, cls.base_url, timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH
        )
23
24
25

    @classmethod
    def tearDownClass(cls):
26
        kill_process_tree(cls.process.pid)
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51

    def run_decode(self):
        response = requests.post(
            self.base_url + "/generate",
            json={
                "text": "The capital of France is",
                "sampling_params": {
                    "temperature": 0,
                    "max_new_tokens": 32,
                },
            },
        )
        print(json.dumps(response.json()))
        print("=" * 100)
        text = response.json()["text"]
        return text

    def get_model_info(self):
        response = requests.get(self.base_url + "/get_model_info")
        model_path = response.json()["model_path"]
        print(json.dumps(response.json()))
        return model_path

    def run_update_weights(self, model_path):
        response = requests.post(
Chayenne's avatar
Chayenne committed
52
            self.base_url + "/update_weights_from_disk",
53
54
55
56
            json={
                "model_path": model_path,
            },
        )
Lianmin Zheng's avatar
Lianmin Zheng committed
57
        ret = response.json()
58
        print(json.dumps(response.json()))
Lianmin Zheng's avatar
Lianmin Zheng committed
59
        return ret
60

61
    def test_update_weights(self):
62
63
64
65
66
        origin_model_path = self.get_model_info()
        print(f"origin_model_path: {origin_model_path}")
        origin_response = self.run_decode()

        # update weights
Lianmin Zheng's avatar
Lianmin Zheng committed
67
        new_model_path = DEFAULT_SMALL_MODEL_NAME_FOR_TEST.replace("-Instruct", "")
Lianmin Zheng's avatar
Lianmin Zheng committed
68
69
        ret = self.run_update_weights(new_model_path)
        assert ret["success"]
70
71
72
73
74
75
76
77
78
79

        updated_model_path = self.get_model_info()
        print(f"updated_model_path: {updated_model_path}")
        assert updated_model_path == new_model_path
        assert updated_model_path != origin_model_path

        updated_response = self.run_decode()
        assert origin_response[:32] != updated_response[:32]

        # update weights back
Lianmin Zheng's avatar
Lianmin Zheng committed
80
81
82
        ret = self.run_update_weights(origin_model_path)
        assert ret["success"]

83
84
85
86
87
88
        updated_model_path = self.get_model_info()
        assert updated_model_path == origin_model_path

        updated_response = self.run_decode()
        assert origin_response[:32] == updated_response[:32]

89
    def test_update_weights_unexist_model(self):
90
91
92
93
94
        origin_model_path = self.get_model_info()
        print(f"origin_model_path: {origin_model_path}")
        origin_response = self.run_decode()

        # update weights
Lianmin Zheng's avatar
Lianmin Zheng committed
95
        new_model_path = DEFAULT_SMALL_MODEL_NAME_FOR_TEST.replace("-Instruct", "wrong")
Lianmin Zheng's avatar
Lianmin Zheng committed
96
97
        ret = self.run_update_weights(new_model_path)
        assert not ret["success"]
98
99
100
101
102
103
104
105
106
107
108

        updated_model_path = self.get_model_info()
        print(f"updated_model_path: {updated_model_path}")
        assert updated_model_path == origin_model_path

        updated_response = self.run_decode()
        assert origin_response[:32] == updated_response[:32]


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