From 702556e5ed216912a7f2625eab345c0338eca19c Mon Sep 17 00:00:00 2001 From: Walter Simson Date: Tue, 29 Oct 2024 12:54:49 -0700 Subject: [PATCH] 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 --- kwave/executor.py | 8 +-- kwave/options/simulation_execution_options.py | 28 +++----- tests/test_executor.py | 6 +- tests/test_simulation_execution_options.py | 71 +++++++++++++++++++ 4 files changed, 87 insertions(+), 26 deletions(-) create mode 100644 tests/test_simulation_execution_options.py diff --git a/kwave/executor.py b/kwave/executor.py index b0799a01..766a6810 100644 --- 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 index 702b0789..3dcf42c6 100644 --- 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 index b92b5859..976d8058 100644 --- 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 new file mode 100644 index 00000000..2da09e72 --- /dev/null +++ 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()