Spaces:
Running
Running
File size: 2,848 Bytes
41e79e2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 |
import os
import signal
import fcntl
import time
import subprocess
from typing import List
MAX_BYTES_PER_READ = 1024
SLEEP_BETWEEN_READS = 0.1
class Result:
timeout: int
exit_code: int
stdout: str
stderr: str
def __init__(self, timeout, exit_code, stdout, stderr):
self.timeout = timeout
self.exit_code = exit_code
self.stdout = stdout
self.stderr = stderr
def set_nonblocking(reader):
fd = reader.fileno()
fl = fcntl.fcntl(fd, fcntl.F_GETFL)
fcntl.fcntl(fd, fcntl.F_SETFL, fl | os.O_NONBLOCK)
def run(
args: List[str],
timeout_seconds: int = 15,
max_output_size: int = 2048,
env = None,
cwd: str | None = None
) -> Result:
"""
Runs the given program with arguments. After the timeout elapses, kills the process
and all other processes in the process group. Captures at most max_output_size bytes
of stdout and stderr each, and discards any output beyond that.
"""
p = subprocess.Popen(
args,
env=env,
stdin=subprocess.DEVNULL,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
start_new_session=True,
bufsize=MAX_BYTES_PER_READ,
cwd=cwd
)
set_nonblocking(p.stdout)
set_nonblocking(p.stderr)
process_group_id = os.getpgid(p.pid)
# We sleep for 0.1 seconds in each iteration.
max_iterations = timeout_seconds * 10
stdout_saved_bytes = []
stderr_saved_bytes = []
stdout_bytes_read = 0
stderr_bytes_read = 0
for _ in range(max_iterations):
this_stdout_read = p.stdout.read(MAX_BYTES_PER_READ)
this_stderr_read = p.stderr.read(MAX_BYTES_PER_READ)
# this_stdout_read and this_stderr_read may be None if stdout or stderr
# are closed. Without these checks, test_close_output fails.
if this_stdout_read is not None and stdout_bytes_read < max_output_size:
stdout_saved_bytes.append(this_stdout_read)
stdout_bytes_read += len(this_stdout_read)
if this_stderr_read is not None and stderr_bytes_read < max_output_size:
stderr_saved_bytes.append(this_stderr_read)
stderr_bytes_read += len(this_stderr_read)
exit_code = p.poll()
if exit_code is not None:
break
time.sleep(SLEEP_BETWEEN_READS)
try:
# Kills the process group. Without this line, test_fork_once fails.
os.killpg(process_group_id, signal.SIGKILL)
except ProcessLookupError:
pass
timeout = exit_code is None
exit_code = exit_code if exit_code is not None else -1
stdout = b"".join(stdout_saved_bytes).decode("utf-8", errors="ignore")
stderr = b"".join(stderr_saved_bytes).decode("utf-8", errors="ignore")
return Result(timeout=timeout, exit_code=exit_code, stdout=stdout, stderr=stderr)
|