restructure

Author: jan-janssen

executorlib/executor/flux.py

@@ -360,14 +360,6 @@ def __init__(
                 create_file_executor,
             )
 
-            from executorlib.task_scheduler.file.queue_spawner import (
-                get_terminate_function,
-            )
-
-            terminate_function = get_terminate_function(
-                terminate_tasks_on_shutdown=terminate_tasks_on_shutdown
-            )
-
             super().__init__(
                 executor=create_file_executor(
                     max_workers=max_workers,
@@ -384,7 +376,7 @@ def __init__(
                     block_allocation=block_allocation,
                     init_function=init_function,
                     disable_dependencies=disable_dependencies,
-                    terminate_function=terminate_function,
+                    terminate_tasks_on_shutdown=terminate_tasks_on_shutdown,
                 )
             )
         else:

executorlib/executor/slurm.py

@@ -166,14 +166,6 @@ def __init__(
                 create_file_executor,
             )
 
-            from executorlib.task_scheduler.file.queue_spawner import (
-                get_terminate_function,
-            )
-
-            terminate_function = get_terminate_function(
-                terminate_tasks_on_shutdown=terminate_tasks_on_shutdown
-            )
-
             super().__init__(
                 executor=create_file_executor(
                     max_workers=max_workers,
@@ -190,7 +182,7 @@ def __init__(
                     block_allocation=block_allocation,
                     init_function=init_function,
                     disable_dependencies=disable_dependencies,
-                    terminate_function=terminate_function,
+                    terminate_tasks_on_shutdown=terminate_tasks_on_shutdown,
                 )
             )
         else:

executorlib/task_scheduler/file/queue_spawner.py

@@ -138,19 +138,3 @@ 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
-    else:
-        return None

executorlib/task_scheduler/file/task_scheduler.py

@@ -12,14 +12,15 @@
 from executorlib.task_scheduler.base import TaskSchedulerBase
 from executorlib.task_scheduler.file.shared import execute_tasks_h5
 from executorlib.task_scheduler.file.subprocess_spawner import (
-    execute_in_subprocess,
+    execute_in_subprocess, terminate_subprocess
 )
 
 try:
-    from executorlib.task_scheduler.file.queue_spawner import execute_with_pysqa
+    from executorlib.task_scheduler.file.queue_spawner import execute_with_pysqa, terminate_with_pysqa
 except ImportError:
     # If pysqa is not available fall back to executing tasks in a subprocess
     execute_with_pysqa = execute_in_subprocess  # type: ignore
+    terminate_with_pysqa = None
 
 
 class FileTaskScheduler(TaskSchedulerBase):
@@ -90,7 +91,7 @@ def create_file_executor(
     init_function: Optional[Callable] = None,
     disable_dependencies: bool = False,
     execute_function: Callable = execute_with_pysqa,
-    terminate_function: Optional[Callable] = None,
+    terminate_tasks_on_shutdown: bool = True,
 ):
     if block_allocation:
         raise ValueError(
@@ -108,6 +109,12 @@ def create_file_executor(
     check_executor(executor=flux_executor)
     check_nested_flux_executor(nested_flux_executor=flux_executor_nesting)
     check_flux_log_files(flux_log_files=flux_log_files)
+    if terminate_tasks_on_shutdown and execute_function != execute_in_subprocess:
+        terminate_function = terminate_with_pysqa
+    elif terminate_tasks_on_shutdown and execute_function == execute_in_subprocess:
+        terminate_function = terminate_subprocess
+    else:
+        terminate_function = None
     return FileTaskScheduler(
         resource_dict=resource_dict,
         pysqa_config_directory=pysqa_config_directory,