Skip to content

Commit

Permalink
Fix (non-container) pytest stdout buffering-related lock-up
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
danieldk committed May 28, 2024
1 parent 9231098 commit 54e3e2f
Showing 1 changed file with 18 additions and 14 deletions.
32 changes: 18 additions & 14 deletions integration-tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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"]
Expand Down

0 comments on commit 54e3e2f

Please sign in to comment.