Fabio
c36510c21e
Some checks failed
Linux Build / Linux cpu (push) Has been cancelled
Linux Build / Linux cpu-clang (push) Has been cancelled
Linux Build / Linux vulkan (push) Has been cancelled
mac-build / macOS CPU (push) Has been cancelled
Nix Build / Nix (cpu) (push) Has been cancelled
Nix Build / Nix (vulkan) (push) Has been cancelled
windows-build / Windows CPU (push) Has been cancelled
103 lines
4.1 KiB
Python
103 lines
4.1 KiB
Python
"""Standard-library HTTP test server lifecycle helper."""
|
|
|
|
import argparse
|
|
import base64
|
|
import concurrent.futures as cf
|
|
import hashlib
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import socket
|
|
import subprocess
|
|
import threading
|
|
import time
|
|
import urllib.error
|
|
import urllib.request
|
|
|
|
|
|
class Server:
|
|
def __init__(self, exe, models, backend, device, out, parallel_jobs=False, extra_args=()):
|
|
self.exe, self.models, self.backend, self.device, self.out = exe, models, backend, device, out
|
|
self.parallel_jobs = parallel_jobs
|
|
self.extra_args = tuple(extra_args)
|
|
|
|
def __enter__(self):
|
|
self.out.mkdir(parents=True, exist_ok=True)
|
|
with socket.socket() as sock:
|
|
sock.bind(('127.0.0.1', 0))
|
|
port = sock.getsockname()[1]
|
|
self.url = f'http://127.0.0.1:{port}'
|
|
config = {'host': '127.0.0.1', 'port': port, 'backend': self.backend,
|
|
'threads': 8, 'device': self.device, 'lazy_load': True, 'max_loaded_models': 1,
|
|
'models': self.models}
|
|
path = self.out / 'config.json'
|
|
path.write_text(json.dumps(config, indent=2), encoding='utf-8')
|
|
self.log = (self.out / 'server.log').open('w', encoding='utf-8')
|
|
try:
|
|
self.proc = subprocess.Popen(
|
|
[str(self.exe), '--config', str(path), '--no-ui'] + (['--parallel-jobs'] if self.parallel_jobs else []) + list(self.extra_args),
|
|
cwd=self.exe.parent, stdout=self.log, stderr=subprocess.STDOUT,
|
|
creationflags=subprocess.CREATE_NO_WINDOW if os.name == 'nt' else 0)
|
|
except BaseException:
|
|
self.log.close()
|
|
raise
|
|
try:
|
|
deadline = time.monotonic() + 120
|
|
while True:
|
|
if self.proc.poll() is not None:
|
|
raise RuntimeError(f'server exited; see {self.out / "server.log"}')
|
|
try:
|
|
if self.api('/health')[0] == 200:
|
|
return self
|
|
except (urllib.error.URLError, TimeoutError):
|
|
if time.monotonic() > deadline:
|
|
raise
|
|
time.sleep(.05)
|
|
except BaseException:
|
|
self.__exit__(None, None, None)
|
|
raise
|
|
|
|
def __exit__(self, *_):
|
|
if self.proc.poll() is None:
|
|
self.proc.terminate()
|
|
try:
|
|
self.proc.wait(timeout=30)
|
|
except subprocess.TimeoutExpired:
|
|
self.proc.kill()
|
|
self.proc.wait(timeout=30)
|
|
finally:
|
|
self.log.close()
|
|
|
|
def api(self, path, body=None):
|
|
request = urllib.request.Request(self.url + path,
|
|
data=json.dumps(body).encode() if body is not None else None,
|
|
headers={'Content-Type': 'application/json'})
|
|
try:
|
|
with urllib.request.urlopen(request, timeout=180) as response:
|
|
return response.status, json.loads(response.read())
|
|
except urllib.error.HTTPError as error:
|
|
return error.code, json.loads(error.read())
|
|
|
|
def status(self, model='higgs'):
|
|
return next(m for m in self.api('/v1/models')[1]['data'] if m['id'] == model)
|
|
|
|
def run(self, request, name, model='higgs', timeout=None):
|
|
body = {'model': model, 'request': request}
|
|
if timeout is not None:
|
|
body['busy_timeout_ms'] = timeout
|
|
started = time.perf_counter()
|
|
code, result = self.api('/v1/tasks/run', body)
|
|
elapsed = time.perf_counter() - started
|
|
audios = {}
|
|
if 'audio' in result:
|
|
audios['audio'] = base64.b64decode(result['audio'])
|
|
for audio in result.get('named_audio_outputs', []):
|
|
audios[audio['id']] = base64.b64decode(audio['audio'])
|
|
for key, audio in audios.items():
|
|
(self.out / f'{name}-{key}.wav').write_bytes(audio)
|
|
record = {'name': name, 'status': code, 'http_s': elapsed, 'timing': result.get('timing'),
|
|
'sha256': {k: hashlib.sha256(v).hexdigest() for k, v in audios.items()}}
|
|
if code != 200:
|
|
record['error'] = result
|
|
(self.out / f'{name}.json').write_text(json.dumps(record, indent=2))
|
|
return record, audios
|