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

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

@@ -1,5 +1,6 @@
 from datetime import datetime, timezone
 from pathlib import Path
+from typing import Optional
 from unittest.mock import MagicMock, Mock, patch
 
 import pytest
@@ -8,11 +9,13 @@ from sqlalchemy.ext.asyncio import AsyncSession
 from dstack._internal.core.errors import SSHError
 from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.common import NetworkMode
+from dstack._internal.core.models.configurations import DevEnvironmentConfiguration
 from dstack._internal.core.models.instances import InstanceStatus
 from dstack._internal.core.models.runs import (
     JobRuntimeData,
     JobStatus,
     JobTerminationReason,
+    RunStatus,
 )
 from dstack._internal.core.models.volumes import (
     InstanceMountPoint,
@@ -100,6 +103,13 @@ class TestProcessRunningJobs:
             repo=repo,
             user=user,
         )
+        pool = await create_pool(session=session, project=project)
+        instance = await create_instance(
+            session=session,
+            project=project,
+            pool=pool,
+            status=InstanceStatus.BUSY,
+        )
         job_provisioning_data = get_job_provisioning_data(dockerized=False)
         job = await create_job(
             session=session,
@@ -107,6 +117,8 @@ class TestProcessRunningJobs:
             status=JobStatus.PROVISIONING,
             submitted_at=datetime(2023, 1, 2, 5, 12, 30, 5, tzinfo=timezone.utc),
             job_provisioning_data=job_provisioning_data,
+            instance=instance,
+            instance_assigned=True,
         )
         with (
             patch("dstack._internal.server.services.runner.ssh.SSHTunnel") as SSHTunnelMock,
@@ -141,12 +153,21 @@ class TestProcessRunningJobs:
             repo=repo,
             user=user,
         )
+        pool = await create_pool(session=session, project=project)
+        instance = await create_instance(
+            session=session,
+            project=project,
+            pool=pool,
+            status=InstanceStatus.BUSY,
+        )
         job_provisioning_data = get_job_provisioning_data(dockerized=False)
         job = await create_job(
             session=session,
             run=run,
             status=JobStatus.PROVISIONING,
             job_provisioning_data=job_provisioning_data,
+            instance=instance,
+            instance_assigned=True,
         )
         with (
             patch("dstack._internal.server.services.runner.ssh.SSHTunnel") as SSHTunnelMock,
@@ -183,12 +204,21 @@ class TestProcessRunningJobs:
             repo=repo,
             user=user,
         )
+        pool = await create_pool(session=session, project=project)
+        instance = await create_instance(
+            session=session,
+            project=project,
+            pool=pool,
+            status=InstanceStatus.BUSY,
+        )
         job_provisioning_data = get_job_provisioning_data(dockerized=False)
         job = await create_job(
             session=session,
             run=run,
             status=JobStatus.RUNNING,
             job_provisioning_data=job_provisioning_data,
+            instance=instance,
+            instance_assigned=True,
         )
         with (
             patch("dstack._internal.server.services.runner.ssh.SSHTunnel") as SSHTunnelMock,
@@ -274,6 +304,13 @@ class TestProcessRunningJobs:
             run_name="test-run",
             run_spec=run_spec,
         )
+        pool = await create_pool(session=session, project=project)
+        instance = await create_instance(
+            session=session,
+            project=project,
+            pool=pool,
+            status=InstanceStatus.BUSY,
+        )
         job_provisioning_data = get_job_provisioning_data(dockerized=True)
 
         with patch(
@@ -285,6 +322,8 @@ class TestProcessRunningJobs:
                 run=run,
                 status=JobStatus.PROVISIONING,
                 job_provisioning_data=job_provisioning_data,
+                instance=instance,
+                instance_assigned=True,
             )
 
         await process_running_jobs()
@@ -334,12 +373,21 @@ class TestProcessRunningJobs:
             repo=repo,
             user=user,
         )
+        pool = await create_pool(session=session, project=project)
+        instance = await create_instance(
+            session=session,
+            project=project,
+            pool=pool,
+            status=InstanceStatus.BUSY,
+        )
         job = await create_job(
             session=session,
             run=run,
             status=JobStatus.PULLING,
             job_provisioning_data=get_job_provisioning_data(dockerized=True),
             job_runtime_data=get_job_runtime_data(network_mode="bridge", ports=None),
+            instance=instance,
+            instance_assigned=True,
         )
         shim_client_mock.get_task.return_value.status = TaskStatus.RUNNING
         shim_client_mock.get_task.return_value.ports = [
@@ -382,6 +430,13 @@ class TestProcessRunningJobs:
             repo=repo,
             user=user,
         )
