[fix] Fix test_llm_api_connector

- Reduce call_counts when using the overlap scheduler.

Signed-off-by: Robin Kobus <[email protected]>

Author: Funatiq

tests/integration/defs/llmapi/test_llm_api_connector.py

@@ -88,8 +88,7 @@ def test_connector_simple(enforce_single_worker, model_with_connector,
     assert len(scheduler.update_state_after_alloc.call_args.args[1]) == 1
 
     # With the overlap scheduler, we generate one extra token.
-    assert scheduler.build_connector_meta.call_count == NUM_TOKENS + int(
-        use_overlap_scheduler)
+    assert scheduler.build_connector_meta.call_count == NUM_TOKENS
 
     # We should have a single `SchedulerOutput` per forward pass.
     for i, call in enumerate(scheduler.build_connector_meta.call_args_list):
@@ -109,8 +108,7 @@ def test_connector_simple(enforce_single_worker, model_with_connector,
             assert len(scheduler_output.cached_requests[0].new_tokens) == 1
 
     # We call `start_load_kv` once at the beginning of each forward pass.
-    assert worker.start_load_kv.call_count == NUM_TOKENS + int(
-        use_overlap_scheduler)
+    assert worker.start_load_kv.call_count == NUM_TOKENS
 
     # Only called once when the request is received.
     assert scheduler.get_num_new_matched_tokens.call_count == 1
@@ -119,19 +117,16 @@ def test_connector_simple(enforce_single_worker, model_with_connector,
                      for call in worker.wait_for_layer_load.call_args_list) + 1
 
     # Called num_layers * num_forward_passes times.
-    assert worker.wait_for_layer_load.call_count == num_layers * (
-        NUM_TOKENS + int(use_overlap_scheduler))
-    assert worker.save_kv_layer.call_count == num_layers * (
-        NUM_TOKENS + int(use_overlap_scheduler))
+    assert worker.wait_for_layer_load.call_count == num_layers * (NUM_TOKENS)
+    assert worker.save_kv_layer.call_count == num_layers * (NUM_TOKENS)
 
     for i, call in enumerate(worker.wait_for_layer_load.call_args_list):
         assert call.args[0] == i % num_layers
 
     for i, call in enumerate(worker.save_kv_layer.call_args_list):
         assert call.args[0] == i % num_layers
 
-    assert worker.wait_for_save.call_count == NUM_TOKENS + int(
-        use_overlap_scheduler)
+    assert worker.wait_for_save.call_count == NUM_TOKENS
 
     assert scheduler.request_finished.call_count == 1
 
@@ -238,9 +233,7 @@ def test_connector_scheduler_output(enforce_single_worker, model_with_connector,
         scheduler.update_state_after_alloc.call_args.args[1]) == math.ceil(
             NUM_INPUT_TOKENS / BLOCK_SIZE)
 
-    # Additional token when using the overlap scheduler.
-    assert scheduler.build_connector_meta.call_count == NUM_TOKENS + int(
-        use_overlap_scheduler)
+    assert scheduler.build_connector_meta.call_count == NUM_TOKENS
 
     for i, call in enumerate(scheduler.build_connector_meta.call_args_list):
         sched_output = call.args[0]