cohere-transcribe 0.1.2 → 0.1.3

data/lib/cohere/transcribe/runtime/engine.rb

@@ -264,7 +264,13 @@ module Cohere
             requested_workers: resolved_options.preprocess_workers,
             enabled: resolved_options.pipeline_preparation,
             worker_limit: vad_file_concurrency_limit(resolved_options),
-            estimate_bytes: estimate_bytes
+            estimate_bytes: estimate_bytes,
+            exclusive_retry: lambda do |prepared|
+              prepared.error.is_a?(Audio::DecodedAudioLimitError)
+            end,
+            retained_bytes: lambda do |prepared|
+              prepared.decoded ? prepared.decoded.samples.byte_size : 0
+            end
           ) do |item, decoded_byte_limit, worker_slot|
             prepare_entry(
               item,
@@ -371,7 +377,12 @@ module Cohere
         end
 
         def preflight_outcome(index, entry, plan, resolved_options, measurements)
-          Output::Publication.with_plan_lock(plan) do
+          if resolved_options.publication&.existing == "skip" && plan.skipped
+            optimistic = Output::Publication.revalidate(plan, resolved_options)
+            return skipped_result(entry, plan) if optimistic.action == :skip
+          end
+
+          Output::Publication.with_plan_lock(plan) do |lock|
             decision = Output::Publication.revalidate(plan, resolved_options)
             case decision.action
             when :skip
@@ -383,7 +394,8 @@ module Cohere
                 plan,
                 decision,
                 resolved_options,
-                measurements
+                measurements,
+                lock
               )
             else
               PreparationItem.new(index: index, entry: entry, plan: plan)
@@ -391,7 +403,7 @@ module Cohere
           end
         end
 
-        def preflight_resumed_entry(index, entry, plan, decision, resolved_options, measurements)
+        def preflight_resumed_entry(index, entry, plan, decision, resolved_options, measurements, lock)
           prepared = Preparation::PreparedEntry.new(
             item: PreparationItem.new(index: index, entry: entry, plan: plan),
             snapshot: source_snapshot(entry.path),
@@ -410,6 +422,7 @@ module Cohere
             decision.generation_id,
             resolved_options,
             measurements,
+            lock: lock,
             defer_word_alignment: resolved_options.alignment == "word"
           )
         rescue FatalRuntimeError, ProgressCallbackError
@@ -569,7 +582,7 @@ module Cohere
           measurements.decode_seconds += prepared.decode_seconds
           measurements.vad_seconds += prepared.vad_seconds
           record_vad_measurements(measurements, prepared.vad_details)
-          Output::Publication.with_plan_lock(plan) do
+          Output::Publication.with_plan_lock(plan) do |lock|
             decision = Output::Publication.revalidate(plan, resolved_options)
             return skipped_result(entry, plan) if decision.action == :skip
 
@@ -587,11 +600,12 @@ module Cohere
                 decision.checkpoint,
                 decision.generation_id,
                 resolved_options,
-                measurements
+                measurements,
+                lock: lock
               )
             end
 
-            process_fresh_entry(prepared, resolved_options, measurements)
+            process_fresh_entry(prepared, resolved_options, measurements, lock: lock)
           end
         rescue FatalRuntimeError, ProgressCallbackError
           raise
@@ -609,7 +623,7 @@ module Cohere
           measurements.decode_seconds += prepared.decode_seconds
           measurements.vad_seconds += prepared.vad_seconds
           record_vad_measurements(measurements, prepared.vad_details)
-          outcome = Output::Publication.with_plan_lock(plan) do
+          outcome = Output::Publication.with_plan_lock(plan) do |lock|
             decision = Output::Publication.revalidate(plan, resolved_options)
             if decision.action == :skip
               WordPipeline::Final.new(index: item.index, result: skipped_result(entry, plan))
@@ -635,6 +649,7 @@ module Cohere
                   decision.generation_id,
                   resolved_options,
                   measurements,
+                  lock: lock,
                   defer_word_alignment: true
                 )
               else
@@ -642,6 +657,7 @@ module Cohere
                   prepared,
                   resolved_options,
                   measurements,
+                  lock: lock,
                   defer_word_alignment: true
                 )
               end
@@ -661,7 +677,7 @@ module Cohere
         def process_word_alignment(work, reload_result, resolved_options, measurements, alignment_state)
           entry = work.entry
           plan = work.plan
-          Output::Publication.with_plan_lock(plan) do
+          Output::Publication.with_plan_lock(plan) do |lock|
             decision = Output::Publication.revalidate(plan, resolved_options)
             return skipped_result(entry, plan) if decision.action == :skip
             if (alignment_error = alignment_state.fetch(:error))
@@ -679,6 +695,7 @@ module Cohere
               resolved_options,
               measurements,
               work.generation_id,
+              lock: lock,
               asr_evicted: true
             )
             ensure_source_unchanged!(entry.path, work.source_snapshot)
@@ -710,7 +727,7 @@ module Cohere
           record_alignment_measurements(measurements, before, aligner) if before && aligner
         end
 
-        def process_fresh_entry(prepared, resolved_options, measurements, defer_word_alignment: false)
+        def process_fresh_entry(prepared, resolved_options, measurements, lock:, defer_word_alignment: false)
           item = prepared.item
           entry = item.entry
           plan = item.plan
@@ -787,7 +804,8 @@ module Cohere
                                 speech_spans: speech_spans,
                                 vad_provider_options: vad_provider_options(vad_details),
                                 directory_binding: plan.directory_bindings.last,
-                                guard_bindings: plan.directory_bindings
+                                guard_bindings: plan.directory_bindings,
+                                lock: lock
                               )
                             ensure
                               measurements.checkpoint_seconds += monotonic - checkpoint_started
@@ -812,12 +830,13 @@ module Cohere
             plan,
             resolved_options,
             measurements,
-            generation_id
+            generation_id,
+            lock: lock
           )
         end
 
         def process_resumed_entry(prepared, checkpoint, generation_id, resolved_options, measurements,
-                                  defer_word_alignment: false)
+                                  lock:, defer_word_alignment: false)
           item = prepared.item
           entry = item.entry
           plan = item.plan
@@ -866,7 +885,8 @@ module Cohere
             plan,
             resolved_options,
             measurements,
-            generation_id
+            generation_id,
+            lock: lock
           )
         end
 
@@ -888,7 +908,7 @@ module Cohere
         end
 
         def render_and_publish(result, samples, segment_times, speech_spans, plan, resolved_options,
-                               measurements, generation_id, asr_evicted: false)
+                               measurements, generation_id, lock:, asr_evicted: false)
           progressive_started = monotonic unless resolved_options.alignment == "word"
           words, fallback_alignment_segments = words_for_segments(
             samples,
@@ -922,7 +942,8 @@ module Cohere
             result,
             resolved_options,
             generation_id: generation_id,
-            speech_spans: speech_spans
+            speech_spans: speech_spans,
+            lock: lock
           )
           if outputs.any?
             result = result.with(

data/lib/cohere/transcribe/runtime/preparation.rb

@@ -31,15 +31,85 @@ module Cohere
         #
         # A pipelined group receives at most half of the configured decoded-PCM
         # budget (and never more than 512 MiB). Estimated decode sizes determine
-        # group membership and per-file ceilings. A file that cannot fit that cap
-        # is prepared alone with the full configured ceiling and without overlap.
+        # group membership and per-file ceilings. Unknown sizes receive an equal
+        # share initially. A ceiling failure first retries within the remaining
+        # configured PCM budget while successful group entries stay retained. If
+        # that is insufficient, later retained audio is released and only those
+        # entries are prepared again after the full-ceiling retry. A known file
+        # that cannot fit the group cap starts with that exclusive path.
         # Native decoder implementation transients are outside this retained-PCM
         # accounting, as they are in the Python path.
         class Pipeline
           include Enumerable
 
-          Group = Data.define(:items, :limits, :exclusive)
-          private_constant :Group
+          Group = Data.define(:items, :limits, :exclusive) do
+            def self.exclusive(item, limit)
+              new(items: [item].freeze, limits: [limit].freeze, exclusive: true)
+            end
+          end
+
+          # Explicit cursor keeps metadata probes on the preparation caller's
+          # native stack while preserving lazy current/next-group discovery.
+          class GroupCursor
+            def initialize(items:, workers:, group_byte_limit:, memory_byte_limit:,
+                           estimated_reservation:, unknown_reservation:, bounded_group:)
+              @items = items
+              @workers = workers
+              @group_byte_limit = group_byte_limit
+              @memory_byte_limit = memory_byte_limit
+              @estimated_reservation = estimated_reservation
+              @unknown_reservation = unknown_reservation
+              @bounded_group = bounded_group
+              @index = 0
+              @pending = nil
+            end
+
+            def next
+              items = []
+              estimates = []
+              total = 0
+              while items.length < @workers
+                pair = take
+                break unless pair
+
+                item, estimate = pair
+                if estimate > @group_byte_limit
+                  return Group.exclusive(item, @memory_byte_limit) if items.empty?
+
+                  @pending = pair
+                  break
+                end
+                if !items.empty? && total + estimate > @group_byte_limit
+                  @pending = pair
+                  break
+                end
+
+                items << item
+                estimates << estimate
+                total += estimate
+              end
+              raise StopIteration if items.empty?
+
+              @bounded_group.call(items, estimates)
+            end
+
+            private
+
+            def take
+              if @pending
+                pair = @pending
+                @pending = nil
+                return pair
+              end
+              return if @index >= @items.length
+
+              item = @items.fetch(@index)
+              @index += 1
+              estimate = @estimated_reservation.call(item) || @unknown_reservation
+              [item, estimate]
+            end
+          end
+          private_constant :Group, :GroupCursor
 
           class WorkerPool
             STOP = Object.new.freeze
@@ -94,7 +164,7 @@ module Cohere
           attr_reader :effective_workers, :group_byte_limit, :memory_byte_limit, :wait_seconds
 
           def initialize(items, memory_byte_limit:, requested_workers:, enabled:, worker_limit: nil,
-                         estimate_bytes: nil, &prepare)
+                         estimate_bytes: nil, exclusive_retry: nil, retained_bytes: nil, &prepare)
             raise ArgumentError, "prepare block is required" unless prepare
 
             @items = items.to_a.freeze
@@ -103,6 +173,8 @@ module Cohere
 
             @prepare = prepare
             @estimate_bytes = estimate_bytes
+            @exclusive_retry = exclusive_retry
+            @retained_bytes = retained_bytes
             @wait_seconds = 0.0
             @enabled = enabled && @items.length > 1
             @worker_limit = worker_limit.nil? ? nil : Integer(worker_limit)
@@ -131,14 +203,24 @@ module Cohere
 
             groups = build_groups
             pool = WorkerPool.new(@effective_workers, &@prepare)
-            pending = submit(groups.first, 0, pool)
-            groups.each_with_index do |group, group_index|
+            group = groups.next
+            group_index = 0
+            pending = submit(group, group_index, pool)
+            loop do
               prepared = resolve(pending)
               pending = nil
-              next_group = groups[group_index + 1]
-              pending = submit(next_group, group_index + 1, pool) if next_group && !group.exclusive && !next_group.exclusive
-              consume(prepared, &block)
-              pending ||= submit(next_group, group_index + 1, pool) if next_group
+              next_group = next_group(groups)
+              if retry_exclusively?(group, prepared)
+                consume_with_exclusive_retry(group, prepared, pool, &block)
+              else
+                pending = submit(next_group, group_index + 1, pool) if next_group && !group.exclusive && !next_group.exclusive
+                consume(prepared, &block)
+              end
+              break unless next_group
+
+              pending ||= submit(next_group, group_index + 1, pool)
+              group = next_group
+              group_index += 1
             end
             completed = true
           ensure
@@ -176,7 +258,7 @@ module Cohere
           def submit(group, group_index, pool)
             Thread.new do
               Thread.current.name = "cohere-audio-group-#{group_index}" if Thread.current.respond_to?(:name=)
-              prepare_group(group, pool)
+              pool.prepare(group)
             end.tap { |thread| thread.report_on_exception = false }
           end
 
@@ -187,36 +269,119 @@ module Cohere
             @wait_seconds += monotonic - started if started
           end
 
-          def prepare_group(group, pool)
-            pool.prepare(group)
+          def next_group(groups)
+            groups.next
+          rescue StopIteration
+            nil
           end
 
-          def build_groups
-            return equal_groups unless @estimate_bytes
-
-            groups = []
-            items = []
-            estimates = []
-            @items.each do |item|
-              estimate = estimated_reservation(item)
-              if estimate.nil? || estimate > @group_byte_limit
-                groups << bounded_group(items, estimates) unless items.empty?
-                groups << Group.new(items: [item].freeze, limits: [@memory_byte_limit].freeze, exclusive: true)
-                items = []
-                estimates = []
+          def retry_exclusively?(group, prepared)
+            return false if group.exclusive || !@exclusive_retry
+
+            prepared.any? { |entry| @exclusive_retry.call(entry) }
+          end
+
+          def consume_with_exclusive_retry(group, prepared, pool, &block)
+            prepared.each_index do |index|
+              entry = prepared.fetch(index)
+              if entry.nil?
+                consume_single_retry(group.items.fetch(index), @memory_byte_limit, pool, &block)
                 next
               end
+              unless @exclusive_retry.call(entry)
+                begin
+                  block.call(entry)
+                ensure
+                  prepared[index] = nil
+                end
+                next
+              end
+
+              prepared[index] = nil
+              failed_limit = group.limits.fetch(index)
+              if failed_limit >= @memory_byte_limit
+                block.call(entry)
+                next
+              end
+              unless @retained_bytes
+                release_retained_entries(prepared)
+                collect_released_audio
+              end
+              limit = available_retry_limit(prepared)
+              retried = prepare_single(group.items.fetch(index), limit, pool) if limit > failed_limit
+              full_retry = retried.nil? ||
+                           (limit < @memory_byte_limit && @exclusive_retry.call(retried.fetch(0)))
+              if full_retry
+                retried&.clear
+                release_retained_entries(prepared)
+                collect_released_audio
+                retried = prepare_single(group.items.fetch(index), @memory_byte_limit, pool)
+              end
+              consume_single_result(retried, &block)
+            end
+          ensure
+            release(prepared)
+          end
+
+          def consume_single_retry(item, limit, pool, &)
+            consume_single_result(prepare_single(item, limit, pool), &)
+          end
+
+          def prepare_single(item, limit, pool)
+            pool.prepare(Group.exclusive(item, limit))
+          end
+
+          def consume_single_result(retained)
+            yield retained.fetch(0)
+          ensure
+            retained&.clear
+            collect_released_audio
+          end
+
+          def available_retry_limit(prepared)
+            return @memory_byte_limit unless @retained_bytes
+
+            retained = prepared.compact.sum { |entry| retained_bytes(entry) }
+            [@memory_byte_limit - retained, 0].max
+          end
+
+          def release_retained_entries(prepared)
+            unless @retained_bytes
+              prepared.each_index do |index|
+                next unless prepared[index]
 
-              if items.length >= @effective_workers || estimates.sum + estimate > @group_byte_limit
-                groups << bounded_group(items, estimates)
-                items = []
-                estimates = []
+                prepared[index] = nil
               end
-              items << item
-              estimates << estimate
+              return
             end
-            groups << bounded_group(items, estimates) unless items.empty?
-            groups
+
+            prepared.each_index do |index|
+              entry = prepared[index]
+              next unless entry && retained_bytes(entry).positive?
+
+              prepared[index] = nil
+            end
+          end
+
+          def retained_bytes(entry)
+            bytes = Integer(@retained_bytes.call(entry))
+            raise ArgumentError, "retained_bytes must return a non-negative integer" if bytes.negative?
+
+            bytes
+          end
+
+          def build_groups
+            return equal_groups.each unless @estimate_bytes
+
+            GroupCursor.new(
+              items: @items,
+              workers: @effective_workers,
+              group_byte_limit: @group_byte_limit,
+              memory_byte_limit: @memory_byte_limit,
+              estimated_reservation: method(:estimated_reservation),
+              unknown_reservation: unknown_reservation,
+              bounded_group: method(:bounded_group)
+            )
           end
 
           def equal_groups
@@ -238,6 +403,10 @@ module Cohere
             nil
           end
 
+          def unknown_reservation
+            [@group_byte_limit / @effective_workers, 1].max
+          end
+
           def bounded_group(items, estimates)
             limits = estimates.dup
             remaining = @group_byte_limit - limits.sum
@@ -261,7 +430,15 @@ module Cohere
               end
             end
           ensure
+            release(prepared)
+          end
+
+          def release(prepared)
             prepared&.clear
+            collect_released_audio
+          end
+
+          def collect_released_audio
             GC.start(full_mark: false, immediate_mark: true, immediate_sweep: true)
           end
 

data/lib/cohere/transcribe/runtime/resources.rb

@@ -3,6 +3,8 @@
 require "monitor"
 require "weakref"
 
+require_relative "../internal/session_ownership"
+
 module Cohere
   module Transcribe
     module Runtime
@@ -12,47 +14,6 @@ module Cohere
       class ModelResources
         OWNER_GUARD = Monitor.new
 
-        # The global owner record and the ObjectSpace finalizer retain this
-        # state, never the ModelResources instance itself. It also lets a new
-        # owner synchronously evict a collected predecessor before loading.
-        class SessionOwnership
-          def initialize
-            @mutex = Mutex.new
-            @session = nil
-          end
-
-          def session
-            @mutex.synchronize { @session }
-          end
-
-          def install(session)
-            @mutex.synchronize do
-              raise TranscriptionRuntimeError, "ASR session ownership is already installed" if @session
-
-              @session = session
-            end
-          end
-
-          def close
-            session = @mutex.synchronize do
-              current = @session
-              @session = nil
-              current
-            end
-            return unless session
-
-            session.close
-            nil
-          end
-
-          def finalize
-            close
-          rescue Exception # rubocop:disable Lint/RescueException -- finalizers must not escape during GC or shutdown
-            nil
-          end
-        end
-        private_constant :SessionOwnership
-
         class << self
           def evict_current_asr_owner
             OWNER_GUARD.synchronize do
@@ -118,7 +79,9 @@ module Cohere
 
         def initialize
           @asr_key = nil
-          @asr_ownership = SessionOwnership.new
+          @asr_ownership = Internal::SessionOwnership.new(
+            installed_error: "ASR session ownership is already installed"
+          )
           @batch_controller = nil
           @closed = false
           ObjectSpace.define_finalizer(
@@ -129,8 +92,8 @@ module Cohere
 
         # Returns [session, loaded]. A key change evicts this instance's prior
         # session, while acquisition also evicts a different process owner.
-        def acquire_asr(key)
-          raise ArgumentError, "an ASR loader block is required" unless block_given?
+        def acquire_asr(key, &loader)
+          raise ArgumentError, "an ASR loader block is required" unless loader
 
           self.class::OWNER_GUARD.synchronize do
             ensure_open!
@@ -139,36 +102,39 @@ module Cohere
 
             owner = self.class.send(:current_owner_locked)
             owner.send(:evict_asr_locked) if owner && !owner.equal?(self)
-            self.class.send(:claim_locked, self)
 
             loaded = false
-            unless @asr_ownership.session
-              installed = false
+            if @asr_ownership.session
+              self.class.send(:claim_locked, self)
+            else
+              completed = false
               session = nil
               begin
-                session = yield
-                raise TranscriptionRuntimeError, "ASR loader returned no native session" if session.nil?
+                Thread.handle_interrupt(Object => :never) do
+                  self.class.send(:claim_locked, self)
+                  session = Thread.handle_interrupt(Object => :on_blocking, &loader)
+                  raise TranscriptionRuntimeError, "ASR loader returned no native session" if session.nil?
 
-                Thread.handle_interrupt(Exception => :never) do
                   @asr_ownership.install(session)
                   @asr_key = owned_key
                   @batch_controller = nil
                   loaded = true
-                  installed = true
                 end
-              rescue Exception # rubocop:disable Lint/RescueException -- roll back asynchronous loader interruption
-                begin
-                  if @asr_ownership.session.equal?(session)
-                    evict_asr_locked
-                  else
-                    session&.close
+                completed = true
+              ensure
+                unless completed
+                  Thread.handle_interrupt(Object => :never) do
+                    if @asr_ownership.session.equal?(session)
+                      evict_asr_locked
+                    else
+                      session&.close
+                    end
+                  rescue Exception # rubocop:disable Lint/RescueException -- preserve loader termination
+                    nil
+                  ensure
+                    self.class.send(:release_locked, self)
                   end
-                rescue Exception # rubocop:disable Lint/RescueException -- preserve the loader failure
-                  nil
                 end
-                raise
-              ensure
-                self.class.send(:release_locked, self) unless installed
               end
             end
             [@asr_ownership.session, loaded].freeze

data/lib/cohere/transcribe/state/checkpoint.rb

@@ -32,7 +32,7 @@ module Cohere
 
       def write_asr_checkpoint(path:, result:, source_snapshot:, asr_contract_key:,
                                speech_spans:, vad_provider_options: nil, generation_id: nil,
-                               directory_binding: nil, guard_bindings: nil)
+                               directory_binding: nil, guard_bindings: nil, lock: nil)
         generation_id = generation_id.to_s
         generation_id = SecureRandom.hex(16) if generation_id.empty?
         payload = asr_checkpoint_payload(
@@ -48,7 +48,8 @@ module Cohere
           payload,
           source_snapshot: source_snapshot,
           directory_binding: directory_binding,
-          guard_bindings: guard_bindings
+          guard_bindings: guard_bindings,
+          commit_guard: lock&.method(:verify!)
         )
         generation_id.freeze
       end