+        pool = await create_pool(session=session, project=project)
+        instance = await create_instance(
+            session=session,
+            project=project,
+            pool=pool,
+            status=InstanceStatus.BUSY,
+        )
         job_provisioning_data = get_job_provisioning_data(dockerized=True)
         job = await create_job(
             session=session,
@@ -389,6 +444,8 @@ class TestProcessRunningJobs:
             status=JobStatus.PULLING,
             job_provisioning_data=job_provisioning_data,
             job_runtime_data=get_job_runtime_data(network_mode="bridge", ports=None),
+            instance=instance,
+            instance_assigned=True,
         )
         shim_client_mock.get_task.return_value.status = TaskStatus.RUNNING
         shim_client_mock.get_task.return_value.ports = None
@@ -467,12 +524,21 @@ class TestProcessRunningJobs:
             run_name="test-run",
             run_spec=run_spec,
         )
+        pool = await create_pool(session=session, project=project)
+        instance = await create_instance(
+            session=session,
+            project=project,
+            pool=pool,
+            status=InstanceStatus.BUSY,
+        )
         job = await create_job(
             session=session,
             run=run,
             status=JobStatus.PROVISIONING,
             job_provisioning_data=get_job_provisioning_data(dockerized=True),
             submitted_at=get_current_datetime(),
+            instance=instance,
+            instance_assigned=True,
         )
         monkeypatch.setattr(
             "dstack._internal.server.services.runner.ssh.SSHTunnel", Mock(return_value=MagicMock())
@@ -495,3 +561,129 @@ class TestProcessRunningJobs:
         shim_client_mock.stop.assert_called_once_with(force=True)
         await session.refresh(job)
         assert job.status == JobStatus.PROVISIONING
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
+    @pytest.mark.parametrize(
+        (
+            "inactivity_duration",
+            "no_connections_secs",
+            "expected_status",
+            "expected_termination_reason",
+            "expected_inactivity_secs",
+        ),
+        [
+            pytest.param(
+                "1h",
+                60 * 60 - 1,
+                JobStatus.RUNNING,
+                None,
+                60 * 60 - 1,
+                id="duration-not-exceeded",
+            ),
+            pytest.param(
+                "1h",
+                60 * 60,
+                JobStatus.TERMINATING,
+                JobTerminationReason.TERMINATED_BY_SERVER,
+                60 * 60,
+                id="duration-exceeded-exactly",
+            ),
+            pytest.param(
+                "1h",
+                60 * 60 + 1,
+                JobStatus.TERMINATING,
+                JobTerminationReason.TERMINATED_BY_SERVER,
+                60 * 60 + 1,
+                id="duration-exceeded",
+            ),
+            pytest.param("off", 60 * 60, JobStatus.RUNNING, None, None, id="duration-off"),
+            pytest.param(False, 60 * 60, JobStatus.RUNNING, None, None, id="duration-false"),
+            pytest.param(None, 60 * 60, JobStatus.RUNNING, None, None, id="duration-none"),
+            pytest.param(
+                "1h",
+                None,
+                JobStatus.TERMINATING,
+                JobTerminationReason.INTERRUPTED_BY_NO_CAPACITY,
+                None,
+                id="legacy-runner",
+            ),
+            pytest.param(
+                None,
+                None,
+                JobStatus.RUNNING,
+                None,
+                None,
+                id="legacy-runner-without-duration",
+            ),
+        ],
+    )
+    async def test_inactivity_duration(
+        self,
+        test_db,
+        session: AsyncSession,
+        inactivity_duration,
+        no_connections_secs: Optional[int],
+        expected_status: JobStatus,
+        expected_termination_reason: Optional[JobTerminationReason],
+        expected_inactivity_secs: Optional[int],
+    ) -> None:
+        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,
+            status=RunStatus.RUNNING,
+            run_name="test-run",
+            run_spec=get_run_spec(
+                run_name="test-run",
+                repo_id=repo.name,
+                configuration=DevEnvironmentConfiguration(
+                    name="test-run",
+                    ide="vscode",
+                    inactivity_duration=inactivity_duration,
+                ),
+            ),
+        )
+        pool = await create_pool(session=session, project=project)
+        instance = await create_instance(
+            session=session,
+            project=project,
+            pool=pool,
+            status=InstanceStatus.BUSY,
+        )
+        job = await create_job(
+            session=session,
+            run=run,
+            status=JobStatus.RUNNING,
+            job_provisioning_data=get_job_provisioning_data(),
+            instance=instance,
+            instance_assigned=True,
+        )
+        with (
+            patch("dstack._internal.server.services.runner.ssh.SSHTunnel") as SSHTunnelMock,
+            patch(
+                "dstack._internal.server.services.runner.client.RunnerClient"
+            ) as RunnerClientMock,
+        ):
+            runner_client_mock = RunnerClientMock.return_value
+            runner_client_mock.pull.return_value = PullResponse(
+                job_states=[],
+                job_logs=[],
+                runner_logs=[],
+                last_updated=0,
+                no_connections_secs=no_connections_secs,
+            )
+            await process_running_jobs()
+            SSHTunnelMock.assert_called_once()
+            runner_client_mock.pull.assert_called_once()
+        await session.refresh(job)
+        assert job.status == expected_status
+        assert job.termination_reason == expected_termination_reason
+        assert job.inactivity_secs == expected_inactivity_secs

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

