[mypy] Misc. typing improvements (#7417)

tests/tensorizer_loader/conftest.py

@@ -1,10 +1,12 @@
 import contextlib
 import functools
 import gc
+from typing import Callable, TypeVar
 
 import pytest
 import ray
 import torch
+from typing_extensions import ParamSpec
 
 from vllm.distributed import (destroy_distributed_environment,
                               destroy_model_parallel)
@@ -22,12 +24,16 @@ def cleanup():
     torch.cuda.empty_cache()
 
 
-def retry_until_skip(n):
+_P = ParamSpec("_P")
+_R = TypeVar("_R")
 
-    def decorator_retry(func):
+
+def retry_until_skip(n: int):
+
+    def decorator_retry(func: Callable[_P, _R]) -> Callable[_P, _R]:
 
         @functools.wraps(func)
-        def wrapper_retry(*args, **kwargs):
+        def wrapper_retry(*args: _P.args, **kwargs: _P.kwargs) -> _R:
             for i in range(n):
                 try:
                     return func(*args, **kwargs)
@@ -35,7 +41,9 @@ def retry_until_skip(n):
                     gc.collect()
                     torch.cuda.empty_cache()
                     if i == n - 1:
-                        pytest.skip("Skipping test after attempts..")
+                        pytest.skip(f"Skipping test after {n} attempts.")
+
+            raise AssertionError("Code should not be reached")
 
         return wrapper_retry
 

tests/test_utils.py

@@ -1,10 +1,8 @@
 import asyncio
 import os
 import socket
-import sys
 from functools import partial
-from typing import (TYPE_CHECKING, Any, AsyncIterator, Awaitable, Protocol,
-                    Tuple, TypeVar)
+from typing import AsyncIterator, Tuple
 
 import pytest
 
@@ -13,26 +11,11 @@ from vllm.utils import (FlexibleArgumentParser, deprecate_kwargs,
 
 from .utils import error_on_warning
 
-if sys.version_info < (3, 10):
-    if TYPE_CHECKING:
-        _AwaitableT = TypeVar("_AwaitableT", bound=Awaitable[Any])
-        _AwaitableT_co = TypeVar("_AwaitableT_co",
-                                 bound=Awaitable[Any],
-                                 covariant=True)
-
-        class _SupportsSynchronousAnext(Protocol[_AwaitableT_co]):
-
-            def __anext__(self) -> _AwaitableT_co:
-                ...
-
-    def anext(i: "_SupportsSynchronousAnext[_AwaitableT]", /) -> "_AwaitableT":
-        return i.__anext__()
-
 
 @pytest.mark.asyncio
 async def test_merge_async_iterators():
 
-    async def mock_async_iterator(idx: int) -> AsyncIterator[str]:
+    async def mock_async_iterator(idx: int):
         try:
             while True:
                 yield f"item from iterator {idx}"
@@ -41,8 +24,10 @@ async def test_merge_async_iterators():
             print(f"iterator {idx} cancelled")
 
     iterators = [mock_async_iterator(i) for i in range(3)]
-    merged_iterator: AsyncIterator[Tuple[int, str]] = merge_async_iterators(
-        *iterators, is_cancelled=partial(asyncio.sleep, 0, result=False))
+    merged_iterator = merge_async_iterators(*iterators,
+                                            is_cancelled=partial(asyncio.sleep,
+                                                                 0,
+                                                                 result=False))
 
     async def stream_output(generator: AsyncIterator[Tuple[int, str]]):
         async for idx, output in generator:
@@ -56,7 +41,8 @@ async def test_merge_async_iterators():
 
     for iterator in iterators:
         try:
-            await asyncio.wait_for(anext(iterator), 1)
+            # Can use anext() in python >= 3.10
+            await asyncio.wait_for(iterator.__anext__(), 1)
         except StopAsyncIteration:
             # All iterators should be cancelled and print this message.
             print("Iterator was cancelled normally")

tests/utils.py

@@ -7,12 +7,13 @@ import time
 import warnings
 from contextlib import contextmanager
 from pathlib import Path
-from typing import Any, Dict, List, Optional
+from typing import Any, Callable, Dict, List, Optional
 
 import openai
 import ray
 import requests
 from transformers import AutoTokenizer
+from typing_extensions import ParamSpec
 
 from vllm.distributed import (ensure_model_parallel_initialized,
                               init_distributed_environment)
@@ -360,13 +361,17 @@ def wait_for_gpu_memory_to_clear(devices: List[int],
         time.sleep(5)
 
 
-def fork_new_process_for_each_test(f):
+_P = ParamSpec("_P")
+
+
+def fork_new_process_for_each_test(
+        f: Callable[_P, None]) -> Callable[_P, None]:
     """Decorator to fork a new process for each test function.
     See https://github.com/vllm-project/vllm/issues/7053 for more details.
     """
 
     @functools.wraps(f)
-    def wrapper(*args, **kwargs):
+    def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> None:
         # Make the process the leader of its own process group
         # to avoid sending SIGTERM to the parent process
         os.setpgrp()