fix

Makefile

@@ -21,11 +21,11 @@ install:
 
 .PHONY: lint
 lint: check_pylint_installed check_pytest_installed
-	@pylint --rcfile=.pylintrc -s n  --jobs=32 ./llumnix
+	@pylint --rcfile=.pylintrc -s n  --jobs=128 ./llumnix
 
 	@pylint --rcfile=.pylintrc \
 			--disable=protected-access,super-init-not-called,unused-argument,redefined-outer-name,invalid-name \
-			-s n --jobs=32 ./tests
+			-s n --jobs=128 ./tests
 
 .PHONY: test
 test: check_pytest_installed

llumnix/arg_utils.py

@@ -199,10 +199,6 @@ def check_args(cls, args: 'EngineManagerArgs', parser: argparse.ArgumentParser):
             if hasattr(action, 'choices') and action.choices is not None and hasattr(args, action.dest):
                 assert getattr(args, action.dest) in action.choices, f"{action.dest} should be one of {action.choices}."
 
-        if args.migration_backend == 'nccl' and args.migration_internal_cache_size != 1:
-            logger.warning("The NCCL migration backend does not support concurrency. Set migration_internal_cache_size to 1.")
-            args.migration_internal_cache_size = 1
-
         assert args.migration_backend != 'gloo' or (args.migration_backend == 'gloo' \
             and not args.disable_init_instance_by_manager and not args.disable_fixed_node_init_instance), \
             ("When using gloo as migration backend, "

tests/unit_test/backends/vllm/test_migration_backend.py

@@ -80,11 +80,12 @@ def get_gpu_cache(self):
 @pytest.mark.parametrize("backend", ['rpc', 'gloo', 'nccl'])
 def test_one_to_many_migrate_cache(setup_ray_env, backend):
     engine_config = EngineArgs(model='facebook/opt-125m', max_model_len=8, enforce_eager=True).create_engine_config()
-    migraiton_config = EngineManagerArgs(migration_cache_blocks=3, migration_num_layers=5).create_migration_config()
+    migraiton_config = EngineManagerArgs(migration_cache_blocks=3, migration_num_layers=5,
+                                         migration_internal_cache_size=2).create_migration_config()
     migraiton_config.migration_backend = backend
 
-    num_worker = 3 if backend != 'nccl' else 2
-    num_gpu_blocks = 300
+    num_worker = 3
+    num_gpu_blocks = 6000
     workers, _ = get_ready_workers(num_worker, num_gpu_blocks, engine_config, migraiton_config)
 
     num_layers = engine_config.model_config.get_num_layers(engine_config.parallel_config)
@@ -125,11 +126,12 @@ def test_one_to_many_migrate_cache(setup_ray_env, backend):
 @pytest.mark.parametrize("backend", ['rpc', 'gloo', 'nccl'])
 def test_many_to_one_migrate_cache(setup_ray_env, backend):
     engine_config = EngineArgs(model='facebook/opt-125m', max_model_len=8, enforce_eager=True).create_engine_config()
-    migraiton_config = EngineManagerArgs(migration_cache_blocks=3, migration_num_layers=5).create_migration_config()
+    migraiton_config = EngineManagerArgs(migration_cache_blocks=3, migration_num_layers=5,
+                                        migration_internal_cache_size=2).create_migration_config()
     migraiton_config.migration_backend = backend
 
-    num_worker = 3 if backend != 'nccl' else 2
-    num_gpu_blocks = 300
+    num_worker = 3
+    num_gpu_blocks = 6000
     workers, _ = get_ready_workers(num_worker, num_gpu_blocks, engine_config, migraiton_config)
 
     num_layers = engine_config.model_config.get_num_layers(engine_config.parallel_config)