Resolved Minor Review Comments

Author: Bihan Rana

src/dstack/_internal/core/models/routers.py

@@ -1,27 +1,16 @@
 from enum import Enum
-from typing import Union
-
-from pydantic import Field
-from typing_extensions import Annotated, Literal
+from typing import Literal
 
 from dstack._internal.core.models.common import CoreModel
 
 
 class RouterType(str, Enum):
     SGLANG = "sglang"
-    VLLM = "vllm"
 
 
 class SGLangRouterConfig(CoreModel):
     type: Literal["sglang"] = "sglang"
-    policy: str = "cache_aware"
-
-
-class VLLMRouterConfig(CoreModel):
-    type: Literal["vllm"] = "vllm"
-    policy: str = "cache_aware"
+    policy: Literal["random", "round_robin", "cache_aware", "power_of_two"] = "cache_aware"
 
 
-AnyRouterConfig = Annotated[
-    Union[SGLangRouterConfig, VLLMRouterConfig], Field(discriminator="type")
-]
+AnyRouterConfig = SGLangRouterConfig

src/dstack/_internal/proxy/gateway/model_routers/__init__.py

@@ -2,15 +2,15 @@
 
 from dstack._internal.core.models.routers import AnyRouterConfig
 from dstack._internal.proxy.gateway.model_routers.sglang import SglangRouter
+from dstack._internal.proxy.lib.errors import ProxyError
 
 from .base import Router, RouterContext
 
 
 def get_router(router: AnyRouterConfig, context: Optional[RouterContext] = None) -> Router:
-    """Factory function to create a router instance from router configuration."""
     if router.type == "sglang":
         return SglangRouter(router=router, context=context)
-    raise ValueError(f"Router type '{router.type}' is not available")
+    raise ProxyError(f"Router type '{router.type}' is not available")
 
 
 __all__ = [

src/dstack/_internal/proxy/gateway/model_routers/sglang.py

@@ -5,8 +5,9 @@
 import urllib.parse
 from typing import List, Optional
 
-from dstack._internal.core.models.routers import SGLangRouterConfig
+from dstack._internal.core.models.routers import RouterType, SGLangRouterConfig
 from dstack._internal.proxy.gateway.const import DSTACK_DIR_ON_GATEWAY
+from dstack._internal.proxy.lib.errors import UnexpectedProxyError
 from dstack._internal.utils.logging import get_logger
 
 from .base import Router, RouterContext
@@ -17,7 +18,7 @@
 class SglangRouter(Router):
     """SGLang router implementation with 1:1 service-to-router."""
 
-    TYPE = "sglang"
+    TYPE = RouterType.SGLANG
 
     def __init__(self, router: SGLangRouterConfig, context: Optional[RouterContext] = None):
         """Initialize SGLang router.
@@ -68,7 +69,9 @@ def start(self) -> None:
             time.sleep(2)
 
             if not self.is_running():
-                raise Exception(f"Failed to start sglang router on port {self.context.port}")
+                raise UnexpectedProxyError(
+                    f"Failed to start sglang router on port {self.context.port}"
+                )
 
             logger.info(
                 "Sglang router started successfully on port %s (prometheus on %s)",

src/dstack/_internal/proxy/gateway/services/nginx.py

@@ -10,7 +10,7 @@
 from pydantic import BaseModel
 from typing_extensions import Literal
 
-from dstack._internal.core.models.routers import AnyRouterConfig
+from dstack._internal.core.models.routers import AnyRouterConfig, RouterType
 from dstack._internal.proxy.gateway.const import PROXY_PORT_ON_GATEWAY
 from dstack._internal.proxy.gateway.model_routers import (
     Router,
@@ -107,7 +107,7 @@ async def register(self, conf: SiteConfig, acme: ACMESettings) -> None:
                 await run_async(self.run_certbot, conf.domain, acme)
 
             if isinstance(conf, ServiceConfig) and conf.router:
-                if conf.router.type == "sglang":
+                if conf.router.type == RouterType.SGLANG:
                     # Check if router already exists for this domain
                     if conf.domain in self._domain_to_router:
                         # Router already exists, reuse it
@@ -314,7 +314,7 @@ def _is_port_available(port: int) -> bool:
             return False
 
     def _allocate_router_port(self) -> int:
-        """Allocate next available router port in fixed range (20000-24999).
+        """Allocate next available router port in fixed range.
 
         Checks both our internal allocation map and actual port availability
         to avoid conflicts with other services. Range chosen to avoid ephemeral ports.
@@ -354,7 +354,7 @@ def _allocate_router_port(self) -> int:
         )
 
     def _allocate_worker_ports(self, num_ports: int) -> list[int]:
-        """Allocate worker ports globally in fixed range (10001-11999).
+        """Allocate worker ports globally in fixed range.
 
         Worker ports are used by nginx to listen and proxy to worker sockets.
         They must be unique across all router instances. Range chosen to avoid ephemeral ports.