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

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_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)