fix: Reorganize click.option decorators in the `session create-from…

…-template` command (#1890)

Backported-from: main (24.03)
Backported-to: 23.09
Co-authored-by: Joongi Kim <joongi@lablup.com>

changes/1890.fix.md

@@ -0,0 +1 @@
+Fix malfunctioning CLI command `session create-from-template` by reorganizing `click.option` decorators

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=(

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

@@ -265,6 +265,7 @@ def create(
 
 def _create_from_template_cmd(docs: str = None):
     @click.argument("template_id")
+    @click_start_option()
     @click.option(
         "-o",
         "--owner",
@@ -275,7 +276,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",
@@ -294,30 +302,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(
@@ -344,35 +369,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,
@@ -387,7 +411,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
@@ -404,31 +428,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)

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