Checkpoint speaker transcriptions incrementally

internal/stage/transcribe.go

@@ -8,7 +8,6 @@ import (
 	"path/filepath"
 	"sort"
 	"strings"
-	"sync"
 
 	"golang.org/x/sync/errgroup"
 
@@ -28,11 +27,14 @@ func (transcribeStage) Run(ctx context.Context, env *Env, m *manifest.Manifest)
 	if env.ArtifactStore == nil {
 		return nil, fmt.Errorf("transcribe: artifact store is required")
 	}
+	if env.ManifestStore == nil {
+		return nil, fmt.Errorf("transcribe: manifest store is required")
+	}
 	if env.Config.Pipeline == nil || env.Config.Session == nil {
 		return nil, fmt.Errorf("transcribe: resolved config must include pipeline and session")
 	}
-	if env.WhisperX == nil {
-		return nil, fmt.Errorf("transcribe: whisperx adapter is required")
+	if m == nil {
+		return nil, fmt.Errorf("transcribe: session manifest is required")
 	}
 
 	var sessionID string
@@ -47,21 +49,29 @@ func (transcribeStage) Run(ctx context.Context, env *Env, m *manifest.Manifest)
 	}
 
 	paths := sessionPathsForEnv(env, sessionID)
-	runLayout, err := resolveRunStageLayout(env, m, paths, sessionID, "transcribe")
+	reconciliation, err := reconcileTranscribeSpeakers(ctx, env, m)
 	if err != nil {
-		return nil, fmt.Errorf("transcribe: resolve run-stage layout: %w", err)
+		return nil, fmt.Errorf("transcribe: reconcile speakers: %w", err)
 	}
-	jobs, err := buildTranscribeJobs(m, paths)
-	if err != nil {
-		return nil, fmt.Errorf("transcribe: resolve audio inputs: %w", err)
+	state := newTranscribeExecutionState(env, m, paths, reconciliation.current)
+	if err := state.checkpointStartingCollection(ctx); err != nil {
+		return nil, fmt.Errorf("transcribe: checkpoint reconciled speaker collection before execution: %w", err)
 	}
-	if len(jobs) == 0 {
-		return nil, fmt.Errorf("transcribe: no prepared audio files found")
+
+	jobs := append([]transcribeReconciliationJob(nil), reconciliation.work...)
+	if len(jobs) > 0 && env.WhisperX == nil {
+		return nil, fmt.Errorf("transcribe: whisperx adapter is required")
 	}
-	for index := range jobs {
-		jobs[index].runOutputPath, err = runLocalPathForCanonical(runLayout, paths, jobs[index].canonicalOutputPath)
+	if len(jobs) > 0 {
+		runLayout, err := resolveRunStageLayout(env, m, paths, sessionID, "transcribe")
 		if err != nil {
-			return nil, fmt.Errorf("transcribe: resolve run-local output path for %q: %w", jobs[index].speakerID, err)
+			return nil, fmt.Errorf("transcribe: resolve run-stage layout: %w", err)
+		}
+		for index := range jobs {
+			jobs[index].runOutputPath, err = runLocalPathForCanonical(runLayout, paths, jobs[index].canonicalOutputPath)
+			if err != nil {
+				return nil, fmt.Errorf("transcribe: resolve run-local output path for %q: %w", jobs[index].speakerID, err)
+			}
 		}
 	}
 
@@ -69,15 +79,10 @@ func (transcribeStage) Run(ctx context.Context, env *Env, m *manifest.Manifest)
 	if env.Config.Pipeline.WhisperX.Concurrency != nil && *env.Config.Pipeline.WhisperX.Concurrency > 0 {
 		concurrency = *env.Config.Pipeline.WhisperX.Concurrency
 	}
-	if concurrency > len(jobs) {
+	if len(jobs) > 0 && concurrency > len(jobs) {
 		concurrency = len(jobs)
 	}
 
-	var (
-		mu        sync.Mutex
-		completed int
-		perFile   = map[string]map[string]any{}
-	)
 	group, groupCtx := errgroup.WithContext(ctx)
 	group.SetLimit(concurrency)
 	for _, job := range jobs {
@@ -86,6 +91,9 @@ func (transcribeStage) Run(ctx context.Context, env *Env, m *manifest.Manifest)
 			if err := groupCtx.Err(); err != nil {
 				return err
 			}
+			if err := state.checkpointFailure(); err != nil {
+				return err
+			}
 			res, runErr := env.WhisperX.Transcribe(groupCtx, whisperx.TranscribeRequest{
 				SpeakerID:               job.speakerID,
 				AudioPath:               job.audioPath,
@@ -95,13 +103,6 @@ func (transcribeStage) Run(ctx context.Context, env *Env, m *manifest.Manifest)
 				return fmt.Errorf("speaker %q (%s): %w", job.speakerID, filepath.Base(job.audioPath), runErr)
 			}
 
-			if _, err := authoritativeOutputPath(job.runOutputPath, res.OutputRawTranscriptPath); err != nil {
-				return fmt.Errorf("speaker %q: %w", job.speakerID, err)
-			}
-			if err := validateTranscriptJSONFile(job.runOutputPath); err != nil {
-				return fmt.Errorf("speaker %q output %q invalid: %w", job.speakerID, job.runOutputPath, err)
-			}
-
 			meta := map[string]any{
 				"attempts":    res.Attempts,
 				"http_status": res.HTTPStatus,
@@ -109,17 +110,17 @@ func (transcribeStage) Run(ctx context.Context, env *Env, m *manifest.Manifest)
 				"output_path": job.runOutputPath,
 			}
 
-			mu.Lock()
-			perFile[job.speakerID] = meta
-			completed++
-			mu.Unlock()
+			if err := state.acceptSpeaker(ctx, job, res.OutputRawTranscriptPath, meta); err != nil {
+				return fmt.Errorf("speaker %q: %w", job.speakerID, err)
+			}
 			return nil
 		})
 	}
 	groupErr := group.Wait()
-	mu.Lock()
-	completedCount := completed
-	mu.Unlock()
+	if err := state.checkpointFailure(); err != nil {
+		return nil, err
+	}
+	completedCount := state.completedCount()
 	if err := ctx.Err(); err != nil {
 		return nil, fmt.Errorf("transcribe: canceled after completing %d of %d audio jobs: %w", completedCount, len(jobs), err)
 	}
@@ -129,34 +130,9 @@ func (transcribeStage) Run(ctx context.Context, env *Env, m *manifest.Manifest)
 	if completedCount != len(jobs) {
 		return nil, fmt.Errorf("transcribe: completed %d of %d planned audio jobs", completedCount, len(jobs))
 	}
-
-	speakers := make([]string, 0, len(perFile))
-	for speaker := range perFile {
-		speakers = append(speakers, speaker)
-	}
-	sort.Strings(speakers)
-
-	outputs := make([]artifacts.Ref, 0, len(speakers))
-	runOutputPaths := make([]string, 0, len(speakers))
-	outputPaths := make([]string, 0, len(speakers))
-	orderedPerFile := make(map[string]any, len(speakers))
-	for _, speaker := range speakers {
-		var selected transcribeJob
-		for _, job := range jobs {
-			if job.speakerID == speaker {
-				selected = job
-				break
-			}
-		}
-		runOutputPaths = append(runOutputPaths, selected.runOutputPath)
-		ref := artifacts.Ref{Kind: "transcript_raw", Category: "transcripts", SessionID: sessionID, AbsolutePath: selected.runOutputPath}
-		materialized, err := materializeRunLocalOutput(env.ArtifactStore, selected.runOutputPath, selected.canonicalOutputPath, ref)
-		if err != nil {
-			return nil, fmt.Errorf("transcribe: materialize %q output: %w", speaker, err)
-		}
-		outputs = append(outputs, materialized)
-		outputPaths = append(outputPaths, selected.canonicalOutputPath)
-		orderedPerFile[speaker] = perFile[speaker]
+	outputs, outputPaths, runOutputPaths, orderedPerFile, err := state.successDetails()
+	if err != nil {
+		return nil, fmt.Errorf("transcribe: assemble completed outputs: %w", err)
 	}
 
 	retries := 0
@@ -168,7 +144,7 @@ func (transcribeStage) Run(ctx context.Context, env *Env, m *manifest.Manifest)
 		Outputs: outputs,
 		Metadata: map[string]any{
 			"stage":             "transcribe",
-			"audio_files_count": len(jobs),
+			"audio_files_count": len(reconciliation.jobs),
 			"language":          env.Config.Pipeline.WhisperX.Language,
 			"concurrency":       concurrency,
 			"retries":           retries,