From 977871cccb51344544554efd9ea540555e2de663 Mon Sep 17 00:00:00 2001 From: Oliver Wissett Date: Fri, 4 Oct 2024 19:19:20 +0100 Subject: [PATCH] Adds non-blocking mode to hydra_submitit_launcher (#2479) --- .../hydra_plugins/hydra_submitit_launcher/config.py | 3 +++ .../hydra_submitit_launcher/submitit_launcher.py | 7 ++++++- .../tests/test_submitit_launcher.py | 13 +++++++++++++ 3 files changed, 22 insertions(+), 1 deletion(-) diff --git a/plugins/hydra_submitit_launcher/hydra_plugins/hydra_submitit_launcher/config.py b/plugins/hydra_submitit_launcher/hydra_plugins/hydra_submitit_launcher/config.py index 6da5658a86..e4da9957bd 100644 --- a/plugins/hydra_submitit_launcher/hydra_plugins/hydra_submitit_launcher/config.py +++ b/plugins/hydra_submitit_launcher/hydra_plugins/hydra_submitit_launcher/config.py @@ -28,6 +28,9 @@ class BaseQueueConf: # redirect stderr to stdout stderr_to_stdout: bool = False + # If True, the launcher will not wait for the job to finish (useful for very long runs) + no_block: bool = False + @dataclass class SlurmQueueConf(BaseQueueConf): diff --git a/plugins/hydra_submitit_launcher/hydra_plugins/hydra_submitit_launcher/submitit_launcher.py b/plugins/hydra_submitit_launcher/hydra_plugins/hydra_submitit_launcher/submitit_launcher.py index 1efc8e4ce8..e97e1cad2a 100644 --- a/plugins/hydra_submitit_launcher/hydra_plugins/hydra_submitit_launcher/submitit_launcher.py +++ b/plugins/hydra_submitit_launcher/hydra_plugins/hydra_submitit_launcher/submitit_launcher.py @@ -19,6 +19,7 @@ class BaseSubmititLauncher(Launcher): _EXECUTOR = "abstract" def __init__(self, **params: Any) -> None: + self.no_block = params.pop("no_block", False) self.params = {} for k, v in params.items(): if OmegaConf.is_config(v): @@ -142,7 +143,11 @@ def launch( ) jobs = executor.map_array(self, *zip(*job_params)) - return [j.results()[0] for j in jobs] + + if self.no_block: + return [] + else: + return [j.results()[0] for j in jobs] class LocalLauncher(BaseSubmititLauncher): diff --git a/plugins/hydra_submitit_launcher/tests/test_submitit_launcher.py b/plugins/hydra_submitit_launcher/tests/test_submitit_launcher.py index 65f7a0e6a4..7d9bd943bc 100644 --- a/plugins/hydra_submitit_launcher/tests/test_submitit_launcher.py +++ b/plugins/hydra_submitit_launcher/tests/test_submitit_launcher.py @@ -61,3 +61,16 @@ def test_example(tmpdir: Path) -> None: ], allow_warnings=True, ) + + +def test_example_no_block(tmpdir: Path) -> None: + run_python_script( + [ + "example/my_app.py", + "-m", + f"hydra.sweep.dir={tmpdir}", + "hydra/launcher=submitit_local", + "hydra.launcher.no_block=True", + ], + allow_warnings=True, + )