Skip to content

Commit f6a2c71

Browse files
committed
Re-apply lost changes
1 parent 700bfb0 commit f6a2c71

File tree

3 files changed

+75
-35
lines changed

3 files changed

+75
-35
lines changed

simvue/executor.py

Lines changed: 55 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
import os
1818
import psutil
1919
import subprocess
20+
import contextlib
2021
import pathlib
2122
import time
2223
import typing
@@ -229,6 +230,24 @@ def callback_function(status_code: int, std_out: str, std_err: str) -> None:
229230
name=f"{identifier}_exit_status", source="user"
230231
)
231232

233+
@property
234+
def processes(self) -> list[psutil.Process]:
235+
"""Create an array containing a list of processes"""
236+
if not self._processes:
237+
return []
238+
239+
_all_processes: list[psutil.Process] = [
240+
psutil.Process(process.pid) for process in self._processes.values()
241+
]
242+
243+
with contextlib.suppress(psutil.NoSuchProcess, psutil.ZombieProcess):
244+
for process in _all_processes:
245+
for child in process.children(recursive=True):
246+
if child not in _all_processes:
247+
_all_processes.append(child)
248+
249+
return list(set(_all_processes))
250+
232251
@property
233252
def success(self) -> int:
234253
"""Return whether all attached processes completed successfully"""
@@ -320,15 +339,37 @@ def _save_output(self) -> None:
320339
f"{self._runner.name}_{proc_id}.out", category="output"
321340
)
322341

323-
def kill_process(self, process_id: str) -> None:
324-
"""Kill a running process by ID"""
325-
if not (process := self._processes.get(process_id)):
326-
logger.error(
327-
f"Failed to terminate process '{process_id}', no such identifier."
328-
)
329-
return
342+
def kill_process(
343+
self, process_id: typing.Union[int, str], kill_children_only: bool = False
344+
) -> None:
345+
"""Kill a running process by ID
330346
331-
parent = psutil.Process(process.pid)
347+
If argument is a string this is a process handled by the client,
348+
else it is a PID of a external monitored process
349+
350+
Parameters
351+
----------
352+
process_id : typing.Union[int, str]
353+
either the identifier for a client created process or the PID
354+
of an external process
355+
kill_children_only : bool, optional
356+
if process_id is an integer, whether to kill only its children
357+
"""
358+
if isinstance(process_id, str):
359+
if not (process := self._processes.get(process_id)):
360+
logger.error(
361+
f"Failed to terminate process '{process_id}', no such identifier."
362+
)
363+
return
364+
try:
365+
parent = psutil.Process(process.pid)
366+
except psutil.NoSuchProcess:
367+
return
368+
elif isinstance(process_id, int):
369+
try:
370+
parent = psutil.Process(process_id)
371+
except psutil.NoSuchProcess:
372+
return
332373

333374
for child in parent.children(recursive=True):
334375
logger.debug(f"Terminating child process {child.pid}: {child.name()}")
@@ -337,17 +378,13 @@ def kill_process(self, process_id: str) -> None:
337378
for child in parent.children(recursive=True):
338379
child.wait()
339380

340-
logger.debug(f"Terminating child process {process.pid}: {process.args}")
341-
process.kill()
342-
process.wait()
343-
344-
if trigger := self._completion_triggers.get(process_id):
345-
trigger.set()
346-
347-
if trigger_process := self._completion_processes.get(process_id):
348-
trigger_process.join()
381+
if not kill_children_only:
382+
logger.debug(f"Terminating process {process.pid}: {process.args}")
383+
process.kill()
384+
process.wait()
349385

350-
self._execute_callback(process_id)
386+
if isinstance(process_id, str):
387+
self._execute_callback(process_id)
351388

352389
def kill_all(self) -> None:
353390
"""Kill all running processes"""

simvue/run.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -298,6 +298,7 @@ def _heartbeat(
298298
and self._abort_on_alert
299299
and self._simvue.get_abort_status()
300300
):
301+
logger.debug("Received abort request from server")
301302
self._alert_raised_trigger.set()
302303
self.kill_all_processes()
303304
if self._dispatcher and self._shutdown_event:

tests/refactor/test_run_class.py

Lines changed: 19 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -487,23 +487,6 @@ def testing_exit(status: int) -> None:
487487
run.kill_all_processes()
488488
raise AssertionError("Run was not terminated")
489489

490-
@pytest.mark.run
491-
def test_abort_on_alert_raise(create_plain_run: typing.Tuple[sv_run.Run, dict], mocker: pytest_mock.MockerFixture) -> None:
492-
def testing_exit(status: int) -> None:
493-
raise SystemExit(status)
494-
mocker.patch("os._exit", testing_exit)
495-
run, _ = create_plain_run
496-
run.config(resources_metrics_interval=1)
497-
run._heartbeat_interval = 1
498-
run._testing = True
499-
alert_id = run.create_alert("abort_test", source="user", trigger_abort=True)
500-
run.add_process(identifier="forever_long", executable="bash", c="sleep 10")
501-
time.sleep(2)
502-
run.log_alert(alert_id, "critical")
503-
time.sleep(4)
504-
if not run._status == "terminated":
505-
run.kill_all_processes()
506-
raise AssertionError("Run was not terminated")
507490

508491
@pytest.mark.run
509492
def test_abort_on_alert_python(create_plain_run: typing.Tuple[sv_run.Run, dict], mocker: pytest_mock.MockerFixture) -> None:
@@ -530,6 +513,25 @@ def testing_exit(status: int) -> None:
530513
assert run._status == "terminated"
531514

532515

516+
@pytest.mark.run
517+
def test_abort_on_alert_raise(create_plain_run: typing.Tuple[sv_run.Run, dict], mocker: pytest_mock.MockerFixture) -> None:
518+
def testing_exit(status: int) -> None:
519+
raise SystemExit(status)
520+
mocker.patch("os._exit", testing_exit)
521+
run, _ = create_plain_run
522+
run.config(resources_metrics_interval=1)
523+
run._heartbeat_interval = 1
524+
run._testing = True
525+
alert_id = run.create_alert("abort_test", source="user", trigger_abort=True)
526+
run.add_process(identifier="forever_long", executable="bash", c="sleep 10")
527+
time.sleep(2)
528+
run.log_alert(alert_id, "critical")
529+
time.sleep(4)
530+
if not run._status == "terminated":
531+
run.kill_all_processes()
532+
raise AssertionError("Run was not terminated")
533+
534+
533535
@pytest.mark.run
534536
def test_kill_all_processes(create_plain_run: typing.Tuple[sv_run.Run, dict]) -> None:
535537
run, _ = create_plain_run

0 commit comments

Comments
 (0)