Reuse cached snapshot in beta fallback

Author: ealdent

docs/modernization-handoff.md

@@ -97,6 +97,7 @@ Delivered:
 - unified Rust session orchestration API (`run_em_on_session`) that applies settings + runs EM in one call inside Rust session orchestration
 - `Lda::Backends::Rust` now routes cached-corpus EM through managed Rust orchestration (`run_em_on_session_with_corpus`), leaving session reuse/recovery decisions in Rust and preferring that managed path even when no active session id is cached locally; when the managed API is unavailable it still prefers direct Rust non-session orchestration (`run_em_with_start_seed`) before legacy Ruby-side beta-input fallback (`run_em`)
 - direct non-session Rust orchestration now reuses the backend's cached Rust corpus snapshot instead of rebuilding corpus arrays from `@corpus` on each fallback invocation
+- legacy Rust beta-input compatibility fallback now also reuses the backend's cached Rust corpus snapshot, only asking the pure-Ruby backend to synthesize the initial beta matrix
 - Rust managed-session orchestration API (`run_em_on_session_with_corpus`) added to recreate missing sessions and run EM in one Rust call, and now directly falls back to start-aware array execution inside Rust if session-backed execution cannot be used
 - Rust session lifecycle replacement API (`replace_corpus_session`) added so corpus reassignment can update existing Rust sessions in place (config reset + corpus swap) instead of Ruby-side drop/recreate
 - `Lda::Backends::Rust` now keeps session-based orchestration on the managed Rust path (`run_em_on_session_with_corpus`) even when sessions are dropped externally

docs/porting-strategy.md

@@ -68,6 +68,7 @@ Completed in `codex/experiment-ruby3-modernization`:
 - Rust session lifecycle replacement API added (`replace_corpus_session`) so corpus reassignment can update existing Rust sessions in place (config reset + corpus swap) instead of Ruby-side drop/recreate.
 - `Lda::Backends::Rust` now routes cached-corpus EM through `run_em_on_session_with_corpus`, leaving session reuse/recovery decisions in Rust, preferring that managed path even when no active session id is cached locally, and reducing Ruby-side fallback branching when sessions are externally dropped.
 - `run_em_on_session_with_corpus` now acts as a unified Rust managed-corpus entrypoint: it attempts session-backed execution first, then falls back to direct start-aware array execution inside Rust when a managed session cannot be used.
+- Legacy `run_em(initial_beta, ...)` compatibility fallback now reuses the Rust backend's cached corpus snapshot and only relies on the pure-Ruby backend to synthesize the initial beta matrix.
 - Dockerized rust runtime workflow added for local parity with CI (`Dockerfile.rust`, `bin/docker-test-rust`).
 - Gem packaging now excludes local Rust cargo build artifacts (`target/**`) for clean release builds.
 - Backend benchmark driver added (`bin/benchmark-backends`) to track pure/native/rust runtime deltas.

docs/rust-orchestration-guardrails.md

@@ -16,6 +16,7 @@ Current parity expectations:
 - `Lda::Backends::Rust` cached-corpus EM should prefer the managed Rust session entrypoint (`run_em_on_session_with_corpus`) even when no active session id is cached locally, rather than branching in Ruby between session-only, recovery, and direct paths.
 - `Lda::Backends::Rust` non-session fallback should prefer Rust start-aware orchestration (`run_em_with_start_seed`) before legacy beta-input orchestration (`run_em`).
 - Direct non-session fallback should reuse the backend's cached Rust corpus snapshot rather than rebuilding corpus arrays from `@corpus` for each invocation.
+- Legacy beta-input compatibility fallback should also reuse the backend's cached Rust corpus snapshot rather than rebuilding full EM corpus input in Ruby.
 - Rust backend corpus/session lifecycle must not leak session count across corpus replacement.
 - Missing-session recovery in managed session orchestration (`run_em_on_session_with_corpus`) must recreate a usable session and keep parity with direct orchestration.
 - Managed Rust corpus orchestration (`run_em_on_session_with_corpus`) must keep parity with direct orchestration even when it falls back internally from session-backed execution to start-seeded array execution.

ext/lda-ruby-rust/README.md

@@ -21,6 +21,7 @@ Current scope:
   - `Lda::RustBackend.seeded_topic_term_probabilities(document_words, document_counts, topics, terms, min_probability)`
   - `Lda::RustBackend.random_topic_term_probabilities(topics, terms, min_probability, random_seed)`
   - `Lda::RustBackend.create_corpus_session(document_words, document_counts, terms)`
+  - `Lda::RustBackend.replace_corpus_session(session_id, document_words, document_counts, terms)`
   - `Lda::RustBackend.drop_corpus_session(session_id)`
   - `Lda::RustBackend.configure_corpus_session(session_id, topics, max_iter, convergence, em_max_iter, em_convergence, init_alpha, min_probability)`
   - `Lda::RustBackend.run_em(initial_beta, document_words, document_counts, max_iter, convergence, em_max_iter, em_convergence, init_alpha, min_probability)`
@@ -29,6 +30,7 @@ Current scope:
   - `Lda::RustBackend.run_em_on_session(session_id, start, topics, max_iter, convergence, em_max_iter, em_convergence, init_alpha, min_probability, random_seed)`
   - `Lda::RustBackend.run_em_on_session_start(session_id, start, random_seed)`
   - `Lda::RustBackend.run_em_on_session_with_start_seed(session_id, start, topics, max_iter, convergence, em_max_iter, em_convergence, init_alpha, min_probability, random_seed)`
+  - `Lda::RustBackend.run_em_on_session_with_corpus(session_id, document_words, document_counts, terms, start, topics, max_iter, convergence, em_max_iter, em_convergence, init_alpha, min_probability, random_seed)`
 
 Hot-path kernels currently executed in Rust when `backend: :rust` is active:
 - topic weights for a word across topics
@@ -46,7 +48,9 @@ Hot-path kernels currently executed in Rust when `backend: :rust` is active:
 - unified session-settings orchestration (`run_em_on_session`) that applies settings and executes EM in one call
 - session-based EM orchestration against Rust-managed corpus lifecycle (`create_corpus_session` + `run_em_on_session_with_start_seed`)
 - settings-aware session orchestration (`configure_corpus_session` + `run_em_on_session_start`)
-- `Lda::Backends::Rust` prefers `run_em_with_start_seed` for direct non-session orchestration when session orchestration is unavailable
+- managed corpus orchestration (`run_em_on_session_with_corpus`) that can recreate missing sessions and, if session-backed execution cannot be used, falls back internally to direct start-aware execution inside Rust
+- `Lda::Backends::Rust` prefers `run_em_on_session_with_corpus` whenever a cached Rust corpus snapshot is available, even if no session id is currently cached locally
+- direct and legacy beta-input compatibility fallbacks both reuse the backend's cached Rust corpus snapshot instead of rebuilding corpus arrays in Ruby
 - unknown EM start modes in seed-aware orchestration follow Ruby's non-seeded fallback behavior (seeded by explicit `random_seed`)
 
 Remaining numeric LDA kernels are still provided by the pure Ruby backend and will move incrementally.

lib/lda-ruby/backends/pure_ruby.rb

@@ -59,6 +59,23 @@ def rust_em_input(start)
         build_em_input(start)
       end
 
+      # Returns only the initial beta matrix for Rust compatibility paths that
+      # already hold a cached corpus snapshot.
+      def rust_initial_beta_probabilities(start, document_words, document_counts, topics, terms)
+        start_mode = start.to_s
+
+        if start_mode.strip.casecmp("seeded").zero? || start_mode.strip.casecmp("deterministic").zero?
+          seeded_topic_term_probabilities(
+            Integer(topics),
+            Integer(terms),
+            document_words,
+            document_counts
+          )
+        else
+          initial_topic_term_probabilities(Integer(topics), Integer(terms))
+        end
+      end
+
       def em_from_input(em_input)
         return nil if em_input.nil?
 
@@ -129,21 +146,20 @@ def build_em_input(start)
         document_words = @corpus.documents.map { |document| document.words.map(&:to_i) }
         document_counts = @corpus.documents.map { |document| document.counts.map(&:to_f) }
 
-        initial_beta_probabilities =
-          if start.to_s.strip.casecmp("seeded").zero? || start.to_s.strip.casecmp("deterministic").zero?
-            seeded_topic_term_probabilities(topics, terms, document_words, document_counts)
-          else
-            initial_topic_term_probabilities(topics, terms)
-          end
-
         {
           topics: topics,
           terms: terms,
           document_words: document_words,
           document_counts: document_counts,
           document_totals: document_counts.map { |counts| counts.sum.to_f },
           document_lengths: document_words.map(&:length),
-          initial_beta_probabilities: initial_beta_probabilities,
+          initial_beta_probabilities: rust_initial_beta_probabilities(
+            start,
+            document_words,
+            document_counts,
+            topics,
+            terms
+          ),
           min_probability: MIN_PROBABILITY
         }
       end

lib/lda-ruby/backends/rust.rb

@@ -221,7 +221,32 @@ def rust_orchestrated_em_with_beta(start)
         return false unless defined?(::Lda::RustBackend)
         return false unless ::Lda::RustBackend.respond_to?(:run_em)
 
