Removed unnecessary changes and refactor implementations

Author: haydnli-shopify

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

@@ -31,86 +31,74 @@ async def process_idle_volumes():
                 .limit(10)
                 .with_for_update(skip_locked=True)
             )
-            volume_models = list(res.unique().scalars().all())
-            if not volume_models:
+            volumes = list(res.unique().scalars().all())
+            if not volumes:
                 return
-            for volume_model in volume_models:
-                await session.refresh(volume_model, ["project", "attachments"])
-                lockset.add(volume_model.id)
+            for volume in volumes:
+                await session.refresh(volume, ["project", "attachments"])
+                lockset.add(volume.id)
 
         try:
-            volumes_to_delete = []
-            for volume_model in volume_models:
-                if _should_delete_idle_volume(volume_model):
-                    volumes_to_delete.append(volume_model)
+            to_delete = []
+            for volume in volumes:
+                if _should_delete_volume(volume):
+                    to_delete.append(volume)
 
-            if volumes_to_delete:
-                await _delete_idle_volumes(session, volumes_to_delete)
+            if to_delete:
+                await _delete_volumes(session, to_delete)
 
         finally:
-            for volume_model in volume_models:
-                lockset.discard(volume_model.id)
+            for volume in volumes:
+                lockset.discard(volume.id)
 
 
-def _should_delete_idle_volume(volume_model: VolumeModel) -> bool:
-    configuration = get_volume_configuration(volume_model)
+def _should_delete_volume(volume: VolumeModel) -> bool:
+    config = get_volume_configuration(volume)
 
-    if configuration.idle_duration is None:
+    if not config.idle_duration:
         return False
 
-    if isinstance(configuration.idle_duration, int) and configuration.idle_duration < 0:
+    if isinstance(config.idle_duration, int) and config.idle_duration < 0:
         return False
 
-    idle_duration_seconds = parse_duration(configuration.idle_duration)
-    if idle_duration_seconds is None or idle_duration_seconds <= 0:
+    duration_seconds = parse_duration(config.idle_duration)
+    if not duration_seconds or duration_seconds <= 0:
         return False
 
-    if len(volume_model.attachments) > 0:
-        logger.debug("Volume %s is still attached to instances, not deleting", volume_model.name)
+    if volume.attachments:
         return False
 
-    idle_duration = _get_volume_idle_duration(volume_model)
-    idle_threshold = datetime.timedelta(seconds=idle_duration_seconds)
+    idle_time = _get_idle_time(volume)
+    threshold = datetime.timedelta(seconds=duration_seconds)
 
-    if idle_duration > idle_threshold:
+    if idle_time > threshold:
         logger.info(
-            "Volume %s idle duration expired: idle time %.1f hours, threshold %.1f hours. Marking for deletion",
-            volume_model.name,
-            idle_duration.total_seconds() / 3600,
-            idle_threshold.total_seconds() / 3600,
+            "Deleting idle volume %s (idle %.1fh)", volume.name, idle_time.total_seconds() / 3600
         )
         return True
 
     return False
 
 
-def _get_volume_idle_duration(volume_model: VolumeModel) -> datetime.timedelta:
-    reference_time = volume_model.created_at
-    if volume_model.last_job_processed_at is not None:
-        reference_time = volume_model.last_job_processed_at
+def _get_idle_time(volume: VolumeModel) -> datetime.timedelta:
+    last_used = volume.last_job_processed_at or volume.created_at
+    last_used_utc = last_used.replace(tzinfo=datetime.timezone.utc)
+    now = get_current_datetime()
 
-    reference_time_utc = reference_time.replace(tzinfo=datetime.timezone.utc)
-    current_time = get_current_datetime()
+    idle_time = now - last_used_utc
+    return max(idle_time, datetime.timedelta(0))
 
-    idle_duration = current_time - reference_time_utc
 
-    if idle_duration.total_seconds() < 0:
-        return datetime.timedelta(0)
+async def _delete_volumes(session: AsyncSession, volumes: List[VolumeModel]):
+    by_project = {}
+    for volume in volumes:
+        project = volume.project
+        if project not in by_project:
+            by_project[project] = []
+        by_project[project].append(volume.name)
 
