feat: Accept multiple paths in the vfolder mkdir API (#1803)

Co-authored-by: Joongi Kim <[email protected]>
Co-authored-by: SangwonYoon <[email protected]>

changes/1803.feature.md

@@ -0,0 +1 @@
+Add support for multi directory mkdir by fixing cli to accept multiple arguments and by adding list type annotation to accept multiple directories

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

@@ -447,8 +447,9 @@ def cp(filenames):
 
 
 @vfolder.command()
+@pass_ctx_obj
 @click.argument("name", type=str)
-@click.argument("path", type=str)
+@click.argument("paths", type=str, nargs=-1)
 @click.option(
     "-p",
     "--parents",
@@ -461,22 +462,30 @@ def cp(filenames):
     "--exist-ok",
     default=False,
     is_flag=True,
-    help="Skip an error caused by file not found",
+    help="Allow specifying already existing directories",
 )
-def mkdir(name, path, parents, exist_ok):
+def mkdir(
+    ctx: CLIContext,
+    name: str,
+    paths: list[str],
+    parents: bool,
+    exist_ok: bool,
+) -> None:
     """Create an empty directory in the virtual folder.
 
     \b
     NAME: Name of a virtual folder.
-    PATH: The name or path of directory. Parent directories are created automatically
-          if they do not exist.
+    PATHS: Relative directory paths to create in the vfolder.
+          Use '-p' option to auto-create parent directories.
+
+    Example: backend.ai vfolder mkdir my_vfolder "dir1" "dir2" "dir3"
     """
     with Session() as session:
         try:
-            session.VFolder(name).mkdir(path, parents=parents, exist_ok=exist_ok)
-            print_done("Done.")
+            results = session.VFolder(name).mkdir(paths, parents=parents, exist_ok=exist_ok)
+            ctx.output.print_result_set(results)
         except Exception as e:
-            print_error(e)
+            ctx.output.print_error(e)
             sys.exit(ExitCode.FAILURE)
 
 

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

@@ -1,7 +1,9 @@
+from __future__ import annotations
+
 import asyncio
 import uuid
 from pathlib import Path
-from typing import Any, Dict, Mapping, Optional, Sequence, Union
+from typing import Any, Mapping, Optional, Sequence, TypeAlias, TypeVar, Union
 
 import aiohttp
 import janus
@@ -20,6 +22,7 @@
 
 from ai.backend.client.output.fields import vfolder_fields
 from ai.backend.client.output.types import FieldSpec, PaginatedResult
+from ai.backend.common.types import ResultSet
 
 from ..compat import current_loop
 from ..config import DEFAULT_CHUNK_SIZE, MAX_INFLIGHT_CHUNKS
@@ -42,6 +45,9 @@
     vfolder_fields["status"],
 )
 
+T = TypeVar("T")
+list_: TypeAlias = list[T]
+
 
 class ResponseFailed(Exception):
     pass
