Update environment variable generation for safer execution (#493)

* Update executor.py to handle spaces in binary path * Update test_executor.py for new path * Update test_executor.py * Update test_executor.py to pass test * Update test_executor.py * update environment variables to remove shell=True * pass binary path as a string * add testing * specify system for default initialization * update gpu default argument test --------- Co-authored-by: David Sinden <[email protected]>
waltsims · Oct 29, 2024 · 702556e · 702556e
1 parent 1a4bba9
commit 702556e
Show file tree

Hide file tree

Showing 4 changed files with 87 additions and 26 deletions.
diff --git a/kwave/executor.py b/kwave/executor.py
@@ -27,12 +27,12 @@ def _make_binary_executable(self):
                 raise e
 
     def run_simulation(self, input_filename: str, output_filename: str, options: str):
-        command = (
-            f"{self.execution_options.system_string}" f'"{self.execution_options.binary_path}" -i {input_filename} -o {output_filename} {options}'
-        )
+        command = [str(self.execution_options.binary_path), "-i", input_filename, "-o", output_filename, options]
 
         try:
-            with subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True, text=True) as proc:
+            with subprocess.Popen(
+                command, env=self.execution_options.env_vars, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True
+            ) as proc:
                 stdout, stderr = "", ""
                 if self.execution_options.show_sim_log:
                     # Stream stdout in real-time

diff --git a/kwave/options/simulation_execution_options.py b/kwave/options/simulation_execution_options.py
@@ -1,9 +1,9 @@
 from dataclasses import dataclass
+import os
 from typing import Optional, Union
 
 from kwave import PLATFORM, BINARY_PATH
 from kwave.ksensor import kSensor
-from kwave.utils.checks import is_unix
 
 
 @dataclass
@@ -117,34 +117,22 @@ def get_options_string(self, sensor: kSensor) -> str:
         return options_string
 
     @property
-    def system_string(self):
-        # set OS string for setting environment variables
-        if is_unix():
-            env_set_str = ""
-            sys_sep_str = " "
-        else:
-            env_set_str = "set "
-            sys_sep_str = " & "
+    def env_vars(self) -> dict:
+        env = os.environ
 
-        # set system string to define domain for thread migration
-        system_string = env_set_str
         if PLATFORM != "darwin":
-            system_string += "OMP_PLACES=cores" + sys_sep_str
+            env.update({"OMP_PLACES": "cores"})
 
         if self.thread_binding is not None:
             if PLATFORM == "darwin":
                 raise ValueError("Thread binding is not supported in MacOS.")
             # read the parameters and update the system options
             if self.thread_binding:
-                system_string = system_string + " " + env_set_str + "OMP_PROC_BIND=SPREAD" + sys_sep_str
+                env.update({"OMP_PROC_BIND": "SPREAD"})
             else:
-                system_string = system_string + " " + env_set_str + "OMP_PROC_BIND=CLOSE" + sys_sep_str
+                env.update({"OMP_PROC_BIND": "CLOSE"})
         else:
             if PLATFORM != "darwin":
-                # set to round-robin over places
-                system_string = system_string + " " + env_set_str + "OMP_PROC_BIND=SPREAD" + sys_sep_str
-
-        if self.system_call:
-            system_string = system_string + " " + self.system_call + sys_sep_str
+                env.update({"OMP_PROC_BIND": "SPREAD"})
 
-        return system_string
+        return env
diff --git a/tests/test_executor.py b/tests/test_executor.py
@@ -99,9 +99,11 @@ def test_run_simulation_success(self):
             sensor_data = executor.run_simulation("input.h5", "output.h5", "options")
 
         normalized_path = os.path.normpath(self.execution_options.binary_path)
-        expected_command = f"{self.execution_options.system_string}" f'"{normalized_path}" -i input.h5 ' f"-o output.h5 " f"options"
+        expected_command = [normalized_path, "-i", "input.h5", "-o", "output.h5", "options"]
 
