Unify managed rust corpus orchestration

Author: ealdent

docs/modernization-handoff.md

@@ -95,9 +95,9 @@ Delivered:
 - Rust-managed session settings lifecycle (`configure_corpus_session`) with settings-aware orchestration (`run_em_on_session_start`) wired in `Lda::Backends::Rust`
 - Rust session execution refactor to shared session corpus storage + borrowed orchestration helpers, eliminating deep corpus clone overhead on each session EM run
 - 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 session-path EM through managed Rust session orchestration (`run_em_on_session_with_corpus`), leaving session reuse/recovery decisions in Rust; when session orchestration is unavailable it still prefers direct Rust non-session orchestration (`run_em_with_start_seed`) before legacy Ruby-side beta-input fallback (`run_em`)
+- `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
-- Rust managed-session orchestration API (`run_em_on_session_with_corpus`) added to recreate missing sessions and run EM in one Rust call
+- 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
 - parity/compatibility test coverage and rust runtime CI

docs/porting-strategy.md

@@ -66,7 +66,8 @@ Completed in `codex/experiment-ruby3-modernization`:
 - 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.
 - Rust managed-session orchestration API added (`run_em_on_session_with_corpus`) to recreate missing sessions and execute EM in one Rust call.
 - 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 session-path EM through `run_em_on_session_with_corpus`, leaving session reuse/recovery decisions in Rust and reducing fallback to non-session orchestration when sessions are externally dropped.
+- `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.
 - 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

@@ -13,11 +13,12 @@ Current parity expectations:
 
 - Rust vs pure backend fixture parity remains exact within existing tolerances used by tests.
 - Session-based orchestration paths (`run_em_on_session`, `run_em_on_session_with_start_seed`, `run_em_on_session_start`, `run_em_on_session_with_corpus`) must match direct non-session orchestration for equivalent settings/seeds.
-- `Lda::Backends::Rust` session-path EM should prefer the managed Rust session entrypoint (`run_em_on_session_with_corpus`) rather than branching in Ruby between session-only and recovery paths.
+- `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.
 - 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.
 - Corpus reassignment through Rust session replacement lifecycle (`replace_corpus_session`) must preserve stable session count and route subsequent EM runs over updated corpus data.
 - Unknown start-mode handling in seed-aware Rust orchestration must match Ruby's non-seeded fallback behavior when given the same explicit seed.
 

ext/lda-ruby-rust/src/lib.rs

@@ -407,20 +407,20 @@ fn random_topic_term_probabilities(
 }
 
 fn corpus_session_data(
-    document_words: Vec<Vec<usize>>,
-    document_counts: Vec<Vec<f64>>,
+    document_words: &[Vec<usize>],
+    document_counts: &[Vec<f64>],
     terms: usize,
 ) -> Arc<CorpusSessionData> {
     Arc::new(CorpusSessionData {
-        document_words,
-        document_counts,
+        document_words: document_words.to_vec(),
+        document_counts: document_counts.to_vec(),
         terms,
     })
 }
 
-fn create_corpus_session(
-    document_words: Vec<Vec<usize>>,
-    document_counts: Vec<Vec<f64>>,
+fn create_corpus_session_internal(
+    document_words: &[Vec<usize>],
+    document_counts: &[Vec<f64>],
     terms: usize,
 ) -> i64 {
     let session_id = NEXT_CORPUS_SESSION_ID.fetch_add(1, Ordering::Relaxed);
@@ -438,11 +438,19 @@ fn create_corpus_session(
     }
 }
 
-fn replace_corpus_session(
-    session_id: i64,
+fn create_corpus_session(
     document_words: Vec<Vec<usize>>,
     document_counts: Vec<Vec<f64>>,
     terms: usize,
+) -> i64 {
+    create_corpus_session_internal(document_words.as_slice(), document_counts.as_slice(), terms)
+}
+
+fn replace_corpus_session_internal(
+    session_id: i64,
+    document_words: &[Vec<usize>],
+    document_counts: &[Vec<f64>],
+    terms: usize,
 ) -> i64 {
     if terms == 0 {
         return 0;
@@ -474,11 +482,25 @@ fn replace_corpus_session(
     }
 }
 
-fn ensure_corpus_session(
+fn replace_corpus_session(
     session_id: i64,
     document_words: Vec<Vec<usize>>,
     document_counts: Vec<Vec<f64>>,
     terms: usize,
+) -> i64 {
+    replace_corpus_session_internal(
+        session_id,
+        document_words.as_slice(),
+        document_counts.as_slice(),
+        terms,
+    )
+}
+
+fn ensure_corpus_session(
+    session_id: i64,
+    document_words: &[Vec<usize>],
+    document_counts: &[Vec<f64>],
+    terms: usize,
 ) -> i64 {
     if terms == 0 {
         return 0;
@@ -488,7 +510,7 @@ fn ensure_corpus_session(
         return session_id;
     }
 
-    create_corpus_session(document_words, document_counts, terms)
+    create_corpus_session_internal(document_words, document_counts, terms)
 }
 
 fn drop_corpus_session(session_id: i64) -> bool {
@@ -667,16 +689,42 @@ fn run_em_on_session_with_corpus(
         return empty_managed_session_em_output();
     }
 
-    let active_session_id =
-        ensure_corpus_session(session_id, document_words, document_counts, terms);
-    if active_session_id <= 0 {
-        return empty_managed_session_em_output();
+    let active_session_id = ensure_corpus_session(
+        session_id,
+        document_words.as_slice(),
+        document_counts.as_slice(),
+        terms,
+    );
+
+    if active_session_id > 0 {
+        let (beta_probabilities, beta_log, gamma, phi) = run_em_on_session(
+            active_session_id,
+            start.clone(),
+            topics,
+            max_iter,
+            convergence,
+            em_max_iter,
+            em_convergence,
+            init_alpha,
+            min_probability,
+            random_seed,
+        );
+
+        if !(beta_probabilities.is_empty()
+            && beta_log.is_empty()
+            && gamma.is_empty()
+            && phi.is_empty())
+        {
+            return (active_session_id, beta_probabilities, beta_log, gamma, phi);
+        }
     }
 
-    let (beta_probabilities, beta_log, gamma, phi) = run_em_on_session(
-        active_session_id,
-        start,
+    let (beta_probabilities, beta_log, gamma, phi) = run_em_with_start_seed_internal(
+        start.as_str(),
+        document_words.as_slice(),
+        document_counts.as_slice(),
         topics,
+        terms,
         max_iter,
         convergence,
         em_max_iter,

lib/lda-ruby/backends/rust.rb

@@ -106,18 +106,18 @@ def model
       private
 
       def rust_orchestrated_em(start)
-        session_orchestrated = rust_orchestrated_em_with_session(start)
-        return true if session_orchestrated
+        managed_orchestrated = rust_orchestrated_em_with_managed_corpus(start)
+        return true if managed_orchestrated
 
         direct_orchestrated = rust_orchestrated_em_with_start_seed(start)
         return true if direct_orchestrated
 
         rust_orchestrated_em_with_beta(start)
       end
 
-      def rust_orchestrated_em_with_session(start)
+      def rust_orchestrated_em_with_managed_corpus(start)
         return false unless defined?(::Lda::RustBackend)
-        return false unless ensure_rust_corpus_session
+        return false unless ensure_rust_corpus_snapshot
 
         random_seed = Integer(next_random_seed)
         if ::Lda::RustBackend.respond_to?(:run_em_on_session_with_corpus)
@@ -134,12 +134,13 @@ def rust_orchestrated_em_with_session(start)
           return false unless managed_output.is_a?(Array) && managed_output.size == 5
 
           session_id, beta_probabilities, beta_log, gamma, phi = managed_output
-          return false unless session_id.is_a?(Numeric) && session_id.positive?
-
           output = [beta_probabilities, beta_log, gamma, phi]
           return false unless valid_rust_em_output?(output, @rust_document_lengths, Integer(num_topics), Integer(@rust_corpus_terms))
 
-          @rust_corpus_session_id = Integer(session_id)
+          @rust_corpus_session_id =
+            if session_id.is_a?(Numeric) && session_id.positive?
+              Integer(session_id)
+            end
           @fallback.apply_em_state(
             beta_probabilities: beta_probabilities,
             beta_log: beta_log,
@@ -349,12 +350,6 @@ def register_rust_corpus_session(previous_session_id = nil)
         drop_rust_corpus_session_by_id(previous_session_id)
       end
 
-      def ensure_rust_corpus_session
-        ensure_rust_corpus_snapshot
-      rescue StandardError
-        false
-      end
-
       def ensure_rust_corpus_snapshot
         has_session_data = @rust_corpus_terms && @rust_document_lengths && @rust_document_words && @rust_document_counts
         return true if has_session_data

test/rust_orchestration_test.rb

@@ -623,8 +623,8 @@ def test_rust_backend_non_session_fallback_prefers_run_em_with_start_seed
     backend.em_convergence = @em_convergence
     backend.init_alpha = @init_alpha
 
-    # Force the direct non-session orchestration path.
-    backend.define_singleton_method(:ensure_rust_corpus_session) { false }
+    # Force the direct non-managed orchestration path.
+    backend.define_singleton_method(:rust_orchestrated_em_with_managed_corpus) { |_start| false }
 
     rust_singleton = Lda::RustBackend.singleton_class
     run_em_calls = 0
@@ -678,7 +678,7 @@ def test_rust_backend_direct_non_session_path_reuses_cached_corpus_snapshot
     backend.em_convergence = @em_convergence
     backend.init_alpha = @init_alpha
 
-    backend.define_singleton_method(:rust_orchestrated_em_with_session) { |_start| false }
+    backend.define_singleton_method(:rust_orchestrated_em_with_managed_corpus) { |_start| false }
     backend.define_singleton_method(:rust_em_corpus_input) do
       raise "direct non-session path should reuse cached corpus snapshot"
     end
@@ -689,14 +689,17 @@ def test_rust_backend_direct_non_session_path_reuses_cached_corpus_snapshot
     backend&.corpus = nil
   end
 
-  def test_rust_backend_session_path_prefers_managed_session_entrypoint
+  def test_rust_backend_prefers_managed_corpus_entrypoint_without_active_session
     backend = nil
     rust_singleton = nil
     run_em_on_session_alias = :__test_original_run_em_on_session_for_managed_preference__
     run_em_on_session_with_corpus_alias = :__test_original_run_em_on_session_with_corpus_for_managed_preference__
+    run_em_with_start_seed_alias = :__test_original_run_em_with_start_seed_for_managed_preference__
 
     omit("run_em_on_session unavailable") unless Lda::RustBackend.respond_to?(:run_em_on_session)
     omit("run_em_on_session_with_corpus unavailable") unless Lda::RustBackend.respond_to?(:run_em_on_session_with_corpus)
+    omit("run_em_with_start_seed unavailable") unless Lda::RustBackend.respond_to?(:run_em_with_start_seed)
+    omit("drop_corpus_session unavailable") unless Lda::RustBackend.respond_to?(:drop_corpus_session)
 
     backend = Lda::Backends::Rust.new(random_seed: 1234)
     backend.corpus = Lda::TextCorpus.new(FIXTURE_DOCUMENTS)
@@ -708,13 +711,20 @@ def test_rust_backend_session_path_prefers_managed_session_entrypoint
     backend.em_convergence = @em_convergence
     backend.init_alpha = @init_alpha
 
+    dropped_session_id = backend.instance_variable_get(:@rust_corpus_session_id)
+    assert_operator dropped_session_id, :>, 0
+    assert_equal true, Lda::RustBackend.drop_corpus_session(dropped_session_id)
+    backend.instance_variable_set(:@rust_corpus_session_id, nil)
+
     rust_singleton = Lda::RustBackend.singleton_class
     run_em_on_session_calls = 0
     run_em_on_session_with_corpus_calls = 0
+    run_em_with_start_seed_calls = 0
 
     silence_redefinition_warnings do
       rust_singleton.send(:alias_method, run_em_on_session_alias, :run_em_on_session)
       rust_singleton.send(:alias_method, run_em_on_session_with_corpus_alias, :run_em_on_session_with_corpus)
+      rust_singleton.send(:alias_method, run_em_with_start_seed_alias, :run_em_with_start_seed)
 
       rust_singleton.send(:define_method, :run_em_on_session) do |*args|
         run_em_on_session_calls += 1
@@ -725,14 +735,30 @@ def test_rust_backend_session_path_prefers_managed_session_entrypoint
         run_em_on_session_with_corpus_calls += 1
         public_send(run_em_on_session_with_corpus_alias, *args)
       end
+
+      rust_singleton.send(:define_method, :run_em_with_start_seed) do |*args|
+        run_em_with_start_seed_calls += 1
+        public_send(run_em_with_start_seed_alias, *args)
+      end
     end
 
     backend.em("seeded")
     assert_equal 0, run_em_on_session_calls
     assert_equal 1, run_em_on_session_with_corpus_calls
+    assert_equal 0, run_em_with_start_seed_calls
     assert_equal @topics, backend.gamma.first.size
+
+    recreated_session_id = backend.instance_variable_get(:@rust_corpus_session_id)
+    assert_operator recreated_session_id, :>, 0
+    assert_not_equal dropped_session_id, recreated_session_id
   ensure
     silence_redefinition_warnings do
+      if defined?(rust_singleton) && rust_singleton.method_defined?(run_em_with_start_seed_alias)
+        rust_singleton.send(:remove_method, :run_em_with_start_seed)
+        rust_singleton.send(:alias_method, :run_em_with_start_seed, run_em_with_start_seed_alias)
+        rust_singleton.send(:remove_method, run_em_with_start_seed_alias)
+      end
+
       if defined?(rust_singleton) && rust_singleton.method_defined?(run_em_on_session_with_corpus_alias)
         rust_singleton.send(:remove_method, :run_em_on_session_with_corpus)
         rust_singleton.send(:alias_method, :run_em_on_session_with_corpus, run_em_on_session_with_corpus_alias)