diff --git a/make_post_sell/lib/karaoke.py b/make_post_sell/lib/karaoke.py index b30945d..0580e06 100644 --- a/make_post_sell/lib/karaoke.py +++ b/make_post_sell/lib/karaoke.py @@ -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:" + ) + """ + 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):