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

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

@@ -4,6 +4,7 @@ from typing import Optional
 from unittest.mock import MagicMock, Mock, patch
 
 import pytest
+from freezegun import freeze_time
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from dstack._internal.core.errors import SSHError
@@ -11,6 +12,7 @@ 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.profiles import UtilizationPolicy
 from dstack._internal.core.models.runs import (
     JobRuntimeData,
     JobStatus,
@@ -39,7 +41,7 @@ from dstack._internal.server.services.volumes import (
 from dstack._internal.server.testing.common import (
     create_instance,
     create_job,
-    create_pool,
+    create_job_metrics_point,
     create_project,
     create_repo,
     create_run,
@@ -103,11 +105,9 @@ 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)
@@ -153,11 +153,9 @@ 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)
@@ -204,11 +202,9 @@ 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)
@@ -304,11 +300,9 @@ 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)
@@ -374,11 +368,9 @@ 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(
@@ -431,11 +423,9 @@ 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)
@@ -476,11 +466,9 @@ class TestProcessRunningJobs:
             repo=repo,
             user=user,
         )
-        pool = await create_pool(session, project)
         instance = await create_instance(
             session=session,
             project=project,
-            pool=pool,
             status=InstanceStatus.IDLE,
         )
         job_provisioning_data = get_job_provisioning_data(dockerized=True)
@@ -525,11 +513,9 @@ 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(
@@ -652,11 +638,9 @@ class TestProcessRunningJobs:
                 ),
             ),
         )
-        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(
@@ -688,3 +672,123 @@ class TestProcessRunningJobs:
         assert job.status == expected_status
         assert job.termination_reason == expected_termination_reason
         assert job.inactivity_secs == expected_inactivity_secs
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
+    @pytest.mark.parametrize(
+        ["samples", "expected_status"],
+        [
+            pytest.param(
+                [
+                    (datetime(2023, 1, 1, 12, 25, 20, tzinfo=timezone.utc), 30),
+                    (datetime(2023, 1, 1, 12, 25, 30, tzinfo=timezone.utc), 30),
+                    (datetime(2023, 1, 1, 12, 29, 50, tzinfo=timezone.utc), 40),
+                ],
+                JobStatus.RUNNING,
+                id="not-enough-points",
+            ),
+            pytest.param(
+                [
+                    (datetime(2023, 1, 1, 12, 20, 10, tzinfo=timezone.utc), 30),
+                    (datetime(2023, 1, 1, 12, 20, 20, tzinfo=timezone.utc), 30),
+                    (datetime(2023, 1, 1, 12, 29, 50, tzinfo=timezone.utc), 80),
+                ],
+                JobStatus.RUNNING,
+                id="any-above-min",
+            ),
+            pytest.param(
+                [
+                    (datetime(2023, 1, 1, 12, 10, 10, tzinfo=timezone.utc), 80),  # outside window
+                    (datetime(2023, 1, 1, 12, 10, 20, tzinfo=timezone.utc), 80),  # outside window
+                    (datetime(2023, 1, 1, 12, 20, 10, tzinfo=timezone.utc), 30),
+                    (datetime(2023, 1, 1, 12, 20, 20, tzinfo=timezone.utc), 30),
+                    (datetime(2023, 1, 1, 12, 29, 50, tzinfo=timezone.utc), 40),
+                ],
+                JobStatus.TERMINATING,
+                id="all-below-min",
+            ),
+        ],
+    )
+    @freeze_time(datetime(2023, 1, 1, 12, 30, tzinfo=timezone.utc))
+    async def test_gpu_utilization(
+        self,
+        test_db,
+        session: AsyncSession,
+        samples: list[tuple[datetime, int]],
+        expected_status: JobStatus,
+    ) -> 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",
+                    utilization_policy=UtilizationPolicy(
+                        min_gpu_utilization=80,
+                        time_window=600,
+                    ),
+                ),
+            ),
+        )
+        instance = await create_instance(
+            session=session,
+            project=project,
+            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,
+        )
+        for timestamp, gpu_util in samples:
+            # two GPUs, the second one always 100% utilized
+            await create_job_metrics_point(
+                session=session,
+                job_model=job,
+                timestamp=timestamp,
+                gpus_memory_usage_bytes=[1024, 1024],
+                gpus_util_percent=[gpu_util, 100],
+            )
+        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=0,
+            )
+            await process_running_jobs()
+            SSHTunnelMock.assert_called_once()
+            runner_client_mock.pull.assert_called_once()
+        await session.refresh(job)
+        assert job.status == expected_status
+        if expected_status == JobStatus.TERMINATING:
+            assert job.termination_reason == JobTerminationReason.TERMINATED_BY_SERVER
+            assert job.termination_reason_message == (
+                "The job GPU utilization below 80% for 600 seconds"
+            )
+        else:
+            assert job.termination_reason is None
+            assert job.termination_reason_message is None

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

@@ -21,7 +21,6 @@ from dstack._internal.server.models import RunModel
 from dstack._internal.server.testing.common import (
     create_instance,
     create_job,
-    create_pool,
     create_project,
     create_repo,
     create_run,
@@ -42,9 +41,6 @@ async def make_run(
         session=session,
         project_id=project.id,
     )
-    project.default_pool = await create_pool(
-        session=session, project=project, pool_name="default-pool"
-    )
     run_name = "test-run"
     profile = Profile(
         name="test-profile",
@@ -60,7 +56,7 @@ async def make_run(
             replicas=parse_obj_as(Range[int], replicas),
         ),
     )
-    return await create_run(
+    run = await create_run(
         session=session,
         project=project,
         repo=repo,
@@ -69,6 +65,8 @@ async def make_run(
         run_spec=run_spec,
         status=status,
     )
+    run.project = project
+    return run
 
 
 class TestProcessRuns:
@@ -117,11 +115,9 @@ 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(
@@ -146,9 +142,7 @@ class TestProcessRuns:
     @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
     async def test_retry_running_to_pending(self, test_db, session: AsyncSession):
         run = await make_run(session, status=RunStatus.RUNNING)
-        instance = await create_instance(
-            session, project=run.project, pool=run.project.default_pool, spot=True
-        )
+        instance = await create_instance(session, project=run.project, spot=True)
         await create_job(
             session=session,
             run=run,
@@ -169,9 +163,7 @@ class TestProcessRuns:
     @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
     async def test_retry_running_to_failed(self, test_db, session: AsyncSession):
         run = await make_run(session, status=RunStatus.RUNNING)
-        instance = await create_instance(
-            session, project=run.project, pool=run.project.default_pool, spot=True
-        )
+        instance = await create_instance(session, project=run.project, spot=True)
         # job exited with non-zero code
         await create_job(
             session=session,
@@ -237,9 +229,7 @@ class TestProcessRunsReplicas:
             submitted_at=run.submitted_at,
             last_processed_at=run.submitted_at,
             replica_num=0,
-            instance=await create_instance(
-                session, project=run.project, pool=run.project.default_pool, spot=True
-            ),
+            instance=await create_instance(session, project=run.project, spot=True),
             job_provisioning_data=get_job_provisioning_data(),
         )
         await create_job(
@@ -250,9 +240,7 @@ class TestProcessRunsReplicas:
             submitted_at=run.submitted_at,
             last_processed_at=run.submitted_at,
             replica_num=1,
-            instance=await create_instance(
-                session, project=run.project, pool=run.project.default_pool, spot=True
-            ),
+            instance=await create_instance(session, project=run.project, spot=True),
             job_provisioning_data=get_job_provisioning_data(),
         )
         with patch("dstack._internal.utils.common.get_current_datetime") as datetime_mock:
@@ -273,9 +261,7 @@ class TestProcessRunsReplicas:
             submitted_at=run.submitted_at,
             last_processed_at=run.last_processed_at,
             replica_num=0,
-            instance=await create_instance(
-                session, project=run.project, pool=run.project.default_pool, spot=True
-            ),
+            instance=await create_instance(session, project=run.project, spot=True),
             job_provisioning_data=get_job_provisioning_data(),
         )
         await create_job(

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

@@ -14,7 +14,7 @@ from dstack._internal.core.models.instances import (
     InstanceType,
     Resources,
 )
-from dstack._internal.core.models.profiles import Profile, ProfileRetryPolicy
+from dstack._internal.core.models.profiles import Profile
 from dstack._internal.core.models.runs import (
     JobProvisioningData,
     JobStatus,
@@ -29,10 +29,10 @@ from dstack._internal.core.models.volumes import (
 from dstack._internal.server.background.tasks.process_submitted_jobs import process_submitted_jobs
 from dstack._internal.server.models import InstanceModel, JobModel, VolumeAttachmentModel
 from dstack._internal.server.testing.common import (
+    ComputeMockSpec,
     create_fleet,
     create_instance,
     create_job,
-    create_pool,
     create_project,
     create_repo,
     create_run,
@@ -52,7 +52,6 @@ class TestProcessSubmittedJobs:
     async def test_fails_job_when_no_backends(self, test_db, session: AsyncSession):
         project = await create_project(session=session)
         user = await create_user(session=session)
-        await create_pool(session=session, project=project)
         repo = await create_repo(
             session=session,
             project_id=project.id,
@@ -93,7 +92,6 @@ class TestProcessSubmittedJobs:
     ):
         project = await create_project(session=session)
         user = await create_user(session=session)
-        pool = await create_pool(session=session, project=project)
         repo = await create_repo(
             session=session,
             project_id=project.id,
@@ -151,12 +149,6 @@ class TestProcessSubmittedJobs:
         assert job is not None
         assert job.status == JobStatus.PROVISIONING
 
-        await session.refresh(pool)
-        instance_offer = InstanceOfferWithAvailability.parse_raw(pool.instances[0].offer)
-        assert offer == instance_offer
-        pool_job_provisioning_data = pool.instances[0].job_provisioning_data
-        assert pool_job_provisioning_data == job.job_provisioning_data
-
     @pytest.mark.asyncio
     @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
     async def test_fails_job_when_privileged_true_and_no_offers_with_create_instance_support(
@@ -166,7 +158,6 @@ class TestProcessSubmittedJobs:
     ):
         project = await create_project(session=session)
         user = await create_user(session=session)
-        pool = await create_pool(session=session, project=project)
         repo = await create_repo(
             session=session,
             project_id=project.id,
@@ -227,9 +218,6 @@ class TestProcessSubmittedJobs:
         assert job.status == JobStatus.TERMINATING
         assert job.termination_reason == JobTerminationReason.FAILED_TO_START_DUE_TO_NO_CAPACITY
 
-        await session.refresh(pool)
-        assert not pool.instances
-
     @pytest.mark.asyncio
     @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
     async def test_fails_job_when_instance_mounts_and_no_offers_with_create_instance_support(
@@ -239,7 +227,6 @@ class TestProcessSubmittedJobs:
     ):
         project = await create_project(session=session)
         user = await create_user(session=session)
-        pool = await create_pool(session=session, project=project)
         repo = await create_repo(
             session=session,
             project_id=project.id,
@@ -300,9 +287,6 @@ class TestProcessSubmittedJobs:
         assert job.status == JobStatus.TERMINATING
         assert job.termination_reason == JobTerminationReason.FAILED_TO_START_DUE_TO_NO_CAPACITY
 
-        await session.refresh(pool)
-        assert not pool.instances
-
     @pytest.mark.asyncio
     @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
     async def test_provisions_job_with_optional_instance_volume_not_attached(
@@ -312,7 +296,6 @@ class TestProcessSubmittedJobs:
     ):
         project = await create_project(session=session)
         user = await create_user(session=session)
-        pool = await create_pool(session=session, project=project)
         repo = await create_repo(
             session=session,
             project_id=project.id,
@@ -369,18 +352,11 @@ class TestProcessSubmittedJobs:
         assert job is not None
         assert job.status == JobStatus.PROVISIONING
 
-        await session.refresh(pool)
-        instance_offer = InstanceOfferWithAvailability.parse_raw(pool.instances[0].offer)
-        assert offer == instance_offer
-        pool_job_provisioning_data = pool.instances[0].job_provisioning_data
-        assert pool_job_provisioning_data == job.job_provisioning_data
-
     @pytest.mark.asyncio
     @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
     async def test_fails_job_when_no_capacity(self, test_db, session: AsyncSession):
         project = await create_project(session=session)
         user = await create_user(session=session)
-        pool = await create_pool(session=session, project=project)
         repo = await create_repo(
             session=session,
             project_id=project.id,
@@ -396,7 +372,6 @@ class TestProcessSubmittedJobs:
                 repo_id=repo.name,
                 profile=Profile(
                     name="default",
-                    retry_policy=ProfileRetryPolicy(retry=True, duration=3600),
                 ),
             ),
         )
@@ -414,15 +389,12 @@ class TestProcessSubmittedJobs:
         assert job is not None
         assert job.status == JobStatus.TERMINATING
         assert job.termination_reason == JobTerminationReason.FAILED_TO_START_DUE_TO_NO_CAPACITY
-        await session.refresh(pool)
-        assert not pool.instances
 
     @pytest.mark.asyncio
     @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
     async def test_assignes_job_to_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,
@@ -430,10 +402,8 @@ class TestProcessSubmittedJobs:
         instance = await create_instance(
             session=session,
             project=project,
-            pool=pool,
             status=InstanceStatus.IDLE,
         )
-        await session.refresh(pool)
         run = await create_run(
             session=session,
             project=project,
@@ -459,7 +429,6 @@ class TestProcessSubmittedJobs:
     async def test_assigns_job_to_instance_with_volumes(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,
@@ -476,12 +445,10 @@ class TestProcessSubmittedJobs:
         instance = await create_instance(
             session=session,
             project=project,
-            pool=pool,
             status=InstanceStatus.IDLE,
             backend=BackendType.AWS,
             region="us-east-1",
         )
-        await session.refresh(pool)
         run_spec = get_run_spec(run_name="test-run", repo_id=repo.name)
         run_spec.configuration.volumes = [
             VolumeMountPoint(name=volume.name, path="/volume"),
@@ -506,6 +473,7 @@ class TestProcessSubmittedJobs:
             backend_mock = Mock()
             m.return_value = backend_mock
             backend_mock.TYPE = BackendType.AWS
+            backend_mock.compute.return_value = Mock(spec=ComputeMockSpec)
             backend_mock.compute.return_value.attach_volume.return_value = VolumeAttachmentData()
             # Submitted jobs processing happens in two steps
             await process_submitted_jobs()
@@ -532,7 +500,6 @@ class TestProcessSubmittedJobs:
     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,
@@ -541,13 +508,11 @@ class TestProcessSubmittedJobs:
         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,
@@ -577,12 +542,10 @@ class TestProcessSubmittedJobs:
         project = await create_project(session)
         user = await create_user(session)
         repo = await create_repo(session=session, project_id=project.id)
-        pool = await create_pool(session=session, project=project)
         fleet = await create_fleet(session=session, project=project)
         instance = await create_instance(
             session=session,
             project=project,
-            pool=pool,
             instance_num=0,
             status=InstanceStatus.BUSY,
         )

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

@@ -7,6 +7,7 @@ from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.volumes import VolumeProvisioningData, VolumeStatus
 from dstack._internal.server.background.tasks.process_volumes import process_submitted_volumes
 from dstack._internal.server.testing.common import (
+    ComputeMockSpec,
     create_project,
     create_user,
     create_volume,
@@ -40,6 +41,7 @@ class TestProcessSubmittedVolumes:
         ) as m:
             aws_mock = Mock()
             m.return_value = aws_mock
+            aws_mock.compute.return_value = Mock(spec=ComputeMockSpec)
             aws_mock.compute.return_value.create_volume.return_value = VolumeProvisioningData(
                 backend=BackendType.AWS,
                 volume_id="1234",