音源分離をAWSで本番スケールさせる：GPUバッチ処理基盤（SQS × ECS/Batch × S3）

            ┌─────────────┐   ObjectCreated   ┌──────────────┐
  upload →  │ S3 入力bucket │ ───イベント通知──→ │   SQS キュー   │ ──┐
            └─────────────┘                   └──────────────┘   │ long poll
                                                     │ 失敗N回      │
                                                     ▼            ▼
                                              ┌──────────┐   ┌──────────────────┐
                                              │   DLQ    │   │  GPUワーカー群     │
                                              └──────────┘   │ (AWS Batch / ECS) │
                                                             │  分離 + 冪等チェック │
                                                             └──────────────────┘
                                                                     │ stems
                                                                     ▼
                                                            ┌────────────────┐
                                                            │  S3 出力bucket   │
                                                            └────────────────┘

# 冪等キー：音声の内容ではなく"入力S3キー + モデル"で出力先を決定論的に決める
def output_prefix(input_key: str, model: str) -> str:
    """同じ入力×同じモデルなら必ず同じ出力prefix。存在すれば再処理しない。"""
    safe_model = model.replace("/", "_")
    return f"stems/{safe_model}/{input_key}"

# heartbeat.py — 処理中だけ可視性を延長し続けるコンテキストマネージャ
import threading
from contextlib import contextmanager

@contextmanager
def visibility_heartbeat(sqs, queue_url: str, receipt_handle: str,
                         *, extend_to: int = 300, interval: int = 120):
    """interval秒ごとに可視性をextend_to秒へ延長。長時間ジョブの二重配信を防ぐ。"""
    stop = threading.Event()

    def beat() -> None:
        while not stop.wait(interval):           # interval待つ。stopが立てば即終了
            try:
                sqs.change_message_visibility(
                    QueueUrl=queue_url,
                    ReceiptHandle=receipt_handle,
                    VisibilityTimeout=extend_to,  # 最初の受信から最大12hの制約に注意
                )
            except Exception:                     # 延長失敗は致命的でない。次の周期で再試行
                pass

    t = threading.Thread(target=beat, daemon=True)
    t.start()
    try:
        yield
    finally:
        stop.set()                                # 処理完了で確実に停止（リーク防止）
        t.join(timeout=1)

{
  "resourceRequirements": [
    { "type": "GPU",    "value": "1" },
    { "type": "VCPU",   "value": "4" },
    { "type": "MEMORY", "value": "16384" }
  ]
}

# worker.py — 本番GPU音源分離ワーカー（SQS駆動・冪等・回復性）
from __future__ import annotations

import json
import logging
import os
import signal
import sys
import tempfile
from pathlib import Path
from urllib.parse import unquote_plus

import boto3
from audio_separator.separator import Separator

from heartbeat import visibility_heartbeat   # 前掲

# --- 構造化ログ（相関IDを必ず通す。音声内容＝PIIは出さない）---
logging.basicConfig(level=logging.INFO, format='%(message)s')
log = logging.getLogger("worker")

def emit(event: str, **fields) -> None:
    log.info(json.dumps({"event": event, **fields}, ensure_ascii=False))

QUEUE_URL = os.environ["QUEUE_URL"]
OUT_BUCKET = os.environ["OUTPUT_BUCKET"]
MODEL = os.environ.get("SEPARATION_MODEL", "UVR-MDX-NET-Inst_HQ_3.onnx")
MODEL_CACHE = os.environ.get("MODEL_CACHE_DIR", "/models")   # イメージに焼き込み済み推奨

sqs = boto3.client("sqs")
s3 = boto3.client("s3")

_shutdown = False
def _on_sigterm(*_):                 # Spot中断(2分前通知)/ECS停止で届く
    global _shutdown
    _shutdown = True
    emit("shutdown_requested")       # 新規受信を止め、処理中ジョブを安全に終える
signal.signal(signal.SIGTERM, _on_sigterm)

# モデルは一度だけVRAMへ。リクエスト間で再利用（ロードは高コスト）
_separator: Separator | None = None
def get_separator() -> Separator:
    global _separator
    if _separator is None:
        _separator = Separator(output_format="flac", model_file_dir=MODEL_CACHE)
        _separator.load_model(model_filename=MODEL)
        emit("model_loaded", model=MODEL)
    return _separator


