Fix CloudRift API compatibility

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
 

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:

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")
+        )

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,