-        self.mock_popen.assert_called_once_with(expected_command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True, text=True)
+        self.mock_popen.assert_called_once_with(
+            expected_command, env=self.execution_options.env_vars, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True
+        )
         self.mock_proc.communicate.assert_called_once()
         self.assertEqual(sensor_data, dotdict())
 

diff --git a/tests/test_simulation_execution_options.py b/tests/test_simulation_execution_options.py
@@ -0,0 +1,71 @@
+import unittest
+from kwave.options.simulation_execution_options import SimulationExecutionOptions
+from kwave.ksensor import kSensor
+from unittest.mock import patch
+
+
+class TestSimulationExecutionOptions(unittest.TestCase):
+    def setUp(self):
+        # Set up a default kSensor object for testing
+        self.sensor = kSensor()
+        self.sensor.record = ["p", "u"]
+        self.sensor.record_start_index = 10
+
+    def test_default_initialization(self):
+        with patch("kwave.options.simulation_execution_options.PLATFORM", "linux"):
+            options = SimulationExecutionOptions()
+            self.assertFalse(options.is_gpu_simulation)
+            self.assertEqual(options.binary_name, "kspaceFirstOrder-OMP")
+            self.assertTrue(options.delete_data)
+            self.assertEqual(options.num_threads, None)  # TODO: confusing logic here
+            self.assertEqual(options.verbose_level, 0)
+
+    def test_gpu_simulation_initialization(self):
+        with patch("kwave.options.simulation_execution_options.PLATFORM", "linux"):
+            options = SimulationExecutionOptions(is_gpu_simulation=True)
+            self.assertTrue(options.is_gpu_simulation)
+            self.assertEqual(options.binary_name, "kspaceFirstOrder-CUDA")
+
+    def test_binary_name_extension_on_windows(self):
+        with patch("kwave.options.simulation_execution_options.PLATFORM", "windows"):
+            options = SimulationExecutionOptions()
+            self.assertTrue(options.binary_name.endswith(".exe"))
+
+    def test_get_options_string(self):
+        options = SimulationExecutionOptions()
+        options_string = options.get_options_string(self.sensor)
+        self.assertIn("--p_raw", options_string)
+        self.assertIn("--u_raw", options_string)
+        self.assertIn("-s 10", options_string)
+
+    def test_env_vars_linux(self):
+        with patch("kwave.options.simulation_execution_options.PLATFORM", "linux"):
+            options = SimulationExecutionOptions()
+            env_vars = options.env_vars
+            self.assertIn("OMP_PLACES", env_vars)
+            self.assertEqual(env_vars["OMP_PLACES"], "cores")
+            self.assertIn("OMP_PROC_BIND", env_vars)
+            self.assertEqual(env_vars["OMP_PROC_BIND"], "SPREAD")
+
+    def test_thread_binding_linux(self):
+        with patch("kwave.options.simulation_execution_options.PLATFORM", "linux"):
+            options = SimulationExecutionOptions(thread_binding=True)
+            env_vars = options.env_vars
+            self.assertEqual(env_vars["OMP_PROC_BIND"], "SPREAD")
+
+    def test_thread_binding_darwin(self):
+        with patch("kwave.options.simulation_execution_options.PLATFORM", "darwin"):
+            options = SimulationExecutionOptions(thread_binding=True)
+            with self.assertRaises(ValueError, msg="Thread binding is not supported in MacOS."):
+                _ = options.env_vars
+
+    def test_env_vars_darwin(self):
+        with patch("kwave.options.simulation_execution_options.PLATFORM", "darwin"):
+            options = SimulationExecutionOptions()
+            env_vars = options.env_vars
+            self.assertNotIn("OMP_PLACES", env_vars)
+            self.assertNotIn("OMP_PROC_BIND", env_vars)
+
+
+if __name__ == "__main__":
+    unittest.main()