@@ -8,6 +8,7 @@ from sqlalchemy.ext.asyncio import AsyncSession
 
 import dstack._internal.server.background.tasks.process_runs as process_runs
 from dstack._internal.core.models.configurations import ServiceConfiguration
+from dstack._internal.core.models.instances import InstanceStatus
 from dstack._internal.core.models.profiles import Profile
 from dstack._internal.core.models.resources import Range
 from dstack._internal.core.models.runs import (
@@ -116,11 +117,20 @@ class TestProcessRuns:
     async def test_terminate_run_jobs(self, test_db, session: AsyncSession):
         run = await make_run(session, status=RunStatus.TERMINATING)
         run.termination_reason = RunTerminationReason.JOB_FAILED
+        pool = await create_pool(session=session, project=run.project)
+        instance = await create_instance(
+            session=session,
+            project=run.project,
+            pool=pool,
+            status=InstanceStatus.BUSY,
+        )
         job = await create_job(
             session=session,
             run=run,
             job_provisioning_data=get_job_provisioning_data(),
             status=JobStatus.RUNNING,
+            instance=instance,
+            instance_assigned=True,
         )
 
         with patch("dstack._internal.server.services.jobs._stop_runner") as stop_runner:
@@ -288,13 +298,27 @@ class TestProcessRunsReplicas:
 
     @pytest.mark.asyncio
     @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
-    async def test_some_failed_to_terminating(self, test_db, session: AsyncSession):
+    @pytest.mark.parametrize(
+        ("job_status", "job_termination_reason"),
+        [
+            (JobStatus.FAILED, JobTerminationReason.CONTAINER_EXITED_WITH_ERROR),
+            (JobStatus.TERMINATING, JobTerminationReason.TERMINATED_BY_SERVER),
+            (JobStatus.TERMINATED, JobTerminationReason.TERMINATED_BY_SERVER),
+        ],
+    )
+    async def test_some_failed_to_terminating(
+        self,
+        test_db,
+        session: AsyncSession,
+        job_status: JobStatus,
+        job_termination_reason: JobTerminationReason,
+    ) -> None:
         run = await make_run(session, status=RunStatus.RUNNING, replicas=2)
         await create_job(
             session=session,
             run=run,
-            status=JobStatus.FAILED,
-            termination_reason=JobTerminationReason.CONTAINER_EXITED_WITH_ERROR,
+            status=job_status,
+            termination_reason=job_termination_reason,
             replica_num=0,
         )
         await create_job(session=session, run=run, status=JobStatus.RUNNING, replica_num=1)

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

@@ -38,6 +38,7 @@ from dstack._internal.server.testing.common import (
     create_run,
     create_user,
     create_volume,
+    get_instance_offer_with_availability,
     get_run_spec,
     get_volume_provisioning_data,
 )
@@ -447,7 +448,7 @@ class TestProcessSubmittedJobs:
         await process_submitted_jobs()
         await session.refresh(job)
         res = await session.execute(select(JobModel).options(joinedload(JobModel.instance)))
-        job = res.scalar_one()
+        job = res.unique().scalar_one()
         assert job.status == JobStatus.SUBMITTED
         assert (
             job.instance_assigned and job.instance is not None and job.instance.id == instance.id
@@ -517,13 +518,57 @@ class TestProcessSubmittedJobs:
                 joinedload(JobModel.instance).selectinload(InstanceModel.volumes)
             )
         )
