musicmouse/python-backend/musicmouse/library/workers.py

"""Where analysis actually burns CPU, and how it is kept from taking the machine over.

Analyzing one track is a few seconds of single-threaded DSP that the GIL will not let
another Python thread overlap with - librosa's work is numba-jitted and numpy glue, not
long C calls that release it. So a library-sized pass parallelizes across *processes*:
:func:`analysis_pool` hands :meth:`~musicmouse.library.MusicLibrary.analyze_pending` an
executor whose workers are separate interpreters, and a 12-core desktop chews through a
first-time scan roughly an order of magnitude faster than the Raspberry Pi this also
has to stay polite on.

Polite means three things, all of them set here rather than at the call site:

* **One core stays free** (:func:`default_worker_count`), so the audio thread and the
  web server never have to fight a full house of analyzers for a timeslice.
* **Workers run niced**, so even the cores they do own yield to playback instantly.
* **Each worker stays single-threaded** - numpy's BLAS and numba would each happily
  start one thread per core *inside* every worker, and N x N threads on a 4-core Pi is
  slower than N, not faster.

Whatever an analyzer returns has to survive the trip back from a worker process, which
is what keeps :class:`~musicmouse.library.analysis.TrackAnalysis` and friends plain
frozen dataclasses of floats. An analyzer that records state in its own instance -
a test double counting calls, say - only sees that state in the worker, so such an
analyzer must be run with ``workers=1``, where everything stays in this process on a
thread.
"""

from __future__ import annotations

import contextlib
import logging
import multiprocessing
import os
from collections.abc import Iterator
from concurrent.futures import Executor, ProcessPoolExecutor
from pathlib import Path
from typing import Final

from musicmouse.library.analysis import Analyzer, BeatGrid, TrackAnalysis, TrackCurves

_log = logging.getLogger(__name__)

__all__ = ["analysis_pool", "analyze_one", "default_worker_count"]

#: Ceiling on the automatic worker count. Every worker is a fresh interpreter with its
#: own librosa, numpy and a decoded track in memory - a few hundred MB each - so on a
#: big machine the limit that bites first is RAM, not cores. An explicit
#: ``analysis_workers`` in the config overrides this; the default stays conservative.
_MAX_AUTO_WORKERS: Final = 8

#: How much worse than everything else analysis schedules. Niceness clamps at the OS
#: maximum (19) and only ever increases, so re-applying it to a reused worker is
#: harmless.
_NICENESS: Final = 5

#: Forced into every worker *before* it imports numpy, which reads these once at import
#: time. Without them each worker opens its own BLAS thread pool sized for the whole
#: machine and the pool oversubscribes every core several times over.
_SINGLE_THREADED: Final = {
    "OMP_NUM_THREADS": "1",
    "OPENBLAS_NUM_THREADS": "1",
    "MKL_NUM_THREADS": "1",
    "NUMEXPR_NUM_THREADS": "1",
    "NUMBA_NUM_THREADS": "1",
}


def default_worker_count() -> int:
    """One worker per core bar one, capped at :data:`_MAX_AUTO_WORKERS`.

    The core left over is for the rest of the app: audio decoding, the web server and
    the serial link all have to stay responsive while a first-time pass runs for hours.
    Single-core machines get 1, which :func:`analysis_pool` turns into the in-process
    path rather than a pool of one.
    """
    return max(1, min(_MAX_AUTO_WORKERS, (os.cpu_count() or 1) - 1))


def analyze_one(
    analyzer: Analyzer, path: Path
) -> tuple[TrackAnalysis, BeatGrid | None, TrackCurves | None]:
    """Analyze one file off the event loop - in a worker process, or on a thread here.

    Lowers the caller's own scheduling priority first. On Linux ``os.nice`` affects only
    the calling *thread*, so on the in-process path this makes idle-time analysis yield
    CPU without touching threads doing other work; in a worker process there is nothing
    else in the process to slow down anyway.
    """
    with contextlib.suppress(OSError):
        os.nice(_NICENESS)
    return analyzer.analyze(path)


def _init_worker() -> None:
    """Runs once per worker process, before it imports librosa or numpy.

    A spawned worker starts from a bare interpreter and pulls the analyzer in when it
    unpickles its first task, so setting the thread-count variables here still lands
    ahead of numpy reading them.
    """
    os.environ.update(_SINGLE_THREADED)
    with contextlib.suppress(OSError):
        os.nice(_NICENESS)


@contextlib.contextmanager
def analysis_pool(workers: int) -> Iterator[Executor | None]:
    """A pool of `workers` analyzer processes, or ``None`` for "stay in this process".

    ``None`` - for ``workers <= 1``, and as the fallback when a pool cannot be started
    at all - means the caller should run each track on a thread instead. That path is
    what the tests and single-core devices use, and it is the only one where an
    analyzer holding state in its own instance behaves as written.

    Workers are *spawned*, never forked: this process has an asyncio loop, a serial
    reader and libVLC's own threads running, and forking that is a well-known way to
    inherit a held lock and deadlock in a child. The price is one librosa import per
    worker, a few seconds paid once per pass - nothing next to the hours of DSP a
    first-time pass over a real library costs.
    """
    if workers <= 1:
        yield None
        return
    try:
        executor = ProcessPoolExecutor(
            max_workers=workers,
            mp_context=multiprocessing.get_context("spawn"),
            initializer=_init_worker,
        )
    except (OSError, ValueError):
        _log.warning(
            "Could not start %d analysis workers; analyzing in this process instead",
            workers,
            exc_info=True,
        )
        yield None
        return
    try:
        yield executor
    finally:
        # `wait=False`: shutdown happens on cancellation too (the app is stopping), and
        # waiting there would hold it up for however long the tracks in flight take.
        executor.shutdown(wait=False, cancel_futures=True)