feat(router): add load-aware generate paths

Signed-off-by: Bruno Volpato <brunocvcunha@gmail.com>

Author: bvolpato

README.md

@@ -77,6 +77,20 @@ vllm-router \
     --intra-node-data-parallel-size 8
 ```
 
+#### Additional Generate Paths
+
+Register extension endpoints that use the `/inference/v1/generate` request
+schema. Configured paths get typed routing, active-load accounting, retries,
+and circuit breaking while preserving extension fields in requests and
+responses. Workers must expose the same paths.
+
+```bash
+vllm-router \
+  --worker-urls http://worker1:8000 http://worker2:8000 \
+  --policy cache_aware \
+  --extra-generate-paths /custom/v1/generate
+```
+
 #### Prefill-Decode Disaggregation
 ```bash
 # When vLLM runs the NIXL connector, prefill/decode URLs are required.

py_src/vllm_router/router.py

@@ -26,6 +26,8 @@ class Router:
     Args:
         worker_urls: List of URLs for worker nodes that will handle requests. Each URL should include
             the protocol, host, and port (e.g., ['http://worker1:8000', 'http://worker2:8000'])
+        extra_generate_paths: Additional HTTP paths that use the inference-generate request schema.
+            Requests on these paths use typed routing, retries, circuit breaking, and load tracking.
         policy: Load balancing policy to use. Options:
             - PolicyType.Random: Randomly select workers
             - PolicyType.RoundRobin: Distribute requests in round-robin fashion

py_src/vllm_router/router_args.py

@@ -5,6 +5,30 @@
 
 logger = logging.getLogger(__name__)
 
+_RESERVED_GENERATE_PATHS = {
+    "/generate",
+    "/inference/v1/generate",
+    "/v1/chat/completions",
+    "/v1/completions",
+    "/rerank",
+    "/v1/rerank",
+    "/v1/responses",
+    "/v1/embeddings",
+    "/liveness",
+    "/readiness",
+    "/health",
+    "/health_generate",
+    "/v1/models",
+    "/get_model_info",
+    "/get_server_info",
+    "/add_worker",
+    "/remove_worker",
+    "/list_workers",
+    "/flush_cache",
+    "/get_loads",
+    "/workers",
+}
+
 
 @dataclasses.dataclass
 class RouterArgs:
@@ -89,6 +113,8 @@ class RouterArgs:
     cb_timeout_duration_secs: int = 60
     cb_window_duration_secs: int = 120
     disable_circuit_breaker: bool = False
