selection function

Author: jan-janssen

executorlib/executor/flux.py

@@ -360,14 +360,13 @@ def __init__(
                 create_file_executor,
             )
 
-            if terminate_tasks_on_shutdown:
-                from executorlib.task_scheduler.file.queue_spawner import (
-                    terminate_with_pysqa,
-                )
+            from executorlib.task_scheduler.file.queue_spawner import (
+                get_terminate_function,
+            )
 
-                terminate_function = terminate_with_pysqa
-            else:
-                terminate_function = None
+            terminate_function = get_terminate_function(
+                terminate_tasks_on_shutdown=terminate_tasks_on_shutdown
+            )
 
             super().__init__(
                 executor=create_file_executor(

executorlib/executor/slurm.py

@@ -166,14 +166,13 @@ def __init__(
                 create_file_executor,
             )
 
-            if terminate_tasks_on_shutdown:
-                from executorlib.task_scheduler.file.queue_spawner import (
-                    terminate_with_pysqa,
-                )
+            from executorlib.task_scheduler.file.queue_spawner import (
+                get_terminate_function,
+            )
 
-                terminate_function = terminate_with_pysqa
-            else:
-                terminate_function = None
+            terminate_function = get_terminate_function(
+                terminate_tasks_on_shutdown=terminate_tasks_on_shutdown
+            )
 
             super().__init__(
                 executor=create_file_executor(

executorlib/task_scheduler/file/queue_spawner.py

@@ -138,3 +138,17 @@ def _pysqa_execute_command(
         return out.split("\n")
     else:
         return out
+
+
+def get_terminate_function(terminate_tasks_on_shutdown: bool) -> Optional[callable]:
+    """
+    Get the appropriate terminate function based on the shutdown configuration.
+
+    Args:
+        terminate_tasks_on_shutdown (bool): Flag indicating whether to terminate tasks on shutdown.
+
+    Returns:
+        Optional[callable]: The terminate function to use, or None if no termination is needed.
+    """
+    if terminate_tasks_on_shutdown:
+        return terminate_with_pysqa