Spaces:
Running
Running
| 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) | |