Skip to content

Commit 80e11d3

Browse files
karajan1001pmrowla
andcommitted
Automatically get group id by pid, and change some default value
When using killpg, we need to call with group id instead of pid 1. Get group id by pid for `killpg` 2. Modify `kill` and `terminate` 's default group to `False` for better compatibility. Co-authored-by: Peter Rowlands (변기호) <[email protected]>
1 parent b5cc7d4 commit 80e11d3

File tree

2 files changed

+36
-15
lines changed

2 files changed

+36
-15
lines changed

src/dvc_task/proc/manager.py

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -111,7 +111,7 @@ def run_signature(
111111
immutable=immutable,
112112
)
113113

114-
def send_signal(self, name: str, sig: int, group: bool = True):
114+
def send_signal(self, name: str, sig: int, group: bool = False):
115115
"""Send `signal` to the specified named process."""
116116
try:
117117
process_info = self[name]
@@ -135,9 +135,10 @@ def handle_closed_process():
135135
if process_info.returncode is None:
136136
try:
137137
if sys.platform != "win32" and group:
138-
os.killpg( # pylint: disable=no-member
139-
process_info.pid, sig
138+
pgid = os.getpgid( # pylint: disable=no-member
139+
process_info.pid
140140
)
141+
os.killpg(pgid, sig) # pylint: disable=no-member
141142
else:
142143
os.kill(process_info.pid, sig)
143144
except ProcessLookupError:
@@ -154,13 +155,18 @@ def handle_closed_process():
154155

155156
def interrupt(self, name: str, group: bool = True):
156157
"""Send interrupt signal to specified named process"""
157-
self.send_signal(name, signal.SIGINT, group)
158+
if sys.platform == "win32":
159+
self.send_signal(
160+
name, signal.CTRL_C_EVENT, group # pylint: disable=no-member
161+
)
162+
else:
163+
self.send_signal(name, signal.SIGINT, group)
158164

159-
def terminate(self, name: str, group: bool = True):
165+
def terminate(self, name: str, group: bool = False):
160166
"""Terminate the specified named process."""
161167
self.send_signal(name, signal.SIGTERM, group)
162168

163-
def kill(self, name: str, group: bool = True):
169+
def kill(self, name: str, group: bool = False):
164170
"""Kill the specified named process."""
165171
if sys.platform == "win32":
166172
self.send_signal(name, signal.SIGTERM, group)

tests/proc/test_manager.py

Lines changed: 24 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -29,13 +29,25 @@ def test_send_signal(
2929
mock_kill.assert_called_once_with(PID_RUNNING, signal.SIGTERM)
3030

3131
if sys.platform != "win32":
32+
gid = 100
33+
mocker.patch("os.getpgid", return_value=gid)
3234
mock_killpg = mocker.patch("os.killpg")
3335
process_manager.send_signal(running_process, signal.SIGINT, True)
34-
mock_killpg.assert_called_once_with(PID_RUNNING, signal.SIGINT)
36+
mock_killpg.assert_called_once_with(gid, signal.SIGINT)
37+
else:
38+
mock_kill.reset_mock()
39+
process_manager.send_signal(
40+
running_process,
41+
signal.CTRL_C_EVENT, # pylint: disable=no-member
42+
True,
43+
)
44+
mock_kill.assert_called_once_with(
45+
PID_RUNNING, signal.CTRL_C_EVENT # pylint: disable=no-member
46+
)
3547

3648
mock_kill.reset_mock()
3749
with pytest.raises(ProcessLookupError):
38-
process_manager.send_signal(finished_process, signal.SIGTERM)
50+
process_manager.send_signal(finished_process, signal.SIGTERM, False)
3951
mock_kill.assert_not_called()
4052

4153
if sys.platform == "win32":
@@ -59,39 +71,42 @@ def side_effect(*args):
5971

6072
mocker.patch("os.kill", side_effect=side_effect)
6173
with pytest.raises(ProcessLookupError):
62-
process_manager.send_signal(running_process, signal.SIGTERM)
74+
process_manager.send_signal(running_process, signal.SIGTERM, False)
6375
assert process_manager[running_process].returncode == -1
6476

6577
with pytest.raises(ProcessLookupError):
66-
process_manager.send_signal("nonexists", signal.SIGTERM)
78+
process_manager.send_signal("nonexists", signal.SIGTERM, False)
6779

6880

6981
if sys.platform == "win32":
7082
SIGKILL = signal.SIGTERM
83+
SIGINT = signal.CTRL_C_EVENT # pylint: disable=no-member
7184
else:
7285
SIGKILL = signal.SIGKILL # pylint: disable=no-member
86+
SIGINT = signal.SIGINT
7387

7488

7589
@pytest.mark.parametrize(
76-
"method, sig",
90+
"method, sig, group",
7791
[
78-
("kill", SIGKILL),
79-
("terminate", signal.SIGTERM),
80-
("interrupt", signal.SIGINT),
92+
("kill", SIGKILL, False),
93+
("terminate", signal.SIGTERM, False),
94+
("interrupt", SIGINT, True),
8195
],
8296
)
8397
def test_kill_commands(
8498
mocker: MockerFixture,
8599
process_manager: ProcessManager,
86100
method: str,
87101
sig: signal.Signals,
102+
group: bool,
88103
):
89104
"""Test shortcut for different signals."""
90105
name = "process"
91106
mock_kill = mocker.patch.object(process_manager, "send_signal")
92107
func = getattr(process_manager, method)
93108
func(name)
94-
mock_kill.assert_called_once_with(name, sig, True)
109+
mock_kill.assert_called_once_with(name, sig, group)
95110

96111

97112
def test_remove(

0 commit comments

Comments
 (0)