fixed mypy

octoml · Dec 7, 2023 · 389d999 · 389d999
1 parent 1c983df
commit 389d999
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 18 deletions.
diff --git a/serve/mlc_serve/engine/engine_common.py b/serve/mlc_serve/engine/engine_common.py
@@ -115,8 +115,8 @@ def update_sequence(
 
 def get_requests_to_process(
     current_states: list[RequestState], cache_manager: KVCacheManager
-) -> Tuple[Union[list[PrefillRequest], list[DecodeRequest]], bool]:
-    requests = []
+) -> Tuple[list[Union[PrefillRequest, DecodeRequest]], bool]:
+    requests : list[Union[PrefillRequest, DecodeRequest]] = []
     # TODO: consider having hybrid batch if the underlying attention kernel supports
     # mixing prefill and decode.
     is_prompt_batch = any(

diff --git a/serve/mlc_serve/engine/sync_engine.py b/serve/mlc_serve/engine/sync_engine.py
@@ -4,7 +4,7 @@
 
 import logging
 from collections import defaultdict
-from typing import Set
+from typing import Set, DefaultDict
 
 from .base import (
     FinishReason,
@@ -108,19 +108,17 @@ def step(self) -> InferenceStepResult:
                 finish_reason = FinishReason.Length
 
             if finish_reason is not None:
-                seq_outputs = [
-                    SequenceOutput(
-                        i,
-                        finish_reason=finish_reason,
-                        num_generated_tokens=len(gen_seq.generated_token_ids),
-                    )
-                    for i, gen_seq in enumerate(state.generation_sequences)
-                ]
-
                 outputs.append(
                     RequestOutput(
                         state.request_id,
-                        seq_outputs,
+                        [
+                            SequenceOutput(
+                                i,
+                                finish_reason=finish_reason,
+                                num_generated_tokens=len(gen_seq.generated_token_ids),
+                            )
+                            for i, gen_seq in enumerate(state.generation_sequences)
+                        ],
                         num_prompt_tokens=state.prompt_len,
                     )
                 )
@@ -158,7 +156,7 @@ def step(self) -> InferenceStepResult:
             return InferenceStepResult(outputs)
 
         requests, _ = get_requests_to_process(
-            self.current_batch.values(), self.cache_manager
+            list(self.current_batch.values()), self.cache_manager
         )
         results = self.text_generator.generate(requests, self.cache_manager.get_cache())
         logger.debug("Finished text generation.")
@@ -182,7 +180,7 @@ def step(self) -> InferenceStepResult:
             else:
                 valid_results.append(res)
 
-        seq_outputs = defaultdict(list)
+        seq_outputs: DefaultDict[RequestId, list[SequenceOutput]] = defaultdict(list)
 
         for res in valid_results:
             request_id = res.sequence_id.request_id
@@ -216,12 +214,12 @@ def step(self) -> InferenceStepResult:
                 )
             )
 
-        for request_id, seq_outputs in seq_outputs.items():
+        for request_id, out_seqs in seq_outputs.items():
             state = self.current_batch[request_id]
             outputs.append(
                 RequestOutput(
                     request_id,
-                    sequences=seq_outputs,
+                    sequences=out_seqs,
                     num_prompt_tokens=state.prompt_len,
                 )
             )

diff --git a/serve/mlc_serve/model/paged_cache_model.py b/serve/mlc_serve/model/paged_cache_model.py
@@ -661,7 +661,7 @@ def generate(
                     outputs.append(
                         TextGenerationResult(
                             sequence_id=sequence_id,
-                            generated_tokens=[maybe_new_token[0]],
+                            generated_tokens=[maybe_new_token[0]],  # type: ignore
                             error=None,
                         )
                     )