Improve type annotation (#1029)

2024-08-11 02:44:59 -07:00
parent fcc0f5ed99
commit 9dae407812
6 changed files with 83 additions and 41 deletions
--- a/python/sglang/srt/mem_cache/memory_pool.py
+++ b/python/sglang/srt/mem_cache/memory_pool.py
@@ -16,7 +16,7 @@ limitations under the License.
 """Memory pool."""

 import logging
-from typing import List
+from typing import List, Union

 import torch

@@ -42,7 +42,7 @@ class ReqToTokenPool:

        return select_index

-    def free(self, free_index):
+    def free(self, free_index: Union[int, List[int]]):
        if isinstance(free_index, (int,)):
            self.free_slots.append(free_index)
        else: