Add Rust corpus-session replacement lifecycle API

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, now preferred by `Lda::Backends::Rust`
 - `Lda::Backends::Rust` now routes all start modes through unified session orchestration when available (`run_em_on_session`); when sessions are unavailable it now prefers direct Rust non-session orchestration (`run_em_with_start_seed`) before legacy Ruby-side beta-input fallback (`run_em`)
 - Rust managed-session orchestration API (`run_em_on_session_with_corpus`) added to recreate missing sessions and run EM in one Rust call
+- 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 retries missing-session runs through `run_em_on_session_with_corpus` to preserve session-based orchestration when sessions are dropped externally
 - parity/compatibility test coverage and rust runtime CI
 

docs/porting-strategy.md

@@ -64,6 +64,7 @@ Completed in `codex/experiment-ruby3-modernization`:
 - Unified Rust session API added (`run_em_on_session`) to apply settings and execute EM in one call; `Lda::Backends::Rust` now prefers this single-call session path before non-session fallbacks.
 - `Lda::Backends::Rust` now prefers direct Rust non-session orchestration (`run_em_with_start_seed`) before legacy `run_em(initial_beta, ...)` compatibility fallback when a session path is unavailable.
 - 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 retries missing-session runs through `run_em_on_session_with_corpus`, reducing fallback to non-session orchestration when sessions are externally dropped.
 - 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.

docs/rust-orchestration-guardrails.md

@@ -16,6 +16,7 @@ Current parity expectations:
 - `Lda::Backends::Rust` non-session fallback should prefer Rust start-aware orchestration (`run_em_with_start_seed`) before legacy beta-input orchestration (`run_em`).
 - 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.
+- 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.
 
 ## Benchmark guardrail

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

@@ -406,18 +406,26 @@ fn random_topic_term_probabilities(
     matrix
 }
 
+fn corpus_session_data(
+    document_words: Vec<Vec<usize>>,
+    document_counts: Vec<Vec<f64>>,
+    terms: usize,
+) -> Arc<CorpusSessionData> {
+    Arc::new(CorpusSessionData {
+        document_words,
+        document_counts,
+        terms,
+    })
+}
+
 fn create_corpus_session(
     document_words: Vec<Vec<usize>>,
     document_counts: Vec<Vec<f64>>,
     terms: usize,
 ) -> i64 {
     let session_id = NEXT_CORPUS_SESSION_ID.fetch_add(1, Ordering::Relaxed);
     let session = CorpusSession {
-        data: Arc::new(CorpusSessionData {
-            document_words,
-            document_counts,
-            terms,
-        }),
+        data: corpus_session_data(document_words, document_counts, terms),
         config: None,
     };
 
@@ -430,6 +438,42 @@ fn create_corpus_session(
     }
 }
 
+fn replace_corpus_session(
+    session_id: i64,
+    document_words: Vec<Vec<usize>>,
+    document_counts: Vec<Vec<f64>>,
+    terms: usize,
+) -> i64 {
+    if terms == 0 {
+        return 0;
+    }
+
+    let replacement_data = corpus_session_data(document_words, document_counts, terms);
+    match corpus_sessions().lock() {
+        Ok(mut sessions) => {
+            if session_id > 0 {
+                let session_key = session_id as u64;
+                if let Some(session) = sessions.get_mut(&session_key) {
+                    session.data = replacement_data;
+                    session.config = None;
+                    return session_id;
+                }
+            }
+
+            let new_session_id = NEXT_CORPUS_SESSION_ID.fetch_add(1, Ordering::Relaxed);
+            sessions.insert(
+                new_session_id,
+                CorpusSession {
+                    data: replacement_data,
+                    config: None,
+                },
+            );
+            new_session_id as i64
+        }
+        Err(_) => 0,
+    }
+}
+
 fn ensure_corpus_session(
     session_id: i64,
     document_words: Vec<Vec<usize>>,
@@ -1146,6 +1190,8 @@ fn init() -> Result<(), Error> {
     )?;
     rust_backend_module
         .define_singleton_method("create_corpus_session", function!(create_corpus_session, 3))?;
+    rust_backend_module
+        .define_singleton_method("replace_corpus_session", function!(replace_corpus_session, 4))?;
     rust_backend_module
         .define_singleton_method("drop_corpus_session", function!(drop_corpus_session, 1))?;
     rust_backend_module

lib/lda-ruby/backends/rust.rb

@@ -56,10 +56,10 @@ def name
       end
 
       def corpus=(corpus)
-        release_rust_corpus_session
+        previous_session_id = @rust_corpus_session_id
         @corpus = corpus
         @fallback.corpus = corpus
-        register_rust_corpus_session
+        register_rust_corpus_session(previous_session_id)
         true
       end
 
@@ -291,31 +291,52 @@ def max_term_index
           .max || -1
       end
 
-      def register_rust_corpus_session
+      def register_rust_corpus_session(previous_session_id = nil)
         @rust_corpus_session_id = nil
         @rust_corpus_terms = nil
         @rust_document_lengths = nil
         @rust_document_words = nil
         @rust_document_counts = nil
 
+        if @corpus.nil?
+          drop_rust_corpus_session_by_id(previous_session_id)
+          return
+        end
+
         return unless defined?(::Lda::RustBackend)
 
         em_input = rust_em_corpus_input
-        return if em_input.nil?
+        if em_input.nil?
+          drop_rust_corpus_session_by_id(previous_session_id)
+          return
+        end
 
         @rust_corpus_terms = Integer(em_input.fetch(:terms))
         @rust_document_lengths = em_input.fetch(:document_lengths)
         @rust_document_words = em_input.fetch(:document_words)
         @rust_document_counts = em_input.fetch(:document_counts)
-        return unless ::Lda::RustBackend.respond_to?(:create_corpus_session)
 
-        session_id = ::Lda::RustBackend.create_corpus_session(
-          @rust_document_words,
-          @rust_document_counts,
-          Integer(@rust_corpus_terms)
-        )
-        return unless session_id.is_a?(Numeric)
-        return unless session_id.positive?
+        session_id =
+          if ::Lda::RustBackend.respond_to?(:replace_corpus_session)
+            ::Lda::RustBackend.replace_corpus_session(
+              Integer(previous_session_id || 0),
+              @rust_document_words,
+              @rust_document_counts,
+              Integer(@rust_corpus_terms)
+            )
+          elsif ::Lda::RustBackend.respond_to?(:create_corpus_session)
+            drop_rust_corpus_session_by_id(previous_session_id)
+            ::Lda::RustBackend.create_corpus_session(
+              @rust_document_words,
+              @rust_document_counts,
+              Integer(@rust_corpus_terms)
+            )
+          end
+
+        unless session_id.is_a?(Numeric) && session_id.positive?
+          drop_rust_corpus_session_by_id(previous_session_id)
+          return
+        end
 
         @rust_corpus_session_id = Integer(session_id)
       rescue StandardError
@@ -324,13 +345,14 @@ def register_rust_corpus_session
         @rust_document_lengths = nil
         @rust_document_words = nil
         @rust_document_counts = nil
+        drop_rust_corpus_session_by_id(previous_session_id)
       end
 
       def ensure_rust_corpus_session
         has_session_data = @rust_corpus_terms && @rust_document_lengths && @rust_document_words && @rust_document_counts
         return true if has_session_data
 
-        register_rust_corpus_session
+        register_rust_corpus_session(@rust_corpus_session_id)
         @rust_corpus_terms && @rust_document_lengths && @rust_document_words && @rust_document_counts
       rescue StandardError
         false
@@ -345,6 +367,12 @@ def release_rust_corpus_session
         @rust_document_words = nil
         @rust_document_counts = nil
 
+        drop_rust_corpus_session_by_id(session_id)
+      rescue StandardError
+        nil
+      end
+
+      def drop_rust_corpus_session_by_id(session_id)
         return unless session_id
         return unless defined?(::Lda::RustBackend)
         return unless ::Lda::RustBackend.respond_to?(:drop_corpus_session)

test/rust_orchestration_test.rb

@@ -508,6 +508,55 @@ def test_run_em_on_session_with_corpus_recreates_missing_session
     end
   end
 
+  def test_replace_corpus_session_updates_existing_session_in_place
+    omit("create_corpus_session unavailable") unless Lda::RustBackend.respond_to?(:create_corpus_session)
+    omit("replace_corpus_session unavailable") unless Lda::RustBackend.respond_to?(:replace_corpus_session)
+    omit("drop_corpus_session unavailable") unless Lda::RustBackend.respond_to?(:drop_corpus_session)
+    omit("run_em_on_session unavailable") unless Lda::RustBackend.respond_to?(:run_em_on_session)
+
+    session_id = Lda::RustBackend.create_corpus_session(@document_words, @document_counts, @terms)
+    assert_operator session_id, :>, 0
+
+    replacement_words = [[0, 1, 2], [2, 3]]
+    replacement_counts = [[2.0, 1.0, 1.0], [1.0, 4.0]]
+    replacement_terms = 5
+    starting_count = Lda::RustBackend.corpus_session_count if Lda::RustBackend.respond_to?(:corpus_session_count)
+
+    replaced_session_id = Lda::RustBackend.replace_corpus_session(
+      session_id,
+      replacement_words,
+      replacement_counts,
+      replacement_terms
+    )
+
+    assert_equal session_id, replaced_session_id
+    if !starting_count.nil?
+      assert_equal starting_count, Lda::RustBackend.corpus_session_count
+    end
+
+    output = Lda::RustBackend.run_em_on_session(
+      replaced_session_id,
+      "seeded",
+      @topics,
+      @max_iter,
+      @convergence,
+      @em_max_iter,
+      @em_convergence,
+      @init_alpha,
+      @min_probability,
+      91_919
+    )
+
+    assert_equal replacement_terms, output[0].first.size
+    assert_equal replacement_words.size, output[2].size
+  ensure
+    if defined?(replaced_session_id) && replaced_session_id.is_a?(Numeric) && replaced_session_id.positive?
+      Lda::RustBackend.drop_corpus_session(replaced_session_id)
+    elsif defined?(session_id) && session_id.is_a?(Numeric) && session_id.positive?
+      Lda::RustBackend.drop_corpus_session(session_id)
+    end
+  end
+
   def test_rust_backend_corpus_session_lifecycle_no_leak
     omit("corpus_session_count unavailable") unless Lda::RustBackend.respond_to?(:corpus_session_count)
 
@@ -516,9 +565,16 @@ def test_rust_backend_corpus_session_lifecycle_no_leak
 
     backend.corpus = Lda::TextCorpus.new(FIXTURE_DOCUMENTS)
     assert_equal starting_count + 1, Lda::RustBackend.corpus_session_count
+    first_session_id = backend.instance_variable_get(:@rust_corpus_session_id)
+    assert_operator first_session_id, :>, 0
 
     backend.corpus = Lda::TextCorpus.new(FIXTURE_DOCUMENTS.reverse)
     assert_equal starting_count + 1, Lda::RustBackend.corpus_session_count
+    second_session_id = backend.instance_variable_get(:@rust_corpus_session_id)
+    assert_operator second_session_id, :>, 0
+    if Lda::RustBackend.respond_to?(:replace_corpus_session)
+      assert_equal first_session_id, second_session_id
+    end
 
     backend.corpus = nil
     assert_equal starting_count, Lda::RustBackend.corpus_session_count