From 99b3ab226959352bb81513557c970b4ed9555033 Mon Sep 17 00:00:00 2001 From: Andrey Cheptsov Date: Sat, 17 Oct 2026 10:57:01 +0200 Subject: [PATCH] Fix CloudRift API compatibility --- .../core/backends/cloudrift/api_client.py | 24 +-- .../background/pipeline_tasks/common.py | 3 + .../backends/cloudrift/test_api_client.py | 195 ++++++++++++++++++ .../test_instances/test_check.py | 68 +++++- .../pipeline_tasks/test_running_jobs.py | 48 +++++ 5 files changed, 321 insertions(+), 17 deletions(-) create mode 100644 src/tests/_internal/core/backends/cloudrift/test_api_client.py diff --git a/src/dstack/_internal/core/backends/cloudrift/api_client.py b/src/dstack/_internal/core/backends/cloudrift/api_client.py index d3bb425e9b..55601135d6 100644 --- a/src/dstack/_internal/core/backends/cloudrift/api_client.py +++ b/src/dstack/_internal/core/backends/cloudrift/api_client.py @@ -13,7 +13,7 @@ CLOUDRIFT_SERVER_ADDRESS = "https://api.cloudrift.ai" -CLOUDRIFT_API_VERSION = "2025-05-29" +CLOUDRIFT_API_VERSION = "2026-10-12" class RiftClient: @@ -138,8 +138,14 @@ def deploy_instance( def list_instances(self, instance_ids: Optional[List[str]] = None) -> List[Dict]: request_data = { "selector": { - "ByStatus": ["Initializing", "Active", "Deactivating"], + "ById": instance_ids, } + if instance_ids is not None + else { + "ByStatus": {"statuses": ["Initializing", "Active", "Deactivating"]}, + }, + # SSH instructions require both masks and include the port when mappings are empty. + "mask": {"with_connection_info": True, "with_credentials": True}, } logger.debug("Listing instances with request data: %s", request_data) response_data = self._make_request("instances/list", request_data) @@ -149,15 +155,8 @@ def list_instances(self, instance_ids: Optional[List[str]] = None) -> List[Dict] return [] def get_instance_by_id(self, instance_id: str) -> Optional[Dict]: - request_data = {"selector": {"ById": [instance_id]}} - logger.debug("Getting instance with request data: %s", request_data) - response_data = self._make_request("instances/list", request_data) - if isinstance(response_data, dict): - instances = response_data.get("instances", []) - if isinstance(instances, list) and len(instances) > 0: - return instances[0] - - return None + instances = self.list_instances(instance_ids=[instance_id]) + return instances[0] if instances else None def terminate_instance(self, instance_id: str) -> bool: request_data = {"selector": {"ById": [instance_id]}} @@ -172,10 +171,9 @@ def terminate_instance(self, instance_id: str) -> bool: instance_info = self.get_instance_by_id(instance_id) is_terminated = instance_info is None or instance_info.get("status") == "Inactive" logger.debug( - "Instance %s is already terminated: %s response: %s", + "Instance %s is already terminated: %s", instance_id, is_terminated, - instance_info, ) return is_terminated diff --git a/src/dstack/_internal/server/background/pipeline_tasks/common.py b/src/dstack/_internal/server/background/pipeline_tasks/common.py index d6f9564452..bff427502a 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/common.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/common.py @@ -10,6 +10,9 @@ def get_provisioning_timeout(backend_type: BackendType, instance_type_name: str) """ if backend_type == BackendType.LAMBDA: return timedelta(minutes=30) + if backend_type == BackendType.CLOUDRIFT: + # CloudRift may download a large VM image before booting and starting the shim. + return timedelta(minutes=30) if backend_type == BackendType.RUNPOD: return timedelta(minutes=20) if backend_type == BackendType.DAYTONA: diff --git a/src/tests/_internal/core/backends/cloudrift/test_api_client.py b/src/tests/_internal/core/backends/cloudrift/test_api_client.py new file mode 100644 index 0000000000..0452bd98d7 --- /dev/null +++ b/src/tests/_internal/core/backends/cloudrift/test_api_client.py @@ -0,0 +1,195 @@ +import pytest +import requests + +from dstack._internal.core.backends.cloudrift.api_client import ( + CLOUDRIFT_SERVER_ADDRESS, + RiftClient, +) +from dstack._internal.core.errors import BackendError, BackendInvalidCredentialsError + +API_ROOT = f"{CLOUDRIFT_SERVER_ADDRESS}/api/v1" + + +def _response(data, version="2025-01-29"): + return {"version": version, "data": data} + + +def _recipes(): + return { + "groups": [ + { + "name": "Linux", + "tags": ["VM"], + "recipes": [ + { + "name": "Ubuntu 24.04 Server", + "tags": ["nvidia-driver"], + "details": {"VirtualMachine": {"image_url": "https://images/nvidia.img"}}, + }, + { + "name": "Ubuntu 24.04 Server (AMD)", + "tags": ["amd-driver"], + "details": {"VirtualMachine": {"image_url": "https://images/amd.img"}}, + }, + ], + } + ] + } + + +class TestRiftClient: + @pytest.mark.parametrize( + ("endpoint", "data"), + [ + ("auth/me", {"email": "test@example.com"}), + ("instance-types/list", {"instance_types": []}), + ("recipes/list", {"groups": []}), + ("instances/list", {"instances": []}), + ("instances/rent", {"instance_ids": []}), + ("instances/terminate", {"terminated": []}), + ], + ) + def test_uses_current_production_protocol(self, requests_mock, endpoint, data): + # CloudRift rejects the old pin on provisioning endpoints, even though + # auth and instance-type discovery still accept it. + requests_mock.post( + f"{API_ROOT}/{endpoint}", + additional_matcher=lambda request: request.json()["version"] == "2026-10-12", + json=_response(data), + ) + + assert RiftClient("test-key")._make_request(endpoint, {}) == data + assert requests_mock.last_request.headers["X-API-Key"] == "test-key" + + @pytest.mark.parametrize("status_code", [401, 403]) + def test_rejects_invalid_credentials(self, requests_mock, status_code): + requests_mock.post(f"{API_ROOT}/auth/me", status_code=status_code, text="invalid key") + + assert RiftClient("test-key").validate_api_key() is False + with pytest.raises(BackendInvalidCredentialsError, match="invalid key"): + RiftClient("test-key")._make_request("auth/me") + + def test_preserves_http_errors(self, requests_mock): + requests_mock.post(f"{API_ROOT}/recipes/list", status_code=400, text="unsupported version") + + with pytest.raises(requests.HTTPError) as exc_info: + RiftClient("test-key").list_recipes() + + assert exc_info.value.response.text == "unsupported version" + + +class TestRiftClientDeployInstance: + @pytest.mark.parametrize("gpu_vendor", ["nvidia", "amd"]) + def test_discovers_image_and_rents_vm(self, requests_mock, gpu_vendor): + requests_mock.post(f"{API_ROOT}/recipes/list", json=_response(_recipes())) + requests_mock.post( + f"{API_ROOT}/instances/rent", + json=_response({"instance_ids": ["instance-id"]}, version="2026-09-09"), + ) + + instance_ids = RiftClient("test-key").deploy_instance( + instance_type="rtx59-7-24-800-ig.1", + region="ap_northeast_kr_se_1", + ssh_keys=["ssh-rsa test"], + cmd="start-shim", + gpu_vendor=gpu_vendor, + ) + + assert instance_ids == ["instance-id"] + assert requests_mock.last_request.json() == { + "version": "2026-10-12", + "data": { + "config": { + "VirtualMachine": { + "cloudinit_commands": "start-shim", + "image_url": f"https://images/{gpu_vendor}.img", + "ssh_key": {"PublicKeys": ["ssh-rsa test"]}, + } + }, + "selector": { + "ByInstanceTypeAndLocation": { + "datacenters": ["ap_northeast_kr_se_1"], + "instance_type": "rtx59-7-24-800-ig.1", + } + }, + "with_public_ip": True, + }, + } + + def test_does_not_rent_without_suitable_image(self, requests_mock): + requests_mock.post(f"{API_ROOT}/recipes/list", json=_response({"groups": []})) + + with pytest.raises(BackendError, match="No suitable VM image"): + RiftClient("test-key").deploy_instance("type.1", "region", [], "start-shim") + + assert requests_mock.call_count == 1 + + +class TestRiftClientListInstances: + def test_selects_active_statuses_and_requests_connection_info(self, requests_mock): + instances = [{"id": "instance-id", "host_address": "192.0.2.1"}] + requests_mock.post( + f"{API_ROOT}/instances/list", + json=_response({"instances": instances}, version="2026-08-05"), + ) + + assert RiftClient("test-key").list_instances() == instances + assert requests_mock.last_request.json()["data"] == { + "selector": {"ByStatus": {"statuses": ["Initializing", "Active", "Deactivating"]}}, + "mask": {"with_connection_info": True, "with_credentials": True}, + } + + @pytest.mark.parametrize("instance_ids", [[], ["instance-id", "other-id"]]) + def test_selects_requested_ids(self, requests_mock, instance_ids): + requests_mock.post(f"{API_ROOT}/instances/list", json=_response({"instances": []})) + + assert RiftClient("test-key").list_instances(instance_ids) == [] + assert requests_mock.last_request.json()["data"]["selector"] == {"ById": instance_ids} + + @pytest.mark.parametrize("found", [True, False]) + def test_gets_instance_by_id_with_connection_info(self, requests_mock, found): + instance = { + "id": "instance-id", + "node_mode": "VirtualMachine", + "host_address": "192.0.2.1", + "port_mappings": [[22, 57001]], + "virtual_machines": [{"ready": True}], + } + requests_mock.post( + f"{API_ROOT}/instances/list", + json=_response({"instances": [instance] if found else []}, version="2026-08-05"), + ) + + assert RiftClient("test-key").get_instance_by_id("instance-id") == ( + instance if found else None + ) + assert requests_mock.last_request.json()["data"] == { + "selector": {"ById": ["instance-id"]}, + "mask": {"with_connection_info": True, "with_credentials": True}, + } + + +class TestRiftClientTerminateInstance: + def test_terminates_requested_instance(self, requests_mock): + requests_mock.post( + f"{API_ROOT}/instances/terminate", + json=_response({"terminated": [{"id": "instance-id"}]}, version="2026-08-05"), + ) + + assert RiftClient("test-key").terminate_instance("instance-id") is True + assert requests_mock.last_request.json() == { + "version": "2026-10-12", + "data": {"selector": {"ById": ["instance-id"]}}, + } + + @pytest.mark.parametrize("status", [None, "Inactive", "Deactivating", "Active"]) + def test_checks_already_terminated_instance(self, requests_mock, status): + requests_mock.post(f"{API_ROOT}/instances/terminate", json=_response({"terminated": []})) + requests_mock.post( + f"{API_ROOT}/instances/list", + json=_response({"instances": [{"status": status}] if status else []}), + ) + + assert RiftClient("test-key").terminate_instance("instance-id") is ( + status in (None, "Inactive") + ) 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 e3d9e401da..e4d08473dd 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,7 +1,7 @@ import datetime as dt import logging from typing import ClassVar, Optional -from unittest.mock import Mock +from unittest.mock import AsyncMock, Mock import pytest import pytest_asyncio @@ -11,6 +11,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from dstack._internal import settings +from dstack._internal.core.models.backends.base import BackendType 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 @@ -92,20 +93,32 @@ async def test_check_shim_transitions_provisioning_on_ready( assert instance.status == InstanceStatus.IDLE assert instance.termination_deadline is None - async def test_check_shim_transitions_provisioning_on_terminating( + @pytest.mark.parametrize( + ("backend", "age_minutes", "expected_status"), + [ + (BackendType.VERDA, 20, InstanceStatus.TERMINATING), + (BackendType.CLOUDRIFT, 20, InstanceStatus.PROVISIONING), + (BackendType.CLOUDRIFT, 31, InstanceStatus.TERMINATING), + ], + ) + async def test_check_shim_respects_backend_provisioning_timeout( self, test_db, session: AsyncSession, worker: InstanceWorker, monkeypatch: pytest.MonkeyPatch, + backend: BackendType, + age_minutes: int, + expected_status: InstanceStatus, ): project = await create_project(session=session) instance = await create_instance( session=session, project=project, status=InstanceStatus.PROVISIONING, + backend=backend, ) - instance.started_at = get_current_datetime() + dt.timedelta(minutes=-20) + instance.started_at = get_current_datetime() - dt.timedelta(minutes=age_minutes) await session.commit() monkeypatch.setattr( @@ -117,8 +130,55 @@ async def test_check_shim_transitions_provisioning_on_terminating( await session.refresh(instance) - assert instance.status == InstanceStatus.TERMINATING + assert instance.status == expected_status assert instance.termination_deadline is not None + if expected_status == InstanceStatus.PROVISIONING: + assert instance.termination_deadline > get_current_datetime() + + @pytest.mark.parametrize( + ("age_minutes", "expected_status"), + [(20, InstanceStatus.PROVISIONING), (31, InstanceStatus.TERMINATING)], + ) + async def test_cloudrift_waits_for_backend_within_timeout( + self, + test_db, + session: AsyncSession, + worker: InstanceWorker, + monkeypatch: pytest.MonkeyPatch, + age_minutes: int, + expected_status: InstanceStatus, + ): + project = await create_project(session=session) + provisioning_data = get_job_provisioning_data( + dockerized=True, backend=BackendType.CLOUDRIFT + ) + provisioning_data.hostname = None + provisioning_data.ssh_port = None + instance = await create_instance( + session=session, + project=project, + status=InstanceStatus.PROVISIONING, + backend=BackendType.CLOUDRIFT, + job_provisioning_data=provisioning_data, + ) + instance.started_at = get_current_datetime() - dt.timedelta(minutes=age_minutes) + await session.commit() + backend = Mock() + monkeypatch.setattr( + instances_check, + "_get_backend_for_provisioning_wait", + AsyncMock(return_value=backend), + ) + + await process_instance(session, worker, instance) + await session.refresh(instance) + + assert instance.status == expected_status + if expected_status == InstanceStatus.TERMINATING: + assert instance.termination_reason == InstanceTerminationReason.PROVISIONING_TIMEOUT + backend.compute.assert_not_called() + else: + backend.compute().update_provisioning_data.assert_called_once() async def test_check_shim_transitions_provisioning_on_busy( self, diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py b/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py index 9493cc1104..08940cc544 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py +++ b/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py @@ -610,6 +610,54 @@ async def test_runs_provisioning_job( assert job_runtime_data.working_dir == "/dstack/run" assert job_runtime_data.username == "dstack" + @pytest.mark.parametrize( + ("age_minutes", "expected_status"), + [(20, JobStatus.PROVISIONING), (31, JobStatus.TERMINATING)], + ) + async def test_cloudrift_provisioning_waits_for_shim_within_timeout( + self, + test_db, + session: AsyncSession, + worker: JobRunningWorker, + ssh_tunnel_mock: Mock, + shim_client_mock: Mock, + age_minutes: int, + expected_status: JobStatus, + ): + project = await create_project(session=session) + user = await create_user(session=session) + repo = await create_repo(session=session, project_id=project.id) + run = await create_run(session=session, project=project, repo=repo, user=user) + instance = await create_instance( + session=session, + project=project, + status=InstanceStatus.BUSY, + backend=BackendType.CLOUDRIFT, + ) + job = await create_job( + session=session, + run=run, + status=JobStatus.PROVISIONING, + submitted_at=get_current_datetime() - timedelta(minutes=age_minutes), + job_provisioning_data=get_job_provisioning_data( + dockerized=True, backend=BackendType.CLOUDRIFT + ), + instance=instance, + instance_assigned=True, + ) + shim_client_mock.healthcheck.return_value = None + + await _process_job(session, worker, job) + await session.refresh(job) + + assert job.status == expected_status + assert job.termination_reason == ( + JobTerminationReason.WAITING_RUNNER_LIMIT_EXCEEDED + if expected_status == JobStatus.TERMINATING + else None + ) + shim_client_mock.submit_task.assert_not_called() + @pytest.mark.parametrize("sshproxy_enforced", [False, True]) async def test_provisioning_shim( self,