-    return idle_duration
-
-
-async def _delete_idle_volumes(session: AsyncSession, volume_models: List[VolumeModel]):
-    volumes_by_project = {}
-    for volume_model in volume_models:
-        project = volume_model.project
-        if project not in volumes_by_project:
-            volumes_by_project[project] = []
-        volumes_by_project[project].append(volume_model.name)
-
-    for project, volume_names in volumes_by_project.items():
-        logger.info("Deleting idle volumes for project %s: %s", project.name, volume_names)
+    for project, names in by_project.items():
         try:
-            await delete_volumes(session, project, volume_names)
+            await delete_volumes(session, project, names)
         except Exception as e:
-            logger.error("Failed to delete idle volumes for project %s: %s", project.name, str(e))
+            logger.error("Failed to delete volumes for project %s: %s", project.name, str(e))

src/dstack/_internal/server/services/jobs/__init__.py

@@ -281,15 +281,18 @@ async def process_terminating_job(
     job_model.instance_id = None
     instance_model.last_job_processed_at = common.get_current_datetime()
 
-    if jrd is not None and jrd.volume_names is not None:
-        volume_names = jrd.volume_names
-    else:
-        volume_names = [va.volume.name for va in instance_model.volume_attachments]
-    volume_models = await list_project_volume_models(
-        session=session, project=instance_model.project, names=volume_names
+    # Update volume timestamps
+    volume_names = (
+        jrd.volume_names
+        if jrd and jrd.volume_names
+        else [va.volume.name for va in instance_model.volume_attachments]
     )
-    for volume_model in volume_models:
-        volume_model.last_job_processed_at = common.get_current_datetime()
+    if volume_names:
+        volumes = await list_project_volume_models(
+            session=session, project=instance_model.project, names=volume_names
+        )
+        for volume in volumes:
+            volume.last_job_processed_at = common.get_current_datetime()
 
     logger.info(
         "%s: instance '%s' has been released, new status is %s",

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

@@ -7,8 +7,8 @@
 from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.volumes import VolumeStatus
 from dstack._internal.server.background.tasks.process_idle_volumes import (
-    _get_volume_idle_duration,
-    _should_delete_idle_volume,
+    _get_idle_time,
+    _should_delete_volume,
     process_idle_volumes,
 )
 from dstack._internal.server.models import VolumeAttachmentModel
@@ -25,7 +25,7 @@
 @pytest.mark.asyncio
 @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
 class TestProcessIdleVolumes:
-    async def test_no_idle_duration_configured(self, test_db, session: AsyncSession):
+    async def test_no_idle_duration(self, test_db, session: AsyncSession):
         project = await create_project(session=session)
         user = await create_user(session=session)
 
@@ -39,106 +39,82 @@ async def test_no_idle_duration_configured(self, test_db, session: AsyncSession)
             volume_provisioning_data=get_volume_provisioning_data(),
         )
 
-        should_delete = _should_delete_idle_volume(volume)
-        assert not should_delete
+        assert not _should_delete_volume(volume)
 
     async def test_idle_duration_disabled(self, test_db, session: AsyncSession):
         project = await create_project(session=session)
         user = await create_user(session=session)
 
-        volume_config = get_volume_configuration(name="test-volume")
-        volume_config.idle_duration = -1
+        config = get_volume_configuration(name="test-volume")
+        config.idle_duration = -1
 
         volume = await create_volume(
             session=session,
             project=project,
             user=user,
             status=VolumeStatus.ACTIVE,
             backend=BackendType.AWS,
-            configuration=volume_config,
+            configuration=config,
             volume_provisioning_data=get_volume_provisioning_data(),
         )
 
-        should_delete = _should_delete_idle_volume(volume)
-        assert not should_delete
+        assert not _should_delete_volume(volume)
 
-    async def test_volume_still_attached(self, test_db, session: AsyncSession):
+    async def test_volume_attached(self, test_db, session: AsyncSession):
         project = await create_project(session=session)
         user = await create_user(session=session)
 
-        volume_config = get_volume_configuration(name="test-volume")
-        volume_config.idle_duration = "1h"
+        config = get_volume_configuration(name="test-volume")
+        config.idle_duration = "1h"
 
         volume = await create_volume(
             session=session,
             project=project,
             user=user,
             status=VolumeStatus.ACTIVE,
             backend=BackendType.AWS,
-            configuration=volume_config,
+            configuration=config,
             volume_provisioning_data=get_volume_provisioning_data(),
         )
 
         instance = await create_instance(session=session, project=project)
-        attachment = VolumeAttachmentModel(
-            volume_id=volume.id,
-            instance_id=instance.id,
+        volume.attachments.append(
+            VolumeAttachmentModel(volume_id=volume.id, instance_id=instance.id)
         )
-        volume.attachments.append(attachment)
         await session.commit()
 
-        should_delete = _should_delete_idle_volume(volume)
-        assert not should_delete
+        assert not _should_delete_volume(volume)
 
-    async def test_idle_duration_not_exceeded(self, test_db, session: AsyncSession):
+    async def test_idle_duration_threshold(self, test_db, session: AsyncSession):
         project = await create_project(session=session)
         user = await create_user(session=session)
 
-        volume_config = get_volume_configuration(name="test-volume")
-        volume_config.idle_duration = "1h"
+        config = get_volume_configuration(name="test-volume")
+        config.idle_duration = "1h"
 
         volume = await create_volume(
             session=session,
             project=project,
             user=user,
             status=VolumeStatus.ACTIVE,
             backend=BackendType.AWS,
-            configuration=volume_config,
+            configuration=config,
             volume_provisioning_data=get_volume_provisioning_data(),
         )
 
+        # Not exceeded - 30 minutes ago
         volume.last_job_processed_at = (
             datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(minutes=30)
         ).replace(tzinfo=None)
+        assert not _should_delete_volume(volume)
 
-        should_delete = _should_delete_idle_volume(volume)
-        assert not should_delete
-
-    async def test_idle_duration_exceeded(self, test_db, session: AsyncSession):
-        project = await create_project(session=session)
-        user = await create_user(session=session)
-
-        volume_config = get_volume_configuration(name="test-volume")
-        volume_config.idle_duration = "1h"
-
-        volume = await create_volume(
-            session=session,
-            project=project,
-            user=user,
-            status=VolumeStatus.ACTIVE,
-            backend=BackendType.AWS,
-            configuration=volume_config,
-            volume_provisioning_data=get_volume_provisioning_data(),
-        )
-
+        # Exceeded - 2 hours ago
         volume.last_job_processed_at = (
             datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(hours=2)
         ).replace(tzinfo=None)
+        assert _should_delete_volume(volume)
 
-        should_delete = _should_delete_idle_volume(volume)
-        assert should_delete
-
-    async def test_volume_never_used_by_job(self, test_db, session: AsyncSession):
+    async def test_never_used_volume(self, test_db, session: AsyncSession):
         project = await create_project(session=session)
         user = await create_user(session=session)
 
@@ -154,38 +130,34 @@ async def test_volume_never_used_by_job(self, test_db, session: AsyncSession):
         )
 
         volume.last_job_processed_at = None
+        idle_time = _get_idle_time(volume)
+        assert idle_time.total_seconds() >= 7000
 
-        idle_duration = _get_volume_idle_duration(volume)
-        assert idle_duration.total_seconds() >= 7000
-
-    async def test_process_idle_volumes_integration(self, test_db, session: AsyncSession):
+    async def test_integration(self, test_db, session: AsyncSession):
         project = await create_project(session=session)
         user = await create_user(session=session)
 
-        volume_config = get_volume_configuration(name="test-volume")
-        volume_config.idle_duration = "1h"
+        config = get_volume_configuration(name="test-volume")
+        config.idle_duration = "1h"
 
         volume = await create_volume(
             session=session,
             project=project,
             user=user,
             status=VolumeStatus.ACTIVE,
             backend=BackendType.AWS,
-            configuration=volume_config,
+            configuration=config,
             volume_provisioning_data=get_volume_provisioning_data(),
         )
 
         volume.last_job_processed_at = (
             datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(hours=2)
         ).replace(tzinfo=None)
-
         await session.commit()
 
         with patch(
             "dstack._internal.server.background.tasks.process_idle_volumes.delete_volumes"
-        ) as mock_delete:
+        ) as mock:
             await process_idle_volumes()
-
-            mock_delete.assert_called_once()
-            call_args = mock_delete.call_args
-            assert call_args[0][2] == ["test-volume"]
+            mock.assert_called_once()
+            assert mock.call_args[0][2] == ["test-volume"]