75 lines
6.0 KiB
Python
75 lines
6.0 KiB
Python
"""Bounded WAV tests with coordinated GPU ownership and disposable input."""
|
|
import io,json,os,signal,subprocess,threading,time,uuid,wave
|
|
from pathlib import Path
|
|
from image_test import probe
|
|
from separator_runtime import MODELS
|
|
class SeparatorTests:
|
|
def __init__(self,root,runtime,scheduler):
|
|
self.root=Path(root);self.runtime=runtime;self.scheduler=scheduler;self.lock=threading.RLock();self.job=None;self.process=None;self.cancel=threading.Event();self.thread=None
|
|
def status(self):
|
|
with self.lock:return dict(job=dict(self.job) if self.job else None)
|
|
def start(self,model,audio,target=None,full=False):
|
|
if full and self.scheduler.gpu_mode!='separator':raise ValueError('Bitte unter Steuerung Audio-Trennung starten.')
|
|
targets={'vocals':'model_bs_roformer_ep_317_sdr_12.9755.ckpt','drums':'htdemucs_ft.yaml','bass':'htdemucs_ft.yaml','guitar':'htdemucs_6s.yaml','piano':'htdemucs_6s.yaml','other':'htdemucs_6s.yaml'}
|
|
if target is not None and targets.get(target)!=model:raise ValueError('Zielspur und Modellpaket passen nicht zusammen.')
|
|
if not isinstance(model,str) or model not in MODELS or not self.runtime.status()['installed'] or not self.runtime.model_ready(model):raise ValueError('Laufzeit und Modellpaket zuerst unter Audio Separator einrichten.')
|
|
if not isinstance(audio,bytes) or not 44<=len(audio)<=(256 if full else 8)*1024**2:raise ValueError('WAV überschreitet die erlaubte Upload-Größe.')
|
|
try:
|
|
with wave.open(io.BytesIO(audio)) as wav:
|
|
if wav.getnchannels() not in (1,2) or wav.getsampwidth()!=2 or not 8000<=wav.getframerate()<=48000 or not 0<wav.getnframes()/wav.getframerate()<=(1200 if full else 30):raise ValueError('PCM-WAV, 16 Bit, Mono/Stereo, maximal 30 Sekunden erforderlich.')
|
|
if len(wav.readframes(wav.getnframes()))!=wav.getnframes()*wav.getnchannels()*2:raise ValueError('WAV-Datei ist unvollständig.')
|
|
except (wave.Error,EOFError):raise ValueError('Ungültige WAV-Datei.') from None
|
|
with self.lock:
|
|
if self.job and self.job['state']=='running':raise ValueError('Trenntest läuft bereits.')
|
|
self.target=target;self.full=full;self.cancel.clear();self.root.mkdir(parents=True,exist_ok=True,mode=0o700);self.job=dict(id=uuid.uuid4().hex,state='running',phase='Wartet auf GPU-Freigabe',started_at=time.time(),files=[])
|
|
self.thread=threading.Thread(target=self._run,args=(model,audio),daemon=True);self.thread.start();return dict(self.job)
|
|
def _run(self,model,audio):
|
|
release=None;directory=self.root/self.job['id'];state='failed';phase='Trennung fehlgeschlagen.'
|
|
try:
|
|
release=self.scheduler.image_reservation(kind='separator')
|
|
gpu=next((g for term in ('RTX 3060','RTX 5080') for g in probe() if term in g['name'] and not g['processes'] and g['free_mib']>4000),None)
|
|
if not gpu:raise ValueError('Keine ausreichend freie GPU. Fremde Dienste werden nicht beendet.')
|
|
directory.mkdir(mode=0o700);(directory/'input.wav').write_bytes(audio);diagnostics=(directory/'worker.tmp').open('wb')
|
|
with self.lock:
|
|
if self.cancel.is_set():raise InterruptedError()
|
|
self.job.update(phase='Modell laden und Audio trennen',gpu=gpu['name'])
|
|
self.process=subprocess.Popen([str(self.runtime.paths()[0]),str(Path(__file__).with_name('separator_worker.py')),str(self.runtime.root/'models'),model,str(directory/'input.wav'),str(directory),getattr(self,'target',None) or ''],stdout=subprocess.DEVNULL,stderr=diagnostics,start_new_session=True,env=dict(os.environ,CUDA_VISIBLE_DEVICES=gpu['uuid'],OMP_NUM_THREADS='6'))
|
|
# Drain diagnostics into a bounded temporary file rather than retaining user data.
|
|
deadline=time.monotonic()+(7200 if getattr(self,'full',False) else 600)
|
|
while self.process.poll() is None:
|
|
if self.cancel.wait(.2):raise InterruptedError()
|
|
if time.monotonic()>deadline:raise ValueError('Trennauftrag überschreitet das Zeitlimit.')
|
|
with self.lock:self.job['elapsed_seconds']=round(time.time()-self.job['started_at'])
|
|
diagnostics.close();error=(directory/'worker.tmp').read_bytes()[-65536:].decode(errors='replace')
|
|
if self.process.returncode:raise ValueError('GPU-Speicher reicht nicht (OOM).' if 'out of memory' in error.lower() else 'Audio Separator hat den Test abgebrochen. Modellpaket und Laufzeit prüfen.')
|
|
names=json.loads((directory/'results.json').read_text())
|
|
if not names or any(Path(n).name!=n or not n.endswith('.wav') or not (directory/n).is_file() for n in names):raise ValueError('Ungültige Ergebnisdateien.')
|
|
with self.lock:self.job['files']=[dict(index=i,name=n) for i,n in enumerate(names)]
|
|
(directory/'complete').touch();state='complete';phase='Spuren fertig · Modell entladen'
|
|
except InterruptedError:state='cancelled';phase='Trenntest abgebrochen'
|
|
except Exception as exc:phase=str(exc) if isinstance(exc,ValueError) else 'Trenntest fehlgeschlagen.'
|
|
finally:
|
|
process=self.process
|
|
if process and process.poll() is None:
|
|
try:os.killpg(process.pid,signal.SIGTERM);process.wait(5)
|
|
except subprocess.TimeoutExpired:os.killpg(process.pid,signal.SIGKILL);process.wait()
|
|
except ProcessLookupError:pass
|
|
if 'diagnostics' in locals():diagnostics.close()
|
|
(directory/'input.wav').unlink(missing_ok=True);(directory/'worker.tmp').unlink(missing_ok=True)
|
|
with self.lock:self.process=None;self.job.update(state=state,phase=phase,finished_at=time.time())
|
|
if release:release()
|
|
def stop(self):
|
|
self.cancel.set()
|
|
if self.thread:self.thread.join(7)
|
|
return {'cancellation_requested':True}
|
|
def audio(self,ident,index):
|
|
if not isinstance(ident,str) or len(ident)!=32 or any(c not in '0123456789abcdef' for c in ident):raise ValueError('Ungültige Auftrag-ID.')
|
|
directory=self.root/ident
|
|
if not (directory/'complete').is_file():raise ValueError('Auftrag nicht fertig.')
|
|
names=json.loads((directory/'results.json').read_text())
|
|
if type(index) is not int or not 0<=index<len(names) or Path(names[index]).name!=names[index]:raise ValueError('Ungültige Spur.')
|
|
path=directory/names[index]
|
|
if path.is_symlink() or path.resolve().parent!=directory.resolve():raise ValueError('Ungültiger Ergebnispfad.')
|
|
if path.stat().st_size>512*1024**2:raise ValueError('Spur zu groß.')
|
|
return path.read_bytes()
|