Simplify vectorization task to sequential processing
Drop the threaded worker pool (queue, mutexes, connection pooling). The task now processes packages one by one via find_each - plenty fast since the external embedding service dominates latency, and much easier to reason about.
This commit is contained in:
parent
3550a4e2e2
commit
2b79470feb
1 changed files with 15 additions and 52 deletions
|
|
@ -5,21 +5,12 @@
|
||||||
# rubocop:disable Metrics/BlockLength
|
# rubocop:disable Metrics/BlockLength
|
||||||
namespace :debshots do
|
namespace :debshots do
|
||||||
desc 'Compute embedding vectors for all packages based on their description ' \
|
desc 'Compute embedding vectors for all packages based on their description ' \
|
||||||
'and long description (missing ones only unless FORCE=1). ' \
|
'and long description (missing ones only unless FORCE=1)'
|
||||||
'Set CONCURRENCY=n to control parallel HTTP requests.'
|
|
||||||
task compute_vectors: :environment do
|
task compute_vectors: :environment do
|
||||||
logger = Logger.new($stdout)
|
logger = Logger.new($stdout)
|
||||||
logger.level = Logger::INFO
|
logger.level = Logger::INFO
|
||||||
|
|
||||||
force = ENV['FORCE'].present?
|
force = ENV['FORCE'].present?
|
||||||
concurrency = [ENV.fetch('CONCURRENCY', '4').to_i, 1].max
|
|
||||||
|
|
||||||
# Never use more threads than the connection pool has connections.
|
|
||||||
# Keep one connection reserved for the main thread.
|
|
||||||
pool_size = ActiveRecord::Base.connection_pool.size
|
|
||||||
concurrency = [concurrency, pool_size - 1].min
|
|
||||||
concurrency = 0 if concurrency.negative?
|
|
||||||
|
|
||||||
scope = force ? Package.all : Package.where(embedding: nil)
|
scope = force ? Package.all : Package.where(embedding: nil)
|
||||||
total = scope.count
|
total = scope.count
|
||||||
|
|
||||||
|
|
@ -30,61 +21,33 @@ namespace :debshots do
|
||||||
|
|
||||||
logger.info "Computing vectors for #{total} packages " \
|
logger.info "Computing vectors for #{total} packages " \
|
||||||
"(#{force ? 'recomputing' : 'missing only'}, " \
|
"(#{force ? 'recomputing' : 'missing only'}, " \
|
||||||
"service: #{Vectorizer.service_url}, " \
|
"service: #{Vectorizer.service_url})"
|
||||||
"concurrency: #{concurrency.zero? ? 1 : concurrency})"
|
|
||||||
|
|
||||||
stats_mutex = Mutex.new
|
done = 0
|
||||||
stats = { done: 0, failed: 0 }
|
failed = 0
|
||||||
|
|
||||||
process_package = lambda do |package|
|
scope.find_each do |package|
|
||||||
begin
|
begin
|
||||||
raise ArgumentError, 'no description text' if package.embedding_text.blank?
|
raise ArgumentError, 'no description text' if package.embedding_text.blank?
|
||||||
|
|
||||||
package.update_embedding!
|
package.update_embedding!
|
||||||
stats_mutex.synchronize { stats[:done] += 1 }
|
done += 1
|
||||||
rescue StandardError => e
|
rescue StandardError => e
|
||||||
# A single broken package must never take down a whole worker
|
# A single broken package must never stop the whole run - log
|
||||||
# thread - log enough detail to diagnose it and move on.
|
# enough detail to diagnose it and move on.
|
||||||
stats_mutex.synchronize { stats[:failed] += 1 }
|
failed += 1
|
||||||
logger.error "Failed for package #{package.name}: #{e.class}: #{e.message}"
|
logger.error "Failed for package #{package.name}: #{e.class}: #{e.message}"
|
||||||
logger.error e.backtrace.first(3).join("\n") if e.backtrace
|
logger.error e.backtrace.first(3).join("\n") if e.backtrace
|
||||||
end
|
end
|
||||||
|
|
||||||
stats_mutex.synchronize do
|
processed = done + failed
|
||||||
processed = stats[:done] + stats[:failed]
|
next unless (processed % 100).zero? || processed == total
|
||||||
return unless (processed % 100).zero?
|
|
||||||
|
|
||||||
percent = (processed.to_f / total * 100).round(1)
|
percent = (processed.to_f / total * 100).round(1)
|
||||||
logger.info "Progress: #{processed}/#{total} (#{percent}%)"
|
logger.info "Progress: #{processed}/#{total} (#{percent}%)"
|
||||||
end
|
end
|
||||||
end
|
|
||||||
|
|
||||||
if concurrency.zero?
|
logger.info "Done. #{done} vectors computed, #{failed} failures."
|
||||||
# Sequential fallback when the connection pool has no room for workers
|
|
||||||
scope.find_each { |package| process_package.call(package) }
|
|
||||||
else
|
|
||||||
# Feed package IDs into a queue that the worker threads pull from.
|
|
||||||
work_queue = Queue.new
|
|
||||||
scope.pluck(:id).each { |id| work_queue << id }
|
|
||||||
concurrency.times { work_queue << nil } # one poison pill per worker
|
|
||||||
|
|
||||||
workers = Array.new(concurrency) do
|
|
||||||
Thread.new do
|
|
||||||
# Each thread gets its own database connection from the pool.
|
|
||||||
ActiveRecord::Base.connection_pool.with_connection do
|
|
||||||
while (package_id = work_queue.pop)
|
|
||||||
package = Package.find_by(id: package_id)
|
|
||||||
process_package.call(package) if package
|
|
||||||
end
|
|
||||||
end
|
|
||||||
rescue StandardError => e
|
|
||||||
logger.error "Worker crashed: #{e.message}"
|
|
||||||
end
|
|
||||||
end
|
|
||||||
|
|
||||||
workers.each(&:join)
|
|
||||||
end
|
|
||||||
|
|
||||||
logger.info "Done. #{stats[:done]} vectors computed, #{stats[:failed]} failures."
|
|
||||||
end
|
end
|
||||||
end
|
end
|
||||||
|
# rubocop:enable Metrics/BlockLength
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue