From 54e3e2faa83de987f050a1c04e53337241ab6915 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 | 32 ++++++++++++++++++-------------- 1 file changed, 18 insertions(+), 14 deletions(-) diff --git a/integration-tests/conftest.py b/integration-tests/conftest.py index ae3f977b416..d81b8736345 100644 --- a/integration-tests/conftest.py +++ b/integration-tests/conftest.py @@ -7,9 +7,10 @@ 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 local_launcher( 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) - - 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() + 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) + + tmp.seek(0) + shutil.copyfileobj(tmp, sys.stderr) if not use_flash_attention: del env["USE_FLASH_ATTENTION"]