105 lines
3.0 KiB
Python
105 lines
3.0 KiB
Python
#!/usr/bin/env python3
|
|
# client-batch.py: test batching via ace-server
|
|
#
|
|
# POST /lm (lm_batch_size=2 in JSON) -> 2 enriched requests
|
|
# POST /synth (JSON array of 2 requests) -> 2 MP3s in one GPU batch
|
|
#
|
|
# Start the server first: ./server.sh
|
|
|
|
import json
|
|
import sys
|
|
import urllib.error
|
|
import urllib.request
|
|
|
|
URL = "http://127.0.0.1:8085"
|
|
|
|
|
|
def post_json(endpoint, data):
|
|
body = json.dumps(data).encode()
|
|
req = urllib.request.Request(
|
|
URL + endpoint,
|
|
data=body,
|
|
headers={"Content-Type": "application/json"},
|
|
)
|
|
with urllib.request.urlopen(req) as resp:
|
|
return resp.read(), resp.headers
|
|
|
|
|
|
def parse_multipart_mixed(data, content_type):
|
|
"""Parse multipart/mixed response into list of body bytes."""
|
|
boundary = None
|
|
for part in content_type.split(";"):
|
|
part = part.strip()
|
|
if part.startswith("boundary="):
|
|
boundary = part[len("boundary="):].strip().encode()
|
|
break
|
|
if not boundary:
|
|
raise ValueError("no boundary in content-type: " + content_type)
|
|
|
|
delimiter = b"--" + boundary
|
|
parts = []
|
|
|
|
for chunk in data.split(delimiter):
|
|
if not chunk or chunk.startswith(b"--"):
|
|
continue
|
|
chunk = chunk.strip(b"\r\n")
|
|
if not chunk:
|
|
continue
|
|
|
|
sep = chunk.find(b"\r\n\r\n")
|
|
if sep < 0:
|
|
continue
|
|
body = chunk[sep + 4:]
|
|
if body.endswith(b"\r\n"):
|
|
body = body[:-2]
|
|
parts.append(body)
|
|
|
|
return parts
|
|
|
|
|
|
# Phase 1: LM generates N variations
|
|
try:
|
|
with open("simple-batch.json") as f:
|
|
request_json = json.load(f)
|
|
except FileNotFoundError:
|
|
print("ERROR: simple-batch.json not found (run from the examples/ directory)")
|
|
sys.exit(1)
|
|
|
|
try:
|
|
lm_batch_size = request_json.get("lm_batch_size", 1)
|
|
print("POST /lm (lm_batch_size=%d)..." % lm_batch_size)
|
|
lm_data, _ = post_json("/lm", request_json)
|
|
except urllib.error.URLError as e:
|
|
print("ERROR: cannot connect to %s (%s)" % (URL, e.reason))
|
|
print("Start the server first: ./server.sh")
|
|
sys.exit(1)
|
|
lm_results = json.loads(lm_data)
|
|
print(" -> %d enriched requests" % len(lm_results))
|
|
|
|
# Phase 2: synth all in one GPU batch (send JSON array)
|
|
print("POST /synth (batch=%d, JSON array)..." % len(lm_results))
|
|
body = json.dumps(lm_results).encode()
|
|
req = urllib.request.Request(
|
|
URL + "/synth",
|
|
data=body,
|
|
headers={"Content-Type": "application/json"},
|
|
)
|
|
with urllib.request.urlopen(req) as resp:
|
|
resp_data = resp.read()
|
|
content_type = resp.headers.get("Content-Type", "")
|
|
|
|
if "multipart/mixed" in content_type:
|
|
parts = parse_multipart_mixed(resp_data, content_type)
|
|
for i, mp3_data in enumerate(parts):
|
|
path = "server-batch%d.mp3" % i
|
|
with open(path, "wb") as f:
|
|
f.write(mp3_data)
|
|
print(" -> %s (%d bytes)" % (path, len(mp3_data)))
|
|
else:
|
|
path = "server-batch0.mp3"
|
|
with open(path, "wb") as f:
|
|
f.write(resp_data)
|
|
print(" -> %s (%d bytes)" % (path, len(resp_data)))
|
|
|
|
print("Done: %d MP3(s)" % len(lm_results))
|