Skip to content

Commit 4055fcf

Browse files
pintaoz-awspintaoz
andauthored
Fix get_child_process_ids parsing issue (#5416)
Co-authored-by: pintaoz <pintaoz@amazon.com>
1 parent fd37adc commit 4055fcf

File tree

2 files changed

+13
-6
lines changed

2 files changed

+13
-6
lines changed

sagemaker-core/src/sagemaker/core/local/utils.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -137,7 +137,11 @@ def get_child_process_ids(pid):
137137
Returns:
138138
(List[int]): Child process ids
139139
"""
140-
cmd = f"pgrep -P {pid}".split()
140+
if not str(pid).isdigit():
141+
raise ValueError("Invalid PID")
142+
143+
cmd = ["pgrep", "-P", str(pid)]
144+
141145
process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
142146
output, err = process.communicate()
143147
if err:

sagemaker-core/tests/unit/local/test_local_utils.py

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -103,21 +103,24 @@ def test_recursive_copy(copy_tree, m_os_path):
103103
@patch("sagemaker.core.local.utils.os")
104104
@patch("sagemaker.core.local.utils.get_child_process_ids")
105105
def test_kill_child_processes(m_get_child_process_ids, m_os):
106-
m_get_child_process_ids.return_value = ["child_pids"]
107-
kill_child_processes("pid")
108-
m_os.kill.assert_called_with("child_pids", 15)
106+
m_get_child_process_ids.return_value = ["345"]
107+
kill_child_processes("123")
108+
m_os.kill.assert_called_with("345", 15)
109109

110110

111111
@patch("sagemaker.core.local.utils.subprocess")
112112
def test_get_child_process_ids(m_subprocess):
113-
cmd = "pgrep -P pid".split()
113+
cmd = "pgrep -P 123".split()
114114
process_mock = Mock()
115115
attrs = {"communicate.return_value": (b"\n", False), "returncode": 0}
116116
process_mock.configure_mock(**attrs)
117117
m_subprocess.Popen.return_value = process_mock
118-
get_child_process_ids("pid")
118+
get_child_process_ids("123")
119119
m_subprocess.Popen.assert_called_with(cmd, stdout=m_subprocess.PIPE, stderr=m_subprocess.PIPE)
120120

121+
def test_get_child_process_ids_exception():
122+
with pytest.raises(ValueError, match="Invalid PID"):
123+
get_child_process_ids("abc")
121124

122125
@patch("sagemaker.core.local.utils.subprocess")
123126
def test_get_docker_host(m_subprocess):

0 commit comments

Comments
 (0)