Merge branch 'main' into Py314u

Author: jan-janssen

src/executorlib/standalone/interactive/spawner.py

@@ -112,6 +112,7 @@ def bootup(
         """
         if self._cwd is not None:
             os.makedirs(self._cwd, exist_ok=True)
+        set_current_directory_in_environment()
         self._process = subprocess.Popen(
             args=self.generate_command(command_lst=command_lst),
             cwd=self._cwd,
@@ -195,3 +196,15 @@ def generate_mpiexec_command(
         if openmpi_oversubscribe:
             command_prepend_lst += ["--oversubscribe"]
         return command_prepend_lst
+
+
+def set_current_directory_in_environment():
+    """
+    Add the current directory to the PYTHONPATH to be able to access local Python modules.
+    """
+    environment = os.environ
+    current_path = os.getcwd()
+    if "PYTHONPATH" in environment and current_path not in environment["PYTHONPATH"]:
+        environment["PYTHONPATH"] = os.getcwd() + ":" + environment["PYTHONPATH"]
+    elif "PYTHONPATH" not in environment:
+        environment["PYTHONPATH"] = os.getcwd()

src/executorlib/task_scheduler/file/shared.py

@@ -90,39 +90,22 @@ def execute_tasks_h5(
         if task_dict is not None and "shutdown" in task_dict and task_dict["shutdown"]:
             if task_dict["wait"] and wait:
                 while len(memory_dict) > 0:
-                    memory_dict = {
-                        key: _check_task_output(
-                            task_key=key,
-                            future_obj=value,
-                            cache_directory=cache_dir_dict[key],
-                        )
-                        for key, value in memory_dict.items()
-                        if not value.done()
-                    }
+                    memory_dict = _refresh_memory_dict(
+                        memory_dict=memory_dict,
+                        cache_dir_dict=cache_dir_dict,
+                    )
             if not task_dict["cancel_futures"] and wait:
-                if (
-                    terminate_function is not None
-                    and terminate_function == terminate_subprocess
-                ):
-                    for task in process_dict.values():
-                        terminate_function(task=task)
-                elif terminate_function is not None:
-                    for queue_id in process_dict.values():
-                        terminate_function(
-                            queue_id=queue_id,
-                            config_directory=pysqa_config_directory,
-                            backend=backend,
-                        )
+                _cancel_processes(
+                    terminate_function=terminate_function,
+                    process_dict=process_dict,
+                    pysqa_config_directory=pysqa_config_directory,
+                    backend=backend,
+                )
             else:
-                memory_dict = {
-                    key: _check_task_output(
-                        task_key=key,
-                        future_obj=value,
-                        cache_directory=cache_dir_dict[key],
-                    )
-                    for key, value in memory_dict.items()
-                    if not value.done()
-                }
+                memory_dict = _refresh_memory_dict(
+                    memory_dict=memory_dict,
+                    cache_dir_dict=cache_dir_dict,
+                )
                 for value in memory_dict.values():
                     if not value.done():
                         value.cancel()
@@ -193,15 +176,10 @@ def execute_tasks_h5(
                 cache_dir_dict[task_key] = cache_directory
             future_queue.task_done()
         else:
-            memory_dict = {
-                key: _check_task_output(
-                    task_key=key,
-                    future_obj=value,
-                    cache_directory=cache_dir_dict[key],
-                )
-                for key, value in memory_dict.items()
-                if not value.done()
-            }
+            memory_dict = _refresh_memory_dict(
+                memory_dict=memory_dict,
+                cache_dir_dict=cache_dir_dict,
+            )
 
 
 def _check_task_output(
@@ -275,3 +253,52 @@ def _convert_args_and_kwargs(
         else:
             task_kwargs[key] = arg
     return task_args, task_kwargs, future_wait_key_lst
+
+
+def _refresh_memory_dict(memory_dict: dict, cache_dir_dict: dict) -> dict:
+    """
+    Refresh memory dictionary
+
+    Args:
+        memory_dict (dict): dictionary with task keys and future objects
+        cache_dir_dict (dict): dictionary with task keys and cache directories
+
+    Returns:
+        dict: Updated memory dictionary
+    """
+    return {
+        key: _check_task_output(
+            task_key=key,
+            future_obj=value,
+            cache_directory=cache_dir_dict[key],
+        )
+        for key, value in memory_dict.items()
+        if not value.done()
+    }
+
+
+def _cancel_processes(
+    process_dict: dict,
+    terminate_function: Optional[Callable] = None,
+    pysqa_config_directory: Optional[str] = None,
+    backend: Optional[str] = None,
+):
+    """
+    Cancel processes
+
+    Args:
+        process_dict (dict): dictionary with task keys and process reference.
+        terminate_function (callable): The function to terminate the tasks.
+        pysqa_config_directory (str): path to the pysqa config directory (only for pysqa based backend).
+        backend (str): name of the backend used to spawn tasks.
+    """
+    if terminate_function is not None and terminate_function == terminate_subprocess:
+        for task in process_dict.values():
+            terminate_function(task=task)
+    elif terminate_function is not None and backend is not None:
+        for queue_id in process_dict.values():
+            terminate_function(
+                queue_id=queue_id,
+                config_directory=pysqa_config_directory,
+                backend=backend,
+            )

src/executorlib/task_scheduler/file/spawner_pysqa.py

@@ -5,6 +5,9 @@
 
 from executorlib.standalone.hdf import dump, get_queue_id
 from executorlib.standalone.inputcheck import check_file_exists
+from executorlib.standalone.interactive.spawner import (
+    set_current_directory_in_environment,
+)
 from executorlib.standalone.scheduler import pysqa_execute_command, terminate_with_pysqa
 
 
@@ -85,6 +88,7 @@ def execute_with_pysqa(
                 os.path.dirname(os.path.abspath(cwd))
             )
         submit_kwargs.update(resource_dict)
+        set_current_directory_in_environment()
         queue_id = qa.submit_job(**submit_kwargs)
         dump(file_name=file_name, data_dict={"queue_id": queue_id})
     return queue_id

src/executorlib/task_scheduler/file/spawner_subprocess.py

@@ -5,6 +5,9 @@
 
 from executorlib.standalone.hdf import dump
 from executorlib.standalone.inputcheck import check_file_exists
+from executorlib.standalone.interactive.spawner import (
+    set_current_directory_in_environment,
+)
 
 
 def execute_in_subprocess(
@@ -53,11 +56,12 @@ def execute_in_subprocess(
         )
     if backend is not None:
         raise ValueError("backend parameter is not supported for subprocess spawner.")
-    if resource_dict is None:
-        resource_dict = {}
-    cwd = resource_dict.get("cwd", cache_directory)
+    cwd = _get_working_directory(
+        cache_directory=cache_directory, resource_dict=resource_dict
+    )
     if cwd is not None:
         os.makedirs(cwd, exist_ok=True)
+    set_current_directory_in_environment()
     return subprocess.Popen(command, universal_newlines=True, cwd=cwd)
 
 
@@ -71,3 +75,14 @@ def terminate_subprocess(task):
     task.terminate()
     while task.poll() is None:
         time.sleep(0.1)
+
+
+def _get_working_directory(
+    cache_directory: Optional[str] = None, resource_dict: Optional[dict] = None
+):
+    if resource_dict is None:
+        resource_dict = {}
+    if "cwd" in resource_dict and resource_dict["cwd"] is not None:
+        return resource_dict["cwd"]
+    else:
+        return cache_directory

src/executorlib/task_scheduler/interactive/spawner_flux.py

@@ -5,7 +5,10 @@
 import flux
 import flux.job
 
-from executorlib.standalone.interactive.spawner import BaseSpawner
+from executorlib.standalone.interactive.spawner import (
+    BaseSpawner,
+    set_current_directory_in_environment,
+)
 
 
 def validate_max_workers(max_workers: int, cores: int, threads_per_core: int):
@@ -118,6 +121,7 @@ def bootup(
                 num_nodes=self._num_nodes,
                 exclusive=self._exclusive,
             )
+        set_current_directory_in_environment()
         jobspec.environment = dict(os.environ)
         if self._pmi_mode is not None:
             jobspec.setattr_shell_option("pmi", self._pmi_mode)

src/executorlib/task_scheduler/interactive/spawner_pysqa.py

@@ -6,7 +6,10 @@
 from pysqa import QueueAdapter
 
 from executorlib.standalone.inputcheck import validate_number_of_cores
-from executorlib.standalone.interactive.spawner import BaseSpawner
+from executorlib.standalone.interactive.spawner import (
+    BaseSpawner,
+    set_current_directory_in_environment,
+)
 from executorlib.standalone.scheduler import pysqa_execute_command, terminate_with_pysqa
 from executorlib.task_scheduler.interactive.blockallocation import (
     BlockAllocationTaskScheduler,
@@ -183,6 +186,7 @@ def _start_process_helper(
             working_directory = os.path.join(self._cwd, hash)
         else:
             working_directory = os.path.abspath(hash)
+        set_current_directory_in_environment()
         return queue_adapter.submit_job(
             command=" ".join(self.generate_command(command_lst=command_lst)),
             working_directory=working_directory,

tests/unit/executor/test_flux_cluster.py

@@ -30,6 +30,11 @@ def echo(i):
     return i
 
 
+def long_running_function(i):
+    sleep(10)
+    return i
+
+
 def mpi_funct(i):
     from mpi4py import MPI
 

tests/unit/standalone/interactive/test_spawner.py

@@ -3,13 +3,14 @@
 from queue import Queue
 from time import sleep
 import shutil
+import os
 from sys import platform
 import unittest
 
 import numpy as np
 
 from executorlib.task_scheduler.base import TaskSchedulerBase
-from executorlib.standalone.interactive.spawner import MpiExecSpawner
+from executorlib.standalone.interactive.spawner import MpiExecSpawner, set_current_directory_in_environment
 from executorlib.task_scheduler.interactive.blockallocation import BlockAllocationTaskScheduler, _execute_multiple_tasks
 from executorlib.task_scheduler.interactive.onetoone import OneProcessTaskScheduler
 from executorlib.standalone.interactive.backend import call_funct
@@ -544,3 +545,22 @@ def test_execute_task_cache_failed_no_argument(self):
         with self.assertRaises(TypeError):
             f.result()
         q.join()
+
+
+class TestEnvManipulation(unittest.TestCase):
+    def test_set_current_directory_in_environment(self):
+        env = os.environ
+        if "PYTHONPATH" in env:
+            python_path = env["PYTHONPATH"]
+            del env["PYTHONPATH"]
+        else:
+            python_path = None
+        self.assertFalse("PYTHONPATH" in env)
+        set_current_directory_in_environment()
+        self.assertTrue("PYTHONPATH" in env)
+        self.assertEqual(env["PYTHONPATH"], os.getcwd())
+        env["PYTHONPATH"] = "/my/special/path"
+        set_current_directory_in_environment()
+        self.assertEqual(env["PYTHONPATH"], os.getcwd() + ":/my/special/path")
+        if python_path is not None:
+            env["PYTHONPATH"] = python_path

tests/unit/standalone/test_command.py

@@ -2,8 +2,10 @@
 import importlib.util
 import sys
 import unittest
-from executorlib.standalone.command import get_cache_execute_command, get_interactive_execute_command
-
+from executorlib.standalone.command import (
+    get_cache_execute_command,
+    get_interactive_execute_command,
+)
 
 skip_mpi4py_test = importlib.util.find_spec("mpi4py") is None