Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: Enhance VFolder mount with additional explicit and verbose options #1838

Merged
merged 39 commits into from
Mar 31, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
39 commits
Select commit Hold shift + click to select a range
243a9e3
feat: Add vfolder mount resolver for extra options
rapsealk Jan 14, 2024
3f9974f
fix: Apply mypy
rapsealk Jan 14, 2024
a7d2bdf
fix: Rename ParsedVolume to MountPoint
rapsealk Jan 15, 2024
22eaf1e
feat: Upgrade prepare_mount_arg_v2 to embrace v1 inputs
rapsealk Jan 15, 2024
7211059
Merge branch 'main' into feature/enhance-vfolder-mount-command
rapsealk Jan 15, 2024
a46cbe7
fix: Leave VFolderPermission to enum.Enum for backward compatibility
rapsealk Jan 15, 2024
47965a4
doc: Add news fragment
rapsealk Jan 15, 2024
de816d2
docs: Rename the news fragment with the PR number
rapsealk Jan 15, 2024
64622fd
feat: Update cli and api handler
rapsealk Jan 15, 2024
278d073
fix: Set default readonly to None to distinguish explicit value
rapsealk Jan 15, 2024
5150f66
fix: Update parameter type hints
rapsealk Jan 15, 2024
d00a4dd
docs: Update news fragment
rapsealk Jan 15, 2024
27be8e6
fix: Move MountPoint type from cli to common
rapsealk Jan 15, 2024
2307cdc
fix: Use explicit permission field
rapsealk Jan 16, 2024
ec3575f
Merge branch 'main' into feature/enhance-vfolder-mount-command
rapsealk Jan 16, 2024
efb658b
Merge branch 'main' into feature/enhance-vfolder-mount-command
rapsealk Jan 16, 2024
e0ac6ad
feat: Add lark-based parser
rapsealk Jan 18, 2024
e060d11
Cannot distinguish purpose of comma
rapsealk Jan 18, 2024
69ce19b
feat: Handle escape comma and etc
rapsealk Jan 18, 2024
9ec7c31
feat: Clear unused comments
rapsealk Jan 19, 2024
a3a8e9f
feat: Migrate parser to client
rapsealk Jan 19, 2024
a207b5e
Merge branch 'main' into feature/enhance-vfolder-mount-command
rapsealk Jan 19, 2024
948135d
feat: Add BUILD file
rapsealk Jan 19, 2024
457f39a
Merge branch 'main' into feature/enhance-vfolder-mount-command
rapsealk Jan 26, 2024
354a5be
fix: Use explicit names for replace target
rapsealk Jan 26, 2024
d374270
feat: Add escape option
rapsealk Jan 27, 2024
e58c0f4
feat: Update test cases
rapsealk Jan 29, 2024
1cc4376
doc: Elaborate news fragment
achimnol Jan 29, 2024
f88a318
fix: Wrap enums in tx.Enum
rapsealk Jan 30, 2024
ee1135d
Merge branch 'main' into feature/enhance-vfolder-mount-command
rapsealk Jan 30, 2024
c237620
Merge branch 'main' into feature/enhance-vfolder-mount-command
achimnol Feb 5, 2024
ae09b9b
feat: Add support for `session create` command
rapsealk Feb 5, 2024
48d8361
Merge branch 'main' into feature/enhance-vfolder-mount-command
achimnol Feb 23, 2024
94a5e18
Merge branch 'main' into feature/enhance-vfolder-mount-command
rapsealk Mar 7, 2024
4d6e912
feat: Sunset old `prepare_mount_arg()`
rapsealk Mar 29, 2024
07a3e74
Merge branch 'main' into feature/enhance-vfolder-mount-command
rapsealk Mar 29, 2024
f899045
chore: update GraphQL schema dump
rapsealk Mar 29, 2024
f3aaa6f
Merge branch 'main' into feature/enhance-vfolder-mount-command
kyujin-cho Mar 31, 2024
62ed6e9
chore: update GraphQL schema dump
rapsealk Mar 31, 2024
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions changes/1838.feature.md
Original file line number Diff line number Diff line change
@@ -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
2 changes: 1 addition & 1 deletion src/ai/backend/client/cli/session/args.py
Original file line number Diff line number Diff line change
Expand Up @@ -76,7 +76,7 @@
"-m",
"--mount",
"mount",
metavar="NAME[=PATH]",
metavar="NAME[=PATH] or NAME[:PATH]",
type=str,
multiple=True,
help=(
Expand Down
42 changes: 26 additions & 16 deletions src/ai/backend/client/cli/session/execute.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down
12 changes: 9 additions & 3 deletions src/ai/backend/client/cli/session/lifecycle.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down Expand Up @@ -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
Expand All @@ -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,
Expand Down Expand Up @@ -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 = {
Expand Down
7 changes: 7 additions & 0 deletions src/ai/backend/client/func/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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.
Expand All @@ -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:
Expand Down Expand Up @@ -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"):
Expand Down Expand Up @@ -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,
Expand Down
1 change: 1 addition & 0 deletions src/ai/backend/common/models/minilang/BUILD
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
python_sources(name="src")
60 changes: 60 additions & 0 deletions src/ai/backend/common/models/minilang/mount.py
Original file line number Diff line number Diff line change
@@ -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
43 changes: 42 additions & 1 deletion src/ai/backend/common/types.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -73,6 +75,7 @@
"MountPermission",
"MountPermissionLiteral",
"MountTypes",
"MountPoint",
"VFolderID",
"QuotaScopeID",
"VFolderUsageMode",
Expand Down Expand Up @@ -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 = {
"\\,": ",",
"\\:": ":",
"\\=": "=",
}
achimnol marked this conversation as resolved.
Show resolved Hide resolved
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
Expand Down
5 changes: 5 additions & 0 deletions src/ai/backend/manager/api/exceptions.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."
Expand Down
10 changes: 5 additions & 5 deletions src/ai/backend/manager/api/schema.graphql
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -1595,4 +1595,4 @@ input ImageRefType {
name: String!
registry: String
architecture: String
}
}
18 changes: 17 additions & 1 deletion src/ai/backend/manager/api/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down
Loading
Loading