Lints, fixes prints

Signed-off-by: Rafael Vasquez <[email protected]>

Author: rafvasq

vllm_spyre/v1/worker/spyre_model_runner.py

@@ -437,7 +437,6 @@ def pad_input_ids(
 
         return input_ids, position_ids, mask
 
-
     def get_kv_cache_spec(self) -> KVCacheSpec:
         """
         This method should generate the KVCache spec by parsing the kv cache

vllm_spyre/v1/worker/spyre_worker.py

@@ -15,7 +15,8 @@
 from vllm.model_executor import set_random_seed
 from vllm.platforms import current_platform
 from vllm.sampling_params import SamplingParams
-from vllm.v1.core.scheduler import CachedRequestData, NewRequestData, SchedulerOutput
+from vllm.v1.core.scheduler import (CachedRequestData, NewRequestData,
+                                    SchedulerOutput)
 from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec
 from vllm.v1.outputs import ModelRunnerOutput
 from vllm.v1.worker.worker_base import WorkerBase as WorkerBaseV1
@@ -81,9 +82,8 @@ def compile_or_warm_up_model(self) -> None:
             "combinations finished. Total warmup time %.3fs.",
             len(wup_new_tokens), all_warmup_total_t)
 
-
     def _warmup_spyre_fixed_size(self, prompt_len, num_decode_tokens,
-                  special_token_ids, batch_size):
+                                 special_token_ids, batch_size):
 
         warmup_start_t = time.time()
         # NOTE(ngl): empty tensor causes spyre to hang, so using
@@ -97,7 +97,9 @@ def _warmup_spyre_fixed_size(self, prompt_len, num_decode_tokens,
             i for i in range(1, vocab_size) if i not in set(special_token_ids)
         ]
         # Convert to tensor for sampling
-        valid_token_ids_tensor = torch.tensor(valid_token_ids, dtype=torch.long, device="cpu")
+        valid_token_ids_tensor = torch.tensor(valid_token_ids,
+                                              dtype=torch.long,
+                                              device="cpu")
 
         # Sample from the valid token ids
         warmup_tokens_tensor = valid_token_ids_tensor[torch.randint(
@@ -106,10 +108,12 @@ def _warmup_spyre_fixed_size(self, prompt_len, num_decode_tokens,
         # Create requests to be used for prefill steps
         dummy_requests = [
             NewRequestData(
-                req_id=f"warmup",
+                req_id="warmup",
                 prompt_token_ids=warmup_tokens_tensor[i].tolist(),
                 prompt="test",
-                mm_inputs=[], mm_hashes=[], mm_positions=[],
+                mm_inputs=[],
+                mm_hashes=[],
+                mm_positions=[],
                 sampling_params=SamplingParams(max_tokens=num_decode_tokens),
                 block_ids=[0],
                 num_computed_tokens=0,
@@ -122,8 +126,10 @@ def _warmup_spyre_fixed_size(self, prompt_len, num_decode_tokens,
             CachedRequestData(
                 req_id=req.req_id,
                 resumed_from_preemption=False,
-                new_token_ids=[valid_token_ids_tensor[torch.randint(
-                    0, len(valid_token_ids_tensor), (1,)).item()]],  # placeholder token
+                new_token_ids=[
+                    valid_token_ids_tensor[torch.randint(
+                        0, len(valid_token_ids_tensor), (1, )).item()]
+                ],  # placeholder token
                 new_block_ids=req.block_ids,
                 num_computed_tokens=req.num_computed_tokens,
             ) for req in dummy_requests
@@ -134,8 +140,10 @@ def _warmup_spyre_fixed_size(self, prompt_len, num_decode_tokens,
         scheduler_output = SchedulerOutput(
             scheduled_new_reqs=dummy_requests,
             scheduled_cached_reqs=[],
-            num_scheduled_tokens={i: prompt_len for i in range(batch_size)},
-            total_num_scheduled_tokens=sum(prompt_len for _ in range(batch_size)),
+            num_scheduled_tokens={i: prompt_len
+                                  for i in range(batch_size)},
+            total_num_scheduled_tokens=sum(prompt_len
+                                           for _ in range(batch_size)),
             scheduled_spec_decode_tokens={},
             scheduled_encoder_inputs={},
             num_common_prefix_blocks=0,
@@ -144,14 +152,14 @@ def _warmup_spyre_fixed_size(self, prompt_len, num_decode_tokens,
         )
 
         # First full forward pass
-        logger.info("[SpyreWorker] Warmup 1/2: Prefill...")
+        logger.info("Warmup 1/2: Prefill...")
         self.execute_model(scheduler_output)  # Prefill step
 
         # Switch to cached requests to trigger decoding steps
         scheduler_output.scheduled_new_reqs = []
         scheduler_output.scheduled_cached_reqs = cached_requests
 
-        logger.info("[SpyreWorker] Warmup 1/2: Decoding...")
+        logger.info("Warmup 1/2: Decoding...")
         for _ in range(num_decode_tokens - 1):
             self.execute_model(scheduler_output)
 
@@ -161,10 +169,11 @@ def _warmup_spyre_fixed_size(self, prompt_len, num_decode_tokens,
             ul_start_time = time.time()
             torch_sendnn.update_lazyhandle()
             ul_stop_time = time.time()
-            logger.info(f"update_lazyhandle() done (duration: {ul_stop_time - ul_start_time}s)")
+            logger.info("update_lazyhandle() done (duration: %.3fs",
+                        ul_stop_time - ul_start_time)
 
         # Second full forward pass
-        logger.info("[SpyreWorker] Warmup 2/2: Prefill step...")
+        logger.info("Warmup 2/2: Prefill step...")
         scheduler_output.scheduled_new_reqs = dummy_requests
         scheduler_output.scheduled_cached_reqs = []
         self.execute_model(scheduler_output)
@@ -173,16 +182,16 @@ def _warmup_spyre_fixed_size(self, prompt_len, num_decode_tokens,
         scheduler_output.scheduled_new_reqs = []
         scheduler_output.scheduled_cached_reqs = cached_requests
 
-        logger.info("[SpyreWorker] Warmup 2/2: Decoding steps...")
+        logger.info("[Warmup 2/2: Decoding steps...")
         for _ in range(num_decode_tokens - 1):
             self.execute_model(scheduler_output)
 
         warmup_end_t = time.time()
         warmup_total_t = warmup_end_t - warmup_start_t
-        logger.info("[SpyreWorker] ... warmup finished.")
-        logger.info(f"\twarmup took {warmup_total_t}s (for prompt length"
-              f"{prompt_len} and max output tokens {num_decode_tokens})")
-
+        logger.info("Warmup finished.")
+        logger.info(
+            "Warmup took %.3fs (for prompt length %d and max output tokens %d)",
+            warmup_total_t, prompt_len, num_decode_tokens)
 
     def check_health(self) -> None:
         """Basic health check (override for device-specific checks)."""