music-assistant-server

1"""Model/base for an Audio Analysis Provider implementation."""
2
3from __future__ import annotations
4
5import asyncio
6import functools
7import sqlite3
8import time
9from abc import abstractmethod
10from concurrent.futures import ThreadPoolExecutor
11from dataclasses import dataclass
12from typing import TYPE_CHECKING, TypeVar
13
14from music_assistant.helpers.util import join_task
15from music_assistant.models.audio_analysis import AudioAnalysisError
16
17from .provider import Provider
18
19if TYPE_CHECKING:
20    from collections.abc import Callable
21    from datetime import datetime
22    from typing import Any, Literal
23
24    from music_assistant_models.config_entries import ProviderConfig
25    from music_assistant_models.enums import ProviderFeature
26    from music_assistant_models.media_items import AudioFormat
27    from music_assistant_models.provider import ProviderManifest
28    from music_assistant_models.streamdetails import StreamDetails
29
30    from music_assistant.mass import MusicAssistant
31    from music_assistant.models.audio_analysis import AudioAnalysisData
32
33_T = TypeVar("_T")
34
35# DEBUG-log offloads that wait longer than this (seconds) for a permit -- time
36# spent queued behind the concurrency cap rather than computing.
37_PERMIT_WAIT_DEBUG_THRESHOLD = 0.5
38
39# Providers that accumulate whole-track feature state (beat grids, per-chunk VQT, CLAP windows)
40# cap analysis at this duration. Multi-hour mixes/audiobooks would hold hundreds of MB of
41# feature state and tie up the scan for hours, and smart crossfade/similarity carry no meaning
42# at that length. Providers opt in by setting max_analysis_duration.
43ACCUMULATING_ANALYSIS_MAX_DURATION_SECONDS = 1800.0
44
45
46# Subclass rather than wrap so existing isinstance(..., asyncio.Semaphore) checks keep
47# working, and keep the counters here so callers read contention state without touching
48# asyncio internals (_value/_waiters).
49class InstrumentedSemaphore(asyncio.Semaphore):
50    """asyncio.Semaphore that also exposes live in_flight/capacity/waiters counts."""
51
52    def __init__(self, value: int) -> None:
53        """Initialize with ``value`` permits."""
54        super().__init__(value)
55        self.capacity = value
56        self.in_flight = 0
57        self.waiters = 0
58
59    async def acquire(self) -> Literal[True]:
60        """Acquire a permit."""
61        self.waiters += 1
62        try:
63            acquired = await super().acquire()
64        finally:
65            self.waiters -= 1
66        self.in_flight += 1
67        return acquired
68
69    def release(self) -> None:
70        """Release a permit."""
71        self.in_flight -= 1
72        super().release()
73
74
75@dataclass
76class AnalysisSessionData:
77    """Base session data stored per analysis session."""
78
79    streamdetails: StreamDetails
80    audio_format: AudioFormat
81
82
83class AudioAnalysisProvider(Provider):
84    """
85    Base representation of an Audio Analysis Provider.
86
87    Receives PCM audio chunks during streaming and produces analysis results
88    such as beat tracking, key detection, or loudness. The same hooks drive
89    both live playback and background scans; providers do not need to know
90    which context they are running in.
91    """
92
93    # Version of the analysis algorithm. Providers should increment this when
94    # their algorithm changes significantly. The base class compares this against
95    # the stored version to decide whether to re-analyze a track.
96    analysis_version: int = 1
97
98    # Maximum track duration (seconds) this provider will analyze; longer tracks are skipped by
99    # start_analysis. None means no limit. Providers that accumulate whole-track state set it.
100    max_analysis_duration: float | None = None
101
102    # Whether this provider holds heavy ML models that can be unloaded while idle and reloaded
103    # on demand. Providers that load such models set this True (see _load_models/_free_models).
104    has_unloadable_models: bool = False
105
106    def __init__(
107        self,
108        mass: MusicAssistant,
109        manifest: ProviderManifest,
110        config: ProviderConfig,
111        supported_features: set[ProviderFeature] | None = None,
112    ) -> None:
113        """Initialize AudioAnalysisProvider."""
114        super().__init__(mass, manifest, config, supported_features)
115        self._sessions: dict[str, AnalysisSessionData] = {}
116        # Serializes (re)loading of heavy models so concurrent session starts load them once.
117        self._models_lock = asyncio.Lock()
118        self._models_loaded = False
119
120    async def start_analysis(
121        self,
122        session_id: str,
123        streamdetails: StreamDetails,
124        audio_format: AudioFormat,
125    ) -> bool:
126        """
127        Start analysis for a new session.
128
129        Returns True if the provider accepted the session.
130
131        :param session_id: Session ID created by the AudioAnalysisController.
132        :param streamdetails: The stream details for the item being analyzed.
133        :param audio_format: PCM format of the audio stream.
134        """
135        if (
136            self.max_analysis_duration is not None
137            and streamdetails.duration
138            and streamdetails.duration > self.max_analysis_duration
139        ):
140            self.logger.debug(
141                "Skipping analysis for %s: %.0fs exceeds the %.0fs limit",
142                streamdetails.uri,
143                streamdetails.duration,
144                self.max_analysis_duration,
145            )
146            return False
147        stored_version = await self.mass.streams.audio_analysis.get_audio_analysis_version(
148            streamdetails.item_id,
149            streamdetails.provider,
150            self.domain,
151            media_type=streamdetails.media_type,
152        )
153        if stored_version is not None and stored_version >= self.analysis_version:
154            return False
155        if self.has_unloadable_models and not await self.ensure_models_loaded():
156            return False
157        self._sessions[session_id] = AnalysisSessionData(
158            streamdetails=streamdetails,
159            audio_format=audio_format,
160        )
161        session = self._sessions[session_id]
162        try:
163            accepted = await self._start_analysis(session_id, streamdetails, audio_format)
164        except AudioAnalysisError as err:
165            await self._record_failure(session, err.reason, err.retry_at)
166            self._sessions.pop(session_id, None)
167            return False
168        except asyncio.CancelledError:
169            # Cancellation is not an analysis failure: let it propagate without recording.
170            raise
171        except Exception as err:
172            # _start_analysis is provider-implemented (ffmpeg/torch/numpy decode); its
173            # failure surface is open-ended, so the catch stays broad. Any unexpected
174            # error is logged with a traceback and recorded as a failure.
175            self.logger.error(
176                "_start_analysis raised for session %s: %s", session_id, err, exc_info=err
177            )
178            await self._record_failure(session, str(err), None)
179            self._sessions.pop(session_id, None)
180            return False
181        if not accepted:
182            self._sessions.pop(session_id, None)
183            return False
184        return True
185
186    @abstractmethod
187    async def process_pcm_chunk(
188        self,
189        session_id: str,
190        pcm_chunk: bytes,
191    ) -> None:
192        """
193        Process a PCM audio chunk.
194
195        Implementations MUST `await` all chunk-processing work; the controller
196        relies on this to backpressure the audio source.
197
198        :param session_id: The analysis session ID.
199        :param pcm_chunk: Raw PCM audio data.
200        """
201
202    async def finalize(self, session_id: str) -> None:
203        """Finalize analysis, persist the result, fire post_analysis, then clean up."""
204        analysis: AudioAnalysisData | None = None
205        session = self._sessions.get(session_id)
206        try:
207            analysis = await self._finalize(session_id)
208        except AudioAnalysisError as err:
209            if session is not None:
210                await self._record_failure(session, err.reason, err.retry_at)
211        except asyncio.CancelledError:
212            # Cancellation is not an analysis failure: let it propagate without recording.
213            raise
214        except Exception as err:
215            # _finalize is provider-implemented (torch/ffmpeg inference); its failure
216            # surface is open-ended, so the catch stays broad — logged and recorded.
217            self.logger.error("_finalize raised for session %s: %s", session_id, err, exc_info=err)
218            if session is not None:
219                await self._record_failure(session, str(err), None)
220        if analysis is not None and session is not None:
221            try:
222                await self.mass.streams.audio_analysis.set_audio_analysis(
223                    item_id=session.streamdetails.item_id,
224                    provider_instance_id_or_domain=session.streamdetails.provider,
225                    aa_provider_domain=self.domain,
226                    analysis=analysis,
227                    analysis_version=self.analysis_version,
228                    media_type=session.streamdetails.media_type,
229                )
230            except Exception as err:
231                # Persisting (DB write + provider lookup) must never break session
232                # cleanup below, so the catch stays broad — logged and skipped.
233                self.logger.warning(
234                    "set_audio_analysis raised for %s: %s", self.domain, err, exc_info=err
235                )
236            else:
237                try:
238                    await self.post_analysis(session.streamdetails, analysis)
239                except Exception as err:
240                    # post_analysis is a provider-implemented hook with an open-ended
241                    # failure surface; a failing side effect must not break cleanup.
242                    self.logger.warning(
243                        "post_analysis raised for %s: %s", self.domain, err, exc_info=err
244                    )
245        self._sessions.pop(session_id, None)
246
247    async def post_analysis(
248        self,
249        streamdetails: StreamDetails,
250        analysis: AudioAnalysisData,
251    ) -> None:
252        """
253        Run side effects after analysis is finalized and persisted.
254
255        Default is a no-op. Implementations MUST self-gate on whether
256        `streamdetails.path` is a writable filesystem path, since this hook
257        fires for both live and background-scan analyses.
258
259        :param streamdetails: The stream details for the analyzed item.
260        :param analysis: The analysis data that was persisted by `_finalize`.
261        """
262        return
263
264    async def cancel(self, session_id: str) -> None:
265        """Cancel an in-progress analysis session."""
266        self._sessions.pop(session_id, None)
267
268    async def unload(self, is_removed: bool = False) -> None:
269        """Handle unload, cancelling any active analysis sessions and freeing models."""
270        for session_id in list(self._sessions):
271            await self.cancel(session_id)
272        async with self._models_lock:
273            self._free_models()
274            self._models_loaded = False
275        await super().unload(is_removed)
276
277    async def ensure_models_loaded(self) -> bool:
278        """
279        Load this provider's heavy models if they are not resident, returning success.
280
281        Safe to call from concurrent sessions: the models are loaded once. Returns False
282        when loading fails, so the caller can decline the session.
283        """
284        async with self._models_lock:
285            if self._models_loaded:
286                return True
287            try:
288                await self._load_models()
289            except Exception as err:
290                self.logger.error("Failed to load analysis models: %s", err, exc_info=err)
291                return False
292            self._models_loaded = True
293            return True
294
295    async def unload_idle_models(self) -> None:
296        """Free heavy models to reclaim memory while idle; the next analysis reloads them."""
297        async with self._models_lock:
298            if not self._models_loaded:
299                return
300            self._free_models()
301            self._models_loaded = False
302            self.logger.debug("Unloaded idle analysis models")
303
304    @abstractmethod
305    async def _start_analysis(
306        self,
307        session_id: str,
308        streamdetails: StreamDetails,
309        audio_format: AudioFormat,
310    ) -> bool:
311        """
312        Provider-specific initialization for a new analysis session.
313
314        Return False to reject the session (e.g. unsupported format).
315
316        :param session_id: The analysis session ID.
317        :param streamdetails: The stream details for the item being analyzed.
318        :param audio_format: PCM format of the audio stream.
319        """
320
321    @abstractmethod
322    async def _finalize(self, session_id: str) -> AudioAnalysisData | None:
323        """
324        Compute and return the analysis for this session (or None to skip).
325
326        The base class persists the returned value via set_audio_analysis() and
327        then fires post_analysis(). Return None to skip both.
328
329        :param session_id: The analysis session ID.
330        """
331
332    async def _load_models(self) -> None:
333        """Load heavy models into memory. Override when has_unloadable_models is True."""
334        return
335
336    def _free_models(self) -> None:
337        """Drop references to heavy models so memory can be reclaimed. Override as needed."""
338        return
339
340    async def _run_offloaded(self, func: Callable[..., _T], /, *args: Any, **kwargs: Any) -> _T:
341        """
342        Run a blocking analysis function on the niced analysis thread pool.
343
344        Route all CPU-heavy analysis work through this. The controller's semaphore caps
345        concurrency, and while any player is streaming the solo lock holds analysis to one
346        offload at a time so it stays below playback. Falls back to asyncio.to_thread when the
347        controller has no pool configured.
348
349        The slot (semaphore permit and solo lock) is held until the worker thread finishes,
350        even if the awaiting coroutine is cancelled: the thread keeps running, so it must keep
351        counting against the caps.
352
353        :param func: The blocking callable to run off the event loop.
354        :param args: Positional arguments passed to func.
355        :param kwargs: Keyword arguments passed to func.
356        """
357        controller = self.mass.streams.audio_analysis
358        semaphore = controller.analysis_semaphore
359        if not isinstance(semaphore, asyncio.Semaphore):
360            return await asyncio.to_thread(func, *args, **kwargs)
361
362        # While a player streams, take the solo lock so analysis runs one offload at a time:
363        # the GIL serializes Python execution, so concurrent inference starves the playback
364        # loop. Idle, the semaphore alone caps concurrency and background work uses spare cores.
365        solo = getattr(controller, "analysis_solo_lock", None)
366        solo_lock = solo if isinstance(solo, asyncio.Lock) else None
367        take_solo = False
368        if solo_lock is not None:
369            try:
370                take_solo = bool(controller.playback_active())
371            except Exception:
372                take_solo = False
373
374        solo_held = False
375
376        def _release(done: asyncio.Future[_T]) -> None:
377            if solo_held and solo_lock is not None:
378                solo_lock.release()
379            semaphore.release()
380            # Retrieve any exception so a cancelled awaiter doesn't leave it unretrieved.
381            if not done.cancelled():
382                done.exception()
383
384        wait_start = time.monotonic()
385        await semaphore.acquire()
386        try:
387            if take_solo and solo_lock is not None:
388                await solo_lock.acquire()
389                solo_held = True
390        except BaseException:
391            # Cancelled (or failed) waiting for the solo lock — give the permit back.
392            semaphore.release()
393            raise
394        wait_seconds = time.monotonic() - wait_start
395        if wait_seconds >= _PERMIT_WAIT_DEBUG_THRESHOLD and isinstance(
396            semaphore, InstrumentedSemaphore
397        ):
398            self.logger.debug(
399                "Analysis offload waited %.1fs for a slot (%d/%d permits in use, %d queued)",
400                wait_seconds,
401                semaphore.in_flight,
402                semaphore.capacity,
403                semaphore.waiters,
404            )
405        executor = getattr(controller, "analysis_executor", None)
406        try:
407            if isinstance(executor, ThreadPoolExecutor):
408                loop = asyncio.get_running_loop()
409                future: asyncio.Future[_T] = asyncio.ensure_future(
410                    loop.run_in_executor(executor, functools.partial(func, *args, **kwargs))
411                )
412            else:
413                future = asyncio.ensure_future(asyncio.to_thread(func, *args, **kwargs))
414        except RuntimeError:
415            # Scheduling the worker failed (e.g. the loop is shutting down) — release the
416            # permit we just took so it isn't leaked, which would shrink the cap over time.
417            if solo_held and solo_lock is not None:
418                solo_lock.release()
419            semaphore.release()
420            raise
421        future.add_done_callback(_release)
422        # join: a cancelled awaiter leaves the thread and its slot held until the work finishes.
423        return await join_task(future)
424
425    async def _run_offloaded_timed(
426        self, func: Callable[..., _T], /, *args: Any, **kwargs: Any
427    ) -> tuple[_T, float]:
428        """
429        Run a blocking analysis function offloaded, returning its result and execution seconds.
430
431        Same slot semantics as ``_run_offloaded``; the returned seconds cover only the
432        callable's own execution, never the time spent queued for a slot.
433
434        :param func: The blocking callable to run off the event loop.
435        :param args: Positional arguments passed to func.
436        :param kwargs: Keyword arguments passed to func.
437        """
438
439        def _timed() -> tuple[_T, float]:
440            start = time.perf_counter()
441            return func(*args, **kwargs), time.perf_counter() - start
442
443        return await self._run_offloaded(_timed)
444
445    async def _record_failure(
446        self,
447        session: AnalysisSessionData,
448        reason: str,
449        retry_at: datetime | None,
450    ) -> None:
451        """Record a failure for this session's track."""
452        # Swallow DB write errors so a failed recorder write never breaks the session
453        # lifecycle (record_analysis_failure only performs a sqlite insert).
454        try:
455            sd = session.streamdetails
456            await self.mass.streams.audio_analysis.record_analysis_failure(
457                item_id=sd.item_id,
458                provider_instance_id_or_domain=sd.provider,
459                aa_provider_domain=self.domain,
460                reason=reason,
461                retry_at=retry_at,
462                analysis_version=self.analysis_version,
463                media_type=sd.media_type,
464            )
465        except sqlite3.Error as err:
466            self.logger.warning(
467                "record_analysis_failure raised for %s: %s", self.domain, err, exc_info=err
468            )
469