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