feat(BA-519): Update SDK to retrieve and use IDs for VFolder API operations instead of names (#3471)

changes/3471.feature.md

@@ -0,0 +1 @@
+Update SDK to retrieve and use IDs for VFolder API operations instead of names

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

@@ -31,7 +31,7 @@
 from ..request import Request
 from .base import BaseFunction, api_function
 
-__all__ = ("VFolder",)
+__all__ = ("VFolderByName",)
 
 _default_list_fields = (
     vfolder_fields["host"],
@@ -53,7 +53,10 @@ class ResponseFailed(Exception):
     pass
 
 
-class VFolder(BaseFunction):
+class VFolderByName(BaseFunction):
+    name: str
+    id: Optional[uuid.UUID] = None
+
     def __init__(self, name: str, id: Optional[uuid.UUID] = None):
         self.name = name
         self.id = id
@@ -242,35 +245,33 @@ async def list_allowed_types(cls):
 
     @api_function
     async def info(self):
-        rqst = Request("GET", "/folders/{0}".format(self.name))
+        await self.update_id_by_name()
+        rqst = Request("GET", "/folders/{0}".format(self.request_key))
         async with rqst.fetch() as resp:
             return await resp.json()
 
     @api_function
     async def delete(self):
-        rqst = Request("DELETE", "/folders/{0}".format(self.name))
+        await self.update_id_by_name()
+        rqst = Request("DELETE", "/folders/{0}".format(self.request_key))
         async with rqst.fetch():
             return {}
 
     @api_function
     async def purge(self) -> Mapping[str, Any]:
-        if self.id is None:
-            vfolder_id = await self._get_id_by_name()
-            self.id = vfolder_id
+        await self.update_id_by_name()
         rqst = Request("POST", "/folders/purge")
         rqst.set_json({
-            "id": self.id.hex,
+            "id": self.request_key,
         })
         async with rqst.fetch():
             return {}
 
     async def _restore(self) -> Mapping[str, Any]:
-        if self.id is None:
-            vfolder_id = await self._get_id_by_name()
-            self.id = vfolder_id
+        await self.update_id_by_name()
         rqst = Request("POST", "/folders/restore-from-trash-bin")
         rqst.set_json({
-            "id": self.id.hex,
+            "id": self.request_key,
         })
         async with rqst.fetch():
             return {}
@@ -285,19 +286,18 @@ async def restore(self):
 
     @api_function
     async def delete_trash(self) -> Mapping[str, Any]:
-        if self.id is None:
-            vfolder_id = await self._get_id_by_name()
-            self.id = vfolder_id
+        await self.update_id_by_name()
         rqst = Request("POST", "/folders/delete-from-trash-bin")
         rqst.set_json({
-            "id": self.id.hex,
+            "id": self.request_key,
         })
         async with rqst.fetch():
             return {}
 
     @api_function
     async def rename(self, new_name):
-        rqst = Request("POST", "/folders/{0}/rename".format(self.name))
+        await self.update_id_by_name()
+        rqst = Request("POST", "/folders/{0}/rename".format(self.request_key))
         rqst.set_json({
             "new_name": new_name,
         })
@@ -427,11 +427,13 @@ async def download(
         max_retries: int = 20,
     ) -> None:
         base_path = Path.cwd() if basedir is None else Path(basedir).resolve()
+        await self.update_id_by_name()
         for relpath in relative_paths:
             file_path = base_path / relpath
             if file_path.exists():
                 raise RuntimeError("The target file already exists", file_path.name)
-            rqst = Request("POST", "/folders/{}/request-download".format(self.name))
+
+            rqst = Request("POST", "/folders/{}/request-download".format(self.request_key))
             rqst.set_json({
                 "path": str(relpath),
             })
@@ -464,13 +466,14 @@ async def _upload_files(
         address_map: Optional[Mapping[str, str]] = None,
     ) -> None:
         base_path = Path.cwd() if basedir is None else Path(basedir).resolve()
+        await self.update_id_by_name()
         for file_path in file_paths:
             if file_path.is_dir():
                 raise BackendClientError(
                     f"Failed to upload {file_path}. Use recursive option to upload directories."
                 )
             file_size = Path(file_path).stat().st_size
-            rqst = Request("POST", "/folders/{}/request-upload".format(self.name))
+            rqst = Request("POST", "/folders/{}/request-upload".format(self.request_key))
             rqst.set_json({
                 "path": "{}".format(str(Path(file_path).relative_to(base_path))),
                 "size": int(file_size),
@@ -556,7 +559,8 @@ async def _mkdir(
         parents: Optional[bool] = False,
         exist_ok: Optional[bool] = False,
     ) -> ResultSet:
-        rqst = Request("POST", "/folders/{}/mkdir".format(self.name))
+        await self.update_id_by_name()
+        rqst = Request("POST", "/folders/{}/mkdir".format(self.request_key))
         rqst.set_json({
             "path": path,
             "parents": parents,
@@ -577,7 +581,8 @@ async def mkdir(
 
     @api_function
     async def rename_file(self, target_path: str, new_name: str):
-        rqst = Request("POST", "/folders/{}/rename-file".format(self.name))
+        await self.update_id_by_name()
+        rqst = Request("POST", "/folders/{}/rename-file".format(self.request_key))
         rqst.set_json({
             "target_path": target_path,
             "new_name": new_name,
@@ -587,7 +592,8 @@ async def rename_file(self, target_path: str, new_name: str):
 
     @api_function
     async def move_file(self, src_path: str, dst_path: str):
-        rqst = Request("POST", "/folders/{}/move-file".format(self.name))
+        await self.update_id_by_name()
+        rqst = Request("POST", "/folders/{}/move-file".format(self.request_key))
         rqst.set_json({
             "src": src_path,
             "dst": dst_path,
@@ -597,7 +603,8 @@ async def move_file(self, src_path: str, dst_path: str):
 
     @api_function
     async def delete_files(self, files: Sequence[Union[str, Path]], recursive: bool = False):
-        rqst = Request("DELETE", "/folders/{}/delete-files".format(self.name))
+        await self.update_id_by_name()
+        rqst = Request("DELETE", "/folders/{}/delete-files".format(self.request_key))
         rqst.set_json({
             "files": files,
             "recursive": recursive,
@@ -607,13 +614,17 @@ async def delete_files(self, files: Sequence[Union[str, Path]], recursive: bool
 
     @api_function
     async def list_files(self, path: Union[str, Path] = "."):
-        rqst = Request("GET", "/folders/{}/files".format(self.name), params={"path": str(path)})
+        await self.update_id_by_name()
+        rqst = Request(
+            "GET", "/folders/{}/files".format(self.request_key), params={"path": str(path)}
+        )
         async with rqst.fetch() as resp:
             return await resp.json()
 
     @api_function
     async def invite(self, perm: str, emails: Sequence[str]):
-        rqst = Request("POST", "/folders/{}/invite".format(self.name))
+        await self.update_id_by_name()
+        rqst = Request("POST", "/folders/{}/invite".format(self.request_key))
         rqst.set_json({
             "perm": perm,
             "user_ids": emails,
@@ -697,7 +708,8 @@ async def umount_host(cls, name: str, edit_fstab: bool = False):
 
     @api_function
     async def share(self, perm: str, emails: Sequence[str]):
-        rqst = Request("POST", "/folders/{}/share".format(self.name))
+        await self.update_id_by_name()
+        rqst = Request("POST", "/folders/{}/share".format(self.request_key))
         rqst.set_json({
             "permission": perm,
             "emails": emails,
@@ -707,7 +719,8 @@ async def share(self, perm: str, emails: Sequence[str]):
 
     @api_function
     async def unshare(self, emails: Sequence[str]):
-        rqst = Request("DELETE", "/folders/{}/unshare".format(self.name))
+        await self.update_id_by_name()
+        rqst = Request("DELETE", "/folders/{}/unshare".format(self.request_key))
         rqst.set_json({
             "emails": emails,
         })
@@ -716,7 +729,8 @@ async def unshare(self, emails: Sequence[str]):
 
     @api_function
     async def leave(self, shared_user_uuid=None):
-        rqst = Request("POST", "/folders/{}/leave".format(self.name))
+        await self.update_id_by_name()
+        rqst = Request("POST", "/folders/{}/leave".format(self.request_key))
         rqst.set_json({"shared_user_uuid": shared_user_uuid})
         async with rqst.fetch() as resp:
             return await resp.json()
@@ -729,7 +743,8 @@ async def clone(
         usage_mode: str = "general",
         permission: str = "rw",
     ):
-        rqst = Request("POST", "/folders/{}/clone".format(self.name))
+        await self.update_id_by_name()
+        rqst = Request("POST", "/folders/{}/clone".format(self.request_key))
         rqst.set_json({
             "target_name": target_name,
             "target_host": target_host,
@@ -743,7 +758,8 @@ async def clone(
     async def update_options(
         self, name: str, permission: Optional[str] = None, cloneable: Optional[bool] = None
     ):
-        rqst = Request("POST", "/folders/{}/update-options".format(self.name))
+        await self.update_id_by_name()
+        rqst = Request("POST", "/folders/{}/update-options".format(self.request_key))
         rqst.set_json({
             "cloneable": cloneable,
             "permission": permission,
@@ -787,3 +803,13 @@ async def change_vfolder_ownership(cls, vfolder: str, user_email: str):
         })
         async with rqst.fetch() as resp:
             return await resp.json()
+
+    async def update_id_by_name(self) -> None:
+        if self.id is None:
+            self.id = await self._get_id_by_name()
+
+    @property
+    def request_key(self) -> str:
+        if self.id is not None:
+            return self.id.hex
+        return self.name

src/ai/backend/client/session.py

@@ -316,7 +316,7 @@ def __init__(
         from .func.storage import Storage
         from .func.system import System
         from .func.user import User
-        from .func.vfolder import VFolder
+        from .func.vfolder import VFolderByName
 
         self.System = System
         self.Admin = Admin
@@ -337,7 +337,7 @@ def __init__(
         self.User = User
         self.ScalingGroup = ScalingGroup
         self.SessionTemplate = SessionTemplate
-        self.VFolder = VFolder
+        self.VFolder = VFolderByName
         self.Dotfile = Dotfile
         self.ServerLog = ServerLog
         self.Permission = Permission

tests/client/test_vfolder.py

@@ -1,5 +1,6 @@
 from typing import Mapping, Optional, Union
 from unittest import mock
+from uuid import UUID
 
 import pytest
 from aioresponses import aioresponses
@@ -77,7 +78,15 @@ def test_list_vfolders():
 def test_delete_vfolder():
     with Session() as session, aioresponses() as m:
         vfolder_name = "fake-vfolder-name"
-        m.delete(build_url(session.config, "/folders/{}".format(vfolder_name)), status=204)
+        source_vfolder_uuid: UUID = UUID("c59395cd-ac91-4cd3-a1b0-3d2568aa2d04")
+        m.get(
+            build_url(session.config, "/folders/_/id"),
+            status=200,
+            payload={"id": source_vfolder_uuid.hex},
+        )
+        m.delete(
+            build_url(session.config, "/folders/{}".format(source_vfolder_uuid.hex)), status=204
+        )
         resp = session.VFolder(vfolder_name).delete()
         assert resp == {}
 
@@ -94,8 +103,14 @@ def test_vfolder_get_info():
             "is_owner": True,
             "permission": "wd",
         }
+        source_vfolder_uuid: UUID = UUID("c59395cd-ac91-4cd3-a1b0-3d2568aa2d04")
         m.get(
-            build_url(session.config, "/folders/{}".format(vfolder_name)),
+            build_url(session.config, "/folders/_/id"),
+            status=200,
+            payload={"id": source_vfolder_uuid.hex},
+        )
+        m.get(
+            build_url(session.config, "/folders/{}".format(source_vfolder_uuid.hex)),
             status=200,
             payload=payload,
         )
@@ -107,8 +122,14 @@ def test_vfolder_delete_files():
     with Session() as session, aioresponses() as m:
         vfolder_name = "fake-vfolder-name"
         files = ["fake-file1", "fake-file2"]
+        source_vfolder_uuid: UUID = UUID("c59395cd-ac91-4cd3-a1b0-3d2568aa2d04")
+        m.get(
+            build_url(session.config, "/folders/_/id"),
+            status=200,
+            payload={"id": source_vfolder_uuid.hex},
+        )
         m.delete(
-            build_url(session.config, "/folders/{}/delete-files".format(vfolder_name)),
+            build_url(session.config, "/folders/{}/delete-files".format(source_vfolder_uuid.hex)),
             status=200,
             payload={},
         )
@@ -140,9 +161,17 @@ def test_vfolder_list_files():
             ],
             "folder_path": "/mnt/local/1f6bd27fde1248cabfb50306ea83fc0a",
         }
+        source_vfolder_uuid: UUID = UUID("c59395cd-ac91-4cd3-a1b0-3d2568aa2d04")
+        m.get(
+            build_url(session.config, "/folders/_/id"),
+            status=200,
+            payload={"id": source_vfolder_uuid.hex},
+        )
         m.get(
             build_url(
-                session.config, "/folders/{}/files".format(vfolder_name), params={"path": "."}
+                session.config,
+                "/folders/{}/files".format(source_vfolder_uuid.hex),
+                params={"path": "."},
             ),
             status=200,
             payload=payload,
@@ -156,8 +185,14 @@ def test_vfolder_invite():
         vfolder_name = "fake-vfolder-name"
         user_ids = ["[email protected]", "[email protected]"]
         payload = {"invited_ids": user_ids}
+        source_vfolder_uuid: UUID = UUID("c59395cd-ac91-4cd3-a1b0-3d2568aa2d04")
+        m.get(
+            build_url(session.config, "/folders/_/id"),
+            status=200,
+            payload={"id": source_vfolder_uuid.hex},
+        )
         m.post(
-            build_url(session.config, "/folders/{}/invite".format(vfolder_name)),
+            build_url(session.config, "/folders/{}/invite".format(source_vfolder_uuid.hex)),
             status=201,
             payload=payload,
         )
@@ -214,8 +249,14 @@ def test_vfolder_clone():
             "permission": "rw",
             "usage_mode": "general",
         }
+        source_vfolder_uuid: UUID = UUID("c59395cd-ac91-4cd3-a1b0-3d2568aa2d04")
+        m.get(
+            build_url(session.config, "/folders/_/id"),
+            status=200,
+            payload={"id": source_vfolder_uuid.hex},
+        )
         m.post(
-            build_url(session.config, "/folders/{}/clone".format(source_vfolder_name)),
+            build_url(session.config, "/folders/{}/clone".format(source_vfolder_uuid.hex)),
             status=201,
             payload=payload,
         )

tests/client/test_vfolder_tus.py

@@ -2,6 +2,7 @@
 from pathlib import Path
 from time import time
 from unittest import mock
+from uuid import UUID
 
 import pytest
 from aioresponses import aioresponses
@@ -159,10 +160,17 @@ async def test_vfolder_download(mocker):
                 ),
                 "url": storage_path,
             }
-
+            source_vfolder_uuid: UUID = UUID("c59395cd-ac91-4cd3-a1b0-3d2568aa2d04")
+            m.get(
+                build_url(session.config, "/folders/_/id"),
+                status=200,
+                payload={"id": source_vfolder_uuid.hex},
+            )
             # 1. Client to Manager throught Request
             m.post(
-                build_url(session.config, "/folders/{}/request-download".format(vfolder_name)),
+                build_url(
+                    session.config, "/folders/{}/request-download".format(source_vfolder_uuid.hex)
+                ),
                 payload=payload,
                 status=200,
                 headers={