From cf9bc255ab1ec1b9ebac03fc47e1bb31595f0ee7 Mon Sep 17 00:00:00 2001 From: Eric Jelli Date: Wed, 24 Jul 2024 14:57:14 +0200 Subject: [PATCH] feature: Expose squeue polling interval to SlurmExecutor via parameter and env variable --- cluster_tools/cluster_tools/schedulers/cluster_executor.py | 3 ++- cluster_tools/cluster_tools/schedulers/slurm.py | 4 +++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/cluster_tools/cluster_tools/schedulers/cluster_executor.py b/cluster_tools/cluster_tools/schedulers/cluster_executor.py index b8ad1ab81..0a13aaa66 100644 --- a/cluster_tools/cluster_tools/schedulers/cluster_executor.py +++ b/cluster_tools/cluster_tools/schedulers/cluster_executor.py @@ -95,6 +95,7 @@ def __init__( job_resources: Optional[Dict[str, Any]] = None, job_name: Optional[str] = None, additional_setup_lines: Optional[List[str]] = None, + interval: Optional[int] = 2, **kwargs: Any, ): """ @@ -133,7 +134,7 @@ def __init__( self.keep_logs = keep_logs self.is_shutting_down = False - self.wait_thread = FileWaitThread(self._completion, self) + self.wait_thread = FileWaitThread(self._completion, self, interval) self.wait_thread.start() os.makedirs(self.cfut_dir, exist_ok=True) diff --git a/cluster_tools/cluster_tools/schedulers/slurm.py b/cluster_tools/cluster_tools/schedulers/slurm.py index 8c72eecf6..94c2bb7db 100644 --- a/cluster_tools/cluster_tools/schedulers/slurm.py +++ b/cluster_tools/cluster_tools/schedulers/slurm.py @@ -61,7 +61,7 @@ "Unclear": ["SUSPENDED", "REVOKED", "SIGNALING", "SPECIAL_EXIT", "STAGE_OUT"], } -SLURM_QUEUE_CHECK_INTERVAL = 1 if "pytest" in sys.modules else 60 +SLURM_QUEUE_CHECK_INTERVAL = 1 if "pytest" in sys.modules else os.environ.get("SLURM_QUEUE_CHECK_INTERVAL", 60) T = TypeVar("T") @@ -84,6 +84,7 @@ def __init__( job_resources: Optional[Dict[str, Any]] = None, job_name: Optional[str] = None, additional_setup_lines: Optional[List[str]] = None, + interval: Optional[int] = None, **kwargs: Any, ): super().__init__( @@ -93,6 +94,7 @@ def __init__( job_resources=job_resources, job_name=job_name, additional_setup_lines=additional_setup_lines, + interval=interval if interval else SLURM_QUEUE_CHECK_INTERVAL, **kwargs, ) self.submit_threads: List["_JobSubmitThread"] = []