Commit f20463e4 authored by Daniël de Kok's avatar Daniël de Kok Committed by Daniël de Kok
Browse files

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.
parent e76b9824
......@@ -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 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=subprocess.PIPE, stderr=subprocess.PIPE, env=env
args,
stdout=tmp,
stderr=subprocess.STDOUT,
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()
tmp.seek(0)
shutil.copyfileobj(tmp, sys.stderr)
if not use_flash_attention:
del env["USE_FLASH_ATTENTION"]
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment