Test idle volumes with pipelines

Author: r4victor

src/dstack/_internal/server/background/scheduled_tasks/idle_volumes.py

@@ -36,7 +36,7 @@ async def process_idle_volumes():
                 select(VolumeModel.id)
                 .where(
                     VolumeModel.status == VolumeStatus.ACTIVE,
-                    VolumeModel.auto_cleanup_enabled != False,
+                    VolumeModel.auto_cleanup_enabled.is_not(False),
                     VolumeModel.deleted == False,
                     VolumeModel.lock_expires_at.is_(None),
                     VolumeModel.id.not_in(lockset),

src/tests/_internal/server/background/scheduled_tasks/test_idle_volumes.py

@@ -22,12 +22,25 @@
     get_volume_provisioning_data,
     list_events,
 )
+from dstack._internal.settings import FeatureFlags
 from dstack._internal.utils.common import get_current_datetime
 
 
+@pytest.fixture
+def patch_pipeline_processing_flag(monkeypatch: pytest.MonkeyPatch):
+    def _apply(enabled: bool):
+        monkeypatch.setattr(FeatureFlags, "PIPELINE_PROCESSING_ENABLED", enabled)
+
+    return _apply
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
-class TestProcessIdleVolumes:
+class TestProcessIdleVolumesScheduledTask:
+    @pytest.fixture(autouse=True)
+    def _patch_feature_flag(self, patch_pipeline_processing_flag):
+        patch_pipeline_processing_flag(False)
+
     async def test_deletes_idle_volumes(self, test_db, session: AsyncSession):
         project = await create_project(session=session)
         user = await create_user(session=session)
@@ -71,17 +84,169 @@ async def test_deletes_idle_volumes(self, test_db, session: AsyncSession):
             m.return_value = aws_mock
             aws_mock.compute.return_value = Mock(spec=ComputeMockSpec)
             await process_idle_volumes()
+            m.assert_called_once()
 
         await session.refresh(volume1)
         await session.refresh(volume2)
         events = await list_events(session)
+        assert not volume1.to_be_deleted
         assert volume1.deleted
         assert volume1.deleted_at is not None
+        assert not volume2.to_be_deleted
         assert not volume2.deleted
         assert volume2.deleted_at is None
         assert len(events) == 1
         assert events[0].message == "Volume deleted due to exceeding auto_cleanup_duration"
 