-        job = res.scalar_one()
+        job = res.unique().scalar_one()
         assert job.status == JobStatus.PROVISIONING
         assert (
             job.instance_assigned and job.instance is not None and job.instance.id == instance.id
         )
         assert job.instance.volumes == [volume]
 
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
+    async def test_assigns_job_to_shared_instance(self, test_db, session: AsyncSession):
+        project = await create_project(session)
+        user = await create_user(session)
+        pool = await create_pool(session=session, project=project)
+        repo = await create_repo(
+            session=session,
+            project_id=project.id,
+        )
+        offer = get_instance_offer_with_availability(gpu_count=8, cpu_count=64, memory_gib=128)
+        instance = await create_instance(
+            session=session,
+            project=project,
+            pool=pool,
+            status=InstanceStatus.IDLE,
+            offer=offer,
+            total_blocks=4,
+            busy_blocks=1,
+        )
+        await session.refresh(pool)
+        run = await create_run(
+            session=session,
+            project=project,
+            repo=repo,
+            user=user,
+        )
+        job = await create_job(
+            session=session,
+            run=run,
+            instance_assigned=False,
+        )
+        await process_submitted_jobs()
+        await session.refresh(job)
+        await session.refresh(instance)
+        res = await session.execute(select(JobModel).options(joinedload(JobModel.instance)))
+        job = res.unique().scalar_one()
+        assert job.status == JobStatus.SUBMITTED
+        assert (
+            job.instance_assigned and job.instance is not None and job.instance.id == instance.id
+        )
+        assert instance.total_blocks == 4
+        assert instance.busy_blocks == 2
+
     @pytest.mark.asyncio
     @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
     async def test_creates_new_instance_in_existing_fleet(self, test_db, session: AsyncSession):
@@ -590,7 +635,7 @@ class TestProcessSubmittedJobs:
 
         await session.refresh(job)
         res = await session.execute(select(JobModel).options(joinedload(JobModel.instance)))
-        job = res.scalar_one()
+        job = res.unique().scalar_one()
         assert job.status == JobStatus.PROVISIONING
         assert job.instance is not None
         assert job.instance.instance_num == 1

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

@@ -24,17 +24,19 @@ from dstack._internal.server.testing.common import (
     create_run,
     create_user,
     create_volume,
+    get_instance_offer_with_availability,
     get_job_provisioning_data,
+    get_job_runtime_data,
+    get_volume_configuration,
     get_volume_provisioning_data,
 )
 
-pytestmark = pytest.mark.usefixtures("image_config_mock")
-
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
+@pytest.mark.usefixtures("test_db", "image_config_mock")
 class TestProcessTerminatingJobs:
-    @pytest.mark.asyncio
-    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
-    async def test_terminates_job(self, test_db, session: AsyncSession):
+    async def test_terminates_job(self, session: AsyncSession):
         project = await create_project(session=session)
         user = await create_user(session=session)
         pool = await create_pool(session=session, project=project)
@@ -73,9 +75,7 @@ class TestProcessTerminatingJobs:
         assert job is not None
         assert job.status == JobStatus.TERMINATED
 
-    @pytest.mark.asyncio
-    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
-    async def test_detaches_job_volumes(self, test_db, session: AsyncSession):
+    async def test_detaches_job_volumes(self, session: AsyncSession):
         project = await create_project(session=session)
         user = await create_user(session=session)
         pool = await create_pool(session=session, project=project)
@@ -122,9 +122,7 @@ class TestProcessTerminatingJobs:
         await session.refresh(job)
         assert job.status == JobStatus.TERMINATED
 
-    @pytest.mark.asyncio
-    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
-    async def test_force_detaches_job_volumes(self, test_db, session: AsyncSession):
+    async def test_force_detaches_job_volumes(self, session: AsyncSession):
         project = await create_project(session=session)
         user = await create_user(session=session)
         pool = await create_pool(session=session, project=project)
@@ -172,7 +170,7 @@ class TestProcessTerminatingJobs:
             backend_mock.compute.return_value.is_volume_detached.assert_called_once()
         await session.refresh(job)
         res = await session.execute(select(JobModel).options(joinedload(JobModel.instance)))
-        job = res.scalar_one()
+        job = res.unique().scalar_one()
         assert job.status == JobStatus.TERMINATING
         # The instance should be released even if detach fails
         # so that stuck volumes don't prevent the instance from terminating.
@@ -214,10 +212,125 @@ class TestProcessTerminatingJobs:
         await session.refresh(job)
         await session.refresh(instance)
         res = await session.execute(select(JobModel).options(joinedload(JobModel.instance)))
-        job = res.scalar_one()
+        job = res.unique().scalar_one()
         res = await session.execute(
             select(InstanceModel).options(joinedload(InstanceModel.volumes))
         )
         instance = res.unique().scalar_one()
         assert job.status == JobStatus.TERMINATED
         assert len(instance.volumes) == 0
+
+    async def test_terminates_job_on_shared_instance(self, session: AsyncSession):
+        project = await create_project(session)
+        user = await create_user(session)
+        pool = await create_pool(session=session, project=project)
+        repo = await create_repo(
+            session=session,
+            project_id=project.id,
+        )
+        instance = await create_instance(
+            session=session,
+            project=project,
+            pool=pool,
+            status=InstanceStatus.BUSY,
+            total_blocks=4,
+            busy_blocks=3,
+        )
+        await session.refresh(pool)
+        run = await create_run(
+            session=session,
+            project=project,
+            repo=repo,
+            user=user,
+        )
+        shared_offer = get_instance_offer_with_availability(blocks=2, total_blocks=4)
+        jrd = get_job_runtime_data(offer=shared_offer)
+        job = await create_job(
+            session=session,
+            run=run,
+            instance_assigned=True,
+            instance=instance,
+            job_runtime_data=jrd,
+            status=JobStatus.TERMINATING,
+            termination_reason=JobTerminationReason.TERMINATED_BY_USER,
+        )
+
+        await process_terminating_jobs()
+
+        await session.refresh(job)
+        await session.refresh(instance)
+        res = await session.execute(select(JobModel).options(joinedload(JobModel.instance)))
+        job = res.unique().scalar_one()
+        assert job.status == JobStatus.TERMINATED
+        assert job.instance_assigned
+        assert job.instance is None
+        assert instance.total_blocks == 4
+        assert instance.busy_blocks == 1
+
+    async def test_detaches_job_volumes_on_shared_instance(self, session: AsyncSession):
+        project = await create_project(session=session)
+        user = await create_user(session=session)
+        pool = await create_pool(session=session, project=project)
+        volume_conf_1 = get_volume_configuration(name="vol-1")
+        volume_1 = await create_volume(
+            session=session,
+            project=project,
+            user=user,
+            status=VolumeStatus.ACTIVE,
+            backend=BackendType.AWS,
+            configuration=volume_conf_1,
+            volume_provisioning_data=get_volume_provisioning_data(),
+        )
+        volume_conf_2 = get_volume_configuration(name="vol-2")
+        volume_2 = await create_volume(
+            session=session,
+            project=project,
+            user=user,
+            status=VolumeStatus.ACTIVE,
+            backend=BackendType.AWS,
+            configuration=volume_conf_2,
+            volume_provisioning_data=get_volume_provisioning_data(),
+        )
+        instance = await create_instance(
+            session=session,
+            project=project,
+            pool=pool,
+            status=InstanceStatus.BUSY,
+            volumes=[volume_1, volume_2],
+        )
+        repo = await create_repo(session=session, project_id=project.id)
+        run = await create_run(
+            session=session,
+            project=project,
+            repo=repo,
+            user=user,
+        )
+        job_provisioning_data = get_job_provisioning_data(dockerized=False)
+        job = await create_job(
+            session=session,
+            run=run,
+            status=JobStatus.TERMINATING,
+            termination_reason=JobTerminationReason.TERMINATED_BY_USER,
+            submitted_at=datetime(2023, 1, 2, 5, 12, 30, 5, tzinfo=timezone.utc),
+            job_provisioning_data=job_provisioning_data,
+            job_runtime_data=get_job_runtime_data(volume_names=["vol-1"]),
+            instance=instance,
+        )
+        with patch("dstack._internal.server.services.backends.get_project_backend_by_type") as m:
+            backend_mock = Mock()
+            m.return_value = backend_mock
+            backend_mock.compute.return_value.is_volume_detached.return_value = True
+
+            await process_terminating_jobs()
+
+        m.assert_awaited_once()
+        backend_mock.compute.return_value.detach_volume.assert_called_once()
+        backend_mock.compute.return_value.is_volume_detached.assert_called_once()
+        await session.refresh(job)
+        await session.refresh(instance)
+        assert job.status == JobStatus.TERMINATED
+        res = await session.execute(
+            select(InstanceModel).options(joinedload(InstanceModel.volumes))
+        )
+        instance = res.unique().scalar_one()
+        assert instance.volumes == [volume_2]