@@ -322,7 +328,7 @@ async def _download_file(
         if_range: str | None = None
         file_unit = "bytes"
         file_mode = "wb"
-        file_req_hdrs: Dict[str, str] = {}
+        file_req_hdrs: dict[str, str] = {}
         try:
             async for session_attempt in AsyncRetrying(
                 wait=wait_exponential(multiplier=0.02, min=0.02, max=5.0),
@@ -514,7 +520,7 @@ async def _upload_recursively(
             if path.is_file():
                 file_list.append(path)
             else:
-                await self._mkdir(path.relative_to(base_path))
+                await self._mkdir([path.relative_to(base_path)])
                 dir_list.append(path)
         await self._upload_files(file_list, basedir, dst_dir, chunk_size, address_map)
         for dir in dir_list:
@@ -545,26 +551,27 @@ async def upload(
 
     async def _mkdir(
         self,
-        path: Union[str, Path],
+        path: str | Path | list_[str | Path],
         parents: Optional[bool] = False,
         exist_ok: Optional[bool] = False,
-    ) -> str:
+    ) -> ResultSet:
         rqst = Request("POST", "/folders/{}/mkdir".format(self.name))
         rqst.set_json({
             "path": path,
             "parents": parents,
             "exist_ok": exist_ok,
         })
         async with rqst.fetch() as resp:
-            return await resp.text()
+            reply = await resp.json()
+            return reply["results"]
 
     @api_function
     async def mkdir(
         self,
-        path: Union[str, Path],
+        path: str | Path | list_[str | Path],
         parents: Optional[bool] = False,
         exist_ok: Optional[bool] = False,
-    ) -> str:
+    ) -> ResultSet:
         return await self._mkdir(path, parents, exist_ok)
 
     @api_function

src/ai/backend/client/output/console.py

@@ -6,7 +6,8 @@
 from tabulate import tabulate
 
 from ai.backend.client.cli.pagination import echo_via_pager, get_preferred_page_size, tabulate_items
-from ai.backend.client.cli.pretty import print_error, print_fail
+from ai.backend.client.cli.pretty import print_done, print_error, print_fail
+from ai.backend.common.types import ResultSet
 
 from .types import BaseOutputHandler, FieldSpec, PaginatedResult
 
@@ -64,6 +65,27 @@ def print_items(
                 )
             )
 
+    def print_result_set(
+        self,
+        result_set: ResultSet,
+    ) -> None:
+        if result_set["success"]:
+            print_done("Successfully created:")
+            print(
+                tabulate(
+                    map(lambda item: [item["item"]], result_set["success"]),
+                    tablefmt="plain",
+                )
+            )
+        if result_set["failed"]:
+            print_fail("Failed to create:")
+            print(
+                tabulate(
+                    map(lambda item: [item["item"], item["msg"]], result_set["failed"]),
+                    tablefmt="plain",
+                )
+            )
+
     def print_list(
         self,
         items: Sequence[_Item],

src/ai/backend/client/output/json.py

@@ -3,6 +3,9 @@
 import json
 from typing import Any, Callable, Mapping, Optional, Sequence, TypeVar
 
+from ai.backend.client.exceptions import BackendAPIError
+from ai.backend.common.types import ResultSet
+
 from .types import BaseOutputHandler, FieldSpec, PaginatedResult
 
 _json_opts: Mapping[str, Any] = {"indent": 2}
@@ -71,6 +74,12 @@ def print_items(
             )
         )
 
+    def print_result_set(
+        self,
+        result_set: ResultSet,
+    ) -> None:
+        print(json.dumps(result_set))
+
     def print_list(
         self,
         items: Sequence[Mapping[str, Any]],
@@ -189,14 +198,30 @@ def print_error(
         self,
         error: Exception,
     ) -> None:
-        print(
-            json.dumps(
-                {
-                    "error": str(error),
-                },
-                **_json_opts,
-            )
-        )
+        match error:
+            case BackendAPIError():
+                print(
+                    json.dumps(
+                        {
+                            "error": error.data["title"],
+                            "api": {
+                                "status": error.status,
+                                "reason": error.reason,
+                                **error.data,
+                            },
+                        },
+                        **_json_opts,
+                    )
+                )
+            case _:
+                print(
+                    json.dumps(
+                        {
+                            "error": str(error),
+                        },
+                        **_json_opts,
+                    )
+                )
 
     def print_fail(
         self,

src/ai/backend/client/output/types.py

@@ -6,6 +6,8 @@
 
 import attr
 
+from ai.backend.common.types import ResultSet
+
 if TYPE_CHECKING:
     from ai.backend.client.cli.types import CLIContext
 
@@ -139,6 +141,13 @@ def print_items(
     ) -> None:
         raise NotImplementedError
 
+    @abstractmethod
+    def print_result_set(
+        self,
+        result_set: ResultSet,
+    ) -> None:
+        raise NotImplementedError
+
     @abstractmethod
     def print_list(
         self,

src/ai/backend/common/types.py

@@ -299,6 +299,16 @@ class CommitStatus(str, enum.Enum):
     ONGOING = "ongoing"
 
 
+class ItemResult(TypedDict):
+    msg: Optional[str]
+    item: Optional[str]
+
+
+class ResultSet(TypedDict):
+    success: list[ItemResult]
+    failed: list[ItemResult]
+
+
 class AbuseReportValue(str, enum.Enum):
     DETECTED = "detected"
     CLEANING = "cleaning"

src/ai/backend/common/validators.py

@@ -245,7 +245,11 @@ def __init__(
         self._relative_only = relative_only
 
     def check_and_return(self, value: Any) -> _PurePath:
-        p = _PurePath(value)
+        try:
+            p = _PurePath(value)
+        except (TypeError, ValueError):
+            self._failure("cannot parse value as a path", value=value)
+
         if self._relative_only and p.is_absolute():
             self._failure("expected relative path but the value is absolute", value=value)
         if self._base_path is not None:

src/ai/backend/manager/api/vfolder.py

@@ -1249,18 +1249,20 @@ async def update_vfolder_options(
 @vfolder_permission_required(VFolderPermission.READ_WRITE)
 @check_api_params(
     t.Dict({
-        t.Key("path"): t.String,
+        t.Key("path"): t.String | t.List(t.String),
         t.Key("parents", default=True): t.ToBool,
         t.Key("exist_ok", default=False): t.ToBool,
     })
 )
 async def mkdir(request: web.Request, params: Any, row: VFolderRow) -> web.Response:
+    if isinstance(params["path"], list) and len(params["path"]) > 50:
+        raise InvalidAPIParameters("Too many directories specified.")
     await ensure_vfolder_status(request, VFolderAccessStatus.UPDATABLE, request.match_info["name"])
     root_ctx: RootContext = request.app["_root.context"]
     folder_name = request.match_info["name"]
     access_key = request["keypair"]["access_key"]
     log.info(
-        "VFOLDER.MKDIR (email:{}, ak:{}, vf:{}, path:{})",
+        "VFOLDER.MKDIR (email:{}, ak:{}, vf:{}, paths:{})",
         request["user"]["email"],
         access_key,
         folder_name,
@@ -1278,9 +1280,14 @@ async def mkdir(request: web.Request, params: Any, row: VFolderRow) -> web.Respo
             "parents": params["parents"],
             "exist_ok": params["exist_ok"],
         },
-    ):
-        pass
-    return web.Response(status=201)
+    ) as (_, storage_resp):
+        storage_reply = await storage_resp.json()
+        match storage_resp.status:
+            case 200 | 207:
+                return web.json_response(storage_reply, status=storage_resp.status)
+            # 422 will be wrapped as VFolderOperationFailed by storage_manager
+            case _:
+                raise RuntimeError("should not reach here")
 
 
 @auth_required

src/ai/backend/manager/models/storage.py

@@ -191,7 +191,10 @@ async def request(
                         extra_data=None,
                     )
                 except VFolderOperationFailed as e:
-                    raise InvalidAPIParameters(e.extra_msg, e.extra_data)
+                    if client_resp.status // 100 == 5:
+                        raise InvalidAPIParameters(e.extra_msg, e.extra_data)
+                    # Raise as-is for semantic failures, not server errors.
+                    raise
             yield proxy_info.client_api_url, client_resp
 
 

src/ai/backend/manager/models/vfolder.py

@@ -768,7 +768,7 @@ async def prepare_vfolder_mounts(
                 params={
                     "volume": storage_manager.split_host(vfolder["host"])[1],
                     "vfid": str(VFolderID(vfolder["quota_scope_id"], vfolder["id"])),
-                    "relpath": str(user_scope.user_uuid.hex),
+                    "relpaths": [str(user_scope.user_uuid.hex)],
                     "exist_ok": True,
                 },
             ):