diff --git a/gateway/src/dstack/gateway/resources/systemd/update.sh b/gateway/src/dstack/gateway/resources/systemd/update.sh deleted file mode 100755 index 99b0a48f2e..0000000000 --- a/gateway/src/dstack/gateway/resources/systemd/update.sh +++ /dev/null @@ -1,42 +0,0 @@ -#!/bin/sh -set -e -root="$( cd -- "$(dirname "$0")" >/dev/null 2>&1; pwd -P )" - -# usage: ./update.sh -if [ "$#" -eq 0 ]; then - echo "Error: Missing wheel URL" - exit 1 -elif [ "$#" -eq 1 ]; then - echo "Error: Missing build version" - exit 1 -fi - -if [ -f "$root/version" ]; then - version=$(cat "$root/version") # blue/green -else - version="blue" -fi - -# check the current build version -current_build=$($root/$version/bin/pip show dstack-gateway | grep Version | awk '{print $2}') -if [ "$current_build" = "$2" ]; then - echo "The build $2 is already installed. Skipping..." - exit 0 -fi - -# flip the version -if [ "$version" = "blue" ]; then - version="green" -else - version="blue" -fi - -"$root/$version/bin/pip" uninstall -y dstack-gateway dstack -"$root/$version/bin/pip" cache remove dstack -"$root/$version/bin/pip" install "$1" -sudo "$root/$version/bin/python" -m dstack.gateway.systemd install -echo "$version" > "$root/version" -sudo systemctl daemon-reload -sudo systemctl restart dstack.gateway - -echo "Update successfully completed" diff --git a/gateway/src/dstack/gateway/systemd/__main__.py b/gateway/src/dstack/gateway/systemd/__main__.py index 2514fe7b90..793bf876f7 100644 --- a/gateway/src/dstack/gateway/systemd/__main__.py +++ b/gateway/src/dstack/gateway/systemd/__main__.py @@ -35,7 +35,7 @@ def install_action(args): ) service_path.write_text(service_file.format(working_dir=working_dir.as_posix())) - for script_name in ["start.sh", "update.sh"]: + for script_name in ["start.sh"]: print(f"Writing {script_name} script...") script = importlib.resources.read_text("dstack.gateway.resources.systemd", script_name) script_path = working_dir / script_name diff --git a/src/dstack/_internal/server/services/gateways/__init__.py b/src/dstack/_internal/server/services/gateways/__init__.py index e76f530f9a..339cfc2177 100644 --- a/src/dstack/_internal/server/services/gateways/__init__.py +++ b/src/dstack/_internal/server/services/gateways/__init__.py @@ -1,6 +1,7 @@ import asyncio import datetime import itertools +import shlex import uuid from collections.abc import AsyncGenerator from contextlib import asynccontextmanager @@ -827,19 +828,53 @@ async def _update_gateway_replica(gateway_replica_model: GatewayReplicaModel, bu # Build package spec with extras and wheel URL gateway_package = get_dstack_gateway_wheel(build) - commands = [ - # prevent update.sh from overwriting itself during execution - "cp dstack/update.sh dstack/_update.sh", - f'sh dstack/_update.sh "{gateway_package}" {build}', - "rm dstack/_update.sh", - ] - stdout = await connection.tunnel.aexec("/bin/sh -c '" + " && ".join(commands) + "'") + command = ( + "/bin/sh -c " + + shlex.quote(_GATEWAY_UPDATE_SCRIPT) + + " sh " # $0 placeholder + + shlex.quote(gateway_package) + + " " + + shlex.quote(build) + ) + stdout = await connection.tunnel.aexec(command) if "Update successfully completed" in stdout: logger.info("Gateway replica %s updated", connection.ip_address) return True return False +# Blue/green: install the new build into the currently inactive venv and flip to it +_GATEWAY_UPDATE_SCRIPT = """\ +set -e +gateway_package="$1" +build="$2" +root=/home/ubuntu/dstack +if [ -f "$root/version" ]; then + version=$(cat "$root/version") +else + version=blue +fi +current_build=$("$root/$version/bin/pip" show dstack-gateway | grep Version | awk '{print $2}') +if [ "$current_build" = "$build" ]; then + echo "The build $build is already installed. Skipping..." + exit 0 +fi +if [ "$version" = blue ]; then + version=green +else + version=blue +fi +"$root/$version/bin/pip" uninstall -y dstack-gateway dstack +"$root/$version/bin/pip" cache remove dstack +"$root/$version/bin/pip" install "$gateway_package" +sudo "$root/$version/bin/python" -m dstack.gateway.systemd install +echo "$version" > "$root/version" +sudo systemctl daemon-reload +sudo systemctl restart dstack.gateway +echo "Update successfully completed" +""" + + def _recently_updated(gateway_replica_model: GatewayReplicaModel) -> bool: return gateway_replica_model.app_updated_at.replace( tzinfo=datetime.timezone.utc