crfm-helm 0.3.0__py3-none-any.whl → 0.5.0__py3-none-any.whl

helm/benchmark/model_deployment_registry.py

@@ -1,4 +1,3 @@
-import os
 from typing import Dict, Optional, List
 from dataclasses import dataclass
 
@@ -7,10 +6,12 @@ import yaml
 
 from helm.common.hierarchical_logger import hlog
 from helm.common.object_spec import ObjectSpec
-from helm.proxy.models import ALL_MODELS, FULL_FUNCTIONALITY_TEXT_MODEL_TAG, MODEL_NAME_TO_MODEL, TEXT_MODEL_TAG, Model
-
-
-MODEL_DEPLOYMENTS_FILE = "model_deployments.yaml"
+from helm.benchmark.model_metadata_registry import (
+    ModelMetadata,
+    get_model_metadata,
+    get_unknown_model_metadata,
+    register_model_metadata,
+)
 
 
 class ClientSpec(ObjectSpec):
@@ -23,65 +24,99 @@ class WindowServiceSpec(ObjectSpec):
 
 @dataclass(frozen=True)
 class ModelDeployment:
-    """A model deployment is an accessible instance of this model (e.g. a hosted endpoint).
-
-    A model can have multiple model deployments."""
+    """
+    A model deployment is an accessible instance of this model (e.g., a hosted endpoint).
+    A model can have multiple model deployments.
+    """
 
     name: str
-    """Name of the model deployment."""
+    """Name of the model deployment. Usually formatted as "<hosting_group>/<engine_name>".
+    Example: "huggingface/t5-11b"."""
 
     client_spec: ClientSpec
     """Specification for instantiating the client for this model deployment."""
 
     model_name: Optional[str] = None
-    """Name of the model that this model deployment is for.
-
-    If unset, defaults to the the same value as `name`."""
+    """Name of the model that this model deployment is for. Refers to the field "name" in the Model class.
+    If unset, defaults to the same value as `name`."""
 
     tokenizer_name: Optional[str] = None
-    """Tokenizer for this model deployment.
-
-    If unset, auto-inferred by the WindowService."""
+    """Tokenizer for this model deployment. If unset, auto-inferred by the WindowService."""
 
     window_service_spec: Optional[WindowServiceSpec] = None
-    """Specification for instantiating the window service for this model deployment"""
+    """Specification for instantiating the window service for this model deployment."""
 
     max_sequence_length: Optional[int] = None
     """Maximum sequence length for this model deployment."""
 
     max_request_length: Optional[int] = None
     """Maximum request length for this model deployment.
-
     If unset, defaults to the same value as max_sequence_length."""
 
+    max_sequence_and_generated_tokens_length: Optional[int] = None
+    """The max length of the model input and output tokens.
+    Some models (like Anthropic/Claude and Megatron) have a specific limit sequence length + max_token.
+    If unset, defaults to INT_MAX (i.e., no limit)."""
+
+    deprecated: bool = False
+    """Whether this model deployment is deprecated."""
+
+    @property
+    def host_organization(self) -> str:
+        """
+        Extracts the host group from the model deployment name.
+        Example: "huggingface" from "huggingface/t5-11b"
+        This can be different from the creator organization (for example "together")
+        """
+        return self.name.split("/")[0]
+
+    @property
+    def engine(self) -> str:
+        """
+        Extracts the model engine from the model deployment name.
+        Example: 'ai21/j1-jumbo' => 'j1-jumbo'
+        """
+        return self.name.split("/")[1]
+
+    def __post_init__(self):
+        if not self.model_name:
+            object.__setattr__(self, "model_name", self.name)
+
 
 @dataclass(frozen=True)
 class ModelDeployments:
     model_deployments: List[ModelDeployment]
 
 
-_name_to_model_deployment: Dict[str, ModelDeployment] = {}
+ALL_MODEL_DEPLOYMENTS: List[ModelDeployment] = []
+DEPLOYMENT_NAME_TO_MODEL_DEPLOYMENT: Dict[str, ModelDeployment] = {
+    deployment.name: deployment for deployment in ALL_MODEL_DEPLOYMENTS
+}
 
 
 def register_model_deployment(model_deployment: ModelDeployment) -> None:
-    hlog(f"Registered model deployment {model_deployment.name}")
-    _name_to_model_deployment[model_deployment.name] = model_deployment
-
-    # Auto-register a model with this name if none exists
-    model_name = model_deployment.model_name or model_deployment.name
-    if model_name not in MODEL_NAME_TO_MODEL:
-        model = Model(
-            group="unknown",
-            name=model_name,
-            tags=[TEXT_MODEL_TAG, FULL_FUNCTIONALITY_TEXT_MODEL_TAG],
+    DEPLOYMENT_NAME_TO_MODEL_DEPLOYMENT[model_deployment.name] = model_deployment
+    ALL_MODEL_DEPLOYMENTS.append(model_deployment)
+
+    model_name: str = model_deployment.model_name or model_deployment.name
+
+    model_metadata: ModelMetadata
+    try:
+        model_metadata = get_model_metadata(model_name)
+    except ValueError:
+        hlog(
+            f"WARNING: Could not find model metadata for model {model_name} of model deployment {model_deployment.name}"
         )
-        MODEL_NAME_TO_MODEL[model_name] = model
-        ALL_MODELS.append(model)
-        hlog(f"Registered default metadata for model {model_name}")
+        model_metadata = get_unknown_model_metadata(model_name)
+        register_model_metadata(model_metadata)
+    deployment_names: List[str] = model_metadata.deployment_names or [model_metadata.name]
+    if model_deployment.name not in deployment_names:
+        if model_metadata.deployment_names is None:
+            model_metadata.deployment_names = []
+        model_metadata.deployment_names.append(model_deployment.name)
 
 
 def register_model_deployments_from_path(path: str) -> None:
-    global _name_to_model_deployment
     hlog(f"Reading model deployments from {path}...")
     with open(path, "r") as f:
         raw = yaml.safe_load(f)
@@ -90,12 +125,100 @@ def register_model_deployments_from_path(path: str) -> None:
         register_model_deployment(model_deployment)
 
 
-def maybe_register_model_deployments_from_base_path(base_path: str) -> None:
-    """Register model deployments from prod_env/model_deployments.yaml"""
-    path = os.path.join(base_path, MODEL_DEPLOYMENTS_FILE)
-    if os.path.exists(path):
-        register_model_deployments_from_path(path)
-
-
-def get_model_deployment(name: str) -> Optional[ModelDeployment]:
-    return _name_to_model_deployment.get(name)
+def get_model_deployment(name: str, warn_deprecated: bool = False) -> ModelDeployment:
+    if name not in DEPLOYMENT_NAME_TO_MODEL_DEPLOYMENT:
+        raise ValueError(f"Model deployment {name} not found")
+    deployment: ModelDeployment = DEPLOYMENT_NAME_TO_MODEL_DEPLOYMENT[name]
+    if deployment.deprecated and warn_deprecated:
+        hlog(f"WARNING: DEPLOYMENT Model deployment {name} is deprecated")
+    return deployment
+
+
+def get_model_deployment_host_organization(name: str) -> str:
+    """Return the host organization name based on the model deployment name.
+
+    Example: "huggingface/t5-11b" -> "huggingface"""
+    deployment: ModelDeployment = get_model_deployment(name)
+    return deployment.host_organization
+
+
+def get_model_names_with_tokenizer(tokenizer_name: str) -> List[str]:
+    """Return the names of all models with the given tokenizer."""
+    deployments: List[ModelDeployment] = [
+        deployment for deployment in ALL_MODEL_DEPLOYMENTS if deployment.tokenizer_name == tokenizer_name
+    ]
+    return [deployment.model_name or deployment.name for deployment in deployments]
+
+
+def get_default_model_deployment_for_model(
+    model_name: str, warn_arg_deprecated: bool = False, ignore_deprecated: bool = False
+) -> Optional[str]:
+    """Returns a valid model deployment name corresponding to the given model arg.
+    This is used as a backwards compatibility layer for model names that are now moved to model deployments.
+    Example: "anthropic/claude-v1.3" => "anthropic/claude-v1.3"
+    Example: "meta/llama-7b" => "together/llama-7b"
+
+    The process to find a model deployment name is as follows:
+    1. If there is a model deployment with the same name as the model arg, use it.
+    2. If there is at least one deployment for the model, use the first one that is available.
+    3. If there are no deployments for the model, returns None.
+
+    This function will also try to find a model deployment name that is not deprecated.
+    If there are no non-deprecated deployments, it will return the first deployment (even if it's deprecated).
+    If ignore_deprecated is True, this function will return None if the model deployment is deprecated.
+
+    If warn_arg_deprecated is True, this function will print a warning if the model deployment name is not the same
+    as the model arg. This is to remind the user that the model name is deprecated and should be replaced with
+    the model deployment name (in their config).
+
+    Args:
+        model_arg: The model arg to convert to a model deployment name.
+        warn_arg_deprecated: Whether to print a warning if the model deployment name is not the same as the model arg.
+        ignore_deprecated: Whether to return None if the model deployment is deprecated.
+    """
+
+    # If there is a model deployment with the same name as the model arg, use it.
+    if model_name in DEPLOYMENT_NAME_TO_MODEL_DEPLOYMENT:
+        deployment: ModelDeployment = DEPLOYMENT_NAME_TO_MODEL_DEPLOYMENT[model_name]
+        if deployment.deprecated and ignore_deprecated:
+            if warn_arg_deprecated:
+                hlog(f"WARNING: Model deployment {model_name} is deprecated")
+            return None
+        return deployment.name
+
+    # If there is at least one deployment for the model, use the first one that is available.
+    available_deployments: List[ModelDeployment] = [
+        deployment for deployment in ALL_MODEL_DEPLOYMENTS if deployment.model_name == model_name
+    ]
+    if len(available_deployments) > 0:
+        available_deployment_names: List[str] = [deployment.name for deployment in available_deployments]
+        if warn_arg_deprecated:
+            hlog("WARNING: Model name is deprecated. Please use the model deployment name instead.")
+            hlog(f"Available model deployments for model {model_name}: {available_deployment_names}")
+
+        # Additionally, if there is a non-deprecated deployment, use it.
+        non_deprecated_deployments: List[ModelDeployment] = [
+            deployment for deployment in available_deployments if not deployment.deprecated
+        ]
+        if len(non_deprecated_deployments) > 0:
+            chosen_deployment = non_deprecated_deployments[0]
+        # There are no non-deprecated deployments, so there are two options:
+        # 1. If we can return an empty string, return it. (no model deployment is available)
+        # 2. If we can't return an empty string, return the first deployment (even if it's deprecated).
+        elif ignore_deprecated:
+            return None
+        else:
+            chosen_deployment = available_deployments[0]
+            if warn_arg_deprecated:
+                hlog(f"WARNING: All model deployments for model {model_name} are deprecated.")
+        if warn_arg_deprecated:
+            hlog(
+                f"Choosing {chosen_deployment.name} (the first one) as "
+                f"the default model deployment for model {model_name}"
+            )
+            hlog("If you want to use a different model deployment, please specify it explicitly.")
+        return chosen_deployment.name
+
+    # Some models are added but have no deployments yet.
+    # In this case, we return None.
+    return None

helm/benchmark/model_metadata_registry.py

@@ -1,50 +1,142 @@
-import os
-from typing import Optional, List
+from typing import Dict, Optional, List
 from dataclasses import dataclass, field
 from datetime import date
 
 import dacite
 import yaml
 
-from helm.proxy.models import ALL_MODELS, MODEL_NAME_TO_MODEL, Model
 
+# Different modalities
+TEXT_MODEL_TAG: str = "TEXT_MODEL_TAG"
+IMAGE_MODEL_TAG: str = "IMAGE_MODEL_TAG"
+CODE_MODEL_TAG: str = "CODE_MODEL_TAG"
+EMBEDDING_MODEL_TAG: str = "EMBEDDING_MODEL_TAG"
 
-MODEL_METADATA_FILE = "model_metadata.yaml"
+# Some model APIs have limited functionalities
+FULL_FUNCTIONALITY_TEXT_MODEL_TAG: str = "FULL_FUNCTIONALITY_TEXT_MODEL_TAG"
+LIMITED_FUNCTIONALITY_TEXT_MODEL_TAG: str = "LIMITED_FUNCTIONALITY_TEXT_MODEL_TAG"
 
+# ChatML format
+CHATML_MODEL_TAG: str = "CHATML_MODEL_TAG"
 
-@dataclass(frozen=True)
+# OpenAI Chat format
+OPENAI_CHATGPT_MODEL_TAG: str = "OPENAI_CHATGPT_MODEL_TAG"
+
+# Mistral instruction-following format
+MISTRAL_MODEL_TAG: str = "MISTRAL_MODEL_TAG"
+
+# For Anthropic models
+ANTHROPIC_CLAUDE_1_MODEL_TAG: str = "ANTHROPIC_CLAUDE_1_MODEL_TAG"
+ANTHROPIC_CLAUDE_2_MODEL_TAG: str = "ANTHROPIC_CLAUDE_2_MODEL_TAG"
+ANTHROPIC_CLAUDE_3_MODEL_TAG: str = "ANTHROPIC_CLAUDE_3_MODEL_TAG"
+
+GOOGLE_PALM_2_MODEL_TAG: str = "GOOGLE_PALM_2_MODEL_TAG"
+GOOGLE_GEMINI_MODEL_TAG: str = "GOOGLE_GEMINI_MODEL_TAG"
+GOOGLE_GEMMA_INSTRUCT_MODEL_TAG: str = "GOOGLE_GEMMA_INSTRUCT_MODEL_TAG"
+
+# Models which emit garbage tokens when temperature=0.
+BUGGY_TEMP_0_TAG: str = "BUGGY_TEMP_0_TAG"
+
+# Models that are used for ablations and fine-grained analyses.
+# These models are selected specifically because of their low marginal cost to evaluate.
+ABLATION_MODEL_TAG: str = "ABLATION_MODEL_TAG"
+
+# Some models (e.g., T5) have stripped newlines.
+# So we cannot use \n as a stop sequence for these models.
+NO_NEWLINES_TAG: str = "NO_NEWLINES_TAG"
+
+# Some models (e.g., UL2) require a prefix (e.g., [NLG]) in the
+# prompts to indicate the mode before doing inference.
+NLG_PREFIX_TAG: str = "NLG_PREFIX_TAG"
+
+# Some models can follow instructions.
+INSTRUCTION_FOLLOWING_MODEL_TAG: str = "INSTRUCTION_FOLLOWING_MODEL_TAG"
+
+# For text-to-image models
+TEXT_TO_IMAGE_MODEL_TAG: str = "TEXT_TO_IMAGE_MODEL_TAG"
+
+# For Vision-langauge models (VLMs)
+VISION_LANGUAGE_MODEL_TAG: str = "VISION_LANGUAGE_MODEL_TAG"
+# IDEFICS require a special prompt format (see `IDEFICSInstructRunExpander`)
+IDEFICS_INSTRUCT_MODEL_TAG: str = "IDEFICS_INSTRUCT_MODEL_TAG"
+IDEFICS_MODEL_TAG: str = "IDEFICS_MODEL_TAG"
+# Llava should use a special prompt format (see `LlavaRunExpander`)
+LLAVA_MODEL_TAG: str = "LLAVA_MODEL_TAG"
+# OpenFlamingo has a special prompt format (see `OpenFlamingoRunExpander`)
+OPEN_FLAMINGO_MODEL_TAG: str = "OPEN_FLAMINGO_MODEL_TAG"
+# Some VLMs do not support multiple images in the prompt
+LIMITED_FUNCTIONALITY_VLM_TAG: str = "LIMITED_FUNCTIONALITY_VLM_TAG"
+FULL_FUNCTIONALITY_VLM_TAG: str = "FULL_FUNCTIONALITY_VLM_TAG"
+
+
+# Frozen is set to false as the model_deployment_registry.py file
+# might populate the deployment_names field.
+
+
+@dataclass(frozen=False)
 class ModelMetadata:
     name: str
-    """Name of the model e.g. "meta/llama-2"."""
+    """Name of the model group (e.g., "openai/davinci"). This is the name of the model,
+    not the name of the deployment.
+    Usually formatted as "<creator_organization>/<engine_name>". Example: "ai21/j1-jumbo"."""
 
-    creator_organization: Optional[str] = None
-    """Organization that originally created the model (e.g. "meta")."""
+    creator_organization_name: str
+    """Name of the organization that created the model."""
 
-    access: Optional[str] = None
-    """How this model is available (e.g., limited).
+    display_name: str
+    """Name that is going to be displayed to the user (on the website, etc.)."""
 
-    If there are multiple deployments, this should be the most permissive access across
-    all deployments."""
+    description: str
+    """Description of the model, to be displayed on the website."""
 
-    todo: bool = False
-    """Whether we have yet to evaluate this model."""
+    access: str
+    """Description of the access level of the model. Should be one of the following:
+    - "open": the model is open-source and can be downloaded from the internet.
+    - "closed": not accessible
+    - "limited": accessible with an API key.
+    If there are multiple deployments, this should be the most permissive access across all deployments."""
 
-    release_date: Optional[date] = None
-    """When the model was released."""
+    release_date: Optional[date]
+    """Release date of the model."""
 
-    num_parameters: Optional[int] = None
-    """The number of model parameters.
+    tags: List[str] = field(default_factory=list)
+    """Tags corresponding to the properties of the model."""
 
+    num_parameters: Optional[int] = None
+    """Number of parameters in the model.
     This should be a string as the number of parameters is usually a round number (175B),
     but we set it as an int for plotting purposes."""
 
-    tags: List[str] = field(default_factory=list)
-    """"""
+    deployment_names: Optional[List[str]] = None
+    """List of the model deployments for this model. Should at least contain one model deployment.
+    Refers to the field "name" in the ModelDeployment class. Defaults to a single model deployment
+    with the same name as the model."""
+
+    @property
+    def creator_organization(self) -> str:
+        """
+        Extracts the creator organization from the model name.
+        Example: 'ai21/j1-jumbo' => 'ai21'
+        This can be different from the hosting organization.
+        """
+        return self.name.split("/")[0]
+
+    @property
+    def engine(self) -> str:
+        """
+        Extracts the model engine from the model name.
+        Example: 'ai21/j1-jumbo' => 'j1-jumbo'
+        """
+        return self.name.split("/")[1]
 
 
 @dataclass(frozen=True)
 class ModelMetadataList:
-    models: List[ModelMetadata]
+    models: List[ModelMetadata] = field(default_factory=list)
+
+
+ALL_MODELS_METADATA: List[ModelMetadata] = []
+MODEL_NAME_TO_MODEL_METADATA: Dict[str, ModelMetadata] = {model.name: model for model in ALL_MODELS_METADATA}
 
 
 def register_model_metadata_from_path(path: str) -> None:
@@ -55,17 +147,71 @@ def register_model_metadata_from_path(path: str) -> None:
     # serialization format for dates
     model_metadata_list = dacite.from_dict(ModelMetadataList, raw)
     for model_metadata in model_metadata_list.models:
-        model = Model(
-            group="none",  # TODO: Group should be part of model deployment, not model
-            name=model_metadata.name,
-            tags=model_metadata.tags,
-        )
-        MODEL_NAME_TO_MODEL[model_metadata.name] = model
-        ALL_MODELS.append(model)
-
-
-def maybe_register_model_metadata_from_base_path(base_path: str) -> None:
-    """Register model metadata from prod_env/model_metadata.yaml"""
-    path = os.path.join(base_path, MODEL_METADATA_FILE)
-    if os.path.exists(path):
-        register_model_metadata_from_path(path)
+        register_model_metadata(model_metadata)
+
+
+def register_model_metadata(model_metadata: ModelMetadata) -> None:
+    """Register a single model configuration."""
+    ALL_MODELS_METADATA.append(model_metadata)
+    MODEL_NAME_TO_MODEL_METADATA[model_metadata.name] = model_metadata
+
+
+def get_model_metadata(model_name: str) -> ModelMetadata:
+    """Return the `ModelMetadata` for the model name."""
+    if model_name not in MODEL_NAME_TO_MODEL_METADATA:
+        raise ValueError(f"No model with name: {model_name}")
+
+    return MODEL_NAME_TO_MODEL_METADATA[model_name]
+
+
+def get_all_models() -> List[str]:
+    """Return all model names."""
+    return list(MODEL_NAME_TO_MODEL_METADATA.keys())
+
+
+def get_model_names_with_tag(tag: str) -> List[str]:
+    """Return all model names of models with the given tag."""
+    return [model.name for model in ALL_MODELS_METADATA if tag in model.tags]
+
+
+def model_has_tag(model_name: str, tag: str) -> bool:
+    """Return True if the model has the given tag. False otherwise."""
+    return tag in get_model_metadata(model_name).tags
+
+
+def get_all_text_models() -> List[str]:
+    """Return all model names of text models."""
+    return get_model_names_with_tag(TEXT_MODEL_TAG)
+
+
+def get_all_code_models() -> List[str]:
+    """Return all model names of code models."""
+    return get_model_names_with_tag(CODE_MODEL_TAG)
+
+
+def get_all_instruction_following_models() -> List[str]:
+    """Return all model names of instruction following models."""
+    return get_model_names_with_tag(INSTRUCTION_FOLLOWING_MODEL_TAG)
+
+
+def is_text_to_image_model(model_name: str) -> bool:
+    """Returns True if the model is a text-to-image model. False otherwise."""
+    return model_has_tag(model_name, TEXT_TO_IMAGE_MODEL_TAG)
+
+
+def is_vlm(model_name: str) -> bool:
+    """Returns True if the model is a vision-language model (VLM). False otherwise."""
+    return model_has_tag(model_name, VISION_LANGUAGE_MODEL_TAG)
+
+
+def get_unknown_model_metadata(helm_model_name: str) -> ModelMetadata:
+    """Return placeholder ModelMetadata for an unknown model."""
+    return ModelMetadata(
+        name=helm_model_name,
+        creator_organization_name="Unknown",
+        display_name=helm_model_name,
+        description=helm_model_name,
+        access="open",
+        release_date=date.today(),
+        tags=[TEXT_MODEL_TAG, FULL_FUNCTIONALITY_TEXT_MODEL_TAG],
+    )

helm/benchmark/multi_gpu_runner.py

@@ -0,0 +1,133 @@
+import signal
+import threading
+import traceback
+from typing import List
+import os
+import time
+import torch
+import torch.multiprocessing as multiprocessing
+from concurrent.futures import ProcessPoolExecutor as Pool
+from tqdm import tqdm
+
+from helm.benchmark.config_registry import (
+    register_configs_from_directory,
+    register_builtin_configs_from_helm_package,
+)
+from helm.benchmark.executor import ExecutionSpec
+from helm.benchmark.runner import Runner, RunSpec, RunnerError
+from helm.common.hierarchical_logger import hlog, htrack_block
+from helm.benchmark.runner_config_registry import RUNNER_CONFIG
+
+_MAX_CONCURRENT_WORKERS_ENV_NAME = "HELM_MAX_CONCURRENT_WORKERS"
+
+
+# From
+# https://stackoverflow.com/questions/71300294/how-to-terminate-pythons-processpoolexecutor-when-parent-process-dies
+def start_thread_to_terminate_when_parent_process_dies(ppid):
+    pid = os.getpid()
+
+    def f():
+        while True:
+            try:
+                os.kill(ppid, 0)
+            except OSError:
+                os.kill(pid, signal.SIGTERM)
+            time.sleep(1)
+
+    thread = threading.Thread(target=f, daemon=True)
+    thread.start()
+
+
+def initialize_worker(gpu_id: int):
+    hlog(f"Worker {gpu_id} initializing")
+
+    # Wait for 0.1 seconds to ensure all workers are initialized with different CUDA_VISIBLE_DEVICES
+    time.sleep(0.1)
+
+    # Pin GPU to worker process
+    os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
+
+    # Necessary for code_metrics in humaneval to work properly
+    multiprocessing.set_start_method("fork", force=True)
+
+
+class MultiGPURunner(Runner):
+    """Runner that runs the entire benchmark on multiple GPUs.
+
+    This is a thin wrapper around `Runner` that runs the entire benchmark on
+    multiple GPUs using `multiprocessing`.
+
+    Note that this runner will load multiple models into memory at the same
+    time if your running configuration specifies that, similar to the `Runner`
+    class. `SlurmRunner` on the other hand will load at most one model on a
+    GPU"""
+
+    def __init__(
+        self,
+        execution_spec: ExecutionSpec,
+        output_path: str,
+        suite: str,
+        skip_instances: bool,
+        cache_instances: bool,
+        cache_instances_only: bool,
+        skip_completed_runs: bool,
+        exit_on_error: bool,
+    ):
+        super().__init__(
+            execution_spec=execution_spec,
+            output_path=output_path,
+            suite=suite,
+            skip_instances=skip_instances,
+            cache_instances=cache_instances,
+            cache_instances_only=cache_instances_only,
+            skip_completed_runs=skip_completed_runs,
+            exit_on_error=exit_on_error,
+        )
+        # Configure max concurrent worker jobs from the environment variable.
+        env_max_concurrent_workers = os.getenv(_MAX_CONCURRENT_WORKERS_ENV_NAME)
+        self.max_concurrent_workers = (
+            int(env_max_concurrent_workers)
+            if env_max_concurrent_workers
+            else (
+                RUNNER_CONFIG.helm_max_concurrent_workers
+                if RUNNER_CONFIG.helm_max_concurrent_workers > 0
+                else torch.cuda.device_count()
+            )
+        )
+
+    def safe_run_one(self, run_spec: RunSpec):
+        register_builtin_configs_from_helm_package()
+        if self.executor.execution_spec.local_path is not None:
+            register_configs_from_directory(self.executor.execution_spec.local_path)
+
+        try:
+            with htrack_block(f"Running {run_spec.name}"):
+                self.run_one(run_spec)
+        except Exception as e:
+            hlog(f"Error when running {run_spec.name}:\n{traceback.format_exc()}")
+            return e
+
+    def run_all(self, run_specs: List[RunSpec]):
+        """Run the entire benchmark on multiple GPU"""
+
+        # Set the start method to forkserver to avoid issues with CUDA.
+        multiprocessing.set_start_method("forkserver")
+
+        with Pool(
+            max_workers=self.max_concurrent_workers,
+            initializer=start_thread_to_terminate_when_parent_process_dies,
+            initargs=(os.getpid(),),
+        ) as pool:
+            # Pin GPUs to each worker process
+            pool.map(initialize_worker, [i for i in range(self.max_concurrent_workers)])
+
+            # Run all queued tasks
+            error_msgs = list(tqdm(pool.map(self.safe_run_one, run_specs), total=len(run_specs), disable=None))
+
+        # Raise exception for failed runs, if any.
+        failed_run_names = [
+            run_spec.name for error_msg, run_spec in zip(error_msgs, run_specs) if error_msg is not None
+        ]
+        if failed_run_names:
+            failed_runs_str = ", ".join([f'"{run_name}"' for run_name in failed_run_names])
+            raise RunnerError(f"Failed runs: [{failed_runs_str}]")

helm/benchmark/presentation/contamination.py

@@ -2,10 +2,10 @@ from dataclasses import dataclass
 from typing import List, Optional
 import dacite
 import importlib_resources as resources
-import yaml  # type: ignore
+import yaml
 
 from helm.common.hierarchical_logger import htrack, hlog
-from helm.proxy.models import MODEL_NAME_TO_MODEL
+from helm.benchmark.model_metadata_registry import MODEL_NAME_TO_MODEL_METADATA
 from helm.benchmark.presentation.schema import Schema
 
 
@@ -70,7 +70,7 @@ def validate_contamination(contamination: Contamination, schema: Schema):
     """Make sure models and groups in contamination are defined according to `schema`."""
     for point in contamination.points:
         for model in point.models:
-            if model not in MODEL_NAME_TO_MODEL:
+            if model not in MODEL_NAME_TO_MODEL_METADATA:
                 hlog(f"WARNING: model {model} not defined in schema")
         for group in point.groups:
             if group not in schema.name_to_run_group: