42 lines
1.8 KiB
Python
42 lines
1.8 KiB
Python
"""CPU-only YuE2 worker regression tests; never import torch or run inference."""
|
|
import importlib.util
|
|
import os
|
|
from pathlib import Path
|
|
import unittest
|
|
from types import SimpleNamespace
|
|
|
|
spec = importlib.util.spec_from_file_location('yue2_worker', Path(__file__).parents[1] / 'scripts/yue2-worker.py')
|
|
worker = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(worker)
|
|
|
|
|
|
class WorkerTests(unittest.TestCase):
|
|
def test_lyrics_allow_multiple_sections_and_wrap_plain_text(self):
|
|
multi = '[Verse 1]\nHello\n\n[Chorus]\nSing it'
|
|
self.assertEqual(worker.normalize_lyrics(multi), multi + '\n')
|
|
self.assertEqual(worker.normalize_lyrics('[Verse 1]\nHello'), '[Verse 1]\nHello\n')
|
|
self.assertEqual(worker.normalize_lyrics('No heading yet'), '[song]\nNo heading yet\n\n')
|
|
with self.assertRaises(ValueError):
|
|
worker.normalize_lyrics('')
|
|
with self.assertRaises(ValueError):
|
|
worker.normalize_lyrics(' ')
|
|
|
|
def test_attention_backend_never_selects_flash(self):
|
|
torch_mod = SimpleNamespace(
|
|
cuda=SimpleNamespace(is_available=lambda: True),
|
|
backends=SimpleNamespace(cudnn=SimpleNamespace(is_available=lambda: True)),
|
|
)
|
|
attention, pipeline = worker.resolve_attention_backend(torch_mod)
|
|
self.assertEqual(attention, 'sdpa')
|
|
self.assertEqual(pipeline, 'torch-eager')
|
|
self.assertNotEqual(attention, 'flash')
|
|
self.assertTrue(os.environ.get('PYTORCH_CUDA_ALLOC_CONF', '').startswith('expandable_segments'))
|
|
|
|
def test_cuda_oom_detection(self):
|
|
self.assertTrue(worker.is_cuda_oom(RuntimeError('CUDA out of memory. Tried to allocate 2.49 GiB')))
|
|
self.assertFalse(worker.is_cuda_oom(RuntimeError('bad lyrics')))
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|