Add --migration-limit to vllm trtllm sglang llama_cpp backends

Author: kthui

components/backends/llama_cpp/src/dynamo/llama_cpp/main.py

@@ -29,6 +29,7 @@ class Config:
     model_path: str
     model_name: Optional[str]
     context_length: int
+    migration_limit: int
 
 
 @dynamo_worker(static=False)
@@ -40,7 +41,13 @@ async def worker(runtime: DistributedRuntime):
 
     model_type = ModelType.Chat  # llama.cpp does the pre-processing
     endpoint = component.endpoint(config.endpoint)
-    await register_llm(model_type, endpoint, config.model_path, config.model_name)
+    await register_llm(
+        model_type,
+        endpoint,
+        config.model_path,
+        config.model_name,
+        migration_limit=config.migration_limit,
+    )
 
     # Initialize the engine
     # For more parameters see:
@@ -100,6 +107,12 @@ def cmd_line_args():
         default=None,
         help="Max model context length. Defaults to models max, usually model_max_length from tokenizer_config.json. Reducing this reduces VRAM requirements.",
     )
+    parser.add_argument(
+        "--migration-limit",
+        type=int,
+        default=0,
+        help="Maximum number of times a request may be migrated to a different engine worker. The number may be overridden by the engine.",
+    )
     args = parser.parse_args()
 
     config = Config()
@@ -124,6 +137,7 @@ def cmd_line_args():
     config.component = parsed_component_name
     config.endpoint = parsed_endpoint_name
     config.context_length = args.context_length
+    config.migration_limit = args.migration_limit
     return config
 
 

components/backends/sglang/src/dynamo/sglang/worker/main.py

@@ -311,11 +311,23 @@ def signal_handler():
 
     logging.info("Signal handlers set up for graceful shutdown")
 
-    server_args = parse_sglang_args_inc(sys.argv[1:])
-    await init(runtime, server_args)
-
-
-async def init(runtime: DistributedRuntime, server_args: ServerArgs):
+    # TODO: Better handle non-sglang args
+    sys_argv = sys.argv[1:]
+    migration_limit = 0
+    try:
+        idx = sys_argv.index("--migration-limit")
+        migration_limit = int(sys_argv[idx + 1])
+        del sys_argv[idx : idx + 2]  # Remove the args from sys_argv
+    except Exception:
+        pass
+
+    server_args = parse_sglang_args_inc(sys_argv)
+    await init(runtime, server_args, migration_limit)
+
+
+async def init(
+    runtime: DistributedRuntime, server_args: ServerArgs, migration_limit: int
+):
     """Initialize worker (either prefill or aggregated)"""
 
     engine = sgl.Engine(server_args=server_args)
@@ -330,6 +342,7 @@ async def init(runtime: DistributedRuntime, server_args: ServerArgs):
         server_args.model_path,
         server_args.served_model_name,
         kv_cache_block_size=server_args.page_size,
+        migration_limit=migration_limit,
     )
 
     if server_args.disaggregation_mode != "null":

components/backends/trtllm/src/dynamo/trtllm/main.py

@@ -140,6 +140,7 @@ async def init(runtime: DistributedRuntime, config: Config):
                 config.model_path,
                 config.served_model_name,
                 kv_cache_block_size=config.kv_block_size,
+                migration_limit=config.migration_limit,
             )
 
         # publisher will be set later if publishing is enabled.

components/backends/trtllm/src/dynamo/trtllm/utils/trtllm_utils.py

@@ -28,6 +28,7 @@ def __init__(self) -> None:
         self.served_model_name: Optional[str] = None
         self.tensor_parallel_size: int = 1
         self.kv_block_size: int = 32
+        self.migration_limit: int = 0
         self.extra_engine_args: str = ""
         self.publish_events_and_metrics: bool = False
         self.disaggregation_mode: DisaggregationMode = DEFAULT_DISAGGREGATION_MODE
@@ -46,6 +47,7 @@ def __str__(self) -> str:
             f"tensor_parallel_size={self.tensor_parallel_size}, "
             f"kv_block_size={self.kv_block_size}, "
             f"extra_engine_args={self.extra_engine_args}, "
+            f"migration_limit={self.migration_limit}, "
             f"publish_events_and_metrics={self.publish_events_and_metrics}, "
             f"disaggregation_mode={self.disaggregation_mode}, "
             f"disaggregation_strategy={self.disaggregation_strategy}, "
@@ -113,6 +115,12 @@ def cmd_line_args():
     parser.add_argument(
         "--kv-block-size", type=int, default=32, help="Size of a KV cache block."
     )
+    parser.add_argument(
+        "--migration-limit",
+        type=int,
+        default=0,
+        help="Maximum number of times a request may be migrated to a different engine worker. The number may be overridden by the engine.",
+    )
 
     parser.add_argument(
         "--extra-engine-args",
@@ -188,6 +196,7 @@ def cmd_line_args():
 
     config.tensor_parallel_size = args.tensor_parallel_size
     config.kv_block_size = args.kv_block_size
+    config.migration_limit = args.migration_limit
     config.extra_engine_args = args.extra_engine_args
     config.publish_events_and_metrics = args.publish_events_and_metrics
 

components/backends/vllm/src/dynamo/vllm/args.py

@@ -31,6 +31,7 @@ class Config:
     component: str
     endpoint: str
     is_prefill_worker: bool
+    migration_limit: int = 0
     kv_port: Optional[int] = None
     side_channel_port: Optional[int] = None
 
@@ -57,6 +58,12 @@ def parse_args() -> Config:
         action="store_true",
         help="Enable prefill functionality for this worker. Currently overwrites the --endpoint to be a specially chosen dyn://dynamo.prefill.generate",
     )
+    parser.add_argument(
+        "--migration-limit",
+        type=int,
+        default=0,
+        help="Maximum number of times a request may be migrated to a different engine worker. The number may be overridden by the engine.",
+    )
 
     parser = AsyncEngineArgs.add_cli_args(parser)
     args = parser.parse_args()
@@ -97,6 +104,7 @@ def parse_args() -> Config:
     config.endpoint = parsed_endpoint_name
     config.engine_args = engine_args
     config.is_prefill_worker = args.is_prefill_worker
+    config.migration_limit = args.migration_limit
 
     if config.engine_args.block_size is None:
         config.engine_args.block_size = 16

components/backends/vllm/src/dynamo/vllm/main.py

@@ -148,6 +148,7 @@ async def init(runtime: DistributedRuntime, config: Config):
             config.model,
             config.served_model_name,
             kv_cache_block_size=config.engine_args.block_size,
+            migration_limit=config.migration_limit,
         )
 
     factory = StatLoggerFactory(component, config.engine_args.data_parallel_rank or 0)