[Fix] Set Pythonpath for TestClusterExecutor (#907)

* [Fix] Set Pythonpath for TestClusterExecutor

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* external function

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Add test for set_current_directory_in_environment

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: jan-janssen

src/executorlib/standalone/command.py

@@ -162,3 +162,15 @@ def generate_slurm_command(
     if slurm_cmd_args is not None and len(slurm_cmd_args) > 0:
         command_prepend_lst += slurm_cmd_args
     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/spawner_pysqa.py

@@ -3,6 +3,7 @@
 
 from pysqa import QueueAdapter
 
+from executorlib.standalone.command import set_current_directory_in_environment
 from executorlib.standalone.hdf import dump, get_queue_id
 from executorlib.standalone.inputcheck import check_file_exists
 from executorlib.standalone.scheduler import pysqa_execute_command, terminate_with_pysqa
@@ -85,6 +86,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

@@ -3,6 +3,7 @@
 import time
 from typing import Optional
 
+from executorlib.standalone.command import set_current_directory_in_environment
 from executorlib.standalone.hdf import dump
 from executorlib.standalone.inputcheck import check_file_exists
 
@@ -53,11 +54,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 +73,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

tests/unit/standalone/test_command.py

@@ -2,8 +2,11 @@
 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, 
+    set_current_directory_in_environment,
+)
 
 skip_mpi4py_test = importlib.util.find_spec("mpi4py") is None
 
@@ -92,3 +95,18 @@ def test_get_cache_execute_command_parallel(self):
             get_cache_execute_command(cores=2, file_name=file_name, backend="flux", openmpi_oversubscribe=True)
         with self.assertRaises(ValueError):
             get_cache_execute_command(cores=2, file_name=file_name, backend="flux", exclusive=True)
+
+    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")