Optimize idle volume cleanup implementation

Author: haydnli-shopify

src/dstack/_internal/server/background/tasks/process_idle_volumes.py

@@ -16,10 +16,7 @@
 logger = get_logger(__name__)
 
 
-async def process_idle_volumes(batch_size: int = 10):
-    """
-    Process volumes to check if they have exceeded their idle_duration and delete them.
-    """
+async def process_idle_volumes():
     lock, lockset = get_locker().get_lockset(VolumeModel.__tablename__)
     async with get_session_ctx() as session:
         async with lock:
@@ -31,36 +28,31 @@ async def process_idle_volumes(batch_size: int = 10):
                     VolumeModel.id.not_in(lockset),
                 )
                 .order_by(VolumeModel.last_processed_at.asc())
-                .limit(batch_size)
+                .limit(10)
                 .with_for_update(skip_locked=True)
             )
             volume_models = list(res.unique().scalars().all())
-            for volume_model in volume_models:
-                await session.refresh(volume_model, ["project", "attachments"])
             if not volume_models:
                 return
-
             for volume_model in volume_models:
+                await session.refresh(volume_model, ["project", "attachments"])
                 lockset.add(volume_model.id)
 
         try:
             volumes_to_delete = []
             for volume_model in volume_models:
-                if await _should_delete_idle_volume(volume_model):
+                if _should_delete_idle_volume(volume_model):
                     volumes_to_delete.append(volume_model)
 
             if volumes_to_delete:
                 await _delete_idle_volumes(session, volumes_to_delete)
 
         finally:
             for volume_model in volume_models:
-                lockset.difference_update([volume_model.id])
+                lockset.discard(volume_model.id)
 
 
-async def _should_delete_idle_volume(volume_model: VolumeModel) -> bool:
-    """
-    Check if a volume should be deleted based on its idle duration.
-    """
+def _should_delete_idle_volume(volume_model: VolumeModel) -> bool:
     configuration = get_volume_configuration(volume_model)
 
     if configuration.idle_duration is None:
@@ -82,32 +74,33 @@ async def _should_delete_idle_volume(volume_model: VolumeModel) -> bool:
 
     if idle_duration > idle_threshold:
         logger.info(
-            "Volume %s idle duration expired: idle time %s seconds, threshold %s seconds. Marking for deletion",
+            "Volume %s idle duration expired: idle time %.1f hours, threshold %.1f hours. Marking for deletion",
             volume_model.name,
-            idle_duration.total_seconds(),
-            idle_threshold.total_seconds(),
+            idle_duration.total_seconds() / 3600,
+            idle_threshold.total_seconds() / 3600,
         )
         return True
 
     return False
 
 
 def _get_volume_idle_duration(volume_model: VolumeModel) -> datetime.timedelta:
-    """
-    Calculate how long a volume has been idle.
-    A volume is considered idle from the time it was last processed by a job.
-    If it was never used by a job, use the created_at time.
-    """
-    last_time = volume_model.created_at.replace(tzinfo=datetime.timezone.utc)
+    reference_time = volume_model.created_at
     if volume_model.last_job_processed_at is not None:
-        last_time = volume_model.last_job_processed_at.replace(tzinfo=datetime.timezone.utc)
-    return get_current_datetime() - last_time
+        reference_time = volume_model.last_job_processed_at
+
+    reference_time_utc = reference_time.replace(tzinfo=datetime.timezone.utc)
+    current_time = get_current_datetime()
+
+    idle_duration = current_time - reference_time_utc
+
+    if idle_duration.total_seconds() < 0:
+        return datetime.timedelta(0)
+
+    return idle_duration
 
 
 async def _delete_idle_volumes(session: AsyncSession, volume_models: List[VolumeModel]):
-    """
-    Delete volumes that have exceeded their idle duration.
-    """
     volumes_by_project = {}
     for volume_model in volume_models:
         project = volume_model.project

src/tests/_internal/server/background/tasks/test_process_idle_volumes.py

@@ -22,15 +22,13 @@
 )
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
 class TestProcessIdleVolumes:
-    @pytest.mark.asyncio
-    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
     async def test_no_idle_duration_configured(self, test_db, session: AsyncSession):
-        """Test that volumes without idle_duration configured are not deleted."""
         project = await create_project(session=session)
         user = await create_user(session=session)
 
-        # Create volume without idle_duration
         volume = await create_volume(
             session=session,
             project=project,
@@ -41,17 +39,13 @@ async def test_no_idle_duration_configured(self, test_db, session: AsyncSession)
             volume_provisioning_data=get_volume_provisioning_data(),
         )
 
-        should_delete = await _should_delete_idle_volume(volume)
+        should_delete = _should_delete_idle_volume(volume)
         assert not should_delete
 
-    @pytest.mark.asyncio
-    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
     async def test_idle_duration_disabled(self, test_db, session: AsyncSession):
-        """Test that volumes with idle_duration set to -1 (disabled) are not deleted."""
         project = await create_project(session=session)
         user = await create_user(session=session)
 
-        # Create volume with idle_duration disabled
         volume_config = get_volume_configuration(name="test-volume")
         volume_config.idle_duration = -1
 
@@ -65,19 +59,15 @@ async def test_idle_duration_disabled(self, test_db, session: AsyncSession):
             volume_provisioning_data=get_volume_provisioning_data(),
         )
 
-        should_delete = await _should_delete_idle_volume(volume)
+        should_delete = _should_delete_idle_volume(volume)
         assert not should_delete
 
-    @pytest.mark.asyncio
-    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
     async def test_volume_still_attached(self, test_db, session: AsyncSession):
-        """Test that volumes still attached to instances are not deleted."""
         project = await create_project(session=session)
         user = await create_user(session=session)
 
-        # Create volume with idle_duration
         volume_config = get_volume_configuration(name="test-volume")
-        volume_config.idle_duration = "1h"  # 1 hour
+        volume_config.idle_duration = "1h"
 
         volume = await create_volume(
             session=session,
@@ -89,7 +79,6 @@ async def test_volume_still_attached(self, test_db, session: AsyncSession):
             volume_provisioning_data=get_volume_provisioning_data(),
         )
 
-        # Create an instance and attach the volume to it
         instance = await create_instance(session=session, project=project)
         attachment = VolumeAttachmentModel(
             volume_id=volume.id,
@@ -98,19 +87,15 @@ async def test_volume_still_attached(self, test_db, session: AsyncSession):
         volume.attachments.append(attachment)
         await session.commit()
 
-        should_delete = await _should_delete_idle_volume(volume)
+        should_delete = _should_delete_idle_volume(volume)
         assert not should_delete
 
-    @pytest.mark.asyncio
-    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
     async def test_idle_duration_not_exceeded(self, test_db, session: AsyncSession):
-        """Test that volumes within idle duration threshold are not deleted."""
         project = await create_project(session=session)
         user = await create_user(session=session)
 
-        # Create volume with idle_duration
         volume_config = get_volume_configuration(name="test-volume")
-        volume_config.idle_duration = "1h"  # 1 hour
+        volume_config.idle_duration = "1h"
 
         volume = await create_volume(
             session=session,
@@ -122,24 +107,19 @@ async def test_idle_duration_not_exceeded(self, test_db, session: AsyncSession):
             volume_provisioning_data=get_volume_provisioning_data(),
         )
 
-        # Set last_job_processed_at to 30 minutes ago (less than 1 hour)
         volume.last_job_processed_at = (
             datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(minutes=30)
         ).replace(tzinfo=None)
 
-        should_delete = await _should_delete_idle_volume(volume)
+        should_delete = _should_delete_idle_volume(volume)
         assert not should_delete
 
-    @pytest.mark.asyncio
-    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
     async def test_idle_duration_exceeded(self, test_db, session: AsyncSession):
-        """Test that volumes exceeding idle duration threshold are marked for deletion."""
         project = await create_project(session=session)
         user = await create_user(session=session)
 
-        # Create volume with idle_duration
         volume_config = get_volume_configuration(name="test-volume")
-        volume_config.idle_duration = "1h"  # 1 hour
+        volume_config.idle_duration = "1h"
 
         volume = await create_volume(
             session=session,
@@ -151,22 +131,17 @@ async def test_idle_duration_exceeded(self, test_db, session: AsyncSession):
             volume_provisioning_data=get_volume_provisioning_data(),
         )
 
-        # Set last_job_processed_at to 2 hours ago (more than 1 hour)
         volume.last_job_processed_at = (
             datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(hours=2)
         ).replace(tzinfo=None)
 
-        should_delete = await _should_delete_idle_volume(volume)
+        should_delete = _should_delete_idle_volume(volume)
         assert should_delete
 
-    @pytest.mark.asyncio
-    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
     async def test_volume_never_used_by_job(self, test_db, session: AsyncSession):
-        """Test idle duration calculation for volumes never used by jobs."""
         project = await create_project(session=session)
         user = await create_user(session=session)
 
-        # Create volume with old created_at time
         volume = await create_volume(
             session=session,
             project=project,
@@ -178,23 +153,17 @@ async def test_volume_never_used_by_job(self, test_db, session: AsyncSession):
             created_at=datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(hours=2),
         )
 
-        # last_job_processed_at is None, so it should use created_at
         volume.last_job_processed_at = None
 
         idle_duration = _get_volume_idle_duration(volume)
-        # Should be approximately 2 hours
-        assert idle_duration.total_seconds() >= 7000  # ~2 hours in seconds
+        assert idle_duration.total_seconds() >= 7000
 
-    @pytest.mark.asyncio
-    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
     async def test_process_idle_volumes_integration(self, test_db, session: AsyncSession):
-        """Integration test for the full process_idle_volumes function."""
         project = await create_project(session=session)
         user = await create_user(session=session)
 
-        # Create volume that should be deleted (exceeded idle duration)
         volume_config = get_volume_configuration(name="test-volume")
-        volume_config.idle_duration = "1h"  # 1 hour
+        volume_config.idle_duration = "1h"
 
         volume = await create_volume(
             session=session,
@@ -206,21 +175,17 @@ async def test_process_idle_volumes_integration(self, test_db, session: AsyncSes
             volume_provisioning_data=get_volume_provisioning_data(),
         )
 
-        # Set as idle for more than threshold
         volume.last_job_processed_at = (
             datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(hours=2)
         ).replace(tzinfo=None)
 
         await session.commit()
 
-        # Mock the delete_volumes function to avoid actual deletion
         with patch(
             "dstack._internal.server.background.tasks.process_idle_volumes.delete_volumes"
         ) as mock_delete:
             await process_idle_volumes()
 
-            # Should have called delete_volumes with the volume
             mock_delete.assert_called_once()
             call_args = mock_delete.call_args
-            # The function is called with (session, project, volume_names_list)
             assert call_args[0][2] == ["test-volume"]