diff --git a/src/ai/backend/client/cli/admin/domain.py b/src/ai/backend/client/cli/admin/domain.py index 505dbcd80a..15e1149301 100644 --- a/src/ai/backend/client/cli/admin/domain.py +++ b/src/ai/backend/client/cli/admin/domain.py @@ -1,12 +1,15 @@ import sys +from typing import Sequence import click from ai.backend.cli.interaction import ask_yn from ai.backend.cli.types import ExitCode -from ai.backend.client.func.domain import _default_detail_fields, _default_list_fields -from ai.backend.client.session import Session +from ...cli.params import BoolExprType, CommaSeparatedListType, OptionalType +from ...func.domain import _default_detail_fields, _default_list_fields +from ...session import Session +from ...types import Undefined, undefined from ..extensions import pass_ctx_obj from ..pretty import print_info from ..types import CLIContext @@ -59,28 +62,37 @@ def list(ctx: CLIContext) -> None: @pass_ctx_obj @click.argument("name", type=str, metavar="NAME") @click.option("-d", "--description", type=str, default="", help="Description of new domain") -@click.option("-i", "--inactive", is_flag=True, help="New domain will be inactive.") -@click.option("--total-resource-slots", type=str, default="{}", help="Set total resource slots.") +@click.option("--inactive", is_flag=True, help="New domain will be inactive.") +@click.option( + "--total-resource-slots", + type=str, + default="{}", + help="Set total resource slots as a JSON string.", +) @click.option( "--allowed-vfolder-hosts", + "--vfolder-host-permissions", + "--vfhost-perms", type=str, default="{}", help=( - "Allowed virtual folder hosts. It must be JSON string (e.g:" + "Allowed virtual folder hosts and permissions for them. It must be JSON string (e.g:" ' --allowed-vfolder-hosts=\'{"HOST_NAME": ["create-vfolder", "modify-vfolder"]}\')' ), ) @click.option( - "--allowed-docker-registries", type=str, multiple=True, help="Allowed docker registries." + "--allowed-docker-registries", + type=CommaSeparatedListType(), + help="Allowed docker registries.", ) def add( ctx: CLIContext, - name, - description, - inactive, - total_resource_slots, - allowed_vfolder_hosts, - allowed_docker_registries, + name: str, + description: str, + inactive: bool, + total_resource_slots: str, + allowed_vfolder_hosts: str, + allowed_docker_registries: Sequence[str], ): """ Add a new domain. @@ -120,30 +132,54 @@ def add( @domain.command() @pass_ctx_obj @click.argument("name", type=str, metavar="NAME") -@click.option("--new-name", type=str, help="New name of the domain") -@click.option("--description", type=str, help="Description of the domain") -@click.option("--is-active", type=bool, help="Set domain inactive.") -@click.option("--total-resource-slots", type=str, help="Update total resource slots.") +@click.option( + "--new-name", + type=OptionalType(str), + default=undefined, + help="New name of the domain", +) +@click.option( + "--description", + type=OptionalType(str), + default=undefined, + help="Set the description of the domain", +) +@click.option( + "--is-active", + type=OptionalType(BoolExprType), + default=undefined, + help="Change the active/inactive status if specified.", +) +@click.option( + "--total-resource-slots", + type=OptionalType(str), + default=undefined, + help="Update total resource slots.", +) @click.option( "--allowed-vfolder-hosts", - type=str, + type=OptionalType(str), + default=undefined, help=( "Allowed virtual folder hosts. It must be JSON string (e.g:" ' --allowed-vfolder-hosts=\'{"HOST_NAME": ["create-vfolder", "modify-vfolder"]}\')' ), ) @click.option( - "--allowed-docker-registries", type=str, multiple=True, help="Allowed docker registries." + "--allowed-docker-registries", + type=OptionalType(CommaSeparatedListType), + default=undefined, + help="Allowed docker registries.", ) def update( ctx: CLIContext, - name, - new_name, - description, - is_active, - total_resource_slots, - allowed_vfolder_hosts, - allowed_docker_registries, + name: str, + new_name: str | Undefined, + description: str | Undefined, + is_active: bool | Undefined, + total_resource_slots: str | Undefined, + allowed_vfolder_hosts: str | Undefined, + allowed_docker_registries: Sequence[str] | Undefined, ): """ Update an existing domain. diff --git a/src/ai/backend/client/cli/admin/scaling_group.py b/src/ai/backend/client/cli/admin/scaling_group.py index 4b23abc3ba..ffcdc2f345 100644 --- a/src/ai/backend/client/cli/admin/scaling_group.py +++ b/src/ai/backend/client/cli/admin/scaling_group.py @@ -7,8 +7,9 @@ from ai.backend.client.output.fields import scaling_group_fields from ai.backend.client.session import Session +from ...types import Undefined, undefined from ..extensions import pass_ctx_obj -from ..params import JSONParamType +from ..params import BoolExprType, JSONParamType, OptionalType from ..types import CLIContext from . import admin @@ -69,10 +70,16 @@ def list(ctx: CLIContext) -> None: @scaling_group.command() @pass_ctx_obj @click.argument("name", type=str, metavar="NAME") -@click.option("-d", "--description", type=str, default="", help="Description of new scaling group.") -@click.option("-i", "--inactive", is_flag=True, help="New scaling group will be inactive.") @click.option( - "-p", + "-d", + "--description", + "--desc", + type=str, + default="", + help="Description of new scaling group.", +) +@click.option("--inactive", is_flag=True, help="New scaling group will be inactive.") +@click.option( "--private", is_flag=True, help=( @@ -82,7 +89,10 @@ def list(ctx: CLIContext) -> None: ) @click.option("--driver", type=str, default="static", help="Set driver.") @click.option( - "--driver-opts", type=JSONParamType(), default="{}", help="Set driver options as a JSON string." + "--driver-opts", + type=JSONParamType(), + default="{}", + help="Set driver options as a JSON string.", ) @click.option("--scheduler", type=str, default="fifo", help="Set scheduler.") @click.option( @@ -92,23 +102,25 @@ def list(ctx: CLIContext) -> None: help="Set scheduler options as a JSON string.", ) @click.option( - "--use-host-network", is_flag=True, help="If true, run containers on host networking mode." + "--use-host-network", + is_flag=True, + help="If true, run containers on host networking mode.", ) @click.option("--wsproxy-addr", type=str, default=None, help="Set app proxy address.") @click.option("--wsproxy-api-token", type=str, default=None, help="Set app proxy API token.") def add( ctx: CLIContext, - name, - description, - inactive, - private, - driver, - driver_opts, - scheduler, - scheduler_opts, - use_host_network, - wsproxy_addr, - wsproxy_api_token, + name: str, + description: str, + inactive: bool, + private: bool, + driver: str, + driver_opts: dict[str, str] | Undefined, + scheduler: str, + scheduler_opts: dict[str, str] | Undefined, + use_host_network: bool, + wsproxy_addr: str, + wsproxy_api_token: str, ): """ Add a new scaling group. @@ -153,46 +165,82 @@ def add( @scaling_group.command() @pass_ctx_obj @click.argument("name", type=str, metavar="NAME") -@click.option("-d", "--description", type=str, default="", help="Description of new scaling group.") -@click.option("-i", "--inactive", is_flag=True, help="New scaling group will be inactive.") @click.option( - "-p", + "-d", + "--description", + "--desc", + type=OptionalType(str), + default=undefined, + help="Description of new scaling group.", +) +@click.option( + "-a", + "--active", + type=OptionalType(BoolExprType), + default=undefined, + help="Change the active/inactive status if specified.", +) +@click.option( "--private", - is_flag=True, - help=( - "The scaling group will be private. " - "Private scaling groups cannot be used when users create new sessions." - ), + type=OptionalType(BoolExprType), + default=undefined, + help="Change the private status if specified", ) -@click.option("--driver", type=str, default="static", help="Set driver.") @click.option( - "--driver-opts", type=JSONParamType(), default=None, help="Set driver options as a JSON string." + "--driver", + type=OptionalType(str), + default=undefined, + help="Set driver.", +) +@click.option( + "--driver-opts", + type=OptionalType(JSONParamType), + default=undefined, + help="Set driver options as a JSON string.", +) +@click.option( + "--scheduler", + type=OptionalType(str), + default=undefined, + help="Set scheduler.", ) -@click.option("--scheduler", type=str, default="fifo", help="Set scheduler.") @click.option( "--scheduler-opts", - type=JSONParamType(), - default=None, + type=OptionalType(JSONParamType), + default=undefined, help="Set scheduler options as a JSON string.", ) @click.option( - "--use-host-network", is_flag=True, help="If true, run containers on host networking mode." + "--use-host-network", + type=OptionalType(BoolExprType), + default=undefined, + help="Change the host-networking mode if specified.", +) +@click.option( + "--wsproxy-addr", + type=OptionalType(str), + default=undefined, + help="Set app proxy address.", +) +@click.option( + "--wsproxy-api-token", + type=OptionalType(str), + default=undefined, + help="Set app proxy API token.", ) -@click.option("--wsproxy-addr", type=str, default=None, help="Set app proxy address.") -@click.option("--wsproxy-api-token", type=str, default=None, help="Set app proxy API token.") def update( ctx: CLIContext, - name, - description, - inactive, - private, - driver, - driver_opts, - scheduler, - scheduler_opts, - use_host_network, - wsproxy_addr, - wsproxy_api_token, + name: str, + description: str | Undefined, + active: bool | Undefined, + private: bool | Undefined, + driver: str | Undefined, + driver_opts: dict | Undefined, + scheduler: str | Undefined, + scheduler_opts: dict | Undefined, + use_host_network: bool | Undefined, + wsproxy_addr: str | Undefined, + wsproxy_api_token: str | Undefined, ): """ Update existing scaling group. @@ -204,8 +252,8 @@ def update( data = session.ScalingGroup.update( name, description=description, - is_active=not inactive, - is_public=not private, + is_active=active, + is_public=not private if private is not undefined else undefined, driver=driver, driver_opts=driver_opts, scheduler=scheduler, diff --git a/src/ai/backend/client/cli/params.py b/src/ai/backend/client/cli/params.py index c9586596d3..2f85388bfc 100644 --- a/src/ai/backend/client/cli/params.py +++ b/src/ai/backend/client/cli/params.py @@ -1,11 +1,24 @@ import json import re from decimal import Decimal -from typing import Any, Mapping, Optional, Union +from typing import Any, Generic, Mapping, Optional, Protocol, TypeVar, Union import click +import trafaret -from ..types import undefined +from ..types import Undefined, undefined + + +class BoolExprType(click.ParamType): + name = "boolean" + + def convert(self, value, param, ctx): + if isinstance(value, bool): + return value + try: + return trafaret.ToBool().check(value) + except trafaret.DataError: + self.fail(f"Cannot parser/convert {value!r} as a boolean.", param, ctx) class ByteSizeParamType(click.ParamType): @@ -115,7 +128,6 @@ def convert( return json.loads(value) except json.JSONDecodeError: self.fail(f"cannot parse {value!r} as JSON", param, ctx) - return value def drange(start: Decimal, stop: Decimal, num: int): @@ -170,17 +182,29 @@ def convert(self, arg, param, ctx): self.fail(repr(e), param, ctx) -class OptionalType(click.ParamType): +T = TypeVar("T") + + +class SingleValueConstructorType(Protocol): + def __init__(self, value: Any) -> None: ... + + +TScalar = TypeVar("TScalar", bound=SingleValueConstructorType) + + +class OptionalType(click.ParamType, Generic[TScalar]): name = "Optional Type Wrapper" - def __init__(self, type_: type) -> None: + def __init__(self, type_: type[TScalar] | type[click.ParamType]) -> None: super().__init__() self.type_ = type_ - def convert(self, value: Any, param, ctx): + def convert(self, value: Any, param, ctx) -> TScalar | Undefined: try: - if value is None or value is undefined: - return value + if value is undefined: + return undefined + if issubclass(self.type_, click.ParamType): + return self.type_()(value) return self.type_(value) except ValueError: self.fail(f"{value!r} is not valid `{self.type_}` or `undefined`", param, ctx) diff --git a/src/ai/backend/client/func/domain.py b/src/ai/backend/client/func/domain.py index 5153e78b7a..f06a52651e 100644 --- a/src/ai/backend/client/func/domain.py +++ b/src/ai/backend/client/func/domain.py @@ -95,8 +95,8 @@ async def create( description: str = "", is_active: bool = True, total_resource_slots: str | Undefined = undefined, - allowed_vfolder_hosts: str | Undefined = undefined, - allowed_docker_registries: Iterable[str] | Undefined = undefined, + allowed_vfolder_hosts: Sequence[str] | Undefined = undefined, + allowed_docker_registries: Sequence[str] | Undefined = undefined, integration_id: str | Undefined = undefined, fields: Iterable[FieldSpec | str] | None = None, ) -> dict: