Warm up pyannote and FunASR runtimes

2026-05-28 13:05:27 +02:00
parent 7ff93b73b3
commit c46c587e65
14 changed files with 696 additions and 20 deletions
@@ -0,0 +1,125 @@
+using MeetingAssistant.LaunchProfiles;
+
+namespace MeetingAssistant.Transcription;
+
+public sealed class PyannoteDiarizationWarmupHostedService : IHostedService
+{
+    private readonly PyannoteTranscriptFinalizer finalizer;
+    private readonly ILaunchProfileOptionsProvider launchProfiles;
+    private readonly ILogger<PyannoteDiarizationWarmupHostedService> logger;
+    private CancellationTokenSource? startupCancellation;
+    private Task? startupTask;
+
+    public PyannoteDiarizationWarmupHostedService(
+        PyannoteTranscriptFinalizer finalizer,
+        ILaunchProfileOptionsProvider launchProfiles,
+        ILogger<PyannoteDiarizationWarmupHostedService> logger)
+    {
+        this.finalizer = finalizer;
+        this.launchProfiles = launchProfiles;
+        this.logger = logger;
+    }
+
+    public Task StartAsync(CancellationToken cancellationToken)
+    {
+        var cancellation = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
+        startupCancellation = cancellation;
+        startupTask = Task.Run(
+            () => WarmUpEnabledRuntimesAsync(cancellation.Token),
+            CancellationToken.None);
+        return Task.CompletedTask;
+    }
+
+    public async Task StopAsync(CancellationToken cancellationToken)
+    {
+        var cancellation = startupCancellation;
+        var task = startupTask;
+        if (cancellation is null || task is null)
+        {
+            return;
+        }
+
+        startupCancellation = null;
+        startupTask = null;
+        await cancellation.CancelAsync();
+        try
+        {
+            await task.WaitAsync(cancellationToken);
+        }
+        catch (OperationCanceledException)
+        {
+        }
+        finally
+        {
+            cancellation.Dispose();
+        }
+    }
+
+    private async Task WarmUpEnabledRuntimesAsync(CancellationToken cancellationToken)
+    {
+        foreach (var diarization in GetEnabledDiarizationOptions())
+        {
+            try
+            {
+                logger.LogInformation(
+                    "Starting pyannote warm-up for image {Image} and model {Model}",
+                    diarization.Image,
+                    diarization.Model);
+                await finalizer.WarmUpAsync(diarization, cancellationToken);
+                logger.LogInformation(
+                    "Finished pyannote warm-up for image {Image} and model {Model}",
+                    diarization.Image,
+                    diarization.Model);
+            }
+            catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
+            {
+                logger.LogInformation("Pyannote warm-up was cancelled during application shutdown");
+                throw;
+            }
+            catch (Exception exception)
+            {
+                logger.LogError(
+                    exception,
+                    "Pyannote warm-up failed for image {Image} and model {Model}; diarization can still retry on demand",
+                    diarization.Image,
+                    diarization.Model);
+            }
+        }
+    }
+
+    private IEnumerable<PyannoteDiarizationOptions> GetEnabledDiarizationOptions()
+    {
+        return launchProfiles.GetProfiles()
+            .SelectMany(profile => GetEnabledDiarizationOptions(profile.Options))
+            .DistinctBy(CreateWarmUpKey);
+    }
+
+    private static IEnumerable<PyannoteDiarizationOptions> GetEnabledDiarizationOptions(
+        MeetingAssistantOptions options)
+    {
+        if (options.Recording.TranscriptionProvider.Equals("whisper-local", StringComparison.OrdinalIgnoreCase) &&
+            options.WhisperLocal.Diarization.Enabled)
+        {
+            yield return options.WhisperLocal.Diarization;
+        }
+
+        if (options.SpeakerIdentification.PyannoteValidation.Enabled &&
+            options.SpeakerIdentification.PyannoteValidation.Diarization.Enabled)
+        {
+            yield return options.SpeakerIdentification.PyannoteValidation.Diarization;
+        }
+    }
+
+    private static string CreateWarmUpKey(PyannoteDiarizationOptions diarization)
+    {
+        return string.Join(
+            '\u001f',
+            diarization.DockerCommand,
+            diarization.Image,
+            diarization.Model,
+            VaultPath.Resolve(diarization.ModelsFolder),
+            diarization.Token,
+            diarization.TokenEnv,
+            diarization.BuildImage.ToString());
+    }
+}