emb 0.3.0 → 0.4.0.pre3

data/lib/emb/batch.rb

@@ -1,69 +1,123 @@
 # frozen_string_literal: true
 
 require 'batch_loader'
+require 'redis_client'
+require_relative 'batch_dispatch'
+require_relative 'values_batch'
 
 module Emb
+  extend BatchDispatch
+
   BATCH_KEY = :emb
 
+  # Raised by resolve_slice when the server reply carries fewer entries than
+  # the slice's texts. Subclasses RedisClient::ProtocolError so the existing
+  # fail-closed rescue/wrap paths treat it as a client error, but stays
+  # distinct so transient_error? does not count it as a transport retry: the
+  # reply shape was wrong on the single send — nothing was re-sent.
+  ShortReplyError = Class.new(RedisClient::ProtocolError)
+
   BATCH_BLOCK = lambda do |items, loader, _args|
     items.group_by(&:first).each do |client, client_items|
       chunk = client.respond_to?(:batch_size) && client.batch_size ? client.batch_size : Emb.configuration.batch_size
 
-      # Chunk at batch_size pairs per EMB.MULTI so a single command stays well
-      # under the server's max_pairs cap and within client read timeouts.
-      client_items.each_slice(chunk) do |slice|
-        pairs = slice.flat_map { |_, model, text| Array(text).flat_map { |t| [model.to_s, t] } }
-        results = Array(client.send_command('EMB.MULTI', *pairs))
-
-        offset = 0
-        slice.each do |item|
-          _, _, text = item
-          texts = Array(text)
-          values = results[offset, texts.size].map { |entry| entry&.unpack('e*') }
-          offset += texts.size
-
-          # eager Proxy#[] shape: vector for a single text, vectors for many
-          loader.call(item, values.size == 1 ? values.first : values)
+      slices = pack_slices(client_items, chunk)
+      if client.respond_to?(:parallel_batch?) && client.parallel_batch? && slices.size > 1
+        dispatch_parallel(client, slices, loader)
+      else
+        # Serial dispatch: one share in flight at a time. Redis errors fail
+        # closed with context; anything else is a local bug and is re-raised
+        # unchanged after the pending set is dropped.
+        slices.each do |slice|
+          resolve_slice(loader, slice, dispatch_slice(client, slice))
+        rescue RedisClient::Error => e
+          fail_batch!(e, slice: slice, budget: retry_budget(client))
+        rescue StandardError
+          clear_batch_pending!
+          raise
         end
       end
     end
   end
 
   class << self
-    def build_batch_loader(client, model, text)
-      BatchLoader.for([client, model, text]).batch(key: BATCH_KEY, &BATCH_BLOCK)
-    end
-  end
+    # BatchDispatch mechanics (pack_slices, dispatch_parallel, resolve_slice,
+    # ...) are internal to BATCH_BLOCK; nothing outside Emb calls them with an
+    # explicit receiver.
+    private(*BatchDispatch.instance_methods(false))
 
-  class BatchProxy
-    def initialize(client)
-      @client = client
-      @models = {}
+    def build_batch_loader(client, model, text, format: :binary)
+      unless %i[binary values].include?(format)
+        raise ArgumentError, "unknown format #{format.inspect} (expected :binary or :values)"
+      end
+
+      # default_value []: an item whose batch failed (fail-closed) resolves to
+      # an empty vector collection instead of nil, so resolver methods like
+      # `loader.sum` do not blow up with NoMethodError-on-nil.
+      # Items carry [client, model, text, format] so the dispatch knows how to
+      # shape and parse the reply.
+      BatchLoader.for([client, model, text, format]).batch(default_value: [], key: BATCH_KEY, &BATCH_BLOCK)
     end
 
-    def [](name)
-      @models[name.to_sym] ||= BatchModelProxy.new(@client, name.to_sym)
+    # Removes every pending item of the batch scope. batch-loader prunes
+    # pending items only after a successful batch block, so failed batches
+    # would otherwise stay queued: retries would re-run the whole batch and
+    # stale items would be re-sent by later batches in the same scope. Guarded
+    # for a nil executor (no batch scope).
+    def clear_batch_pending!
+      key = [BATCH_BLOCK.source_location, BATCH_KEY]
+      BatchLoader::Executor.current&.items_by_block&.delete(key)
     end
 
-    def inspect
-      "#<Emb::BatchProxy client=#{@client.inspect}>"
+    # Fail-closed tail for a failed batch: clear the pending set, then raise
+    # Emb::ServerError carrying the cause and the models/texts/attempts
+    # context. The raise happens inside a rescue of the original error so Ruby
+    # attaches it as `cause` regardless of whether this runs while a rescue is
+    # active (serial path) or from the forcing thread after parallel workers
+    # captured the error (parallel path). `attempts` counts the error's retry
+    # class: connection/protocol errors (the ones redis-client actually
+    # re-sends) report `budget + 1`; operation errors and read timeouts (never
+    # re-sent — a timeout may already have executed server work) report 1. A
+    # pre-send connection refusal is retried across instances by the
+    # connection router before it can reach here as a terminal error.
+    def fail_batch!(error, slice:, budget:)
+      clear_batch_pending!
+      attempts = transient_error?(error) ? budget + 1 : 1
+      models = slice.map { |_, model, _| model }.uniq.join(', ')
+      texts = slice.sum { |_, _, text| Array(text).size }
+      message = "batch failed after #{attempts} attempt(s) " \
+                "(models: #{models}, #{texts} text(s)) #{error.class}: #{error.message}"
+      begin
+        raise error
+      rescue StandardError
+        raise ServerError.new(message, attempts: attempts)
+      end
     end
-  end
 
-  class BatchModelProxy
-    attr_reader :name
+    # How many additional re-sends redis-client performs for transient
+    # failures: the per-client option when set, else the global default.
+    # Normalized to an Integer (redis-client also accepts a delay Array, whose
+    # truthy slots each grant one retry).
+    def retry_budget(client)
+      value = client.reconnect_attempts if client.respond_to?(:reconnect_attempts)
+      value = Emb.configuration.reconnect_attempts if value.nil?
+      return value if value.is_a?(Integer)
 
-    def initialize(client, name)
-      @client = client
-      @name = name
+      value.is_a?(Array) ? value.count(&:itself) : 0
     end
 
-    def [](text, *texts)
-      Emb.build_batch_loader(@client, @name, texts.empty? ? text : [text, *texts])
+    # The error classes redis-client actually re-dispatches under
+    # reconnect_attempts: ConnectionError (connect/transport breaks) and
+    # ProtocolError. ReadTimeoutError is intercepted before the retry loop
+    # (a timed-out command may already have executed server work), so it is
+    # terminal and counts a single attempt. Emb::ShortReplyError is raised
+    # locally by resolve_slice (reply shape wrong on the single send), so it
+    # is likewise terminal and counts one attempt.
+    def transient_error?(error)
+      (error.is_a?(RedisClient::ConnectionError) && !error.is_a?(RedisClient::ReadTimeoutError)) ||
+        (error.is_a?(RedisClient::ProtocolError) && !error.is_a?(ShortReplyError))
     end
 
-    def inspect
-      "#<Emb::BatchModelProxy #{@name}>"
-    end
+    private :clear_batch_pending!, :fail_batch!, :retry_budget, :transient_error?
   end
 end

data/lib/emb/batch_dispatch.rb

@@ -0,0 +1,171 @@
+# frozen_string_literal: true
+
+module Emb
+  # Share-dispatch mechanics for the deferred batch path: EMB/MULTI wire
+  # shaping, per-slice result mapping, and the bounded concurrent fan-out used
+  # by `lazy: :batch`. Extended into Emb by batch.rb.
+  module BatchDispatch
+    # Sends one chunk share and returns the raw reply entries. Single-model
+    # shares use plain `EMB <model> <text>...` (one inference, model once);
+    # mixed-model shares keep EMB.MULTI per-pair nil semantics. Errors after
+    # the command may have been sent are terminal and propagate so the forcing
+    # thread can fail closed.
+    def dispatch_slice(client, slice)
+      models = slice.map { |_, model, _, _| model }.uniq
+      args = models.size == 1 ? same_model_args(slice, models.first) : mixed_model_args(slice)
+      Array(client.send_command(*args))
+    end
+
+    def same_model_args(slice, model)
+      texts = slice.flat_map { |_, _, text, _| Array(text) }
+      if slice_format(slice) == :values
+        ['EMB', model.to_s, 'VALUES', *texts]
+      else
+        ['EMB', model.to_s, *texts]
+      end
+    end
+
+    def mixed_model_args(slice)
+      pairs = slice.flat_map { |_, model, text, _| Array(text).flat_map { |t| [model.to_s, t] } }
+      if slice_format(slice) == :values
+        ['EMB.MULTI', 'VALUES', *pairs]
+      else
+        ['EMB.MULTI', *pairs]
+      end
+    end
+
+    # slice_format: the batch item layout is [client, model, text, format];
+    # items created before the format slot existed default to :binary.
+    def slice_format(slice)
+      formats = slice.map { |item| item[3] }.uniq
+      if formats.size > 1
+        raise ArgumentError, "cannot mix :binary and :values formats in one batch"
+      end
+      formats.first || :binary
+    end
+
+    # Maps a slice's reply entries onto its items in deferral order. Runs on
+    # the forcing thread (batch-loader's executor is per-thread), so workers
+    # never resolve loaders. A reply shorter than the slice's texts is a
+    # protocol violation: fail the batch via Emb::ShortReplyError (distinct
+    # from a client-raised ProtocolError so it is not counted as a transport
+    # retry).
+    def resolve_slice(loader, slice, results)
+      if slice_format(slice) == :values
+        return ValuesBatch.resolve(loader, slice, results)
+      end
+
+      expected = slice.sum { |_, _, text, _| Array(text).size }
+      unless results.size >= expected
+        raise ShortReplyError, "expected #{expected} reply entries, got #{results.size}"
+      end
+
+      offset = 0
+      slice.each do |item|
+        _, _, text, _ = item
+        texts = Array(text)
+        values = entry_values(results, offset, texts)
+        offset += texts.size
+
+        # eager Proxy#[] shape: vector for a single text, vectors for many
+        loader.call(item, values)
+      end
+    end
+
+    def entry_values(results, offset, texts)
+      values = results[offset, texts.size].map { |entry| entry&.unpack('e*') }
+      values.size == 1 ? values.first : values
+    end
+
+    # The worker captures failures as outcomes; only the forcing thread
+    # resolves loaders (see resolve_slice).
+    def dispatch_share(client, slice)
+      [:ok, slice, dispatch_slice(client, slice)]
+    rescue StandardError => e
+      [:error, slice, e]
+    end
+
+    # Dispatches all shares concurrently over at most the client's connection
+    # capacity workers, then fails closed on the first terminal error.
+    def dispatch_parallel(client, slices, loader)
+      workers = slices.size.clamp(1, worker_capacity(client))
+      queue = share_queue(slices, workers)
+      outcomes = Array.new(slices.size)
+      run_worker_threads(client, queue, outcomes, workers)
+      resolve_outcomes(client, outcomes, loader)
+    end
+
+    def share_queue(slices, workers)
+      queue = Queue.new
+      slices.each_with_index { |slice, index| queue << [index, slice] }
+      workers.times { queue << nil }
+      queue
+    end
+
+    def worker_capacity(client)
+      return Float::INFINITY unless client.respond_to?(:pools)
+
+      client.pools.sum { |pool| pool.respond_to?(:size) ? pool.size : 1 }
+    end
+
+    def run_worker_threads(client, queue, outcomes, workers)
+      workers.times.map do
+        Thread.new do
+          while (job = queue.pop)
+            index, slice = job
+            outcomes[index] = dispatch_share(client, slice)
+          end
+        end
+      end.each(&:join)
+    end
+
+    # Redis errors fail closed with context (Emb::ServerError, matching the
+    # serial path); non-redis errors are local bugs and re-raise unchanged.
+    def resolve_outcomes(client, outcomes, loader)
+      first_error, failed_slice = collect_outcomes(outcomes, loader)
+      return unless first_error
+
+      raise first_error unless first_error.is_a?(RedisClient::Error)
+
+      fail_batch!(first_error, slice: failed_slice, budget: retry_budget(client))
+    rescue StandardError
+      clear_batch_pending!
+      raise
+    end
+
+    def collect_outcomes(outcomes, loader)
+      first_error = nil
+      failed_slice = nil
+      outcomes.each do |status, slice, result|
+        if status == :ok
+          resolve_slice(loader, slice, result)
+        else
+          first_error ||= result
+          failed_slice ||= slice
+        end
+      end
+      [first_error, failed_slice]
+    end
+
+    # Packs items into shares by accumulated text count so one command stays
+    # within `chunk` texts. An item larger than the chunk goes alone; the
+    # server truncates it with null reply slots, as in the eager path. `used`
+    # carries the running text count of the current slice instead of
+    # re-summing it for every item (O(n) per item would be O(n²)).
+    def pack_slices(items, chunk)
+      slices = []
+      used = 0
+      items.each do |item|
+        size = Array(item[2]).size
+        if slices.empty? || used + size > chunk
+          slices << [item]
+          used = size
+        else
+          slices.last << item
+          used += size
+        end
+      end
+      slices
+    end
+  end
+end

data/lib/emb/batch_scope.rb

@@ -0,0 +1,20 @@
+# frozen_string_literal: true
+
+require 'batch_loader'
+
+module Emb
+  # Runs a block of work and clears the per-thread batch scope afterwards,
+  # even when the block raises. Clearing is unconditional: batch-loader scopes
+  # are thread-global and shared by every client, so a deferred-mode client
+  # (global `lazy` config or a per-call override) inside the block would
+  # otherwise leak cached values and pending loaders into the next request/job.
+  # Under the eager default (`lazy: false`) embed calls never create a scope,
+  # so the clear is a no-op — the middleware is observably inert there.
+  module BatchScope
+    def self.wrap
+      yield
+    ensure
+      BatchLoader::Executor.clear_current
+    end
+  end
+end

data/lib/emb/client.rb

@@ -1,52 +1,41 @@
 # frozen_string_literal: true
 
-require 'connection_pool'
+require_relative 'connection_router'
 require 'redis_client'
 
 module Emb
   class Client
     include Commands
 
-    attr_reader :pool, :batch_size
-
-    def initialize(pool: nil, batch: nil, **redis_options)
+    # Accepts a single redis URL (String) or several (Array of Strings) naming
+    # interchangeable emb instances serving the same model set; each url gets
+    # its own pool of `pool` connections (see ConnectionRouter).
+    def initialize(pool: nil, lazy: nil, **redis_options)
       cfg = Emb.configuration
-      @batch_enabled = batch.nil? ? cfg.batch : batch
+      @lazy_mode = lazy.nil? ? cfg.lazy : validate_lazy!(lazy)
       @batch_size = redis_options.delete(:batch_size) || cfg.batch_size
-      size = pool.nil? ? cfg.pool : pool
       url = extract_url!(redis_options, cfg)
+      # Captured before ConnectionRouter consumes the merged options, so
+      # fail-closed batches can report the retry budget (Emb::ServerError).
       redis_options = merged_redis_options(redis_options, cfg, url)
-
-      @pool = ConnectionPool.new(size: size) do
-        RedisClient.new(url: url, **redis_options)
-      end
-
+      @reconnect_attempts = redis_options.fetch(:reconnect_attempts, cfg.reconnect_attempts)
+      @router = ConnectionRouter.new(pool || cfg.pool, instance_urls(url), redis_options)
       @registry = {}
     end
 
-    def send_command(*args)
-      return @pool.with { |r| r.call(*args) } unless Emb.debug?
+    def send_command(...) = @router.call(...)
 
-      start = Process.clock_gettime(Process::CLOCK_MONOTONIC)
-      result = @pool.with { |r| r.call(*args) }
-      elapsed = (Process.clock_gettime(Process::CLOCK_MONOTONIC) - start) * 1000
+    def pools = @router.pools
 
-      $stdout.puts "[EMB] #{args.map(&:inspect).join(' ')} (#{format('%.2f', elapsed)}ms)"
-
-      result
-    end
+    attr_reader :batch_size, :lazy_mode, :reconnect_attempts
 
     def [](name)
       @registry[name] ||= Proxy.new(self, name.to_sym)
     end
 
-    def batch
-      @batch ||= BatchProxy.new(self)
-    end
+    def lazy? = @lazy_mode != false
 
-    def batch?
-      @batch_enabled
-    end
+    def parallel_batch? = @lazy_mode == :batch
 
     # Live view of the server's runtime configuration (CONFIG GET/SET).
     def config
@@ -103,9 +92,23 @@ module Emb
 
     private
 
+    def validate_lazy!(value)
+      unless Configuration::LAZY_MODES.include?(value)
+        raise ArgumentError, "lazy must be false, :multi, or :batch (got #{value.inspect})"
+      end
+
+      value
+    end
+
+    def instance_urls(url)
+      raise ArgumentError, 'url array must not be empty' if url.is_a?(Array) && url.empty?
+
+      url.nil? ? [nil] : Array(url)
+    end
+
     def merged_redis_options(opts, cfg, url)
       defaults = cfg.to_h
-      keys = defaults.keys - %i[url pool batch batch_size]
+      keys = defaults.keys - %i[url pool lazy batch_size]
       keys -= %i[host port] if url
 
       keys.each do |key|

data/lib/emb/commands.rb

@@ -1,10 +1,14 @@
 # frozen_string_literal: true
 
+require_relative 'script_reply_decode'
+
 module Emb
   # Server-status and runtime-config command wrappers. They only depend on
   # #send_command, so they live in this module (included into Emb::Client)
   # rather than growing the client class.
   module Commands
+    include ScriptReplyDecode
+
     # EMB.STATS as a Symbol-keyed Hash. No client-side type layer: values are
     # exactly what the RESP decoder returned (Integer where the server sends
     # RESP integers, String otherwise).
@@ -19,8 +23,101 @@ module Emb
       parse_info(send_command('INFO', *sections.map(&:to_s)))
     end
 
+    # Evaluate a script against a model (EMB.EVAL). texts become the script's
+    # KEYS, args its ARGV. Replies are typed per the reply grammar: hash
+    # replies (flat field/value pair arrays under RESP2) become Ruby Hashes,
+    # nested values recurse. decode: opt-in float decoding — see
+    # parse_script_reply (decode: :f32 or decode: {field => :f32}).
+    def eval(model, script, texts, args = [], decode: nil)
+      texts = Array(texts)
+      schema = normalize_decode(decode)
+      reply = send_command('EMB.EVAL', model.to_s, script, texts.size, *texts.map(&:to_s), *args.map(&:to_s))
+      parse_script_reply(reply, multi: texts.size > 1, decode: schema)
+    end
+
+    # Evaluate a previously loaded script by SHA1 (EMB.EVSHA). See #eval.
+    def evalsha(model, sha, texts, args = [], decode: nil)
+      texts = Array(texts)
+      schema = normalize_decode(decode)
+      reply = send_command('EMB.EVSHA', model.to_s, sha, texts.size, *texts.map(&:to_s), *args.map(&:to_s))
+      parse_script_reply(reply, multi: texts.size > 1, decode: schema)
+    end
+
+    # EMB.SCRIPT subcommands: a small command object so the surface reads
+    # naturally (script.load / script.exists / script.flush).
+    def script
+      @script ||= ScriptCommands.new(self)
+    end
+
+    # EMB.SCRIPT LOAD/EXISTS/FLUSH.
+    class ScriptCommands
+      def initialize(client)
+        @client = client
+      end
+
+      # Compile, cache, and return the script's SHA1 for the model.
+      def load(model, source)
+        @client.send_command('EMB.SCRIPT', 'LOAD', model.to_s, source)
+      end
+
+      # Whether each SHA is cached for the model.
+      def exists(model, *shas)
+        @client.send_command('EMB.SCRIPT', 'EXISTS', model.to_s, *shas.map(&:to_s)).map { |flag| flag == 1 }
+      end
+
+      # Clear the model's cached scripts, or all models when omitted.
+      def flush(model = nil)
+        args = ['EMB.SCRIPT', 'FLUSH']
+        args << model.to_s if model
+        @client.send_command(*args)
+      end
+    end
+
+    # Remove all cached embeddings, or only entries belonging to +model+.
+    # Returns the server's integer removed-entry count.
+    def cache_flush(model = nil)
+      args = ['EMB.CACHE.FLUSH']
+      args << model.to_s unless model.nil?
+      send_command(*args)
+    end
+
+    # Ask the server to start an asynchronous cache snapshot. Completion and
+    # failures are observable through #stats or #server_info(:cache).
+    def save_cache
+      send_command('EMB.SAVE')
+    end
+
     private
 
