fix: use upload path instead of inline base64 to prevent API OOM
The inline path loaded entire base64 payloads into API memory via Plug.Parsers, then decoded them again for UploadStore — 4 concurrent 45MB files consumed ~420MB on a 768MB droplet, pushing BEAM into swap and timing out all RPCs. Now uses POST /upload (streaming 64KB chunks, constant memory) then references upload_ids in the execute call. The JSON body drops from ~60MB to ~500 bytes. Zero base64 encoding on the request side.
This commit is contained in:
parent
3fbef65b20
commit
2e73c09e63
1 changed files with 90 additions and 91 deletions
|
|
@ -1,11 +1,11 @@
|
||||||
"""Karaoke mode — vocal isolation via unsandbox zerotrust containers.
|
"""Karaoke mode — vocal isolation via unsandbox zerotrust containers.
|
||||||
|
|
||||||
Downloads media from S3, sends it + voxsplit.c into a zerotrust container
|
Downloads media from S3, uploads it + voxsplit.c via the /upload endpoint
|
||||||
(no network, destroyed after use), gets back both instrumentals and vocals,
|
(streaming, constant memory), then kicks off /execute referencing upload_ids.
|
||||||
uploads them to S3.
|
Gets back both instrumentals and vocals, uploads them to S3.
|
||||||
|
|
||||||
Disk-backed pipeline: every stage streams through tmpfiles so peak memory
|
Upload-based pipeline: files stream directly from disk to the API in 64KB
|
||||||
stays ~64KB per worker on the request side instead of ~1GB.
|
chunks — never base64-encoded, never held in memory on either side.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
import base64
|
import base64
|
||||||
|
|
@ -25,83 +25,30 @@ from .un import API_BASE, _resolve_credentials, validate_keys
|
||||||
|
|
||||||
log = logging.getLogger(__name__)
|
log = logging.getLogger(__name__)
|
||||||
|
|
||||||
# 49152 = 3 * 16384 — multiple of 3, so intermediate base64 chunks produce
|
|
||||||
# no '=' padding; concatenation equals whole-file encoding.
|
|
||||||
_B64_CHUNK_RAW = 3 * 16384
|
|
||||||
_READ_CHUNK = 65536 # 64KB streaming buffer
|
_READ_CHUNK = 65536 # 64KB streaming buffer
|
||||||
|
|
||||||
|
|
||||||
VOXSPLIT_SOURCE = os.environ.get(
|
VOXSPLIT_SOURCE = os.environ.get(
|
||||||
"VOXSPLIT_SOURCE_PATH",
|
"VOXSPLIT_SOURCE_PATH",
|
||||||
os.path.join(os.path.dirname(__file__), "voxsplit.c"),
|
os.path.join(os.path.dirname(__file__), "voxsplit.c"),
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
def _stream_b64_encode(src_path, out_file):
|
def _sign_body(secret_key, timestamp, method, path, body_bytes):
|
||||||
"""Stream base64-encode *src_path* into *out_file* in aligned chunks.
|
"""HMAC-SHA256 matching un._sign_request for small payloads."""
|
||||||
|
h = hmac.new(secret_key.encode(), digestmod=hashlib.sha256)
|
||||||
Each intermediate chunk (49152 raw bytes = 3*16384) encodes to base64
|
h.update(f"{timestamp}:{method}:{path}:{body_bytes}".encode())
|
||||||
with no padding, so the concatenation matches whole-file encoding.
|
return h.hexdigest()
|
||||||
"""
|
|
||||||
with open(src_path, "rb") as f:
|
|
||||||
while True:
|
|
||||||
chunk = f.read(_B64_CHUNK_RAW)
|
|
||||||
if not chunk:
|
|
||||||
break
|
|
||||||
out_file.write(base64.b64encode(chunk))
|
|
||||||
|
|
||||||
|
|
||||||
def _build_json_on_disk(tmpdir, media_path, voxsplit_bin_path, script):
|
def _sign_from_disk(secret_key, timestamp, method, path, file_path):
|
||||||
"""Build the /execute JSON payload on disk without holding media in RAM.
|
|
||||||
|
|
||||||
Writes prefix → streaming base64 media → middle → voxsplit b64 → suffix.
|
|
||||||
Returns path to the request.json file.
|
|
||||||
"""
|
|
||||||
media_ph = "MEDIA_B64_SENTINEL_48291"
|
|
||||||
voxsplit_ph = "VOXSPLIT_B64_SENTINEL_48291"
|
|
||||||
|
|
||||||
skeleton = {
|
|
||||||
"language": "bash",
|
|
||||||
"code": script,
|
|
||||||
"network_mode": "zerotrust",
|
|
||||||
"ttl": 300,
|
|
||||||
"vcpu": 2,
|
|
||||||
"input_files": [
|
|
||||||
{"filename": "media", "content": media_ph},
|
|
||||||
{"filename": "voxsplit.c", "content": voxsplit_ph},
|
|
||||||
],
|
|
||||||
}
|
|
||||||
skeleton_json = json.dumps(skeleton)
|
|
||||||
|
|
||||||
# Locate the quoted sentinel strings and split around them
|
|
||||||
media_quoted = json.dumps(media_ph) # '"MEDIA_B64_SENTINEL_48291"'
|
|
||||||
voxsplit_quoted = json.dumps(voxsplit_ph)
|
|
||||||
|
|
||||||
idx1 = skeleton_json.index(media_quoted)
|
|
||||||
idx2 = skeleton_json.index(voxsplit_quoted)
|
|
||||||
|
|
||||||
prefix = skeleton_json[:idx1] + '"'
|
|
||||||
middle = '"' + skeleton_json[idx1 + len(media_quoted):idx2] + '"'
|
|
||||||
suffix = '"' + skeleton_json[idx2 + len(voxsplit_quoted):]
|
|
||||||
|
|
||||||
json_path = os.path.join(tmpdir, "request.json")
|
|
||||||
with open(json_path, "wb") as out:
|
|
||||||
out.write(prefix.encode())
|
|
||||||
_stream_b64_encode(media_path, out)
|
|
||||||
out.write(middle.encode())
|
|
||||||
_stream_b64_encode(voxsplit_bin_path, out)
|
|
||||||
out.write(suffix.encode())
|
|
||||||
|
|
||||||
return json_path
|
|
||||||
|
|
||||||
|
|
||||||
def _sign_from_disk(secret_key, timestamp, method, path, json_path):
|
|
||||||
"""Incremental HMAC-SHA256 matching un._sign_request but reading from disk.
|
"""Incremental HMAC-SHA256 matching un._sign_request but reading from disk.
|
||||||
|
|
||||||
Produces: HMAC(key, "ts:METHOD:path:" + <file contents>)
|
Produces: HMAC(key, "ts:METHOD:path:" + <file contents>)
|
||||||
"""
|
"""
|
||||||
h = hmac.new(secret_key.encode(), digestmod=hashlib.sha256)
|
h = hmac.new(secret_key.encode(), digestmod=hashlib.sha256)
|
||||||
h.update(f"{timestamp}:{method}:{path}:".encode())
|
h.update(f"{timestamp}:{method}:{path}:".encode())
|
||||||
with open(json_path, "rb") as f:
|
with open(file_path, "rb") as f:
|
||||||
while True:
|
while True:
|
||||||
chunk = f.read(_READ_CHUNK)
|
chunk = f.read(_READ_CHUNK)
|
||||||
if not chunk:
|
if not chunk:
|
||||||
|
|
@ -110,16 +57,61 @@ def _sign_from_disk(secret_key, timestamp, method, path, json_path):
|
||||||
return h.hexdigest()
|
return h.hexdigest()
|
||||||
|
|
||||||
|
|
||||||
def _make_streaming_request(json_path, public_key, secret_key, tmpdir):
|
def _upload_file(file_path, filename, public_key, secret_key):
|
||||||
"""POST the on-disk JSON to /execute, stream response to disk.
|
"""Upload a file via POST /upload (streaming, constant memory).
|
||||||
|
|
||||||
Returns path to response.json.
|
Returns the upload_id string.
|
||||||
"""
|
"""
|
||||||
timestamp = int(time.time())
|
timestamp = int(time.time())
|
||||||
method = "POST"
|
method = "POST"
|
||||||
|
path = "/upload"
|
||||||
|
|
||||||
|
signature = _sign_from_disk(secret_key, timestamp, method, path, file_path)
|
||||||
|
|
||||||
|
headers = {
|
||||||
|
"Authorization": f"Bearer {public_key}",
|
||||||
|
"X-Timestamp": str(timestamp),
|
||||||
|
"X-Signature": signature,
|
||||||
|
"X-Filename": filename,
|
||||||
|
"Content-Type": "application/octet-stream",
|
||||||
|
}
|
||||||
|
|
||||||
|
url = f"{API_BASE}{path}"
|
||||||
|
|
||||||
|
with open(file_path, "rb") as body:
|
||||||
|
response = requests.post(
|
||||||
|
url, data=body, headers=headers, stream=False, timeout=300,
|
||||||
|
)
|
||||||
|
|
||||||
|
response.raise_for_status()
|
||||||
|
result = response.json()
|
||||||
|
return result["upload_id"]
|
||||||
|
|
||||||
|
|
||||||
|
def _execute_with_uploads(upload_ids, script, public_key, secret_key, tmpdir):
|
||||||
|
"""POST /execute referencing upload_ids instead of inline content.
|
||||||
|
|
||||||
|
The JSON body contains only metadata — no file content. Returns path
|
||||||
|
to response.json on disk.
|
||||||
|
"""
|
||||||
|
payload = {
|
||||||
|
"language": "bash",
|
||||||
|
"code": script,
|
||||||
|
"network_mode": "zerotrust",
|
||||||
|
"ttl": 300,
|
||||||
|
"vcpu": 2,
|
||||||
|
"input_files": [
|
||||||
|
{"upload_id": uid, "filename": fname}
|
||||||
|
for fname, uid in upload_ids
|
||||||
|
],
|
||||||
|
}
|
||||||
|
|
||||||
|
body_str = json.dumps(payload)
|
||||||
|
timestamp = int(time.time())
|
||||||
|
method = "POST"
|
||||||
path = "/execute"
|
path = "/execute"
|
||||||
|
|
||||||
signature = _sign_from_disk(secret_key, timestamp, method, path, json_path)
|
signature = _sign_body(secret_key, timestamp, method, path, body_str)
|
||||||
|
|
||||||
headers = {
|
headers = {
|
||||||
"Authorization": f"Bearer {public_key}",
|
"Authorization": f"Bearer {public_key}",
|
||||||
|
|
@ -129,11 +121,9 @@ def _make_streaming_request(json_path, public_key, secret_key, tmpdir):
|
||||||
}
|
}
|
||||||
|
|
||||||
url = f"{API_BASE}{path}"
|
url = f"{API_BASE}{path}"
|
||||||
|
response = requests.post(
|
||||||
with open(json_path, "rb") as body:
|
url, data=body_str.encode(), headers=headers, stream=True, timeout=300,
|
||||||
response = requests.post(
|
)
|
||||||
url, data=body, headers=headers, stream=True, timeout=300,
|
|
||||||
)
|
|
||||||
|
|
||||||
response.raise_for_status()
|
response.raise_for_status()
|
||||||
|
|
||||||
|
|
@ -213,8 +203,9 @@ def process_karaoke(s3_client, bucket, s3_key, s3_path,
|
||||||
skip_validate=False):
|
skip_validate=False):
|
||||||
"""Download from S3 -> process in unsandbox zerotrust -> upload both tracks.
|
"""Download from S3 -> process in unsandbox zerotrust -> upload both tracks.
|
||||||
|
|
||||||
Disk-backed pipeline: streams through tmpfiles so peak memory stays ~64KB
|
Upload-based pipeline: files stream to the API via POST /upload (64KB
|
||||||
per worker on the request side.
|
chunks, constant memory on both sides), then /execute references them
|
||||||
|
by upload_id. Zero base64 encoding on the request side.
|
||||||
|
|
||||||
Produces two files under s3_path:
|
Produces two files under s3_path:
|
||||||
{s3_path}/instrumentals
|
{s3_path}/instrumentals
|
||||||
|
|
@ -234,9 +225,7 @@ def process_karaoke(s3_client, bucket, s3_key, s3_path,
|
||||||
log.info("process_karaoke: bucket=%s key=%s s3_path=%s", bucket, s3_key, s3_path)
|
log.info("process_karaoke: bucket=%s key=%s s3_path=%s", bucket, s3_key, s3_path)
|
||||||
|
|
||||||
with tempfile.TemporaryDirectory(prefix="karaoke_") as tmpdir:
|
with tempfile.TemporaryDirectory(prefix="karaoke_") as tmpdir:
|
||||||
# 1. Stream S3 download to disk (get_object, not download_file —
|
# 1. Stream S3 download to disk
|
||||||
# download_file uses boto3 TransferManager which spawns threads
|
|
||||||
# that break under uWSGI worker recycling)
|
|
||||||
media_path = os.path.join(tmpdir, "media.bin")
|
media_path = os.path.join(tmpdir, "media.bin")
|
||||||
resp = s3_client.get_object(Bucket=bucket, Key=s3_key)
|
resp = s3_client.get_object(Bucket=bucket, Key=s3_key)
|
||||||
with open(media_path, "wb") as f:
|
with open(media_path, "wb") as f:
|
||||||
|
|
@ -247,12 +236,27 @@ def process_karaoke(s3_client, bucket, s3_key, s3_path,
|
||||||
break
|
break
|
||||||
f.write(chunk)
|
f.write(chunk)
|
||||||
|
|
||||||
# 2. Prepare voxsplit binary (text -> UTF-8 bytes, matching original)
|
# 2. Prepare voxsplit binary
|
||||||
voxsplit_bin = os.path.join(tmpdir, "voxsplit.bin")
|
voxsplit_bin = os.path.join(tmpdir, "voxsplit.bin")
|
||||||
with open(VOXSPLIT_SOURCE, "r") as src, open(voxsplit_bin, "wb") as dst:
|
with open(VOXSPLIT_SOURCE, "r") as src, open(voxsplit_bin, "wb") as dst:
|
||||||
dst.write(src.read().encode())
|
dst.write(src.read().encode())
|
||||||
|
|
||||||
# 3. Build execution script — extract audio to WAV first, then split
|
# 3. Upload both files via /upload (streaming, constant memory)
|
||||||
|
try:
|
||||||
|
media_upload_id = _upload_file(
|
||||||
|
media_path, "media", public_key, secret_key,
|
||||||
|
)
|
||||||
|
os.unlink(media_path)
|
||||||
|
|
||||||
|
voxsplit_upload_id = _upload_file(
|
||||||
|
voxsplit_bin, "voxsplit.c", public_key, secret_key,
|
||||||
|
)
|
||||||
|
os.unlink(voxsplit_bin)
|
||||||
|
except Exception as e:
|
||||||
|
log.warning("Unsandbox upload failed: %s", e)
|
||||||
|
return None
|
||||||
|
|
||||||
|
# 4. Build execution script
|
||||||
extract_wav = (
|
extract_wav = (
|
||||||
"gcc -O2 -o /tmp/voxsplit /tmp/input/voxsplit.c -lm && "
|
"gcc -O2 -o /tmp/voxsplit /tmp/input/voxsplit.c -lm && "
|
||||||
"ffmpeg -y -i /tmp/input/media -vn -acodec pcm_s16le -ar 44100 -ac 2 /tmp/audio.wav && "
|
"ffmpeg -y -i /tmp/input/media -vn -acodec pcm_s16le -ar 44100 -ac 2 /tmp/audio.wav && "
|
||||||
|
|
@ -276,20 +280,15 @@ def process_karaoke(s3_client, bucket, s3_key, s3_path,
|
||||||
"cp /tmp/split-vocal.wav /tmp/artifacts/vocals.wav"
|
"cp /tmp/split-vocal.wav /tmp/artifacts/vocals.wav"
|
||||||
)
|
)
|
||||||
|
|
||||||
# 4. Build JSON payload on disk (streaming base64)
|
# 5. Execute with upload_id references (tiny JSON body, no file content)
|
||||||
json_path = _build_json_on_disk(tmpdir, media_path, voxsplit_bin, script)
|
|
||||||
os.unlink(media_path)
|
|
||||||
os.unlink(voxsplit_bin)
|
|
||||||
|
|
||||||
# 5. Send request, stream response to disk
|
|
||||||
try:
|
try:
|
||||||
resp_path = _make_streaming_request(
|
resp_path = _execute_with_uploads(
|
||||||
json_path, public_key, secret_key, tmpdir,
|
[("media", media_upload_id), ("voxsplit.c", voxsplit_upload_id)],
|
||||||
|
script, public_key, secret_key, tmpdir,
|
||||||
)
|
)
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
log.warning("Unsandbox execute failed: %s", e)
|
log.warning("Unsandbox execute failed: %s", e)
|
||||||
return None
|
return None
|
||||||
os.unlink(json_path)
|
|
||||||
|
|
||||||
# 6. Process response, upload artifacts to S3
|
# 6. Process response, upload artifacts to S3
|
||||||
return _process_response_from_disk(
|
return _process_response_from_disk(
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue