Resolve comments and kubernetes sglang router intregration

Test gateway package update

Test gateway package update

Test gateway package update

Test gateway package update

Test gateway package update

Resolve rate limits and location issue

Resolve rate limits and location issue

Resolve rate limits and location issue

Resolve all major comments

Resolve all major comments

Resolve kubernetes gateway issue with sglang intregration

Author: Bihan Rana

gateway/pyproject.toml

@@ -11,7 +11,7 @@ requires-python = ">=3.10"
 dynamic = ["version"]
 dependencies = [
     # release builds of dstack-gateway depend on a PyPI version of dstack instead
-    "dstack[gateway] @ git+https://github.com/Bihan/dstack.git@add_sglang_router_support",
+    "dstack[gateway] @ https://github.com/Bihan/dstack/archive/refs/heads/add_sglang_router_support.tar.gz",
 ]
 
 [project.optional-dependencies]

src/dstack/_internal/core/backends/base/compute.py

@@ -1014,26 +1014,25 @@ def get_latest_runner_build() -> Optional[str]:
     return None
 
 
-def get_dstack_gateway_wheel(build: str) -> str:
+def get_dstack_gateway_wheel(build: str, router: Optional[AnyRouterConfig] = None) -> str:
     channel = "release" if settings.DSTACK_RELEASE else "stgn"
     base_url = f"https://dstack-gateway-downloads.s3.amazonaws.com/{channel}"
     if build == "latest":
         r = requests.get(f"{base_url}/latest-version", timeout=5)
         r.raise_for_status()
         build = r.text.strip()
         logger.debug("Found the latest gateway build: %s", build)
-    # return f"{base_url}/dstack_gateway-{build}-py3-none-any.whl"
-    return "https://bihan-test-bucket.s3.eu-west-1.amazonaws.com/dstack_gateway-0.0.1-py3-none-any.whl"
+    # wheel = f"{base_url}/dstack_gateway-{build}-py3-none-any.whl"
+    wheel = "https://bihan-test-bucket.s3.eu-west-1.amazonaws.com/dstack_gateway-0.0.1-py3-none-any.whl"
+    # Build package spec with extras if router is specified
+    if router:
+        return f"dstack-gateway[{router.type}] @ {wheel}"
+    return f"dstack-gateway @ {wheel}"
 
 
 def get_dstack_gateway_commands(router: Optional[AnyRouterConfig] = None) -> List[str]:
     build = get_dstack_runner_version()
-    wheel = get_dstack_gateway_wheel(build)
-    # Build package spec with extras if router is specified
-    if router:
-        gateway_package = f"dstack-gateway[{router.type}] @ {wheel}"
-    else:
-        gateway_package = f"dstack-gateway @ {wheel}"
+    gateway_package = get_dstack_gateway_wheel(build, router)
     return [
         "mkdir -p /home/ubuntu/dstack",
         "python3 -m venv /home/ubuntu/dstack/blue",

src/dstack/_internal/core/backends/kubernetes/compute.py

@@ -1,3 +1,4 @@
+import shlex
 import subprocess
 import tempfile
 import threading
@@ -51,6 +52,7 @@
     SSHConnectionParams,
 )
 from dstack._internal.core.models.resources import CPUSpec, Memory
+from dstack._internal.core.models.routers import AnyRouterConfig
 from dstack._internal.core.models.runs import Job, JobProvisioningData, Requirements, Run
 from dstack._internal.core.models.volumes import Volume
 from dstack._internal.utils.common import parse_memory
