[ci][distributed] try to fix pp test (#7054)

2024-08-01 22:03:12 -07:00
parent 3bb4b1e4cd
commit 252357793d
4 changed files with 45 additions and 3 deletions
--- a/vllm/utils.py
+++ b/vllm/utils.py
@@ -928,7 +928,8 @@ def error_on_invalid_device_count_status():
    with contextlib.suppress(Exception):
        # future pytorch will fix the issue, device_count will not be cached
        # at that time, `.cache_info().currsize` will error out
-        cache_entries = torch.cuda.device_count.cache_info().currsize
+        cache_entries = torch.cuda.device_count.cache_info(  # type: ignore
+        ).currsize
    if cache_entries != 0:
        # the function is already called, and the result is cached
        remembered = torch.cuda.device_count()