Extend Modal runner timeout

Author: msaroufim

src/runners/modal_runner.py

@@ -85,6 +85,8 @@
     "modal_runner_archs",
 )
 
+MODAL_RUN_TIMEOUT_SECONDS = 60 * 60
+
 
 class TimeoutException(Exception):
     pass
@@ -111,7 +113,7 @@ def timeout_handler(signum, frame):
 
 def modal_run_config(  # noqa: C901
     config: dict,
-    timeout_seconds: int = 600,
+    timeout_seconds: int = MODAL_RUN_TIMEOUT_SECONDS,
 ) -> FullResult:
     """Modal version of run_pytorch_script, handling timeouts"""
     try:

src/runners/modal_runner_archs.py

@@ -1,13 +1,21 @@
 # This file contains wrapper functions for running
 # Modal apps on specific devices. We will fix this later.
-from modal_runner import app, cuda_image, modal_run_config
+from modal_runner import MODAL_RUN_TIMEOUT_SECONDS, app, cuda_image, modal_run_config
 
 gpus = ["T4", "L4", "L4:4", "A100-80GB", "H100!", "B200"]
 for gpu in gpus:
     gpu_slug = gpu.lower().split("-")[0].strip("!").replace(":", "x")
-    app.function(gpu=gpu, image=cuda_image, name=f"run_cuda_script_{gpu_slug}", serialized=True)(
-        modal_run_config
-    )
-    app.function(gpu=gpu, image=cuda_image, name=f"run_pytorch_script_{gpu_slug}", serialized=True)(
-        modal_run_config
-    )
+    app.function(
+        gpu=gpu,
+        image=cuda_image,
+        name=f"run_cuda_script_{gpu_slug}",
+        serialized=True,
+        timeout=MODAL_RUN_TIMEOUT_SECONDS,
+    )(modal_run_config)
+    app.function(
+        gpu=gpu,
+        image=cuda_image,
+        name=f"run_pytorch_script_{gpu_slug}",
+        serialized=True,
+        timeout=MODAL_RUN_TIMEOUT_SECONDS,
+    )(modal_run_config)