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

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

@@ -28,9 +28,9 @@ from dstack._internal.server.background.tasks.process_instances import (
     process_instances,
 )
 from dstack._internal.server.testing.common import (
+    ComputeMockSpec,
     create_instance,
     create_job,
-    create_pool,
     create_project,
     create_repo,
     create_run,
@@ -49,10 +49,10 @@ class TestCheckShim:
         self, test_db, session: AsyncSession
     ):
         project = await create_project(session=session)
-        pool = await create_pool(session, project)
-
         instance = await create_instance(
-            session, project, pool, status=InstanceStatus.PROVISIONING
+            session=session,
+            project=project,
+            status=InstanceStatus.PROVISIONING,
         )
         instance.termination_deadline = get_current_datetime() + dt.timedelta(days=1)
         instance.health_status = "ssh connect problem"
@@ -78,10 +78,10 @@ class TestCheckShim:
         self, test_db, session: AsyncSession
     ):
         project = await create_project(session=session)
-        pool = await create_pool(session, project)
-
         instance = await create_instance(
-            session, project, pool, status=InstanceStatus.PROVISIONING
+            session=session,
+            project=project,
+            status=InstanceStatus.PROVISIONING,
         )
         instance.started_at = get_current_datetime() + dt.timedelta(minutes=-20)
         instance.health_status = "ssh connect problem"
@@ -110,7 +110,6 @@ class TestCheckShim:
     ):
         user = await create_user(session=session)
         project = await create_project(session=session, owner=user)
-        pool = await create_pool(session, project)
         repo = await create_repo(
             session=session,
             project_id=project.id,
@@ -121,9 +120,10 @@ class TestCheckShim:
             repo=repo,
             user=user,
         )
-
         instance = await create_instance(
-            session, project, pool, status=InstanceStatus.PROVISIONING
+            session=session,
+            project=project,
+            status=InstanceStatus.PROVISIONING,
         )
         instance.termination_deadline = get_current_datetime().replace(
             tzinfo=dt.timezone.utc
@@ -158,10 +158,11 @@ class TestCheckShim:
     @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
     async def test_check_shim_start_termination_deadline(self, test_db, session: AsyncSession):
         project = await create_project(session=session)
-        pool = await create_pool(session, project)
-
-        instance = await create_instance(session, project, pool, status=InstanceStatus.IDLE)
-
+        instance = await create_instance(
+            session=session,
+            project=project,
+            status=InstanceStatus.IDLE,
+        )
         health_status = "SSH connection fail"
         with patch(
             "dstack._internal.server.background.tasks.process_instances._instance_healthcheck"
@@ -183,9 +184,11 @@ class TestCheckShim:
     @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
     async def test_check_shim_stop_termination_deadline(self, test_db, session: AsyncSession):
         project = await create_project(session=session)
-        pool = await create_pool(session, project)
-
-        instance = await create_instance(session, project, pool, status=InstanceStatus.IDLE)
+        instance = await create_instance(
+            session=session,
+            project=project,
+            status=InstanceStatus.IDLE,
+        )
         instance.termination_deadline = get_current_datetime() + dt.timedelta(minutes=19)
         await session.commit()
 
@@ -206,9 +209,11 @@ class TestCheckShim:
     @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
     async def test_check_shim_terminate_instance_by_dedaline(self, test_db, session: AsyncSession):
         project = await create_project(session=session)
-        pool = await create_pool(session, project)
-
-        instance = await create_instance(session, project, pool, status=InstanceStatus.IDLE)
+        instance = await create_instance(
+            session=session,
+            project=project,
+            status=InstanceStatus.IDLE,
+        )
         termination_deadline_time = get_current_datetime() + dt.timedelta(minutes=-19)
         instance.termination_deadline = termination_deadline_time
         await session.commit()
@@ -251,7 +256,6 @@ class TestCheckShim:
     ):
         # see https://github.com/dstackai/dstack/issues/2041
         project = await create_project(session=session)
-        pool = await create_pool(session, project)
         if has_job:
             user = await create_user(session=session)
             repo = await create_repo(
@@ -272,9 +276,8 @@ class TestCheckShim:
         else:
             job = None
         instance = await create_instance(
-            session,
-            project,
-            pool,
+            session=session,
+            project=project,
             created_at=get_current_datetime(),
             termination_policy=termination_policy,
             status=InstanceStatus.IDLE,
@@ -302,8 +305,9 @@ class TestTerminateIdleTime:
     @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
     async def test_terminate_by_idle_timeout(self, test_db, session: AsyncSession):
         project = await create_project(session=session)
-        pool = await create_pool(session, project)
-        instance = await create_instance(session, project, pool, status=InstanceStatus.IDLE)
+        instance = await create_instance(
+            session=session, project=project, status=InstanceStatus.IDLE
+        )
         instance.termination_idle_time = 300
         instance.termination_policy = TerminationPolicy.DESTROY_AFTER_IDLE
         instance.last_job_processed_at = get_current_datetime() + dt.timedelta(minutes=-19)
@@ -320,11 +324,9 @@ class TestSSHInstanceTerminateProvisionTimeoutExpired:
     @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
     async def test_terminate_by_idle_timeout(self, test_db, session: AsyncSession):
         project = await create_project(session=session)
-        pool = await create_pool(session, project)
         instance = await create_instance(
-            session,
-            project,
-            pool,
+            session=session,
+            project=project,
             status=InstanceStatus.PENDING,
             created_at=get_current_datetime() - dt.timedelta(days=100),
         )
@@ -357,10 +359,9 @@ class TestTerminate:
     @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
     async def test_terminate(self, test_db, session: AsyncSession):
         project = await create_project(session=session)
-        pool = await create_pool(session, project)
-
-        instance = await create_instance(session, project, pool, status=InstanceStatus.TERMINATING)
-
+        instance = await create_instance(
+            session=session, project=project, status=InstanceStatus.TERMINATING
+        )
         reason = "some reason"
         instance.termination_reason = reason
         instance.last_job_processed_at = get_current_datetime() + dt.timedelta(minutes=-19)
@@ -384,8 +385,9 @@ class TestTerminate:
     @pytest.mark.parametrize("error", [BackendError("err"), RuntimeError("err")])
     async def test_terminate_retry(self, test_db, session: AsyncSession, error: Exception):
         project = await create_project(session=session)
-        pool = await create_pool(session, project)
-        instance = await create_instance(session, project, pool, status=InstanceStatus.TERMINATING)
+        instance = await create_instance(
+            session=session, project=project, status=InstanceStatus.TERMINATING
+        )
         instance.termination_reason = "some reason"
         initial_time = dt.datetime(2025, 1, 1, tzinfo=dt.timezone.utc)
         instance.last_job_processed_at = initial_time
@@ -415,8 +417,9 @@ class TestTerminate:
     @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
     async def test_terminate_not_retries_if_too_early(self, test_db, session: AsyncSession):
         project = await create_project(session=session)
-        pool = await create_pool(session, project)
-        instance = await create_instance(session, project, pool, status=InstanceStatus.TERMINATING)
+        instance = await create_instance(
+            session=session, project=project, status=InstanceStatus.TERMINATING
+        )
         instance.termination_reason = "some reason"
         initial_time = dt.datetime(2025, 1, 1, tzinfo=dt.timezone.utc)
         instance.last_job_processed_at = initial_time
@@ -446,8 +449,9 @@ class TestTerminate:
     @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
     async def test_terminate_on_termination_deadline(self, test_db, session: AsyncSession):
         project = await create_project(session=session)
-        pool = await create_pool(session, project)
-        instance = await create_instance(session, project, pool, status=InstanceStatus.TERMINATING)
+        instance = await create_instance(
+            session=session, project=project, status=InstanceStatus.TERMINATING
+        )
         instance.termination_reason = "some reason"
         initial_time = dt.datetime(2025, 1, 1, tzinfo=dt.timezone.utc)
         instance.last_job_processed_at = initial_time
@@ -505,11 +509,9 @@ class TestCreateInstance:
         expected_blocks: int,
     ):
         project = await create_project(session=session)
-        pool = await create_pool(session, project)
         instance = await create_instance(
-            session,
-            project,
-            pool,
+            session=session,
+            project=project,
             status=InstanceStatus.PENDING,
             total_blocks=requested_blocks,
             busy_blocks=0,
@@ -531,6 +533,7 @@ class TestCreateInstance:
                 price=1.0,
                 availability=InstanceAvailability.AVAILABLE,
             )
+            backend_mock.compute.return_value = Mock(spec=ComputeMockSpec)
             backend_mock.compute.return_value.get_offers_cached.return_value = [offer]
             backend_mock.compute.return_value.create_instance.return_value = JobProvisioningData(
                 backend=offer.backend,
@@ -611,11 +614,9 @@ class TestAddSSHInstance:
         host_info["cpus"] = cpus
         host_info["gpu_count"] = gpus
         project = await create_project(session=session)
-        pool = await create_pool(session, project)
         instance = await create_instance(
-            session,
-            project,
-            pool,
+            session=session,
+            project=project,
             status=InstanceStatus.PENDING,
             created_at=get_current_datetime(),
             remote_connection_info=get_remote_connection_info(),

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

@@ -21,7 +21,6 @@ from dstack._internal.server.testing.common import (
     create_instance,
     create_job,
     create_job_metrics_point,
-    create_pool,
     create_project,
     create_repo,
     create_run,
@@ -45,11 +44,9 @@ class TestCollectMetrics:
             session=session,
             project_id=project.id,
         )
-        pool = await create_pool(session=session, project=project)
         instance = await create_instance(
             session=session,
             project=project,
-            pool=pool,
             status=InstanceStatus.BUSY,
         )
         run = await create_run(

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

@@ -7,6 +7,7 @@ from dstack._internal.server.background.tasks.process_placement_groups import (
     process_placement_groups,
 )
 from dstack._internal.server.testing.common import (
+    ComputeMockSpec,
     create_fleet,
     create_placement_group,
     create_project,
@@ -34,6 +35,7 @@ class TestProcessPlacementGroups:
         with patch("dstack._internal.server.services.backends.get_project_backend_by_type") as m:
             aws_mock = Mock()
             m.return_value = aws_mock
+            aws_mock.compute.return_value = Mock(spec=ComputeMockSpec)
             await process_placement_groups()
             aws_mock.compute.return_value.delete_placement_group.assert_called_once()
         await session.refresh(placement_group1)

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

@@ -0,0 +1,186 @@
+from collections.abc import Generator
+from datetime import datetime, timezone
+from unittest.mock import Mock, patch
+
+import pytest
+import pytest_asyncio
+from freezegun import freeze_time
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from dstack._internal.core.models.instances import InstanceStatus
+from dstack._internal.core.models.runs import JobStatus
+from dstack._internal.core.models.users import GlobalRole, ProjectRole
+from dstack._internal.server.background.tasks.process_prometheus_metrics import (
+    collect_prometheus_metrics,
+    delete_prometheus_metrics,
+)
+from dstack._internal.server.models import JobModel, JobPrometheusMetrics
+from dstack._internal.server.services.projects import add_project_member
+from dstack._internal.server.testing.common import (
+    create_instance,
+    create_job,
+    create_job_prometheus_metrics,
+    create_project,
+    create_repo,
+    create_run,
+    create_user,
+    get_job_provisioning_data,
+)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
+@pytest.mark.usefixtures("test_db", "image_config_mock")
+class TestCollectPrometheusMetrics:
+    @pytest_asyncio.fixture
+    async def job(self, session: AsyncSession) -> JobModel:
+        user = await create_user(session=session, global_role=GlobalRole.USER)
+        project = await create_project(session=session, owner=user)
+        await add_project_member(
+            session=session, project=project, user=user, project_role=ProjectRole.USER
+        )
+        repo = await create_repo(
+            session=session,
+            project_id=project.id,
+        )
+        instance = await create_instance(
+            session=session,
+            project=project,
+            status=InstanceStatus.BUSY,
+        )
+        run = await create_run(
+            session=session,
+            project=project,
+            repo=repo,
+            user=user,
+        )
+        job = await create_job(
+            session=session,
+            run=run,
+            status=JobStatus.RUNNING,
+            job_provisioning_data=get_job_provisioning_data(),
+            instance_assigned=True,
+            instance=instance,
+        )
+        return job
+
+    @pytest.fixture
+    def ssh_tunnel_mock(self) -> Generator[Mock, None, None]:
+        with patch("dstack._internal.server.services.runner.ssh.SSHTunnel") as SSHTunnelMock:
+            yield SSHTunnelMock
+
+    @pytest.fixture
+    def shim_client_mock(self) -> Generator[Mock, None, None]:
+        with patch("dstack._internal.server.services.runner.client.ShimClient") as ShimClientMock:
+            yield ShimClientMock.return_value
+
+    @freeze_time(datetime(2023, 1, 2, 3, 5, 20, tzinfo=timezone.utc))
+    async def test_inserts_new_record(
+        self, session: AsyncSession, job: JobModel, ssh_tunnel_mock: Mock, shim_client_mock: Mock
+    ):
+        shim_client_mock.get_task_metrics.return_value = "# prom response"
+
+        await collect_prometheus_metrics()
+
+        ssh_tunnel_mock.assert_called_once()
+        shim_client_mock.get_task_metrics.assert_called_once()
+        res = await session.execute(
+            select(JobPrometheusMetrics).where(JobPrometheusMetrics.job_id == job.id)
+        )
+        metrics = res.scalar_one()
+        assert metrics.text == "# prom response"
+        assert metrics.collected_at == datetime(2023, 1, 2, 3, 5, 20)
+
+    @freeze_time(datetime(2023, 1, 2, 3, 5, 20, tzinfo=timezone.utc))
+    async def test_updates_record(
+        self, session: AsyncSession, job: JobModel, ssh_tunnel_mock: Mock, shim_client_mock: Mock
+    ):
+        metrics = await create_job_prometheus_metrics(
+            session=session,
+            job=job,
+            collected_at=datetime(2023, 1, 2, 3, 5, 0),
+            text="# prom old response",
+        )
+        shim_client_mock.get_task_metrics.return_value = "# prom new response"
+
+        await collect_prometheus_metrics()
+
+        ssh_tunnel_mock.assert_called_once()
+        shim_client_mock.get_task_metrics.assert_called_once()
+        res = await session.execute(
+            select(JobPrometheusMetrics)
+            .where(JobPrometheusMetrics.job_id == job.id)
+            .execution_options(populate_existing=True)
+        )
+        metrics = res.scalar_one()
+        assert metrics.text == "# prom new response"
+        assert metrics.collected_at == datetime(2023, 1, 2, 3, 5, 20)
+
+    @freeze_time(datetime(2023, 1, 2, 3, 5, 20, tzinfo=timezone.utc))
+    async def test_skips_recently_updated(
+        self, session: AsyncSession, job: JobModel, ssh_tunnel_mock: Mock, shim_client_mock: Mock
+    ):
+        metrics = await create_job_prometheus_metrics(
+            session=session,
+            job=job,
+            collected_at=datetime(2023, 1, 2, 3, 5, 15),
+            text="# prom old response",
+        )
+        shim_client_mock.get_task_metrics.return_value = "# prom new response"
+
+        await collect_prometheus_metrics()
+
+        ssh_tunnel_mock.assert_not_called()
+        shim_client_mock.get_task_metrics.assert_not_called()
+        res = await session.execute(
+            select(JobPrometheusMetrics)
+            .where(JobPrometheusMetrics.job_id == job.id)
+            .execution_options(populate_existing=True)
+        )
+        metrics = res.scalar_one()
+        assert metrics.text == "# prom old response"
+        assert metrics.collected_at == datetime(2023, 1, 2, 3, 5, 15)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
+@pytest.mark.usefixtures("test_db", "image_config_mock")
+class TestDeletePrometheusMetrics:
+    @freeze_time(datetime(2023, 1, 2, 3, 5, 20, tzinfo=timezone.utc))
+    async def test_deletes_old_metrics(self, session: AsyncSession):
+        user = await create_user(session=session, global_role=GlobalRole.USER)
+        project = await create_project(session=session, owner=user)
+        await add_project_member(
+            session=session, project=project, user=user, project_role=ProjectRole.USER
+        )
+        repo = await create_repo(session=session, project_id=project.id)
+        run_1 = await create_run(
+            session=session, project=project, repo=repo, user=user, run_name="run-1"
+        )
+        job_1 = await create_job(session=session, run=run_1)
+        # old metrics
+        await create_job_prometheus_metrics(
+            session=session,
+            job=job_1,
+            collected_at=datetime(2023, 1, 2, 2, 3, 30),
+        )
+        run_2 = await create_run(
+            session=session, project=project, repo=repo, user=user, run_name="run-2"
+        )
+        job_2 = await create_job(session=session, run=run_2)
+        # recent metrics
+        metrics_2 = await create_job_prometheus_metrics(
+            session=session,
+            job=job_2,
+            collected_at=datetime(2023, 1, 2, 3, 5, 0),
+        )
+
+        await delete_prometheus_metrics()
+
+        res = await session.execute(
+            select(JobPrometheusMetrics).join(JobModel).where(JobModel.project_id == project.id)
+        )
+        all_metrics = res.scalars().all()
+        assert len(all_metrics) == 1
+        assert all_metrics[0] == metrics_2