From 9a1475d816071dace5c75ae39d7f2af5090c2023 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Danie=CC=88l=20de=20Kok?= Date: Tue, 28 May 2024 07:25:14 +0000 Subject: [PATCH] Fix (non-container) pytest stdout buffering-related lock-up Two issues: 1. When one of the stdout/stderr pipe buffers of a process started with `subprocess.Popen` is full, the process can get blocked until the buffer is drained. 2. Calling `Popen.wait` can deadlock when called before draining the pipe buffers (if they are full). This avoids the issue altogether by giving the child process a temporary file to write to. --- integration-tests/conftest.py | 28 ++++++++++++++++------------ 1 file changed, 16 insertions(+), 12 deletions(-) diff --git a/integration-tests/conftest.py b/integration-tests/conftest.py index ae3f977b4..d81b87363 100644 --- a/integration-tests/conftest.py +++ b/integration-tests/conftest.py @@ -7,9 +7,10 @@ import os import docker import json import math +import shutil +import tempfile import time import random -import re from docker.errors import NotFound from typing import Optional, List, Dict @@ -347,19 +348,22 @@ def launcher(event_loop): if not use_flash_attention: env["USE_FLASH_ATTENTION"] = "false" - with subprocess.Popen( - args, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=env - ) as process: - yield ProcessLauncherHandle(process, port) + with tempfile.TemporaryFile("w+") as tmp: + # We'll output stdout/stderr to a temporary file. Using a pipe + # cause the process to block until stdout is read. + with subprocess.Popen( + args, + stdout=tmp, + stderr=subprocess.STDOUT, + env=env, + ) as process: + yield ProcessLauncherHandle(process, port) - process.terminate() - process.wait(60) + process.terminate() + process.wait(60) - launcher_output = process.stdout.read().decode("utf-8") - print(launcher_output, file=sys.stderr) - - process.stdout.close() - process.stderr.close() + tmp.seek(0) + shutil.copyfileobj(tmp, sys.stderr) if not use_flash_attention: del env["USE_FLASH_ATTENTION"]