Move cli helpers to separate file and add click Union type back

Author: sjmonson

src/guidellm/__main__.py

@@ -1,8 +1,7 @@
 import asyncio
 import codecs
-import json
 from pathlib import Path
-from typing import Any, get_args
+from typing import get_args
 
 import click
 from pydantic import ValidationError
@@ -14,34 +13,13 @@
 from guidellm.config import print_config
 from guidellm.preprocess.dataset import ShortPromptStrategy, process_dataset
 from guidellm.scheduler import StrategyType
+from guidellm.utils import cli as cli_tools
 
 STRATEGY_PROFILE_CHOICES = set(
     list(get_args(ProfileType)) + list(get_args(StrategyType))
 )
 
 
-def parse_json(ctx, param, value):  # noqa: ARG001
-    if value is None:
-        return None
-    try:
-        return json.loads(value)
-    except json.JSONDecodeError as err:
-        raise click.BadParameter(f"{param.name} must be a valid JSON string.") from err
-
-
-def set_if_not_default(ctx: click.Context, **kwargs) -> dict[str, Any]:
-    """
-    Set the value of a click option if it is not the default value.
-    This is useful for setting options that are not None by default.
-    """
-    values = {}
-    for k, v in kwargs.items():
-        if ctx.get_parameter_source(k) != click.core.ParameterSource.DEFAULT:
-            values[k] = v
-
-    return values
-
-
 @click.group()
 def cli():
     pass
@@ -52,7 +30,10 @@ def cli():
 )
 @click.option(
     "--scenario",
-    type=str,
+    type=cli_tools.Union(
+        click.Path(exists=True, readable=True, file_okay=True, dir_okay=False),
+        click.STRING
+    ),
     default=None,
     help=("TODO: A scenario or path to config"),
 )
@@ -72,7 +53,7 @@ def cli():
 )
 @click.option(
     "--backend-args",
-    callback=parse_json,
+    callback=cli_tools.parse_json,
     default=GenerativeTextScenario.get_default("backend_args"),
     help=(
         "A JSON string containing any arguments to pass to the backend as a "
@@ -101,7 +82,7 @@ def cli():
 @click.option(
     "--processor-args",
     default=GenerativeTextScenario.get_default("processor_args"),
-    callback=parse_json,
+    callback=cli_tools.parse_json,
     help=(
         "A JSON string containing any arguments to pass to the processor constructor "
         "as a dict with **kwargs."
@@ -119,7 +100,7 @@ def cli():
 @click.option(
     "--data-args",
     default=GenerativeTextScenario.get_default("data_args"),
-    callback=parse_json,
+    callback=cli_tools.parse_json,
     help=(
         "A JSON string containing any arguments to pass to the dataset creation "
         "as a dict with **kwargs."
@@ -220,7 +201,7 @@ def cli():
 )
 @click.option(
     "--output-extras",
-    callback=parse_json,
+    callback=cli_tools.parse_json,
     help="A JSON string of extra data to save with the output benchmarks",
 )
 @click.option(
@@ -265,7 +246,7 @@ def benchmark(
 ):
     click_ctx = click.get_current_context()
 
-    overrides = set_if_not_default(
+    overrides = cli_tools.set_if_not_default(
         click_ctx,
         target=target,
         backend_type=backend_type,
@@ -370,15 +351,15 @@ def preprocess():
 @click.option(
     "--processor-args",
     default=None,
-    callback=parse_json,
+    callback=cli_tools.parse_json,
     help=(
         "A JSON string containing any arguments to pass to the processor constructor "
         "as a dict with **kwargs."
     ),
 )
 @click.option(
     "--data-args",
-    callback=parse_json,
+    callback=cli_tools.parse_json,
     help=(
         "A JSON string containing any arguments to pass to the dataset creation "
         "as a dict with **kwargs."

src/guidellm/utils/cli.py

@@ -0,0 +1,63 @@
+import json
+from typing import Any
+
+import click
+
+
+def parse_json(ctx, param, value):  # noqa: ARG001
+    if value is None:
+        return None
+    try:
+        return json.loads(value)
+    except json.JSONDecodeError as err:
+        raise click.BadParameter(f"{param.name} must be a valid JSON string.") from err
+
+
+def set_if_not_default(ctx: click.Context, **kwargs) -> dict[str, Any]:
+    """
+    Set the value of a click option if it is not the default value.
+    This is useful for setting options that are not None by default.
+    """
+    values = {}
+    for k, v in kwargs.items():
+        if ctx.get_parameter_source(k) != click.core.ParameterSource.DEFAULT:
+            values[k] = v
+
+    return values
+
+
+class Union(click.ParamType):
+    """
+    A custom click parameter type that allows for multiple types to be accepted.
+    """
+
+    def __init__(self, *types: click.ParamType):
+        self.types = types
+        self.name = "".join(t.name for t in types)
+
+    def convert(self, value, param, ctx):
+        fails = []
+        for t in self.types:
+            try:
+                return t.convert(value, param, ctx)
+            except click.BadParameter as e:
+                fails.append(str(e))
+                continue
+
+        self.fail("; ".join(fails) or f"Invalid value: {value}")  # noqa: RET503
+
+
+    def get_metavar(self, param: click.Parameter) -> str:
+        def get_choices(t: click.ParamType) -> str:
+            meta = t.get_metavar(param)
+            return meta if meta is not None else t.name
+
+        # Get the choices for each type in the union.
+        choices_str = "|".join(map(get_choices, self.types))
+
+        # Use curly braces to indicate a required argument.
+        if param.required and param.param_type_name == "argument":
+            return f"{{{choices_str}}}"
+
+        # Use square braces to indicate an option or optional argument.
+        return f"[{choices_str}]"