Resolve gateway installation command

Author: Bihan Rana

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

@@ -1029,17 +1029,16 @@ def get_dstack_gateway_wheel(build: str) -> str:
 def get_dstack_gateway_commands(router: Optional[AnyRouterConfig] = None) -> List[str]:
     build = get_dstack_runner_version()
     wheel = get_dstack_gateway_wheel(build)
-    # Use router type directly as pip extra
+    # Build package spec with extras if router is specified
     if router:
-        gateway_package = f"dstack-gateway[{router.type}]"
+        gateway_package = f"dstack-gateway[{router.type}] @ {wheel}"
     else:
-        gateway_package = "dstack-gateway"
+        gateway_package = f"dstack-gateway @ {wheel}"
     return [
         "mkdir -p /home/ubuntu/dstack",
         "python3 -m venv /home/ubuntu/dstack/blue",
         "python3 -m venv /home/ubuntu/dstack/green",
-        f"/home/ubuntu/dstack/blue/bin/pip install {wheel}",
-        f"/home/ubuntu/dstack/blue/bin/pip install --upgrade '{gateway_package}'",
+        f"/home/ubuntu/dstack/blue/bin/pip install '{gateway_package}'",
         "sudo /home/ubuntu/dstack/blue/bin/python -m dstack.gateway.systemd install --run",
     ]
 

src/dstack/_internal/server/services/gateways/__init__.py

@@ -454,20 +454,19 @@ async def _update_gateway(gateway_compute_model: GatewayComputeModel, build: str
         gateway_compute_model.configuration
     )
 
-    # Determine gateway package with router extras (similar to compute.py)
+    # Build package spec with extras and wheel URL (similar to compute.py)
+    wheel = get_dstack_gateway_wheel(build)
     if compute_config.router:
-        gateway_package = f"dstack-gateway[{compute_config.router.type}]"
+        gateway_package = f"dstack-gateway[{compute_config.router.type}] @ {wheel}"
     else:
-        gateway_package = "dstack-gateway"
+        gateway_package = f"dstack-gateway @ {wheel}"
 
     commands = [
         # prevent update.sh from overwriting itself during execution
         "cp dstack/update.sh dstack/_update.sh",
-        f"sh dstack/_update.sh {get_dstack_gateway_wheel(build)} {build}",
+        # Pass the full package spec (with extras) to update.sh instead of just the wheel
+        f"sh dstack/_update.sh '{gateway_package}' {build}",
         "rm dstack/_update.sh",
-        # Install gateway package with router extras to the active venv (blue or green)
-        # update.sh writes the active version to dstack/version
-        f"version=$(cat /home/ubuntu/dstack/version) && /home/ubuntu/dstack/$version/bin/pip install --upgrade '{gateway_package}'",
     ]
     stdout = await connection.tunnel.aexec("/bin/sh -c '" + " && ".join(commands) + "'")
     if "Update successfully completed" in stdout: