diff --git a/mkdocs/docs/reference/env.md b/mkdocs/docs/reference/env.md index 37b5c0a6f..4a169fc22 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 60ee0d6d0..9bf202d3f 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 8185621f8..0bd31d06c 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 7a37ba5cf..b04b94b85 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 bbb48134b..c80449aca 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 4ff57e0fd..e3d9e401d 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")