def _already_done(prefix: str) -> bool:
    """出力S3キーが存在するなら処理済み＝冪等にスキップ（二重GPU実行を防ぐ）。"""
    resp = s3.list_objects_v2(Bucket=OUT_BUCKET, Prefix=prefix, MaxKeys=1)
    return resp.get("KeyCount", 0) > 0


def process_message(msg: dict) -> None:
    body = json.loads(msg["Body"])
    rec = body["Records"][0]["s3"]                 # S3イベント通知の形
    in_bucket = rec["bucket"]["name"]
    in_key = unquote_plus(rec["object"]["key"])    # S3通知のkeyはURLエンコード済（空白→'+'・特殊文字→%XX）
    prefix = f"stems/{MODEL.replace('/', '_')}/{in_key}"   # 決定論的な出力先
    cid = msg["MessageId"]

    if _already_done(prefix):
        emit("skip_idempotent", cid=cid, key=in_key)
        return                                     # 既に成果物あり＝成功扱い

    with visibility_heartbeat(sqs, QUEUE_URL, msg["ReceiptHandle"]):
        with tempfile.TemporaryDirectory() as tmp:
            src = Path(tmp) / Path(in_key).name
            s3.download_file(in_bucket, in_key, str(src))
            sep = get_separator()
            sep.output_dir = tmp
            stems = sep.separate(str(src))         # GPU処理（ハートビートが守る）
            for stem in stems:
                key = f"{prefix}/{Path(stem).name}"
                s3.upload_file(stem, OUT_BUCKET, key)
            emit("separated", cid=cid, key=in_key, stems=len(stems))


def main() -> None:
    get_separator()                                # 起動時ウォームアップ
    while not _shutdown:
        resp = sqs.receive_message(
            QueueUrl=QUEUE_URL, MaxNumberOfMessages=1,
            WaitTimeSeconds=20,                    # ロングポーリング（空ポーリング課金を削減）
        )
        for msg in resp.get("Messages", []):
            try:
                process_message(msg)
                sqs.delete_message(QueueUrl=QUEUE_URL,
                                   ReceiptHandle=msg["ReceiptHandle"])  # 成功時のみ削除
            except Exception:
                # 失敗は削除しない → 可視性切れで再配信 → maxReceiveCount超でDLQへ
                emit("process_failed", cid=msg.get("MessageId"))
                log.exception("processing error")
    emit("drained_and_exit")
    sys.exit(0)


if __name__ == "__main__":
    main()

# main.tf — SQS + DLQ + S3イベント通知
resource "aws_sqs_queue" "dlq" {
  name                      = "separation-dlq"
  message_retention_seconds = 1209600          # 14日。原因調査の猶予を確保
}

resource "aws_sqs_queue" "jobs" {
  name                       = "separation-jobs"
  visibility_timeout_seconds = 300             # ハートビート前提の初期値
  receive_wait_time_seconds  = 20              # ロングポーリング既定

  redrive_policy = jsonencode({
    deadLetterTargetArn = aws_sqs_queue.dlq.arn
    maxReceiveCount     = 5                     # 5回失敗で毒メッセージをDLQへ
  })
}

# S3 → SQS（ObjectCreatedのみ。入力prefixに限定して再帰を防ぐ）
resource "aws_s3_bucket_notification" "ingest" {
  bucket = aws_s3_bucket.input.id
  queue {
    queue_arn     = aws_sqs_queue.jobs.arn
    events        = ["s3:ObjectCreated:*"]
    filter_prefix = "uploads/"
  }
}

# test_idempotency.py — GPU不要。冪等ロジックだけを検証
from worker import output_prefix  # 純関数として切り出しておく

def test_output_prefix_is_deterministic():
    a = output_prefix("uploads/song.wav", "UVR-MDX-NET-Inst_HQ_3.onnx")
    b = output_prefix("uploads/song.wav", "UVR-MDX-NET-Inst_HQ_3.onnx")
    assert a == b                       # 同じ入力×モデルなら必ず同じ出力先＝冪等の土台

def test_model_namespaced():
    p1 = output_prefix("uploads/x.wav", "model_a")
    p2 = output_prefix("uploads/x.wav", "model_b")
    assert p1 != p2                     # モデルが違えば出力先も分かれる