diff --git a/src/runners/modal_runner.py b/src/runners/modal_runner.py index f093ae9ba..488abf62e 100644 --- a/src/runners/modal_runner.py +++ b/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: diff --git a/src/runners/modal_runner_archs.py b/src/runners/modal_runner_archs.py index f1557f5b5..e768020ca 100644 --- a/src/runners/modal_runner_archs.py +++ b/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)