feat: improve reliability & usability of VFolder removal API

changes/1884.feature.md

@@ -0,0 +1,2 @@
+* Enable to track VFolder deletion tasks.
+* Prevent concurrent tree-deletion using semaphore & file-lock.

configs/storage-proxy/sample.toml

@@ -24,6 +24,10 @@ event-loop = "uvloop"
 # Settings it zero means no limit.
 scandir-limit = 1000
 
+# The maximum number of directory(tree) delete concurrency that can be deleted from one storage-proxy process.
+# Should be larger than 1.
+directory-delete-concurrency = 20
+
 # The maximum allowed size of a single upload session.
 max-upload-size = "100g"
 

src/ai/backend/client/cli/vfolder.py

@@ -148,10 +148,10 @@ def create(name, host, group, host_path, usage_mode, permission, cloneable):
             sys.exit(ExitCode.FAILURE)
 
 
-@vfolder.command()
+@vfolder.command(aliases=["trash", "move-to-trash"])
 @click.argument("name", type=str)
 def delete(name):
-    """Delete the given virtual folder. The virtual folder will be under `delete-pending` status, which means trash-bin.
+    """Move the given virtual folder to trash-bin. The virtual folder will be under `delete-pending` status.
     This operation can be retracted by
     calling `restore()`.
 
@@ -160,7 +160,7 @@ def delete(name):
     """
     with Session() as session:
         try:
-            session.VFolder(name).delete()
+            session.VFolder(name).move_to_trash()
             print_done("Deleted.")
         except Exception as e:
             print_error(e)
@@ -186,40 +186,41 @@ def purge(name):
 @vfolder.command()
 @click.argument("name", type=str)
 def delete_trash(name):
-    """Delete the given virtual folder's real data. The virtual folder should be under `delete-pending` status, which means trash-bin.
+    """This command is deprecated, use `delete_forever`. Delete the given virtual folder's real data. The virtual folder should be under `delete-pending` status, which means trash-bin.
     This operation is irreversible!
 
     NAME: Name of a virtual folder.
     """
     with Session() as session:
         try:
-            session.VFolder(name).delete_trash()
-            print_done("Delete completed.")
+            session.VFolder(name).delete_forever()
+            print_done("Delete task started.")
         except Exception as e:
             print_error(e)
             sys.exit(ExitCode.FAILURE)
 
 
 @vfolder.command()
 @click.argument("name", type=str)
-def recover(name):
-    """Restore the given virtual folder from deleted status, Deprecated since 24.03.1; use `restore`
+def delete_forever(name):
+    """Delete the given virtual folder's real data. The virtual folder should be under `delete-pending` status, which means trash-bin.
+    This operation is irreversible!
 
     NAME: Name of a virtual folder.
     """
     with Session() as session:
         try:
-            session.VFolder(name).restore()
-            print_done("Restored.")
+            session.VFolder(name).delete_forever()
+            print_done("Delete task started.")
         except Exception as e:
             print_error(e)
             sys.exit(ExitCode.FAILURE)
 
 
-@vfolder.command()
+@vfolder.command(aliases=["recover"])
 @click.argument("name", type=str)
 def restore(name):
-    """Restore the given virtual folder from deleted status, from trash bin.
+    """Restore the given virtual folder from `delete-pending` status(trash-bin) to `ready`.
 
     NAME: Name of a virtual folder.
     """

src/ai/backend/client/func/vfolder.py

@@ -246,12 +246,25 @@ async def info(self):
         async with rqst.fetch() as resp:
             return await resp.json()
 
-    @api_function
-    async def delete(self):
-        rqst = Request("DELETE", "/folders/{0}".format(self.name))
+    async def _trash(self) -> Mapping[str, Any]:
+        if self.id is None:
+            vfolder_id = await self._get_id_by_name()
+            self.id = vfolder_id
+        rqst = Request("DELETE", "/folders")
+        rqst.set_json({
+            "id": self.id.hex,
+        })
         async with rqst.fetch():
             return {}
 
+    @api_function
+    async def delete(self):
+        return await self._trash()
+
+    @api_function
+    async def move_to_trash(self):
+        return await self._trash()
+
     @api_function
     async def purge(self) -> Mapping[str, Any]:
         if self.id is None:
@@ -268,7 +281,7 @@ async def _restore(self) -> Mapping[str, Any]:
         if self.id is None:
             vfolder_id = await self._get_id_by_name()
             self.id = vfolder_id
-        rqst = Request("POST", "/folders/restore-from-trash-bin")
+        rqst = Request("POST", "/folders/restore")
         rqst.set_json({
             "id": self.id.hex,
         })
@@ -283,18 +296,28 @@ async def recover(self):
     async def restore(self):
         return await self._restore()
 
-    @api_function
-    async def delete_trash(self) -> Mapping[str, Any]:
+    async def _delete_forever(self) -> Mapping[str, Any]:
         if self.id is None:
             vfolder_id = await self._get_id_by_name()
             self.id = vfolder_id
-        rqst = Request("POST", "/folders/delete-from-trash-bin")
+        rqst = Request("POST", "/folders/delete-forever")
         rqst.set_json({
             "id": self.id.hex,
         })
         async with rqst.fetch():
             return {}
 
+    @api_function
+    async def delete_trash(self) -> Mapping[str, Any]:
+        """
+        Deprecated, use `delete_forever()`.
+        """
+        return await self._delete_forever()
+
+    @api_function
+    async def delete_forever(self) -> Mapping[str, Any]:
+        return await self._delete_forever()
+
     @api_function
     async def rename(self, new_name):
         rqst = Request("POST", "/folders/{0}/rename".format(self.name))

src/ai/backend/common/bgtask.py

@@ -27,6 +27,7 @@
 from redis.asyncio.client import Pipeline
 
 from . import redis_helper
+from .defs import BackgroundTaskLogLevel as LogLevel
 from .events import (
     BgtaskCancelledEvent,
     BgtaskDoneEvent,
@@ -70,6 +71,7 @@ async def update(
         self,
         increment: Union[int, float] = 0,
         message: str | None = None,
+        log_level: LogLevel = LogLevel.INFO,
     ) -> None:
         self.current_progress += increment
         # keep the state as local variables because they might be changed
@@ -86,6 +88,7 @@ async def _pipe_builder(r: Redis) -> Pipeline:
                     "current": str(current),
                     "total": str(total),
                     "msg": message or "",
+                    "log_level": str(log_level),
                     "last_update": str(time.time()),
                 },
             )
@@ -99,6 +102,7 @@ async def _pipe_builder(r: Redis) -> Pipeline:
                 message=message,
                 current_progress=current,
                 total_progress=total,
+                log_level=log_level,
             ),
         )
 
@@ -154,10 +158,16 @@ async def push_bgtask_events(
                         "message": event.message,
                     }
                     match event:
-                        case BgtaskUpdatedEvent():
-                            body["current_progress"] = event.current_progress
-                            body["total_progress"] = event.total_progress
-                            await resp.send(json.dumps(body), event=event.name, retry=5)
+                        case BgtaskUpdatedEvent(
+                            name=name,
+                            current_progress=current_progress,
+                            total_progress=total_progress,
+                            log_level=log_level,
+                        ):
+                            body["current_progress"] = current_progress
+                            body["total_progress"] = total_progress
+                            body["log_level"] = str(log_level)
+                            await resp.send(json.dumps(body), event=name, retry=5)
                         case BgtaskDoneEvent():
                             if extra_data:
                                 body.update(extra_data)
@@ -232,6 +242,7 @@ async def start(
         self,
         func: BackgroundTask,
         name: str | None = None,
+        total_progress: int = 0,
         **kwargs,
     ) -> uuid.UUID:
         task_id = uuid.uuid4()
@@ -246,7 +257,7 @@ async def _pipe_builder(r: Redis) -> Pipeline:
                 mapping={
                     "status": "started",
                     "current": "0",
-                    "total": "0",
+                    "total": total_progress,
                     "msg": "",
                     "started_at": now,
                     "last_update": now,
@@ -257,7 +268,9 @@ async def _pipe_builder(r: Redis) -> Pipeline:
 
         await redis_helper.execute(redis_producer, _pipe_builder)
 
-        task = asyncio.create_task(self._wrapper_task(func, task_id, name, **kwargs))
+        task = asyncio.create_task(
+            self._wrapper_task(func, task_id, name, total_progress, **kwargs)
+        )
         self.ongoing_tasks.add(task)
         return task_id
 
@@ -266,24 +279,30 @@ async def _wrapper_task(
         func: BackgroundTask,
         task_id: uuid.UUID,
         task_name: str | None,
+        task_total_progress: int = 0,
         **kwargs,
     ) -> None:
         task_status: TaskStatus = "bgtask_started"
-        reporter = ProgressReporter(self.event_producer, task_id)
+        reporter = ProgressReporter(
+            self.event_producer, task_id, total_progress=task_total_progress
+        )
         message = ""
         event_cls: Type[BgtaskDoneEvent] | Type[BgtaskCancelledEvent] | Type[BgtaskFailedEvent] = (
             BgtaskDoneEvent
         )
+        log_level = LogLevel.INFO
         try:
             message = await func(reporter, **kwargs) or ""
             task_status = "bgtask_done"
         except asyncio.CancelledError:
             task_status = "bgtask_cancelled"
             event_cls = BgtaskCancelledEvent
+            log_level = LogLevel.WARNING
         except Exception as e:
             task_status = "bgtask_failed"
             event_cls = BgtaskFailedEvent
             message = repr(e)
+            log_level = LogLevel.ERROR
             log.exception("Task {} ({}): unhandled error", task_id, task_name)
         finally:
             redis_producer = self.event_producer.redis_client
@@ -297,6 +316,7 @@ async def _pipe_builder(r: Redis):
                         "status": task_status.removeprefix("bgtask_"),
                         "msg": message,
                         "last_update": str(time.time()),
+                        "log_level": str(log_level),
                     },
                 )
                 await pipe.expire(tracker_key, MAX_BGTASK_ARCHIVE_PERIOD)

src/ai/backend/common/defs.py

@@ -1,3 +1,4 @@
+from enum import StrEnum, auto
 from typing import Final
 
 # Redis database IDs depending on purposes
@@ -10,3 +11,10 @@
 
 
 DEFAULT_FILE_IO_TIMEOUT: Final = 10
+
+
+class BackgroundTaskLogLevel(StrEnum):
+    INFO = auto()
+    WARNING = auto()
+    ERROR = auto()
+    DEBUG = auto()

src/ai/backend/common/events.py

@@ -34,6 +34,7 @@
 from redis.asyncio import ConnectionPool
 
 from . import msgpack, redis_helper
+from .defs import BackgroundTaskLogLevel
 from .logging import BraceStyleAdapter
 from .types import (
     AgentId,
@@ -555,13 +556,15 @@ class BgtaskUpdatedEvent(AbstractEvent):
     current_progress: float = attrs.field()
     total_progress: float = attrs.field()
     message: Optional[str] = attrs.field(default=None)
+    log_level: BackgroundTaskLogLevel = attrs.field(default=BackgroundTaskLogLevel.INFO)
 
     def serialize(self) -> tuple:
         return (
             str(self.task_id),
             self.current_progress,
             self.total_progress,
             self.message,
+            str(self.log_level),
         )
 
     @classmethod
@@ -571,6 +574,7 @@ def deserialize(cls, value: tuple):
             value[1],
             value[2],
             value[3],
+            BackgroundTaskLogLevel(value[4]),
         )