Dummy for more regions/vendors.

src/backend/compute_memory_requirements.py

@@ -4,30 +4,36 @@ from src.logging import setup_logger
 logger = setup_logger(__name__)
 
 
-def get_instance_needs(model_name: str, access_token: str):
+def get_instance_needs(model_name: str, access_token: str, region='us-east-1', vendor='aws'):
     """Scales up compute based on size and price."""
     needed_space = get_size(model_name, access_token)
     if needed_space:
-        if needed_space < 20:
-            # Cheapest
-            return 'x1', 'nvidia-a10g'
-        elif needed_space < 60:
-            return 'x4', 'nvidia-t4'
-        elif needed_space < 80:
-            return 'x1', 'nvidia-a100'
-        elif needed_space < 95:
-            return 'x4', 'nvidia-a10g'
-        elif needed_space < 150:
-            return 'x2', 'nvidia-a100'
-        # Not doing any higher (for now) as that would start costing a lot.
+        # AWS is the only thing I've implemented this for for now.
+        if region =='us-east-1' and vendor == 'aws':
+            if needed_space < 20:
+                # Cheapest
+                return 'x1', 'nvidia-a10g'
+            elif needed_space < 60:
+                return 'x4', 'nvidia-t4'
+            elif needed_space < 80:
+                return 'x1', 'nvidia-a100'
+            elif needed_space < 95:
+                return 'x4', 'nvidia-a10g'
+            elif needed_space < 150:
+                return 'x2', 'nvidia-a100'
+            # Not doing any higher (for now) as that would start costing a lot.
+        else:
+            logger.warning("Not implemented for region %s vendor %s" % (region, vendor))
+            logger.warning("Only implemented for aws us-east-1. Pretending that's what you asked for.")
+            return get_instance_needs(model_name=model_name, access_token=access_token)
     else:
         # A default size to start trying to scale up from.
         return 'x4', 'nvidia-l4'
 
 
 # Code based in part on https://huggingface.co/spaces/hf-accelerate/model-memory-usage
-def get_size(model_name: str, access_token: str, library="auto",
-             dtype="float32"):
+def get_size(model_name: str, access_token: str, library='auto',
+             dtype='float32'):
     """
     This is just to get a size estimate of the model.
     Assuming dtype float32, which isn't always true.
@@ -54,6 +60,6 @@ if __name__ == '__main__':
     # Debugging here
     import os
 
-    num_gigs_debug = get_size("upstage/SOLAR-10.7B-v1.0",
-                              access_token=os.environ.get("HF_TOKEN"))
+    num_gigs_debug = get_size('upstage/SOLAR-10.7B-v1.0',
+                              access_token=os.environ.get('HF_TOKEN'))
     print(num_gigs_debug)