fix

src/handler.py

@@ -34,8 +34,8 @@ engine_args = AsyncEngineArgs(
     tensor_parallel_size=NUM_GPU_SHARD,
     dtype="auto",
     seed=0,
-    max_num_batched_tokens=8192,
-    max_num_seqs=4096,
+    #max_num_batched_tokens=8192,
+    #max_num_seqs=4096,
     disable_log_stats=False
 )
 
@@ -43,12 +43,12 @@ engine_args = AsyncEngineArgs(
 llm = AsyncLLMEngine.from_engine_args(engine_args)
 
 # Incorporate metrics tracking
-llm.engine._log_system_stats = vllm_log_system_stats
+# llm.engine._log_system_stats = vllm_log_system_stats
 
 def concurrency_controller() -> bool:
     # Compute pending sequences
     total_pending_sequences = len(llm.engine.scheduler.waiting) + len(llm.engine.scheduler.swapped)
-    print("vLLM has {total_pending_sequences} pending sequences in its internal queue.")
+    print("vLLM has {} pending sequences in its internal queue.".format(total_pending_sequences))
 
     # If we have over 30 pending sequences, then we'll start auto-scaling.
     return total_pending_sequences > 30
@@ -153,8 +153,10 @@ async def handler_streaming(job: dict) -> Generator[dict[str, list], None, None]
     results_generator = llm.generate(prompt, sampling_params, request_id)
 
     # Streaming case
+    print("Phase B")
     positions = None
     async for request_output in results_generator:
+        print("Phase C")
         prompt = request_output.prompt
         text_outputs = []
 
@@ -164,6 +166,7 @@ async def handler_streaming(job: dict) -> Generator[dict[str, list], None, None]
                 'token_pos': 0
             }] * len(request_output.outputs)
 
+        print("Phase D")
         for idx, output in enumerate(request_output.outputs):
             # Extract the chunk position
             text_pos = positions[idx]['text_pos']
@@ -173,6 +176,7 @@ async def handler_streaming(job: dict) -> Generator[dict[str, list], None, None]
                 text_chunk = " ".join(output.text.split(" ")[text_pos:])
                 text_outputs.append(text_chunk)
 
+        print("Phase E")
         # Metrics for the vLLM serverless worker
         runpod_metrics = prepare_metrics()
         metrics = {}
@@ -183,6 +187,7 @@ async def handler_streaming(job: dict) -> Generator[dict[str, list], None, None]
         # The input tokens is the prompt. For each 'num_seqs' we'll have that many of them.
         metrics['input_tokens'] = len(request_output.prompt_token_ids)
 
+        print("Phase F")
         metrics['output_tokens'] = []
         for output in request_output.outputs:
             token_pos = positions[idx]['token_pos']
@@ -190,6 +195,7 @@ async def handler_streaming(job: dict) -> Generator[dict[str, list], None, None]
 
             metrics['output_tokens'].append(num_output_tokens)
 
+        print("Phase G")
         # Update positions
         for idx, output in enumerate(request_output.outputs):
             positions[idx] = {
@@ -197,6 +203,7 @@ async def handler_streaming(job: dict) -> Generator[dict[str, list], None, None]
                 'token_pos': len(output.token_ids)
             }
 
+        print("Phase H")
         ret = {
             "text": text_outputs,
             "metrics": metrics,