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.
This commit is contained in:
Daniël de Kok 2024-05-28 07:25:14 +00:00 committed by Daniël de Kok
parent e76b9824ae
commit f20463e4e3
1 changed files with 16 additions and 12 deletions

View File

@ -7,9 +7,10 @@ import os
import docker import docker
import json import json
import math import math
import shutil
import tempfile
import time import time
import random import random
import re
from docker.errors import NotFound from docker.errors import NotFound
from typing import Optional, List, Dict from typing import Optional, List, Dict
@ -347,19 +348,22 @@ def launcher(event_loop):
if not use_flash_attention: if not use_flash_attention:
env["USE_FLASH_ATTENTION"] = "false" env["USE_FLASH_ATTENTION"] = "false"
with subprocess.Popen( with tempfile.TemporaryFile("w+") as tmp:
args, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=env # We'll output stdout/stderr to a temporary file. Using a pipe
) as process: # cause the process to block until stdout is read.
yield ProcessLauncherHandle(process, port) with subprocess.Popen(
args,
stdout=tmp,
stderr=subprocess.STDOUT,
env=env,
) as process:
yield ProcessLauncherHandle(process, port)
process.terminate() process.terminate()
process.wait(60) process.wait(60)
launcher_output = process.stdout.read().decode("utf-8") tmp.seek(0)
print(launcher_output, file=sys.stderr) shutil.copyfileobj(tmp, sys.stderr)
process.stdout.close()
process.stderr.close()
if not use_flash_attention: if not use_flash_attention:
del env["USE_FLASH_ATTENTION"] del env["USE_FLASH_ATTENTION"]