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

fix: Reorganize click.option decorators in the session create-from-template command #1890

Merged
merged 6 commits into from
Feb 28, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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/1890.fix.md
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Fix malfunctioning CLI command `session create-from-template` by reorganizing `click.option` decorators
12 changes: 10 additions & 2 deletions src/ai/backend/client/cli/session/args.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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",
Expand All @@ -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(
Expand All @@ -82,6 +89,7 @@
click.option(
"--scaling-group",
"--sgroup",
metavar="SCALING_GROUP",
type=str,
default=None,
help=(
Expand Down
116 changes: 72 additions & 44 deletions src/ai/backend/client/cli/session/lifecycle.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand All @@ -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",
Expand All @@ -293,30 +301,47 @@ 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",
type=OptionalType(int),
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(
Expand All @@ -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,
Expand All @@ -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
Expand All @@ -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)
Expand Down
3 changes: 1 addition & 2 deletions src/ai/backend/manager/models/session_template.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading