family_home_lab/portal/tasks.py

"""Background workers (Celery + Redis).

Handles async work that shouldn't block the web request: S3 bucket
operations, backup notifications, and (future) media processing.

Run with:
    celery -A portal.tasks.celery_app worker --loglevel=info
"""
from __future__ import annotations

import logging

from celery import Celery

from .config import settings
from .s3 import s3_bucket_name, s3_client

logger = logging.getLogger(__name__)

celery_app = Celery(
    "family_home_lab",
    broker=settings.CELERY_BROKER_URL,
    backend=settings.REDIS_URL,
    include=["portal.tasks"],
)

celery_app.conf.update(
    task_serializer="json",
    result_serializer="json",
    accept_content=["json"],
    timezone="UTC",
    enable_utc=True,
    broker_connection_retry_on_startup=True,
)


@celery_app.task(bind=True, max_retries=3, default_retry_delay=30)
def ensure_user_bucket(self, username: str) -> dict:
    """Create (idempotently) the per-user S3 bucket in Garage."""
    try:
        client = s3_client()
        for bucket in ("shared-media", s3_bucket_name(username)):
            try:
                client.head_bucket(Bucket=bucket)
                logger.info("Bucket %s exists", bucket)
            except Exception:
                client.create_bucket(Bucket=bucket)
                logger.info("Created bucket %s", bucket)
        return {"ok": True, "buckets": ["shared-media", s3_bucket_name(username)]}
    except Exception as exc:  # noqa: BLE001
        logger.exception("ensure_user_bucket failed")
        raise self.retry(exc=exc) from exc


@celery_app.task(bind=True, queue="transcription-mus", max_retries=1, default_retry_delay=120)
def transcribe_sheetmusic(self, bucket: str, key: str, user: str = "") -> dict:
    """Transcribe audio to MIDI + MusicXML + engraved PDF / tabs (Kyutai MuScriptor).

    Runs on the dedicated `transcriber-mus` worker. Requires HUGGINGFACE_TOKEN (the
    CC BY-NC model weights are gated) and is slow on CPU. Uses the muscriptor CLI:
        muscriptor transcribe <audio> --format sheets --output <dir>
    Uploads .mid/.musicxml/.pdf back to the bucket under sheetmusic/<user>/<tag>.
    """
    import os, pathlib, shutil, subprocess, tempfile, uuid

    client = s3_client()
    workdir = tempfile.mkdtemp(prefix="sheets_")
    try:
        base = pathlib.Path(key).stem or "score"
        audio_path = os.path.join(workdir, base + pathlib.Path(key).suffix)
        client.download_file(bucket, key, audio_path)

        if not os.environ.get("HUGGINGFACE_TOKEN"):
            return {"ok": False, "error": "HUGGINGFACE_TOKEN not set (CC BY-NC model is gated)"}

        outdir = os.path.join(workdir, "score")
        subprocess.run(
            ["muscriptor", "transcribe", audio_path, "--format", "sheets", "--output", outdir],
            check=True, capture_output=True,
        )

        tag = uuid.uuid4().hex[:8]
        prefix = f"sheetmusic/{user}/{tag}" if user else f"sheetmusic/{tag}"
        uploaded = []
        for p in pathlib.Path(outdir).rglob("*"):
            if p.is_file() and p.suffix.lower() in {".mid", ".musicxml", ".pdf"}:
                target = f"{prefix}/{p.name}"
                client.upload_file(str(p), bucket, target)
                uploaded.append(target)
        logger.info("Sheet-music transcription complete %s/%s -> %s", bucket, key, uploaded)
        return {"ok": True, "bucket": bucket, "files": uploaded, "tag": tag}
    except Exception as exc:  # noqa: BLE001
        logger.exception("muScriptor transcription failed")
        raise self.retry(exc=exc) from exc
    finally:
        shutil.rmtree(workdir, ignore_errors=True)
    """Hook point for backup/health notifications (currently a no-op stub)."""
    logger.info("Backup notify stub fired for %s", scope)
    return {"ok": True, "scope": scope}


@celery_app.task(bind=True, queue="transcription", max_retries=2, default_retry_delay=60)
def transcribe_audio(self, bucket: str, key: str, user: str = "") -> dict:
    """Transcribe an audio file (in Garage) to MIDI using Spotify Basic Pitch.

    Runs on the dedicated `transcriber` worker (phish: has basic-pitch + tensorflow
    installed). Downloads audio from Garage, runs basic-pitch, uploads the MIDI
    back to the same bucket under `transcriptions/<user>/...`.
    """
    import os
    import pathlib
    import shutil
    import tempfile
    import uuid

    client = s3_client()
    workdir = tempfile.mkdtemp(prefix="transcribe_")
    try:
        base = pathlib.Path(key).stem or "input"
        audio_path = os.path.join(workdir, base + pathlib.Path(key).suffix)
        client.download_file(bucket, key, audio_path)

        outdir = os.path.join(workdir, "out")
        os.makedirs(outdir, exist_ok=True)
        try:
            import basic_pitch
            import pathlib
            from basic_pitch.inference import predict_and_save
        except Exception as exc:  # noqa: BLE001
            logger.exception("basic-pitch missing on this worker")
            return {"ok": False, "error": f"basic-pitch not installed: {exc}"}

        # This basic-pitch ships its default model as TFLite (no SavedModel .pb).
        bpdir = list(basic_pitch.__path__)[0]
        model_path = pathlib.Path(bpdir) / "saved_models" / "icassp_2022" / "nmp.tflite"
        predict_and_save(
            [audio_path],
            outdir,
            save_midi=True,
            sonify_midi=False,
            save_model_outputs=True,
            save_notes=True,
            model_or_model_path=model_path,
        )

        tag = uuid.uuid4().hex[:8]
        prefix = f"transcriptions/{user}/{tag}" if user else f"transcriptions/{tag}"
        uploaded = []
        # basic-pitch may lay files out under different subpaths; walk the output
        # dir and upload every .mid / .csv we find.
        for p in pathlib.Path(outdir).rglob("*.mid"):
            target = f"{prefix}/{p.name}"
            client.upload_file(str(p), bucket, target)
            uploaded.append(target)
        for p in pathlib.Path(outdir).rglob("*.csv"):
            target = f"{prefix}/{p.name}"
            client.upload_file(str(p), bucket, target)
            uploaded.append(target)
        logger.info("Transcription complete for %s/%s -> %s", bucket, key, uploaded)
        return {"ok": True, "bucket": bucket, "files": uploaded, "tag": tag}
    except Exception as exc:  # noqa: BLE001
        logger.exception("transcription failed")
        raise self.retry(exc=exc) from exc
    finally:
        shutil.rmtree(workdir, ignore_errors=True)

@celery_app.task
def notify_backup(scope: str = "family-home-lab") -> dict:
    """Hook point for backup/health notifications (currently a no-op stub)."""
    logger.info("Backup notify stub fired for %s", scope)
    return {"ok": True, "scope": scope}