Warm up pyannote and FunASR runtimes

MeetingAssistant/MeetingAssistantOptions.cs

@@ -181,7 +181,7 @@ public sealed class AzureSpeechOptions
 
     public string KeyEnv { get; set; } = "AZURE_SPEECH_KEY";
 
-    public TimeSpan RecognitionStopTimeout { get; set; } = TimeSpan.FromSeconds(10);
+    public TimeSpan RecognitionStopTimeout { get; set; } = TimeSpan.FromMinutes(3);
 
     public bool DiarizeIntermediateResults { get; set; } = true;
 
@@ -321,6 +321,7 @@ public sealed class SpeakerIdentityPyannoteValidationOptions
     public PyannoteDiarizationOptions Diarization { get; set; } = new()
     {
         Enabled = true,
+        CommandTimeout = TimeSpan.FromHours(1),
         AlignmentMode = PyannoteAlignmentMode.PyannoteTurns
     };
 }

MeetingAssistant/Program.cs

@@ -89,6 +89,7 @@ builder.Services.AddSingleton<MeetingRecordingCoordinator>();
 builder.Services.AddHostedService<TemporaryRecordingCleanupHostedService>();
 builder.Services.AddHostedService<SpeakerIdentityDatabaseInitializer>();
 builder.Services.AddHostedService<SpeechRecognitionPipelineHostedService>();
+builder.Services.AddHostedService<PyannoteDiarizationWarmupHostedService>();
 #if WINDOWS
 builder.Services.AddHostedService<GlobalHotkeyService>();
 builder.Services.AddHostedService<UnoTaskbarIconService>();

MeetingAssistant/Transcription/PyannoteDiarizationWarmupHostedService.cs

@@ -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());
+    }
+}

MeetingAssistant/Transcription/PyannoteTranscriptFinalizer.cs

@@ -82,6 +82,58 @@ public sealed class PyannoteTranscriptFinalizer
         return segments;
     }
 
+    public async Task WarmUpAsync(
+        PyannoteDiarizationOptions diarization,
+        CancellationToken cancellationToken)
+    {
+        if (!diarization.Enabled)
+        {
+            return;
+        }
+
+        var token = ResolveToken(diarization);
+        if (string.IsNullOrWhiteSpace(token))
+        {
+            logger.LogWarning(
+                "Pyannote warm-up is enabled but no Hugging Face token is configured directly or via {TokenEnv}",
+                diarization.TokenEnv);
+            return;
+        }
+
+        var modelsFolder = VaultPath.Resolve(diarization.ModelsFolder);
+        Directory.CreateDirectory(modelsFolder);
+        using var timeoutSource = diarization.CommandTimeout > TimeSpan.Zero
+            ? new CancellationTokenSource(diarization.CommandTimeout)
+            : null;
+        using var linkedSource = timeoutSource is null
+            ? null
+            : CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, timeoutSource.Token);
+
+        try
+        {
+            if (diarization.BuildImage)
+            {
+                await EnsureDockerImageAsync(diarization, modelsFolder, linkedSource?.Token ?? cancellationToken);
+            }
+
+            var result = await commandRunner.RunAsync(
+                diarization.DockerCommand,
+                BuildWarmUpDockerArguments(diarization, modelsFolder),
+                linkedSource?.Token ?? cancellationToken,
+                new Dictionary<string, string> { ["HF_TOKEN"] = token });
+            if (result.ExitCode != 0)
+            {
+                throw new InvalidOperationException(
+                    $"pyannote warm-up failed with exit code {result.ExitCode}: {result.StandardError}");
+            }
+        }
+        catch (OperationCanceledException) when (!cancellationToken.IsCancellationRequested && timeoutSource?.IsCancellationRequested == true)
+        {
+            throw new TimeoutException(
+                $"pyannote warm-up timed out after {diarization.CommandTimeout}.");
+        }
+    }
+
     private async Task<CommandResult> RunDiarizationAsync(
         string fullAudioPath,
         string token,
@@ -189,6 +241,33 @@ public sealed class PyannoteTranscriptFinalizer
         ];
     }
 
+    private string[] BuildWarmUpDockerArguments(
+        PyannoteDiarizationOptions diarization,
+        string modelsFolder)
+    {
+        return
+        [
+            "run",
+            "--rm",
+            "-e",
+            "HF_TOKEN",
+            "-v",
+            $"{modelsFolder}:/workspace/cache",
+            "-e",
+            "HF_HOME=/workspace/cache/huggingface",
+            "-e",
+            "XDG_CACHE_HOME=/workspace/cache",
+            "-e",
+            "PIP_CACHE_DIR=/workspace/cache/pip",
+            "-e",
+            "TORCH_HOME=/workspace/cache/torch",
+            diarization.Image,
+            "sh",
+            "-lc",
+            BuildWarmUpPythonCommand(diarization)
+        ];
+    }
+
     private static string BuildPythonCommand(
         PyannoteDiarizationOptions diarization,
         SpeechRecognitionPipelineOptions pipelineOptions)
@@ -217,6 +296,21 @@ public sealed class PyannoteTranscriptFinalizer
             + "python /tmp/meeting_assistant_pyannote.py";
     }
 
+    private static string BuildWarmUpPythonCommand(PyannoteDiarizationOptions diarization)
+    {
+        var model = diarization.Model;
+        return
+            "cat > /tmp/meeting_assistant_pyannote_warmup.py <<'PY'\n"
+            + "import os\n"
+            + "import torch\n"
+            + "from pyannote.audio import Pipeline\n"
+            + $"pipeline = Pipeline.from_pretrained({JsonSerializer.Serialize(model)}, token=os.environ.get('HF_TOKEN'))\n"
+            + "pipeline.to(torch.device('cpu'))\n"
+            + "print('pyannote warm-up complete')\n"
+            + "PY\n"
+            + "python /tmp/meeting_assistant_pyannote_warmup.py";
+    }
+
     private static string BuildPipelineInvocation(SpeechRecognitionPipelineOptions pipelineOptions)
     {
         return pipelineOptions.NumSpeakers is > 0

MeetingAssistant/Transcription/SpeechRecognitionPipelineHostedService.cs

@@ -1,27 +1,41 @@
+using MeetingAssistant.LaunchProfiles;
+
 namespace MeetingAssistant.Transcription;
 
 public sealed class SpeechRecognitionPipelineHostedService : IHostedService
 {
     private readonly ISpeechRecognitionPipelineFactory pipelineFactory;
+    private readonly ILaunchProfileOptionsProvider? launchProfiles;
     private readonly ILogger<SpeechRecognitionPipelineHostedService> logger;
     private CancellationTokenSource? startupCancellation;
     private Task? startupTask;
-    private ISpeechRecognitionPipeline? startupPipeline;
+    private List<ISpeechRecognitionPipeline> startupPipelines = [];
 
     public SpeechRecognitionPipelineHostedService(
         ISpeechRecognitionPipelineFactory pipelineFactory,
         ILogger<SpeechRecognitionPipelineHostedService> logger)
+        : this(pipelineFactory, null, logger)
+    {
+    }
+
+    public SpeechRecognitionPipelineHostedService(
+        ISpeechRecognitionPipelineFactory pipelineFactory,
+        ILaunchProfileOptionsProvider? launchProfiles,
+        ILogger<SpeechRecognitionPipelineHostedService> logger)
     {
         this.pipelineFactory = pipelineFactory;
+        this.launchProfiles = launchProfiles;
         this.logger = logger;
     }
 
     public Task StartAsync(CancellationToken cancellationToken)
     {
-        startupCancellation = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
-        startupPipeline = pipelineFactory.Create();
+        var cancellation = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
+        var pipelines = CreateStartupPipelines();
+        startupCancellation = cancellation;
+        startupPipelines = pipelines;
         startupTask = Task.Run(
-            () => WarmUpPipelineAsync(startupPipeline, startupCancellation.Token),
+            () => WarmUpPipelinesAsync(pipelines, cancellation.Token),
             CancellationToken.None);
         return Task.CompletedTask;
     }
@@ -30,15 +44,15 @@ public sealed class SpeechRecognitionPipelineHostedService : IHostedService
     {
         var cancellation = startupCancellation;
         var task = startupTask;
-        var pipeline = startupPipeline;
-        if (cancellation is null || task is null || pipeline is null)
+        var pipelines = startupPipelines;
+        if (cancellation is null || task is null)
         {
             return;
         }
 
         startupCancellation = null;
         startupTask = null;
-        startupPipeline = null;
+        startupPipelines = [];
         await cancellation.CancelAsync();
         try
         {
@@ -49,11 +63,43 @@ public sealed class SpeechRecognitionPipelineHostedService : IHostedService
         }
         finally
         {
-            await pipeline.DisposeAsync();
+            foreach (var pipeline in pipelines)
+            {
+                await pipeline.DisposeAsync();
+            }
+
             cancellation.Dispose();
         }
     }
 
+    private List<ISpeechRecognitionPipeline> CreateStartupPipelines()
+    {
+        if (launchProfiles is null)
+        {
+            return [pipelineFactory.Create()];
+        }
+
+        return launchProfiles.GetProfiles()
+            .Where(ShouldWarmUpProfile)
+            .Select(profile => pipelineFactory.Create(profile.Name))
+            .ToList();
+    }
+
+    private static bool ShouldWarmUpProfile(LaunchProfile profile)
+    {
+        return profile.Name.Equals(
+                ConfigurationLaunchProfileOptionsProvider.DefaultProfileName,
+                StringComparison.OrdinalIgnoreCase) ||
+            profile.Options.Recording.TranscriptionProvider.Equals("funasr", StringComparison.OrdinalIgnoreCase);
+    }
+
+    private async Task WarmUpPipelinesAsync(
+        IReadOnlyList<ISpeechRecognitionPipeline> pipelines,
+        CancellationToken cancellationToken)
+    {
+        await Task.WhenAll(pipelines.Select(pipeline => WarmUpPipelineAsync(pipeline, cancellationToken)));
+    }
+
     private async Task WarmUpPipelineAsync(
         ISpeechRecognitionPipeline pipeline,
         CancellationToken cancellationToken)

MeetingAssistant/appsettings.json

@@ -89,7 +89,7 @@
         "de-DE"
       ],
       "KeyEnv": "AZURE_SPEECH_KEY",
-      "RecognitionStopTimeout": "00:00:10",
+      "RecognitionStopTimeout": "00:03:00",
       "DiarizeIntermediateResults": true,
       "PostProcessingOption": "",
       "PhraseListWeight": 1.5
@@ -124,7 +124,7 @@
           "AnnotationSource": "SpeakerDiarization",
           "AlignmentMode": "PyannoteTurns",
           "TokenEnv": "HF_TOKEN",
-          "CommandTimeout": "00:30:00"
+          "CommandTimeout": "01:00:00"
         }
       }
     },