|
| 1 | +import datetime |
| 2 | +from typing import List |
| 3 | + |
| 4 | +from sqlalchemy import select |
| 5 | +from sqlalchemy.ext.asyncio import AsyncSession |
| 6 | +from sqlalchemy.orm import joinedload |
| 7 | + |
| 8 | +from dstack._internal.core.backends.base.compute import ComputeWithVolumeSupport |
| 9 | +from dstack._internal.core.errors import BackendNotAvailable |
| 10 | +from dstack._internal.core.models.profiles import parse_duration |
| 11 | +from dstack._internal.core.models.volumes import VolumeStatus |
| 12 | +from dstack._internal.server.db import get_db, get_session_ctx |
| 13 | +from dstack._internal.server.models import ProjectModel, VolumeModel |
| 14 | +from dstack._internal.server.services import backends as backends_services |
| 15 | +from dstack._internal.server.services.locking import get_locker |
| 16 | +from dstack._internal.server.services.volumes import ( |
| 17 | + get_volume_configuration, |
| 18 | + volume_model_to_volume, |
| 19 | +) |
| 20 | +from dstack._internal.utils import common |
| 21 | +from dstack._internal.utils.common import get_current_datetime |
| 22 | +from dstack._internal.utils.logging import get_logger |
| 23 | + |
| 24 | +logger = get_logger(__name__) |
| 25 | + |
| 26 | + |
| 27 | +async def process_idle_volumes(): |
| 28 | + lock, lockset = get_locker(get_db().dialect_name).get_lockset(VolumeModel.__tablename__) |
| 29 | + async with get_session_ctx() as session: |
| 30 | + async with lock: |
| 31 | + res = await session.execute( |
| 32 | + select(VolumeModel.id) |
| 33 | + .where( |
| 34 | + VolumeModel.status == VolumeStatus.ACTIVE, |
| 35 | + VolumeModel.deleted == False, |
| 36 | + VolumeModel.id.not_in(lockset), |
| 37 | + ) |
| 38 | + .order_by(VolumeModel.last_processed_at.asc()) |
| 39 | + .limit(10) |
| 40 | + .with_for_update(skip_locked=True, key_share=True) |
| 41 | + ) |
| 42 | + volume_ids = list(res.scalars().all()) |
| 43 | + if not volume_ids: |
| 44 | + return |
| 45 | + for volume_id in volume_ids: |
| 46 | + lockset.add(volume_id) |
| 47 | + |
| 48 | + res = await session.execute( |
| 49 | + select(VolumeModel) |
| 50 | + .where(VolumeModel.id.in_(volume_ids)) |
| 51 | + .options(joinedload(VolumeModel.project).joinedload(ProjectModel.backends)) |
| 52 | + .options(joinedload(VolumeModel.user)) |
| 53 | + .options(joinedload(VolumeModel.attachments)) |
| 54 | + .execution_options(populate_existing=True) |
| 55 | + ) |
| 56 | + volume_models = list(res.unique().scalars().all()) |
| 57 | + try: |
| 58 | + volumes_to_delete = [v for v in volume_models if _should_delete_volume(v)] |
| 59 | + if not volumes_to_delete: |
| 60 | + return |
| 61 | + await _delete_idle_volumes(session, volumes_to_delete) |
| 62 | + finally: |
| 63 | + lockset.difference_update(volume_ids) |
| 64 | + |
| 65 | + |
| 66 | +def _should_delete_volume(volume: VolumeModel) -> bool: |
| 67 | + if volume.attachments: |
| 68 | + return False |
| 69 | + |
| 70 | + config = get_volume_configuration(volume) |
| 71 | + if not config.auto_cleanup_duration: |
| 72 | + return False |
| 73 | + |
| 74 | + duration_seconds = parse_duration(config.auto_cleanup_duration) |
| 75 | + if not duration_seconds or duration_seconds <= 0: |
| 76 | + return False |
| 77 | + |
| 78 | + idle_time = _get_idle_time(volume) |
| 79 | + threshold = datetime.timedelta(seconds=duration_seconds) |
| 80 | + return idle_time > threshold |
| 81 | + |
| 82 | + |
| 83 | +def _get_idle_time(volume: VolumeModel) -> datetime.timedelta: |
| 84 | + last_used = volume.last_job_processed_at or volume.created_at |
| 85 | + last_used_utc = last_used.replace(tzinfo=datetime.timezone.utc) |
| 86 | + idle_time = get_current_datetime() - last_used_utc |
| 87 | + return max(idle_time, datetime.timedelta(0)) |
| 88 | + |
| 89 | + |
| 90 | +async def _delete_idle_volumes(session: AsyncSession, volumes: List[VolumeModel]): |
| 91 | + # Note: Multiple volumes are deleted in the same transaction, |
| 92 | + # so long deletion of one volume may block processing other volumes. |
| 93 | + for volume_model in volumes: |
| 94 | + logger.info("Deleting idle volume %s", volume_model.name) |
| 95 | + try: |
| 96 | + await _delete_idle_volume(session, volume_model) |
| 97 | + except Exception: |
| 98 | + logger.exception("Error when deleting idle volume %s", volume_model.name) |
| 99 | + |
| 100 | + volume_model.deleted = True |
| 101 | + volume_model.deleted_at = get_current_datetime() |
| 102 | + |
| 103 | + logger.info("Deleted idle volume %s", volume_model.name) |
| 104 | + |
| 105 | + await session.commit() |
| 106 | + |
| 107 | + |
| 108 | +async def _delete_idle_volume(session: AsyncSession, volume_model: VolumeModel): |
| 109 | + volume = volume_model_to_volume(volume_model) |
| 110 | + |
| 111 | + if volume.provisioning_data is None: |
| 112 | + logger.error( |
| 113 | + f"Failed to delete volume {volume_model.name}. volume.provisioning_data is None." |
| 114 | + ) |
| 115 | + return |
| 116 | + |
| 117 | + if volume.provisioning_data.backend is None: |
| 118 | + logger.error( |
| 119 | + f"Failed to delete volume {volume_model.name}. volume.provisioning_data.backend is None." |
| 120 | + ) |
| 121 | + return |
| 122 | + |
| 123 | + try: |
| 124 | + backend = await backends_services.get_project_backend_by_type_or_error( |
| 125 | + project=volume_model.project, |
| 126 | + backend_type=volume.provisioning_data.backend, |
| 127 | + ) |
| 128 | + except BackendNotAvailable: |
| 129 | + logger.error( |
| 130 | + f"Failed to delete volume {volume_model.name}. Backend {volume.configuration.backend} not available." |
| 131 | + ) |
| 132 | + return |
| 133 | + |
| 134 | + compute = backend.compute() |
| 135 | + assert isinstance(compute, ComputeWithVolumeSupport) |
| 136 | + await common.run_async( |
| 137 | + compute.delete_volume, |
| 138 | + volume=volume, |
| 139 | + ) |
0 commit comments