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}")