dstack 0.18.41__py3-none-any.whl → 0.18.43__py3-none-any.whl

tests/_internal/server/background/tasks/test_process_terminating_jobs.py

@@ -13,7 +13,7 @@ from dstack._internal.core.models.volumes import VolumeStatus
 from dstack._internal.server.background.tasks.process_terminating_jobs import (
     process_terminating_jobs,
 )
-from dstack._internal.server.models import InstanceModel, JobModel
+from dstack._internal.server.models import InstanceModel, JobModel, VolumeAttachmentModel
 from dstack._internal.server.services.volumes import volume_model_to_volume
 from dstack._internal.server.testing.common import (
     create_instance,
@@ -176,7 +176,7 @@ class TestProcessTerminatingJobs:
         # so that stuck volumes don't prevent the instance from terminating.
         assert job.instance is None
         assert job.volumes_detached_at is not None
-        assert len(instance.volumes) == 1
+        assert len(instance.volume_attachments) == 1
 
         # Force detach called
         with (
@@ -214,11 +214,11 @@ class TestProcessTerminatingJobs:
         res = await session.execute(select(JobModel).options(joinedload(JobModel.instance)))
         job = res.unique().scalar_one()
         res = await session.execute(
-            select(InstanceModel).options(joinedload(InstanceModel.volumes))
+            select(InstanceModel).options(joinedload(InstanceModel.volume_attachments))
         )
         instance = res.unique().scalar_one()
         assert job.status == JobStatus.TERMINATED
-        assert len(instance.volumes) == 0
+        assert len(instance.volume_attachments) == 0
 
     async def test_terminates_job_on_shared_instance(self, session: AsyncSession):
         project = await create_project(session)
@@ -330,7 +330,12 @@ class TestProcessTerminatingJobs:
         await session.refresh(instance)
         assert job.status == JobStatus.TERMINATED
         res = await session.execute(
-            select(InstanceModel).options(joinedload(InstanceModel.volumes))
+            select(InstanceModel).options(
+                joinedload(InstanceModel.volume_attachments).joinedload(
+                    VolumeAttachmentModel.volume
+                )
+            )
         )
         instance = res.unique().scalar_one()
-        assert instance.volumes == [volume_2]
+        assert len(instance.volume_attachments) == 1
+        assert instance.volume_attachments[0].volume == volume_2

tests/_internal/server/conftest.py

@@ -8,6 +8,7 @@ from dstack._internal.server.main import app
 from dstack._internal.server.services import encryption as encryption  # import for side-effect
 from dstack._internal.server.services import logs as logs_services
 from dstack._internal.server.services.docker import ImageConfig, ImageConfigObject
+from dstack._internal.server.services.logs.filelog import FileLogStorage
 from dstack._internal.server.testing.conf import postgres_container, session, test_db  # noqa: F401
 
 
@@ -18,13 +19,11 @@ def client(event_loop):
 
 
 @pytest.fixture
-def test_log_storage(
-    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
-) -> logs_services.FileLogStorage:
+def test_log_storage(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> FileLogStorage:
     root = tmp_path / "test_logs"
     root.mkdir()
-    storage = logs_services.FileLogStorage(root)
-    monkeypatch.setattr(logs_services, "_default_log_storage", storage)
+    storage = FileLogStorage(root)
+    monkeypatch.setattr(logs_services, "_log_storage", storage)
     return storage
 
 

tests/_internal/server/routers/test_backends.py

@@ -1335,6 +1335,7 @@ class TestGetConfigInfo:
             "vpc_ids": None,
             "default_vpcs": None,
             "public_ips": None,
+            "iam_instance_profile": None,
             "tags": None,
             "os_images": None,
             "creds": json.loads(backend.auth.plaintext),

tests/_internal/server/routers/test_logs.py

@@ -3,7 +3,7 @@ from httpx import AsyncClient
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from dstack._internal.core.models.users import GlobalRole, ProjectRole
-from dstack._internal.server.services.logs import FileLogStorage
+from dstack._internal.server.services.logs.filelog import FileLogStorage
 from dstack._internal.server.services.projects import add_project_member
 from dstack._internal.server.testing.common import create_project, create_user, get_auth_headers
 

tests/_internal/server/routers/test_runs.py

@@ -180,7 +180,7 @@ def get_dev_env_run_plan_dict(
                     ],
                     "env": {},
                     "home_dir": "/root",
-                    "image_name": "dstackai/base:py3.13-0.6-cuda-12.1",
+                    "image_name": "dstackai/base:py3.13-0.7-cuda-12.1",
                     "user": None,
                     "privileged": privileged,
                     "job_name": f"{run_name}-0-0",
@@ -337,7 +337,7 @@ def get_dev_env_run_dict(
                     ],
                     "env": {},
                     "home_dir": "/root",
-                    "image_name": "dstackai/base:py3.13-0.6-cuda-12.1",
+                    "image_name": "dstackai/base:py3.13-0.7-cuda-12.1",
                     "user": None,
                     "privileged": privileged,
                     "job_name": f"{run_name}-0-0",

tests/_internal/server/routers/test_volumes.py

@@ -11,7 +11,7 @@ from sqlalchemy.ext.asyncio import AsyncSession
 
 from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.users import GlobalRole, ProjectRole
-from dstack._internal.server.models import VolumeModel
+from dstack._internal.server.models import VolumeAttachmentModel, VolumeModel
 from dstack._internal.server.services.projects import add_project_member
 from dstack._internal.server.testing.common import (
     create_instance,
@@ -76,6 +76,7 @@ class TestListVolumes:
                 "deleted": False,
                 "volume_id": None,
                 "provisioning_data": None,
+                "attachments": [],
                 "attachment_data": None,
             },
             {
@@ -91,6 +92,7 @@ class TestListVolumes:
                 "deleted": False,
                 "volume_id": None,
                 "provisioning_data": None,
+                "attachments": [],
                 "attachment_data": None,
             },
         ]
@@ -117,6 +119,7 @@ class TestListVolumes:
                 "deleted": False,
                 "volume_id": None,
                 "provisioning_data": None,
+                "attachments": [],
                 "attachment_data": None,
             },
         ]
@@ -170,6 +173,7 @@ class TestListVolumes:
                 "deleted": False,
                 "volume_id": None,
                 "provisioning_data": None,
+                "attachments": [],
                 "attachment_data": None,
             },
         ]
@@ -217,6 +221,7 @@ class TestListProjectVolumes:
                 "deleted": False,
                 "volume_id": None,
                 "provisioning_data": None,
+                "attachments": [],
                 "attachment_data": None,
             }
         ]
@@ -264,6 +269,7 @@ class TestGetVolume:
             "deleted": False,
             "volume_id": None,
             "provisioning_data": None,
+            "attachments": [],
             "attachment_data": None,
         }
 
@@ -325,6 +331,7 @@ class TestCreateVolume:
             "deleted": False,
             "volume_id": None,
             "provisioning_data": None,
+            "attachments": [],
             "attachment_data": None,
         }
         res = await session.execute(select(VolumeModel))
@@ -391,7 +398,7 @@ class TestDeleteVolumes:
             project=project,
             pool=pool,
         )
-        volume.instances.append(instance)
+        volume.attachments.append(VolumeAttachmentModel(instance=instance))
         await session.commit()
         response = await client.post(
             f"/api/project/{project.name}/volumes/delete",

tests/_internal/server/services/runner/test_client.py

@@ -9,7 +9,13 @@ from dstack._internal.core.consts import DSTACK_SHIM_HTTP_PORT
 from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.common import NetworkMode
 from dstack._internal.core.models.resources import Memory
-from dstack._internal.core.models.volumes import InstanceMountPoint, VolumeMountPoint
+from dstack._internal.core.models.volumes import (
+    InstanceMountPoint,
+    VolumeAttachment,
+    VolumeAttachmentData,
+    VolumeInstance,
+    VolumeMountPoint,
+)
 from dstack._internal.server.schemas.runner import (
     HealthcheckResponse,
     JobResult,
@@ -133,7 +139,12 @@ class TestShimClientV1(BaseShimClientTest):
             volume_id="vol-id",
             configuration=get_volume_configuration(backend=BackendType.GCP),
             external=False,
-            device_name="/dev/sdv",
+            attachments=[
+                VolumeAttachment(
+                    instance=VolumeInstance(name="instance", instance_num=0, instance_id="i-1"),
+                    attachment_data=VolumeAttachmentData(device_name="/dev/sdv"),
+                )
+            ],
         )
 
         submitted = client.submit(
@@ -150,6 +161,7 @@ class TestShimClientV1(BaseShimClientTest):
             mounts=[VolumeMountPoint(name="vol", path="/vol")],
             volumes=[volume],
             instance_mounts=[InstanceMountPoint(instance_path="/mnt/nfs/home", path="/home")],
+            instance_id="i-1",
         )
 
         assert submitted is True
@@ -198,6 +210,7 @@ class TestShimClientV1(BaseShimClientTest):
             mounts=[],
             volumes=[],
             instance_mounts=[],
+            instance_id="",
         )
 
         assert submitted is False
@@ -294,7 +307,12 @@ class TestShimClientV2(BaseShimClientTest):
             volume_id="vol-id",
             configuration=get_volume_configuration(backend=BackendType.GCP),
             external=False,
-            device_name="/dev/sdv",
+            attachments=[
+                VolumeAttachment(
+                    instance=VolumeInstance(name="instance", instance_num=0, instance_id="i-1"),
+                    attachment_data=VolumeAttachmentData(device_name="/dev/sdv"),
+                )
+            ],
         )
 
         client.submit_task(
@@ -316,6 +334,7 @@ class TestShimClientV2(BaseShimClientTest):
             host_ssh_user="dstack",
             host_ssh_keys=["host_key"],
             container_ssh_keys=["project_key", "user_key"],
+            instance_id="i-1",
         )
 
         assert adapter.call_count == 2

tests/_internal/server/services/test_logs.py

@@ -17,11 +17,11 @@ from dstack._internal.core.models.logs import LogEvent, LogEventSource
 from dstack._internal.server.models import ProjectModel
 from dstack._internal.server.schemas.logs import PollLogsRequest
 from dstack._internal.server.schemas.runner import LogEvent as RunnerLogEvent
-from dstack._internal.server.services.logs import (
+from dstack._internal.server.services.logs.aws import (
     CloudWatchLogStorage,
-    FileLogStorage,
-    LogStorageError,
 )
+from dstack._internal.server.services.logs.base import LogStorageError
+from dstack._internal.server.services.logs.filelog import FileLogStorage
 from dstack._internal.server.testing.common import create_project
 
 

tests/_internal/server/services/test_offers.py

@@ -0,0 +1,167 @@
+from unittest.mock import Mock, patch
+
+import pytest
+
+from dstack._internal.core.models.backends.base import BackendType
+from dstack._internal.core.models.profiles import Profile
+from dstack._internal.core.models.resources import ResourcesSpec
+from dstack._internal.core.models.runs import Requirements
+from dstack._internal.server.services.offers import get_offers_by_requirements
+from dstack._internal.server.testing.common import (
+    get_instance_offer_with_availability,
+    get_volume,
+    get_volume_configuration,
+)
+
+
+class TestGetOffersByRequirements:
+    @pytest.mark.asyncio
+    async def test_returns_all_offers(self):
+        profile = Profile(name="test")
+        requirements = Requirements(resources=ResourcesSpec())
+        with patch("dstack._internal.server.services.backends.get_project_backends") as m:
+            aws_backend_mock = Mock()
+            aws_backend_mock.TYPE = BackendType.AWS
+            aws_offer = get_instance_offer_with_availability(backend=BackendType.AWS)
+            aws_backend_mock.compute.return_value.get_offers_cached.return_value = [aws_offer]
+            runpod_backend_mock = Mock()
+            runpod_backend_mock.TYPE = BackendType.RUNPOD
+            runpod_offer = get_instance_offer_with_availability(backend=BackendType.RUNPOD)
+            runpod_backend_mock.compute.return_value.get_offers_cached.return_value = [
+                runpod_offer
+            ]
+            m.return_value = [aws_backend_mock, runpod_backend_mock]
+            res = await get_offers_by_requirements(
+                project=Mock(),
+                profile=profile,
+                requirements=requirements,
+            )
+            m.assert_awaited_once()
+            assert res == [(aws_backend_mock, aws_offer), (runpod_backend_mock, runpod_offer)]
+
+    @pytest.mark.asyncio
+    async def test_returns_multinode_offers(self):
+        profile = Profile(name="test")
+        requirements = Requirements(resources=ResourcesSpec())
+        with patch("dstack._internal.server.services.backends.get_project_backends") as m:
+            aws_backend_mock = Mock()
+            aws_backend_mock.TYPE = BackendType.AWS
+            aws_offer = get_instance_offer_with_availability(backend=BackendType.AWS)
+            aws_backend_mock.compute.return_value.get_offers_cached.return_value = [aws_offer]
+            runpod_backend_mock = Mock()
+            runpod_backend_mock.TYPE = BackendType.RUNPOD
+            runpod_offer = get_instance_offer_with_availability(backend=BackendType.RUNPOD)
+            runpod_backend_mock.compute.return_value.get_offers_cached.return_value = [
+                runpod_offer
+            ]
+            m.return_value = [aws_backend_mock, runpod_backend_mock]
+            res = await get_offers_by_requirements(
+                project=Mock(),
+                profile=profile,
+                requirements=requirements,
+                multinode=True,
+            )
+            m.assert_awaited_once()
+            assert res == [(aws_backend_mock, aws_offer)]
+
+    @pytest.mark.asyncio
+    async def test_returns_volume_offers(self):
+        profile = Profile(name="test")
+        requirements = Requirements(resources=ResourcesSpec())
+        with patch("dstack._internal.server.services.backends.get_project_backends") as m:
+            aws_backend_mock = Mock()
+            aws_backend_mock.TYPE = BackendType.AWS
+            aws_offer = get_instance_offer_with_availability(backend=BackendType.AWS)
+            aws_backend_mock.compute.return_value.get_offers_cached.return_value = [aws_offer]
+            runpod_backend_mock = Mock()
+            runpod_backend_mock.TYPE = BackendType.RUNPOD
+            runpod_offer1 = get_instance_offer_with_availability(
+                backend=BackendType.RUNPOD, region="eu"
+            )
+            runpod_offer2 = get_instance_offer_with_availability(
+                backend=BackendType.RUNPOD, region="us"
+            )
+            runpod_backend_mock.compute.return_value.get_offers_cached.return_value = [
+                runpod_offer1,
+                runpod_offer2,
+            ]
+            m.return_value = [aws_backend_mock, runpod_backend_mock]
+            res = await get_offers_by_requirements(
+                project=Mock(),
+                profile=profile,
+                requirements=requirements,
+                volumes=[
+                    [
+                        get_volume(
+                            configuration=get_volume_configuration(
+                                backend=BackendType.RUNPOD, region="us"
+                            )
+                        )
+                    ]
+                ],
+            )
+            m.assert_awaited_once()
+            assert res == [(runpod_backend_mock, runpod_offer2)]
+
+    @pytest.mark.asyncio
+    async def test_returns_az_offers(self):
+        profile = Profile(name="test", availability_zones=["az1", "az3"])
+        requirements = Requirements(resources=ResourcesSpec())
+        with patch("dstack._internal.server.services.backends.get_project_backends") as m:
+            aws_backend_mock = Mock()
+            aws_backend_mock.TYPE = BackendType.AWS
+            aws_offer1 = get_instance_offer_with_availability(
+                backend=BackendType.AWS, availability_zones=["az1"]
+            )
+            aws_offer2 = get_instance_offer_with_availability(
+                backend=BackendType.AWS, availability_zones=["az2"]
+            )
+            aws_offer3 = get_instance_offer_with_availability(
+                backend=BackendType.AWS, availability_zones=["az2", "az3"]
+            )
+            expected_aws_offer3 = aws_offer3.copy()
+            expected_aws_offer3.availability_zones = ["az3"]
+            aws_offer4 = get_instance_offer_with_availability(
+                backend=BackendType.AWS, availability_zones=None
+            )
+            aws_backend_mock.compute.return_value.get_offers_cached.return_value = [
+                aws_offer1,
+                aws_offer2,
+                aws_offer3,
+                aws_offer4,
+            ]
+            m.return_value = [aws_backend_mock]
+            res = await get_offers_by_requirements(
+                project=Mock(),
+                profile=profile,
+                requirements=requirements,
+            )
+            m.assert_awaited_once()
+            assert res == [(aws_backend_mock, aws_offer1), (aws_backend_mock, expected_aws_offer3)]
+
+    @pytest.mark.asyncio
+    async def test_returns_no_offers_for_multinode_instance_mounts_and_non_multinode_backend(self):
+        # Regression test for https://github.com/dstackai/dstack/issues/2211
+        profile = Profile(name="test", backends=[BackendType.RUNPOD])
+        requirements = Requirements(resources=ResourcesSpec())
+        with patch("dstack._internal.server.services.backends.get_project_backends") as m:
+            aws_backend_mock = Mock()
+            aws_backend_mock.TYPE = BackendType.AWS
+            aws_offer = get_instance_offer_with_availability(backend=BackendType.AWS)
+            aws_backend_mock.compute.return_value.get_offers_cached.return_value = [aws_offer]
+            runpod_backend_mock = Mock()
+            runpod_backend_mock.TYPE = BackendType.RUNPOD
+            runpod_offer = get_instance_offer_with_availability(backend=BackendType.RUNPOD)
+            runpod_backend_mock.compute.return_value.get_offers_cached.return_value = [
+                runpod_offer
+            ]
+            m.return_value = [aws_backend_mock, runpod_backend_mock]
+            res = await get_offers_by_requirements(
+                project=Mock(),
+                profile=profile,
+                requirements=requirements,
+                multinode=True,
+                instance_mounts=True,
+            )
+            m.assert_awaited_once()
+            assert res == []

tests/_internal/server/services/test_pools.py

@@ -7,11 +7,115 @@ import dstack._internal.server.services.pools as services_pools
 from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.instances import InstanceStatus, InstanceType, Resources
 from dstack._internal.core.models.pools import Instance
+from dstack._internal.core.models.profiles import Profile
 from dstack._internal.server.models import InstanceModel
-from dstack._internal.server.testing.common import create_project, create_user
+from dstack._internal.server.testing.common import (
+    create_instance,
+    create_pool,
+    create_project,
+    create_user,
+    get_volume,
+    get_volume_configuration,
+)
 from dstack._internal.utils.common import get_current_datetime
 
 
+class TestFilterPoolInstances:
+    # TODO: Refactor filter_pool_instances to not depend on InstanceModel and simplify tests
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
+    async def test_returns_all_instances(self, test_db, session: AsyncSession):
+        user = await create_user(session=session)
+        project = await create_project(session=session, owner=user)
+        pool = await create_pool(session=session, project=project)
+        aws_instance = await create_instance(
+            session=session,
+            project=project,
+            pool=pool,
+            backend=BackendType.AWS,
+        )
+        runpod_instance = await create_instance(
+            session=session,
+            project=project,
+            pool=pool,
+            backend=BackendType.RUNPOD,
+        )
+        instances = [aws_instance, runpod_instance]
+        res = services_pools.filter_pool_instances(
+            pool_instances=instances,
+            profile=Profile(name="test"),
+        )
+        assert res == instances
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
+    async def test_returns_multinode_instances(self, test_db, session: AsyncSession):
+        user = await create_user(session=session)
+        project = await create_project(session=session, owner=user)
+        pool = await create_pool(session=session, project=project)
+        aws_instance = await create_instance(
+            session=session,
+            project=project,
+            pool=pool,
+            backend=BackendType.AWS,
+        )
+        runpod_instance = await create_instance(
+            session=session,
+            project=project,
+            pool=pool,
+            backend=BackendType.RUNPOD,
+        )
+        instances = [aws_instance, runpod_instance]
+        res = services_pools.filter_pool_instances(
+            pool_instances=instances,
+            profile=Profile(name="test"),
+            multinode=True,
+        )
+        assert res == [aws_instance]
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
+    async def test_returns_volume_instances(self, test_db, session: AsyncSession):
+        user = await create_user(session=session)
+        project = await create_project(session=session, owner=user)
+        pool = await create_pool(session=session, project=project)
+        aws_instance = await create_instance(
+            session=session,
+            project=project,
+            pool=pool,
+            backend=BackendType.AWS,
+        )
+        runpod_instance1 = await create_instance(
+            session=session,
+            project=project,
+            pool=pool,
+            backend=BackendType.RUNPOD,
+            region="eu",
+        )
+        runpod_instance2 = await create_instance(
+            session=session,
+            project=project,
+            pool=pool,
+            backend=BackendType.RUNPOD,
+            region="us",
+        )
+        instances = [aws_instance, runpod_instance1, runpod_instance2]
+        res = services_pools.filter_pool_instances(
+            pool_instances=instances,
+            profile=Profile(name="test"),
+            volumes=[
+                [
+                    get_volume(
+                        configuration=get_volume_configuration(
+                            backend=BackendType.RUNPOD, region="us"
+                        )
+                    )
+                ]
+            ],
+        )
+        assert res == [runpod_instance2]
+
+
 class TestGenerateInstanceName:
     @pytest.mark.asyncio
     @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)