From 42aa0fdd4abaf82fd3ade972946b0b882129799e Mon Sep 17 00:00:00 2001 From: Dmitry Meyer Date: Wed, 16 Sep 2026 13:17:27 +0000 Subject: [PATCH] Do not downgrade runner and shim by default When several server replicas of different versions run at the same time, as during a rolling deployment, each of them tried to install the component version matching its own, fighting over the instance and reinstalling the binaries in a loop. The server now installs a component only if it can prove that the installed version is older than the expected one. Versions are compared as PyPA versions rather than as strings, so equivalent spellings such as `0.20.1` and `0.20.1.0` no longer trigger a reinstall. An unparseable version never leads to an install: on the installed side it means a dev build, which is assumed to be the newest one; on the expected side it's a server misconfiguration, which is logged as a warning. Downgrading is still possible via `DSTACK_RUNNER_ALLOW_DOWNGRADE` and `DSTACK_SHIM_ALLOW_DOWNGRADE`. They are meant for a server downgrade, which requires manual actions anyway, and are incompatible with rolling deployment. Also: * Drop the `DSTACK_SHIM_VERSION` fallback to `DSTACK_RUNNER_VERSION`, which was never documented and warned about its removal for a long time. * Document `DSTACK_SHIM_VERSION` and `DSTACK_{RUNNER,SHIM}_VERSION_URL`, and correct the `DSTACK_RUNNER_VERSION` description: it never defaulted to `latest`, that is only the download URL placeholder fallback when the version cannot be determined. Co-Authored-By: Claude Opus 5 (1M context) --- mkdocs/docs/reference/env.md | 13 +- .../_internal/core/backends/base/compute.py | 7 - .../pipeline_tasks/instances/check.py | 201 +++++++++--- src/dstack/_internal/settings.py | 2 + .../pipeline_tasks/test_instances/conftest.py | 6 + .../test_instances/test_check.py | 309 ++++++++++-------- 6 files changed, 352 insertions(+), 186 deletions(-) diff --git a/mkdocs/docs/reference/env.md b/mkdocs/docs/reference/env.md index 37b5c0a6fa..4a169fc22f 100644 --- a/mkdocs/docs/reference/env.md +++ b/mkdocs/docs/reference/env.md @@ -161,6 +161,8 @@ For more details on the options below, refer to the [server deployment](../guide - `DSTACK_SERVER_SSH_CONNECT_TIMEOUT`{ #DSTACK_SERVER_SSH_CONNECT_TIMEOUT } – The SSH `ConnectTimeout` for server-instance connections, in seconds. Defaults to `3`. Increase if there are high-latency links between the server and instances. - `DSTACK_SERVER_SSH_POOL_DISABLED`{ #DSTACK_SERVER_SSH_POOL_DISABLED } – Disables the reuse of server SSH connections to instances. If set, significantly decreases server RAM usage, but slows down processing and may cause CPU spikes due to frequent SSH-connection establishment. +- `DSTACK_RUNNER_ALLOW_DOWNGRADE`{ #DSTACK_RUNNER_ALLOW_DOWNGRADE } – When set to any value, allows the server to install an older `dstack-runner` version over a newer one. By default, the server skips installation if a newer version is already installed, so that replicas running different server versions don't reinstall the runner over each other during a rolling deployment. Set this variable only while downgrading the server, and unset it afterwards — leaving it set brings the reinstall loop back. +- `DSTACK_SHIM_ALLOW_DOWNGRADE`{ #DSTACK_SHIM_ALLOW_DOWNGRADE } – Same as `DSTACK_RUNNER_ALLOW_DOWNGRADE` but for `dstack-shim`. ??? info "Internal environment variables" The following environment variables are intended for development purposes: @@ -169,10 +171,15 @@ slows down processing and may cause CPU spikes due to frequent SSH-connection es * `DSTACK_SERVER_UVICORN_LOG_LEVEL` – Sets uvicorn logger log level. Defaults to `ERROR`. * `DSTACK_SERVER_MAX_OFFERS_TRIED` - Sets how many instance offers to try when starting a job. Setting a high value can degrade server performance. - * `DSTACK_RUNNER_VERSION` – Sets exact runner version for debug. Defaults to `latest`. Ignored if `DSTACK_RUNNER_DOWNLOAD_URL` is set. + * `DSTACK_RUNNER_VERSION` – Overrides the `dstack-runner` version the server installs on instances. Release builds of the server + default to their own version; dev builds have no default version. Must be a valid [PyPA version](https://packaging.python.org/en/latest/specifications/version-specifiers/), for example, `0.20.1`. + * `DSTACK_RUNNER_VERSION_URL` – URL to fetch the `dstack-runner` version from. The response body must be the version string, + see `DSTACK_RUNNER_VERSION` for the format. Used only if `DSTACK_RUNNER_VERSION` is not set. * `DSTACK_RUNNER_DOWNLOAD_URL` – Overrides `dstack-runner` binary download URL. The URL can contain `{version}` and/or `{arch}` placeholders, - where `{version}` is `dstack` version in the `X.Y.Z` format or `latest`, and `{arch}` is either `amd64` or `arm64`, for example, - `https://dstack.example.com/{arch}/{version}/dstack-runner`. + where `{version}` is the `dstack-runner` version (see `DSTACK_RUNNER_VERSION`), or `latest` if the version cannot be determined, + and `{arch}` is either `amd64` or `arm64`, for example, `https://dstack.example.com/{arch}/{version}/dstack-runner`. + * `DSTACK_SHIM_VERSION` – Same as `DSTACK_RUNNER_VERSION` but for `dstack-shim`. + * `DSTACK_SHIM_VERSION_URL` – Same as `DSTACK_RUNNER_VERSION_URL` but for `dstack-shim`. * `DSTACK_SHIM_DOWNLOAD_URL` – Overrides `dstack-shim` binary download URL. The URL can contain `{version}` and/or `{arch}` placeholders, see `DSTACK_RUNNER_DOWNLOAD_URL` for the details. * `DSTACK_GATEWAY_PACKAGE_URL` – Overrides the URL the `dstack` package is installed from on gateway instances (used as `dstack[gateway] @ `). diff --git a/src/dstack/_internal/core/backends/base/compute.py b/src/dstack/_internal/core/backends/base/compute.py index 60ee0d6d01..9bf202d3f9 100644 --- a/src/dstack/_internal/core/backends/base/compute.py +++ b/src/dstack/_internal/core/backends/base/compute.py @@ -969,13 +969,6 @@ def get_dstack_shim_version() -> Optional[str]: return version if version := settings.DSTACK_SHIM_VERSION: return version - if version := settings.DSTACK_RUNNER_VERSION: - logger.warning( - "DSTACK_SHIM_VERSION is not set, using DSTACK_RUNNER_VERSION." - " Future versions will not fall back to DSTACK_RUNNER_VERSION." - " Set DSTACK_SHIM_VERSION to supress this warning." - ) - return version if version_url := settings.DSTACK_SHIM_VERSION_URL: return _fetch_version(version_url) if settings.DSTACK_USE_LATEST_FROM_BRANCH: diff --git a/src/dstack/_internal/server/background/pipeline_tasks/instances/check.py b/src/dstack/_internal/server/background/pipeline_tasks/instances/check.py index 8185621f88..0bd31d06c8 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/instances/check.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/instances/check.py @@ -5,11 +5,13 @@ from typing import Optional import gpuhunt +import packaging.version import requests from sqlalchemy import func, select from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import joinedload +from dstack._internal import settings from dstack._internal.core.backends.base.backend import Backend from dstack._internal.core.backends.base.compute import ( get_dstack_runner_download_url, @@ -46,6 +48,7 @@ from dstack._internal.server.schemas.instances import InstanceCheck from dstack._internal.server.schemas.runner import ( ComponentInfo, + ComponentName, ComponentStatus, InstanceHealthResponse, ) @@ -482,86 +485,114 @@ def _maybe_install_components( installation_requested = False if (runner_info := components.runner) is not None: - installation_requested |= _maybe_install_runner(instance_model, shim_client, runner_info) + installation_requested |= _maybe_install_runner( + instance_model=instance_model, + shim_client=shim_client, + runner_info=runner_info, + allow_downgrade=settings.DSTACK_RUNNER_ALLOW_DOWNGRADE, + ) else: - logger.debug("Instance %s: no runner info", instance_model.name) + logger.debug( + "Instance %s: %s: no component info", instance_model.name, ComponentName.RUNNER.value + ) if (shim_info := components.shim) is not None: if shim_info.status == ComponentStatus.INSTALLED: installed_shim_version = shim_info.version - installation_requested |= _maybe_install_shim(instance_model, shim_client, shim_info) + installation_requested |= _maybe_install_shim( + instance_model=instance_model, + shim_client=shim_client, + shim_info=shim_info, + allow_downgrade=settings.DSTACK_SHIM_ALLOW_DOWNGRADE, + ) else: - logger.debug("Instance %s: no shim info", instance_model.name) + logger.debug( + "Instance %s: %s: no component info", instance_model.name, ComponentName.SHIM.value + ) - # old shim without `dstack-shim` component and `/api/shutdown` support - # or the same version is already running - # or we just requested installation of at least one component - # or at least one component is already being installed - # or at least one shim task won't survive restart running_shim_version = shim_client.get_version_string() + # skip the restart if: if ( + # old shim without `dstack-shim` component and `/api/shutdown` support installed_shim_version is None + # the same version is already running -- a string comparison on purpose: both sides come + # from the same `Version` build-time variable, one read off the on-disk binary, the other + # reported by the running process, so the question is whether it's the same build, not + # which version is newer or installed_shim_version == running_shim_version + # we just requested installation of at least one component or installation_requested + # at least one component is already being installed or any(component.status == ComponentStatus.INSTALLING for component in components) + # at least one shim task won't survive restart or not shim_client.is_safe_to_restart() ): return if shim_client.shutdown(force=False): logger.debug( - "Instance %s: restarting shim %s -> %s", + "Instance %s: %s: restarting %r -> %r", instance_model.name, + ComponentName.SHIM.value, running_shim_version, installed_shim_version, ) else: - logger.debug("Instance %s: cannot restart shim", instance_model.name) + logger.debug( + "Instance %s: %s: cannot restart", instance_model.name, ComponentName.SHIM.value + ) def _maybe_install_runner( instance_model: InstanceModel, shim_client: runner_client.ShimClient, runner_info: ComponentInfo, + allow_downgrade: bool, ) -> bool: # For developers: # * To install the latest dev build for the current branch from the CI, # set DSTACK_USE_LATEST_FROM_BRANCH=1. # * To provide your own build, set DSTACK_RUNNER_VERSION_URL and DSTACK_RUNNER_DOWNLOAD_URL. - expected_version = get_dstack_runner_version() - if expected_version is None: - return False - + name = runner_info.name installed_version = runner_info.version + expected_version = get_dstack_runner_version() logger.debug( - "Instance %s: runner status=%s installed_version=%s", + "Instance %s: %s: status=%s installed_version=%r expected_version=%r", instance_model.name, + name, runner_info.status.value, - installed_version or "(no version)", + installed_version, + expected_version, ) - if runner_info.status == ComponentStatus.INSTALLING: - logger.debug("Instance %s: runner is already being installed", instance_model.name) + + if not expected_version: return False - if installed_version and installed_version == expected_version: - logger.debug("Instance %s: expected runner version already installed", instance_model.name) + + if not _should_install_component( + instance_model=instance_model, + component_info=runner_info, + expected_version=expected_version, + allow_downgrade=allow_downgrade, + ): return False - url = get_dstack_runner_download_url( + download_url = get_dstack_runner_download_url( arch=_get_instance_cpu_arch(instance_model), version=expected_version, ) logger.debug( - "Instance %s: installing runner %s -> %s from %s", + "Instance %s: %s: installing %r -> %r from %s", instance_model.name, - installed_version or "(no version)", + name, + installed_version, expected_version, - url, + download_url, ) try: - shim_client.install_runner(url) + shim_client.install_runner(download_url) return True - except requests.RequestException as exc: - logger.warning("Instance %s: shim.install_runner(): %s", instance_model.name, exc) + except requests.RequestException as e: + logger.warning("Instance %s: %s: failed to install: %s", instance_model.name, name, e) return False @@ -569,49 +600,129 @@ def _maybe_install_shim( instance_model: InstanceModel, shim_client: runner_client.ShimClient, shim_info: ComponentInfo, + allow_downgrade: bool, ) -> bool: # For developers: # * To install the latest dev build for the current branch from the CI, # set DSTACK_USE_LATEST_FROM_BRANCH=1. # * To provide your own build, set DSTACK_SHIM_VERSION_URL and DSTACK_SHIM_DOWNLOAD_URL. - expected_version = get_dstack_shim_version() - if expected_version is None: - return False - + name = shim_info.name installed_version = shim_info.version + expected_version = get_dstack_shim_version() logger.debug( - "Instance %s: shim status=%s installed_version=%s running_version=%s", + "Instance %s: %s: status=%s installed_version=%r expected_version=%r running_version=%r", instance_model.name, + name, shim_info.status.value, - installed_version or "(no version)", + installed_version, + expected_version, shim_client.get_version_string(), ) - if shim_info.status == ComponentStatus.INSTALLING: - logger.debug("Instance %s: shim is already being installed", instance_model.name) + + if not expected_version: return False - if installed_version and installed_version == expected_version: - logger.debug("Instance %s: expected shim version already installed", instance_model.name) + + if not _should_install_component( + instance_model=instance_model, + component_info=shim_info, + expected_version=expected_version, + allow_downgrade=allow_downgrade, + ): return False - url = get_dstack_shim_download_url( + download_url = get_dstack_shim_download_url( arch=_get_instance_cpu_arch(instance_model), version=expected_version, ) logger.debug( - "Instance %s: installing shim %s -> %s from %s", + "Instance %s: %s: installing %r -> %r from %s", instance_model.name, - installed_version or "(no version)", + name, + installed_version, expected_version, - url, + download_url, ) try: - shim_client.install_shim(url) + shim_client.install_shim(download_url) return True - except requests.RequestException as exc: - logger.warning("Instance %s: shim.install_shim(): %s", instance_model.name, exc) + except requests.RequestException as e: + logger.warning("Instance %s: %s: failed to install: %s", instance_model.name, name, e) return False +def _should_install_component( + instance_model: InstanceModel, + component_info: ComponentInfo, + expected_version: str, + allow_downgrade: bool, +) -> bool: + """ + Decides whether the component should be installed, comparing the installed and the expected + versions. Logs the reason if the installation is skipped. + + Unless `allow_downgrade` is set, the component is not installed if the installed version is + newer than the expected one. This keeps server replicas running different versions from + reinstalling the component over each other during a rolling deployment. + """ + name = component_info.name + + if component_info.status == ComponentStatus.INSTALLING: + logger.debug("Instance %s: %s: already being installed", instance_model.name, name) + return False + + expected_version_parsed = _parse_pypa_version(expected_version) + if expected_version_parsed is None: + logger.warning( + "Instance %s: %s: failed to parse expected_version: %r", + instance_model.name, + name, + expected_version, + ) + return False + + installed_version = component_info.version + if not installed_version: + return True + + installed_version_parsed = _parse_pypa_version(installed_version) + if installed_version_parsed is None: + # Dev builds report `latest`; treat any unparseable version as the newest one + if not allow_downgrade: + logger.debug( + "Instance %s: %s: cannot parse installed_version, skipping the install", + instance_model.name, + name, + ) + return False + elif installed_version_parsed == expected_version_parsed: + logger.debug( + "Instance %s: %s: expected version already installed", instance_model.name, name + ) + return False + elif installed_version_parsed > expected_version_parsed and not allow_downgrade: + logger.debug("Instance %s: %s: newer version already installed", instance_model.name, name) + return False + + return True + + +def _parse_pypa_version(version_string: str) -> Optional[packaging.version.Version]: + """ + Parses the version for comparing component versions, keeping the pre-release, dev, post-release, + and local segments, that is, `0.20.1rc1` is older than `0.20.1`. + Returns `None` if the version is not PyPA-conformant, e.g., `latest` reported by dev builds. + + Not to be confused with `ShimClient.get_version_tuple()`, which is based on + `dstack._internal.server.services.runner.client._parse_version` -- it truncates the version to + `(major, minor, micro)` and treats unparseable versions as the latest, suiting feature gating + but not version comparison. + """ + try: + return packaging.version.parse(version_string) + except packaging.version.InvalidVersion: + return None + + def _get_instance_cpu_arch(instance_model: InstanceModel) -> Optional[gpuhunt.CPUArchitecture]: job_provisioning_data = get_instance_provisioning_data(instance_model) if job_provisioning_data is None: diff --git a/src/dstack/_internal/settings.py b/src/dstack/_internal/settings.py index 7a37ba5cfb..b04b94b856 100644 --- a/src/dstack/_internal/settings.py +++ b/src/dstack/_internal/settings.py @@ -13,9 +13,11 @@ DSTACK_RUNNER_VERSION = os.getenv("DSTACK_RUNNER_VERSION") DSTACK_RUNNER_VERSION_URL = os.getenv("DSTACK_RUNNER_VERSION_URL") DSTACK_RUNNER_DOWNLOAD_URL = os.getenv("DSTACK_RUNNER_DOWNLOAD_URL") +DSTACK_RUNNER_ALLOW_DOWNGRADE = os.getenv("DSTACK_RUNNER_ALLOW_DOWNGRADE") is not None DSTACK_SHIM_VERSION = os.getenv("DSTACK_SHIM_VERSION") DSTACK_SHIM_VERSION_URL = os.getenv("DSTACK_SHIM_VERSION_URL") DSTACK_SHIM_DOWNLOAD_URL = os.getenv("DSTACK_SHIM_DOWNLOAD_URL") +DSTACK_SHIM_ALLOW_DOWNGRADE = os.getenv("DSTACK_SHIM_ALLOW_DOWNGRADE") is not None DSTACK_USE_LATEST_FROM_BRANCH = os.getenv("DSTACK_USE_LATEST_FROM_BRANCH") is not None DSTACK_GATEWAY_PACKAGE_URL = os.getenv("DSTACK_GATEWAY_PACKAGE_URL") diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_instances/conftest.py b/src/tests/_internal/server/background/pipeline_tasks/test_instances/conftest.py index bbb48134b2..c80449aca2 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_instances/conftest.py +++ b/src/tests/_internal/server/background/pipeline_tasks/test_instances/conftest.py @@ -15,6 +15,12 @@ from dstack._internal.server.schemas.instances import InstanceCheck +def pytest_configure(config: pytest.Config) -> None: + config.addinivalue_line( + "markers", "allow_downgrade: allow the server to install an older component version" + ) + + @pytest.fixture def fetcher() -> InstanceFetcher: return InstanceFetcher( diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_instances/test_check.py b/src/tests/_internal/server/background/pipeline_tasks/test_instances/test_check.py index 4ff57e0fd6..e3d9e401da 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_instances/test_check.py +++ b/src/tests/_internal/server/background/pipeline_tasks/test_instances/test_check.py @@ -1,6 +1,6 @@ import datetime as dt import logging -from typing import Optional +from typing import ClassVar, Optional from unittest.mock import Mock import pytest @@ -10,6 +10,7 @@ from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession +from dstack._internal import settings from dstack._internal.core.models.common import validate_json_extra_ignore from dstack._internal.core.models.fleets import FleetNodesSpec from dstack._internal.core.models.health import HealthStatus @@ -633,31 +634,65 @@ def shim_client_mock( return mock -@pytest.mark.usefixtures("get_dstack_runner_version_mock") -class TestMaybeInstallRunner(BaseTestMaybeInstallComponents): +@pytest.mark.usefixtures("version_mock", "download_url_mock") +class BaseTestMaybeInstallComponent(BaseTestMaybeInstallComponents): + """ + Shared tests for `_maybe_install_runner` and `_maybe_install_shim`. The two must behave + identically -- only the version getter, the download URL getter, and the install call differ. + """ + + COMPONENT_NAME: ClassVar[ComponentName] + DOWNLOAD_URL: ClassVar[str] + ALLOW_DOWNGRADE_SETTING: ClassVar[str] + @pytest.fixture - def component_list(self) -> ComponentList: - components = ComponentList() - components.add( - ComponentInfo( - name=ComponentName.RUNNER, - version=self.EXPECTED_VERSION, - status=ComponentStatus.INSTALLED, - ), + def component_info(self) -> ComponentInfo: + return ComponentInfo( + name=self.COMPONENT_NAME, + version=self.EXPECTED_VERSION, + status=ComponentStatus.INSTALLED, ) + + @pytest.fixture + def component_list(self, component_info: ComponentInfo) -> ComponentList: + components = ComponentList() + components.add(component_info) return components @pytest.fixture - def get_dstack_runner_version_mock(self, monkeypatch: pytest.MonkeyPatch) -> Mock: - mock = Mock(return_value=self.EXPECTED_VERSION) - monkeypatch.setattr(instances_check, "get_dstack_runner_version", mock) - return mock + def version_mock(self, monkeypatch: pytest.MonkeyPatch) -> Mock: + raise NotImplementedError @pytest.fixture - def get_dstack_runner_download_url_mock(self, monkeypatch: pytest.MonkeyPatch) -> Mock: - mock = Mock(return_value="https://example.com/runner") - monkeypatch.setattr(instances_check, "get_dstack_runner_download_url", mock) - return mock + def download_url_mock(self, monkeypatch: pytest.MonkeyPatch) -> Mock: + raise NotImplementedError + + @pytest.fixture + def install_mock(self, shim_client_mock: Mock) -> Mock: + raise NotImplementedError + + @pytest.fixture(autouse=True) + def allow_downgrade_settings( + self, request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch + ) -> None: + # The settings are read from the environment at import time, so they must be pinned in + # both directions -- otherwise the tests fail for developers who have the variables set + # in their environment and run `pytest` directly, without env isolation. + monkeypatch.setattr(settings, "DSTACK_RUNNER_ALLOW_DOWNGRADE", False) + monkeypatch.setattr(settings, "DSTACK_SHIM_ALLOW_DOWNGRADE", False) + if request.node.get_closest_marker("allow_downgrade") is not None: + monkeypatch.setattr(settings, self.ALLOW_DOWNGRADE_SETTING, True) + + def assert_installed(self, install_mock: Mock, download_url_mock: Mock) -> None: + download_url_mock.assert_called_once_with(arch=None, version=self.EXPECTED_VERSION) + install_mock.assert_called_once_with(self.DOWNLOAD_URL) + + def assert_installing_logged(self, log: pytest.LogCaptureFixture, installed_version: str): + expected = ( + f"{self.COMPONENT_NAME.value}: installing {installed_version!r}" + f" -> {self.EXPECTED_VERSION!r} from {self.DOWNLOAD_URL}" + ) + assert expected in log.text async def test_cannot_determine_expected_version( self, @@ -665,206 +700,172 @@ async def test_cannot_determine_expected_version( instance: InstanceModel, debug_task_log: pytest.LogCaptureFixture, shim_client_mock: Mock, - get_dstack_runner_version_mock: Mock, + install_mock: Mock, + version_mock: Mock, ): - get_dstack_runner_version_mock.return_value = None + version_mock.return_value = None instances_check._maybe_install_components(instance, shim_client_mock) shim_client_mock.get_components.assert_called_once() - shim_client_mock.install_runner.assert_not_called() + install_mock.assert_not_called() - async def test_expected_version_already_installed( + async def test_cannot_parse_expected_version( self, test_db, instance: InstanceModel, debug_task_log: pytest.LogCaptureFixture, shim_client_mock: Mock, + install_mock: Mock, + version_mock: Mock, ): - shim_client_mock.get_components.return_value.runner.version = self.EXPECTED_VERSION + version_mock.return_value = "latest" instances_check._maybe_install_components(instance, shim_client_mock) - assert "expected runner version already installed" in debug_task_log.text - shim_client_mock.get_components.assert_called_once() - shim_client_mock.install_runner.assert_not_called() + assert ( + f"{self.COMPONENT_NAME.value}: failed to parse expected_version: 'latest'" + in debug_task_log.text + ) + install_mock.assert_not_called() - @pytest.mark.parametrize("status", [ComponentStatus.NOT_INSTALLED, ComponentStatus.ERROR]) - async def test_install_not_installed_or_error( + @pytest.mark.parametrize("installed_version", ["0.20.1", "0.20.1.0"]) + async def test_expected_version_already_installed( self, test_db, instance: InstanceModel, debug_task_log: pytest.LogCaptureFixture, shim_client_mock: Mock, - get_dstack_runner_download_url_mock: Mock, - status: ComponentStatus, + component_info: ComponentInfo, + install_mock: Mock, + installed_version: str, ): - shim_client_mock.get_components.return_value.runner.version = "" - shim_client_mock.get_components.return_value.runner.status = status + # `0.20.1.0` is the same version as `0.20.1` according to PyPA + component_info.version = installed_version instances_check._maybe_install_components(instance, shim_client_mock) - assert f"installing runner (no version) -> {self.EXPECTED_VERSION}" in debug_task_log.text - get_dstack_runner_download_url_mock.assert_called_once_with( - arch=None, - version=self.EXPECTED_VERSION, + assert ( + f"{self.COMPONENT_NAME.value}: expected version already installed" + in debug_task_log.text ) shim_client_mock.get_components.assert_called_once() - shim_client_mock.install_runner.assert_called_once_with( - get_dstack_runner_download_url_mock.return_value - ) + install_mock.assert_not_called() - @pytest.mark.parametrize("installed_version", ["0.19.40", "0.21.0", "dev"]) - async def test_install_installed( + @pytest.mark.parametrize("status", [ComponentStatus.NOT_INSTALLED, ComponentStatus.ERROR]) + async def test_install_not_installed_or_error( self, test_db, instance: InstanceModel, debug_task_log: pytest.LogCaptureFixture, shim_client_mock: Mock, - get_dstack_runner_download_url_mock: Mock, - installed_version: str, + component_info: ComponentInfo, + install_mock: Mock, + download_url_mock: Mock, + status: ComponentStatus, ): - shim_client_mock.get_components.return_value.runner.version = installed_version + component_info.version = "" + component_info.status = status instances_check._maybe_install_components(instance, shim_client_mock) - assert ( - f"installing runner {installed_version} -> {self.EXPECTED_VERSION}" - in debug_task_log.text - ) - get_dstack_runner_download_url_mock.assert_called_once_with( - arch=None, - version=self.EXPECTED_VERSION, - ) + self.assert_installing_logged(debug_task_log, "") shim_client_mock.get_components.assert_called_once() - shim_client_mock.install_runner.assert_called_once_with( - get_dstack_runner_download_url_mock.return_value - ) + self.assert_installed(install_mock, download_url_mock) - async def test_already_installing( + async def test_install_older_version( self, test_db, instance: InstanceModel, debug_task_log: pytest.LogCaptureFixture, shim_client_mock: Mock, + component_info: ComponentInfo, + install_mock: Mock, + download_url_mock: Mock, ): - shim_client_mock.get_components.return_value.runner.version = "dev" - shim_client_mock.get_components.return_value.runner.status = ComponentStatus.INSTALLING + component_info.version = "0.19.40" instances_check._maybe_install_components(instance, shim_client_mock) - assert "runner is already being installed" in debug_task_log.text + self.assert_installing_logged(debug_task_log, "0.19.40") shim_client_mock.get_components.assert_called_once() - shim_client_mock.install_runner.assert_not_called() - + self.assert_installed(install_mock, download_url_mock) -@pytest.mark.usefixtures("get_dstack_shim_version_mock") -class TestMaybeInstallShim(BaseTestMaybeInstallComponents): - @pytest.fixture - def component_list(self) -> ComponentList: - components = ComponentList() - components.add( - ComponentInfo( - name=ComponentName.SHIM, - version=self.EXPECTED_VERSION, - status=ComponentStatus.INSTALLED, - ), - ) - return components - - @pytest.fixture - def get_dstack_shim_version_mock(self, monkeypatch: pytest.MonkeyPatch) -> Mock: - mock = Mock(return_value=self.EXPECTED_VERSION) - monkeypatch.setattr(instances_check, "get_dstack_shim_version", mock) - return mock - - @pytest.fixture - def get_dstack_shim_download_url_mock(self, monkeypatch: pytest.MonkeyPatch) -> Mock: - mock = Mock(return_value="https://example.com/shim") - monkeypatch.setattr(instances_check, "get_dstack_shim_download_url", mock) - return mock - - async def test_cannot_determine_expected_version( + async def test_skips_newer_version( self, test_db, instance: InstanceModel, debug_task_log: pytest.LogCaptureFixture, shim_client_mock: Mock, - get_dstack_shim_version_mock: Mock, + component_info: ComponentInfo, + install_mock: Mock, ): - get_dstack_shim_version_mock.return_value = None + component_info.version = "0.21.0" instances_check._maybe_install_components(instance, shim_client_mock) + assert ( + f"{self.COMPONENT_NAME.value}: newer version already installed" in debug_task_log.text + ) shim_client_mock.get_components.assert_called_once() - shim_client_mock.install_shim.assert_not_called() + install_mock.assert_not_called() - async def test_expected_version_already_installed( + @pytest.mark.allow_downgrade + async def test_installs_newer_version_if_downgrade_allowed( self, test_db, instance: InstanceModel, debug_task_log: pytest.LogCaptureFixture, shim_client_mock: Mock, + component_info: ComponentInfo, + install_mock: Mock, + download_url_mock: Mock, ): - shim_client_mock.get_components.return_value.shim.version = self.EXPECTED_VERSION + component_info.version = "0.21.0" instances_check._maybe_install_components(instance, shim_client_mock) - assert "expected shim version already installed" in debug_task_log.text - shim_client_mock.get_components.assert_called_once() - shim_client_mock.install_shim.assert_not_called() + self.assert_installing_logged(debug_task_log, "0.21.0") + self.assert_installed(install_mock, download_url_mock) - @pytest.mark.parametrize("status", [ComponentStatus.NOT_INSTALLED, ComponentStatus.ERROR]) - async def test_install_not_installed_or_error( + async def test_skips_unparsable_installed_version( self, test_db, instance: InstanceModel, debug_task_log: pytest.LogCaptureFixture, shim_client_mock: Mock, - get_dstack_shim_download_url_mock: Mock, - status: ComponentStatus, + component_info: ComponentInfo, + install_mock: Mock, ): - shim_client_mock.get_components.return_value.shim.version = "" - shim_client_mock.get_components.return_value.shim.status = status + # dev builds report `latest`, assuming that it's the newest version + component_info.version = "latest" instances_check._maybe_install_components(instance, shim_client_mock) - assert f"installing shim (no version) -> {self.EXPECTED_VERSION}" in debug_task_log.text - get_dstack_shim_download_url_mock.assert_called_once_with( - arch=None, - version=self.EXPECTED_VERSION, + assert ( + f"{self.COMPONENT_NAME.value}: cannot parse installed_version, skipping the install" + in debug_task_log.text ) shim_client_mock.get_components.assert_called_once() - shim_client_mock.install_shim.assert_called_once_with( - get_dstack_shim_download_url_mock.return_value - ) + install_mock.assert_not_called() - @pytest.mark.parametrize("installed_version", ["0.19.40", "0.21.0", "dev"]) - async def test_install_installed( + @pytest.mark.allow_downgrade + async def test_installs_unparsable_installed_version_if_downgrade_allowed( self, test_db, instance: InstanceModel, debug_task_log: pytest.LogCaptureFixture, shim_client_mock: Mock, - get_dstack_shim_download_url_mock: Mock, - installed_version: str, + component_info: ComponentInfo, + install_mock: Mock, + download_url_mock: Mock, ): - shim_client_mock.get_components.return_value.shim.version = installed_version + component_info.version = "latest" instances_check._maybe_install_components(instance, shim_client_mock) - assert ( - f"installing shim {installed_version} -> {self.EXPECTED_VERSION}" - in debug_task_log.text - ) - get_dstack_shim_download_url_mock.assert_called_once_with( - arch=None, - version=self.EXPECTED_VERSION, - ) - shim_client_mock.get_components.assert_called_once() - shim_client_mock.install_shim.assert_called_once_with( - get_dstack_shim_download_url_mock.return_value - ) + self.assert_installing_logged(debug_task_log, "latest") + self.assert_installed(install_mock, download_url_mock) async def test_already_installing( self, @@ -872,15 +873,61 @@ async def test_already_installing( instance: InstanceModel, debug_task_log: pytest.LogCaptureFixture, shim_client_mock: Mock, + component_info: ComponentInfo, + install_mock: Mock, ): - shim_client_mock.get_components.return_value.shim.version = "dev" - shim_client_mock.get_components.return_value.shim.status = ComponentStatus.INSTALLING + component_info.version = "0.19.40" + component_info.status = ComponentStatus.INSTALLING instances_check._maybe_install_components(instance, shim_client_mock) - assert "shim is already being installed" in debug_task_log.text + assert f"{self.COMPONENT_NAME.value}: already being installed" in debug_task_log.text shim_client_mock.get_components.assert_called_once() - shim_client_mock.install_shim.assert_not_called() + install_mock.assert_not_called() + + +class TestMaybeInstallRunner(BaseTestMaybeInstallComponent): + COMPONENT_NAME = ComponentName.RUNNER + DOWNLOAD_URL = "https://example.com/runner" + ALLOW_DOWNGRADE_SETTING = "DSTACK_RUNNER_ALLOW_DOWNGRADE" + + @pytest.fixture + def version_mock(self, monkeypatch: pytest.MonkeyPatch) -> Mock: + mock = Mock(return_value=self.EXPECTED_VERSION) + monkeypatch.setattr(instances_check, "get_dstack_runner_version", mock) + return mock + + @pytest.fixture + def download_url_mock(self, monkeypatch: pytest.MonkeyPatch) -> Mock: + mock = Mock(return_value=self.DOWNLOAD_URL) + monkeypatch.setattr(instances_check, "get_dstack_runner_download_url", mock) + return mock + + @pytest.fixture + def install_mock(self, shim_client_mock: Mock) -> Mock: + return shim_client_mock.install_runner + + +class TestMaybeInstallShim(BaseTestMaybeInstallComponent): + COMPONENT_NAME = ComponentName.SHIM + DOWNLOAD_URL = "https://example.com/shim" + ALLOW_DOWNGRADE_SETTING = "DSTACK_SHIM_ALLOW_DOWNGRADE" + + @pytest.fixture + def version_mock(self, monkeypatch: pytest.MonkeyPatch) -> Mock: + mock = Mock(return_value=self.EXPECTED_VERSION) + monkeypatch.setattr(instances_check, "get_dstack_shim_version", mock) + return mock + + @pytest.fixture + def download_url_mock(self, monkeypatch: pytest.MonkeyPatch) -> Mock: + mock = Mock(return_value=self.DOWNLOAD_URL) + monkeypatch.setattr(instances_check, "get_dstack_shim_download_url", mock) + return mock + + @pytest.fixture + def install_mock(self, shim_client_mock: Mock) -> Mock: + return shim_client_mock.install_shim @pytest.mark.usefixtures("maybe_install_runner_mock", "maybe_install_shim_mock")