+    async def test_deletes_idle_volume_with_null_auto_cleanup_enabled(
+        self, test_db, session: AsyncSession
+    ):
+        project = await create_project(session=session)
+        user = await create_user(session=session)
+        volume = await create_volume(
+            session=session,
+            project=project,
+            user=user,
+            status=VolumeStatus.ACTIVE,
+            backend=BackendType.AWS,
+            configuration=get_volume_configuration(
+                name="test-volume",
+                auto_cleanup_duration="1h",
+            ),
+            volume_provisioning_data=get_volume_provisioning_data(),
+            last_job_processed_at=datetime.datetime.now(datetime.timezone.utc)
+            - datetime.timedelta(hours=2),
+        )
+        volume.auto_cleanup_enabled = None
+        await session.commit()
+
+        with patch(
+            "dstack._internal.server.services.backends.get_project_backend_by_type_or_error"
+        ) as m:
+            aws_mock = Mock()
+            m.return_value = aws_mock
+            aws_mock.compute.return_value = Mock(spec=ComputeMockSpec)
+            await process_idle_volumes()
+            m.assert_called_once()
+
+        await session.refresh(volume)
+        events = await list_events(session)
+        assert not volume.to_be_deleted
+        assert volume.deleted
+        assert volume.deleted_at is not None
+        assert len(events) == 1
+        assert events[0].message == "Volume deleted due to exceeding auto_cleanup_duration"
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
+class TestProcessIdleVolumesPipelineTask:
+    @pytest.fixture(autouse=True)
+    def _patch_feature_flag(self, patch_pipeline_processing_flag):
+        patch_pipeline_processing_flag(True)
+
+    async def test_deletes_idle_volumes(self, test_db, session: AsyncSession):
+        project = await create_project(session=session)
+        user = await create_user(session=session)
+
+        config1 = get_volume_configuration(
+            name="test-volume",
+            auto_cleanup_duration="1h",
+        )
+        config2 = get_volume_configuration(
+            name="test-volume",
+            auto_cleanup_duration="3h",
+        )
+        volume1 = await create_volume(
+            session=session,
+            project=project,
+            user=user,
+            status=VolumeStatus.ACTIVE,
+            backend=BackendType.AWS,
+            configuration=config1,
+            volume_provisioning_data=get_volume_provisioning_data(),
+            last_job_processed_at=datetime.datetime.now(datetime.timezone.utc)
+            - datetime.timedelta(hours=2),
+        )
+        volume2 = await create_volume(
+            session=session,
+            project=project,
+            user=user,
+            status=VolumeStatus.ACTIVE,
+            backend=BackendType.AWS,
+            configuration=config2,
+            volume_provisioning_data=get_volume_provisioning_data(),
+            last_job_processed_at=datetime.datetime.now(datetime.timezone.utc)
+            - datetime.timedelta(hours=2),
+        )
+        await session.commit()
+
+        with patch(
+            "dstack._internal.server.services.backends.get_project_backend_by_type_or_error"
+        ) as m:
+            aws_mock = Mock()
+            m.return_value = aws_mock
+            aws_mock.compute.return_value = Mock(spec=ComputeMockSpec)
+            await process_idle_volumes()
+            m.assert_not_called()
+
+        await session.refresh(volume1)
+        await session.refresh(volume2)
+        events = await list_events(session)
+        assert volume1.to_be_deleted
+        assert not volume1.deleted
+        assert volume1.deleted_at is None
+        assert not volume2.to_be_deleted
+        assert not volume2.deleted
+        assert volume2.deleted_at is None
+        assert len(events) == 1
+        assert (
+            events[0].message
+            == "Volume marked for deletion due to exceeding auto_cleanup_duration"
+        )
+
+    async def test_deletes_idle_volume_with_null_auto_cleanup_enabled(
+        self, test_db, session: AsyncSession
+    ):
+        project = await create_project(session=session)
+        user = await create_user(session=session)
+        volume = await create_volume(
+            session=session,
+            project=project,
+            user=user,
+            status=VolumeStatus.ACTIVE,
+            backend=BackendType.AWS,
+            configuration=get_volume_configuration(
+                name="test-volume",
+                auto_cleanup_duration="1h",
+            ),
+            volume_provisioning_data=get_volume_provisioning_data(),
+            last_job_processed_at=datetime.datetime.now(datetime.timezone.utc)
+            - datetime.timedelta(hours=2),
+        )
+        volume.auto_cleanup_enabled = None
+        await session.commit()
+
+        with patch(
+            "dstack._internal.server.services.backends.get_project_backend_by_type_or_error"
+        ) as m:
+            aws_mock = Mock()
+            m.return_value = aws_mock
+            aws_mock.compute.return_value = Mock(spec=ComputeMockSpec)
+            await process_idle_volumes()
+            m.assert_not_called()
+
+        await session.refresh(volume)
+        events = await list_events(session)
+        assert volume.to_be_deleted
+        assert not volume.deleted
+        assert volume.deleted_at is None
+        assert len(events) == 1
+        assert (
+            events[0].message
+            == "Volume marked for deletion due to exceeding auto_cleanup_duration"
+        )
+
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)

src/tests/_internal/server/background/scheduled_tasks/test_submitted_jobs.py

@@ -561,9 +561,9 @@ async def test_fails_job_when_attaching_volume_marked_for_deletion(
         job = res.unique().scalar_one()
         assert job.status == JobStatus.TERMINATING
         assert job.termination_reason == JobTerminationReason.VOLUME_ERROR
-        assert job.termination_reason_message == "Failed to attach volume"
-        assert job.instance is not None
-        assert len(job.instance.volume_attachments) == 0
+        assert job.termination_reason_message is not None
+        assert "marked for deletion and cannot be attached" in job.termination_reason_message
+        assert job.instance is None
 
     @pytest.mark.asyncio
     @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)