fix: kill runaway code on timeout in container and local executors
Co-authored-by: George Weale <gweale@google.com> PiperOrigin-RevId: 956773467
This commit is contained in:
George Weale
committed by
Copybara-Service
Copybara-Service
parent
b5b27cb074
commit
27548e392f
@@ -33,6 +33,74 @@ from .code_execution_utils import CodeExecutionResult
|
||||
logger = logging.getLogger('google_adk.' + __name__)
|
||||
DEFAULT_IMAGE_TAG = 'adk-code-executor:latest'
|
||||
|
||||
# Reported by the supervisor below when it kills a run that hit the bound.
|
||||
# Follows the convention of coreutils `timeout`; unlike 128 + SIGALRM it is not
|
||||
# something the executed code produces by letting an alarm of its own fire.
|
||||
_TIMEOUT_EXIT_CODE = 124
|
||||
|
||||
# Runs the code under a supervisor that enforces a hard wall-clock bound inside
|
||||
# the container. The code runs in a forked child in its own process group; the
|
||||
# supervisor waits for it and, when the bound expires, SIGKILLs the whole group
|
||||
# and exits with `_TIMEOUT_EXIT_CODE`. The group is also swept once the code
|
||||
# finishes normally, so nothing it started is left running in the shared
|
||||
# container (a leftover process would also hold the exec's output open).
|
||||
#
|
||||
# Two properties matter for code that may be hostile: the deadline lives in a
|
||||
# process the code never runs in, so it cannot be disarmed from inside (an
|
||||
# alarm armed in the executing process could be cancelled with one call), and
|
||||
# SIGKILL to the group reaches what the code spawned, not just its top frame.
|
||||
# The code is passed as an argument rather than inlined so that tracebacks keep
|
||||
# the original line numbers, and argv is restored to what `python3 -c` would
|
||||
# have given so that code parsing arguments still works.
|
||||
#
|
||||
# Not covered: code that deliberately leaves the group (`os.setsid()`) or
|
||||
# double-forks away survives the bound until the container is torn down.
|
||||
_TIMEOUT_WRAPPER = """\
|
||||
import os, signal, sys
|
||||
|
||||
_timeout = int(sys.argv[1])
|
||||
_source = sys.argv[2]
|
||||
del sys.argv[1:]
|
||||
|
||||
_pid = os.fork()
|
||||
if _pid == 0:
|
||||
try:
|
||||
os.setpgid(0, 0)
|
||||
except OSError:
|
||||
pass
|
||||
exec(compile(_source, '<adk_code>', 'exec'), {'__name__': '__main__'})
|
||||
else:
|
||||
|
||||
def _sweep_group():
|
||||
try:
|
||||
os.killpg(_pid, signal.SIGKILL)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
def _expire(_signum, _frame):
|
||||
_sweep_group()
|
||||
try:
|
||||
os.kill(_pid, signal.SIGKILL)
|
||||
except OSError:
|
||||
pass
|
||||
os._exit(%d)
|
||||
|
||||
try:
|
||||
os.setpgid(_pid, _pid)
|
||||
except OSError:
|
||||
pass
|
||||
signal.signal(signal.SIGALRM, _expire)
|
||||
signal.alarm(_timeout)
|
||||
_status = os.waitpid(_pid, 0)[1]
|
||||
signal.alarm(0)
|
||||
_sweep_group()
|
||||
os._exit(
|
||||
128 + os.WTERMSIG(_status)
|
||||
if os.WIFSIGNALED(_status)
|
||||
else os.WEXITSTATUS(_status)
|
||||
)
|
||||
""" % _TIMEOUT_EXIT_CODE
|
||||
|
||||
|
||||
class ContainerCodeExecutor(BaseCodeExecutor):
|
||||
"""A code executor that uses a custom container to execute code.
|
||||
@@ -87,6 +155,23 @@ class ContainerCodeExecutor(BaseCodeExecutor):
|
||||
code must make network requests and you trust it.
|
||||
"""
|
||||
|
||||
# Overrides the BaseCodeExecutor attribute: unlike the base default of None,
|
||||
# the timeout here is always finite and must be positive (0 would mean no
|
||||
# bound at all).
|
||||
timeout_seconds: int = Field(default=300, gt=0)
|
||||
"""The wall-clock timeout in seconds for a single code execution.
|
||||
|
||||
Every execution shares one long-lived container, so an unbounded run (e.g. a
|
||||
loop emitted by the model) would keep burning that container's CPU for every
|
||||
later caller. Defaults to 300, matching ``GkeCodeExecutor``. A computation
|
||||
that legitimately runs longer than the timeout is killed, so raise it rather
|
||||
than removing it; ``None`` is rejected, unlike on the base class.
|
||||
|
||||
When the timeout expires the executed code is killed along with the process
|
||||
group it runs in, so what it spawned goes with it. Code that deliberately
|
||||
detaches from that group keeps running until the container is torn down.
|
||||
"""
|
||||
|
||||
# Overrides the BaseCodeExecutor attribute: this executor cannot be stateful.
|
||||
stateful: bool = Field(default=False, frozen=True, exclude=True)
|
||||
|
||||
@@ -151,7 +236,13 @@ class ContainerCodeExecutor(BaseCodeExecutor):
|
||||
output = ''
|
||||
error = ''
|
||||
exec_result = self._container.exec_run(
|
||||
['python3', '-c', code_execution_input.code],
|
||||
[
|
||||
'python3',
|
||||
'-c',
|
||||
_TIMEOUT_WRAPPER,
|
||||
str(self.timeout_seconds),
|
||||
code_execution_input.code,
|
||||
],
|
||||
demux=True,
|
||||
)
|
||||
logger.debug('Executed code:\n```\n%s\n```', code_execution_input.code)
|
||||
@@ -165,6 +256,15 @@ class ContainerCodeExecutor(BaseCodeExecutor):
|
||||
):
|
||||
error = exec_result.output[1].decode('utf-8')
|
||||
|
||||
if exec_result.exit_code == _TIMEOUT_EXIT_CODE:
|
||||
# Appended rather than assigned: whatever the code managed to write to
|
||||
# stderr before the alarm fired is still the useful diagnostic, but on
|
||||
# its own it would hide the fact that the run was cut short.
|
||||
timed_out = (
|
||||
f'Code execution timed out after {self.timeout_seconds} seconds.'
|
||||
)
|
||||
error = f'{error}\n{timed_out}' if error else timed_out
|
||||
|
||||
# Collect the final result.
|
||||
return CodeExecutionResult(
|
||||
stdout=output,
|
||||
|
||||
@@ -18,8 +18,10 @@ from contextlib import redirect_stdout
|
||||
import io
|
||||
import logging
|
||||
import multiprocessing
|
||||
import os
|
||||
import queue
|
||||
import re
|
||||
import signal
|
||||
import traceback
|
||||
from typing import Any
|
||||
|
||||
@@ -33,11 +35,23 @@ from .code_execution_utils import CodeExecutionResult
|
||||
|
||||
logger = logging.getLogger('google_adk.' + __name__)
|
||||
|
||||
# How long to wait for a timed-out execution to exit after SIGTERM before
|
||||
# escalating to SIGKILL, so that the timeout itself cannot block forever.
|
||||
_TERMINATE_GRACE_SECONDS = 5
|
||||
|
||||
|
||||
def _execute_in_process(
|
||||
code: str, globals_: dict[str, Any], result_queue: multiprocessing.Queue
|
||||
) -> None:
|
||||
"""Executes code in a separate process and puts result in queue."""
|
||||
# Detach into a new session/process group before running anything, so that a
|
||||
# timed-out execution can be killed together with everything it spawned.
|
||||
if hasattr(os, 'setsid'):
|
||||
try:
|
||||
os.setsid()
|
||||
except OSError:
|
||||
logger.debug('Could not detach the execution process group.')
|
||||
|
||||
stdout = io.StringIO()
|
||||
error = None
|
||||
try:
|
||||
@@ -48,6 +62,53 @@ def _execute_in_process(
|
||||
result_queue.put((stdout.getvalue(), error))
|
||||
|
||||
|
||||
def _execution_group(
|
||||
process: multiprocessing.process.BaseProcess,
|
||||
) -> int | None:
|
||||
"""Returns the group the execution detached into, or None if it has not."""
|
||||
if process.pid is None or not hasattr(os, 'killpg'):
|
||||
return None
|
||||
try:
|
||||
group = os.getpgid(process.pid)
|
||||
# Only report the group once the execution has detached into its own;
|
||||
# otherwise the group is still ours and signalling it would take down the
|
||||
# agent along with the code it is running.
|
||||
return group if group != os.getpgid(0) else None
|
||||
except OSError:
|
||||
return None
|
||||
|
||||
|
||||
def _signal_group(group: int, sig: int) -> None:
|
||||
"""Signals every process left in a group, tolerating an empty one."""
|
||||
try:
|
||||
os.killpg(group, sig)
|
||||
except OSError:
|
||||
logger.debug('Could not signal the execution process group.')
|
||||
|
||||
|
||||
def _kill_execution(process: multiprocessing.process.BaseProcess) -> None:
|
||||
"""Kills a timed-out execution along with any process it spawned."""
|
||||
# Resolved up front: once the execution process has been reaped its group can
|
||||
# no longer be looked up through it, and the group is what holds whatever the
|
||||
# code spawned.
|
||||
group = _execution_group(process)
|
||||
|
||||
# SIGTERM first, so the code and its children get the same grace period the
|
||||
# execution process itself gets before anything is killed outright.
|
||||
if group is not None:
|
||||
_signal_group(group, signal.SIGTERM)
|
||||
process.terminate()
|
||||
process.join(_TERMINATE_GRACE_SECONDS)
|
||||
|
||||
# Escalate unconditionally: the execution process exiting says nothing about
|
||||
# a child of it that is ignoring SIGTERM.
|
||||
if group is not None:
|
||||
_signal_group(group, signal.SIGKILL)
|
||||
if process.is_alive():
|
||||
process.kill()
|
||||
process.join()
|
||||
|
||||
|
||||
def _prepare_globals(code: str, globals_: dict[str, Any]) -> None:
|
||||
"""Prepare globals for code execution, injecting __name__ if needed."""
|
||||
if re.search(r"if\s+__name__\s*==\s*['\"]__main__['\"]", code):
|
||||
@@ -102,8 +163,7 @@ class UnsafeLocalCodeExecutor(BaseCodeExecutor):
|
||||
if err:
|
||||
error = err
|
||||
except queue.Empty:
|
||||
process.terminate()
|
||||
process.join()
|
||||
_kill_execution(process)
|
||||
error = f'Code execution timed out after {self.timeout_seconds} seconds.'
|
||||
|
||||
# Collect the final result.
|
||||
|
||||
@@ -14,9 +14,19 @@
|
||||
|
||||
"""Tests for the ContainerCodeExecutor container hardening defaults."""
|
||||
|
||||
import os
|
||||
import signal
|
||||
import subprocess
|
||||
import sys
|
||||
import textwrap
|
||||
import time
|
||||
from unittest import mock
|
||||
|
||||
from google.adk.code_executors import container_code_executor
|
||||
from google.adk.code_executors.code_execution_utils import CodeExecutionInput
|
||||
from google.adk.code_executors.container_code_executor import ContainerCodeExecutor
|
||||
import pydantic
|
||||
import pytest
|
||||
|
||||
|
||||
def _mock_docker_client():
|
||||
@@ -56,3 +66,239 @@ def test_container_network_can_be_explicitly_enabled(mock_docker):
|
||||
|
||||
_, kwargs = client.containers.run.call_args
|
||||
assert not kwargs['network_disabled']
|
||||
|
||||
|
||||
def _executed_command(container) -> list[str]:
|
||||
"""Returns the command of the last `exec_run` call on the container."""
|
||||
args, _ = container.exec_run.call_args
|
||||
return args[0]
|
||||
|
||||
|
||||
@mock.patch('google.adk.code_executors.container_code_executor.docker')
|
||||
def test_execute_code_bounds_execution_by_default(mock_docker):
|
||||
"""Code runs under a finite timeout even when the caller sets none."""
|
||||
client = _mock_docker_client()
|
||||
mock_docker.from_env.return_value = client
|
||||
executor = ContainerCodeExecutor(image='test-image')
|
||||
container = client.containers.run.return_value
|
||||
container.exec_run.return_value = mock.MagicMock(
|
||||
exit_code=0, output=(b'', b'')
|
||||
)
|
||||
|
||||
executor.execute_code(mock.MagicMock(), CodeExecutionInput(code='x = 1'))
|
||||
|
||||
# The container is shared by every invocation, so an unbounded run would pin
|
||||
# it for all later callers.
|
||||
assert executor.timeout_seconds == 300
|
||||
assert _executed_command(container) == [
|
||||
'python3',
|
||||
'-c',
|
||||
container_code_executor._TIMEOUT_WRAPPER,
|
||||
'300',
|
||||
'x = 1',
|
||||
]
|
||||
|
||||
|
||||
@mock.patch('google.adk.code_executors.container_code_executor.docker')
|
||||
def test_execute_code_passes_configured_timeout(mock_docker):
|
||||
"""The inherited `timeout_seconds` bounds the in-container execution."""
|
||||
client = _mock_docker_client()
|
||||
mock_docker.from_env.return_value = client
|
||||
executor = ContainerCodeExecutor(image='test-image', timeout_seconds=7)
|
||||
container = client.containers.run.return_value
|
||||
container.exec_run.return_value = mock.MagicMock(
|
||||
exit_code=0, output=(b'', b'')
|
||||
)
|
||||
|
||||
executor.execute_code(
|
||||
mock.MagicMock(), CodeExecutionInput(code='while True: pass')
|
||||
)
|
||||
|
||||
assert _executed_command(container) == [
|
||||
'python3',
|
||||
'-c',
|
||||
container_code_executor._TIMEOUT_WRAPPER,
|
||||
'7',
|
||||
'while True: pass',
|
||||
]
|
||||
|
||||
|
||||
@mock.patch('google.adk.code_executors.container_code_executor.docker')
|
||||
@pytest.mark.parametrize('timeout', [0, -1, None])
|
||||
def test_non_positive_timeout_is_rejected(mock_docker, timeout):
|
||||
"""A timeout of 0 or None would mean no bound, so it is refused up front."""
|
||||
mock_docker.from_env.return_value = _mock_docker_client()
|
||||
|
||||
with pytest.raises(pydantic.ValidationError):
|
||||
ContainerCodeExecutor(image='test-image', timeout_seconds=timeout)
|
||||
|
||||
|
||||
@mock.patch('google.adk.code_executors.container_code_executor.docker')
|
||||
def test_execute_code_reports_timeout(mock_docker):
|
||||
"""A run the supervisor cut short is reported as a timeout."""
|
||||
client = _mock_docker_client()
|
||||
mock_docker.from_env.return_value = client
|
||||
executor = ContainerCodeExecutor(image='test-image', timeout_seconds=7)
|
||||
container = client.containers.run.return_value
|
||||
container.exec_run.return_value = mock.MagicMock(
|
||||
exit_code=container_code_executor._TIMEOUT_EXIT_CODE, output=(b'', b'')
|
||||
)
|
||||
|
||||
result = executor.execute_code(
|
||||
mock.MagicMock(), CodeExecutionInput(code='while True: pass')
|
||||
)
|
||||
|
||||
assert 'timed out after 7 seconds' in result.stderr
|
||||
|
||||
|
||||
@mock.patch('google.adk.code_executors.container_code_executor.docker')
|
||||
def test_execute_code_reports_timeout_alongside_stderr(mock_docker):
|
||||
"""Output written before the alarm fired does not hide the timeout."""
|
||||
client = _mock_docker_client()
|
||||
mock_docker.from_env.return_value = client
|
||||
executor = ContainerCodeExecutor(image='test-image', timeout_seconds=7)
|
||||
container = client.containers.run.return_value
|
||||
container.exec_run.return_value = mock.MagicMock(
|
||||
exit_code=container_code_executor._TIMEOUT_EXIT_CODE,
|
||||
output=(b'', b'a warning from the code'),
|
||||
)
|
||||
|
||||
result = executor.execute_code(
|
||||
mock.MagicMock(), CodeExecutionInput(code='while True: pass')
|
||||
)
|
||||
|
||||
assert 'a warning from the code' in result.stderr
|
||||
assert 'timed out after 7 seconds' in result.stderr
|
||||
|
||||
|
||||
# The wrapper below is a string of Python that only ever runs inside the
|
||||
# container, so the tests run it directly on this host instead: no docker, no
|
||||
# daemon, and only snippets written here.
|
||||
_POSIX_ONLY = pytest.mark.skipif(
|
||||
not hasattr(os, 'fork') or not hasattr(os, 'killpg'),
|
||||
reason='The in-container bound is enforced with POSIX process groups.',
|
||||
)
|
||||
|
||||
|
||||
def _run_wrapper(timeout: int, code: str) -> subprocess.CompletedProcess:
|
||||
"""Runs the wrapper exactly as the container executor asks the container to."""
|
||||
return subprocess.run(
|
||||
[
|
||||
sys.executable,
|
||||
'-c',
|
||||
container_code_executor._TIMEOUT_WRAPPER,
|
||||
str(timeout),
|
||||
code,
|
||||
],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=30,
|
||||
check=False,
|
||||
)
|
||||
|
||||
|
||||
def _is_alive(pid: int) -> bool:
|
||||
"""Returns whether `pid` is a live (non-zombie) process."""
|
||||
try:
|
||||
with open(f'/proc/{pid}/stat', encoding='utf-8') as stat_file:
|
||||
state = stat_file.read().rsplit(')', 1)[1].split()[0]
|
||||
except OSError:
|
||||
return False
|
||||
return state != 'Z'
|
||||
|
||||
|
||||
@_POSIX_ONLY
|
||||
def test_wrapper_kills_a_run_that_hits_the_bound():
|
||||
"""A loop that never returns is killed, not merely waited on."""
|
||||
started = time.monotonic()
|
||||
|
||||
completed = _run_wrapper(1, 'while True: pass')
|
||||
|
||||
assert completed.returncode == container_code_executor._TIMEOUT_EXIT_CODE
|
||||
assert time.monotonic() - started < 15
|
||||
|
||||
|
||||
@_POSIX_ONLY
|
||||
def test_wrapper_bound_cannot_be_disarmed_by_the_executed_code():
|
||||
"""The deadline is held by a process the executed code never runs in."""
|
||||
completed = _run_wrapper(
|
||||
1,
|
||||
'import signal, time\n'
|
||||
'signal.alarm(0)\n'
|
||||
'signal.signal(signal.SIGALRM, signal.SIG_IGN)\n'
|
||||
'time.sleep(25)\n'
|
||||
'print("outlived the bound")\n',
|
||||
)
|
||||
|
||||
assert completed.returncode == container_code_executor._TIMEOUT_EXIT_CODE
|
||||
assert 'outlived the bound' not in completed.stdout
|
||||
|
||||
|
||||
@_POSIX_ONLY
|
||||
@pytest.mark.skipif(
|
||||
not os.path.isdir('/proc'), reason='Liveness is checked through /proc.'
|
||||
)
|
||||
def test_wrapper_kills_what_the_code_spawned(tmp_path):
|
||||
"""Whatever the run started dies with it, so the container is not pinned."""
|
||||
pid_file = tmp_path / 'spawned.pid'
|
||||
code = textwrap.dedent(f"""
|
||||
import os
|
||||
import time
|
||||
|
||||
spawned = os.fork()
|
||||
if spawned == 0:
|
||||
time.sleep(60)
|
||||
os._exit(0)
|
||||
with open({str(pid_file)!r}, 'w') as f:
|
||||
f.write(str(spawned))
|
||||
time.sleep(60)
|
||||
""")
|
||||
|
||||
completed = _run_wrapper(1, code)
|
||||
|
||||
assert completed.returncode == container_code_executor._TIMEOUT_EXIT_CODE
|
||||
assert pid_file.exists(), 'the code never got as far as spawning a process'
|
||||
spawned_pid = int(pid_file.read_text())
|
||||
try:
|
||||
deadline = time.monotonic() + 10
|
||||
while time.monotonic() < deadline and _is_alive(spawned_pid):
|
||||
time.sleep(0.05)
|
||||
assert not _is_alive(spawned_pid)
|
||||
finally:
|
||||
try:
|
||||
os.kill(spawned_pid, signal.SIGKILL)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
@_POSIX_ONLY
|
||||
def test_wrapper_leaves_argv_as_a_plain_python_c_run_would():
|
||||
"""The timeout and the source do not leak into the code's own arguments."""
|
||||
completed = _run_wrapper(
|
||||
5,
|
||||
'import argparse, sys\n'
|
||||
'argparse.ArgumentParser().parse_args()\n'
|
||||
'print(sys.argv)\n',
|
||||
)
|
||||
|
||||
assert completed.returncode == 0, completed.stderr
|
||||
assert completed.stdout.strip() == "['-c']"
|
||||
|
||||
|
||||
@_POSIX_ONLY
|
||||
def test_wrapper_passes_through_output_and_exit_code():
|
||||
"""A run that finishes on its own is reported exactly as it ended."""
|
||||
completed = _run_wrapper(5, 'import sys\nprint("hello")\nsys.exit(3)\n')
|
||||
|
||||
assert completed.returncode == 3
|
||||
assert completed.stdout.strip() == 'hello'
|
||||
|
||||
|
||||
@_POSIX_ONLY
|
||||
def test_wrapper_reports_an_uncaught_error_against_the_original_line():
|
||||
"""Wrapping the code does not shift the line numbers in its traceback."""
|
||||
completed = _run_wrapper(5, 'x = 1\nraise ValueError("boom")\n')
|
||||
|
||||
assert completed.returncode == 1
|
||||
assert '"<adk_code>", line 2' in completed.stderr
|
||||
assert 'boom' in completed.stderr
|
||||
|
||||
@@ -12,11 +12,16 @@
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import multiprocessing
|
||||
import os
|
||||
import signal
|
||||
import textwrap
|
||||
import time
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from google.adk.agents.base_agent import BaseAgent
|
||||
from google.adk.agents.invocation_context import InvocationContext
|
||||
from google.adk.code_executors import unsafe_local_code_executor
|
||||
from google.adk.code_executors.code_execution_utils import CodeExecutionInput
|
||||
from google.adk.code_executors.code_execution_utils import CodeExecutionResult
|
||||
from google.adk.code_executors.unsafe_local_code_executor import UnsafeLocalCodeExecutor
|
||||
@@ -25,6 +30,25 @@ from google.adk.sessions.session import Session
|
||||
import pytest
|
||||
|
||||
|
||||
def _written_pid(pid_file) -> int | None:
|
||||
"""Returns the pid the executed code recorded, or None if it has not yet."""
|
||||
try:
|
||||
recorded = pid_file.read_text().strip()
|
||||
except OSError:
|
||||
return None
|
||||
return int(recorded) if recorded else None
|
||||
|
||||
|
||||
def _is_alive(pid: int) -> bool:
|
||||
"""Returns whether `pid` is a live (non-zombie) process."""
|
||||
try:
|
||||
with open(f"/proc/{pid}/stat", encoding="utf-8") as stat_file:
|
||||
state = stat_file.read().rsplit(")", 1)[1].split()[0]
|
||||
except OSError:
|
||||
return False
|
||||
return state != "Z"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_invocation_context() -> InvocationContext:
|
||||
"""Provides a mock InvocationContext."""
|
||||
@@ -131,3 +155,93 @@ class TestUnsafeLocalCodeExecutor:
|
||||
|
||||
assert result.stdout == ""
|
||||
assert "Code execution timed out after 1 seconds." in result.stderr
|
||||
|
||||
def test_kill_execution_signals_group_before_killing_it(self, monkeypatch):
|
||||
"""The group gets SIGTERM and its grace period before SIGKILL."""
|
||||
signalled = []
|
||||
monkeypatch.setattr(
|
||||
unsafe_local_code_executor.os,
|
||||
"killpg",
|
||||
lambda group, sig: signalled.append((group, sig)),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
unsafe_local_code_executor,
|
||||
"_execution_group",
|
||||
lambda process: 4321,
|
||||
)
|
||||
process = MagicMock()
|
||||
process.is_alive.return_value = False
|
||||
process.terminate.side_effect = lambda: signalled.append(("child", "term"))
|
||||
|
||||
unsafe_local_code_executor._kill_execution(process)
|
||||
|
||||
assert signalled == [
|
||||
(4321, signal.SIGTERM),
|
||||
("child", "term"),
|
||||
(4321, signal.SIGKILL),
|
||||
]
|
||||
process.join.assert_any_call(
|
||||
unsafe_local_code_executor._TERMINATE_GRACE_SECONDS
|
||||
)
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not hasattr(os, "killpg")
|
||||
or not hasattr(os, "fork")
|
||||
or not os.path.isdir("/proc"),
|
||||
reason="Process-group teardown is checked on POSIX with /proc only.",
|
||||
)
|
||||
def test_kill_execution_kills_what_the_code_spawned(self, tmp_path):
|
||||
"""Killing a live execution takes the processes it spawned with it."""
|
||||
pid_file = tmp_path / "spawned.pid"
|
||||
# Forked rather than spawned through `sys.executable`, so the descendant
|
||||
# exists within milliseconds and the test never waits on interpreter
|
||||
# start-up.
|
||||
code = textwrap.dedent(f"""
|
||||
import os
|
||||
import time
|
||||
|
||||
spawned = os.fork()
|
||||
if spawned == 0:
|
||||
time.sleep(60)
|
||||
os._exit(0)
|
||||
with open({str(pid_file)!r}, 'w') as f:
|
||||
f.write(str(spawned))
|
||||
time.sleep(60)
|
||||
""")
|
||||
ctx = multiprocessing.get_context("spawn")
|
||||
result_queue = ctx.Queue()
|
||||
process = ctx.Process(
|
||||
target=unsafe_local_code_executor._execute_in_process,
|
||||
args=(code, {}, result_queue),
|
||||
daemon=True,
|
||||
)
|
||||
process.start()
|
||||
spawned_pid = None
|
||||
try:
|
||||
# Waiting for the pid to be written rather than for a fixed duration:
|
||||
# the only thing that has to have happened is the fork. The file exists
|
||||
# from the moment it is opened, so its content is what is polled for.
|
||||
deadline = time.time() + 30
|
||||
while time.time() < deadline and not _written_pid(pid_file):
|
||||
time.sleep(0.05)
|
||||
spawned_pid = _written_pid(pid_file)
|
||||
if spawned_pid is None:
|
||||
pytest.skip("this environment could not start the execution process")
|
||||
|
||||
unsafe_local_code_executor._kill_execution(process)
|
||||
|
||||
assert not process.is_alive()
|
||||
deadline = time.time() + 10
|
||||
while time.time() < deadline and _is_alive(spawned_pid):
|
||||
time.sleep(0.05)
|
||||
assert not _is_alive(spawned_pid)
|
||||
finally:
|
||||
if spawned_pid is not None:
|
||||
try:
|
||||
os.kill(spawned_pid, signal.SIGKILL)
|
||||
except OSError:
|
||||
pass
|
||||
if process.is_alive():
|
||||
process.kill()
|
||||
process.join()
|
||||
result_queue.close()
|
||||
|
||||
Reference in New Issue
Block a user