@@ -403,7 +405,9 @@ def create_gateway(
         # Consider deploying an NLB. It seems it requires some extra configuration on the cluster:
         # https://docs.aws.amazon.com/eks/latest/userguide/network-load-balancing.html
         instance_name = generate_unique_gateway_instance_name(configuration)
-        commands = _get_gateway_commands(authorized_keys=[configuration.ssh_key_pub])
+        commands = _get_gateway_commands(
+            authorized_keys=[configuration.ssh_key_pub], router=configuration.router
+        )
         pod = client.V1Pod(
             metadata=client.V1ObjectMeta(
                 name=instance_name,
@@ -940,9 +944,13 @@ def _add_authorized_key_to_jump_pod(
     )
 
 
-def _get_gateway_commands(authorized_keys: List[str]) -> List[str]:
+def _get_gateway_commands(
+    authorized_keys: List[str], router: Optional[AnyRouterConfig] = None
+) -> List[str]:
     authorized_keys_content = "\n".join(authorized_keys).strip()
-    gateway_commands = " && ".join(get_dstack_gateway_commands())
+    gateway_commands = " && ".join(get_dstack_gateway_commands(router=router))
+    quoted_gateway_commands = shlex.quote(gateway_commands)
+
     commands = [
         # install packages
         "apt-get update && apt-get install -y sudo wget openssh-server nginx python3.10-venv libaugeas0",
@@ -971,7 +979,7 @@ def _get_gateway_commands(authorized_keys: List[str]) -> List[str]:
         # start sshd
         "/usr/sbin/sshd -p 22 -o PermitUserEnvironment=yes",
         # run gateway
-        f"su ubuntu -c '{gateway_commands}'",
+        f"su ubuntu -c {quoted_gateway_commands}",
         "sleep infinity",
     ]
     return commands

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

@@ -9,7 +9,7 @@
 
 def get_router(router: AnyRouterConfig, context: Optional[RouterContext] = None) -> Router:
     if router.type == "sglang":
-        return SglangRouter(router=router, context=context)
+        return SglangRouter(config=router, context=context)
     raise ProxyError(f"Router type '{router.type}' is not available")
 
 

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

@@ -28,13 +28,13 @@ class Router(ABC):
 
     def __init__(
         self,
-        router: Optional[AnyRouterConfig] = None,
+        config: Optional[AnyRouterConfig] = None,
         context: Optional[RouterContext] = None,
     ):
         """Initialize router with context.
 
         Args:
-            router: Optional router configuration (implementation-specific)
+            config: Optional router configuration (implementation-specific)
             context: Runtime context for the router (host, port, logging, etc.)
         """
         self.context = context or RouterContext()

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

@@ -1,12 +1,13 @@
-import json
 import shutil
 import subprocess
+import sys
 import time
 import urllib.parse
 from typing import List, Optional
 
+import httpx
+
 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
 
@@ -20,50 +21,43 @@ class SglangRouter(Router):
 
     TYPE = RouterType.SGLANG
 
-    def __init__(self, router: SGLangRouterConfig, context: Optional[RouterContext] = None):
+    def __init__(self, config: SGLangRouterConfig, context: Optional[RouterContext] = None):
         """Initialize SGLang router.
 
         Args:
-            router: SGLang router configuration (policy, cache_threshold, etc.)
+            config: SGLang router configuration (policy, cache_threshold, etc.)
             context: Runtime context for the router (host, port, logging, etc.)
         """
-        super().__init__(router=router, context=context)
-        self.config = router
+        super().__init__(config=config, context=context)
+        self.config = config
 
     def start(self) -> None:
         try:
             logger.info("Starting sglang-router-new on port %s...", self.context.port)
 
-            # Determine active venv (blue or green)
-            version_file = DSTACK_DIR_ON_GATEWAY / "version"
-            if version_file.exists():
-                version = version_file.read_text().strip()
-            else:
-                version = "blue"
-
-            venv_python = DSTACK_DIR_ON_GATEWAY / version / "bin" / "python3"
-
+            # Prometheus port is offset by 10000 from router port to keep it in a separate range
             prometheus_port = self.context.port + 10000
 
             cmd = [
-                str(venv_python),
+                sys.executable,
                 "-m",
                 "sglang_router.launch_router",
                 "--host",
-                "0.0.0.0",
+                self.context.host,
                 "--port",
                 str(self.context.port),
                 "--prometheus-port",
                 str(prometheus_port),
+                "--prometheus-host",
+                self.context.host,
                 "--log-level",
                 self.context.log_level,
                 "--log-dir",
                 str(self.context.log_dir),
+                "--policy",
+                self.config.policy,
             ]
 
-            if hasattr(self.config, "policy") and self.config.policy:
-                cmd.extend(["--policy", self.config.policy])
-
             subprocess.Popen(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
 
             time.sleep(2)
@@ -129,12 +123,9 @@ def stop(self) -> None:
     def is_running(self) -> bool:
         """Check if the SGLang router is running and responding to HTTP requests on the assigned port."""
         try:
-            result = subprocess.run(
-                ["curl", "-s", f"http://{self.context.host}:{self.context.port}/workers"],
-                capture_output=True,
-                timeout=5,
-            )
-            return result.returncode == 0
+            with httpx.Client(timeout=5.0) as client:
+                response = client.get(f"http://{self.context.host}:{self.context.port}/workers")
+                return response.status_code == 200
         except Exception as e:
             logger.error(f"Error checking sglang router status on port {self.context.port}: {e}")
             return False
@@ -151,8 +142,11 @@ def update_replicas(self, replica_urls: List[str]) -> None:
         for worker in current_workers:
             url = worker.get("url")
             if url and isinstance(url, str):
-                current_worker_urls.add(url)
-        target_worker_urls = set(replica_urls)
+                # Normalize URL by removing trailing slashes to avoid path artifacts
+                normalized_url = url.rstrip("/")
+                current_worker_urls.add(normalized_url)
+        # Normalize target URLs to ensure consistent comparison
+        target_worker_urls = {url.rstrip("/") for url in replica_urls}
 
         # Workers to add
         workers_to_add = target_worker_urls - current_worker_urls
@@ -186,87 +180,76 @@ def update_replicas(self, replica_urls: List[str]) -> None:
 
     def _get_router_workers(self) -> List[dict]:
         try:
-            result = subprocess.run(
-                ["curl", "-s", f"http://{self.context.host}:{self.context.port}/workers"],
-                capture_output=True,
-                timeout=5,
-            )
-            if result.returncode == 0:
-                response = json.loads(result.stdout.decode())
-                workers = response.get("workers", [])
-                return workers
-            return []
+            with httpx.Client(timeout=5.0) as client:
+                response = client.get(f"http://{self.context.host}:{self.context.port}/workers")
+                if response.status_code == 200:
+                    response_data = response.json()
+                    workers = response_data.get("workers", [])
+                    return workers
+                return []
         except Exception as e:
             logger.error(f"Error getting sglang router workers: {e}")
             return []
 
     def _add_worker_to_router(self, worker_url: str) -> bool:
         try:
             payload = {"url": worker_url, "worker_type": "regular"}
-            result = subprocess.run(
-                [
-                    "curl",
-                    "-X",
-                    "POST",
+            with httpx.Client(timeout=5.0) as client:
+                response = client.post(
                     f"http://{self.context.host}:{self.context.port}/workers",
-                    "-H",
-                    "Content-Type: application/json",
-                    "-d",
-                    json.dumps(payload),
-                ],
-                capture_output=True,
-                timeout=5,
-            )
-
-            if result.returncode == 0:
-                response = json.loads(result.stdout.decode())
-                if response.get("status") == "accepted":
-                    logger.info(
-                        "Added worker %s to sglang router on port %s",
+                    json=payload,
+                )
+                if response.status_code == 200:
+                    response_data = response.json()
+                    if response_data.get("status") == "accepted":
+                        logger.info(
+                            "Added worker %s to sglang router on port %s",
+                            worker_url,
+                            self.context.port,
+                        )
+                        return True
+                    else:
+                        logger.error("Failed to add worker %s: %s", worker_url, response_data)
+                        return False
+                else:
+                    logger.error(
+                        "Failed to add worker %s: status %d, %s",
                         worker_url,
-                        self.context.port,
+                        response.status_code,
+                        response.text,
                     )
-                    return True
-                else:
-                    logger.error("Failed to add worker %s: %s", worker_url, response)
                     return False
-            else:
-                logger.error("Failed to add worker %s: %s", worker_url, result.stderr.decode())
-                return False
         except Exception as e:
             logger.error(f"Error adding worker {worker_url}: {e}")
             return False
 
     def _remove_worker_from_router(self, worker_url: str) -> bool:
         try:
             encoded_url = urllib.parse.quote(worker_url, safe="")
-
-            result = subprocess.run(
-                [
-                    "curl",
-                    "-X",
-                    "DELETE",
-                    f"http://{self.context.host}:{self.context.port}/workers/{encoded_url}",
-                ],
-                capture_output=True,
-                timeout=5,
-            )
-
-            if result.returncode == 0:
-                response = json.loads(result.stdout.decode())
-                if response.get("status") == "accepted":
-                    logger.info(
-                        "Removed worker %s from sglang router on port %s",
+            with httpx.Client(timeout=5.0) as client:
+                response = client.delete(
+                    f"http://{self.context.host}:{self.context.port}/workers/{encoded_url}"
+                )
+                if response.status_code == 200:
+                    response_data = response.json()
+                    if response_data.get("status") == "accepted":
+                        logger.info(
+                            "Removed worker %s from sglang router on port %s",
+                            worker_url,
+                            self.context.port,
+                        )
+                        return True
+                    else:
+                        logger.error("Failed to remove worker %s: %s", worker_url, response_data)
+                        return False
+                else:
+                    logger.error(
+                        "Failed to remove worker %s: status %d, %s",
                         worker_url,
-                        self.context.port,
+                        response.status_code,
+                        response.text,
                     )
-                    return True
-                else:
-                    logger.error("Failed to remove worker %s: %s", worker_url, response)
                     return False
-            else:
-                logger.error("Failed to remove worker %s: %s", worker_url, result.stderr.decode())
-                return False
         except Exception as e:
             logger.error(f"Error removing worker {worker_url}: {e}")
             return False