+    # Additional typed inference-generate routes
+    extra_generate_paths: List[str] = dataclasses.field(default_factory=list)
 
     @staticmethod
     def add_cli_args(
@@ -128,6 +154,13 @@ def add_cli_args(
             default=[],
             help="List of worker URLs (e.g., http://worker1:8000 http://worker2:8000)",
         )
+        parser.add_argument(
+            f"--{prefix}extra-generate-paths",
+            type=str,
+            nargs="*",
+            default=[],
+            help="Additional HTTP paths using the inference-generate request schema",
+        )
 
         # Routing policy configuration
         parser.add_argument(
@@ -518,6 +551,19 @@ def from_cli_args(
         return cls(**args_dict)
 
     def _validate_router_args(self):
+        if len(set(self.extra_generate_paths)) != len(self.extra_generate_paths):
+            raise ValueError("extra_generate_paths must not contain duplicates")
+        for path in self.extra_generate_paths:
+            if (
+                len(path) < 2
+                or not path.startswith("/")
+                or path in _RESERVED_GENERATE_PATHS
+                or path.startswith("/v1/responses/")
+                or path.startswith("/workers/")
+                or any(character in path for character in "?#{}*")
+            ):
+                raise ValueError(f"invalid extra generate path: {path}")
+
         # Validate configuration based on mode
         if self.vllm_pd_disaggregation:
             # Validate PD configuration - skip URL requirements if using service discovery

py_test/unit/test_arg_parser.py

@@ -24,6 +24,7 @@ def test_default_values(self):
         assert args.port == 30000
         assert args.policy == "cache_aware"
         assert args.worker_urls == []
+        assert args.extra_generate_paths == []
         assert args.vllm_pd_disaggregation is False
         assert args.prefill_urls == []
         assert args.decode_urls == []
@@ -440,6 +441,22 @@ def test_parse_basic_args(self):
         assert router_args.worker_urls == ["http://worker1:8000", "http://worker2:8000"]
         assert router_args.policy == "round_robin"
 
+    def test_parse_extra_generate_paths(self):
+        router_args = parse_router_args(
+            [
+                "--worker-urls",
+                "http://worker1:8000",
+                "--extra-generate-paths",
+                "/custom/v1/generate",
+                "/extension/generate",
+            ]
+        )
+
+        assert router_args.extra_generate_paths == [
+            "/custom/v1/generate",
+            "/extension/generate",
+        ]
+
     def test_parse_pd_args(self):
         """Test parsing PD disaggregated mode arguments."""
         args = [

py_test/unit/test_validation.py

@@ -399,6 +399,29 @@ def test_request_id_headers_validation(self):
 class TestLaunchValidation:
     """Test launch-time validation logic."""
 
+    def test_extra_generate_paths_accept_distinct_extension_routes(self):
+        args = RouterArgs(
+            extra_generate_paths=["/custom/v1/generate", "/extension/generate"]
+        )
+
+        args._validate_router_args()
+
+    @pytest.mark.parametrize(
+        "paths",
+        [
+            ["custom/v1/generate"],
+            ["/inference/v1/generate"],
+            ["/v1/responses/id"],
+            ["/custom/{path}"],
+            ["/custom/v1/generate", "/custom/v1/generate"],
+        ],
+    )
+    def test_extra_generate_paths_reject_invalid_or_conflicting_routes(self, paths):
+        args = RouterArgs(extra_generate_paths=paths)
+
+        with pytest.raises(ValueError, match="extra_generate_paths|extra generate path"):
+            args._validate_router_args()
+
     def test_pd_mode_requires_urls(self):
         """Test that PD mode requires prefill and decode URLs."""
         # PD mode without URLs should fail

src/core/worker.rs

@@ -884,10 +884,6 @@ pub fn start_health_checker(
         let mut interval =
             tokio::time::interval(tokio::time::Duration::from_secs(check_interval_secs));
 
-        // Counter for periodic load reset (every 10 health check cycles)
-        let mut check_count = 0u64;
-        const LOAD_RESET_INTERVAL: u64 = 10;
-
         loop {
             interval.tick().await;
 
@@ -897,8 +893,6 @@ pub fn start_health_checker(
                 break;
             }
 
-            check_count += 1;
-
             // Check health of all workers
             let workers_to_check = match workers.read() {
                 Ok(guard) => guard.clone(),
@@ -908,22 +902,6 @@ pub fn start_health_checker(
                 }
             };
 
-            // Periodically reset load counters to prevent drift
-            // Only do this when we believe all workers should be idle
-            if check_count.is_multiple_of(LOAD_RESET_INTERVAL) {
-                let max_load = workers_to_check.iter().map(|w| w.load()).max().unwrap_or(0);
-                // Only reset if load appears to be very low (likely drift)
-                if max_load <= 2 {
-                    tracing::debug!(
-                        "Resetting load counters to prevent drift (max_load: {})",
-                        max_load
-                    );
-                    for worker in &workers_to_check {
-                        worker.reset_load();
-                    }
-                }
-            }
-
             // Perform health checks concurrently
             let health_checks = workers_to_check.iter().map(|worker| {
                 let worker_url = worker.url().to_string();

src/core/worker_registry.rs

@@ -364,10 +364,6 @@ impl WorkerRegistry {
             let mut interval =
                 tokio::time::interval(tokio::time::Duration::from_secs(check_interval_secs));
 
-            // Counter for periodic load reset (every 10 health check cycles)
-            let mut check_count = 0u64;
-            const LOAD_RESET_INTERVAL: u64 = 10;
-
             loop {
                 interval.tick().await;
 
@@ -387,15 +383,6 @@ impl WorkerRegistry {
                 for worker in &workers {
                     let _ = worker.check_health_async().await; // Use async version directly
                 }
-
-                // Reset loads periodically
-                check_count += 1;
-                if check_count.is_multiple_of(LOAD_RESET_INTERVAL) {
-                    tracing::debug!("Resetting worker loads (cycle {})", check_count);
-                    for worker in &workers {
-                        worker.reset_load();
-                    }
-                }
             }
         });
 

src/lib.rs

@@ -34,6 +34,7 @@ struct Router {
     host: String,
     port: u16,
     worker_urls: Vec<String>,
+    extra_generate_paths: Vec<String>,
     policy: PolicyType,
     worker_startup_timeout_secs: u64,
     worker_startup_check_interval: u64,
@@ -310,6 +311,7 @@ impl Router {
         otlp_traces_endpoint = None,
         // KV connector default (PD disaggregation)
         kv_connector = String::from("nixl"),
+        extra_generate_paths = vec![],
     ))]
     #[allow(clippy::too_many_arguments)]
     fn new(
@@ -372,11 +374,13 @@ impl Router {
         enable_trace: bool,
         otlp_traces_endpoint: Option<String>,
         kv_connector: String,
+        extra_generate_paths: Vec<String>,
     ) -> PyResult<Self> {
         Ok(Router {
             host,
             port,
             worker_urls,
+            extra_generate_paths,
             policy,
             worker_startup_timeout_secs,
             worker_startup_check_interval,
@@ -502,6 +506,7 @@ impl Router {
                 } else {
                     None
                 },
+                extra_generate_paths: self.extra_generate_paths.clone(),
             })
             .await
             .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))

src/main.rs

@@ -106,6 +106,10 @@ struct CliArgs {
     #[arg(long, num_args = 0..)]
     worker_urls: Vec<String>,
 
+    /// Additional HTTP paths using the inference-generate request schema
+    #[arg(long, num_args = 0..)]
+    extra_generate_paths: Vec<String>,
+
     /// Load balancing policy to use
     #[arg(long, default_value = "cache_aware", value_parser = ["random", "round_robin", "cache_aware", "power_of_two", "consistent_hash", "rendezvous_hash"])]
     policy: String,
@@ -614,6 +618,7 @@ impl CliArgs {
             } else {
                 None
             },
+            extra_generate_paths: self.extra_generate_paths.clone(),
         }
     }
 }

src/routers/http/openai_router.rs

@@ -223,6 +223,16 @@ impl super::super::RouterTrait for OpenAIRouter {
             .into_response()
     }
 
+    async fn route_inference_generate_path(
+        &self,
+        headers: Option<&HeaderMap>,
+        body: &InferenceGenerateRequest,
+        _path: &str,
+        model_id: Option<&str>,
+    ) -> Response {
+        self.route_inference_generate(headers, body, model_id).await
+    }
+
     async fn route_chat(
         &self,
         headers: Option<&HeaderMap>,