From 1ea753f9dac7099177d89d83d32b306f959ebcf3 Mon Sep 17 00:00:00 2001 From: Sanghun Lee Date: Fri, 17 Jan 2025 19:32:31 +0900 Subject: [PATCH] fix(BA-461): Filter vfolder deletion info by vfolder status (#3446) Co-authored-by: octodog Backported-from: main (24.12) Backported-to: 24.09 Backport-of: 3446 --- changes/3446.fix.md | 1 + src/ai/backend/manager/models/group.py | 32 ++++++++++++++++---------- src/ai/backend/manager/models/user.py | 31 +++++++++++++++++-------- 3 files changed, 43 insertions(+), 21 deletions(-) create mode 100644 changes/3446.fix.md diff --git a/changes/3446.fix.md b/changes/3446.fix.md new file mode 100644 index 00000000000..5c8dffd2ce6 --- /dev/null +++ b/changes/3446.fix.md @@ -0,0 +1 @@ +Filter vfolders by status before initiating a vfolder deletion task diff --git a/src/ai/backend/manager/models/group.py b/src/ai/backend/manager/models/group.py index a7608cc8a80..32f70f85c9e 100644 --- a/src/ai/backend/manager/models/group.py +++ b/src/ai/backend/manager/models/group.py @@ -678,24 +678,32 @@ async def delete_vfolders( :return: number of deleted rows """ - from . import VFolderDeletionInfo, initiate_vfolder_deletion, vfolders - - query = ( - sa.select([vfolders.c.id, vfolders.c.host]) - .select_from(vfolders) - .where(vfolders.c.group == group_id) + from . import ( + VFolderDeletionInfo, + VFolderRow, + VFolderStatusSet, + initiate_vfolder_deletion, + vfolder_status_map, ) - async with engine.begin_session() as db_conn: - result = await db_conn.execute(query) - target_vfs = result.fetchall() - delete_query = sa.delete(vfolders).where(vfolders.c.group == group_id) - result = await db_conn.execute(delete_query) + + target_vfs: list[VFolderDeletionInfo] = [] + async with engine.begin_session() as db_session: + query = sa.select(VFolderRow).where( + sa.and_( + VFolderRow.group == group_id, + VFolderRow.status.in_(vfolder_status_map[VFolderStatusSet.DELETABLE]), + ) + ) + result = await db_session.scalars(query) + rows = cast(list[VFolderRow], result.fetchall()) + for vf in rows: + target_vfs.append(VFolderDeletionInfo(VFolderID.from_row(vf), vf.host)) storage_ptask_group = aiotools.PersistentTaskGroup() try: await initiate_vfolder_deletion( engine, - [VFolderDeletionInfo(VFolderID.from_row(vf), vf["host"]) for vf in target_vfs], + target_vfs, storage_manager, storage_ptask_group, ) diff --git a/src/ai/backend/manager/models/user.py b/src/ai/backend/manager/models/user.py index 2c7b207f52e..7a0e410d5f9 100644 --- a/src/ai/backend/manager/models/user.py +++ b/src/ai/backend/manager/models/user.py @@ -1107,24 +1107,37 @@ async def delete_vfolders( :return: number of deleted rows """ - from . import VFolderDeletionInfo, initiate_vfolder_deletion, vfolder_permissions, vfolders + from . import ( + VFolderDeletionInfo, + VFolderRow, + VFolderStatusSet, + initiate_vfolder_deletion, + vfolder_permissions, + vfolder_status_map, + ) - async with engine.begin_session() as conn: - await conn.execute( + target_vfs: list[VFolderDeletionInfo] = [] + async with engine.begin_session() as db_session: + await db_session.execute( vfolder_permissions.delete().where(vfolder_permissions.c.user == user_uuid), ) - result = await conn.execute( - sa.select([vfolders.c.id, vfolders.c.host, vfolders.c.quota_scope_id]) - .select_from(vfolders) - .where(vfolders.c.user == user_uuid), + result = await db_session.scalars( + sa.select(VFolderRow).where( + sa.and_( + VFolderRow.user == user_uuid, + VFolderRow.status.in_(vfolder_status_map[VFolderStatusSet.DELETABLE]), + ) + ), ) - target_vfs = result.fetchall() + rows = cast(list[VFolderRow], result.fetchall()) + for vf in rows: + target_vfs.append(VFolderDeletionInfo(VFolderID.from_row(vf), vf.host)) storage_ptask_group = aiotools.PersistentTaskGroup() try: await initiate_vfolder_deletion( engine, - [VFolderDeletionInfo(VFolderID.from_row(vf), vf["host"]) for vf in target_vfs], + target_vfs, storage_manager, storage_ptask_group, )