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 ...@@ -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): ...@@ -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,
process.terminate() stdout=tmp,
process.wait(60) stderr=subprocess.STDOUT,
env=env,
launcher_output = process.stdout.read().decode("utf-8") ) as process:
print(launcher_output, file=sys.stderr) yield ProcessLauncherHandle(process, port)
process.stdout.close() process.terminate()
process.stderr.close() process.wait(60)
tmp.seek(0)
shutil.copyfileobj(tmp, sys.stderr)
if not use_flash_attention: if not use_flash_attention:
del env["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