feat: Enhance VFolder mount with additional explicit and verbose opti…

…ons (#1838)

Co-authored-by: Joongi Kim <joongi@lablup.com>
Co-authored-by: Kyujin Cho <kyujin.cho@lablup.com>

changes/1838.feature.md

@@ -0,0 +1 @@
+Allow overriding vfolder mount permissions in API calls and CLI commands to create new sessions, with addition of a generic parser of comma-separated "key=value" list for CLI args and API params

src/ai/backend/client/cli/session/args.py

@@ -76,7 +76,7 @@
         "-m",
         "--mount",
         "mount",
-        metavar="NAME[=PATH]",
+        metavar="NAME[=PATH] or NAME[:PATH]",
         type=str,
         multiple=True,
         help=(

src/ai/backend/client/cli/session/execute.py

@@ -20,7 +20,7 @@
 from ai.backend.cli.params import CommaSeparatedListType, RangeExprOptionType
 from ai.backend.cli.types import ExitCode
 from ai.backend.common.arch import DEFAULT_IMAGE_ARCH
-from ai.backend.common.types import ClusterMode
+from ai.backend.common.types import ClusterMode, MountExpression
 
 from ...compat import asyncio_run, current_loop
 from ...config import local_cache_path
@@ -245,26 +245,35 @@ def prepare_env_arg(env: Sequence[str]) -> Mapping[str, str]:
 
 
 def prepare_mount_arg(
-    mount_args: Optional[Sequence[str]],
-) -> Tuple[Sequence[str], Mapping[str, str]]:
+    mount_args: Optional[Sequence[str]] = None,
+    *,
+    escape: bool = True,
+) -> Tuple[Sequence[str], Mapping[str, str], Mapping[str, Mapping[str, str]]]:
     """
     Parse the list of mount arguments into a list of
-    vfolder name and in-container mount path pairs.
+    vfolder name and in-container mount path pairs,
+    followed by extra options.
+
+    :param mount_args: A list of mount arguments such as
+        [
+            "type=bind,source=/colon:path/test,target=/data",
+            "type=bind,source=/colon:path/abcd,target=/zxcv,readonly",
+            # simple formats are still supported
+            "vf-abcd:/home/work/zxcv",
+        ]
     """
     mounts = set()
     mount_map = {}
+    mount_options = {}
     if mount_args is not None:
-        for value in mount_args:
-            if "=" in value:
-                sp = value.split("=", maxsplit=1)
-            elif ":" in value:  # docker-like volume mount mapping
-                sp = value.split(":", maxsplit=1)
-            else:
-                sp = [value]
-            mounts.add(sp[0])
-            if len(sp) == 2:
-                mount_map[sp[0]] = sp[1]
-    return list(mounts), mount_map
+        for mount_arg in mount_args:
+            mountpoint = {**MountExpression(mount_arg).parse(escape=escape)}
+            mount = str(mountpoint.pop("source"))
+            mounts.add(mount)
+            if target := mountpoint.pop("target", None):
+                mount_map[mount] = str(target)
+            mount_options[mount] = mountpoint
+    return list(mounts), mount_map, mount_options
 
 
 @main.command()
@@ -448,7 +457,7 @@ def run(
     envs = prepare_env_arg(env)
     resources = prepare_resource_arg(resources)
     resource_opts = prepare_resource_arg(resource_opts)
-    mount, mount_map = prepare_mount_arg(mount)
+    mount, mount_map, mount_options = prepare_mount_arg(mount, escape=True)
 
     if env_range is None:
         env_range = []  # noqa
@@ -628,6 +637,7 @@ async def _run(session, idx, name, envs, clean_cmd, build_cmd, exec_cmd, is_mult
                 cluster_mode=cluster_mode,
                 mounts=mount,
                 mount_map=mount_map,
+                mount_options=mount_options,
                 envs=envs,
                 resources=resources,
                 resource_opts=resource_opts,

src/ai/backend/client/cli/session/lifecycle.py

@@ -37,7 +37,12 @@
 from .. import events
 from ..pretty import print_done, print_error, print_fail, print_info, print_wait, print_warn
 from .args import click_start_option
-from .execute import format_stats, prepare_env_arg, prepare_mount_arg, prepare_resource_arg
+from .execute import (
+    format_stats,
+    prepare_env_arg,
+    prepare_mount_arg,
+    prepare_resource_arg,
+)
 from .ssh import container_ssh_ctx
 
 list_expr = CommaSeparatedListType()
@@ -169,7 +174,7 @@ def create(
         envs = prepare_env_arg(env)
         parsed_resources = prepare_resource_arg(resources)
         parsed_resource_opts = prepare_resource_arg(resource_opts)
-        mount, mount_map = prepare_mount_arg(mount)
+        mount, mount_map, mount_options = prepare_mount_arg(mount, escape=True)
 
         preopen_ports = preopen
         assigned_agent_list = assign_agent
@@ -189,6 +194,7 @@ def create(
                     cluster_mode=cluster_mode,
                     mounts=mount,
                     mount_map=mount_map,
+                    mount_options=mount_options,
                     envs=envs,
                     startup_command=startup_command,
                     resources=parsed_resources,
@@ -425,7 +431,7 @@ def create_from_template(
             if len(resource_opts) > 0 or no_resource
             else undefined
         )
-        prepared_mount, prepared_mount_map = (
+        prepared_mount, prepared_mount_map, _ = (
             prepare_mount_arg(mount) if len(mount) > 0 or no_mount else (undefined, undefined)
         )
         kwargs = {

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

@@ -176,6 +176,7 @@ async def get_or_create(
         callback_url: Optional[str] = None,
         mounts: List[str] = None,
         mount_map: Mapping[str, str] = None,
+        mount_options: Optional[Mapping[str, Mapping[str, str]]] = None,
         envs: Mapping[str, str] = None,
         startup_command: str = None,
         resources: Mapping[str, str | int] = None,
@@ -238,6 +239,7 @@ async def get_or_create(
             If you want different paths, names should be absolute paths.
             The target mount path of vFolders should not overlap with the linux system folders.
             vFolders which has a dot(.) prefix in its name are not affected.
+        :param mount_options: Mapping which contains extra options for vfolder.
         :param envs: The environment variables which always bypasses the jail policy.
         :param resources: The resource specification. (TODO: details)
         :param cluster_size: The number of containers in this compute session.
@@ -264,6 +266,8 @@ async def get_or_create(
             mounts = []
         if mount_map is None:
             mount_map = {}
+        if mount_options is None:
+            mount_options = {}
         if resources is None:
             resources = {}
         if resource_opts is None:
@@ -303,12 +307,14 @@ async def get_or_create(
             if assign_agent is not None:
                 params["config"].update({
                     "mount_map": mount_map,
+                    "mount_options": mount_options,
                     "preopen_ports": preopen_ports,
                     "agentList": assign_agent,
                 })
             else:
                 params["config"].update({
                     "mount_map": mount_map,
+                    "mount_options": mount_options,
                     "preopen_ports": preopen_ports,
                 })
         if api_session.get().api_version >= (4, "20190615"):
@@ -1201,6 +1207,7 @@ async def get_or_create(
         callback_url: Optional[str] = None,
         mounts: Optional[List[str]] = None,
         mount_map: Optional[Mapping[str, str]] = None,
+        mount_options: Optional[Mapping[str, Mapping[str, str]]] = None,
         envs: Optional[Mapping[str, str]] = None,
         startup_command: Optional[str] = None,
         resources: Optional[Mapping[str, str]] = None,

src/ai/backend/common/models/minilang/BUILD

@@ -0,0 +1 @@
+python_sources(name="src")

src/ai/backend/common/models/minilang/mount.py

@@ -0,0 +1,60 @@
+from typing import Annotated, Mapping, Sequence, TypeAlias
+
+from lark import Lark, Transformer, lexer
+from lark.exceptions import LarkError
+
+_grammar = r"""
+    start: pair ("," pair)*
+    pair: key [("="|":") value]
+    key: SLASH? CNAME (SEPARATOR|CNAME|DIGIT)*
+    value: SLASH? CNAME (SEPARATOR|CNAME|DIGIT)* | ESCAPED_STRING
+
+    SEPARATOR: SLASH | "\\," | "\\=" | "\\:" | DASH
+    SLASH: "/"
+    DASH: "-"
+
+    %import common.CNAME
+    %import common.DIGIT
+    %import common.ESCAPED_STRING
+    %import common.WS
+    %ignore WS
+"""
+
+PairType: TypeAlias = tuple[str, str]
+
+
+class DictTransformer(Transformer):
+    reserved_keys = frozenset({"type", "source", "target", "perm", "permission"})
+
+    def start(self, pairs: Sequence[PairType]) -> Mapping[str, str]:
+        if pairs[0][0] not in self.reserved_keys:  # [["vf-000", "/home/work"]]
+            result = {"source": pairs[0][0]}
+            if target := pairs[0][1]:
+                result["target"] = target
+            return result
+        return dict(pairs)  # [("type", "bind"), ("source", "vf-000"), ...]
+
+    def pair(self, token: Annotated[Sequence[str], 2]) -> PairType:
+        return (token[0], token[1])
+
+    def key(self, token: list[lexer.Token]) -> str:
+        return "".join(token)
+
+    def value(self, token: list[lexer.Token]) -> str:
+        return "".join(token)
+
+
+_parser = Lark(_grammar, parser="lalr")
+
+
+class MountPointParser:
+    def __init__(self) -> None:
+        self._parser = _parser
+
+    def parse_mount(self, expr: str) -> Mapping[str, str]:
+        try:
+            ast = self._parser.parse(expr)
+            result = DictTransformer().transform(ast)
+        except LarkError as e:
+            raise ValueError(f"Virtual folder mount parsing error: {e}")
+        return result

src/ai/backend/common/types.py

@@ -14,7 +14,7 @@
 from dataclasses import dataclass
 from decimal import Decimal
 from ipaddress import ip_address, ip_network
-from pathlib import PurePosixPath
+from pathlib import Path, PurePosixPath
 from ssl import SSLContext
 from typing import (
     TYPE_CHECKING,
@@ -43,9 +43,11 @@
 import trafaret as t
 import typeguard
 from aiohttp import Fingerprint
+from pydantic import BaseModel, ConfigDict, Field
 from redis.asyncio import Redis
 
 from .exception import InvalidIpAddressValue
+from .models.minilang.mount import MountPointParser
 
 __all__ = (
     "aobject",
@@ -73,6 +75,7 @@
     "MountPermission",
     "MountPermissionLiteral",
     "MountTypes",
+    "MountPoint",
     "VFolderID",
     "QuotaScopeID",
     "VFolderUsageMode",
@@ -382,6 +385,44 @@ class MountTypes(enum.StrEnum):
     K8S_HOSTPATH = "k8s-hostpath"
 
 
+class MountPoint(BaseModel):
+    type: MountTypes = Field(default=MountTypes.BIND)
+    source: Path
+    target: Path | None = Field(default=None)
+    permission: MountPermission | None = Field(alias="perm", default=None)
+
+    model_config = ConfigDict(populate_by_name=True)
+
+
+class MountExpression:
+    def __init__(self, expression: str, *, escape_map: Optional[Mapping[str, str]] = None) -> None:
+        self.expression = expression
+        self.escape_map = {
+            "\\,": ",",
+            "\\:": ":",
+            "\\=": "=",
+        }
+        if escape_map is not None:
+            self.escape_map.update(escape_map)
+        # self.unescape_map = {v: k for k, v in self.escape_map.items()}
+
+    def __str__(self) -> str:
+        return self.expression
+
+    def __repr__(self) -> str:
+        return self.__str__()
+
+    def parse(self, *, escape: bool = True) -> Mapping[str, str]:
+        parser = MountPointParser()
+        result = {**parser.parse_mount(self.expression)}
+        if escape:
+            for key, value in result.items():
+                for raw, alternative in self.escape_map.items():
+                    if raw in value:
+                        result[key] = value.replace(raw, alternative)
+        return MountPoint(**result).model_dump()  # type: ignore[arg-type]
+
+
 class HostPortPair(namedtuple("HostPortPair", "host port")):
     def as_sockaddr(self) -> Tuple[str, int]:
         return str(self.host), self.port

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

@@ -303,6 +303,11 @@ class VFolderFilterStatusNotAvailable(BackendError, web.HTTPBadRequest):
     error_title = "There is no available virtual folder to filter its status."
 
 
+class VFolderPermissionError(BackendError, web.HTTPBadRequest):
+    error_type = "https://api.backend.ai/probs/vfolder-permission-error"
+    error_title = "The virtual folder does not permit the specified permission."
+
+
 class DotfileCreationFailed(BackendError, web.HTTPBadRequest):
     error_type = "https://api.backend.ai/probs/generic-bad-request"
     error_title = "Dotfile creation has failed."

src/ai/backend/manager/api/schema.graphql

@@ -31,7 +31,7 @@ type Queries {
     """Added since 24.03.0. Available values: GENERAL, MODEL_STORE"""
     type: [String] = ["GENERAL"]
   ): [Group]
-  image(reference: String!, architecture: String = "aarch64"): Image
+  image(reference: String!, architecture: String = "x86_64"): Image
   images(is_installed: Boolean, is_operation: Boolean): [Image]
   user(domain_name: String, email: String): User
   user_from_uuid(domain_name: String, user_id: ID): User
@@ -973,9 +973,9 @@ type Mutations {
   rescan_images(registry: String): RescanImages
   preload_image(references: [String]!, target_agents: [String]!): PreloadImage
   unload_image(references: [String]!, target_agents: [String]!): UnloadImage
-  modify_image(architecture: String = "aarch64", props: ModifyImageInput!, target: String!): ModifyImage
-  forget_image(architecture: String = "aarch64", reference: String!): ForgetImage
-  alias_image(alias: String!, architecture: String = "aarch64", target: String!): AliasImage
+  modify_image(architecture: String = "x86_64", props: ModifyImageInput!, target: String!): ModifyImage
+  forget_image(architecture: String = "x86_64", reference: String!): ForgetImage
+  alias_image(alias: String!, architecture: String = "x86_64", target: String!): AliasImage
   dealias_image(alias: String!): DealiasImage
   clear_images(registry: String): ClearImages
   create_keypair_resource_policy(name: String!, props: CreateKeyPairResourcePolicyInput!): CreateKeyPairResourcePolicy
@@ -1595,4 +1595,4 @@ input ImageRefType {
   name: String!
   registry: String
   architecture: String
-}
+}

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

@@ -55,7 +55,15 @@
 from ai.backend.common.exception import UnknownImageReference
 from ai.backend.common.logging import BraceStyleAdapter
 from ai.backend.common.plugin.monitor import GAUGE
-from ai.backend.common.types import AccessKey, AgentId, ClusterMode, SessionTypes, VFolderID
+from ai.backend.common.types import (
+    AccessKey,
+    AgentId,
+    ClusterMode,
+    MountPermission,
+    MountTypes,
+    SessionTypes,
+    VFolderID,
+)
 
 from ..config import DEFAULT_CHUNK_SIZE
 from ..defs import DEFAULT_IMAGE_ARCH, DEFAULT_ROLE
@@ -186,6 +194,14 @@ def check_and_return(self, value: Any) -> object:
 creation_config_v5 = t.Dict({
     t.Key("mounts", default=None): t.Null | t.List(t.String),
     tx.AliasedKey(["mount_map", "mountMap"], default=None): t.Null | t.Mapping(t.String, t.String),
+    tx.AliasedKey(["mount_options", "mountOptions"], default=None): t.Null
+    | t.Mapping(
+        t.String,
+        t.Dict({
+            t.Key("type", default=MountTypes.BIND): tx.Enum(MountTypes),
+            tx.AliasedKey(["permission", "perm"], default=None): t.Null | tx.Enum(MountPermission),
+        }).ignore_extra("*"),
+    ),
     t.Key("environ", default=None): t.Null | t.Mapping(t.String, t.String),
     # cluster_size is moved to the root-level parameters
     tx.AliasedKey(["scaling_group", "scalingGroup"], default=None): t.Null | t.String,