diff --git a/changes/1890.fix.md b/changes/1890.fix.md new file mode 100644 index 0000000000..a92ec5b287 --- /dev/null +++ b/changes/1890.fix.md @@ -0,0 +1 @@ +Fix malfunctioning CLI command `session create-from-template` by reorganizing `click.option` decorators diff --git a/src/ai/backend/client/cli/session/args.py b/src/ai/backend/client/cli/session/args.py index 9cf9635dfb..6a432324be 100644 --- a/src/ai/backend/client/cli/session/args.py +++ b/src/ai/backend/client/cli/session/args.py @@ -8,6 +8,8 @@ "--name", "--client-token", metavar="NAME", + type=str, + default=None, help="Specify a human-readable session name. If not set, a random hex string is used.", ), # job scheduling options @@ -38,7 +40,9 @@ help="The maximum duration to wait until the session starts.", ), click.option( - "--no-reuse", is_flag=True, help="Do not reuse existing sessions but return an error." + "--no-reuse", + is_flag=True, + help="Do not reuse existing sessions but return an error.", ), click.option( "--callback-url", @@ -58,7 +62,10 @@ ), # extra options click.option( - "--tag", type=str, default=None, help="User-defined tag string to annotate sessions." + "--tag", + type=str, + default=None, + help="User-defined tag string to annotate sessions.", ), # resource spec click.option( @@ -82,6 +89,7 @@ click.option( "--scaling-group", "--sgroup", + metavar="SCALING_GROUP", type=str, default=None, help=( diff --git a/src/ai/backend/client/cli/session/lifecycle.py b/src/ai/backend/client/cli/session/lifecycle.py index 11f72acbde..1ff6f9e1da 100644 --- a/src/ai/backend/client/cli/session/lifecycle.py +++ b/src/ai/backend/client/cli/session/lifecycle.py @@ -264,6 +264,7 @@ def create( def _create_from_template_cmd(docs: str = None): @click.argument("template_id") + @click_start_option() @click.option( "-o", "--owner", @@ -274,7 +275,14 @@ def _create_from_template_cmd(docs: str = None): help="Set the owner of the target session explicitly.", ) # job scheduling options - @click.option("-i", "--image", default=undefined, help="Set compute_session image to run.") + @click.option( + "-i", + "--image", + metavar="IMAGE", + type=OptionalType(str), + default=undefined, + help="Set compute_session image to run.", + ) @click.option( "-c", "--startup-command", @@ -293,17 +301,18 @@ def _create_from_template_cmd(docs: str = None): "The session will get scheduled after all of them successfully finish." ), ) + # resource spec @click.option( - "--depends", - metavar="SESSION_ID", - type=str, - multiple=True, + "--scaling-group", + "--sgroup", + metavar="SCALING_GROUP", + type=OptionalType(str), + default=undefined, help=( - "Set the list of session ID or names that the newly created session depends on. " - "The session will get scheduled after all of them successfully finish." + "The scaling group to execute session. If not specified " + "all available scaling groups are included in the scheduling." ), ) - # resource spec @click.option( "--cluster-size", metavar="NUMBER", @@ -311,12 +320,28 @@ def _create_from_template_cmd(docs: str = None): default=undefined, help="The size of cluster in number of containers.", ) + # resource grouping @click.option( - "--resource-opts", - metavar="KEY=VAL", - type=str, - multiple=True, - help="Resource options for creating compute session (e.g: shmem=64m)", + "-d", + "--domain", + metavar="DOMAIN_NAME", + type=OptionalType(str), + default=undefined, + help=( + "Domain name where the session will be spawned. " + "If not specified, config's domain name will be used." + ), + ) + @click.option( + "-g", + "--group", + metavar="GROUP_NAME", + type=OptionalType(str), + default=undefined, + help=( + "Group name where the session is spawned. " + "User should be a member of the group to execute the code." + ), ) # template overrides @click.option( @@ -343,35 +368,34 @@ def _create_from_template_cmd(docs: str = None): "any resource specified at template," ), ) - @click_start_option() def create_from_template( # base args template_id: str, - name: str | Undefined, # click_start_option + name: str | None, # click_start_option owner: str | Undefined, # job scheduling options - type_: Literal["batch", "interactive"] | Undefined, # click_start_option + type: Literal["batch", "interactive"], # click_start_option starts_at: str | None, # click_start_option image: str | Undefined, startup_command: str | Undefined, enqueue_only: bool, # click_start_option - max_wait: int | Undefined, # click_start_option + max_wait: int, # click_start_option no_reuse: bool, # click_start_option depends: Sequence[str], - callback_url: str, # click_start_option + callback_url: str | None, # click_start_option # execution environment env: Sequence[str], # click_start_option # extra options - tag: str | Undefined, # click_start_option + tag: str | None, # click_start_option # resource spec mount: Sequence[str], # click_start_option - scaling_group: str | Undefined, # click_start_option + scaling_group: str | Undefined, resources: Sequence[str], # click_start_option - cluster_size: int | Undefined, # click_start_option + cluster_size: int | Undefined, resource_opts: Sequence[str], # click_start_option # resource grouping - domain: str | None, # click_start_option - group: str | None, # click_start_option + domain: str | Undefined, + group: str | Undefined, # template overrides no_mount: bool, no_env: bool, @@ -386,7 +410,7 @@ def create_from_template( \b TEMPLATE_ID: The template ID to create a session from. """ - if name is undefined: + if name is None: name = f"pysdk-{secrets.token_hex(5)}" else: name = name @@ -403,31 +427,35 @@ def create_from_template( prepared_mount, prepared_mount_map = ( prepare_mount_arg(mount) if len(mount) > 0 or no_mount else (undefined, undefined) ) + kwargs = { + "name": name, + "type_": type, + "starts_at": starts_at, + "enqueue_only": enqueue_only, + "max_wait": max_wait, + "no_reuse": no_reuse, + "dependencies": depends, + "callback_url": callback_url, + "cluster_size": cluster_size, + "mounts": prepared_mount, + "mount_map": prepared_mount_map, + "envs": envs, + "startup_command": startup_command, + "resources": parsed_resources, + "resource_opts": parsed_resource_opts, + "owner_access_key": owner, + "domain_name": domain, + "group_name": group, + "scaling_group": scaling_group, + "tag": tag, + } + kwargs = {key: value for key, value in kwargs.items() if value is not undefined} with Session() as session: try: compute_session = session.ComputeSession.create_from_template( template_id, image=image, - name=name, - type_=type_, - starts_at=starts_at, - enqueue_only=enqueue_only, - max_wait=max_wait, - no_reuse=no_reuse, - dependencies=depends, - callback_url=callback_url, - cluster_size=cluster_size, - mounts=prepared_mount, - mount_map=prepared_mount_map, - envs=envs, - startup_command=startup_command, - resources=parsed_resources, - resource_opts=parsed_resource_opts, - owner_access_key=owner, - domain_name=domain, - group_name=group, - scaling_group=scaling_group, - tag=tag, + **kwargs, ) except Exception as e: print_error(e) diff --git a/src/ai/backend/manager/models/session_template.py b/src/ai/backend/manager/models/session_template.py index 357189f203..90a1f890c0 100644 --- a/src/ai/backend/manager/models/session_template.py +++ b/src/ai/backend/manager/models/session_template.py @@ -94,8 +94,7 @@ def check_task_template(raw_data: Mapping[str, Any]) -> Mapping[str, Any]: for p in mounts.values(): if p is None: continue - if p.startswith("/home/work/"): - p = p.replace("/home/work/", "") + p = p.removeprefix("/home/work/") if not verify_vfolder_name(p): raise InvalidArgument(f"Path {p} is reserved for internal operations.") return data