feat(BA-532): Add configurable directory permission for vfolders (#3510)

Author: fregataa

changes/3510.feature.md

@@ -0,0 +1 @@
+Add configurable directory permission for vfolders to support mount vfolders on customized UID/GID containers

src/ai/backend/common/defs.py

@@ -41,3 +41,6 @@
 RESERVED_VFOLDER_PATTERNS = [re.compile(x) for x in _RESERVED_VFOLDER_PATTERNS]
 API_VFOLDER_LENGTH_LIMIT: Final[int] = 64
 MODEL_VFOLDER_LENGTH_LIMIT: Final[int] = 128
+
+DEFAULT_VFOLDER_PERMISSION_MODE: Final[int] = 0o755
+VFOLDER_GROUP_PERMISSION_MODE: Final[int] = 0o775

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

@@ -23,6 +23,7 @@
     List,
     Mapping,
     MutableMapping,
+    Optional,
     ParamSpec,
     Sequence,
     Tuple,
@@ -48,6 +49,7 @@
 from ai.backend.common import msgpack, redis_helper
 from ai.backend.common import typed_validators as tv
 from ai.backend.common import validators as tx
+from ai.backend.common.defs import VFOLDER_GROUP_PERMISSION_MODE
 from ai.backend.common.types import (
     QuotaScopeID,
     QuotaScopeType,
@@ -430,6 +432,7 @@ async def create(request: web.Request, params: CreateRequestModel) -> web.Respon
     group_type: ProjectType | None = None
     max_vfolder_count: int
     max_quota_scope_size: int
+    container_uid: Optional[int] = None
 
     async with root_ctx.db.begin_session() as sess:
         match group_id_or_name:
@@ -491,9 +494,14 @@ async def create(request: web.Request, params: CreateRequestModel) -> web.Respon
                     cast(int, user_row.resource_policy_row.max_vfolder_count),
                     cast(int, user_row.resource_policy_row.max_quota_scope_size),
                 )
+                container_uid = cast(Optional[int], user_row.container_uid)
             case _:
                 raise GroupNotFound(extra_data=group_id_or_name)
 
+        vfolder_permission_mode = (
+            VFOLDER_GROUP_PERMISSION_MODE if container_uid is not None else None
+        )
+
         # Check if group exists when it's given a non-empty value.
         if group_id_or_name and group_uuid is None:
             raise GroupNotFound(extra_data=group_id_or_name)
@@ -615,6 +623,7 @@ async def create(request: web.Request, params: CreateRequestModel) -> web.Respon
                         "volume": root_ctx.storage_manager.split_host(folder_host)[1],
                         "vfid": str(vfid),
                         "options": options,
+                        "mode": vfolder_permission_mode,
                     },
                 ):
                     pass

src/ai/backend/storage/abc.py

@@ -15,6 +15,7 @@
     final,
 )
 
+from ai.backend.common.defs import DEFAULT_VFOLDER_PERMISSION_MODE
 from ai.backend.common.etcd import AsyncEtcd
 from ai.backend.common.events import EventDispatcher, EventProducer
 from ai.backend.common.types import BinarySize, HardwareMetadata, QuotaScopeID
@@ -265,7 +266,8 @@ async def get_hwinfo(self) -> HardwareMetadata:
     async def create_vfolder(
         self,
         vfid: VFolderID,
-        exist_ok=False,
+        exist_ok: bool = False,
+        mode: int = DEFAULT_VFOLDER_PERMISSION_MODE,
     ) -> None:
         raise NotImplementedError
 

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

@@ -20,6 +20,7 @@
     Callable,
     Iterator,
     NotRequired,
+    Optional,
     TypedDict,
     cast,
 )
@@ -31,6 +32,7 @@
 from aiohttp import hdrs, web
 
 from ai.backend.common import validators as tx
+from ai.backend.common.defs import DEFAULT_VFOLDER_PERMISSION_MODE
 from ai.backend.common.events import (
     DoVolumeMountEvent,
     DoVolumeUnmountEvent,
@@ -340,6 +342,7 @@ async def create_vfolder(request: web.Request) -> web.Response:
     class Params(TypedDict):
         volume: str
         vfid: VFolderID
+        mode: Optional[int]
         options: dict[str, Any] | None  # deprecated
 
     async with cast(
@@ -350,6 +353,7 @@ class Params(TypedDict):
                 {
                     t.Key("volume"): t.String(),
                     t.Key("vfid"): tx.VFolderID(),
+                    t.Key("mode", default=DEFAULT_VFOLDER_PERMISSION_MODE): t.Null,
                     t.Key("options", default=None): t.Null | t.Dict().allow_extra("*"),
                 },
             ),
@@ -358,9 +362,10 @@ class Params(TypedDict):
         await log_manager_api_entry(log, "create_vfolder", params)
         assert params["vfid"].quota_scope_id is not None
         ctx: RootContext = request.app["ctx"]
+        perm_mode = cast(int, params["mode"])
         async with ctx.get_volume(params["volume"]) as volume:
             try:
-                await volume.create_vfolder(params["vfid"])
+                await volume.create_vfolder(params["vfid"], mode=perm_mode)
             except QuotaScopeNotFoundError:
                 assert params["vfid"].quota_scope_id
                 if initial_max_size_for_quota_scope := (params["options"] or {}).get(
@@ -373,7 +378,7 @@ class Params(TypedDict):
                     params["vfid"].quota_scope_id, options=options
                 )
                 try:
-                    await volume.create_vfolder(params["vfid"])
+                    await volume.create_vfolder(params["vfid"], mode=perm_mode)
                 except QuotaScopeNotFoundError:
                     raise ExternalError("Failed to create vfolder due to quota scope not found.")
             return web.Response(status=204)

src/ai/backend/storage/vfs/__init__.py

@@ -16,6 +16,7 @@
 import janus
 import trafaret as t
 
+from ai.backend.common.defs import DEFAULT_VFOLDER_PERMISSION_MODE
 from ai.backend.common.types import BinarySize, HardwareMetadata, QuotaScopeID
 from ai.backend.logging import BraceStyleAdapter
 
@@ -387,13 +388,18 @@ async def get_hwinfo(self) -> HardwareMetadata:
     async def create_vfolder(
         self,
         vfid: VFolderID,
-        exist_ok=False,
+        exist_ok: bool = False,
+        mode: int = DEFAULT_VFOLDER_PERMISSION_MODE,
     ) -> None:
         qspath = self.quota_model.mangle_qspath(vfid)
         if not qspath.exists():
             raise QuotaScopeNotFoundError
         vfpath = self.mangle_vfpath(vfid)
-        await aiofiles.os.makedirs(vfpath, 0o755, exist_ok=exist_ok)
+        await aiofiles.os.makedirs(vfpath, mode, exist_ok=exist_ok)
+        if mode != DEFAULT_VFOLDER_PERMISSION_MODE:
+            # The mode parameter in os.makedirs() sometimes fails to set directory permissions correctly.
+            # Calling os.chmod() afterward ensures the desired permissions are properly applied.
+            os.chmod(vfpath, mode)
 
     @final
     async def delete_vfolder(self, vfid: VFolderID) -> None: