embedding_util 0.1.1 → 0.1.2

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: f388bd90069b46caca18046622511f7fe2bd5c25f2a56a8d02fc01cc673bc682
-  data.tar.gz: 75b484781ae6689e6bf471007257c1e63f4e8aa2485b558e0d41b2f03643749f
+  metadata.gz: fb8276720e51a283fb4b8c70507ae54d25ea968caffe8142f636a6e6ea948b71
+  data.tar.gz: eccb6e6d11006238a1af1bfc25fb3b5146325bf78a1deb0f447a73541dec128d
 SHA512:
-  metadata.gz: 4bd65f54bc228373056843d55aa866d510660361e294bff33b18ec8440f02ae807f4bf5773e6250eb2a5152e53d3656dadef0f624b50f0c1ec32fab689f78367
-  data.tar.gz: 9f51238b20d2aabfda68aace11f276e79cbdeb291468bd888933330869a5717c52f33db2ceb7a46b870a8f138a1b9fcbdbc94e4e5108f529ba7d4aa5ab299e6a
+  metadata.gz: e625a2389828218a8f39cc4dc69d6ea965f59e032d190407668029df227f3c7ffee32ce30b64141d8a38654573ac591315c08696f051db3f68772efaf1b7bf34
+  data.tar.gz: 5da07a5487e7167128b67d8b6e39009d6a2fd4c0949daf1e61ae2cb9e98fa668beecd618713b1d0545daa47a86ac6f4e77f66d3429093c658852b594b4aa3cc5

data/CHANGELOG.md

@@ -1,3 +1,11 @@
+## [0.1.2] - 2026-06-10
+
+- Add self-hosted reranker recovery for llama.cpp physical batch-size failures
+- Start managed reranker servers with `--ubatch-size 1024`
+- Restart managed reranker servers once with `--ubatch-size 4096` when larger rerank requests require it
+- Add configuration and CLI options for reranker ubatch defaults and maximums
+- Add clearer guidance for app-managed reranker endpoints that need a larger `--ubatch-size`
+
 ## [0.1.1] - 2026-06-08
 
 - Fix local server lifecycle cleanup for Ramalama and direct `llama-server`

data/README.md

@@ -118,7 +118,9 @@ Reranker model:
 
 - repo: `ggml-org/Qwen3-Reranker-0.6B-Q8_0-GGUF`
 - file: `qwen3-reranker-0.6b-q8_0.gguf`
-- server flags: `--reranking`
+- server flags: `--reranking --ubatch-size 1024`
+
+For self-hosted rerankers, `embedding_util` starts with `--ubatch-size 1024`. If a larger rerank request exceeds llama.cpp's physical batch size, the managed reranker is restarted once with `config.reranker_max_ubatch_size`, which defaults to `4096`, and the request is retried.
 
 Do not combine embedding and reranking flags for this profile. Run separate local servers.
 
@@ -133,6 +135,8 @@ EmbeddingUtil.configure do |config|
   config.reranker_port = 18081
   config.startup_timeout = 3600
   config.shutdown_idle = 300
