remove support for multiple SUT UIDs as inputs

Author: rogthefrog

src/modelgauge/command_line.py

@@ -51,6 +51,15 @@ def compact_uid_list(registry) -> str:
     return "\t" + valid_uids_str
 
 
+def listify(value):
+    """Some functions accept a single UID or a list of them."""
+    if isinstance(value, str):
+        return [
+            value,
+        ]
+    return value
+
+
 # Define some reusable options
 DATA_DIR_OPTION = click.option(
     "--data-dir",
@@ -117,13 +126,8 @@ def _validate_sut_uid(ctx, param, value):
     # and that is enforced in the input validation. This function can still
     # handle multiple SUT UIDs at a time to match validate_uid's semantics,
     # and because it doesn't know or care about the logic in the caller.
-    if isinstance(value, str):
-        sut_uids = [value]
-    else:
-        sut_uids = value
-
+    sut_uids = listify(value)
     requested_sut_uids = classify_sut_uids(sut_uids)
-
     valid_sut_uids = requested_sut_uids["known"]
     for sut_uid in requested_sut_uids["dynamic"]:
         dynamic_sut = make_dynamic_sut_for(sut_uid)  # a tuple that can be splatted for SUTS.register
@@ -166,10 +170,7 @@ def validate_uid(ctx, param, value):
         raise ValueError(f"Cannot validate UID for unknown parameter: {param.opts}")
 
     # This function handles multi-values and single values.
-    if isinstance(value, str):
-        values = [value]
-    else:
-        values = value
+    values = listify(value)
 
     unknown_uids = []
     for uid in values:
@@ -184,6 +185,7 @@ def validate_uid(ctx, param, value):
 
 def get_missing_secrets(secrets, registry, uids):
     missing_secrets: List[MissingSecretValues] = []
+    uids = listify(uids)
     for uid in uids:
         missing_secrets.extend(registry.get_missing_dependencies(uid, secrets=secrets))
     return missing_secrets
@@ -210,10 +212,7 @@ def classify_sut_uids(uids):
     """The CLI now accepts dynamic SUT ids (e.g. "deepseek-ai/DeepSeek-V3:together:hfrelay") in addition to
     pre-registered SUT ids (e.g. "phi-3.5-moe-instruct"). SUT creation and validation are different
     between those two types. This function returns the SUT ids organized by type."""
-    if isinstance(uids, str):
-        uids = [
-            uids,
-        ]
+    uids = listify(uids)
     if len(uids) < 1:
         _bad_uid_error(SUTS, "Please provide at least one SUT uid.")
     identified = {"known": [], "dynamic": [], "unknown": []}

src/modelgauge/main.py

@@ -17,6 +17,7 @@
     create_sut_options,
     display_header,
     display_list_item,
+    listify,
     make_suts,
     modelgauge_cli,
     sut_options_options,
@@ -173,7 +174,7 @@ def run_sut(
 @modelgauge_cli.command()
 @click.option("--test", "-t", help="Which registered TEST to run.", required=True, callback=validate_uid)
 @LOCAL_PLUGIN_DIR_OPTION
-@click.option("--sut", "-s", help="Which SUT to run.", required=True, callback=validate_uid)
+@click.option("--sut", "-s", help="Which SUT to run.", required=True, multiple=False, callback=validate_uid)
 @DATA_DIR_OPTION
 @MAX_TEST_ITEMS_OPTION
 @click.option(
@@ -214,25 +215,25 @@ def run_test(
             sut,
         ]
     )
-    sut_obj = suts[0]
+    sut_instance = suts[0]
 
     # Current this only knows how to do prompt response, so assert that is what we have.
-    assert isinstance(sut_obj, PromptResponseSUT)
+    assert isinstance(sut_instance, PromptResponseSUT)
     assert isinstance(test_obj, PromptResponseTest)
 
-    annotator_objs = []
+    annotator_instances = []
     for annotator in test_obj.get_annotators():
         obj = ANNOTATORS.make_instance(annotator, secrets=secrets)
         assert isinstance(obj, CompletionAnnotator)
-        annotator_objs.append(obj)
+        annotator_instances.append(obj)
 
     if output_file is None:
         os.makedirs("output", exist_ok=True)
         output_file = os.path.join("output", normalize_filename(f"record_for_{test}_{sut}.json"))
     test_record = run_prompt_response_test(
         test_obj,
-        sut_obj,
-        annotator_objs,
+        sut_instance,
+        annotator_instances,
         data_dir,
         max_test_items,
         use_caching=not no_caching,
@@ -253,7 +254,7 @@ def run_test(
     "-s",
     "--sut",
     help="Which SUT to run.",
-    multiple=True,
+    multiple=False,
     required=False,
     callback=validate_uid,
 )
@@ -295,7 +296,7 @@ def run_job(
 ):
     """Run rows in a CSV through (a) SUT(s) and/or a set of annotators.
 
-    If running SUTs, the file must have 'UID' and 'Text' columns. The output will be saved to a CSV file.
+    If running a SUT, the file must have 'UID' and 'Text' columns. The output will be saved to a CSV file.
     If running ONLY annotators, the file must have 'UID', 'Prompt', 'SUT', and 'Response' columns. The output will be saved to a json lines file.
     """
     logging.basicConfig(level=logging.DEBUG if debug else logging.INFO)
@@ -312,16 +313,18 @@ def run_job(
     else:
         ensemble = None
 
+    sut_uids = listify(sut_uid)
+
     # Check all objects for missing secrets.
     secrets = load_secrets_from_config()
     if sut_uid:
-        check_secrets(secrets, sut_uids=sut_uid, annotator_uids=annotator_uids)
+        check_secrets(secrets, sut_uids=sut_uids, annotator_uids=annotator_uids)
     else:
         check_secrets(secrets, annotator_uids=annotator_uids)
 
     suts = {}
-    if sut_uid:
-        all_suts = make_suts(sut_uid)
+    if sut_uids:
+        all_suts = make_suts(sut_uids)
         for sut in all_suts:
             if AcceptsTextPrompt not in sut.capabilities:
                 raise click.BadParameter(f"{sut.uid} does not accept text prompts")
@@ -376,11 +379,11 @@ def show_progress(data):
 @modelgauge_cli.command()
 @sut_options_options
 @click.option(
-    "sut_uids",
+    "sut_uid",
     "-s",
     "--sut",
-    help="Which SUT(s) to run.",
-    multiple=True,
+    help="Which SUT to run.",
+    multiple=False,
     required=False,
     callback=validate_uid,
 )
@@ -423,7 +426,7 @@ def show_progress(data):
     type=click.Path(exists=True, path_type=pathlib.Path),
 )
 def run_csv_items(
-    sut_uids,
+    sut_uid,
     annotator_uids,
     ensemble,
     workers,
@@ -455,11 +458,14 @@ def run_csv_items(
         annotator_uids = annotator_uids + tuple(PRIVATE_ANNOTATOR_SET.annotators)
     else:
         ensemble = None
+
+    sut_uids = listify(sut_uid)
+
     # Check all objects for missing secrets.
     secrets = load_secrets_from_config()
     check_secrets(secrets, sut_uids=sut_uids, annotator_uids=annotator_uids)
 
-    if len(sut_uids):
+    if sut_uids:
         all_suts = make_suts(sut_uids)
         suts = {}
         for sut in all_suts:

src/modelgauge/secret_values.py

@@ -2,10 +2,10 @@
 from dataclasses import dataclass
 from typing import Generic, List, Mapping, Optional, Sequence, Type, TypeVar
 
-from pydantic import BaseModel
-
 from modelgauge.general import get_concrete_subclasses
 
+from pydantic import BaseModel
+
 
 def loggable_secret(secret: str) -> str:
     if str is None or not isinstance(secret, str):

tests/modelgauge_tests/test_cli.py

@@ -12,7 +12,7 @@
 from modelgauge import main
 from modelgauge.annotator_registry import ANNOTATORS
 from modelgauge.annotator_set import AnnotatorSet
-from modelgauge.command_line import _validate_sut_uid, check_secrets, classify_sut_uids, validate_uid
+from modelgauge.command_line import _validate_sut_uid, check_secrets, classify_sut_uids, listify, validate_uid
 from modelgauge.config import MissingSecretsFromConfig
 from modelgauge.data_schema import (
     DEFAULT_PROMPT_RESPONSE_SCHEMA as PROMPT_RESPONSE_SCHEMA,
@@ -354,7 +354,6 @@ def test_run_job_sut_only_output_name(caplog, tmp_path, prompts_file):
         ["run-job", "--sut", "demo_yes_no", "--output-dir", tmp_path, str(prompts_file)],
         catch_exceptions=False,
     )
-    print(result.output)
     assert result.exit_code == 0
 
     out_path = Path(re.findall(r"\S+\.csv", caplog.text)[0])
@@ -598,3 +597,24 @@ def test_validate_uid():
         )
         == "my-fake-annotator"
     )
+
+
+def test_listify():
+    assert listify("string") == [
+        "string",
+    ]
+    assert listify(["a", "b"]) == ["a", "b"]
+    assert listify(("a", "b")) == ("a", "b")
+
+    def noop():
+        pass
+
+    class Noop:
+        pass
+
+    assert listify(noop) == noop
+    assert listify(Noop) == Noop
+    n = Noop()
+    assert listify(n) == n
+
+    assert listify(None) is None