-        em_input = @fallback.rust_em_input(start)
+        em_input =
+          if ensure_rust_corpus_snapshot && @fallback.respond_to?(:rust_initial_beta_probabilities)
+            topics = Integer(num_topics)
+            terms = Integer(@rust_corpus_terms)
+            initial_beta_probabilities = @fallback.rust_initial_beta_probabilities(
+              start,
+              @rust_document_words,
+              @rust_document_counts,
+              topics,
+              terms
+            )
+
+            {
+              topics: topics,
+              terms: terms,
+              document_words: @rust_document_words,
+              document_counts: @rust_document_counts,
+              document_totals: @rust_document_counts.map { |counts| counts.sum.to_f },
+              document_lengths: @rust_document_lengths,
+              initial_beta_probabilities: initial_beta_probabilities,
+              min_probability: MIN_PROBABILITY
+            }
+          else
+            @fallback.rust_em_input(start)
+          end
+
         return true if em_input.nil?
 
         output = ::Lda::RustBackend.run_em(

test/pure_ruby_orchestration_test.rb

@@ -36,6 +36,26 @@ def test_em_from_input_matches_seeded_em_output
     assert_nested_close(direct.compute_phi, from_input.compute_phi, 1e-9)
   end
 
+  def test_rust_initial_beta_probabilities_matches_rust_em_input_for_random_start
+    from_helper = build_backend
+    from_input = build_backend
+
+    document_words = from_helper.corpus.documents.map { |document| document.words.map(&:to_i) }
+    document_counts = from_helper.corpus.documents.map { |document| document.counts.map(&:to_f) }
+    terms = from_helper.corpus.documents.flat_map(&:words).max + 1
+
+    helper_beta = from_helper.rust_initial_beta_probabilities(
+      "random",
+      document_words,
+      document_counts,
+      from_helper.num_topics,
+      terms
+    )
+    em_input = from_input.rust_em_input("random")
+
+    assert_nested_close(helper_beta, em_input[:initial_beta_probabilities], 1e-12)
+  end
+
   def test_apply_em_state_sets_outputs
     backend = build_backend
 

test/rust_orchestration_test.rb

@@ -689,6 +689,74 @@ def test_rust_backend_direct_non_session_path_reuses_cached_corpus_snapshot
     backend&.corpus = nil
   end
 
+  def test_rust_backend_beta_fallback_reuses_cached_corpus_snapshot
+    backend = nil
+    fallback = nil
+    fallback_singleton = nil
+    rust_em_input_alias = :__test_original_rust_em_input_for_beta_snapshot__
+    rust_initial_beta_alias = :__test_original_rust_initial_beta_for_beta_snapshot__
+
+    backend = Lda::Backends::Rust.new(random_seed: 1234)
+    backend.corpus = Lda::TextCorpus.new(FIXTURE_DOCUMENTS)
+    backend.verbose = false
+    backend.num_topics = @topics
+    backend.max_iter = @max_iter
+    backend.convergence = @convergence
+    backend.em_max_iter = @em_max_iter
+    backend.em_convergence = @em_convergence
+    backend.init_alpha = @init_alpha
+
+    backend.define_singleton_method(:rust_orchestrated_em_with_managed_corpus) { |_start| false }
+    backend.define_singleton_method(:rust_orchestrated_em_with_start_seed) { |_start| false }
+
+    cached_document_words = backend.instance_variable_get(:@rust_document_words)
+    cached_document_counts = backend.instance_variable_get(:@rust_document_counts)
+    cached_terms = backend.instance_variable_get(:@rust_corpus_terms)
+    expected_topics = @topics
+
+    fallback = backend.instance_variable_get(:@fallback)
+    fallback_singleton = fallback.singleton_class
+    used_cached_snapshot = false
+
+    silence_redefinition_warnings do
+      fallback_singleton.send(:alias_method, rust_em_input_alias, :rust_em_input)
+      fallback_singleton.send(:alias_method, rust_initial_beta_alias, :rust_initial_beta_probabilities)
+
+      fallback_singleton.send(:define_method, :rust_em_input) do |_start|
+        raise "beta fallback should not rebuild full rust_em_input when snapshot is cached"
+      end
+
+      fallback_singleton.send(:define_method, :rust_initial_beta_probabilities) do |start, document_words, document_counts, topics, terms|
+        used_cached_snapshot =
+          document_words.equal?(cached_document_words) &&
+          document_counts.equal?(cached_document_counts) &&
+          topics == expected_topics &&
+          terms == cached_terms
+        public_send(rust_initial_beta_alias, start, document_words, document_counts, topics, terms)
+      end
+    end
+
+    backend.em("random")
+    assert_equal true, used_cached_snapshot
+    assert_equal @topics, backend.gamma.first.size
+  ensure
+    silence_redefinition_warnings do
+      if defined?(fallback_singleton) && fallback_singleton.method_defined?(rust_initial_beta_alias)
+        fallback_singleton.send(:remove_method, :rust_initial_beta_probabilities)
+        fallback_singleton.send(:alias_method, :rust_initial_beta_probabilities, rust_initial_beta_alias)
+        fallback_singleton.send(:remove_method, rust_initial_beta_alias)
+      end
+
+      if defined?(fallback_singleton) && fallback_singleton.method_defined?(rust_em_input_alias)
+        fallback_singleton.send(:remove_method, :rust_em_input)
+        fallback_singleton.send(:alias_method, :rust_em_input, rust_em_input_alias)
+        fallback_singleton.send(:remove_method, rust_em_input_alias)
+      end
+    end
+
+    backend&.corpus = nil
+  end
+
   def test_rust_backend_prefers_managed_corpus_entrypoint_without_active_session
     backend = nil
     rust_singleton = nil