Initial release

engine/examples/client-batch.py

@@ -0,0 +1,104 @@
+#!/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))