fix: disk-backed karaoke pipeline to prevent OOM kills
process_karaoke held 4+ copies of every media file in memory simultaneously (raw bytes, base64, JSON serialization). With 14-16 concurrent workers on large files, memory exploded past 3.2GB RSS. Replace all in-memory buffers with a disk-backed pipeline: - Stream S3 download to tmpfile instead of .read() - Build JSON payload on disk with streaming base64 encoding - Incremental HMAC-SHA256 signing from disk - Stream HTTP request body from file, response to file - Decode artifacts to tmpfiles one at a time for S3 upload Peak request-side memory drops from ~1GB/worker to ~64KB/worker. Function signature unchanged — all callers work without modification.
This commit is contained in:
parent
e87bb0b0c3
commit
5736ff3330
1 changed files with 216 additions and 63 deletions
|
|
@ -3,30 +3,219 @@
|
|||
Downloads media from S3, sends it + voxsplit.c into a zerotrust container
|
||||
(no network, destroyed after use), gets back both instrumentals and vocals,
|
||||
uploads them to S3.
|
||||
|
||||
Disk-backed pipeline: every stage streams through tmpfiles so peak memory
|
||||
stays ~64KB per worker on the request side instead of ~1GB.
|
||||
"""
|
||||
|
||||
import base64
|
||||
import fcntl
|
||||
import hashlib
|
||||
import hmac
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import tempfile
|
||||
import time
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
|
||||
from .un import _make_request, _resolve_credentials, validate_keys
|
||||
import requests
|
||||
|
||||
from .un import API_BASE, _resolve_credentials, validate_keys
|
||||
|
||||
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
|
||||
|
||||
VOXSPLIT_SOURCE = os.environ.get(
|
||||
"VOXSPLIT_SOURCE_PATH",
|
||||
os.path.join(os.path.dirname(__file__), "voxsplit.c"),
|
||||
)
|
||||
|
||||
|
||||
def _stream_b64_encode(src_path, out_file):
|
||||
"""Stream base64-encode *src_path* into *out_file* in aligned chunks.
|
||||
|
||||
Each intermediate chunk (49152 raw bytes = 3*16384) encodes to base64
|
||||
with no padding, so the concatenation matches whole-file encoding.
|
||||
"""
|
||||
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):
|
||||
"""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.
|
||||
|
||||
Produces: HMAC(key, "ts:METHOD:path:" + <file contents>)
|
||||
"""
|
||||
h = hmac.new(secret_key.encode(), digestmod=hashlib.sha256)
|
||||
h.update(f"{timestamp}:{method}:{path}:".encode())
|
||||
with open(json_path, "rb") as f:
|
||||
while True:
|
||||
chunk = f.read(_READ_CHUNK)
|
||||
if not chunk:
|
||||
break
|
||||
h.update(chunk)
|
||||
return h.hexdigest()
|
||||
|
||||
|
||||
def _make_streaming_request(json_path, public_key, secret_key, tmpdir):
|
||||
"""POST the on-disk JSON to /execute, stream response to disk.
|
||||
|
||||
Returns path to response.json.
|
||||
"""
|
||||
timestamp = int(time.time())
|
||||
method = "POST"
|
||||
path = "/execute"
|
||||
|
||||
signature = _sign_from_disk(secret_key, timestamp, method, path, json_path)
|
||||
|
||||
headers = {
|
||||
"Authorization": f"Bearer {public_key}",
|
||||
"X-Timestamp": str(timestamp),
|
||||
"X-Signature": signature,
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
url = f"{API_BASE}{path}"
|
||||
|
||||
with open(json_path, "rb") as body:
|
||||
response = requests.post(
|
||||
url, data=body, headers=headers, stream=True, timeout=300,
|
||||
)
|
||||
|
||||
response.raise_for_status()
|
||||
|
||||
resp_path = os.path.join(tmpdir, "response.json")
|
||||
with open(resp_path, "wb") as f:
|
||||
for chunk in response.iter_content(chunk_size=_READ_CHUNK):
|
||||
if chunk:
|
||||
f.write(chunk)
|
||||
|
||||
return resp_path
|
||||
|
||||
|
||||
def _process_response_from_disk(resp_path, s3_client, bucket, s3_path,
|
||||
is_video, extension, tmpdir):
|
||||
"""Parse response JSON, decode artifacts to tmpfiles, upload to S3.
|
||||
|
||||
Returns {"instrumentals": size, "vocals": size} or None.
|
||||
"""
|
||||
with open(resp_path, "r") as f:
|
||||
result = json.load(f)
|
||||
os.unlink(resp_path)
|
||||
|
||||
if not result.get("success") or result.get("exit_code") != 0:
|
||||
log.warning(
|
||||
"Unsandbox execute non-zero exit: success=%s exit_code=%s "
|
||||
"stdout=%.2000s stderr=%.2000s",
|
||||
result.get("success"), result.get("exit_code"),
|
||||
result.get("stdout", ""), result.get("stderr", ""),
|
||||
)
|
||||
return None
|
||||
|
||||
artifacts = result.get("artifacts", [])
|
||||
if len(artifacts) < 2:
|
||||
return None
|
||||
|
||||
artifact_map = {}
|
||||
for a in artifacts:
|
||||
fname = a.get("filename", "")
|
||||
if fname.startswith("instrumentals"):
|
||||
artifact_map["instrumentals"] = a
|
||||
elif fname.startswith("vocals"):
|
||||
artifact_map["vocals"] = a
|
||||
|
||||
if "instrumentals" not in artifact_map or "vocals" not in artifact_map:
|
||||
return None
|
||||
|
||||
sizes = {}
|
||||
content_type = "video/mp4" if is_video else "audio/wav"
|
||||
|
||||
for track_name in ("instrumentals", "vocals"):
|
||||
artifact = artifact_map[track_name]
|
||||
decoded = base64.b64decode(artifact["content_base64"])
|
||||
|
||||
track_path = os.path.join(tmpdir, f"{track_name}.bin")
|
||||
with open(track_path, "wb") as f:
|
||||
f.write(decoded)
|
||||
size = len(decoded)
|
||||
del decoded
|
||||
del artifact["content_base64"]
|
||||
|
||||
with open(track_path, "rb") as f:
|
||||
s3_client.put_object(
|
||||
Bucket=bucket,
|
||||
Key=f"{s3_path}/{track_name}",
|
||||
Body=f,
|
||||
CacheControl="private, max-age=172800",
|
||||
ContentType=content_type,
|
||||
)
|
||||
sizes[track_name] = size
|
||||
os.unlink(track_path)
|
||||
|
||||
return sizes
|
||||
|
||||
|
||||
def process_karaoke(s3_client, bucket, s3_key, s3_path,
|
||||
is_video, extension, public_key=None, secret_key=None,
|
||||
skip_validate=False):
|
||||
"""Download from S3 -> process in unsandbox zerotrust -> upload both tracks.
|
||||
|
||||
Disk-backed pipeline: streams through tmpfiles so peak memory stays ~64KB
|
||||
per worker on the request side.
|
||||
|
||||
Produces two files under s3_path:
|
||||
{s3_path}/instrumentals
|
||||
{s3_path}/vocals
|
||||
|
|
@ -45,17 +234,23 @@ 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)
|
||||
|
||||
with tempfile.TemporaryDirectory(prefix="karaoke_") as tmpdir:
|
||||
# 1. Download original from S3 (get_object, not download_file —
|
||||
# 1. Stream S3 download to disk (get_object, not download_file —
|
||||
# download_file uses boto3 TransferManager which spawns threads
|
||||
# that break under uWSGI worker recycling)
|
||||
media_path = os.path.join(tmpdir, "media.bin")
|
||||
resp = s3_client.get_object(Bucket=bucket, Key=s3_key)
|
||||
media_bytes = resp["Body"].read()
|
||||
with open(media_path, "wb") as f:
|
||||
body = resp["Body"]
|
||||
while True:
|
||||
chunk = body.read(_READ_CHUNK)
|
||||
if not chunk:
|
||||
break
|
||||
f.write(chunk)
|
||||
|
||||
# 2. Base64 encode
|
||||
media_b64 = base64.b64encode(media_bytes).decode()
|
||||
with open(VOXSPLIT_SOURCE, "r") as f:
|
||||
voxsplit_src = f.read()
|
||||
voxsplit_b64 = base64.b64encode(voxsplit_src.encode()).decode()
|
||||
# 2. Prepare voxsplit binary (text -> UTF-8 bytes, matching original)
|
||||
voxsplit_bin = os.path.join(tmpdir, "voxsplit.bin")
|
||||
with open(VOXSPLIT_SOURCE, "r") as src, open(voxsplit_bin, "wb") as dst:
|
||||
dst.write(src.read().encode())
|
||||
|
||||
# 3. Build execution script — extract audio to WAV first, then split
|
||||
extract_wav = (
|
||||
|
|
@ -75,73 +270,31 @@ def process_karaoke(s3_client, bucket, s3_key, s3_path,
|
|||
f"-c:v copy -map 0:v -map 1:a -shortest /tmp/artifacts/{vocals_name}"
|
||||
)
|
||||
else:
|
||||
inst_name = "instrumentals.wav"
|
||||
vocals_name = "vocals.wav"
|
||||
script = (
|
||||
f"{extract_wav} && "
|
||||
"cp /tmp/split-instrumental.wav /tmp/artifacts/instrumentals.wav && "
|
||||
"cp /tmp/split-vocal.wav /tmp/artifacts/vocals.wav"
|
||||
)
|
||||
|
||||
# 4. Call unsandbox execute API via SDK
|
||||
payload = {
|
||||
"language": "bash",
|
||||
"code": script,
|
||||
"network_mode": "zerotrust",
|
||||
"ttl": 300,
|
||||
"vcpu": 2,
|
||||
"input_files": [
|
||||
{"filename": "media", "content": media_b64},
|
||||
{"filename": "voxsplit.c", "content": voxsplit_b64},
|
||||
],
|
||||
}
|
||||
# 4. Build JSON payload on disk (streaming base64)
|
||||
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:
|
||||
result = _make_request("POST", "/execute", public_key, secret_key, payload)
|
||||
resp_path = _make_streaming_request(
|
||||
json_path, public_key, secret_key, tmpdir,
|
||||
)
|
||||
except Exception as e:
|
||||
log.warning("Unsandbox execute failed: %s", e)
|
||||
return None
|
||||
os.unlink(json_path)
|
||||
|
||||
if not result.get("success") or result.get("exit_code") != 0:
|
||||
log.warning("Unsandbox execute non-zero exit: success=%s exit_code=%s stdout=%.2000s stderr=%.2000s",
|
||||
result.get("success"), result.get("exit_code"),
|
||||
result.get("stdout", ""), result.get("stderr", ""))
|
||||
return None
|
||||
|
||||
# 5. Extract artifacts and upload to S3
|
||||
artifacts = result.get("artifacts", [])
|
||||
if len(artifacts) < 2:
|
||||
return None
|
||||
|
||||
# Map artifact filenames to track names
|
||||
artifact_map = {}
|
||||
for a in artifacts:
|
||||
fname = a.get("filename", "")
|
||||
if fname.startswith("instrumentals"):
|
||||
artifact_map["instrumentals"] = a
|
||||
elif fname.startswith("vocals"):
|
||||
artifact_map["vocals"] = a
|
||||
|
||||
if "instrumentals" not in artifact_map or "vocals" not in artifact_map:
|
||||
return None
|
||||
|
||||
sizes = {}
|
||||
content_type = "video/mp4" if is_video else "audio/wav"
|
||||
|
||||
for track_name in ("instrumentals", "vocals"):
|
||||
artifact = artifact_map[track_name]
|
||||
data = base64.b64decode(artifact["content_base64"])
|
||||
|
||||
s3_client.put_object(
|
||||
Bucket=bucket,
|
||||
Key=f"{s3_path}/{track_name}",
|
||||
Body=data,
|
||||
CacheControl="private, max-age=172800",
|
||||
ContentType=content_type,
|
||||
)
|
||||
sizes[track_name] = len(data)
|
||||
|
||||
return sizes
|
||||
# 6. Process response, upload artifacts to S3
|
||||
return _process_response_from_disk(
|
||||
resp_path, s3_client, bucket, s3_path, is_video, extension, tmpdir,
|
||||
)
|
||||
|
||||
|
||||
def tracks_exist(s3_client, bucket, s3_path):
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue