Class: Cohere::Transcribe::ASR::BatchController

Inherits:
Object
  • Object
show all
Defined in:
lib/cohere/transcribe/asr/batching.rb

Overview

Persistent batch-size controller. OOM learning and fatal circuit state live with the retained native session, not with one audio file.

Constant Summary collapse

DEVICE_DEFAULTS = 
{ "cuda" => 24, "mps" => 8, "cpu" => 4 }.freeze

Instance Attribute Summary collapse

Class Method Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(current_size:, max_size:, initial_size:, audio_budget_seconds:, adaptive:, target_memory_ratio:, total_memory_bytes: 0, memory_budget_bytes: 0) ⇒ BatchController

Returns a new instance of BatchController.

Raises:

  • (ArgumentError) 


158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
 
# File 'lib/cohere/transcribe/asr/batching.rb', line 158

def initialize(current_size:, max_size:, initial_size:, audio_budget_seconds:, adaptive:,
               target_memory_ratio:, total_memory_bytes: 0, memory_budget_bytes: 0)
  @current_size = positive_integer(current_size, "current_size")
  @max_size = positive_integer(max_size, "max_size")
  @initial_size = positive_integer(initial_size, "initial_size")
  raise ArgumentError, "max_size must be at least current_size" if @max_size < @current_size

  @audio_budget_seconds = Float(audio_budget_seconds)
  raise ArgumentError, "audio_budget_seconds must be positive" unless @audio_budget_seconds.positive?

  raise ArgumentError, "adaptive must be a boolean" unless [true, false].include?(adaptive)

  @adaptive = adaptive
  @target_memory_ratio = Float(target_memory_ratio)
  @total_memory_bytes = Integer(total_memory_bytes)
  @memory_budget_bytes = Integer(memory_budget_bytes)
  @growth_cooldown = 0
  @circuit_breaker_error = nil
end

Instance Attribute Details

#adaptive ⇒ Object (readonly)

Returns the value of attribute adaptive.



98
99
100
 
# File 'lib/cohere/transcribe/asr/batching.rb', line 98

def adaptive
  @adaptive
end

#audio_budget_seconds ⇒ Object (readonly)

Returns the value of attribute audio_budget_seconds.



98
99
100
 
# File 'lib/cohere/transcribe/asr/batching.rb', line 98

def audio_budget_seconds
  @audio_budget_seconds
end

#circuit_breaker_error ⇒ Object

Returns the value of attribute circuit_breaker_error.



105
106
107
 
# File 'lib/cohere/transcribe/asr/batching.rb', line 105

def circuit_breaker_error
  @circuit_breaker_error
end

#current_size ⇒ Object

Returns the value of attribute current_size.



105
106
107
 
# File 'lib/cohere/transcribe/asr/batching.rb', line 105

def current_size
  @current_size
end

#growth_cooldown ⇒ Object

Returns the value of attribute growth_cooldown.



105
106
107
 
# File 'lib/cohere/transcribe/asr/batching.rb', line 105

def growth_cooldown
  @growth_cooldown
end

#initial_size ⇒ Object (readonly)

Returns the value of attribute initial_size.



98
99
100
 
# File 'lib/cohere/transcribe/asr/batching.rb', line 98

def initial_size
  @initial_size
end

#max_size ⇒ Object (readonly)

Returns the value of attribute max_size.



98
99
100
 
# File 'lib/cohere/transcribe/asr/batching.rb', line 98

def max_size
  @max_size
end

#memory_budget_bytes ⇒ Object (readonly)

Returns the value of attribute memory_budget_bytes.



98
99
100
 
# File 'lib/cohere/transcribe/asr/batching.rb', line 98

def memory_budget_bytes
  @memory_budget_bytes
end

#target_memory_ratio ⇒ Object (readonly)

Returns the value of attribute target_memory_ratio.



98
99
100
 
# File 'lib/cohere/transcribe/asr/batching.rb', line 98

def target_memory_ratio
  @target_memory_ratio
end

#total_memory_bytes ⇒ Object (readonly)

Returns the value of attribute total_memory_bytes.



98
99
100
 
# File 'lib/cohere/transcribe/asr/batching.rb', line 98

def total_memory_bytes
  @total_memory_bytes
end

Class Method Details

.create(options, device:, durations:, memory: nil, physical_max: nil) ⇒ Object 



107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
 
# File 'lib/cohere/transcribe/asr/batching.rb', line 107

def self.create(options, device:, durations:, memory: nil, physical_max: nil)
  default_initial = DEVICE_DEFAULTS.fetch(device.to_s, DEVICE_DEFAULTS.fetch("cpu"))
  requested_initial = options.batch_size || if options.batch_max_size
                                              [default_initial, options.batch_max_size].min
                                            else
                                              default_initial
                                            end
  initial = physical_max ? [requested_initial, physical_max].min : requested_initial
  free_bytes, total_bytes = memory || [0, 0]
  free_bytes = Integer(free_bytes)
  total_bytes = Integer(total_bytes)
  requested_maximum = if !options.adaptive_batch
                        requested_initial
                      elsif options.batch_max_size
                        options.batch_max_size
                      elsif options.batch_size
                        options.batch_size
                      elsif device.to_s == "cuda"
                        total_gib = total_bytes.fdiv(1024**3)
                        [requested_initial, [128, (total_gib * 4).to_i].min].max
                      else
                        default_initial
                      end
  maximum = [requested_initial, requested_maximum].max
  maximum = [maximum, physical_max].min if physical_max
  maximum = [initial, maximum].max
  longest = durations.map { |duration| Float(duration) }.max || 1.0
  audio_budget = options.batch_audio_seconds || (initial * [longest, 0.25].max)

  used_bytes = [total_bytes - free_bytes, 0].max
  memory_budget = if device.to_s == "cuda" && total_bytes.positive?
                    [
                      (options.batch_vram_target * total_bytes).to_i,
                      used_bytes + (0.95 * free_bytes).to_i
                    ].min
                  else
                    0
                  end

  new(
    current_size: initial,
    max_size: maximum,
    initial_size: initial,
    audio_budget_seconds: audio_budget,
    adaptive: options.adaptive_batch,
    target_memory_ratio: options.batch_vram_target,
    total_memory_bytes: total_bytes,
    memory_budget_bytes: memory_budget
  )
end

Instance Method Details

#circuit_open? ⇒ Boolean

Returns:

  • (Boolean) 


234
235
236
 
# File 'lib/cohere/transcribe/asr/batching.rb', line 234

def circuit_open?
  !@circuit_breaker_error.nil?
end

#configure_group(options, durations) ⇒ Object 



178
179
180
181
182
 
# File 'lib/cohere/transcribe/asr/batching.rb', line 178

def configure_group(options, durations)
  longest = durations.map { |duration| Float(duration) }.max || 1.0
  @audio_budget_seconds = options.batch_audio_seconds || (@initial_size * [longest, 0.25].max)
  self
end

#open_circuit(error) ⇒ Object 



228
229
230
231
232
 
# File 'lib/cohere/transcribe/asr/batching.rb', line 228

def open_circuit(error)
  return @circuit_breaker_error if @circuit_breaker_error

  @circuit_breaker_error = FailurePolicy.fingerprint(error)
end

#record_oom(attempted_rows) ⇒ Object 



192
193
194
195
196
197
198
199
200
201
202
203
 
# File 'lib/cohere/transcribe/asr/batching.rb', line 192

def record_oom(attempted_rows)
  rows = positive_integer(attempted_rows, "attempted_rows")
  @growth_cooldown = [@growth_cooldown, 2].max
  if rows <= 1
    @current_size = 1
    return self
  end

  @max_size = [@max_size, rows - 1].min
  @current_size = [1, [@max_size, rows / 2].min].max
  self
end

#record_success(attempted_rows, baseline_reserved_bytes: 0, peak_reserved_bytes: 0) ⇒ Object 



205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
 
# File 'lib/cohere/transcribe/asr/batching.rb', line 205

def record_success(attempted_rows, baseline_reserved_bytes: 0, peak_reserved_bytes: 0)
  rows = positive_integer(attempted_rows, "attempted_rows")
  if @growth_cooldown.positive?
    @growth_cooldown -= 1
    return self
  end
  return self unless @adaptive && @current_size < @max_size && rows >= @current_size
  return self unless @total_memory_bytes.positive? && peak_reserved_bytes.positive?

  budget = @memory_budget_bytes.positive? ? @memory_budget_bytes : (@target_memory_ratio * @total_memory_bytes).to_i
  headroom = (budget - peak_reserved_bytes).fdiv(@total_memory_bytes)
  return self if headroom <= 0.05

  factor = headroom >= 0.20 ? 1.25 : 1.125
  proposed = [@current_size + 1, (@current_size * factor).ceil].max
  incremental = [1, peak_reserved_bytes - baseline_reserved_bytes].max
  available = [0, budget - baseline_reserved_bytes].max
  memory_estimate = (rows * available).div(incremental)
  proposed = [proposed, [@current_size, memory_estimate].max].min if memory_estimate.positive?
  @current_size = [@max_size, proposed].min
  self
end

#take_count(rows) ⇒ Object 



184
185
186
187
188
189
190
 
# File 'lib/cohere/transcribe/asr/batching.rb', line 184

def take_count(rows)
  return 0 if rows.empty?

  longest = [yield(rows.first), 1.0 / SAMPLE_RATE].max
  frame_limited = [(@audio_budget_seconds / longest).floor, 1].max
  [rows.length, @current_size, frame_limited].min
end