+    # Converts a scripted reply (single value, or an array of per-text values
+    # for multi-text calls) through the hash grammar: a flat field/value pair
+    # array becomes a Hash, recursively. Pure even-length string lists are
+    # indistinguishable from two-field hashes on the RESP2 wire, so scripts
+    # that must return them should nest them (wrap the list in a table).
+    #
+    # decode is the normalized (see normalize_decode) opt-in float decoding:
+    #   nil                        → no decoding, parsing exactly as before
+    #   :f32                       → each value position is unpack('e*')'d when
+    #                                it is a packed float bulk; numeric arrays
+    #                                pass through unchanged (element types kept)
+    #   {field => :f32, ...}       → after hash parsing, the named field(s) of
+    #                                each hash reply decode as :f32 (fields
+    #                                absent from a reply are left untouched)
+    def parse_script_reply(reply, multi:, decode: nil)
+      parsed = (multi ? reply : [reply]).map { |v| script_hash(v) }
+      decoded = apply_decode(parsed, decode)
+      multi ? decoded : decoded.first
+    end
+
+    def script_hash(value)
+      return value unless value.is_a?(Array) && value.size.even?
+
+      pairs = value.each_slice(2).to_a
+      return value unless pairs.all? { |field, _| field.is_a?(String) }
+
+      pairs.to_h { |field, val| [field, script_hash(val)] }
+    end
+
     # Parse Redis INFO section text into a nested Hash:
     #   {Server: {redis_version: "0.2.4", uptime_secs: "7"}, Cache: {…}, …}
     # Section names and keys are Symbols; values pass through as the server

data/lib/emb/configuration.rb

@@ -3,25 +3,51 @@
 module Emb
   class Configuration
     OPTIONS = %i[
-      host port url pool batch batch_size driver protocol
+      host port url pool lazy batch_size driver protocol
       connect_timeout read_timeout write_timeout reconnect_attempts
     ].freeze
 
+    # Execution modes for embed calls. false = eager (default, one EMB round
+    # trip per call); :multi = defer and coalesce into EMB.MULTI, serial;
+    # :batch = defer and execute chunk shares concurrently. Mutually exclusive
+    # by construction.
+    LAZY_MODES = [false, :multi, :batch].freeze
+
     attr_accessor(*OPTIONS)
 
+    def lazy=(value)
+      unless LAZY_MODES.include?(value)
+        raise ArgumentError, "lazy must be false, :multi, or :batch (got #{value.inspect})"
+      end
+
+      @lazy = value
+    end
+
     def initialize
       self.host = 'localhost'
       self.port = 6379
       self.url = nil
       self.pool = 5
-      self.batch = true
+      self.lazy = false
       self.batch_size = 512
       self.driver = nil
       self.protocol = 2
       self.connect_timeout = nil
-      self.read_timeout = nil
-      self.write_timeout = nil
-      self.reconnect_attempts = 3
+      # Read/write timeouts are explicit, NOT nil: nil forwards nothing and
+      # redis-client silently applies its 1.0s default, which makes 512-pair
+      # EMB.MULTI batches fail under load. 10s covers worst-case inference on
+      # shared CPUs; scale up if you raise batch_size.
+      self.read_timeout = 10
+      self.write_timeout = 10
+      # 0 = default: a failing batch fails closed after one attempt and raises
+      # Emb::ServerError. Set > 0 to opt into bounded re-sends: redis-client
+      # retries connection/protocol failures (never read timeouts) up to that
+      # many extra times — EMB.MULTI is not idempotent, so each re-send
+      # duplicates inference — and the batch still terminates in
+      # Emb::ServerError. An Array of per-retry delays is also accepted (one
+      # retry per entry). Operation errors (server error replies) are never
+      # retried.
+      self.reconnect_attempts = 0
     end
 
     def to_h