meg/backend

1from src.backend.model_utils import calculate_memory, get_model2from src.logging import setup_logger3 4logger = setup_logger(__name__)5 6 7def get_instance_needs(model_name: str, access_token: str, region='us-east-1', vendor='aws'):8    """Scales up compute based on size and price."""9    needed_space = get_size(model_name, access_token)10    if needed_space:11        # AWS is the only thing I've implemented this for for now.12        if region =='us-east-1' and vendor == 'aws':13            if needed_space < 20:14                # Cheapest15                return 'x1', 'nvidia-a10g'16            elif needed_space < 60:17                return 'x4', 'nvidia-t4'18            elif needed_space < 80:19                return 'x1', 'nvidia-a100'20            elif needed_space < 95:21                return 'x4', 'nvidia-a10g'22            elif needed_space < 150:23                return 'x2', 'nvidia-a100'24            # Not doing any higher (for now) as that would start costing a lot.25        else:26            logger.warning("Not implemented for region %s vendor %s" % (region, vendor))27            logger.warning("Only implemented for aws us-east-1. Pretending that's what you asked for.")28            return get_instance_needs(model_name=model_name, access_token=access_token)29    else:30        # A default size to start trying to scale up from.31        return 'x4', 'nvidia-l4'32 33 34# Code based in part on https://huggingface.co/spaces/hf-accelerate/model-memory-usage35def get_size(model_name: str, access_token: str, library='auto',36             dtype='float32'):37    """38    This is just to get a size estimate of the model.39    Assuming dtype float32, which isn't always true.40    Only works for transformers and timm models AFAIK.41    """42    model = get_model(model_name, library, access_token)43    data = calculate_memory(model, dtype)44    size = data[0]['Total Size']45    split_size = size.split()46    # Assuming we're working in GB.47    try:48        assert split_size[1] == 'GB'49        num_gigs = float(split_size[0])50    except AssertionError:51        logger.warning(52            "Tried to get model size and it's not GB, it's %s" % size)53        logger.warning(54            "Have not implemented handling for this, just going with 30GB.")55        num_gigs = 3056    return num_gigs57 58 59if __name__ == '__main__':60    # Debugging here61    import os62 63    num_gigs_debug = get_size('upstage/SOLAR-10.7B-v1.0',64                              access_token=os.environ.get('HF_TOKEN'))65    print(num_gigs_debug)66