finished test

Signed-off-by: Sophie du Couédic <[email protected]>

Author: sducouedic

tests/e2e/test_spyre_cb_scheduler_steps.py

@@ -475,24 +475,23 @@ def test_two_sequences_finish_same_time_as_new_arrive(
     )
 
 
-@pytest.mark.a
 @pytest.mark.cb
 @pytest.mark.parametrize("model", get_spyre_model_list())
 @pytest.mark.parametrize("backend", get_spyre_backend_list())
-def test_new_sequence_joins_during_decode(
-        model: str, backend: str, monkeypatch: pytest.MonkeyPatch):
+def test_new_sequence_joins_during_decode(model: str, backend: str,
+                                          monkeypatch: pytest.MonkeyPatch):
     """ Scenario where a new sequence joins while decoding other sequences
 
     Configuration:
         * max_num_seqs: 4
         * number of prompts: 4
             * 1: len = 49, max tokens = 119, step joining = 0
             * 2: len = 14, max tokens = 52, step joining = 0
-            * 3: len = 89, max tokens = 101, step joining = 32
+            * 3: len = 89, max tokens = 104, step joining = 32
             * 4: len = 9, max tokens = 65, step joining = 131
     """
 
-    seqs_max_tokens = [119, 52, 101, 65]
+    seqs_max_tokens = [119, 52, 104, 65]
     prompts_lengths = [49, 14, 89, 9]
     steps_add_reqs = [0, 0, 32, 131]
     available_blocks = -1  # no restriction
@@ -510,7 +509,6 @@ def test_new_sequence_joins_during_decode(
         },
         {
             # Prefill sequence 0
-            # total blocks in use: 1
             "step": 1,
             "tkv": 64,
             "waiting": ["1"],
@@ -521,7 +519,6 @@ def test_new_sequence_joins_during_decode(
         },
         {
             # Prefill sequence 1
-            # total blocks in use: 1 + 1 = 2
             "step": 2,
             "tkv": 64,
             "waiting": [],
@@ -532,7 +529,6 @@ def test_new_sequence_joins_during_decode(
         },
         {
             # Decode sequences 0 and 1
-            # total blocks in use: 2 + 2 = 4
             "step": 3,
             "tkv": 65,
             "waiting": [],
@@ -553,18 +549,16 @@ def test_new_sequence_joins_during_decode(
         },
         {
             # Prefill sequence 2
-            # total blocks in use: 4 + 2 (long prefill) = 6
             "step": 33,
             "tkv": 94,
             "waiting": [],
             "running": ["2", "1", "0"],
             "request_outputs": ["2"],
-            "n_reserved_blocks": 9,  # prefill (2 block) + 100 decode (2 block)
+            "n_reserved_blocks": 9,  # prefill (2 block) + 103 decode (2 block)
             "n_used_blocks": 6
         },
         {
             # Decode sequences 0, 1, and 2
-            # total blocks in use: 6
             "step": 34,
             "tkv": 95,
             "waiting": [],
@@ -576,7 +570,6 @@ def test_new_sequence_joins_during_decode(
         {
             # Sequence 1 finishes at step 54
             # (start step + 2 prefills + 51 decodes - 1) = 2 + 2 + 51 - 1 = 54
-            # total blocks in use: 6
             "step": 54,
             "tkv": 115,
             "waiting": [],
@@ -588,7 +581,6 @@ def test_new_sequence_joins_during_decode(
         },
         {
             # Decode sequences 0 and 2
-            # total blocks in use: 4
             "step": 55,
             "tkv": 116,
             "waiting": [],
@@ -599,7 +591,6 @@ def test_new_sequence_joins_during_decode(
         },
         {
             # Decode sequences 0 and 2, tkv arrives to new block
-            # total blocks in use: 6
             "step": 68,
             "tkv": 129,
             "waiting": [],
@@ -611,7 +602,6 @@ def test_new_sequence_joins_during_decode(
         {
             # Sequence 0 finishes at step 121
             # (start step + 3 prefills + 118 decode - 1) = 1 + 3 + 118 - 1 = 121
-            # total blocks in use: 6
             "step": 121,
             "tkv": 182,
             "waiting": [],
@@ -623,7 +613,6 @@ def test_new_sequence_joins_during_decode(
         },
         {
             # Decode sequence 2
-            # total blocks in use: 3
             "step": 122,
             "tkv": 183,
             "waiting": [],
@@ -634,7 +623,6 @@ def test_new_sequence_joins_during_decode(
         },
         {
             # Sequence 3 joins: one iteration in waiting queue
-            # total blocks in use: 3
             "step": 131,
             "tkv": 192,
             "waiting": ["3"],
@@ -643,6 +631,71 @@ def test_new_sequence_joins_during_decode(
             "n_reserved_blocks": 4,
             "n_used_blocks": 3
         },
+        {
+            # Prefill sequence 3
+            "step": 132,
+            "tkv": 192,
+            "waiting": [],
+            "running": ["3", "2"],
+            "request_outputs": ["3"],
+            "n_reserved_blocks": 8,  # prefill (3 blocks) + 64 decode (1 block)
+            "n_used_blocks": 6  # prefill (3 block)
+        },
+        {
+            # Decode sequences 2 and 3
+            "step": 133,
+            "tkv": 193,
+            "waiting": [],
+            "running": ["3", "2"],
+            "request_outputs": ["3", "2"],
+            "n_reserved_blocks": 8,  # prefill (3 blocks) + 64 decode (1 block)
+            "n_used_blocks": 8  # 2 blocks extended, one for each sequence
+        },
+        {
+            # Sequence 2 finishes at step 137
+            # (start step + 2 prefills + 103 decodes) = 33 + 2 + 103 - 1 = 137
+            "step": 137,
+            "tkv": 197,
+            "waiting": [],
+            "running": ["3"],
+            "request_outputs": ["3", "2"],
+            "finished_requests": ["2"],
+            "n_reserved_blocks": 8,
+            "n_used_blocks": 8
+        },
+        {
+            # Decode sequence 3
+            "step": 138,
+            "tkv": 70,
+            "waiting": [],
+            "running": ["3"],
+            "request_outputs": ["3"],
+            # 6 blocks freed: finished sequence (4) + left padding stripping (2)
+            "n_reserved_blocks": 2,
+            "n_used_blocks": 2
+        },
+        {
+            # Sequence 3 finishes at step 196
+            # (start step + 1 prefills + 103 decodes) = 132 + 1 + 64 - 1 = 196
+            "step": 196,
+            "tkv": 128,
+            "waiting": [],
+            "running": [],
+            "request_outputs": ["3"],
+            "finished_requests": ["3"],
+            "n_reserved_blocks": 2,
+            "n_used_blocks": 2
+        },
+        {
+            # Tkv should be cleared one step later
+            "step": 197,
+            "tkv": 0,
+            "waiting": [],
+            "running": [],
+            "request_outputs": [],
+            "n_reserved_blocks": 0,
+            "n_used_blocks": 0
+        },
     ]
 
     check_scheduler_inference_steps(
@@ -658,6 +711,7 @@ def test_new_sequence_joins_during_decode(
         use_cb=True,
     )
 
+
 @pytest.mark.cb
 @pytest.mark.parametrize("model", get_spyre_model_list())
 @pytest.mark.parametrize("backend", get_spyre_backend_list())