mirror of
https://github.com/multimodal-art-projection/YuE.git
synced 2026-10-17 05:50:07 +08:00
281 lines
10 KiB
Python
281 lines
10 KiB
Python
import asyncio
|
|
import io
|
|
import os
|
|
import random
|
|
import subprocess
|
|
import sys
|
|
import threading
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
import yue2.progress as progress_module
|
|
from yue2.progress import Progress
|
|
|
|
|
|
class Terminal(io.StringIO):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.refreshed = threading.Event()
|
|
self.writes = 0
|
|
|
|
def isatty(self):
|
|
return True
|
|
|
|
def write(self, value):
|
|
result = super().write(value)
|
|
if value.startswith("\r"):
|
|
self.writes += 1
|
|
if self.writes >= 2:
|
|
self.refreshed.set()
|
|
return result
|
|
|
|
|
|
def test_quiet_never_touches_stream_or_starts_thread(monkeypatch):
|
|
class ForbiddenStream:
|
|
def __getattr__(self, name):
|
|
raise AssertionError(f"Quiet progress touched stream.{name}")
|
|
|
|
def forbidden(*args, **kwargs):
|
|
raise AssertionError("Quiet progress started a thread")
|
|
|
|
monkeypatch.setattr(progress_module.threading, "Thread", forbidden)
|
|
with Progress(enabled=False, stream=ForbiddenStream()) as reporter:
|
|
with reporter.stage("Generating tokens", unit="tokens") as stage:
|
|
stage.token("semantic", 123)
|
|
stage.update(2)
|
|
reporter.complete(1.5, 0.2)
|
|
assert stage.completed == 2
|
|
|
|
|
|
def test_loading_heartbeat_refreshes_without_callbacks_and_stops():
|
|
stream = Terminal()
|
|
reporter = Progress(stream=stream, refresh_interval=0.01)
|
|
with reporter.stage("Loading model"):
|
|
assert stream.refreshed.wait(2.0), "No heartbeat during blocking model load"
|
|
thread = reporter._thread
|
|
assert thread.is_alive()
|
|
assert not thread.is_alive()
|
|
assert reporter._thread is None
|
|
assert stream.getvalue().endswith("\n")
|
|
assert "Completed Loading model" in stream.getvalue()
|
|
|
|
|
|
def test_non_tty_is_throttled_and_counts_only_callback_outputs(monkeypatch):
|
|
now = [100.0]
|
|
monkeypatch.setattr(progress_module, "_clock", lambda: now[0])
|
|
stream = io.StringIO()
|
|
with Progress(stream=stream) as reporter:
|
|
with reporter.stage("Generating tokens", unit="tokens") as stage:
|
|
for _ in range(100):
|
|
stage.token("semantic", object())
|
|
assert len(stream.getvalue().splitlines()) == 1
|
|
now[0] += 5.0
|
|
stage.update(100)
|
|
stage.update(100)
|
|
assert len(stream.getvalue().splitlines()) == 2
|
|
output = stream.getvalue()
|
|
assert "100 tokens | 20.0 tokens/s | elapsed 5.0s" in output
|
|
assert "%" not in output
|
|
assert "\r" not in output
|
|
assert len(output.splitlines()) == 3
|
|
|
|
|
|
def test_real_total_can_arrive_after_stage_starts(monkeypatch):
|
|
now = [0.0]
|
|
monkeypatch.setattr(progress_module, "_clock", lambda: now[0])
|
|
stream = io.StringIO()
|
|
with Progress(stream=stream).stage("Synthesizing", unit="steps") as stage:
|
|
now[0] = 5.0
|
|
stage.update(8, total=32)
|
|
stage.set_total(64)
|
|
now[0] = 10.0
|
|
stage.update(16)
|
|
stage.advance(16)
|
|
output = stream.getvalue()
|
|
assert "8/32 steps (25%)" in output
|
|
assert "16/64 steps (25%)" in output
|
|
assert "32/64 steps (50%)" in output
|
|
assert stage.completed == 32 # Exit must not invent completion of all 64 steps.
|
|
|
|
|
|
@pytest.mark.parametrize("exception,status", [
|
|
(RuntimeError("model failed"), "Failed"),
|
|
(InterruptedError("request cancelled"), "Cancelled"),
|
|
(KeyboardInterrupt(), "Cancelled"),
|
|
(asyncio.CancelledError(), "Cancelled"),
|
|
])
|
|
def test_exceptions_propagate_with_correct_status_and_no_thread_leak(exception, status):
|
|
stream = io.StringIO()
|
|
reporter = Progress(stream=stream)
|
|
with pytest.raises(type(exception)):
|
|
with reporter.stage("Working"):
|
|
thread = reporter._thread
|
|
raise exception
|
|
assert status in stream.getvalue()
|
|
assert "Completed" not in stream.getvalue()
|
|
assert not thread.is_alive()
|
|
|
|
|
|
def test_truncated_stage_and_summary_do_not_claim_completion():
|
|
stream = io.StringIO()
|
|
with Progress(stream=stream) as reporter:
|
|
with reporter.stage("Generating tokens", unit="tokens") as stage:
|
|
stage.advance(200)
|
|
thread = reporter._thread
|
|
stage.finish(status="truncated")
|
|
assert not thread.is_alive()
|
|
reporter.complete(8.0, 1.5, truncated=True)
|
|
output = stream.getvalue()
|
|
assert output.count("Finished (generation limit reached)") == 2
|
|
assert "Completed" not in output
|
|
assert "8.0s audio in 1.5s" in output
|
|
|
|
|
|
def test_complete_uses_stderr_and_needs_no_thread(monkeypatch, capsys):
|
|
def forbidden(*args, **kwargs):
|
|
raise AssertionError("Completion summary started a thread")
|
|
|
|
monkeypatch.setattr(progress_module.threading, "Thread", forbidden)
|
|
reporter = Progress()
|
|
reporter.complete(123.45, 6.78)
|
|
reporter.complete(123.45, 6.78)
|
|
capture = capsys.readouterr()
|
|
assert capture.out == ""
|
|
assert capture.err == "[YuE2] Completed: 123.5s audio in 6.8s\n"
|
|
|
|
|
|
def test_concurrent_callbacks_are_counted_exactly_and_do_not_change_rng():
|
|
before = random.getstate()
|
|
stream = io.StringIO()
|
|
with Progress(stream=stream).stage("Generating tokens", unit="tokens") as stage:
|
|
with ThreadPoolExecutor(max_workers=8) as pool:
|
|
list(pool.map(lambda _: stage.token("semantic", None), range(4000)))
|
|
assert stage.completed == 4000
|
|
assert "4000 tokens" in stream.getvalue()
|
|
assert random.getstate() == before
|
|
|
|
|
|
def test_module_import_and_quiet_usage_require_no_torch_or_rng_modules():
|
|
module_file = Path(progress_module.__file__).resolve()
|
|
code = f"""
|
|
import importlib.util, sys
|
|
spec = importlib.util.spec_from_file_location('standalone_progress', {str(module_file)!r})
|
|
module = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(module)
|
|
with module.Progress(enabled=False).stage('CPU test') as stage:
|
|
stage.advance()
|
|
assert not {{'torch', 'numpy', 'random'}}.intersection(sys.modules)
|
|
"""
|
|
subprocess.run([sys.executable, "-I", "-S", "-c", code], check=True, timeout=10)
|
|
|
|
|
|
def test_closed_output_stream_does_not_fail_work():
|
|
stream = io.StringIO()
|
|
stream.close()
|
|
with Progress(stream=stream).stage("Loading model") as stage:
|
|
stage.advance()
|
|
assert stage.completed == 1
|
|
|
|
|
|
def test_outer_context_cleans_up_an_unfinished_stage_on_failure():
|
|
stream = io.StringIO()
|
|
with pytest.raises(RuntimeError, match="load failed"):
|
|
with Progress(stream=stream) as reporter:
|
|
stage = reporter.stage("Loading model").__enter__()
|
|
thread = reporter._thread
|
|
raise RuntimeError("load failed")
|
|
assert not thread.is_alive()
|
|
assert stage.status == "failed"
|
|
assert "Completed" not in stream.getvalue()
|
|
|
|
|
|
def test_invalid_values_fail_without_spawning_threads():
|
|
for interval in [0, -1, float("nan"), float("inf")]:
|
|
with pytest.raises(ValueError):
|
|
Progress(refresh_interval=interval)
|
|
reporter = Progress(enabled=False)
|
|
for total in [-1, 1.5, True]:
|
|
with pytest.raises(ValueError):
|
|
reporter.stage("Invalid", total=total)
|
|
with reporter.stage("Valid", total=0) as stage:
|
|
stage.set_total(None)
|
|
stage.advance(2)
|
|
with pytest.raises(ValueError):
|
|
stage.update(1)
|
|
with pytest.raises(ValueError):
|
|
stage.finish(status="invented")
|
|
assert stage.status == "completed"
|
|
|
|
|
|
def test_tty_output_is_ascii_single_line_and_refresh_is_throttled(monkeypatch):
|
|
now = [0.0]
|
|
monkeypatch.setattr(progress_module, "_clock", lambda: now[0])
|
|
stream = Terminal()
|
|
with Progress(stream=stream, refresh_interval=0.25).stage("Generating\nASCII \x1b text", unit="tokens") as stage:
|
|
for _ in range(50):
|
|
stage.advance()
|
|
assert stream.writes == 1
|
|
now[0] += 0.25
|
|
stage.advance()
|
|
assert stream.writes == 2
|
|
output = stream.getvalue()
|
|
assert output.isascii()
|
|
assert "\x1b" not in output
|
|
assert output.count("\n") == 1
|
|
assert "51 tokens" in output
|
|
|
|
|
|
@pytest.mark.skipif(os.name != "posix", reason="PTY regression requires POSIX")
|
|
def test_zero_width_pty_preserves_stage_label_and_token_rate(monkeypatch):
|
|
import fcntl
|
|
import struct
|
|
import termios
|
|
|
|
monkeypatch.delenv("COLUMNS", raising=False)
|
|
now = [0.0]
|
|
monkeypatch.setattr(progress_module, "_clock", lambda: now[0])
|
|
master, slave = os.openpty()
|
|
try:
|
|
fcntl.ioctl(slave, termios.TIOCSWINSZ, struct.pack("HHHH", 0, 0, 0, 0))
|
|
assert os.get_terminal_size(slave).columns == 0
|
|
with os.fdopen(os.dup(slave), "w", buffering=1) as stream:
|
|
with Progress(stream=stream).stage("Generating song", unit="tokens") as stage:
|
|
now[0] = 10.0
|
|
stage.update(1500)
|
|
output = os.read(master, 8192).decode("ascii")
|
|
finally:
|
|
os.close(slave)
|
|
os.close(master)
|
|
assert "Completed Generating song: 1500 tokens | 150.0 tokens/s | 10.0s" in output
|
|
assert "[YuE2] | :" not in output
|
|
|
|
|
|
@pytest.mark.parametrize("columns", [20, 59])
|
|
def test_narrow_terminal_uses_throttled_unclipped_lines(monkeypatch, columns):
|
|
monkeypatch.delenv("COLUMNS", raising=False)
|
|
monkeypatch.setattr(Terminal, "fileno", lambda self: 2)
|
|
monkeypatch.setattr(progress_module.os, "get_terminal_size", lambda fd: os.terminal_size((columns, 24)))
|
|
now = [0.0]
|
|
monkeypatch.setattr(progress_module, "_clock", lambda: now[0])
|
|
stream = Terminal()
|
|
with Progress(stream=stream).stage("Generating song", unit="tokens") as stage:
|
|
for _ in range(100):
|
|
stage.advance()
|
|
assert len(stream.getvalue().splitlines()) == 1
|
|
now[0] = 5.0
|
|
stage.update(100)
|
|
output = stream.getvalue()
|
|
assert "\r" not in output
|
|
assert len(output.splitlines()) == 3
|
|
assert "Completed Generating song: 100 tokens | 20.0 tokens/s | elapsed 5.0s" in output
|
|
|
|
|
|
@pytest.mark.parametrize("override,expected", [("120", 120), ("0", 80), ("-2", 80), ("invalid", 80)])
|
|
def test_columns_override_requires_positive_integer(monkeypatch, override, expected):
|
|
monkeypatch.setenv("COLUMNS", override)
|
|
monkeypatch.setattr(Terminal, "fileno", lambda self: 2)
|
|
monkeypatch.setattr(progress_module.os, "get_terminal_size", lambda fd: os.terminal_size((0, 0)))
|
|
assert Progress(stream=Terminal())._columns == expected
|