meeting-assistant/MeetingAssistant/Transcription/PyannoteDiarizationWarmupHostedService.cs

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);
                if (cancellationToken.IsCancellationRequested)
                {
                    return;
                }

                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) when (cancellationToken.IsCancellationRequested)
            {
                return;
            }
            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());
    }
}