From 0cf8044c21521caf17dd46467293881db8738a9a Mon Sep 17 00:00:00 2001
From: Russell Bryant <rbryant@redhat.com>
Date: Tue, 8 Oct 2024 18:08:22 -0400
Subject: [PATCH] mypy: check additional directories (#9162)

Signed-off-by: Russell Bryant <rbryant@redhat.com>
---
 .github/workflows/mypy.yaml | 15 ++-------------
 format.sh                   | 12 +-----------
 tools/mypy.sh               | 36 ++++++++++++++++++++++++++++++++++++
 vllm/usage/usage_lib.py     |  4 ++--
 4 files changed, 41 insertions(+), 26 deletions(-)
 create mode 100755 tools/mypy.sh

diff --git a/.github/workflows/mypy.yaml b/.github/workflows/mypy.yaml
index 24f58f88361c8..d578d7c521402 100644
--- a/.github/workflows/mypy.yaml
+++ b/.github/workflows/mypy.yaml
@@ -11,7 +11,7 @@ on:
       - main
 
 jobs:
-  ruff:
+  mypy:
     runs-on: ubuntu-latest
     strategy:
       matrix:
@@ -32,15 +32,4 @@ jobs:
         pip install types-setuptools
     - name: Mypy
       run: |
-        mypy
-        mypy tests --follow-imports skip
-        mypy vllm/attention --follow-imports skip
-        mypy vllm/distributed --follow-imports skip
-        mypy vllm/engine  --follow-imports skip
-        mypy vllm/executor --follow-imports skip
-        mypy vllm/lora --follow-imports skip
-        mypy vllm/model_executor  --follow-imports skip
-        mypy vllm/prompt_adapter --follow-imports skip
-        mypy vllm/spec_decode --follow-imports skip
-        mypy vllm/worker --follow-imports skip
-
+        tools/mypy.sh
diff --git a/format.sh b/format.sh
index a0df92b350133..1ac028d00e3a4 100755
--- a/format.sh
+++ b/format.sh
@@ -96,17 +96,7 @@ echo 'vLLM yapf: Done'
 
 # Run mypy
 echo 'vLLM mypy:'
-mypy --follow-imports skip  # Note that this is less strict than CI
-mypy tests --follow-imports skip
-mypy vllm/attention --follow-imports skip
-mypy vllm/distributed --follow-imports skip
-mypy vllm/engine  --follow-imports skip
-mypy vllm/executor --follow-imports skip
-mypy vllm/lora --follow-imports skip
-mypy vllm/model_executor  --follow-imports skip
-mypy vllm/prompt_adapter --follow-imports skip
-mypy vllm/spec_decode --follow-imports skip
-mypy vllm/worker --follow-imports skip
+tools/mypy.sh
 echo 'vLLM mypy: Done'
 
 
diff --git a/tools/mypy.sh b/tools/mypy.sh
new file mode 100755
index 0000000000000..4e358faafe8de
--- /dev/null
+++ b/tools/mypy.sh
@@ -0,0 +1,36 @@
+#!/bin/bash
+
+CI=${1:-0}
+
+run_mypy() {
+    echo "Running mypy on $1"
+    if [ $CI -eq 1 ] && [ -z "$1" ]; then
+        mypy "$@"
+        return
+    fi
+    mypy --follow-imports skip "$@"
+}
+
+run_mypy # Note that this is less strict than CI
+run_mypy tests
+run_mypy vllm/assets
+run_mypy vllm/attention
+#run_mypy vllm/compilation
+#run_mypy vllm/core
+run_mypy vllm/distributed
+run_mypy vllm/engine
+#run_mypy vllm/entrypoints
+run_mypy vllm/executor
+#run_mypy vllm/inputs
+run_mypy vllm/logging
+run_mypy vllm/lora
+run_mypy vllm/model_executor
+run_mypy vllm/multimodal
+run_mypy vllm/platforms
+run_mypy vllm/plugins
+run_mypy vllm/prompt_adapter
+run_mypy vllm/spec_decode
+run_mypy vllm/transformers_utils
+run_mypy vllm/usage
+#run_mypy vllm/vllm_flash_attn
+run_mypy vllm/worker
diff --git a/vllm/usage/usage_lib.py b/vllm/usage/usage_lib.py
index 7fadfd5dfffb4..9ae46ff43a916 100644
--- a/vllm/usage/usage_lib.py
+++ b/vllm/usage/usage_lib.py
@@ -204,7 +204,7 @@ def _report_continous_usage(self):
             self._write_to_file(data)
             self._send_to_server(data)
 
-    def _send_to_server(self, data):
+    def _send_to_server(self, data: Dict[str, Any]) -> None:
         try:
             global_http_client = global_http_connection.get_sync_client()
             global_http_client.post(_USAGE_STATS_SERVER, json=data)
@@ -212,7 +212,7 @@ def _send_to_server(self, data):
             # silently ignore unless we are using debug log
             logging.debug("Failed to send usage data to server")
 
-    def _write_to_file(self, data):
+    def _write_to_file(self, data: Dict[str, Any]) -> None:
         os.makedirs(os.path.dirname(_USAGE_STATS_JSON_PATH), exist_ok=True)
         Path(_USAGE_STATS_JSON_PATH).touch(exist_ok=True)
         with open(_USAGE_STATS_JSON_PATH, "a") as f: