fix

AlibabaPAI · Nov 8, 2024 · e0ea522 · e0ea522
1 parent 207237d
commit e0ea522
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/tests/unit_test/llumlet/test_local_migration_scheduler.py b/tests/unit_test/llumlet/test_local_migration_scheduler.py
@@ -70,9 +70,11 @@ def test_scheduler_policy():
     engine.add_request(request_id="5", length=4, expected_steps=math.inf)
     scheduler.request_migration_policy = "LCFS"
     request = scheduler.get_migrate_out_request()
+    request.migrating = True
     assert request.request_id == "3"
     assert request.output_len >= request.expected_steps and request.inference_type == RequestInferenceType.DECODE
     request = scheduler.get_migrate_out_request()
+    request.migrating = True
     assert request.request_id == "5"
     request = scheduler.get_migrate_out_request()
     assert request.request_id == "4"