def get_config()

in src/nccl_healthcheck/config.py [0:0]


def get_config(instance_type: str) -> config_pb2.ASeriesConfig:
  """Returns the config for the given instance type.

  Args:
    instance_type: The instance type to get the config for.

  Returns:
    The config for the given instance type.

  Raises:
    ValueError: If the instance type is not supported.
  """
  if instance_type == "a3-highgpu-8g":
    return _create_a3_config()
  elif instance_type == "a3-megagpu-8g":
    return _create_a3plus_config()
  elif instance_type == "a3-megagpu-8g-debian":
    return _create_a3plus_debian_config()
  elif instance_type == "a3-ultragpu-8g":
    return _create_a3ultra_config()
  elif instance_type == "a4-highgpu-8g":
    return _create_a4_config()
  else:
    raise ValueError(f"Unsupported instance type: {instance_type}")