+  config.reranker_ubatch_size = 1024
+  config.reranker_max_ubatch_size = 4096
   config.timeout = 60
 end
 ```
@@ -159,6 +163,8 @@ Environment variables are also supported:
 - `EMBEDDING_UTIL_VERBOSE`
 - `EMBEDDING_UTIL_EMBEDDING_PORT`
 - `EMBEDDING_UTIL_RERANKER_PORT`
+- `EMBEDDING_UTIL_RERANKER_UBATCH_SIZE`
+- `EMBEDDING_UTIL_RERANKER_MAX_UBATCH_SIZE`
 
 ## Development
 

data/lib/embedding_util/cli.rb

@@ -14,6 +14,8 @@ module EmbeddingUtil
       timeout: ->(value) { value },
       startup_timeout: ->(value) { value },
       shutdown_idle: :to_i.to_proc,
+      reranker_ubatch_size: :to_i.to_proc,
+      reranker_max_ubatch_size: :to_i.to_proc,
       verbose: ->(value) { value }
     }.freeze
 
@@ -25,6 +27,8 @@ module EmbeddingUtil
     class_option :timeout, type: :numeric, desc: "HTTP timeout in seconds"
     class_option :startup_timeout, type: :numeric, desc: "Seconds to wait for self-hosted server startup"
     class_option :shutdown_idle, type: :numeric, desc: "Stop self-hosted server after this many seconds without stdout/stderr activity"
+    class_option :reranker_ubatch_size, type: :numeric, desc: "llama.cpp physical batch size for self-hosted reranker servers"
+    class_option :reranker_max_ubatch_size, type: :numeric, desc: "Largest reranker physical batch size for automatic retry"
     class_option :verbose, type: :boolean, desc: "Print self-hosting diagnostics"
 
     desc "support", "Display configured provider support"

data/lib/embedding_util/configuration.rb

@@ -3,7 +3,7 @@
 module EmbeddingUtil
   class Configuration
     attr_accessor :profile, :provider, :endpoint, :embedding_endpoint, :reranker_endpoint, :timeout, :startup_timeout, :shutdown_idle, :host,
-                  :embedding_port, :reranker_port, :state_dir, :verbose
+                  :embedding_port, :reranker_port, :state_dir, :verbose, :reranker_ubatch_size, :reranker_max_ubatch_size
     attr_reader :runtime
 
     def initialize
@@ -19,6 +19,8 @@ module EmbeddingUtil
       @host = ENV.fetch("EMBEDDING_UTIL_HOST", "127.0.0.1")
       @embedding_port = Integer(ENV.fetch("EMBEDDING_UTIL_EMBEDDING_PORT", "18080"))
       @reranker_port = Integer(ENV.fetch("EMBEDDING_UTIL_RERANKER_PORT", "18081"))
+      @reranker_ubatch_size = Integer(ENV.fetch("EMBEDDING_UTIL_RERANKER_UBATCH_SIZE", "1024"))
+      @reranker_max_ubatch_size = Integer(ENV.fetch("EMBEDDING_UTIL_RERANKER_MAX_UBATCH_SIZE", "4096"))
       @state_dir = ENV.fetch("EMBEDDING_UTIL_STATE_DIR", File.expand_path("~/.local/state/embedding_util"))
       @verbose = ENV.fetch("EMBEDDING_UTIL_VERBOSE", "false").match?(/\A(?:1|true|yes|on)\z/i)
     end

data/lib/embedding_util/profiles.rb

@@ -19,7 +19,7 @@ module EmbeddingUtil
         repo: "ggml-org/Qwen3-Reranker-0.6B-Q8_0-GGUF",
         file: "qwen3-reranker-0.6b-q8_0.gguf",
         model: "qwen3-reranker-0.6b",
-        server_flags: ["--reranking"]
+        server_flags: ["--reranking", "--ubatch-size", "1024"]
       }
     )
 

data/lib/embedding_util/providers/endpoint.rb

@@ -118,7 +118,7 @@ module EmbeddingUtil
         end
 
         raise EndpointNotFoundError.new(uri, path: path, body: response.body) if response.code.to_i == 404 && route_missing_response?(response.body)
-        raise EndpointError, "#{uri} returned #{response.code}: #{response.body}" unless response.is_a?(Net::HTTPSuccess)
+        raise EndpointError, endpoint_error_message(uri, response, path) unless response.is_a?(Net::HTTPSuccess)
 
         JSON.parse(response.body)
       rescue JSON::ParserError => e
@@ -148,6 +148,18 @@ module EmbeddingUtil
       def fallback_rerank_not_found?(error)
         error.path == "/v1/rerank"
       end
+
+      def endpoint_error_message(uri, response, path)
+        message = "#{uri} returned #{response.code}: #{response.body}"
+        return message unless reranker_batch_size_error?(path, response.body)
+
+        "#{message}. Restart the reranker server with a larger llama.cpp --ubatch-size; " \
+          "embedding_util-managed reranker servers use --ubatch-size 1024 by default."
+      end
+
+      def reranker_batch_size_error?(path, body)
+        path.end_with?("/rerank") && body.to_s.include?("increase the physical batch size")
+      end
     end
   end
 end

data/lib/embedding_util/providers/self_hosted.rb

@@ -27,7 +27,14 @@ module EmbeddingUtil
       end
 
       def rerank(query, documents, profile: config.resolved_profile)
-        endpoint = ServerManager.new(config: config).ensure_server(:reranker, profile: profile)
+        manager = ServerManager.new(config: config)
+        endpoint = manager.ensure_server(:reranker, profile: profile)
+        endpoint_provider(reranker_endpoint: endpoint).rerank(query, documents, profile: profile)
+      rescue EndpointError => e
+        raise unless reranker_batch_size_error?(e) && can_escalate_reranker_ubatch?
+
+        config.reranker_ubatch_size = config.reranker_max_ubatch_size
+        endpoint = manager.restart_server(:reranker, profile: profile)
         endpoint_provider(reranker_endpoint: endpoint).rerank(query, documents, profile: profile)
       end
 
@@ -39,6 +46,14 @@ module EmbeddingUtil
         endpoint_config.reranker_endpoint = reranker_endpoint
         Endpoint.new(config: endpoint_config)
       end
+
+      def reranker_batch_size_error?(error)
+        error.message.include?("increase the physical batch size")
+      end
+
+      def can_escalate_reranker_ubatch?
+        config.reranker_ubatch_size < config.reranker_max_ubatch_size
+      end
     end
   end
 end

data/lib/embedding_util/runtime_command.rb

@@ -2,13 +2,14 @@
 
 module EmbeddingUtil
   class RuntimeCommand
-    attr_reader :runtime, :server_model, :host, :port
+    attr_reader :runtime, :server_model, :host, :port, :server_flags
 
-    def initialize(runtime:, server_model:, host:, port:)
+    def initialize(runtime:, server_model:, host:, port:, server_flags: nil)
       @runtime = self.class.normalize_runtime(runtime)
       @server_model = server_model
       @host = host
       @port = port
+      @server_flags = server_flags || server_model.settings.fetch(:server_flags)
     end
 
     def self.available?(runtime)
@@ -86,7 +87,7 @@ module EmbeddingUtil
         "--name", server_name,
         "--host", host,
         "--port", port.to_s,
-        "--runtime-args=#{server_model.settings.fetch(:server_flags).join(' ')}",
+        "--runtime-args=#{server_flags.join(' ')}",
         huggingface_model
       ]
     end
@@ -98,7 +99,7 @@ module EmbeddingUtil
         "--port", port.to_s,
         "-hf", server_model.settings.fetch(:repo),
         "-hff", server_model.settings.fetch(:file),
-        *server_model.settings.fetch(:server_flags)
+        *server_flags
       ]
     end
 

data/lib/embedding_util/server_manager.rb

@@ -37,7 +37,7 @@ module EmbeddingUtil
       server_model = model.is_a?(ServerModel) ? model : ServerModel.parse(model)
       resolved_runtime = RuntimeCommand.resolve(runtime)
       selected_port = selected_port_for(server_model, host: host, port: port)
-      command = RuntimeCommand.new(runtime: resolved_runtime, server_model: server_model, host: host, port: selected_port)
+      command = runtime_command(resolved_runtime, server_model, host, selected_port)
       last_output_at = Time.now
 
       FileUtils.mkdir_p(config.state_dir)
@@ -61,6 +61,17 @@ module EmbeddingUtil
       end
     end
 
+    def restart_server(capability, profile: config.resolved_profile)
+      server_model = ServerModel.for(capability, profile)
+
+      with_lock(server_model) do
+        stop_server(server_model)
+        start_background(server_model)
+      end
+
+      wait_for_healthy(server_model, log_path: server_log_path(server_model))
+    end
+
     private
 
     def start_background(server_model)
@@ -75,6 +86,8 @@ module EmbeddingUtil
         "--port", selected_port.to_s
       ]
       argv.push("--shutdown-idle", config.shutdown_idle.to_s) unless config.shutdown_idle.nil?
+      argv.push("--reranker-ubatch-size", config.reranker_ubatch_size.to_s)
+      argv.push("--reranker-max-ubatch-size", config.reranker_max_ubatch_size.to_s)
       warn "starting #{server_model.name} in background: #{argv.join(' ')}" if config.verbose
       warn "#{server_model.name} log: #{log_path}" if config.verbose
       pid = Process.spawn(*argv, out: [log_path, "a"], err: %i[child out], pgroup: true)
@@ -100,6 +113,42 @@ module EmbeddingUtil
       available_port(host, server_model.default_port(config))
     end
 
+    def runtime_command(runtime, server_model, host, port)
+      RuntimeCommand.new(
+        runtime: runtime,
+        server_model: server_model,
+        host: host,
+        port: port,
+        server_flags: server_flags(server_model)
+      )
+    end
+
+    def server_flags(server_model)
+      flags = server_model.settings.fetch(:server_flags)
+      return flags unless server_model.capability == :reranker
+
+      with_ubatch_size(flags, config.reranker_ubatch_size)
+    end
+
+    def with_ubatch_size(flags, size)
+      filtered = []
+      skip_next = false
+      flags.each do |flag|
+        if skip_next
+          skip_next = false
+          next
+        end
+
+        if ["--ubatch-size", "-ub"].include?(flag)
+          skip_next = true
+          next
+        end
+
+        filtered << flag
+      end
+      filtered + ["--ubatch-size", size.to_s]
+    end
+
     def required_port(host, port)
       return port if port_available?(host, port)
 
@@ -213,6 +262,21 @@ module EmbeddingUtil
       end
     end
 
+    def stop_server(server_model)
+      state = read_state(server_model)
+      return delete_state(server_model) unless state
+
+      runtime = state.fetch("runtime", config.runtime)
+      port = state.fetch("port", server_model.default_port(config))
+      command = runtime_command(runtime, server_model, config.host, port)
+      if command.detached_server?
+        stop_detached_server(command)
+      else
+        terminate_runtime_process(command, state["pid"])
+      end
+      delete_state(server_model)
+    end
+
     def cleanup_runtime(command, wait_thread)
       return unless command
 

data/lib/embedding_util/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module EmbeddingUtil
-  VERSION = "0.1.1"
+  VERSION = "0.1.2"
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: embedding_util
 version: !ruby/object:Gem::Version
-  version: 0.1.1
+  version: 0.1.2
 platform: ruby
 authors:
 - hmdne