update

llumnix/llumlet/migration_coordinator.py

@@ -27,16 +27,16 @@ class MigrationStatus(enum.Enum):
     """Status of Migration."""
     RUNNING = enum.auto()
     # aborted by src instance
-    FINISHED_ABORTED_SRC = enum.auto()
+    ABORTED_SRC = enum.auto()
     # aborted by dst instance
-    FINISHED_ABORTED_DST = enum.auto()
+    ABORTED_DST = enum.auto()
     FINISHED_DONE = enum.auto()
 
     @staticmethod
     def is_finished(status: "MigrationStatus") -> bool:
         return status in [
-            MigrationStatus.FINISHED_ABORTED_SRC,
-            MigrationStatus.FINISHED_ABORTED_DST,
+            MigrationStatus.ABORTED_SRC,
+            MigrationStatus.ABORTED_DST,
             MigrationStatus.FINISHED_DONE
         ]
 
@@ -77,7 +77,7 @@ def migrate_out_onestage(self, migrate_in_ray_actor: "ray.actor.ActorHandle",  m
             if is_last_stage:
                 self.backend_engine.add_running_request(migrate_out_request)
                 self.backend_engine.remove_migrating_out_request_last_stage(migrate_out_request)
-            migration_status = MigrationStatus.FINISHED_ABORTED_DST
+            migration_status = MigrationStatus.ABORTED_DST
             return migration_status
         # do stage send/recv
         migrate_out_request.stage_timestamps.append(time.time())
@@ -86,7 +86,7 @@ def migrate_out_onestage(self, migrate_in_ray_actor: "ray.actor.ActorHandle",  m
         self.backend_engine.send_blocks(migrate_in_ray_actor, src_blocks, dst_blocks)
         if not is_last_stage and migrate_out_request.should_abort_migration():
             # migrate-out request abort by scheduler during send/recv
-            migration_status = MigrationStatus.FINISHED_ABORTED_SRC
+            migration_status = MigrationStatus.ABORTED_SRC
 
         return migration_status
 
@@ -102,7 +102,7 @@ def migrate_out_multistage(self, migrate_in_ray_actor: "ray.actor.ActorHandle",
             if MigrationStatus.is_finished(status):
                 return status
         # exceed max stages
-        return MigrationStatus.FINISHED_ABORTED_SRC
+        return MigrationStatus.ABORTED_SRC
 
     def migrate_in_pre_alloc(self, request_id: str, block_num: int) -> List[int]:
         """prev alloc blocks to migrate in request

tests/unit_test/backends/vllm/test_simulator.py

@@ -105,4 +105,4 @@ async def check_output_len():
 
     await check_output_len()
 
-    que.cleanup()
+    que.cleanup()

tests/unit_test/llumlet/test_migration_coordinator.py

@@ -65,7 +65,7 @@ def test_migrate_out_onestage(setup_ray_env):
     migrate_out_request.should_abort_migration.return_value = False
     migrate_in_ray_actor.execute_migration_method.remote.return_value = ray_remote_call.remote(dst_blocks)
     status = coordinator.migrate_out_onestage(migrate_in_ray_actor, migrate_out_request)
-    assert status == MigrationStatus.FINISHED_ABORTED_DST
+    assert status == MigrationStatus.ABORTED_DST
 
     migrate_out_request = MagicMock()
     src_blocks = [1, 2, 3]
@@ -74,7 +74,7 @@ def test_migrate_out_onestage(setup_ray_env):
     migrate_out_request.should_abort_migration.return_value = True
     migrate_in_ray_actor.execute_migration_method.remote.return_value = ray_remote_call.remote(dst_blocks)
     status = coordinator.migrate_out_onestage(migrate_in_ray_actor, migrate_out_request)
-    assert status == MigrationStatus.FINISHED_ABORTED_SRC
+    assert status == MigrationStatus.ABORTED_SRC
 
 # setup_ray_env should be passed after migrate_out_onestage
 @patch.object(MigrationCoordinator, 'migrate_out_onestage')
@@ -104,4 +104,4 @@ def test_migrate_out_multistage(_, setup_ray_env):
                                                     MigrationStatus.RUNNING]
     status = coordinator.migrate_out_multistage(migrate_in_ray_actor, migrate_out_request)
     assert coordinator.migrate_out_onestage.call_count == max_stages + 1
-    assert status == MigrationStatus.FINISHED_ABORTED_SRC
+    assert status == MigrationStatus.ABORTED_SRC