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.
|
||||
|
||||
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.
|
||||
Downloads media from S3, uploads it + voxsplit.c via the /upload endpoint
|
||||
(streaming, constant memory), then kicks off /execute referencing upload_ids.
|
||||
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.
|
||||
Upload-based pipeline: files stream directly from disk to the API in 64KB
|
||||
chunks — never base64-encoded, never held in memory on either side.
|
||||
"""
|
||||
|
||||
import base64
|
||||
|
|
@ -25,83 +25,30 @@ 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 _sign_body(secret_key, timestamp, method, path, body_bytes):
|
||||
"""HMAC-SHA256 matching un._sign_request for small payloads."""
|
||||
h = hmac.new(secret_key.encode(), digestmod=hashlib.sha256)
|
||||
h.update(f"{timestamp}:{method}:{path}:{body_bytes}".encode())
|
||||
return h.hexdigest()
|
||||
|
||||
|
||||
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):
|
||||
def _sign_from_disk(secret_key, timestamp, method, path, file_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:
|
||||
with open(file_path, "rb") as f:
|
||||
while True:
|
||||
chunk = f.read(_READ_CHUNK)
|
||||
if not chunk:
|
||||
|
|
@ -110,16 +57,61 @@ def _sign_from_disk(secret_key, timestamp, method, path, json_path):
|
|||
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.
|
||||
def _upload_file(file_path, filename, public_key, secret_key):
|
||||
"""Upload a file via POST /upload (streaming, constant memory).
|
||||
|
||||
Returns path to response.json.
|
||||
Returns the upload_id string.
|
||||
"""
|
||||
timestamp = int(time.time())
|
||||
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"
|
||||
|
||||
signature = _sign_from_disk(secret_key, timestamp, method, path, json_path)
|
||||
signature = _sign_body(secret_key, timestamp, method, path, body_str)
|
||||
|
||||
headers = {
|
||||
"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}"
|
||||
|
||||
with open(json_path, "rb") as body:
|
||||
response = requests.post(
|
||||
url, data=body, headers=headers, stream=True, timeout=300,
|
||||
)
|
||||
response = requests.post(
|
||||
url, data=body_str.encode(), headers=headers, stream=True, timeout=300,
|
||||
)
|
||||
|
||||
response.raise_for_status()
|
||||
|
||||
|
|
@ -213,8 +203,9 @@ def process_karaoke(s3_client, bucket, s3_key, s3_path,
|
|||
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.
|
||||
Upload-based pipeline: files stream to the API via POST /upload (64KB
|
||||
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:
|
||||
{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)
|
||||
|
||||
with tempfile.TemporaryDirectory(prefix="karaoke_") as tmpdir:
|
||||
# 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)
|
||||
# 1. Stream S3 download to disk
|
||||
media_path = os.path.join(tmpdir, "media.bin")
|
||||
resp = s3_client.get_object(Bucket=bucket, Key=s3_key)
|
||||
with open(media_path, "wb") as f:
|
||||
|
|
@ -247,12 +236,27 @@ def process_karaoke(s3_client, bucket, s3_key, s3_path,
|
|||
break
|
||||
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")
|
||||
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
|
||||
# 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 = (
|
||||
"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 && "
|
||||
|
|
@ -276,20 +280,15 @@ def process_karaoke(s3_client, bucket, s3_key, s3_path,
|
|||
"cp /tmp/split-vocal.wav /tmp/artifacts/vocals.wav"
|
||||
)
|
||||
|
||||
# 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
|
||||
# 5. Execute with upload_id references (tiny JSON body, no file content)
|
||||
try:
|
||||
resp_path = _make_streaming_request(
|
||||
json_path, public_key, secret_key, tmpdir,
|
||||
resp_path = _execute_with_uploads(
|
||||
[("media", media_upload_id), ("voxsplit.c", voxsplit_upload_id)],
|
||||
script, public_key, secret_key, tmpdir,
|
||||
)
|
||||
except Exception as e:
|
||||
log.warning("Unsandbox execute failed: %s", e)
|
||||
return None
|
||||
os.unlink(json_path)
|
||||
|
||||
# 6. Process response, upload artifacts to S3
|
||||
return _process_response_from_disk(
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue