diff --git a/.github/workflows/build-artifacts.yml b/.github/workflows/build-artifacts.yml index 26d4b6d1d..d9e8ae0cb 100644 --- a/.github/workflows/build-artifacts.yml +++ b/.github/workflows/build-artifacts.yml @@ -85,19 +85,9 @@ jobs: with: name: frontend-build path: src/dstack/_internal/server/statics - - name: Run pytest on POSIX - if: matrix.os != 'windows-latest' - # Skip Postgres tests on macos since macos runner doesn't have Docker. - run: | - RUNPOSTGRES="" - if [ "${{ matrix.os }}" != "macos-latest" ]; then - RUNPOSTGRES="--runpostgres" - fi - uv run pytest -n auto src/tests --runui $RUNPOSTGRES - - name: Run pytest on Windows - if: matrix.os == 'windows-latest' - run: | - uv run pytest -n auto src/tests --runui --runpostgres + - name: Run pytest + # Skip Postgres and Docker tests on macos since macos runner doesn't have Docker. + run: uv run pytest -n auto src/tests ${{ matrix.os == 'macos-latest' && '--runui' || '--runall' }} runner-test: defaults: diff --git a/AGENTS.md b/AGENTS.md index c9f86f765..daacac7de 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -36,7 +36,7 @@ Before touching a subsystem, read the relevant notes in `contributing/`: `ARCHIT - Preserve client/server backward compatibility when updating Pydantic models. New fields must be excluded from client requests when unset, for compatibility with older servers (see `core/compatibility/*.py`). If the allowed domain/type of an existing field is extended, server responses may need to be patched for older clients (see `server/compatibility/*.py`). No need to exclude new fields from server responses, since clients rely on `validate_extra_ignore`. ## Testing Guidelines -- Default to `uv run pytest`. Use markers from `src/tests/conftest.py` like `--runpostgres` if need to include specific tests. +- Default to `uv run pytest` for quick testing. Opt-in tests are skipped by default; enable them when the change touches what they cover, with the flags from `src/tests/conftest.py`: `--runpostgres`, `--rundocker`, `--runui`, or `--runall` for all of them. - Scope the run to the change: for trivial or localized edits, run only the affected test modules, `Test*` classes, or `-k` selection instead of the whole suite. Reserve the full suite for broad or cross-cutting changes. - Speed up large runs with `-n auto` (pytest-xdist), e.g. `uv run pytest -n auto`. - Group tests for the same unit (function/class) using `Test*` classes that mirror unit's name. diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index cf398ba0c..6329cf6fe 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -55,11 +55,15 @@ Then you can do: uv run pytest src/tests ``` -(Optionally) By default, tests run against SQLite. -Use the `--runpostgres` flag to run the tests against Postgres as well: +(Optionally) Some tests are skipped by default because they need extra setup. Enable them with these flags: + +* `--runpostgres` – run the tests against Postgres in addition to SQLite. Requires Docker. +* `--rundocker` – run the tests that use Docker containers. Requires Docker. +* `--runui` – run the tests of the bundled frontend. Requires the frontend build in `src/dstack/_internal/server/statics`. +* `--runall` – all of the above. ```shell -uv run pytest src/tests --runpostgres +uv run pytest src/tests --runall ``` Alternatively, you can run tests via [tox](https://tox.wiki/) inside an isolated environment ensuring that the `dstack` package itself is built correctly and all requirements are specified and correct. tox and [just](https://just.systems/) must be already installed. diff --git a/contributing/PROXY.md b/contributing/PROXY.md index a8dc522cc..d053b43d8 100644 --- a/contributing/PROXY.md +++ b/contributing/PROXY.md @@ -136,19 +136,21 @@ Users can create a gateway using the `dstack apply` command. The gateway YAML co Provisioning happens as follows: 1. Launch a non-GPU instance (usually the smallest) with all ports exposed. 2. Install Nginx, Certbot, and patch configs. -3. Create blue-green virtual environments. -4. Install the `dstack[gateway]` Python package. -5. Run the systemd service `dstack.gateway.service`. +3. Run the gateway install script (`core/backends/base/resources/gateway_install.sh`), which installs uv, Python, and the `dstack[gateway]` package and runs the systemd service `dstack.gateway.service`. The same script is used for updates, see below. ### Update -The gateway has a "blue-green deployment"-like configuration: there are two virtual environments to be swapped on update. The systemd service uses the newly installed package after a restart. +The gateway has a "blue-green deployment"-like configuration: the app runs from one of two virtual environments, `blue` or `green`, as recorded in `~/dstack/version`. The systemd service uses the newly installed package after a restart. -The update process looks like this: -1. Install the new package to the unused venv. -2. Update scripts and systemd service config. -3. Swap the active venv name in `~/dstack/version`. -4. Restart the systemd service. +On startup, the `dstack` server runs the install script on every gateway over SSH. The script: +1. Exits if the active venv already has the target `dstack` version on the script's `python_version`. Unversioned packages (master builds, `DSTACK_GATEWAY_PACKAGE_URL`) are always reinstalled. +2. Installs uv (`uv_version`) and Python (`python_version`) if missing. Both are pinned in the script. +3. Recreates the unused venv and installs the new package into it. +4. Updates scripts and systemd service config. +5. Swaps the active venv name in `~/dstack/version`. +6. Restarts the systemd service. + +Changing `python_version` migrates existing gateways to the new Python on their next update. The venvs are created with pip (`uv venv --seed`) because update scripts of older servers call pip directly, e.g. after a server downgrade. ## Gateway development diff --git a/src/dstack/_internal/core/backends/base/compute.py b/src/dstack/_internal/core/backends/base/compute.py index b2fb2c8ff..2efc68ba5 100644 --- a/src/dstack/_internal/core/backends/base/compute.py +++ b/src/dstack/_internal/core/backends/base/compute.py @@ -1,3 +1,4 @@ +import importlib.resources import os import random import re @@ -1275,12 +1276,12 @@ def get_gateway_user_data(authorized_key: str) -> str: package_update=True, packages=[ "nginx", - "python3.10-venv", + "wget", ], snap={"commands": [["install", "--classic", "certbot"]]}, runcmd=[ ["ln", "-s", "/snap/bin/certbot", "/usr/bin/certbot"], - ["su", "ubuntu", "-c", " && ".join(get_dstack_gateway_commands())], + ["su", "ubuntu", "-c", get_dstack_gateway_install_command()], ], ssh_authorized_keys=[authorized_key], ) @@ -1389,15 +1390,30 @@ def get_dstack_gateway_package_and_target_version() -> tuple[str, str | None]: return package, None -def get_dstack_gateway_commands() -> List[str]: - gateway_package, _ = get_dstack_gateway_package_and_target_version() - 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 '{gateway_package}'", - "sudo /home/ubuntu/dstack/blue/bin/python -m dstack._internal.proxy.gateway.systemd install --run", - ] +def get_dstack_gateway_install_command() -> str: + """ + Returns a shell command that installs the gateway app on a fresh gateway host or updates it + on an existing one. Must be run as the `ubuntu` user. + """ + gateway_package, target_version = get_dstack_gateway_package_and_target_version() + return shlex.join( + [ + "/bin/sh", + "-c", + _get_dstack_gateway_install_script(), + "sh", # $0 + gateway_package, + target_version or "", + ] + ) + + +def _get_dstack_gateway_install_script() -> str: + return ( + importlib.resources.files("dstack._internal.core.backends.base") + .joinpath("resources/gateway_install.sh") + .read_text() + ) def merge_tags( diff --git a/src/dstack/_internal/core/backends/base/resources/gateway_install.sh b/src/dstack/_internal/core/backends/base/resources/gateway_install.sh new file mode 100644 index 000000000..e098eb308 --- /dev/null +++ b/src/dstack/_internal/core/backends/base/resources/gateway_install.sh @@ -0,0 +1,71 @@ +# Installs the dstack gateway app on a fresh gateway host or updates it on an existing one. +# Usage: sh -c "$script" sh +# is empty for unversioned packages (master builds, custom URLs). +# +# Blue/green: the app runs from one of two venvs, blue or green, named in the version file. +# The inactive venv is recreated from scratch, the dstack package is installed into it, +# and the gateway service is restarted from it. +# +# Python and the venvs are managed by uv, which is installed by the script, +# so the gateway doesn't depend on the OS Python. Changing python_version migrates existing +# gateways to the new Python on their next update. +# +# On a fresh host there is no version file and no venvs. Blue is considered active but fails +# the version check, and the script installs uv, Python, and the package into green, creates +# an empty blue venv, writes the systemd service, and starts it. +set -e +gateway_package="$1" +build="$2" +gateway_dir=/home/ubuntu/dstack +uv_version=0.13.1 +python_version=3.12 +uv="$gateway_dir/bin/uv" +# To prevent the uv cache grow and avoid pruning the cache, re-download all deps, which takes a few seconds. +export UV_NO_CACHE=1 + +mkdir -p "$gateway_dir" +cd "$gateway_dir" +if [ -f "$gateway_dir/version" ]; then + active=$(cat "$gateway_dir/version") +else + active=blue +fi +if [ -n "$build" ]; then + current=$("$gateway_dir/$active/bin/python" -c 'import sys, importlib.metadata as m; print("%d.%d %s" % (sys.version_info[0], sys.version_info[1], m.version("dstack")))' 2>/dev/null || true) + if [ "$current" = "$python_version $build" ]; then + echo "The build $build on Python $python_version is already installed. Skipping..." + exit 0 + fi +fi +if [ "$active" = blue ]; then + version=green +else + version=blue +fi +# The static binary from the uv release needs no Python to bootstrap and runs on any Linux. +# GitHub is required anyway, as uv downloads Python builds from there. +if [ "$("$uv" --version 2>/dev/null | cut -d ' ' -f 2)" != "$uv_version" ]; then + archive=$(mktemp) + wget -qO "$archive" "https://github.com/astral-sh/uv/releases/download/$uv_version/uv-$(uname -m)-unknown-linux-musl.tar.gz" + mkdir -p "$gateway_dir/bin" + tar -xzf "$archive" --strip-components=1 -C "$gateway_dir/bin" + rm -f "$archive" +fi +"$uv" python install "$python_version" +# The venv may contain root-owned files, e.g. bytecode written when installing the systemd +# service from it with sudo +sudo rm -rf "${gateway_dir:?}/${version:?}" +# --seed installs pip. If the dstack server is downgraded, the older server runs its own update +# script, which installs into the inactive venv with its pip, so both venvs must exist. +"$uv" venv --managed-python --seed --python "$python_version" "$gateway_dir/$version" +if [ ! -d "$gateway_dir/$active" ]; then + "$uv" venv --managed-python --seed --python "$python_version" "$gateway_dir/$active" +fi +# Compile bytecode now, as otherwise the app compiles it on the first start, prolonging +# the downtime during the restart +"$uv" pip install --compile-bytecode --python "$gateway_dir/$version/bin/python" "$gateway_package" +sudo "$gateway_dir/$version/bin/python" -m dstack._internal.proxy.gateway.systemd install +echo "$version" > "$gateway_dir/version" +sudo systemctl daemon-reload +sudo systemctl restart dstack.gateway +echo "Update successfully completed" diff --git a/src/dstack/_internal/core/backends/kubernetes/compute.py b/src/dstack/_internal/core/backends/kubernetes/compute.py index ee4539fe3..4ea57a560 100644 --- a/src/dstack/_internal/core/backends/kubernetes/compute.py +++ b/src/dstack/_internal/core/backends/kubernetes/compute.py @@ -32,7 +32,7 @@ generate_unique_name, generate_unique_volume_name, get_docker_commands, - get_dstack_gateway_commands, + get_dstack_gateway_install_command, merge_tags, ) from dstack._internal.core.backends.base.offers import ( @@ -1423,8 +1423,6 @@ def _wait_for_load_balancer_address( def _get_gateway_commands(authorized_keys: List[str]) -> List[str]: authorized_keys_content = "\n".join(authorized_keys).strip() - gateway_commands = " && ".join(get_dstack_gateway_commands()) - quoted_gateway_commands = shlex.quote(gateway_commands) commands = [ # install packages @@ -1451,7 +1449,7 @@ def _get_gateway_commands(authorized_keys: List[str]) -> List[str]: "rm -rf /etc/ssh/ssh_host_*", "ssh-keygen -A > /dev/null", # install gateway - f"su ubuntu -c {quoted_gateway_commands}", + f"su ubuntu -c {shlex.quote(get_dstack_gateway_install_command())}", # start docker-systemctl-replacement as an init replacement (PID 1), which # - starts and supervises enabled services (sshd, nginx, dstack.gateway) # - stops running services on SIGTERM (graceful shutdown) diff --git a/src/dstack/_internal/server/services/gateways/__init__.py b/src/dstack/_internal/server/services/gateways/__init__.py index 793654ded..e641ff3a9 100644 --- a/src/dstack/_internal/server/services/gateways/__init__.py +++ b/src/dstack/_internal/server/services/gateways/__init__.py @@ -1,7 +1,6 @@ import asyncio import datetime import itertools -import shlex import uuid from collections.abc import AsyncGenerator from contextlib import asynccontextmanager @@ -16,7 +15,7 @@ import dstack._internal.utils.random_names as random_names from dstack._internal.core.backends.base.compute import ( - get_dstack_gateway_package_and_target_version, + get_dstack_gateway_install_command, ) from dstack._internal.core.backends.features import ( BACKENDS_WITH_GATEWAY_SUPPORT, @@ -776,11 +775,11 @@ async def init_gateways(session: AsyncSession): "Skipping gateway replicas update due to DSTACK_SKIP_GATEWAY_UPDATE env variable" ) else: - gateway_package, target_version = get_dstack_gateway_package_and_target_version() + install_command = get_dstack_gateway_install_command() for gateway_replica, res in await gather_map_async( gateway_replicas, - lambda r: _update_gateway_replica(r, gateway_package, target_version), + lambda r: _update_gateway_replica(r, install_command), return_exceptions=True, ): if isinstance(res, Exception): @@ -802,10 +801,12 @@ async def init_gateways(session: AsyncSession): ) +_GATEWAY_UPDATE_TIMEOUT = 600 + + async def _update_gateway_replica( gateway_replica_model: GatewayReplicaModel, - gateway_package: str, - target_version: str | None, + install_command: str, ) -> bool: if gateway_replica_model.ip_address is None: logger.warning( @@ -824,58 +825,13 @@ async def _update_gateway_replica( ) logger.debug("Updating gateway replica %s", connection.ip_address) - command = ( - "/bin/sh -c " - + shlex.quote(_GATEWAY_UPDATE_SCRIPT) - + " sh " # $0 placeholder - + shlex.quote(gateway_package) - + " " - + shlex.quote(target_version or "") - ) - stdout = await connection.tunnel.aexec(command, timeout=_GATEWAY_UPDATE_TIMEOUT) + stdout = await connection.tunnel.aexec(install_command, timeout=_GATEWAY_UPDATE_TIMEOUT) if "Update successfully completed" in stdout: logger.info("Gateway replica %s updated", connection.ip_address) return True return False -_GATEWAY_UPDATE_TIMEOUT = 600 - -# 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 -if [ -n "$build" ]; then - current_build=$("$root/$version/bin/pip" show dstack | grep Version | awk '{print $2}') - if [ "$current_build" = "$build" ]; then - echo "The build $build is already installed. Skipping..." - exit 0 - fi -fi -if [ "$version" = blue ]; then - version=green -else - version=blue -fi -# dstack-gateway is a pre-0.21.4 package that may still be installed and require a conflicting dstack version -"$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._internal.proxy.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 diff --git a/src/tests/_internal/core/backends/base/conftest.py b/src/tests/_internal/core/backends/base/conftest.py new file mode 100644 index 000000000..f4e178c20 --- /dev/null +++ b/src/tests/_internal/core/backends/base/conftest.py @@ -0,0 +1,4 @@ +from tests._internal.core.backends.base.gateway_host import ( # noqa: F401 + gateway_host, + gateway_host_container, +) diff --git a/src/tests/_internal/core/backends/base/gateway_host.py b/src/tests/_internal/core/backends/base/gateway_host.py new file mode 100644 index 000000000..7b0c30b2d --- /dev/null +++ b/src/tests/_internal/core/backends/base/gateway_host.py @@ -0,0 +1,213 @@ +import io +import re +import shlex +import tarfile +from collections.abc import Iterator +from dataclasses import dataclass +from typing import Optional, cast + +import pytest +from testcontainers.core.container import DockerContainer + +from dstack._internal.core.backends.base.compute import _get_dstack_gateway_install_script +from dstack._internal.utils.common import get_or_error + +GATEWAY_DIR = "/home/ubuntu/dstack" + +_STUB_PATH = "/usr/local/bin/gateway-host-stub" +_EVENTS_PATH = "/tmp/gateway-host-events" + +# A single stub dispatching on the name it is invoked by. Every call is logged to the events file. +# - python: a venv's interpreter. Reports the Python and dstack versions from the venv's +# `installed` file to the gateway install script's version check. +# - uv: reports the version from its `uv.version` file, creates venvs, and fails on the +# subcommand named in its `uv.fail_on` file. +# - wget: serves a uv release archive with the stub as uv, for the version in the URL. +# - uname: reports x86_64 regardless of the machine running the tests. +# - sudo: logs systemctl calls with the active venv at the time of the call, logs and runs other +# commands. +_STUB = f"""\ +#!/bin/sh +set -e +log() {{ + printf "%s\\n" "$*" >> {_EVENTS_PATH} +}} +case "${{0##*/}}" in +python) + venv_path=${{0%/bin/python}} + if [ "$1" = -c ]; then + log "${{venv_path##*/}} python -c" + exec cat "$venv_path/installed" + fi + log "${{venv_path##*/}} python $*" + ;; +uv) + log "uv $*" + if [ -f "$0.fail_on" ] && [ "$1" = "$(cat "$0.fail_on")" ]; then + exit 1 + fi + case "$1" in + --version) + echo "uv $(cat "$0.version") (0123456 2026-01-01 x86_64-unknown-linux-musl)" + ;; + venv) + for arg; do venv_path=$arg; done + mkdir -p "$venv_path/bin" + ln -s {_STUB_PATH} "$venv_path/bin/python" + ;; + esac + ;; +wget) + for arg; do url=$arg; done + log "wget $url" + version=${{url#*/download/}} + asset=${{url##*/}} + release=$(mktemp -d) + mkdir "$release/${{asset%.tar.gz}}" + ln -s {_STUB_PATH} "$release/${{asset%.tar.gz}}/uv" + echo "${{version%%/*}}" > "$release/${{asset%.tar.gz}}/uv.version" + tar -czf "$2" -C "$release" "${{asset%.tar.gz}}" + ;; +uname) + echo x86_64 + ;; +sudo) + if [ "$1" = systemctl ]; then + log "$* active=$(cat {GATEWAY_DIR}/version)" + else + log "sudo $*" + exec "$@" + fi + ;; +esac +""" + + +def get_install_script_var(name: str) -> str: + """ + Returns a variable pinned in the gateway install script, e.g. `python_version`. + """ + match = re.search(rf"^{name}=(\S+)$", _get_dstack_gateway_install_script(), re.MULTILINE) + assert match is not None + return match.group(1) + + +@dataclass +class GatewayHostResult: + returncode: int + stdout: str + stderr: str + + +class GatewayHost: + """ + A gateway host emulated in a Docker container for running the gateway install command. + + The command runs for real in a shell, which covers quoting, version checks, and the order of + venv switching and service restart. The commands that need the network, root, or systemd are + stubbed: sudo, wget, uname, uv, and the venvs' python. Everything else, including `rm -rf`, + runs as is. + + Stubs alone cannot contain the script: any command they don't cover would run on the machine + running the tests, and a script change can add one without updating the stubs. The container + bounds what such a change can touch. It also lets the script run unmodified against the real + `/home/ubuntu/dstack` on the same Ubuntu image that gateways use. + """ + + def __init__(self, container: DockerContainer): + self._container = container + + def add_venv(self, name: str, installed: Optional[str], active: bool = False): + """ + Args: + installed: `" "` reported by the venv's python. + `None` emulates a broken venv. + """ + venv_path = f"{GATEWAY_DIR}/{name}" + script = f"mkdir -p {venv_path}/bin && ln -s {_STUB_PATH} {venv_path}/bin/python" + if installed is not None: + script += f" && echo {shlex.quote(installed)} > {venv_path}/installed" + if active: + script += f" && echo {name} > {GATEWAY_DIR}/version" + self._exec(script) + + def add_uv(self, version: str, fail_on: Optional[str] = None): + """ + Args: + fail_on: The uv subcommand to fail, e.g. `"pip"`. + """ + uv_path = f"{GATEWAY_DIR}/bin/uv" + script = ( + f"mkdir -p {GATEWAY_DIR}/bin && ln -s {_STUB_PATH} {uv_path}" + f" && echo {shlex.quote(version)} > {uv_path}.version" + ) + if fail_on is not None: + script += f" && echo {shlex.quote(fail_on)} > {uv_path}.fail_on" + self._exec(script) + + def run(self, command: str) -> GatewayHostResult: + result = self._container.get_wrapped_container().exec_run( + ["sh", "-c", command], demux=True + ) + # With demux=True, the output is a (stdout, stderr) pair + stdout, stderr = cast(tuple[Optional[bytes], Optional[bytes]], result.output) + return GatewayHostResult( + returncode=get_or_error(result.exit_code), + stdout=(stdout or b"").decode(), + stderr=(stderr or b"").decode(), + ) + + def events(self) -> list[str]: + return self._exec(f"cat {_EVENTS_PATH}").splitlines() + + def active_venv(self) -> str: + return self._exec(f"cat {GATEWAY_DIR}/version").strip() + + def uv_version(self) -> str: + return self._exec(f"cat {GATEWAY_DIR}/bin/uv.version").strip() + + def exists(self, path: str) -> bool: + return self.run(f"test -e {shlex.quote(path)}").returncode == 0 + + def reset(self): + self._exec(f"rm -rf {GATEWAY_DIR} && : > {_EVENTS_PATH}") + + def _exec(self, script: str) -> str: + result = self.run(script) + assert result.returncode == 0, result.stderr + return result.stdout + + +@pytest.fixture(scope="session") +def gateway_host_container() -> Iterator[DockerContainer]: + # The same image as Kubernetes gateway pods. `init` lets the container stop without a timeout. + with ( + DockerContainer("ubuntu:22.04") + .with_command("sleep infinity") + .with_kwargs(init=True) as container + ): + container.get_wrapped_container().put_archive("/", _get_stubs_archive()) + yield container + + +@pytest.fixture +def gateway_host(gateway_host_container: DockerContainer) -> GatewayHost: + host = GatewayHost(gateway_host_container) + host.reset() + return host + + +def _get_stubs_archive() -> bytes: + buffer = io.BytesIO() + with tarfile.open(fileobj=buffer, mode="w") as archive: + stub = _STUB.encode() + info = tarfile.TarInfo(_STUB_PATH.lstrip("/")) + info.size = len(stub) + info.mode = 0o755 + archive.addfile(info, io.BytesIO(stub)) + for command in ("sudo", "uname", "wget"): + info = tarfile.TarInfo(f"usr/local/bin/{command}") + info.type = tarfile.SYMTYPE + info.linkname = _STUB_PATH + archive.addfile(info) + return buffer.getvalue() diff --git a/src/tests/_internal/core/backends/base/test_compute.py b/src/tests/_internal/core/backends/base/test_compute.py index 59d963cf4..e22834a7f 100644 --- a/src/tests/_internal/core/backends/base/test_compute.py +++ b/src/tests/_internal/core/backends/base/test_compute.py @@ -4,6 +4,7 @@ import gpuhunt import pytest +from dstack._internal import settings from dstack._internal.core.backends.base.compute import ( ComputeWithAllOffersCached, GoArchType, @@ -11,6 +12,7 @@ generate_unique_gateway_instance_name, generate_unique_instance_name, generate_unique_volume_name, + get_dstack_gateway_install_command, normalize_arch, parse_provider_ref_from_name, ) @@ -23,6 +25,7 @@ get_instance_offer_with_availability, get_volume, ) +from tests._internal.core.backends.base.gateway_host import GatewayHost, get_install_script_var class _ComputeWithCachedOffers(ComputeWithAllOffersCached): @@ -207,3 +210,125 @@ def test_32bit_not_supported(self, arch: str): def test_unknown_arch(self): with pytest.raises(ValueError, match="Unsupported architecture: MIPS"): normalize_arch("MIPS") + + +_UV_VERSION = get_install_script_var("uv_version") +_PYTHON_VERSION = get_install_script_var("python_version") +_UV_URL = ( + f"https://github.com/astral-sh/uv/releases/download/{_UV_VERSION}" + "/uv-x86_64-unknown-linux-musl.tar.gz" +) + + +@pytest.mark.docker +class TestGetDstackGatewayInstallCommand: + @pytest.fixture(autouse=True) + def gateway_package(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(settings, "DSTACK_GATEWAY_PACKAGE_URL", None) + monkeypatch.setattr(settings, "DSTACK_VERSION", "0.22.3") + + def test_installs_on_fresh_host(self, gateway_host: GatewayHost): + result = gateway_host.run(get_dstack_gateway_install_command()) + + assert result.returncode == 0, result.stderr + assert "Update successfully completed" in result.stdout + assert gateway_host.active_venv() == "green" + # The inactive venv is created too for update scripts of older servers that use its pip + assert gateway_host.events() == [ + f"wget {_UV_URL}", + f"uv python install {_PYTHON_VERSION}", + "sudo rm -rf /home/ubuntu/dstack/green", + f"uv venv --managed-python --seed --python {_PYTHON_VERSION} /home/ubuntu/dstack/green", + f"uv venv --managed-python --seed --python {_PYTHON_VERSION} /home/ubuntu/dstack/blue", + "uv pip install --compile-bytecode --python /home/ubuntu/dstack/green/bin/python dstack[gateway]==0.22.3", + "sudo /home/ubuntu/dstack/green/bin/python -m dstack._internal.proxy.gateway.systemd install", + "green python -m dstack._internal.proxy.gateway.systemd install", + "systemctl daemon-reload active=green", + "systemctl restart dstack.gateway active=green", + ] + + def test_skips_if_build_and_python_match(self, gateway_host: GatewayHost): + gateway_host.add_venv("blue", installed=f"{_PYTHON_VERSION} 0.22.3", active=True) + + result = gateway_host.run(get_dstack_gateway_install_command()) + + assert result.returncode == 0, result.stderr + assert "Skipping" in result.stdout + assert gateway_host.active_venv() == "blue" + assert gateway_host.events() == ["blue python -c"] + + @pytest.mark.parametrize( + "installed", + [ + pytest.param("3.10 0.22.3", id="outdated-python"), + pytest.param(f"{_PYTHON_VERSION} 0.22.2", id="outdated-build"), + pytest.param(None, id="broken-venv"), + ], + ) + def test_recreates_inactive_venv_and_switches_to_it( + self, gateway_host: GatewayHost, installed: Optional[str] + ): + gateway_host.add_venv("blue", installed=installed, active=True) + gateway_host.add_venv("green", installed="3.10 0.22.1") + gateway_host.add_uv(_UV_VERSION) + + result = gateway_host.run(get_dstack_gateway_install_command()) + + assert result.returncode == 0, result.stderr + assert gateway_host.active_venv() == "green" + assert not gateway_host.exists("/home/ubuntu/dstack/green/installed") + assert gateway_host.events() == [ + "blue python -c", + "uv --version", + f"uv python install {_PYTHON_VERSION}", + "sudo rm -rf /home/ubuntu/dstack/green", + f"uv venv --managed-python --seed --python {_PYTHON_VERSION} /home/ubuntu/dstack/green", + "uv pip install --compile-bytecode --python /home/ubuntu/dstack/green/bin/python dstack[gateway]==0.22.3", + "sudo /home/ubuntu/dstack/green/bin/python -m dstack._internal.proxy.gateway.systemd install", + "green python -m dstack._internal.proxy.gateway.systemd install", + "systemctl daemon-reload active=green", + "systemctl restart dstack.gateway active=green", + ] + + def test_never_skips_unversioned_package( + self, gateway_host: GatewayHost, monkeypatch: pytest.MonkeyPatch + ): + monkeypatch.setattr(settings, "DSTACK_GATEWAY_PACKAGE_URL", "https://example.com/d.whl") + gateway_host.add_venv("green", installed=f"{_PYTHON_VERSION} 0.22.3", active=True) + gateway_host.add_uv(_UV_VERSION) + + result = gateway_host.run(get_dstack_gateway_install_command()) + + assert result.returncode == 0, result.stderr + assert gateway_host.active_venv() == "blue" + events = gateway_host.events() + assert "green python -c" not in events + assert ( + "uv pip install --compile-bytecode --python /home/ubuntu/dstack/blue/bin/python" + " dstack[gateway] @ https://example.com/d.whl" + ) in events + + def test_replaces_outdated_uv(self, gateway_host: GatewayHost): + gateway_host.add_venv("blue", installed="3.10 0.22.2", active=True) + gateway_host.add_uv("0.1.0") + + result = gateway_host.run(get_dstack_gateway_install_command()) + + assert result.returncode == 0, result.stderr + assert gateway_host.events()[1:4] == [ + "uv --version", + f"wget {_UV_URL}", + f"uv python install {_PYTHON_VERSION}", + ] + assert gateway_host.uv_version() == _UV_VERSION + + def test_keeps_active_venv_if_install_fails(self, gateway_host: GatewayHost): + gateway_host.add_venv("blue", installed="3.10 0.22.2", active=True) + gateway_host.add_uv(_UV_VERSION, fail_on="pip") + + result = gateway_host.run(get_dstack_gateway_install_command()) + + assert result.returncode != 0 + assert "Update successfully completed" not in result.stdout + assert gateway_host.active_venv() == "blue" + assert not any(e.startswith("systemctl") for e in gateway_host.events()) diff --git a/src/tests/_internal/server/services/gateways/test_gateways.py b/src/tests/_internal/server/services/gateways/test_gateways.py index a121f97ed..78cecf903 100644 --- a/src/tests/_internal/server/services/gateways/test_gateways.py +++ b/src/tests/_internal/server/services/gateways/test_gateways.py @@ -1,13 +1,11 @@ -import shlex -import subprocess from datetime import datetime, timedelta, timezone -from pathlib import Path -from unittest.mock import AsyncMock, Mock +from unittest.mock import ANY, AsyncMock, Mock import pytest from sqlalchemy.ext.asyncio import AsyncSession from dstack._internal import settings as core_settings +from dstack._internal.core.backends.base.compute import get_dstack_gateway_install_command from dstack._internal.proxy.gateway.const import SERVICE_SCALING_WINDOWS from dstack._internal.proxy.gateway.schemas.stats import Stat from dstack._internal.server.services import gateways @@ -98,21 +96,32 @@ async def test_no_replicas_returns_empty(self, test_db, session: AsyncSession): @pytest.mark.asyncio class TestInitGateways: - @pytest.mark.parametrize("target_version", ["0.22.2", "0.22.3"]) - async def test_release_version_controls_existing_gateway_update( + @pytest.mark.parametrize( + ("stdout", "updated"), + [ + pytest.param("Update successfully completed\n", True, id="updated"), + pytest.param( + "The build 0.22.3 on Python 3.12 is already installed. Skipping...\n", + False, + id="skipped", + ), + ], + ) + async def test_runs_install_command_on_existing_gateway_replicas( self, session: AsyncSession, monkeypatch: pytest.MonkeyPatch, - gateway_update_sandbox: tuple[Path, Path, Mock], - target_version: str, + stdout: str, + updated: bool, ): - root, events_file, connection = gateway_update_sandbox now = datetime(2026, 1, 1, tzinfo=timezone.utc) previously_updated_at = now - timedelta(minutes=2) monkeypatch.setattr(core_settings, "DSTACK_GATEWAY_PACKAGE_URL", None) - monkeypatch.setattr(core_settings, "DSTACK_VERSION", target_version) + monkeypatch.setattr(core_settings, "DSTACK_VERSION", "0.22.3") monkeypatch.setattr(gateways.settings, "SKIP_GATEWAY_UPDATE", False) monkeypatch.setattr(gateways, "get_current_datetime", lambda: now) + connection = Mock(ip_address="1.1.1.1") + connection.tunnel.aexec = AsyncMock(return_value=stdout) monkeypatch.setattr( gateways.gateway_connections_pool, "get_or_add", @@ -133,89 +142,8 @@ async def test_release_version_controls_existing_gateway_update( await gateways.init_gateways(session) - configure.assert_awaited_once() - events = events_file.read_text().splitlines() - if target_version == "0.22.2": - assert events == ["blue pip show dstack"] - assert (root / "version").read_text().strip() == "blue" - assert replica.app_updated_at == previously_updated_at - return - - assert (root / "version").read_text().strip() == "green" - assert replica.app_updated_at == now - install_package = f"green pip install dstack[gateway]=={target_version}" - install_service = "green python -m dstack._internal.proxy.gateway.systemd install" - reload_service = "systemctl daemon-reload active=green" - restart_service = "systemctl restart dstack.gateway active=green" - assert ( - events.index(install_package) - < events.index(install_service) - < events.index(reload_service) - < events.index(restart_service) + connection.tunnel.aexec.assert_awaited_once_with( + get_dstack_gateway_install_command(), timeout=ANY ) - - -@pytest.fixture -def gateway_update_sandbox(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): - # Execute the real shell: mocking aexec would miss quoting, version gating, and the order - # of service installation / venv switching / restart. Only the remote commands are stubbed; - # no SSH, package installation, sudo, or real systemd operations are performed. - root = tmp_path / "gateway root" - root.mkdir() - (root / "version").write_text("blue\n") - events = tmp_path / "events" - events.touch() - for color in ("blue", "green"): - venv_bin = root / color / "bin" - venv_bin.mkdir(parents=True) - for command in ("pip", "python"): - _write_executable( - venv_bin / command, - f'printf "%s\\n" "{color} {command} $*" >> "$TEST_EVENTS"\n' - + ( - 'if [ "$1" = show ]; then echo "Version: 0.22.2"; fi\n' - if command == "pip" - else "" - ), - ) - command_bin = tmp_path / "commands" - command_bin.mkdir() - _write_executable( - command_bin / "sudo", - 'if [ "$1" = systemctl ]; then\n' - ' printf "%s\\n" "$* active=$(cat "$TEST_ROOT/version")" >> "$TEST_EVENTS"\n' - "else\n" - ' exec "$@"\n' - "fi\n", - ) - # The only script substitution redirects its hard-coded remote root to this temporary dir. - remote_root = "root=/home/ubuntu/dstack" - assert gateways._GATEWAY_UPDATE_SCRIPT.count(remote_root) == 1 - monkeypatch.setattr( - gateways, - "_GATEWAY_UPDATE_SCRIPT", - gateways._GATEWAY_UPDATE_SCRIPT.replace(remote_root, f"root={shlex.quote(str(root))}"), - ) - - def execute(command: str, timeout: float) -> str: - return subprocess.run( - shlex.split(command), - check=True, - capture_output=True, - text=True, - timeout=min(timeout, 5), - env={ - "PATH": f"{command_bin}:/usr/bin:/bin", - "TEST_EVENTS": str(events), - "TEST_ROOT": str(root), - }, - ).stdout - - connection = Mock(ip_address="1.1.1.1") - connection.tunnel.aexec = AsyncMock(side_effect=execute) - return root, events, connection - - -def _write_executable(path: Path, script: str): - path.write_text("#!/bin/sh\nset -e\n" + script) - path.chmod(0o755) + configure.assert_awaited_once() + assert replica.app_updated_at == (now if updated else previously_updated_at) diff --git a/src/tests/conftest.py b/src/tests/conftest.py index 6d225b133..12ab356ac 100644 --- a/src/tests/conftest.py +++ b/src/tests/conftest.py @@ -21,10 +21,17 @@ def pytest_configure(config): config.addinivalue_line( "markers", "postgres: mark test as testing Postgres to run only with --runpostgres" ) + config.addinivalue_line( + "markers", "docker: mark test as requiring Docker to run only with --rundocker" + ) config.addinivalue_line( "markers", "windows: mark test to be run on Windows in addition to POSIX" ) config.addinivalue_line("markers", "windows_only: mark test to be run on Windows only") + if config.getoption("--runall"): + config.option.runui = True + config.option.runpostgres = True + config.option.rundocker = True def pytest_addoption(parser): @@ -32,11 +39,18 @@ def pytest_addoption(parser): parser.addoption( "--runpostgres", action="store_true", default=False, help="Run tests with PostgreSQL" ) + parser.addoption( + "--rundocker", action="store_true", default=False, help="Run tests requiring Docker" + ) + parser.addoption( + "--runall", action="store_true", default=False, help="Run all of the opt-in tests above" + ) def pytest_collection_modifyitems(config, items): skip_ui = pytest.mark.skip(reason="need --runui option to run") skip_postgres = pytest.mark.skip(reason="need --runpostgres option to run") + skip_docker = pytest.mark.skip(reason="need --rundocker option to run") is_windows = os.name == "nt" skip_posix = pytest.mark.skip(reason="requires POSIX") skip_windows = pytest.mark.skip(reason="requires Windows") @@ -45,6 +59,8 @@ def pytest_collection_modifyitems(config, items): item.add_marker(skip_ui) if not config.getoption("--runpostgres") and "postgres" in item.keywords: item.add_marker(skip_postgres) + if not config.getoption("--rundocker") and "docker" in item.keywords: + item.add_marker(skip_docker) for_windows_only = "windows_only" in item.keywords for_windows = for_windows_only or "windows" in item.keywords if for_windows_only and not is_windows: