Warm up pyannote and FunASR runtimes

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