Merge pull request #493 from vyom1611/master-2

Modal benchmarking script updated to replace deprecated calls
PistonDevelopers · May 29, 2024 · 55f3665 · 55f3665
2 parents 2db7716 + 062f096
commit 55f3665
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/dev/cuda/benchmark_on_modal.py b/dev/cuda/benchmark_on_modal.py
@@ -27,7 +27,7 @@
 N_GPUS = int(os.environ.get("N_GPUS", 1))
 GPU_MEM = int(os.environ.get("GPU_MEM", 40))
 GPU_NAME = os.environ.get("GPU_NAME", "A100")
-GPU_CONFIG = GPU_NAME_TO_MODAL_CLASS_MAP[GPU_NAME](count=N_GPUS, memory=GPU_MEM)
+GPU_CONFIG = GPU_NAME_TO_MODAL_CLASS_MAP[GPU_NAME](count=N_GPUS, size=str(GPU_MEM)+'GB')
 
 APP_NAME = "llm.c benchmark run"
 
@@ -51,7 +51,7 @@
     )
 )
 
-stub = Stub(APP_NAME)
+stub = modal.App(APP_NAME)
 
 
 def execute_command(command: str):