Dynamic SUTs refactor (#1141)

* New SUTFactory object. Move classify_sut_uid there

* Remove pre-flight registration of dynamic suts. Flesh out sut factory make_instance

* Dynamic SUTs inject secrets. Non-hf factories return hydrated SUT instances

* SUT factory also makes dynamic suts.

* Update HF factories

* try solve circular imports

* Use SUT factory for UID validation-- dynamic suts default to true

* check secrets bypasses dynamic suts using sut factory

* remove validate_uid --sut callback

Author: bkorycki

src/modelbench/run.py

@@ -12,13 +12,11 @@
 import sys
 from collections import defaultdict
 from datetime import datetime, timezone
-from typing import List
 
 import click
 
 import termcolor
 from click import echo
-from modelgauge.command_line import compact_uid_list, validate_uid
 from modelgauge.config import load_secrets_from_config, write_default_config
 from modelgauge.load_plugins import load_plugins
 from modelgauge.locales import DEFAULT_LOCALE, LOCALES, PUBLISHED_LOCALES, validate_locale
@@ -82,7 +80,7 @@ def at_end(result, **kwargs):
 @cli.command(help="List known suts")
 @local_plugin_dir_option
 def list_suts():
-    print(compact_uid_list(SUTS))
+    print(SUTS.compact_uid_list())
 
 
 @cli.command(help="run a benchmark")
@@ -102,7 +100,6 @@ def list_suts():
     multiple=False,
     help="SUT UID to run",
     required=True,
-    callback=validate_uid,
 )
 @click.option("--anonymize", type=int, help="Randon number seed for consistent anonymization SUTs")
 @click.option("--threads", default=32, help="How many threads to use per stage")

src/modelgauge/command_line.py

@@ -6,8 +6,9 @@
 from modelgauge.annotator_registry import ANNOTATORS
 from modelgauge.config import write_default_config
 from modelgauge.load_plugins import load_plugins
-from modelgauge.preflight import validate_sut_uid, listify
+from modelgauge.preflight import listify
 from modelgauge.sut import SUTOptions
+from modelgauge.sut_factory import SUT_FACTORY
 from modelgauge.test_registry import TESTS
 
 
@@ -42,12 +43,6 @@ def load_local_plugins(_, __, path: pathlib.Path):
         __import__(plugin.name)
 
 
-def compact_uid_list(registry) -> str:
-    valid_uids = sorted(registry.keys(), key=lambda x: x.lower())
-    valid_uids_str = "\n\t".join(valid_uids)
-    return "\t" + valid_uids_str
-
-
 # Define some reusable options
 DATA_DIR_OPTION = click.option(
     "--data-dir",
@@ -105,14 +100,13 @@ def validate_uid(ctx, param, value):
     """Callback function for click.option UID validation.
     Raises a BadParameter exception if the user-supplied arg(s) are not valid UIDs.
     Applicable for parameters '--test', '--sut', and '--annotator'.
-    SUT IDs are validated in validate_sut_uid via this function.
     If no UID is provided (e.g. an empty list or `None`), the value is returned as-is.
     """
     if not value:
         return value
     # Identify what object we are validating UIDs for.
     if "--sut" in param.opts:
-        return validate_sut_uid(value)
+        registry = SUT_FACTORY
     elif "--test" in param.opts:
         registry = TESTS
     elif "--annotator" in param.opts:
@@ -125,7 +119,7 @@ def validate_uid(ctx, param, value):
 
     unknown_uids = []
     for uid in values:
-        if uid not in registry.keys():
+        if not registry.knows(uid):
             unknown_uids.append(uid)
     if len(unknown_uids) == 0:
         return value
@@ -137,6 +131,6 @@ def validate_uid(ctx, param, value):
 # this is used for all types of UIDs, not just SUTs
 def _bad_uid_error(registry, message, hint=""):
     raise click.BadParameter(
-        f"{message}.\nValid options are:\n{compact_uid_list(registry)}",
+        f"{message}.\nValid options are:\n{registry.compact_uid_list()}",
         param_hint=hint,
     )

src/modelgauge/dynamic_sut_factory.py

@@ -1,8 +1,10 @@
 from abc import ABC, abstractmethod
+from typing import Sequence
 
+from modelgauge.dependency_injection import inject_dependencies
+from modelgauge.secret_values import InjectSecret, RawSecrets
 from modelgauge.dynamic_sut_metadata import DynamicSUTMetadata
-
-from modelgauge.secret_values import InjectSecret
+from modelgauge.sut import SUT
 
 
 class ModelNotSupportedError(Exception):
@@ -27,7 +29,18 @@ class UnknownSUTMakerError(Exception):
 
 
 class DynamicSUTFactory(ABC):
+    def __init__(self, raw_secrets: RawSecrets):
+        self.raw_secrets = raw_secrets
+
+    def injected_secrets(self):
+        """Return the injected secrets as specified by `get_secrets`."""
+        return inject_dependencies(self.get_secrets(), {}, secrets=self.raw_secrets)[0]
+
     @staticmethod
     @abstractmethod
-    def make_sut(sut_metadata: DynamicSUTMetadata):
+    def get_secrets() -> list[InjectSecret]:
+        pass
+
+    @abstractmethod
+    def make_sut(self, sut_metadata: DynamicSUTMetadata) -> SUT:
         pass

src/modelgauge/dynamic_sut_finder.py

@@ -84,8 +84,7 @@ def _get_entry(self, uid: str) -> FactoryEntry:
             try:
                 entry = self._lookup[uid]
             except KeyError:
-                known_uids = list(self._lookup.keys())
-                raise KeyError(f"No registration for {uid}. Known uids: {known_uids}")
+                raise ValueError(f"No registration for {uid}. Known uids:\n{self.compact_uid_list()}")
         return entry
 
     def items(self) -> List[Tuple[str, FactoryEntry[_T]]]:
@@ -97,3 +96,12 @@ def keys(self) -> List[str]:
         """List all keys in the registry."""
         with self.lock:
             return list(self._lookup.keys())
+
+    def knows(self, uid: str) -> bool:
+        """Check if the factory knows about a given uid."""
+        return uid in self.keys()
+
+    def compact_uid_list(self) -> str:
+        valid_uids = sorted(list(self._lookup.keys()), key=lambda x: x.lower())
+        valid_uids_str = "\n\t".join(valid_uids)
+        return "\t" + valid_uids_str

src/modelgauge/instance_factory.py

@@ -25,8 +25,8 @@
 from modelgauge.general import normalize_filename
 from modelgauge.instance_factory import FactoryEntry
 from modelgauge.load_plugins import list_plugins
-from modelgauge.pipeline_runner import AnnotatorRunner, build_runner, PromptPlusAnnotatorRunner, PromptRunner
-from modelgauge.preflight import check_secrets, listify, make_sut
+from modelgauge.pipeline_runner import build_runner
+from modelgauge.preflight import check_secrets, make_sut
 from modelgauge.prompt import TextPrompt
 from modelgauge.secret_values import get_all_secrets, RawSecrets
 from modelgauge.simple_test_runner import run_prompt_response_test
@@ -126,7 +126,7 @@ def list_secrets() -> None:
 
 @modelgauge_cli.command()
 @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)
 @sut_options_options
 @click.option("--prompt", help="The full text to send to the SUT.")
 @click.option(
@@ -144,9 +144,6 @@ def run_sut(
     top_k: Optional[int],
 ):
     """Send a prompt from the command line to a SUT."""
-    secrets = load_secrets_from_config()
-    check_secrets(secrets, sut_uids=[sut])
-
     sut_instance = make_sut(sut)
 
     # Current this only knows how to do prompt response, so assert that is what we have.
@@ -168,7 +165,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, multiple=False, callback=validate_uid)
+@click.option("--sut", "-s", help="Which SUT to run.", required=True, multiple=False)
 @DATA_DIR_OPTION
 @MAX_TEST_ITEMS_OPTION
 @click.option(
@@ -245,7 +242,6 @@ def run_test(
     help="Which SUT to run.",
     multiple=False,
     required=False,
-    callback=validate_uid,
 )
 @click.option(
     "annotator_uids",
@@ -309,7 +305,7 @@ def run_job(
         sut = make_sut(sut_uid)
         if AcceptsTextPrompt not in sut.capabilities:
             raise click.BadParameter(f"{sut_uid} does not accept text prompts")
-        check_secrets(secrets, sut_uids=[sut_uid], annotator_uids=annotator_uids)
+        check_secrets(secrets, annotator_uids=annotator_uids)
         sut_options = create_sut_options(max_tokens, temp, top_p, top_k)
     else:
         sut = None

src/modelgauge/main.py

@@ -1,19 +1,12 @@
+from typing import List
+
 from modelgauge.annotator_registry import ANNOTATORS
 from modelgauge.config import load_secrets_from_config, raise_if_missing_from_config
-from modelgauge.dynamic_sut_finder import make_dynamic_sut_for
 from modelgauge.secret_values import MissingSecretValues
-
-
-from typing import List
-
-from modelgauge.sut import SUTNotFoundException
+from modelgauge.sut_factory import SUT_FACTORY
 from modelgauge.sut_registry import SUTS
 from modelgauge.test_registry import TESTS
 
-SUT_TYPE_DYNAMIC = "dynamic"
-SUT_TYPE_KNOWN = "known"
-SUT_TYPE_UNKNOWN = "unknown"
-
 
 def listify(value):
     """Some functions accept a single UID or a list of them. This turns a single one into a list."""
@@ -33,7 +26,7 @@ def check_secrets(secrets, sut_uids=None, test_uids=None, annotator_uids=None):
     """Checks if all secrets are present for the given UIDs. Raises an error and reports all missing secrets."""
     missing_secrets: List[MissingSecretValues] = []
     if sut_uids is not None:
-        missing_secrets.extend(get_missing_secrets(secrets, SUTS, listify(sut_uids)))
+        missing_secrets.extend(get_missing_secrets(secrets, SUT_FACTORY, listify(sut_uids)))
     if test_uids is not None:
         missing_secrets.extend(get_missing_secrets(secrets, TESTS, test_uids))
         # Check secrets for the annotators in the test as well.
@@ -46,39 +39,9 @@ def check_secrets(secrets, sut_uids=None, test_uids=None, annotator_uids=None):
     return True
 
 
-def classify_sut_uid(uid: str) -> str:
-    if uid in SUTS.keys():
-        return SUT_TYPE_KNOWN
-    elif ":" in uid:
-        return SUT_TYPE_DYNAMIC
-    else:
-        return SUT_TYPE_UNKNOWN
-
-
-def validate_sut_uid(sut_uid):
-    # A blank sut uid is OK for some invocations of modelgauge.
-    # Commands where a non-blank sut uid is required must enforce that with click
-    if not sut_uid:
-        return sut_uid
-
-    sut_type = classify_sut_uid(sut_uid)
-    if sut_type == SUT_TYPE_KNOWN:
-        pass
-    elif sut_type == SUT_TYPE_DYNAMIC:
-        dynamic_sut = make_dynamic_sut_for(sut_uid)  # a tuple that can be splatted for SUTS.register
-        if dynamic_sut:
-            SUTS.register(*dynamic_sut)
-        else:
-            raise SUTNotFoundException(f"{sut_uid} is not a valid dynamic SUT UID")
-    else:
-        raise SUTNotFoundException(f"{sut_uid} is not a valid SUT UID")
-
-    return sut_uid
-
-
 def make_sut(sut_uid: str):
     """Checks that user has all required secrets and returns instantiated SUT."""
     secrets = load_secrets_from_config()
     check_secrets(secrets, sut_uids=[sut_uid])
-    sut = SUTS.make_instance(sut_uid, secrets=secrets)
+    sut = SUT_FACTORY.make_instance(sut_uid, secrets=secrets)
     return sut

src/modelgauge/preflight.py

@@ -14,10 +14,6 @@
 REFUSAL_RESPONSE = ""
 
 
-class SUTNotFoundException(Exception):
-    pass
-
-
 class SUTOptions(BaseModel):
     """
     An exhaustive set of options that could potentially be desired by a SUT.