/
/
/
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 # Tracked so unload() can cancel them before freeing the models they infer against.
120 self._finalize_tasks: set[asyncio.Task[None]] = set()
121
122 async def start_analysis(
123 self,
124 session_id: str,
125 streamdetails: StreamDetails,
126 audio_format: AudioFormat,
127 ) -> bool:
128 """
129 Start analysis for a new session.
130
131 Returns True if the provider accepted the session.
132
133 :param session_id: Session ID created by the AudioAnalysisController.
134 :param streamdetails: The stream details for the item being analyzed.
135 :param audio_format: PCM format of the audio stream.
136 """
137 if self.unloading:
138 # The controller snapshots providers up front and awaits between them, so a
139 # session can still arrive for a provider that is already on its way out.
140 return False
141 if (
142 self.max_analysis_duration is not None
143 and streamdetails.duration
144 and streamdetails.duration > self.max_analysis_duration
145 ):
146 self.logger.debug(
147 "Skipping analysis for %s: %.0fs exceeds the %.0fs limit",
148 streamdetails.uri,
149 streamdetails.duration,
150 self.max_analysis_duration,
151 )
152 return False
153 stored_version = await self.mass.streams.audio_analysis.get_audio_analysis_version(
154 streamdetails.item_id,
155 streamdetails.provider,
156 self.domain,
157 media_type=streamdetails.media_type,
158 )
159 if stored_version is not None and stored_version >= self.analysis_version:
160 return False
161 if self.has_unloadable_models and not await self.ensure_models_loaded():
162 return False
163 self._sessions[session_id] = AnalysisSessionData(
164 streamdetails=streamdetails,
165 audio_format=audio_format,
166 )
167 session = self._sessions[session_id]
168 try:
169 accepted = await self._start_analysis(session_id, streamdetails, audio_format)
170 except AudioAnalysisError as err:
171 await self._record_failure(session, err.reason, err.retry_at)
172 self._sessions.pop(session_id, None)
173 return False
174 except asyncio.CancelledError:
175 # Cancellation is not an analysis failure: let it propagate without recording.
176 raise
177 except Exception as err:
178 # _start_analysis is provider-implemented (ffmpeg/torch/numpy decode); its
179 # failure surface is open-ended, so the catch stays broad. Any unexpected
180 # error is logged with a traceback and recorded as a failure.
181 self.logger.error(
182 "_start_analysis raised for session %s: %s", session_id, err, exc_info=err
183 )
184 await self._record_failure(session, str(err), None)
185 self._sessions.pop(session_id, None)
186 return False
187 if not accepted:
188 self._sessions.pop(session_id, None)
189 return False
190 return True
191
192 @abstractmethod
193 async def process_pcm_chunk(
194 self,
195 session_id: str,
196 pcm_chunk: bytes,
197 ) -> None:
198 """
199 Process a PCM audio chunk.
200
201 Implementations MUST `await` all chunk-processing work; the controller
202 relies on this to backpressure the audio source.
203
204 :param session_id: The analysis session ID.
205 :param pcm_chunk: Raw PCM audio data.
206 """
207
208 async def finalize(self, session_id: str) -> None:
209 """Finalize analysis, persist the result, fire post_analysis, then clean up."""
210 if self.unloading:
211 # Too late to be cancelled by unload(); the models are about to be freed.
212 self._sessions.pop(session_id, None)
213 return
214 task = asyncio.current_task()
215 if task is not None:
216 self._finalize_tasks.add(task)
217 try:
218 analysis: AudioAnalysisData | None = None
219 session = self._sessions.get(session_id)
220 try:
221 analysis = await self._finalize(session_id)
222 except AudioAnalysisError as err:
223 if session is not None:
224 await self._record_failure(session, err.reason, err.retry_at)
225 except asyncio.CancelledError:
226 # Cancellation is not an analysis failure: let it propagate without recording.
227 raise
228 except Exception as err:
229 # _finalize is provider-implemented (torch/ffmpeg inference); its failure
230 # surface is open-ended, so the catch stays broad â logged and recorded.
231 self.logger.error(
232 "_finalize raised for session %s: %s", session_id, err, exc_info=err
233 )
234 if session is not None:
235 await self._record_failure(session, str(err), None)
236 if analysis is not None and session is not None:
237 try:
238 await self.mass.streams.audio_analysis.set_audio_analysis(
239 item_id=session.streamdetails.item_id,
240 provider_instance_id_or_domain=session.streamdetails.provider,
241 aa_provider_domain=self.domain,
242 analysis=analysis,
243 analysis_version=self.analysis_version,
244 media_type=session.streamdetails.media_type,
245 )
246 except Exception as err:
247 # Persisting (DB write + provider lookup) must never break session
248 # cleanup below, so the catch stays broad â logged and skipped.
249 self.logger.warning(
250 "set_audio_analysis raised for %s: %s", self.domain, err, exc_info=err
251 )
252 else:
253 try:
254 await self.post_analysis(session.streamdetails, analysis)
255 except Exception as err:
256 # post_analysis is a provider-implemented hook with an open-ended
257 # failure surface; a failing side effect must not break cleanup.
258 self.logger.warning(
259 "post_analysis raised for %s: %s", self.domain, err, exc_info=err
260 )
261 self._sessions.pop(session_id, None)
262 finally:
263 if task is not None:
264 self._finalize_tasks.discard(task)
265
266 async def post_analysis(
267 self,
268 streamdetails: StreamDetails,
269 analysis: AudioAnalysisData,
270 ) -> None:
271 """
272 Run side effects after analysis is finalized and persisted.
273
274 Default is a no-op. Implementations MUST self-gate on whether
275 `streamdetails.path` is a writable filesystem path, since this hook
276 fires for both live and background-scan analyses.
277
278 :param streamdetails: The stream details for the analyzed item.
279 :param analysis: The analysis data that was persisted by `_finalize`.
280 """
281 return
282
283 async def cancel(self, session_id: str) -> None:
284 """Cancel an in-progress analysis session."""
285 self._sessions.pop(session_id, None)
286
287 async def unload(self, is_removed: bool = False) -> None:
288 """Handle unload, cancelling in-flight analysis work and freeing models."""
289 for session_id in list(self._sessions):
290 await self.cancel(session_id)
291 # Cancelled, not awaited: inference runs for minutes and records no failure, so the
292 # track is simply analyzed again on its next play.
293 current_task = asyncio.current_task()
294 # Reaching unload() from inside a finalize must not cancel that finalize itself.
295 finalize_tasks = [task for task in self._finalize_tasks if task is not current_task]
296 for task in finalize_tasks:
297 task.cancel()
298 await asyncio.gather(*finalize_tasks, return_exceptions=True)
299 async with self._models_lock:
300 self._free_models()
301 self._models_loaded = False
302 await super().unload(is_removed)
303
304 async def ensure_models_loaded(self) -> bool:
305 """
306 Load this provider's heavy models if they are not resident, returning success.
307
308 Safe to call from concurrent sessions: the models are loaded once. Returns False
309 when loading fails, or when the provider is unloading, so the caller can decline
310 the session.
311 """
312 async with self._models_lock:
313 # Checked under the lock unload() frees within: the lock alone only serializes
314 # the two, so a loader that wins it afterwards would strand the models.
315 if self.unloading:
316 return False
317 if self._models_loaded:
318 return True
319 try:
320 await self._load_models()
321 except Exception as err:
322 self.logger.error("Failed to load analysis models: %s", err, exc_info=err)
323 return False
324 self._models_loaded = True
325 return True
326
327 async def unload_idle_models(self) -> None:
328 """Free heavy models to reclaim memory while idle; the next analysis reloads them."""
329 async with self._models_lock:
330 if not self._models_loaded:
331 return
332 self._free_models()
333 self._models_loaded = False
334 self.logger.debug("Unloaded idle analysis models")
335
336 @abstractmethod
337 async def _start_analysis(
338 self,
339 session_id: str,
340 streamdetails: StreamDetails,
341 audio_format: AudioFormat,
342 ) -> bool:
343 """
344 Provider-specific initialization for a new analysis session.
345
346 Return False to reject the session (e.g. unsupported format).
347
348 :param session_id: The analysis session ID.
349 :param streamdetails: The stream details for the item being analyzed.
350 :param audio_format: PCM format of the audio stream.
351 """
352
353 @abstractmethod
354 async def _finalize(self, session_id: str) -> AudioAnalysisData | None:
355 """
356 Compute and return the analysis for this session (or None to skip).
357
358 The base class persists the returned value via set_audio_analysis() and
359 then fires post_analysis(). Return None to skip both.
360
361 :param session_id: The analysis session ID.
362 """
363
364 async def _load_models(self) -> None:
365 """Load heavy models into memory. Override when has_unloadable_models is True."""
366 return
367
368 def _free_models(self) -> None:
369 """Drop references to heavy models so memory can be reclaimed. Override as needed."""
370 return
371
372 async def _run_offloaded(self, func: Callable[..., _T], /, *args: Any, **kwargs: Any) -> _T:
373 """
374 Run a blocking analysis function on the niced analysis thread pool.
375
376 Route all CPU-heavy analysis work through this. The controller's semaphore caps
377 concurrency, and while any player is streaming the solo lock holds analysis to one
378 offload at a time so it stays below playback. Falls back to asyncio.to_thread when the
379 controller has no pool configured.
380
381 The slot (semaphore permit and solo lock) is held until the worker thread finishes,
382 even if the awaiting coroutine is cancelled: the thread keeps running, so it must keep
383 counting against the caps.
384
385 :param func: The blocking callable to run off the event loop.
386 :param args: Positional arguments passed to func.
387 :param kwargs: Keyword arguments passed to func.
388 """
389 controller = self.mass.streams.audio_analysis
390 semaphore = controller.analysis_semaphore
391 if not isinstance(semaphore, asyncio.Semaphore):
392 return await asyncio.to_thread(func, *args, **kwargs)
393
394 # While a player streams, take the solo lock so analysis runs one offload at a time:
395 # the GIL serializes Python execution, so concurrent inference starves the playback
396 # loop. Idle, the semaphore alone caps concurrency and background work uses spare cores.
397 solo = getattr(controller, "analysis_solo_lock", None)
398 solo_lock = solo if isinstance(solo, asyncio.Lock) else None
399 take_solo = False
400 if solo_lock is not None:
401 try:
402 take_solo = bool(controller.playback_active())
403 except Exception:
404 take_solo = False
405
406 solo_held = False
407
408 def _release(done: asyncio.Future[_T]) -> None:
409 if solo_held and solo_lock is not None:
410 solo_lock.release()
411 semaphore.release()
412 # Retrieve any exception so a cancelled awaiter doesn't leave it unretrieved.
413 if not done.cancelled():
414 done.exception()
415
416 wait_start = time.monotonic()
417 await semaphore.acquire()
418 try:
419 if take_solo and solo_lock is not None:
420 await solo_lock.acquire()
421 solo_held = True
422 except BaseException:
423 # Cancelled (or failed) waiting for the solo lock â give the permit back.
424 semaphore.release()
425 raise
426 wait_seconds = time.monotonic() - wait_start
427 if wait_seconds >= _PERMIT_WAIT_DEBUG_THRESHOLD and isinstance(
428 semaphore, InstrumentedSemaphore
429 ):
430 self.logger.debug(
431 "Analysis offload waited %.1fs for a slot (%d/%d permits in use, %d queued)",
432 wait_seconds,
433 semaphore.in_flight,
434 semaphore.capacity,
435 semaphore.waiters,
436 )
437 executor = getattr(controller, "analysis_executor", None)
438 try:
439 if isinstance(executor, ThreadPoolExecutor):
440 loop = asyncio.get_running_loop()
441 future: asyncio.Future[_T] = asyncio.ensure_future(
442 loop.run_in_executor(executor, functools.partial(func, *args, **kwargs))
443 )
444 else:
445 future = asyncio.ensure_future(asyncio.to_thread(func, *args, **kwargs))
446 except RuntimeError:
447 # Scheduling the worker failed (e.g. the loop is shutting down) â release the
448 # permit we just took so it isn't leaked, which would shrink the cap over time.
449 if solo_held and solo_lock is not None:
450 solo_lock.release()
451 semaphore.release()
452 raise
453 future.add_done_callback(_release)
454 # join: a cancelled awaiter leaves the thread and its slot held until the work finishes.
455 return await join_task(future)
456
457 async def _run_offloaded_timed(
458 self, func: Callable[..., _T], /, *args: Any, **kwargs: Any
459 ) -> tuple[_T, float]:
460 """
461 Run a blocking analysis function offloaded, returning its result and execution seconds.
462
463 Same slot semantics as ``_run_offloaded``; the returned seconds cover only the
464 callable's own execution, never the time spent queued for a slot.
465
466 :param func: The blocking callable to run off the event loop.
467 :param args: Positional arguments passed to func.
468 :param kwargs: Keyword arguments passed to func.
469 """
470
471 def _timed() -> tuple[_T, float]:
472 start = time.perf_counter()
473 return func(*args, **kwargs), time.perf_counter() - start
474
475 return await self._run_offloaded(_timed)
476
477 async def _record_failure(
478 self,
479 session: AnalysisSessionData,
480 reason: str,
481 retry_at: datetime | None,
482 ) -> None:
483 """Record a failure for this session's track."""
484 # Swallow DB write errors so a failed recorder write never breaks the session
485 # lifecycle (record_analysis_failure only performs a sqlite insert).
486 try:
487 sd = session.streamdetails
488 await self.mass.streams.audio_analysis.record_analysis_failure(
489 item_id=sd.item_id,
490 provider_instance_id_or_domain=sd.provider,
491 aa_provider_domain=self.domain,
492 reason=reason,
493 retry_at=retry_at,
494 analysis_version=self.analysis_version,
495 media_type=sd.media_type,
496 )
497 except sqlite3.Error as err:
498 self.logger.warning(
499 "record_analysis_failure raised for %s: %s", self.domain, err, exc_info=err
500 )
501