feat: Add migration to LLM requests (#1930)

Author: kthui

launch/dynamo-run/src/flags.rs

@@ -162,6 +162,11 @@ pub struct Flags {
     #[arg(long)]
     pub request_template: Option<PathBuf>,
 
+    /// How many times a request can be migrated to another worker if the HTTP server lost
+    /// connection to the current worker.
+    #[arg(long, value_parser = clap::value_parser!(u32).range(0..1024))]
+    pub migration_limit: Option<u32>,
+
     /// Everything after a `--`.
     /// These are the command line arguments to the python engine when using `pystr` or `pytok`.
     #[arg(index = 2, last = true, hide = true, allow_hyphen_values = true)]
@@ -180,6 +185,9 @@ impl Flags {
                 if self.kv_cache_block_size.is_some() {
                     anyhow::bail!("'--kv-cache-block-size' flag should only be used on the worker node, not on the ingress");
                 }
+                if self.migration_limit.is_some() {
+                    anyhow::bail!("'--migration-limit' flag should only be used on the worker node, not on the ingress");
+                }
             }
             Output::EchoFull => {}
             Output::EchoCore => {

launch/dynamo-run/src/lib.rs

@@ -45,7 +45,8 @@ pub async fn run(
         .context_length(flags.context_length)
         .http_port(Some(flags.http_port))
         .router_config(Some(flags.router_config()))
-        .request_template(flags.request_template.clone());
+        .request_template(flags.request_template.clone())
+        .migration_limit(flags.migration_limit);
 
     // If `in=dyn` we want the trtllm/sglang/vllm subprocess to listen on that endpoint.
     // If not, then the endpoint isn't exposed so we let LocalModel invent one.

launch/dynamo-run/src/subprocess.rs

@@ -48,6 +48,8 @@ pub async fn start(
         card.kv_cache_block_size.to_string(),
         "--context-length".to_string(),
         card.context_length.to_string(),
+        "--migration-limit".to_string(),
+        card.migration_limit.to_string(),
     ];
     // TRTLLM only
     // The worker node will only publish events and metrics if the router mode is KV

launch/dynamo-run/src/subprocess/sglang_inc.py

@@ -42,6 +42,7 @@ class Config:
     nnodes: int
     node_rank: int
     dist_init_addr: str
+    migration_limit: int
     extra_engine_args: str
 
 
@@ -202,7 +203,13 @@ async def init(runtime: DistributedRuntime, config: Config):
     model_type = (
         ModelType.Backend if not engine_args.is_embedding else ModelType.Embedding
     )
-    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,
+    )
 
     # the server will gracefully shutdown (i.e., keep opened TCP streams finishes)
     # after the lease is revoked
@@ -268,6 +275,12 @@ def cmd_line_args():
         default="",
         help="Host address (e.g., `192.168.0.2:25000`) of the node with rank 0",
     )
+    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",
         type=str,
@@ -304,6 +317,7 @@ def cmd_line_args():
     config.nnodes = args.nnodes
     config.node_rank = args.node_rank
     config.dist_init_addr = args.dist_init_addr
+    config.migration_limit = args.migration_limit
     config.extra_engine_args = args.extra_engine_args
     return config
 

launch/dynamo-run/src/subprocess/trtllm_inc.py

@@ -122,6 +122,7 @@ class Config:
     model_name: Optional[str] = None
     tensor_parallel_size: int
     kv_block_size: int
+    migration_limit: int
     extra_engine_args: str
     publish_events_and_metrics: bool
     disaggregation_mode: str
@@ -136,6 +137,7 @@ def __str__(self) -> str:
             f"model_name={self.model_name}, "
             f"tensor_parallel_size={self.tensor_parallel_size}, "
             f"kv_block_size={self.kv_block_size}, "
+            f"migration_limit={self.migration_limit}, "
             f"extra_engine_args={self.extra_engine_args}, "
             f"publish_events_and_metrics={self.publish_events_and_metrics}, "
             f"disaggregation_mode={self.disaggregation_mode}, "
@@ -404,6 +406,7 @@ async def init(runtime: DistributedRuntime, config: Config):
                 config.model_path,
                 config.model_name,
                 kv_cache_block_size=config.kv_block_size,
+                migration_limit=config.migration_limit,
             )
 
         # publisher will be set later if publishing is enabled.
@@ -476,6 +479,12 @@ def cmd_line_args():
         default=None,
         help="This argument is not used by TRTLLM. Please provide max_input_len, max_seq_len and max_output_len in yaml file and point --extra-engine-args to the yaml file.",
     )
+    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",
         type=str,
@@ -557,6 +566,7 @@ def cmd_line_args():
     config.endpoint = parsed_endpoint_name
     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
     config.disaggregation_mode = disaggregation_mode

launch/dynamo-run/src/subprocess/vllm_inc.py

@@ -56,6 +56,7 @@ class Config:
     tensor_parallel_size: int
     kv_block_size: int
     context_length: int
+    migration_limit: int
     extra_engine_args: str
 
 
@@ -233,6 +234,7 @@ async def init(runtime: DistributedRuntime, config: Config):
             "max_model_len", None
         ),  # if None, takes length from tokenizer
         kv_cache_block_size=arg_map["block_size"],
+        migration_limit=config.migration_limit,
     )
     handler = RequestHandler(component, engine_client, default_sampling_params)
     handler.setup_kv_metrics()
@@ -276,6 +278,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.",
+    )
     parser.add_argument(
         "--extra-engine-args",
         type=str,
@@ -308,6 +316,7 @@ def cmd_line_args():
     config.tensor_parallel_size = args.tensor_parallel_size
     config.kv_block_size = args.kv_block_size
     config.context_length = args.context_length
+    config.migration_limit = args.migration_limit
     config.extra_engine_args = args.extra_engine_args
 
     return config

launch/dynamo-run/src/subprocess/vllm_v1_inc.py

@@ -65,6 +65,7 @@ class Config:
     tensor_parallel_size: int
     kv_block_size: int
     context_length: int
+    migration_limit: int
     extra_engine_args: str
 
 
@@ -218,6 +219,7 @@ async def init(runtime: DistributedRuntime, config: Config):
         config.model_path,
         config.model_name,
         kv_cache_block_size=config.kv_block_size,
+        migration_limit=config.migration_limit,
     )
 
     arg_map = {
@@ -333,6 +335,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.",
+    )
     parser.add_argument(
         "--extra-engine-args",
         type=str,
@@ -365,6 +373,7 @@ def cmd_line_args():
     config.tensor_parallel_size = args.tensor_parallel_size
     config.kv_block_size = args.kv_block_size
     config.context_length = args.context_length
+    config.migration_limit = args.migration_limit
     config.extra_engine_args = args.extra_engine_args
 
     return config

lib/bindings/python/rust/lib.rs

@@ -131,7 +131,7 @@ fn log_message(level: &str, message: &str, module: &str, file: &str, line: u32)
 }
 
 #[pyfunction]
-#[pyo3(signature = (model_type, endpoint, model_path, model_name=None, context_length=None, kv_cache_block_size=None, router_mode=None))]
+#[pyo3(signature = (model_type, endpoint, model_path, model_name=None, context_length=None, kv_cache_block_size=None, router_mode=None, migration_limit=0))]
 #[allow(clippy::too_many_arguments)]
 fn register_llm<'p>(
     py: Python<'p>,
@@ -142,6 +142,7 @@ fn register_llm<'p>(
     context_length: Option<u32>,
     kv_cache_block_size: Option<u32>,
     router_mode: Option<RouterMode>,
+    migration_limit: u32,
 ) -> PyResult<Bound<'p, PyAny>> {
     let model_type_obj = match model_type {
         ModelType::Chat => llm_rs::model_type::ModelType::Chat,
@@ -162,7 +163,8 @@ fn register_llm<'p>(
             .model_name(model_name)
             .context_length(context_length)
             .kv_cache_block_size(kv_cache_block_size)
-            .router_config(Some(router_config));
+            .router_config(Some(router_config))
+            .migration_limit(Some(migration_limit));
         // Download from HF, load the ModelDeploymentCard
         let mut local_model = builder.build().await.map_err(to_pyerr)?;
         // Advertise ourself on etcd so ingress can find us

lib/llm/src/discovery/watcher.rs

@@ -19,6 +19,7 @@ use dynamo_runtime::{
 use crate::{
     backend::Backend,
     kv_router::{KvPushRouter, KvRouterConfig},
+    migration::Migration,
     model_type::ModelType,
     preprocessor::{OpenAIPreprocessor, PreprocessedEmbeddingRequest, PreprocessedRequest},
     protocols::common::llm_backend::{EmbeddingsEngineOutput, LLMEngineOutput},
@@ -197,12 +198,14 @@ impl ModelWatcher {
                 // function. Needs checking carefully, possibly we need to store it in state.
                 let _cache_dir = Some(card.move_from_nats(self.drt.nats_client()).await?);
 
+                // Chat Completions
                 let frontend = SegmentSource::<
                     SingleIn<NvCreateChatCompletionRequest>,
                     ManyOut<Annotated<NvCreateChatCompletionStreamResponse>>,
                 >::new();
                 let preprocessor = OpenAIPreprocessor::new(card.clone()).await?.into_operator();
                 let backend = Backend::from_mdc(card.clone()).await?.into_operator();
+                let migration = Migration::from_mdc(card.clone()).await?.into_operator();
                 let router =
                     PushRouter::<PreprocessedRequest, Annotated<LLMEngineOutput>>::from_client(
                         client.clone(),
@@ -231,19 +234,23 @@ impl ModelWatcher {
                 let chat_engine = frontend
                     .link(preprocessor.forward_edge())?
                     .link(backend.forward_edge())?
+                    .link(migration.forward_edge())?
                     .link(service_backend)?
+                    .link(migration.backward_edge())?
                     .link(backend.backward_edge())?
                     .link(preprocessor.backward_edge())?
                     .link(frontend)?;
                 self.manager
                     .add_chat_completions_model(&model_entry.name, chat_engine)?;
 
+                // Completions
                 let frontend = SegmentSource::<
                     SingleIn<NvCreateCompletionRequest>,
                     ManyOut<Annotated<NvCreateCompletionResponse>>,
                 >::new();
                 let preprocessor = OpenAIPreprocessor::new(card.clone()).await?.into_operator();
                 let backend = Backend::from_mdc(card.clone()).await?.into_operator();
+                let migration = Migration::from_mdc(card.clone()).await?.into_operator();
                 let router =
                     PushRouter::<PreprocessedRequest, Annotated<LLMEngineOutput>>::from_client(
                         client,
@@ -272,7 +279,9 @@ impl ModelWatcher {
                 let completions_engine = frontend
                     .link(preprocessor.forward_edge())?
                     .link(backend.forward_edge())?
+                    .link(migration.forward_edge())?
                     .link(service_backend)?
+                    .link(migration.backward_edge())?
                     .link(backend.backward_edge())?
                     .link(preprocessor.backward_edge())?
                     .link(frontend)?;

lib/llm/src/lib.rs

@@ -22,6 +22,7 @@ pub mod hub;
 // pub mod key_value_store;
 pub mod kv_router;
 pub mod local_model;
+pub mod migration;
 pub mod mocker;
 pub mod model_card;
 pub mod model_type;