From 375876fde5d974551288ae3c6471c2171134f501 Mon Sep 17 00:00:00 2001 From: Rowan Copley Date: Wed, 14 Jan 2026 14:10:12 -0800 Subject: [PATCH 1/4] Fix Close() race conditions in PooledHugotEmbedder Two critical bugs were fixed: Bug 1: Close() during Embed() caused SIGSEGV - Close() could destroy the ONNX session while Embed() was still using it - Fix: Use WaitGroup to track in-flight operations, Close() waits for them Bug 2: Multiple Close() calls caused panic - Calling Close() multiple times would call session.Destroy() multiple times - Fix: Use sync.Once to ensure Close() only executes once Changes: - Added synchronization fields to PooledHugotEmbedder: - closed (atomic.Bool): prevents new Embed() after Close() - wg (sync.WaitGroup): tracks in-flight Embed() calls - closeOnce (sync.Once): ensures Close() runs exactly once - closeErr (error): stores error from Close() - Updated Embed() to check closed flag and register with WaitGroup - Updated Close() to use sync.Once and wait for in-flight operations Tests added: - close_race_test.go: Tests for both bug scenarios - happy_path_test.go: 7 tests for normal usage patterns - pipeline_collision_test.go: Tests for pipeline selection Co-Authored-By: Claude Opus 4.5 --- pkg/termite/lib/embeddings/batch_test.go | 2 + pkg/termite/lib/embeddings/close_race_test.go | 321 ++++++++++++++++ pkg/termite/lib/embeddings/happy_path_test.go | 294 +++++++++++++++ pkg/termite/lib/embeddings/hugot.go | 67 +++- .../lib/embeddings/pipeline_collision_test.go | 344 ++++++++++++++++++ 5 files changed, 1012 insertions(+), 16 deletions(-) create mode 100644 pkg/termite/lib/embeddings/close_race_test.go create mode 100644 pkg/termite/lib/embeddings/happy_path_test.go create mode 100644 pkg/termite/lib/embeddings/pipeline_collision_test.go diff --git a/pkg/termite/lib/embeddings/batch_test.go b/pkg/termite/lib/embeddings/batch_test.go index 8170603..bb700e9 100644 --- a/pkg/termite/lib/embeddings/batch_test.go +++ b/pkg/termite/lib/embeddings/batch_test.go @@ -191,6 +191,8 @@ func findModelPath(t *testing.T) string { // Check common locations paths := []string{ + filepath.Join(os.Getenv("HOME"), ".termite/models/embedders/BAAI/bge-small-en-v1.5"), + filepath.Join(os.Getenv("HOME"), ".termite/models/embedders/bge-small-en-v1.5"), filepath.Join(os.Getenv("HOME"), ".cache/termite/models/BAAI--bge-small-en-v1.5"), "./models/BAAI--bge-small-en-v1.5", "../../../../../models/embedders/BAAI/bge-small-en-v1.5", diff --git a/pkg/termite/lib/embeddings/close_race_test.go b/pkg/termite/lib/embeddings/close_race_test.go new file mode 100644 index 0000000..73039b4 --- /dev/null +++ b/pkg/termite/lib/embeddings/close_race_test.go @@ -0,0 +1,321 @@ +// Copyright 2025 Antfly, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//go:build onnx && ORT + +package embeddings + +import ( + "context" + "fmt" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/antflydb/antfly-go/libaf/ai" + "go.uber.org/zap" +) + +// TestCloseWhileEmbedding tests the race condition between Close() and Embed(). +// +// Code inspection found that Close() has no synchronization with Embed(): +// +// func (p *PooledHugotEmbedder) Close() error { +// if p.session != nil && !p.sessionShared { +// return p.session.Destroy() // No lock, no WaitGroup +// } +// return nil +// } +// +// Hypothesis: If Close() is called while Embed() is running, the session +// could be destroyed mid-inference, causing undefined behavior. +// +// This test verifies whether: +// 1. Hugot/ONNX has internal protection (blocks Destroy until inference completes) +// 2. Or the race causes crashes/panics +// 3. Or silent corruption occurs +func TestCloseWhileEmbedding(t *testing.T) { + modelPath := findModelPath(t) + if modelPath == "" { + t.Skip("Model not found, skipping close race test") + } + + const poolSize = 2 + logger := zap.NewNop() + + // Create embedder that OWNS its session (sessionShared=false) + // This is important - shared sessions won't trigger the race because + // Close() skips session.Destroy() for shared sessions. + embedder, err := NewPooledHugotEmbedder(modelPath, "model.onnx", poolSize, logger) + if err != nil { + t.Fatalf("Failed to create embedder: %v", err) + } + + // Track what happens + var embedStarted sync.WaitGroup + var embedErr atomic.Value // stores error + var embedPanicked atomic.Bool + var closeErr atomic.Value + + embedStarted.Add(1) + ctx := context.Background() + + // Start a slow embed operation with many texts + go func() { + defer func() { + if r := recover(); r != nil { + embedPanicked.Store(true) + t.Logf("Embed PANICKED: %v", r) + } + }() + + // Use multiple texts to make inference take longer + contents := make([][]ai.ContentPart, 50) + for i := 0; i < len(contents); i++ { + contents[i] = []ai.ContentPart{ + ai.TextContent{Text: fmt.Sprintf("This is test sentence number %d for the close race test. We want inference to take a while.", i)}, + } + } + + embedStarted.Done() // Signal that we're about to start + + _, err := embedder.Embed(ctx, contents) + if err != nil { + embedErr.Store(err) + t.Logf("Embed returned error: %v", err) + } else { + t.Log("Embed completed successfully (no error)") + } + }() + + // Wait for embed to start, then close immediately + embedStarted.Wait() + time.Sleep(10 * time.Millisecond) // Let inference begin + + t.Log("Calling Close() while Embed() is running...") + if err := embedder.Close(); err != nil { + closeErr.Store(err) + t.Logf("Close returned error: %v", err) + } else { + t.Log("Close completed successfully (no error)") + } + + // Wait a bit for embed to finish or crash + time.Sleep(2 * time.Second) + + // Report findings + if embedPanicked.Load() { + t.Error("BUG CONFIRMED: Embed panicked when Close was called during inference") + t.Log("Severity: HIGH - session.Destroy() is immediate and causes crash") + } else if e := embedErr.Load(); e != nil { + t.Logf("Embed returned error after Close: %v", e) + t.Log("This could indicate partial protection or timing-dependent behavior") + } else { + t.Log("No panic or error detected - possible scenarios:") + t.Log(" 1. Hugot's session.Destroy() blocks until inference completes (safe)") + t.Log(" 2. ONNX Runtime has internal reference counting (safe)") + t.Log(" 3. We got lucky with timing (race exists but wasn't triggered)") + } +} + +// TestCloseWhileEmbeddingStress runs many iterations to increase chance +// of triggering the race condition. +func TestCloseWhileEmbeddingStress(t *testing.T) { + modelPath := findModelPath(t) + if modelPath == "" { + t.Skip("Model not found, skipping close race stress test") + } + + const iterations = 20 + const poolSize = 2 + logger := zap.NewNop() + + var panicCount atomic.Int32 + var errorCount atomic.Int32 + var successCount atomic.Int32 + + for iter := 0; iter < iterations; iter++ { + t.Run(fmt.Sprintf("iter_%d", iter), func(t *testing.T) { + embedder, err := NewPooledHugotEmbedder(modelPath, "model.onnx", poolSize, logger) + if err != nil { + t.Fatalf("Failed to create embedder: %v", err) + } + + var embedDone sync.WaitGroup + embedDone.Add(1) + + var panicked atomic.Bool + var embedError atomic.Value + + ctx := context.Background() + + go func() { + defer embedDone.Done() + defer func() { + if r := recover(); r != nil { + panicked.Store(true) + panicCount.Add(1) + } + }() + + contents := make([][]ai.ContentPart, 20) + for i := 0; i < len(contents); i++ { + contents[i] = []ai.ContentPart{ + ai.TextContent{Text: fmt.Sprintf("stress test sentence %d for iteration %d", i, iter)}, + } + } + + _, err := embedder.Embed(ctx, contents) + if err != nil { + embedError.Store(err) + errorCount.Add(1) + } else { + successCount.Add(1) + } + }() + + // Variable delay to hit different points in the inference + delay := time.Duration(iter%10) * time.Millisecond + time.Sleep(delay) + + // Close while embed is (probably) running + _ = embedder.Close() + + embedDone.Wait() + + if panicked.Load() { + t.Errorf("Iteration %d: PANIC detected", iter) + } + }) + } + + t.Logf("Summary: %d panics, %d errors, %d successes out of %d iterations", + panicCount.Load(), errorCount.Load(), successCount.Load(), iterations) + + if panicCount.Load() > 0 { + t.Errorf("BUG CONFIRMED: %d panics detected - Close() race is dangerous", panicCount.Load()) + } +} + +// TestMultipleCloseIsSafe verifies that calling Close() multiple times +// doesn't cause issues (tests assumption A5: session destroy idempotence). +func TestMultipleCloseIsSafe(t *testing.T) { + modelPath := findModelPath(t) + if modelPath == "" { + t.Skip("Model not found, skipping multiple close test") + } + + const poolSize = 2 + logger := zap.NewNop() + + embedder, err := NewPooledHugotEmbedder(modelPath, "model.onnx", poolSize, logger) + if err != nil { + t.Fatalf("Failed to create embedder: %v", err) + } + + // Do one successful embed first + ctx := context.Background() + contents := [][]ai.ContentPart{ + {ai.TextContent{Text: "test before close"}}, + } + _, err = embedder.Embed(ctx, contents) + if err != nil { + t.Fatalf("Initial embed failed: %v", err) + } + + // Now close multiple times + var panicCount atomic.Int32 + var wg sync.WaitGroup + + for i := 0; i < 5; i++ { + wg.Add(1) + go func(attempt int) { + defer wg.Done() + defer func() { + if r := recover(); r != nil { + panicCount.Add(1) + t.Logf("Close attempt %d PANICKED: %v", attempt, r) + } + }() + + err := embedder.Close() + if err != nil { + t.Logf("Close attempt %d returned error: %v", attempt, err) + } else { + t.Logf("Close attempt %d succeeded", attempt) + } + }(i) + } + + wg.Wait() + + if panicCount.Load() > 0 { + t.Errorf("Multiple Close() calls caused %d panics - assumption A5 violated", panicCount.Load()) + } else { + t.Log("Multiple Close() calls are safe (assumption A5 validated)") + } +} + +// TestEmbedAfterClose verifies behavior when Embed is called after Close. +// This tests the "use-after-close" scenario from the TLA+ model. +func TestEmbedAfterClose(t *testing.T) { + modelPath := findModelPath(t) + if modelPath == "" { + t.Skip("Model not found, skipping embed-after-close test") + } + + const poolSize = 2 + logger := zap.NewNop() + + embedder, err := NewPooledHugotEmbedder(modelPath, "model.onnx", poolSize, logger) + if err != nil { + t.Fatalf("Failed to create embedder: %v", err) + } + + // Close first + err = embedder.Close() + if err != nil { + t.Fatalf("Close failed: %v", err) + } + t.Log("Embedder closed") + + // Now try to embed + var panicked atomic.Bool + defer func() { + if r := recover(); r != nil { + panicked.Store(true) + t.Logf("Embed after Close PANICKED: %v", r) + } + }() + + ctx := context.Background() + contents := [][]ai.ContentPart{ + {ai.TextContent{Text: "test after close"}}, + } + + _, err = embedder.Embed(ctx, contents) + if err != nil { + t.Logf("Embed after Close returned error: %v", err) + t.Log("This is expected behavior - the embedder gracefully rejects calls after Close") + } else { + t.Log("Embed after Close succeeded - the session destruction may have been skipped (sessionShared=true) or ONNX allows this") + } + + if panicked.Load() { + t.Error("BUG: Embed after Close caused a panic") + } else { + t.Log("Embed after Close did not panic (some level of protection exists)") + } +} diff --git a/pkg/termite/lib/embeddings/happy_path_test.go b/pkg/termite/lib/embeddings/happy_path_test.go new file mode 100644 index 0000000..87c93ae --- /dev/null +++ b/pkg/termite/lib/embeddings/happy_path_test.go @@ -0,0 +1,294 @@ +//go:build onnx && ORT + +// Happy Path E2E Test for PooledHugotEmbedder +// +// This test verifies normal usage patterns work correctly. +// It should PASS both before and after the bugfixes. +// +// Run: +// export ONNXRUNTIME_ROOT=$PWD/onnxruntime +// export DYLD_LIBRARY_PATH=$ONNXRUNTIME_ROOT/darwin-arm64/lib:$DYLD_LIBRARY_PATH +// go test -v -tags="onnx,ORT" -run TestHappyPath ./pkg/termite/lib/embeddings/ + +package embeddings + +import ( + "context" + "fmt" + "sync" + "testing" + "time" + + "github.com/antflydb/antfly-go/libaf/ai" + "go.uber.org/zap" +) + +// TestHappyPath_SingleEmbed tests basic single-threaded usage. +func TestHappyPath_SingleEmbed(t *testing.T) { + modelPath := findModelPath(t) + if modelPath == "" { + t.Skip("Model not found") + } + logger := zap.NewNop() + + embedder, err := NewPooledHugotEmbedder(modelPath, "model.onnx", 2, logger) + if err != nil { + t.Fatalf("Failed to create embedder: %v", err) + } + defer embedder.Close() + + ctx := context.Background() + contents := [][]ai.ContentPart{ + {ai.TextContent{Text: "Hello world"}}, + {ai.TextContent{Text: "This is a test"}}, + {ai.TextContent{Text: "Embeddings are useful"}}, + } + + result, err := embedder.Embed(ctx, contents) + if err != nil { + t.Fatalf("Embed failed: %v", err) + } + + if len(result) != 3 { + t.Errorf("Expected 3 embeddings, got %d", len(result)) + } + + for i, emb := range result { + if len(emb) == 0 { + t.Errorf("Embedding %d is empty", i) + } + t.Logf("Embedding %d: dim=%d, first_val=%.4f", i, len(emb), emb[0]) + } +} + +// TestHappyPath_MultipleSequentialEmbeds tests multiple sequential calls. +func TestHappyPath_MultipleSequentialEmbeds(t *testing.T) { + modelPath := findModelPath(t) + if modelPath == "" { + t.Skip("Model not found") + } + logger := zap.NewNop() + + embedder, err := NewPooledHugotEmbedder(modelPath, "model.onnx", 2, logger) + if err != nil { + t.Fatalf("Failed to create embedder: %v", err) + } + defer embedder.Close() + + ctx := context.Background() + + for i := 0; i < 5; i++ { + contents := [][]ai.ContentPart{ + {ai.TextContent{Text: fmt.Sprintf("Sequential test %d", i)}}, + } + + result, err := embedder.Embed(ctx, contents) + if err != nil { + t.Fatalf("Embed %d failed: %v", i, err) + } + + if len(result) != 1 { + t.Errorf("Embed %d: expected 1 result, got %d", i, len(result)) + } + } + + t.Log("5 sequential embeds completed successfully") +} + +// TestHappyPath_ConcurrentEmbeds tests concurrent usage within pool limits. +func TestHappyPath_ConcurrentEmbeds(t *testing.T) { + modelPath := findModelPath(t) + if modelPath == "" { + t.Skip("Model not found") + } + logger := zap.NewNop() + + poolSize := 2 + embedder, err := NewPooledHugotEmbedder(modelPath, "model.onnx", poolSize, logger) + if err != nil { + t.Fatalf("Failed to create embedder: %v", err) + } + defer embedder.Close() + + ctx := context.Background() + numWorkers := 10 + embedsPerWorker := 5 + + var wg sync.WaitGroup + errors := make(chan error, numWorkers*embedsPerWorker) + + start := time.Now() + + for w := 0; w < numWorkers; w++ { + wg.Add(1) + go func(workerID int) { + defer wg.Done() + + for i := 0; i < embedsPerWorker; i++ { + contents := [][]ai.ContentPart{ + {ai.TextContent{Text: fmt.Sprintf("Worker %d embed %d", workerID, i)}}, + } + + _, err := embedder.Embed(ctx, contents) + if err != nil { + errors <- fmt.Errorf("worker %d embed %d: %w", workerID, i, err) + } + } + }(w) + } + + wg.Wait() + close(errors) + + duration := time.Since(start) + totalOps := numWorkers * embedsPerWorker + + var errs []error + for err := range errors { + errs = append(errs, err) + } + + if len(errs) > 0 { + for _, err := range errs { + t.Errorf("Error: %v", err) + } + t.Fatalf("%d errors occurred", len(errs)) + } + + t.Logf("%d concurrent embeds completed in %v (%.1f ops/sec)", + totalOps, duration, float64(totalOps)/duration.Seconds()) +} + +// TestHappyPath_CloseAfterAllComplete tests proper close after work is done. +func TestHappyPath_CloseAfterAllComplete(t *testing.T) { + modelPath := findModelPath(t) + if modelPath == "" { + t.Skip("Model not found") + } + logger := zap.NewNop() + + embedder, err := NewPooledHugotEmbedder(modelPath, "model.onnx", 2, logger) + if err != nil { + t.Fatalf("Failed to create embedder: %v", err) + } + + ctx := context.Background() + + // Do some work + for i := 0; i < 3; i++ { + contents := [][]ai.ContentPart{ + {ai.TextContent{Text: fmt.Sprintf("Test %d", i)}}, + } + _, err := embedder.Embed(ctx, contents) + if err != nil { + t.Fatalf("Embed %d failed: %v", i, err) + } + } + + // Close after all work is done - should succeed + err = embedder.Close() + if err != nil { + t.Errorf("Close() returned error: %v", err) + } + + t.Log("Close() after all embeds complete: success") +} + +// TestHappyPath_LargeBatch tests handling of larger batches. +func TestHappyPath_LargeBatch(t *testing.T) { + modelPath := findModelPath(t) + if modelPath == "" { + t.Skip("Model not found") + } + logger := zap.NewNop() + + embedder, err := NewPooledHugotEmbedder(modelPath, "model.onnx", 2, logger) + if err != nil { + t.Fatalf("Failed to create embedder: %v", err) + } + defer embedder.Close() + + ctx := context.Background() + + // Create a batch of 20 texts + batchSize := 20 + contents := make([][]ai.ContentPart, batchSize) + for i := 0; i < batchSize; i++ { + contents[i] = []ai.ContentPart{ + ai.TextContent{Text: fmt.Sprintf("Large batch test sentence number %d with some extra text", i)}, + } + } + + start := time.Now() + result, err := embedder.Embed(ctx, contents) + duration := time.Since(start) + + if err != nil { + t.Fatalf("Large batch embed failed: %v", err) + } + + if len(result) != batchSize { + t.Errorf("Expected %d embeddings, got %d", batchSize, len(result)) + } + + t.Logf("Batch of %d texts embedded in %v", batchSize, duration) +} + +// TestHappyPath_ContextCancellation tests that context cancellation is handled. +func TestHappyPath_ContextCancellation(t *testing.T) { + modelPath := findModelPath(t) + if modelPath == "" { + t.Skip("Model not found") + } + logger := zap.NewNop() + + embedder, err := NewPooledHugotEmbedder(modelPath, "model.onnx", 2, logger) + if err != nil { + t.Fatalf("Failed to create embedder: %v", err) + } + defer embedder.Close() + + // Create already-cancelled context + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + contents := [][]ai.ContentPart{ + {ai.TextContent{Text: "This should fail due to cancelled context"}}, + } + + _, err = embedder.Embed(ctx, contents) + if err == nil { + t.Error("Expected error with cancelled context, got nil") + } else { + t.Logf("Cancelled context correctly returned error: %v", err) + } +} + +// TestHappyPath_EmptyInput tests handling of empty input. +func TestHappyPath_EmptyInput(t *testing.T) { + modelPath := findModelPath(t) + if modelPath == "" { + t.Skip("Model not found") + } + logger := zap.NewNop() + + embedder, err := NewPooledHugotEmbedder(modelPath, "model.onnx", 2, logger) + if err != nil { + t.Fatalf("Failed to create embedder: %v", err) + } + defer embedder.Close() + + ctx := context.Background() + contents := [][]ai.ContentPart{} + + result, err := embedder.Embed(ctx, contents) + if err != nil { + t.Errorf("Empty input should not error: %v", err) + } + + if len(result) != 0 { + t.Errorf("Expected 0 results for empty input, got %d", len(result)) + } + + t.Log("Empty input handled correctly") +} diff --git a/pkg/termite/lib/embeddings/hugot.go b/pkg/termite/lib/embeddings/hugot.go index 22743b2..dda77de 100644 --- a/pkg/termite/lib/embeddings/hugot.go +++ b/pkg/termite/lib/embeddings/hugot.go @@ -20,6 +20,7 @@ import ( "fmt" "math" "runtime" + "sync" "sync/atomic" "github.com/antflydb/antfly-go/libaf/ai" @@ -79,6 +80,12 @@ type PooledHugotEmbedder struct { poolSize int caps embeddings.EmbedderCapabilities batchSize int + + // Synchronization for safe Close() behavior + closed atomic.Bool // Prevents new Embed() calls after Close() + wg sync.WaitGroup // Waits for in-flight Embed() calls to complete + closeOnce sync.Once // Ensures Close() runs exactly once + closeErr error // Stores error from Close() } // NewPooledHugotEmbedder creates a new pooled embedder using the Hugot ONNX runtime. @@ -290,10 +297,27 @@ func (p *PooledHugotEmbedder) Capabilities() embeddings.EmbedderCapabilities { return p.caps } +// ErrEmbedderClosed is returned when Embed is called on a closed embedder. +var ErrEmbedderClosed = errors.New("embedder is closed") + // Embed generates embeddings for the given content. // Thread-safe: uses semaphore to limit concurrent pipeline access. // Processes texts in batches to avoid memory explosion on CoreML. func (p *PooledHugotEmbedder) Embed(ctx context.Context, contents [][]ai.ContentPart) ([][]float32, error) { + // Check if embedder is closed before starting + if p.closed.Load() { + return nil, ErrEmbedderClosed + } + + // Track this in-flight operation so Close() waits for us + p.wg.Add(1) + defer p.wg.Done() + + // Double-check after registration (handles race with Close()) + if p.closed.Load() { + return nil, ErrEmbedderClosed + } + if len(contents) == 0 { return [][]float32{}, nil } @@ -375,23 +399,34 @@ func (p *PooledHugotEmbedder) Embed(ctx context.Context, contents [][]ai.Content // Close releases resources. // Properly closes each pipeline to remove it from the session, then destroys // the session if it was created by this embedder (not shared). +// Thread-safe: waits for in-flight Embed() calls to complete before destroying. +// Safe to call multiple times (only the first call takes effect). func (p *PooledHugotEmbedder) Close() error { - // Close each pipeline to remove it from the session - for _, pipeline := range p.pipelines { - if pipeline != nil { - name := pipeline.PipelineName - if err := khugot.ClosePipeline[*pipelines.FeatureExtractionPipeline](p.session, name); err != nil { - p.logger.Warn("Failed to close pipeline", zap.String("name", name), zap.Error(err)) + p.closeOnce.Do(func() { + // Set closed flag to prevent new Embed() calls + p.closed.Store(true) + + // Wait for all in-flight Embed() calls to complete + p.wg.Wait() + + // Close each pipeline to remove it from the session + for _, pipeline := range p.pipelines { + if pipeline != nil { + name := pipeline.PipelineName + if err := khugot.ClosePipeline[*pipelines.FeatureExtractionPipeline](p.session, name); err != nil { + p.logger.Warn("Failed to close pipeline", zap.String("name", name), zap.Error(err)) + } } } - } - p.pipelines = nil - - if p.session != nil && !p.sessionShared { - p.logger.Info("Destroying Hugot session (owned by this pooled embedder)") - return p.session.Destroy() - } else if p.sessionShared { - p.logger.Debug("Skipping session destruction (shared session)") - } - return nil + p.pipelines = nil + + // Now safe to destroy the session + if p.session != nil && !p.sessionShared { + p.logger.Info("Destroying Hugot session (owned by this pooled embedder)") + p.closeErr = p.session.Destroy() + } else if p.sessionShared { + p.logger.Debug("Skipping session destruction (shared session)") + } + }) + return p.closeErr } diff --git a/pkg/termite/lib/embeddings/pipeline_collision_test.go b/pkg/termite/lib/embeddings/pipeline_collision_test.go new file mode 100644 index 0000000..1af9be6 --- /dev/null +++ b/pkg/termite/lib/embeddings/pipeline_collision_test.go @@ -0,0 +1,344 @@ +// Copyright 2025 Antfly, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//go:build onnx && ORT + +package embeddings + +import ( + "context" + "fmt" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/antflydb/antfly-go/libaf/ai" + "go.uber.org/zap" +) + +// TestPipelineCollision validates the TLA+ counterexample showing two workers +// can be assigned the same pipeline due to round-robin counter wrapping. +// +// Bug mechanism (hugot.go:307-309): +// idx := int(p.nextPipeline.Add(1) % uint64(p.poolSize)) +// pipeline := p.pipelines[idx] +// +// The semaphore limits concurrent users to poolSize, but if a worker finishes +// quickly and another starts before the first slow worker completes, they can +// both get the same pipeline index. +// +// Counterexample trace (poolSize=2, 3 workers): +// 1. w1: acquires sem, nextPipeline=1, idx=1%2=1 +// 2. w2: acquires sem, nextPipeline=2, idx=2%2=0 +// 3. w2: completes quickly, releases sem +// 4. w3: acquires freed sem slot +// 5. w3: nextPipeline=3, idx=3%2=1 <-- COLLISION with w1! +func TestPipelineCollision(t *testing.T) { + modelPath := findModelPath(t) + if modelPath == "" { + t.Skip("Model not found, skipping pipeline collision test") + } + + const poolSize = 2 + const numWorkers = 10 // More workers = more collision opportunities + const iterations = 50 // Run multiple times to catch race + + logger := zap.NewNop() + + for iter := 0; iter < iterations; iter++ { + t.Run(fmt.Sprintf("iteration_%d", iter), func(t *testing.T) { + embedder, err := NewPooledHugotEmbedder(modelPath, "model.onnx", poolSize, logger) + if err != nil { + t.Fatalf("Failed to create embedder: %v", err) + } + defer embedder.Close() + + // Track collision detection + var collisionCount atomic.Int32 + var collisionDetails sync.Map // For debugging: stores collision info + // Note: Direct pipeline tracking is not possible without modifying production code + // The race detector (-race flag) will catch actual concurrent access + + var wg sync.WaitGroup + ctx := context.Background() + + // Start multiple workers concurrently + for w := 0; w < numWorkers; w++ { + wg.Add(1) + workerID := w + go func() { + defer wg.Done() + + // Simple input - just need to trigger pipeline selection + contents := [][]ai.ContentPart{ + {ai.TextContent{Text: "test sentence for collision detection"}}, + } + + // Get the pipeline index that will be selected + // We can predict this from nextPipeline, but that's racy + // Instead, we'll detect collision by checking pipelineUsers + + // Record that we're about to use a pipeline + // The actual pipeline selection happens inside Embed(), which we can't intercept + // So we approximate by checking if another goroutine is also embedding + + // Mark entry + startTime := time.Now() + + // Actually call Embed - this is where collision would manifest + _, err := embedder.Embed(ctx, contents) + if err != nil { + t.Logf("Worker %d: Embed error: %v", workerID, err) + } + + duration := time.Since(startTime) + t.Logf("Worker %d: completed in %v", workerID, duration) + }() + } + + wg.Wait() + + if collisionCount.Load() > 0 { + t.Errorf("Detected %d pipeline collisions!", collisionCount.Load()) + collisionDetails.Range(func(key, value any) bool { + t.Logf("Collision detail: %v", value) + return true + }) + } + }) + } +} + +// TestPipelineCollisionWithInstrumentation uses a more sophisticated approach +// to detect collisions by monitoring the nextPipeline counter and timing. +func TestPipelineCollisionWithInstrumentation(t *testing.T) { + modelPath := findModelPath(t) + if modelPath == "" { + t.Skip("Model not found, skipping instrumented collision test") + } + + const poolSize = 2 + logger := zap.NewNop() + + embedder, err := NewPooledHugotEmbedder(modelPath, "model.onnx", poolSize, logger) + if err != nil { + t.Fatalf("Failed to create embedder: %v", err) + } + defer embedder.Close() + + // Track pipeline usage by monitoring the atomic counter + // Since we can't directly observe which pipeline each goroutine gets, + // we instead track timing to detect overlapping usage + + type usageRecord struct { + workerID int + pipelineIdx int + startTime time.Time + endTime time.Time + } + + var records []usageRecord + var recordsMu sync.Mutex + var wg sync.WaitGroup + ctx := context.Background() + + // Strategy: Capture nextPipeline before and after Embed to determine + // which pipeline index was used. This is racy but gives us insight. + + numWorkers := 20 + for w := 0; w < numWorkers; w++ { + wg.Add(1) + workerID := w + go func() { + defer wg.Done() + + // Capture counter before call + // Note: This is inherently racy - another goroutine could increment between + // our read and the actual Add(1) inside Embed. But it's good enough for testing. + beforeCounter := embedder.nextPipeline.Load() + + startTime := time.Now() + + contents := [][]ai.ContentPart{ + {ai.TextContent{Text: "test sentence for instrumented collision detection"}}, + } + _, err := embedder.Embed(ctx, contents) + if err != nil { + t.Logf("Worker %d: error: %v", workerID, err) + } + + endTime := time.Now() + + // The pipeline index used was likely (beforeCounter+1) % poolSize + // This is approximate due to races + pipelineIdx := int((beforeCounter + 1) % uint64(poolSize)) + + recordsMu.Lock() + records = append(records, usageRecord{ + workerID: workerID, + pipelineIdx: pipelineIdx, + startTime: startTime, + endTime: endTime, + }) + recordsMu.Unlock() + }() + } + + wg.Wait() + + // Analyze records for overlapping usage of the same pipeline + collisions := 0 + for i := 0; i < len(records); i++ { + for j := i + 1; j < len(records); j++ { + r1, r2 := records[i], records[j] + + // Check if same pipeline and overlapping time + if r1.pipelineIdx == r2.pipelineIdx { + // Check for overlap: r1.start < r2.end AND r2.start < r1.end + if r1.startTime.Before(r2.endTime) && r2.startTime.Before(r1.endTime) { + collisions++ + t.Logf("POTENTIAL COLLISION: worker %d (pipeline %d, %v-%v) overlaps with worker %d (pipeline %d, %v-%v)", + r1.workerID, r1.pipelineIdx, r1.startTime.Format("15:04:05.000"), r1.endTime.Format("15:04:05.000"), + r2.workerID, r2.pipelineIdx, r2.startTime.Format("15:04:05.000"), r2.endTime.Format("15:04:05.000")) + } + } + } + } + + if collisions > 0 { + t.Errorf("Detected %d potential pipeline collisions (may include false positives due to timing approximation)", collisions) + } else { + t.Logf("No collisions detected in %d operations", numWorkers) + } +} + +// TestPipelineCollisionStress runs many concurrent operations to try to trigger +// the race condition and relies on the -race flag to detect data races. +// +// Run with: go test -v -race -tags="onnx,ORT" -run TestPipelineCollisionStress +func TestPipelineCollisionStress(t *testing.T) { + modelPath := findModelPath(t) + if modelPath == "" { + t.Skip("Model not found, skipping stress test") + } + + const poolSize = 2 + const numGoroutines = 50 + const opsPerGoroutine = 10 + + logger := zap.NewNop() + + embedder, err := NewPooledHugotEmbedder(modelPath, "model.onnx", poolSize, logger) + if err != nil { + t.Fatalf("Failed to create embedder: %v", err) + } + defer embedder.Close() + + var wg sync.WaitGroup + var errorCount atomic.Int32 + ctx := context.Background() + + start := time.Now() + + for g := 0; g < numGoroutines; g++ { + wg.Add(1) + go func(goroutineID int) { + defer wg.Done() + + for op := 0; op < opsPerGoroutine; op++ { + contents := [][]ai.ContentPart{ + {ai.TextContent{Text: fmt.Sprintf("stress test sentence %d-%d", goroutineID, op)}}, + } + + _, err := embedder.Embed(ctx, contents) + if err != nil { + errorCount.Add(1) + // Don't log every error to avoid spam + } + } + }(g) + } + + wg.Wait() + duration := time.Since(start) + + totalOps := numGoroutines * opsPerGoroutine + t.Logf("Completed %d operations in %v (%.1f ops/sec)", totalOps, duration, float64(totalOps)/duration.Seconds()) + t.Logf("Errors: %d", errorCount.Load()) + + // If we get here without the race detector firing, either: + // 1. The pipelines are actually thread-safe (contrary to assumption A6) + // 2. We didn't trigger the race condition + // 3. The race exists but wasn't detected + // + // The -race flag should catch concurrent access to the same pipeline + // if it's truly not thread-safe. +} + +// TestFirstEmbedUsesPipelineOne verifies the TLA+ finding that the first +// embed uses pipeline 1, not pipeline 0, due to Add(1) returning the new value. +func TestFirstEmbedUsesPipelineOne(t *testing.T) { + modelPath := findModelPath(t) + if modelPath == "" { + t.Skip("Model not found, skipping first-embed test") + } + + const poolSize = 4 // Use larger pool to make index more obvious + logger := zap.NewNop() + + embedder, err := NewPooledHugotEmbedder(modelPath, "model.onnx", poolSize, logger) + if err != nil { + t.Fatalf("Failed to create embedder: %v", err) + } + defer embedder.Close() + + // Check initial counter value + initialCounter := embedder.nextPipeline.Load() + t.Logf("Initial nextPipeline counter: %d", initialCounter) + + if initialCounter != 0 { + t.Errorf("Expected initial counter to be 0, got %d", initialCounter) + } + + // Perform first embed + ctx := context.Background() + contents := [][]ai.ContentPart{ + {ai.TextContent{Text: "first embed test"}}, + } + _, err = embedder.Embed(ctx, contents) + if err != nil { + t.Fatalf("First embed failed: %v", err) + } + + // Check counter after first embed + afterCounter := embedder.nextPipeline.Load() + t.Logf("After first embed, nextPipeline counter: %d", afterCounter) + + if afterCounter != 1 { + t.Errorf("Expected counter to be 1 after first embed, got %d", afterCounter) + } + + // The pipeline index used was: (0 + 1) % poolSize = 1 % 4 = 1 + // This means pipeline 0 is never used on the first call! + expectedPipelineUsed := int((initialCounter + 1) % uint64(poolSize)) + t.Logf("First embed used pipeline index: %d (pipeline 0 was skipped)", expectedPipelineUsed) + + if expectedPipelineUsed != 1 { + t.Errorf("Expected first embed to use pipeline 1, calculated %d", expectedPipelineUsed) + } + + // This confirms the TLA+ finding - it's a quirk but not a bug + t.Log("CONFIRMED: First embed skips pipeline 0 (uses pipeline 1)") +} From 5f5bf1d98897c57293b328b1d6369bb4151f6c8c Mon Sep 17 00:00:00 2001 From: Rowan Copley Date: Wed, 14 Jan 2026 14:10:12 -0800 Subject: [PATCH 2/4] Fix Close() race conditions in PooledHugotEmbedder Two critical bugs were fixed: Bug 1: Close() during Embed() caused SIGSEGV - Close() could destroy the ONNX session while Embed() was still using it - Fix: Use WaitGroup to track in-flight operations, Close() waits for them Bug 2: Multiple Close() calls caused panic - Calling Close() multiple times would call session.Destroy() multiple times - Fix: Use sync.Once to ensure Close() only executes once Changes: - Added synchronization fields to PooledHugotEmbedder: - closed (atomic.Bool): prevents new Embed() after Close() - wg (sync.WaitGroup): tracks in-flight Embed() calls - closeOnce (sync.Once): ensures Close() runs exactly once - closeErr (error): stores error from Close() - Updated Embed() to check closed flag and register with WaitGroup - Updated Close() to use sync.Once and wait for in-flight operations Tests added: - close_race_test.go: Tests for both bug scenarios - happy_path_test.go: 7 tests for normal usage patterns - pipeline_collision_test.go: Tests for pipeline selection --- pkg/termite/lib/embeddings/batch_test.go | 2 + pkg/termite/lib/embeddings/close_race_test.go | 321 ++++++++++++++++ pkg/termite/lib/embeddings/happy_path_test.go | 294 +++++++++++++++ pkg/termite/lib/embeddings/hugot.go | 67 +++- .../lib/embeddings/pipeline_collision_test.go | 344 ++++++++++++++++++ 5 files changed, 1012 insertions(+), 16 deletions(-) create mode 100644 pkg/termite/lib/embeddings/close_race_test.go create mode 100644 pkg/termite/lib/embeddings/happy_path_test.go create mode 100644 pkg/termite/lib/embeddings/pipeline_collision_test.go diff --git a/pkg/termite/lib/embeddings/batch_test.go b/pkg/termite/lib/embeddings/batch_test.go index 8170603..bb700e9 100644 --- a/pkg/termite/lib/embeddings/batch_test.go +++ b/pkg/termite/lib/embeddings/batch_test.go @@ -191,6 +191,8 @@ func findModelPath(t *testing.T) string { // Check common locations paths := []string{ + filepath.Join(os.Getenv("HOME"), ".termite/models/embedders/BAAI/bge-small-en-v1.5"), + filepath.Join(os.Getenv("HOME"), ".termite/models/embedders/bge-small-en-v1.5"), filepath.Join(os.Getenv("HOME"), ".cache/termite/models/BAAI--bge-small-en-v1.5"), "./models/BAAI--bge-small-en-v1.5", "../../../../../models/embedders/BAAI/bge-small-en-v1.5", diff --git a/pkg/termite/lib/embeddings/close_race_test.go b/pkg/termite/lib/embeddings/close_race_test.go new file mode 100644 index 0000000..73039b4 --- /dev/null +++ b/pkg/termite/lib/embeddings/close_race_test.go @@ -0,0 +1,321 @@ +// Copyright 2025 Antfly, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//go:build onnx && ORT + +package embeddings + +import ( + "context" + "fmt" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/antflydb/antfly-go/libaf/ai" + "go.uber.org/zap" +) + +// TestCloseWhileEmbedding tests the race condition between Close() and Embed(). +// +// Code inspection found that Close() has no synchronization with Embed(): +// +// func (p *PooledHugotEmbedder) Close() error { +// if p.session != nil && !p.sessionShared { +// return p.session.Destroy() // No lock, no WaitGroup +// } +// return nil +// } +// +// Hypothesis: If Close() is called while Embed() is running, the session +// could be destroyed mid-inference, causing undefined behavior. +// +// This test verifies whether: +// 1. Hugot/ONNX has internal protection (blocks Destroy until inference completes) +// 2. Or the race causes crashes/panics +// 3. Or silent corruption occurs +func TestCloseWhileEmbedding(t *testing.T) { + modelPath := findModelPath(t) + if modelPath == "" { + t.Skip("Model not found, skipping close race test") + } + + const poolSize = 2 + logger := zap.NewNop() + + // Create embedder that OWNS its session (sessionShared=false) + // This is important - shared sessions won't trigger the race because + // Close() skips session.Destroy() for shared sessions. + embedder, err := NewPooledHugotEmbedder(modelPath, "model.onnx", poolSize, logger) + if err != nil { + t.Fatalf("Failed to create embedder: %v", err) + } + + // Track what happens + var embedStarted sync.WaitGroup + var embedErr atomic.Value // stores error + var embedPanicked atomic.Bool + var closeErr atomic.Value + + embedStarted.Add(1) + ctx := context.Background() + + // Start a slow embed operation with many texts + go func() { + defer func() { + if r := recover(); r != nil { + embedPanicked.Store(true) + t.Logf("Embed PANICKED: %v", r) + } + }() + + // Use multiple texts to make inference take longer + contents := make([][]ai.ContentPart, 50) + for i := 0; i < len(contents); i++ { + contents[i] = []ai.ContentPart{ + ai.TextContent{Text: fmt.Sprintf("This is test sentence number %d for the close race test. We want inference to take a while.", i)}, + } + } + + embedStarted.Done() // Signal that we're about to start + + _, err := embedder.Embed(ctx, contents) + if err != nil { + embedErr.Store(err) + t.Logf("Embed returned error: %v", err) + } else { + t.Log("Embed completed successfully (no error)") + } + }() + + // Wait for embed to start, then close immediately + embedStarted.Wait() + time.Sleep(10 * time.Millisecond) // Let inference begin + + t.Log("Calling Close() while Embed() is running...") + if err := embedder.Close(); err != nil { + closeErr.Store(err) + t.Logf("Close returned error: %v", err) + } else { + t.Log("Close completed successfully (no error)") + } + + // Wait a bit for embed to finish or crash + time.Sleep(2 * time.Second) + + // Report findings + if embedPanicked.Load() { + t.Error("BUG CONFIRMED: Embed panicked when Close was called during inference") + t.Log("Severity: HIGH - session.Destroy() is immediate and causes crash") + } else if e := embedErr.Load(); e != nil { + t.Logf("Embed returned error after Close: %v", e) + t.Log("This could indicate partial protection or timing-dependent behavior") + } else { + t.Log("No panic or error detected - possible scenarios:") + t.Log(" 1. Hugot's session.Destroy() blocks until inference completes (safe)") + t.Log(" 2. ONNX Runtime has internal reference counting (safe)") + t.Log(" 3. We got lucky with timing (race exists but wasn't triggered)") + } +} + +// TestCloseWhileEmbeddingStress runs many iterations to increase chance +// of triggering the race condition. +func TestCloseWhileEmbeddingStress(t *testing.T) { + modelPath := findModelPath(t) + if modelPath == "" { + t.Skip("Model not found, skipping close race stress test") + } + + const iterations = 20 + const poolSize = 2 + logger := zap.NewNop() + + var panicCount atomic.Int32 + var errorCount atomic.Int32 + var successCount atomic.Int32 + + for iter := 0; iter < iterations; iter++ { + t.Run(fmt.Sprintf("iter_%d", iter), func(t *testing.T) { + embedder, err := NewPooledHugotEmbedder(modelPath, "model.onnx", poolSize, logger) + if err != nil { + t.Fatalf("Failed to create embedder: %v", err) + } + + var embedDone sync.WaitGroup + embedDone.Add(1) + + var panicked atomic.Bool + var embedError atomic.Value + + ctx := context.Background() + + go func() { + defer embedDone.Done() + defer func() { + if r := recover(); r != nil { + panicked.Store(true) + panicCount.Add(1) + } + }() + + contents := make([][]ai.ContentPart, 20) + for i := 0; i < len(contents); i++ { + contents[i] = []ai.ContentPart{ + ai.TextContent{Text: fmt.Sprintf("stress test sentence %d for iteration %d", i, iter)}, + } + } + + _, err := embedder.Embed(ctx, contents) + if err != nil { + embedError.Store(err) + errorCount.Add(1) + } else { + successCount.Add(1) + } + }() + + // Variable delay to hit different points in the inference + delay := time.Duration(iter%10) * time.Millisecond + time.Sleep(delay) + + // Close while embed is (probably) running + _ = embedder.Close() + + embedDone.Wait() + + if panicked.Load() { + t.Errorf("Iteration %d: PANIC detected", iter) + } + }) + } + + t.Logf("Summary: %d panics, %d errors, %d successes out of %d iterations", + panicCount.Load(), errorCount.Load(), successCount.Load(), iterations) + + if panicCount.Load() > 0 { + t.Errorf("BUG CONFIRMED: %d panics detected - Close() race is dangerous", panicCount.Load()) + } +} + +// TestMultipleCloseIsSafe verifies that calling Close() multiple times +// doesn't cause issues (tests assumption A5: session destroy idempotence). +func TestMultipleCloseIsSafe(t *testing.T) { + modelPath := findModelPath(t) + if modelPath == "" { + t.Skip("Model not found, skipping multiple close test") + } + + const poolSize = 2 + logger := zap.NewNop() + + embedder, err := NewPooledHugotEmbedder(modelPath, "model.onnx", poolSize, logger) + if err != nil { + t.Fatalf("Failed to create embedder: %v", err) + } + + // Do one successful embed first + ctx := context.Background() + contents := [][]ai.ContentPart{ + {ai.TextContent{Text: "test before close"}}, + } + _, err = embedder.Embed(ctx, contents) + if err != nil { + t.Fatalf("Initial embed failed: %v", err) + } + + // Now close multiple times + var panicCount atomic.Int32 + var wg sync.WaitGroup + + for i := 0; i < 5; i++ { + wg.Add(1) + go func(attempt int) { + defer wg.Done() + defer func() { + if r := recover(); r != nil { + panicCount.Add(1) + t.Logf("Close attempt %d PANICKED: %v", attempt, r) + } + }() + + err := embedder.Close() + if err != nil { + t.Logf("Close attempt %d returned error: %v", attempt, err) + } else { + t.Logf("Close attempt %d succeeded", attempt) + } + }(i) + } + + wg.Wait() + + if panicCount.Load() > 0 { + t.Errorf("Multiple Close() calls caused %d panics - assumption A5 violated", panicCount.Load()) + } else { + t.Log("Multiple Close() calls are safe (assumption A5 validated)") + } +} + +// TestEmbedAfterClose verifies behavior when Embed is called after Close. +// This tests the "use-after-close" scenario from the TLA+ model. +func TestEmbedAfterClose(t *testing.T) { + modelPath := findModelPath(t) + if modelPath == "" { + t.Skip("Model not found, skipping embed-after-close test") + } + + const poolSize = 2 + logger := zap.NewNop() + + embedder, err := NewPooledHugotEmbedder(modelPath, "model.onnx", poolSize, logger) + if err != nil { + t.Fatalf("Failed to create embedder: %v", err) + } + + // Close first + err = embedder.Close() + if err != nil { + t.Fatalf("Close failed: %v", err) + } + t.Log("Embedder closed") + + // Now try to embed + var panicked atomic.Bool + defer func() { + if r := recover(); r != nil { + panicked.Store(true) + t.Logf("Embed after Close PANICKED: %v", r) + } + }() + + ctx := context.Background() + contents := [][]ai.ContentPart{ + {ai.TextContent{Text: "test after close"}}, + } + + _, err = embedder.Embed(ctx, contents) + if err != nil { + t.Logf("Embed after Close returned error: %v", err) + t.Log("This is expected behavior - the embedder gracefully rejects calls after Close") + } else { + t.Log("Embed after Close succeeded - the session destruction may have been skipped (sessionShared=true) or ONNX allows this") + } + + if panicked.Load() { + t.Error("BUG: Embed after Close caused a panic") + } else { + t.Log("Embed after Close did not panic (some level of protection exists)") + } +} diff --git a/pkg/termite/lib/embeddings/happy_path_test.go b/pkg/termite/lib/embeddings/happy_path_test.go new file mode 100644 index 0000000..87c93ae --- /dev/null +++ b/pkg/termite/lib/embeddings/happy_path_test.go @@ -0,0 +1,294 @@ +//go:build onnx && ORT + +// Happy Path E2E Test for PooledHugotEmbedder +// +// This test verifies normal usage patterns work correctly. +// It should PASS both before and after the bugfixes. +// +// Run: +// export ONNXRUNTIME_ROOT=$PWD/onnxruntime +// export DYLD_LIBRARY_PATH=$ONNXRUNTIME_ROOT/darwin-arm64/lib:$DYLD_LIBRARY_PATH +// go test -v -tags="onnx,ORT" -run TestHappyPath ./pkg/termite/lib/embeddings/ + +package embeddings + +import ( + "context" + "fmt" + "sync" + "testing" + "time" + + "github.com/antflydb/antfly-go/libaf/ai" + "go.uber.org/zap" +) + +// TestHappyPath_SingleEmbed tests basic single-threaded usage. +func TestHappyPath_SingleEmbed(t *testing.T) { + modelPath := findModelPath(t) + if modelPath == "" { + t.Skip("Model not found") + } + logger := zap.NewNop() + + embedder, err := NewPooledHugotEmbedder(modelPath, "model.onnx", 2, logger) + if err != nil { + t.Fatalf("Failed to create embedder: %v", err) + } + defer embedder.Close() + + ctx := context.Background() + contents := [][]ai.ContentPart{ + {ai.TextContent{Text: "Hello world"}}, + {ai.TextContent{Text: "This is a test"}}, + {ai.TextContent{Text: "Embeddings are useful"}}, + } + + result, err := embedder.Embed(ctx, contents) + if err != nil { + t.Fatalf("Embed failed: %v", err) + } + + if len(result) != 3 { + t.Errorf("Expected 3 embeddings, got %d", len(result)) + } + + for i, emb := range result { + if len(emb) == 0 { + t.Errorf("Embedding %d is empty", i) + } + t.Logf("Embedding %d: dim=%d, first_val=%.4f", i, len(emb), emb[0]) + } +} + +// TestHappyPath_MultipleSequentialEmbeds tests multiple sequential calls. +func TestHappyPath_MultipleSequentialEmbeds(t *testing.T) { + modelPath := findModelPath(t) + if modelPath == "" { + t.Skip("Model not found") + } + logger := zap.NewNop() + + embedder, err := NewPooledHugotEmbedder(modelPath, "model.onnx", 2, logger) + if err != nil { + t.Fatalf("Failed to create embedder: %v", err) + } + defer embedder.Close() + + ctx := context.Background() + + for i := 0; i < 5; i++ { + contents := [][]ai.ContentPart{ + {ai.TextContent{Text: fmt.Sprintf("Sequential test %d", i)}}, + } + + result, err := embedder.Embed(ctx, contents) + if err != nil { + t.Fatalf("Embed %d failed: %v", i, err) + } + + if len(result) != 1 { + t.Errorf("Embed %d: expected 1 result, got %d", i, len(result)) + } + } + + t.Log("5 sequential embeds completed successfully") +} + +// TestHappyPath_ConcurrentEmbeds tests concurrent usage within pool limits. +func TestHappyPath_ConcurrentEmbeds(t *testing.T) { + modelPath := findModelPath(t) + if modelPath == "" { + t.Skip("Model not found") + } + logger := zap.NewNop() + + poolSize := 2 + embedder, err := NewPooledHugotEmbedder(modelPath, "model.onnx", poolSize, logger) + if err != nil { + t.Fatalf("Failed to create embedder: %v", err) + } + defer embedder.Close() + + ctx := context.Background() + numWorkers := 10 + embedsPerWorker := 5 + + var wg sync.WaitGroup + errors := make(chan error, numWorkers*embedsPerWorker) + + start := time.Now() + + for w := 0; w < numWorkers; w++ { + wg.Add(1) + go func(workerID int) { + defer wg.Done() + + for i := 0; i < embedsPerWorker; i++ { + contents := [][]ai.ContentPart{ + {ai.TextContent{Text: fmt.Sprintf("Worker %d embed %d", workerID, i)}}, + } + + _, err := embedder.Embed(ctx, contents) + if err != nil { + errors <- fmt.Errorf("worker %d embed %d: %w", workerID, i, err) + } + } + }(w) + } + + wg.Wait() + close(errors) + + duration := time.Since(start) + totalOps := numWorkers * embedsPerWorker + + var errs []error + for err := range errors { + errs = append(errs, err) + } + + if len(errs) > 0 { + for _, err := range errs { + t.Errorf("Error: %v", err) + } + t.Fatalf("%d errors occurred", len(errs)) + } + + t.Logf("%d concurrent embeds completed in %v (%.1f ops/sec)", + totalOps, duration, float64(totalOps)/duration.Seconds()) +} + +// TestHappyPath_CloseAfterAllComplete tests proper close after work is done. +func TestHappyPath_CloseAfterAllComplete(t *testing.T) { + modelPath := findModelPath(t) + if modelPath == "" { + t.Skip("Model not found") + } + logger := zap.NewNop() + + embedder, err := NewPooledHugotEmbedder(modelPath, "model.onnx", 2, logger) + if err != nil { + t.Fatalf("Failed to create embedder: %v", err) + } + + ctx := context.Background() + + // Do some work + for i := 0; i < 3; i++ { + contents := [][]ai.ContentPart{ + {ai.TextContent{Text: fmt.Sprintf("Test %d", i)}}, + } + _, err := embedder.Embed(ctx, contents) + if err != nil { + t.Fatalf("Embed %d failed: %v", i, err) + } + } + + // Close after all work is done - should succeed + err = embedder.Close() + if err != nil { + t.Errorf("Close() returned error: %v", err) + } + + t.Log("Close() after all embeds complete: success") +} + +// TestHappyPath_LargeBatch tests handling of larger batches. +func TestHappyPath_LargeBatch(t *testing.T) { + modelPath := findModelPath(t) + if modelPath == "" { + t.Skip("Model not found") + } + logger := zap.NewNop() + + embedder, err := NewPooledHugotEmbedder(modelPath, "model.onnx", 2, logger) + if err != nil { + t.Fatalf("Failed to create embedder: %v", err) + } + defer embedder.Close() + + ctx := context.Background() + + // Create a batch of 20 texts + batchSize := 20 + contents := make([][]ai.ContentPart, batchSize) + for i := 0; i < batchSize; i++ { + contents[i] = []ai.ContentPart{ + ai.TextContent{Text: fmt.Sprintf("Large batch test sentence number %d with some extra text", i)}, + } + } + + start := time.Now() + result, err := embedder.Embed(ctx, contents) + duration := time.Since(start) + + if err != nil { + t.Fatalf("Large batch embed failed: %v", err) + } + + if len(result) != batchSize { + t.Errorf("Expected %d embeddings, got %d", batchSize, len(result)) + } + + t.Logf("Batch of %d texts embedded in %v", batchSize, duration) +} + +// TestHappyPath_ContextCancellation tests that context cancellation is handled. +func TestHappyPath_ContextCancellation(t *testing.T) { + modelPath := findModelPath(t) + if modelPath == "" { + t.Skip("Model not found") + } + logger := zap.NewNop() + + embedder, err := NewPooledHugotEmbedder(modelPath, "model.onnx", 2, logger) + if err != nil { + t.Fatalf("Failed to create embedder: %v", err) + } + defer embedder.Close() + + // Create already-cancelled context + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + contents := [][]ai.ContentPart{ + {ai.TextContent{Text: "This should fail due to cancelled context"}}, + } + + _, err = embedder.Embed(ctx, contents) + if err == nil { + t.Error("Expected error with cancelled context, got nil") + } else { + t.Logf("Cancelled context correctly returned error: %v", err) + } +} + +// TestHappyPath_EmptyInput tests handling of empty input. +func TestHappyPath_EmptyInput(t *testing.T) { + modelPath := findModelPath(t) + if modelPath == "" { + t.Skip("Model not found") + } + logger := zap.NewNop() + + embedder, err := NewPooledHugotEmbedder(modelPath, "model.onnx", 2, logger) + if err != nil { + t.Fatalf("Failed to create embedder: %v", err) + } + defer embedder.Close() + + ctx := context.Background() + contents := [][]ai.ContentPart{} + + result, err := embedder.Embed(ctx, contents) + if err != nil { + t.Errorf("Empty input should not error: %v", err) + } + + if len(result) != 0 { + t.Errorf("Expected 0 results for empty input, got %d", len(result)) + } + + t.Log("Empty input handled correctly") +} diff --git a/pkg/termite/lib/embeddings/hugot.go b/pkg/termite/lib/embeddings/hugot.go index 22743b2..dda77de 100644 --- a/pkg/termite/lib/embeddings/hugot.go +++ b/pkg/termite/lib/embeddings/hugot.go @@ -20,6 +20,7 @@ import ( "fmt" "math" "runtime" + "sync" "sync/atomic" "github.com/antflydb/antfly-go/libaf/ai" @@ -79,6 +80,12 @@ type PooledHugotEmbedder struct { poolSize int caps embeddings.EmbedderCapabilities batchSize int + + // Synchronization for safe Close() behavior + closed atomic.Bool // Prevents new Embed() calls after Close() + wg sync.WaitGroup // Waits for in-flight Embed() calls to complete + closeOnce sync.Once // Ensures Close() runs exactly once + closeErr error // Stores error from Close() } // NewPooledHugotEmbedder creates a new pooled embedder using the Hugot ONNX runtime. @@ -290,10 +297,27 @@ func (p *PooledHugotEmbedder) Capabilities() embeddings.EmbedderCapabilities { return p.caps } +// ErrEmbedderClosed is returned when Embed is called on a closed embedder. +var ErrEmbedderClosed = errors.New("embedder is closed") + // Embed generates embeddings for the given content. // Thread-safe: uses semaphore to limit concurrent pipeline access. // Processes texts in batches to avoid memory explosion on CoreML. func (p *PooledHugotEmbedder) Embed(ctx context.Context, contents [][]ai.ContentPart) ([][]float32, error) { + // Check if embedder is closed before starting + if p.closed.Load() { + return nil, ErrEmbedderClosed + } + + // Track this in-flight operation so Close() waits for us + p.wg.Add(1) + defer p.wg.Done() + + // Double-check after registration (handles race with Close()) + if p.closed.Load() { + return nil, ErrEmbedderClosed + } + if len(contents) == 0 { return [][]float32{}, nil } @@ -375,23 +399,34 @@ func (p *PooledHugotEmbedder) Embed(ctx context.Context, contents [][]ai.Content // Close releases resources. // Properly closes each pipeline to remove it from the session, then destroys // the session if it was created by this embedder (not shared). +// Thread-safe: waits for in-flight Embed() calls to complete before destroying. +// Safe to call multiple times (only the first call takes effect). func (p *PooledHugotEmbedder) Close() error { - // Close each pipeline to remove it from the session - for _, pipeline := range p.pipelines { - if pipeline != nil { - name := pipeline.PipelineName - if err := khugot.ClosePipeline[*pipelines.FeatureExtractionPipeline](p.session, name); err != nil { - p.logger.Warn("Failed to close pipeline", zap.String("name", name), zap.Error(err)) + p.closeOnce.Do(func() { + // Set closed flag to prevent new Embed() calls + p.closed.Store(true) + + // Wait for all in-flight Embed() calls to complete + p.wg.Wait() + + // Close each pipeline to remove it from the session + for _, pipeline := range p.pipelines { + if pipeline != nil { + name := pipeline.PipelineName + if err := khugot.ClosePipeline[*pipelines.FeatureExtractionPipeline](p.session, name); err != nil { + p.logger.Warn("Failed to close pipeline", zap.String("name", name), zap.Error(err)) + } } } - } - p.pipelines = nil - - if p.session != nil && !p.sessionShared { - p.logger.Info("Destroying Hugot session (owned by this pooled embedder)") - return p.session.Destroy() - } else if p.sessionShared { - p.logger.Debug("Skipping session destruction (shared session)") - } - return nil + p.pipelines = nil + + // Now safe to destroy the session + if p.session != nil && !p.sessionShared { + p.logger.Info("Destroying Hugot session (owned by this pooled embedder)") + p.closeErr = p.session.Destroy() + } else if p.sessionShared { + p.logger.Debug("Skipping session destruction (shared session)") + } + }) + return p.closeErr } diff --git a/pkg/termite/lib/embeddings/pipeline_collision_test.go b/pkg/termite/lib/embeddings/pipeline_collision_test.go new file mode 100644 index 0000000..1af9be6 --- /dev/null +++ b/pkg/termite/lib/embeddings/pipeline_collision_test.go @@ -0,0 +1,344 @@ +// Copyright 2025 Antfly, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//go:build onnx && ORT + +package embeddings + +import ( + "context" + "fmt" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/antflydb/antfly-go/libaf/ai" + "go.uber.org/zap" +) + +// TestPipelineCollision validates the TLA+ counterexample showing two workers +// can be assigned the same pipeline due to round-robin counter wrapping. +// +// Bug mechanism (hugot.go:307-309): +// idx := int(p.nextPipeline.Add(1) % uint64(p.poolSize)) +// pipeline := p.pipelines[idx] +// +// The semaphore limits concurrent users to poolSize, but if a worker finishes +// quickly and another starts before the first slow worker completes, they can +// both get the same pipeline index. +// +// Counterexample trace (poolSize=2, 3 workers): +// 1. w1: acquires sem, nextPipeline=1, idx=1%2=1 +// 2. w2: acquires sem, nextPipeline=2, idx=2%2=0 +// 3. w2: completes quickly, releases sem +// 4. w3: acquires freed sem slot +// 5. w3: nextPipeline=3, idx=3%2=1 <-- COLLISION with w1! +func TestPipelineCollision(t *testing.T) { + modelPath := findModelPath(t) + if modelPath == "" { + t.Skip("Model not found, skipping pipeline collision test") + } + + const poolSize = 2 + const numWorkers = 10 // More workers = more collision opportunities + const iterations = 50 // Run multiple times to catch race + + logger := zap.NewNop() + + for iter := 0; iter < iterations; iter++ { + t.Run(fmt.Sprintf("iteration_%d", iter), func(t *testing.T) { + embedder, err := NewPooledHugotEmbedder(modelPath, "model.onnx", poolSize, logger) + if err != nil { + t.Fatalf("Failed to create embedder: %v", err) + } + defer embedder.Close() + + // Track collision detection + var collisionCount atomic.Int32 + var collisionDetails sync.Map // For debugging: stores collision info + // Note: Direct pipeline tracking is not possible without modifying production code + // The race detector (-race flag) will catch actual concurrent access + + var wg sync.WaitGroup + ctx := context.Background() + + // Start multiple workers concurrently + for w := 0; w < numWorkers; w++ { + wg.Add(1) + workerID := w + go func() { + defer wg.Done() + + // Simple input - just need to trigger pipeline selection + contents := [][]ai.ContentPart{ + {ai.TextContent{Text: "test sentence for collision detection"}}, + } + + // Get the pipeline index that will be selected + // We can predict this from nextPipeline, but that's racy + // Instead, we'll detect collision by checking pipelineUsers + + // Record that we're about to use a pipeline + // The actual pipeline selection happens inside Embed(), which we can't intercept + // So we approximate by checking if another goroutine is also embedding + + // Mark entry + startTime := time.Now() + + // Actually call Embed - this is where collision would manifest + _, err := embedder.Embed(ctx, contents) + if err != nil { + t.Logf("Worker %d: Embed error: %v", workerID, err) + } + + duration := time.Since(startTime) + t.Logf("Worker %d: completed in %v", workerID, duration) + }() + } + + wg.Wait() + + if collisionCount.Load() > 0 { + t.Errorf("Detected %d pipeline collisions!", collisionCount.Load()) + collisionDetails.Range(func(key, value any) bool { + t.Logf("Collision detail: %v", value) + return true + }) + } + }) + } +} + +// TestPipelineCollisionWithInstrumentation uses a more sophisticated approach +// to detect collisions by monitoring the nextPipeline counter and timing. +func TestPipelineCollisionWithInstrumentation(t *testing.T) { + modelPath := findModelPath(t) + if modelPath == "" { + t.Skip("Model not found, skipping instrumented collision test") + } + + const poolSize = 2 + logger := zap.NewNop() + + embedder, err := NewPooledHugotEmbedder(modelPath, "model.onnx", poolSize, logger) + if err != nil { + t.Fatalf("Failed to create embedder: %v", err) + } + defer embedder.Close() + + // Track pipeline usage by monitoring the atomic counter + // Since we can't directly observe which pipeline each goroutine gets, + // we instead track timing to detect overlapping usage + + type usageRecord struct { + workerID int + pipelineIdx int + startTime time.Time + endTime time.Time + } + + var records []usageRecord + var recordsMu sync.Mutex + var wg sync.WaitGroup + ctx := context.Background() + + // Strategy: Capture nextPipeline before and after Embed to determine + // which pipeline index was used. This is racy but gives us insight. + + numWorkers := 20 + for w := 0; w < numWorkers; w++ { + wg.Add(1) + workerID := w + go func() { + defer wg.Done() + + // Capture counter before call + // Note: This is inherently racy - another goroutine could increment between + // our read and the actual Add(1) inside Embed. But it's good enough for testing. + beforeCounter := embedder.nextPipeline.Load() + + startTime := time.Now() + + contents := [][]ai.ContentPart{ + {ai.TextContent{Text: "test sentence for instrumented collision detection"}}, + } + _, err := embedder.Embed(ctx, contents) + if err != nil { + t.Logf("Worker %d: error: %v", workerID, err) + } + + endTime := time.Now() + + // The pipeline index used was likely (beforeCounter+1) % poolSize + // This is approximate due to races + pipelineIdx := int((beforeCounter + 1) % uint64(poolSize)) + + recordsMu.Lock() + records = append(records, usageRecord{ + workerID: workerID, + pipelineIdx: pipelineIdx, + startTime: startTime, + endTime: endTime, + }) + recordsMu.Unlock() + }() + } + + wg.Wait() + + // Analyze records for overlapping usage of the same pipeline + collisions := 0 + for i := 0; i < len(records); i++ { + for j := i + 1; j < len(records); j++ { + r1, r2 := records[i], records[j] + + // Check if same pipeline and overlapping time + if r1.pipelineIdx == r2.pipelineIdx { + // Check for overlap: r1.start < r2.end AND r2.start < r1.end + if r1.startTime.Before(r2.endTime) && r2.startTime.Before(r1.endTime) { + collisions++ + t.Logf("POTENTIAL COLLISION: worker %d (pipeline %d, %v-%v) overlaps with worker %d (pipeline %d, %v-%v)", + r1.workerID, r1.pipelineIdx, r1.startTime.Format("15:04:05.000"), r1.endTime.Format("15:04:05.000"), + r2.workerID, r2.pipelineIdx, r2.startTime.Format("15:04:05.000"), r2.endTime.Format("15:04:05.000")) + } + } + } + } + + if collisions > 0 { + t.Errorf("Detected %d potential pipeline collisions (may include false positives due to timing approximation)", collisions) + } else { + t.Logf("No collisions detected in %d operations", numWorkers) + } +} + +// TestPipelineCollisionStress runs many concurrent operations to try to trigger +// the race condition and relies on the -race flag to detect data races. +// +// Run with: go test -v -race -tags="onnx,ORT" -run TestPipelineCollisionStress +func TestPipelineCollisionStress(t *testing.T) { + modelPath := findModelPath(t) + if modelPath == "" { + t.Skip("Model not found, skipping stress test") + } + + const poolSize = 2 + const numGoroutines = 50 + const opsPerGoroutine = 10 + + logger := zap.NewNop() + + embedder, err := NewPooledHugotEmbedder(modelPath, "model.onnx", poolSize, logger) + if err != nil { + t.Fatalf("Failed to create embedder: %v", err) + } + defer embedder.Close() + + var wg sync.WaitGroup + var errorCount atomic.Int32 + ctx := context.Background() + + start := time.Now() + + for g := 0; g < numGoroutines; g++ { + wg.Add(1) + go func(goroutineID int) { + defer wg.Done() + + for op := 0; op < opsPerGoroutine; op++ { + contents := [][]ai.ContentPart{ + {ai.TextContent{Text: fmt.Sprintf("stress test sentence %d-%d", goroutineID, op)}}, + } + + _, err := embedder.Embed(ctx, contents) + if err != nil { + errorCount.Add(1) + // Don't log every error to avoid spam + } + } + }(g) + } + + wg.Wait() + duration := time.Since(start) + + totalOps := numGoroutines * opsPerGoroutine + t.Logf("Completed %d operations in %v (%.1f ops/sec)", totalOps, duration, float64(totalOps)/duration.Seconds()) + t.Logf("Errors: %d", errorCount.Load()) + + // If we get here without the race detector firing, either: + // 1. The pipelines are actually thread-safe (contrary to assumption A6) + // 2. We didn't trigger the race condition + // 3. The race exists but wasn't detected + // + // The -race flag should catch concurrent access to the same pipeline + // if it's truly not thread-safe. +} + +// TestFirstEmbedUsesPipelineOne verifies the TLA+ finding that the first +// embed uses pipeline 1, not pipeline 0, due to Add(1) returning the new value. +func TestFirstEmbedUsesPipelineOne(t *testing.T) { + modelPath := findModelPath(t) + if modelPath == "" { + t.Skip("Model not found, skipping first-embed test") + } + + const poolSize = 4 // Use larger pool to make index more obvious + logger := zap.NewNop() + + embedder, err := NewPooledHugotEmbedder(modelPath, "model.onnx", poolSize, logger) + if err != nil { + t.Fatalf("Failed to create embedder: %v", err) + } + defer embedder.Close() + + // Check initial counter value + initialCounter := embedder.nextPipeline.Load() + t.Logf("Initial nextPipeline counter: %d", initialCounter) + + if initialCounter != 0 { + t.Errorf("Expected initial counter to be 0, got %d", initialCounter) + } + + // Perform first embed + ctx := context.Background() + contents := [][]ai.ContentPart{ + {ai.TextContent{Text: "first embed test"}}, + } + _, err = embedder.Embed(ctx, contents) + if err != nil { + t.Fatalf("First embed failed: %v", err) + } + + // Check counter after first embed + afterCounter := embedder.nextPipeline.Load() + t.Logf("After first embed, nextPipeline counter: %d", afterCounter) + + if afterCounter != 1 { + t.Errorf("Expected counter to be 1 after first embed, got %d", afterCounter) + } + + // The pipeline index used was: (0 + 1) % poolSize = 1 % 4 = 1 + // This means pipeline 0 is never used on the first call! + expectedPipelineUsed := int((initialCounter + 1) % uint64(poolSize)) + t.Logf("First embed used pipeline index: %d (pipeline 0 was skipped)", expectedPipelineUsed) + + if expectedPipelineUsed != 1 { + t.Errorf("Expected first embed to use pipeline 1, calculated %d", expectedPipelineUsed) + } + + // This confirms the TLA+ finding - it's a quirk but not a bug + t.Log("CONFIRMED: First embed skips pipeline 0 (uses pipeline 1)") +} From bdccfbf12d8ec63b1f3d0b02a80d1a72b54a23e6 Mon Sep 17 00:00:00 2001 From: Rowan Copley Date: Thu, 15 Jan 2026 16:00:12 -0800 Subject: [PATCH 3/4] move+rename e2e embeddings test and remove bug fix-specific commment --- .../happy_path_test.go => e2e/embeddings/batch_happy_path.go | 1 - 1 file changed, 1 deletion(-) rename pkg/termite/lib/embeddings/happy_path_test.go => e2e/embeddings/batch_happy_path.go (99%) diff --git a/pkg/termite/lib/embeddings/happy_path_test.go b/e2e/embeddings/batch_happy_path.go similarity index 99% rename from pkg/termite/lib/embeddings/happy_path_test.go rename to e2e/embeddings/batch_happy_path.go index 87c93ae..fc44139 100644 --- a/pkg/termite/lib/embeddings/happy_path_test.go +++ b/e2e/embeddings/batch_happy_path.go @@ -3,7 +3,6 @@ // Happy Path E2E Test for PooledHugotEmbedder // // This test verifies normal usage patterns work correctly. -// It should PASS both before and after the bugfixes. // // Run: // export ONNXRUNTIME_ROOT=$PWD/onnxruntime From 071cc61e8dd9ce2ea9042ad715935a0611552974 Mon Sep 17 00:00:00 2001 From: Rowan Copley Date: Fri, 16 Jan 2026 17:57:54 -0800 Subject: [PATCH 4/4] add fixes for other model race conditions --- pkg/termite/lib/chunking/hugot.go | 49 ++++- pkg/termite/lib/chunking/hugot_test.go | 159 ++++++++++++++++ pkg/termite/lib/classification/hugot.go | 61 ++++++- pkg/termite/lib/classification/hugot_test.go | 157 ++++++++++++++++ pkg/termite/lib/embeddings/close_race_test.go | 79 -------- pkg/termite/lib/generation/close_race_test.go | 166 +++++++++++++++++ pkg/termite/lib/generation/hugot.go | 69 ++++++- pkg/termite/lib/ner/hugot.go | 49 ++++- pkg/termite/lib/ner/hugot_test.go | 154 ++++++++++++++++ pkg/termite/lib/reranking/hugot.go | 49 ++++- pkg/termite/lib/reranking/hugot_test.go | 172 ++++++++++++++++++ 11 files changed, 1051 insertions(+), 113 deletions(-) create mode 100644 pkg/termite/lib/chunking/hugot_test.go create mode 100644 pkg/termite/lib/classification/hugot_test.go create mode 100644 pkg/termite/lib/generation/close_race_test.go create mode 100644 pkg/termite/lib/ner/hugot_test.go create mode 100644 pkg/termite/lib/reranking/hugot_test.go diff --git a/pkg/termite/lib/chunking/hugot.go b/pkg/termite/lib/chunking/hugot.go index a20d6b1..0830567 100644 --- a/pkg/termite/lib/chunking/hugot.go +++ b/pkg/termite/lib/chunking/hugot.go @@ -21,6 +21,7 @@ import ( "runtime" "sort" "strings" + "sync" "sync/atomic" "github.com/antflydb/antfly-go/libaf/chunking" @@ -67,8 +68,17 @@ type PooledHugotChunker struct { logger *zap.Logger sessionShared bool poolSize int + + // Synchronization for safe Close() behavior + closed atomic.Bool // Prevents new operations after Close() + wg sync.WaitGroup // Waits for in-flight operations to complete + closeOnce sync.Once // Ensures Close() runs exactly once + closeErr error // Stores error from Close() } +// ErrChunkerClosed is returned when Chunk is called on a closed chunker. +var ErrChunkerClosed = errors.New("chunker is closed") + // NewPooledHugotChunker creates a new pooled chunker using the Hugot ONNX runtime. // poolSize determines how many concurrent requests can be processed (0 = auto-detect from CPU count). // onnxFilename specifies which ONNX file to load (e.g., "model.onnx", "model_f16.onnx", "model_i8.onnx"). @@ -216,6 +226,20 @@ func newPooledHugotChunkerInternal(config HugotChunkerConfig, modelPath string, // Chunk splits text using neural token classification with per-request config overrides. // Thread-safe: uses semaphore to limit concurrent pipeline access. func (p *PooledHugotChunker) Chunk(ctx context.Context, text string, opts chunking.ChunkOptions) ([]chunking.Chunk, error) { + // Check if closed before starting + if p.closed.Load() { + return nil, ErrChunkerClosed + } + + // Track this in-flight operation so Close() waits for us + p.wg.Add(1) + defer p.wg.Done() + + // Double-check after registration (handles race with Close()) + if p.closed.Load() { + return nil, ErrChunkerClosed + } + if text == "" { p.logger.Debug("Chunk called with empty text") return nil, nil @@ -436,12 +460,23 @@ func (p *PooledHugotChunker) aggregateByTargetTokens(chunks []chunking.Chunk, co // Close releases the Hugot session and resources. // Only destroys the session if it was created by this chunker (not shared). +// Thread-safe: waits for in-flight operations to complete before destroying. +// Safe to call multiple times (only the first call takes effect). func (p *PooledHugotChunker) Close() error { - if p.session != nil && !p.sessionShared { - p.logger.Info("Destroying Hugot session (owned by this pooled chunker)") - return p.session.Destroy() - } else if p.sessionShared { - p.logger.Debug("Skipping session destruction (shared session)") - } - return nil + p.closeOnce.Do(func() { + // Set closed flag to prevent new operations + p.closed.Store(true) + + // Wait for all in-flight operations to complete + p.wg.Wait() + + // Now safe to destroy the session + if p.session != nil && !p.sessionShared { + p.logger.Info("Destroying Hugot session (owned by this pooled chunker)") + p.closeErr = p.session.Destroy() + } else if p.sessionShared { + p.logger.Debug("Skipping session destruction (shared session)") + } + }) + return p.closeErr } diff --git a/pkg/termite/lib/chunking/hugot_test.go b/pkg/termite/lib/chunking/hugot_test.go new file mode 100644 index 0000000..56d9c5f --- /dev/null +++ b/pkg/termite/lib/chunking/hugot_test.go @@ -0,0 +1,159 @@ +// Copyright 2025 Antfly, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//go:build onnx && ORT + +package chunking + +import ( + "context" + "os" + "path/filepath" + "sync" + "testing" + "time" + + "github.com/antflydb/antfly-go/libaf/chunking" + "go.uber.org/zap" +) + +// findModelPath searches for a chunker model in common locations. +func findModelPath(t *testing.T) string { + t.Helper() + + // Check common model locations + homeDir, _ := os.UserHomeDir() + paths := []string{ + filepath.Join(homeDir, ".termite", "models", "chunkers", "mirth", "chonky-mmbert-small-multilingual-1"), + "../../../../testdata/chunkers/chonky-mmbert-small-multilingual-1", + } + + for _, p := range paths { + if _, err := os.Stat(filepath.Join(p, "model.onnx")); err == nil { + t.Logf("Found model at %s", p) + return p + } + } + + return "" +} + +// TestCloseWhileChunking tests the race condition between Close() and Chunk(). +// With the fix in place, the -race detector should not find any races. +func TestCloseWhileChunking(t *testing.T) { + modelPath := findModelPath(t) + if modelPath == "" { + t.Skip("Chunker model not found, skipping close race test") + } + + const poolSize = 2 + logger := zap.NewNop() + config := DefaultHugotChunkerConfig() + + chunker, err := NewPooledHugotChunker(config, modelPath, "model.onnx", poolSize, logger) + if err != nil { + t.Fatalf("Failed to create chunker: %v", err) + } + + ctx := context.Background() + // Use a longer text to make chunking take more time + text := `Machine learning is a subset of artificial intelligence that focuses on building systems that learn from data. + +These systems improve their performance on specific tasks over time without being explicitly programmed. Deep learning, a more specialized form of machine learning, uses artificial neural networks with many layers. + +The field has seen tremendous growth in recent years, with applications ranging from image recognition to natural language processing. Companies across industries are adopting these technologies to automate processes and gain insights from their data.` + + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + // This runs concurrently with Close() - the -race detector catches unsync access + _, _ = chunker.Chunk(ctx, text, chunking.ChunkOptions{}) + }() + + time.Sleep(10 * time.Millisecond) // Let inference begin + _ = chunker.Close() + wg.Wait() +} + +// TestMultipleCloseIsSafe verifies that calling Close() multiple times +// doesn't cause issues (protected by sync.Once). +func TestMultipleCloseIsSafe(t *testing.T) { + modelPath := findModelPath(t) + if modelPath == "" { + t.Skip("Chunker model not found, skipping multiple close test") + } + + const poolSize = 2 + logger := zap.NewNop() + config := DefaultHugotChunkerConfig() + + chunker, err := NewPooledHugotChunker(config, modelPath, "model.onnx", poolSize, logger) + if err != nil { + t.Fatalf("Failed to create chunker: %v", err) + } + + // Do one successful chunk first + ctx := context.Background() + text := "This is a test document for chunking." + _, err = chunker.Chunk(ctx, text, chunking.ChunkOptions{}) + if err != nil { + t.Fatalf("Initial chunk failed: %v", err) + } + + // Now close multiple times concurrently + var wg sync.WaitGroup + for i := 0; i < 5; i++ { + wg.Add(1) + go func() { + defer wg.Done() + _ = chunker.Close() // Should not panic + }() + } + wg.Wait() +} + +// TestChunkAfterClose verifies behavior when Chunk is called after Close. +func TestChunkAfterClose(t *testing.T) { + modelPath := findModelPath(t) + if modelPath == "" { + t.Skip("Chunker model not found, skipping chunk-after-close test") + } + + const poolSize = 2 + logger := zap.NewNop() + config := DefaultHugotChunkerConfig() + + chunker, err := NewPooledHugotChunker(config, modelPath, "model.onnx", poolSize, logger) + if err != nil { + t.Fatalf("Failed to create chunker: %v", err) + } + + // Close first + err = chunker.Close() + if err != nil { + t.Fatalf("Close failed: %v", err) + } + + // Now try to chunk + ctx := context.Background() + text := "Test document after close." + _, err = chunker.Chunk(ctx, text, chunking.ChunkOptions{}) + if err == nil { + t.Error("Expected error when chunking after close, got nil") + } + if err != ErrChunkerClosed { + t.Errorf("Expected ErrChunkerClosed, got: %v", err) + } +} diff --git a/pkg/termite/lib/classification/hugot.go b/pkg/termite/lib/classification/hugot.go index 3c72fcd..6ed3f90 100644 --- a/pkg/termite/lib/classification/hugot.go +++ b/pkg/termite/lib/classification/hugot.go @@ -23,6 +23,7 @@ import ( "path/filepath" "runtime" "strings" + "sync" "sync/atomic" "github.com/antflydb/termite/pkg/termite/lib/hugot" @@ -55,8 +56,17 @@ type PooledHugotClassifier struct { sessionShared bool poolSize int config Config + + // Synchronization for safe Close() behavior + closed atomic.Bool // Prevents new operations after Close() + wg sync.WaitGroup // Waits for in-flight operations to complete + closeOnce sync.Once // Ensures Close() runs exactly once + closeErr error // Stores error from Close() } +// ErrClassifierClosed is returned when Classify is called on a closed classifier. +var ErrClassifierClosed = errors.New("classifier is closed") + // NewHugotClassifier creates a new zero-shot classifier using the Hugot ONNX runtime. func NewHugotClassifier(modelPath string, logger *zap.Logger) (*HugotClassifier, error) { return NewHugotClassifierWithSession(modelPath, nil, logger) @@ -491,6 +501,20 @@ func (p *PooledHugotClassifier) Classify(ctx context.Context, texts []string, la // ClassifyWithHypothesis classifies texts using a custom hypothesis template. func (p *PooledHugotClassifier) ClassifyWithHypothesis(ctx context.Context, texts []string, labels []string, hypothesisTemplate string) ([][]Classification, error) { + // Check if closed before starting + if p.closed.Load() { + return nil, ErrClassifierClosed + } + + // Track this in-flight operation so Close() waits for us + p.wg.Add(1) + defer p.wg.Done() + + // Double-check after registration (handles race with Close()) + if p.closed.Load() { + return nil, ErrClassifierClosed + } + if len(texts) == 0 { return [][]Classification{}, nil } @@ -538,6 +562,20 @@ func (p *PooledHugotClassifier) ClassifyWithHypothesis(ctx context.Context, text // MultiLabelClassify classifies texts allowing multiple labels per text. func (p *PooledHugotClassifier) MultiLabelClassify(ctx context.Context, texts []string, labels []string) ([][]Classification, error) { + // Check if closed before starting + if p.closed.Load() { + return nil, ErrClassifierClosed + } + + // Track this in-flight operation so Close() waits for us + p.wg.Add(1) + defer p.wg.Done() + + // Double-check after registration (handles race with Close()) + if p.closed.Load() { + return nil, ErrClassifierClosed + } + if len(texts) == 0 { return [][]Classification{}, nil } @@ -591,12 +629,25 @@ func (p *PooledHugotClassifier) MultiLabelClassify(ctx context.Context, texts [] } // Close releases resources. +// Thread-safe: waits for in-flight operations to complete before destroying. +// Safe to call multiple times (only the first call takes effect). func (p *PooledHugotClassifier) Close() error { - if p.session != nil && !p.sessionShared { - p.logger.Info("Destroying Hugot session (owned by this pooled ZSC)") - return p.session.Destroy() - } - return nil + p.closeOnce.Do(func() { + // Set closed flag to prevent new operations + p.closed.Store(true) + + // Wait for all in-flight operations to complete + p.wg.Wait() + + // Now safe to destroy the session + if p.session != nil && !p.sessionShared { + p.logger.Info("Destroying Hugot session (owned by this pooled ZSC)") + p.closeErr = p.session.Destroy() + } else if p.sessionShared { + p.logger.Debug("Skipping session destruction (shared session)") + } + }) + return p.closeErr } // Config returns the classifier configuration. diff --git a/pkg/termite/lib/classification/hugot_test.go b/pkg/termite/lib/classification/hugot_test.go new file mode 100644 index 0000000..026d03e --- /dev/null +++ b/pkg/termite/lib/classification/hugot_test.go @@ -0,0 +1,157 @@ +// Copyright 2025 Antfly, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//go:build onnx && ORT + +package classification + +import ( + "context" + "os" + "path/filepath" + "sync" + "testing" + "time" + + "go.uber.org/zap" +) + +// findModelPath searches for a classifier model in common locations. +func findModelPath(t *testing.T) string { + t.Helper() + + // Check common model locations + homeDir, _ := os.UserHomeDir() + paths := []string{ + filepath.Join(homeDir, ".termite", "models", "classifiers", "MoritzLaworski", "mDeBERTa-v3-base-mnli-xnli"), + filepath.Join(homeDir, ".termite", "models", "classifiers", "facebook", "bart-large-mnli"), + "../../../../testdata/classifiers/mDeBERTa-v3-base-mnli-xnli", + } + + for _, p := range paths { + if _, err := os.Stat(filepath.Join(p, "model.onnx")); err == nil { + t.Logf("Found model at %s", p) + return p + } + } + + return "" +} + +// TestCloseWhileClassifying tests the race condition between Close() and Classify(). +// With the fix in place, the -race detector should not find any races. +func TestCloseWhileClassifying(t *testing.T) { + modelPath := findModelPath(t) + if modelPath == "" { + t.Skip("Classifier model not found, skipping close race test") + } + + const poolSize = 2 + logger := zap.NewNop() + + classifier, err := NewPooledHugotClassifier(modelPath, poolSize, logger) + if err != nil { + t.Fatalf("Failed to create classifier: %v", err) + } + + ctx := context.Background() + texts := []string{ + "I love this product, it's amazing!", + "The weather is terrible today.", + } + labels := []string{"positive", "negative", "neutral"} + + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + // This runs concurrently with Close() - the -race detector catches unsync access + _, _ = classifier.Classify(ctx, texts, labels) + }() + + time.Sleep(10 * time.Millisecond) // Let inference begin + _ = classifier.Close() + wg.Wait() +} + +// TestMultipleCloseIsSafe verifies that calling Close() multiple times +// doesn't cause issues (protected by sync.Once). +func TestMultipleCloseIsSafe(t *testing.T) { + modelPath := findModelPath(t) + if modelPath == "" { + t.Skip("Classifier model not found, skipping multiple close test") + } + + const poolSize = 2 + logger := zap.NewNop() + + classifier, err := NewPooledHugotClassifier(modelPath, poolSize, logger) + if err != nil { + t.Fatalf("Failed to create classifier: %v", err) + } + + // Do one successful classify first + ctx := context.Background() + texts := []string{"Test sentence"} + labels := []string{"positive", "negative"} + _, err = classifier.Classify(ctx, texts, labels) + if err != nil { + t.Fatalf("Initial classify failed: %v", err) + } + + // Now close multiple times concurrently + var wg sync.WaitGroup + for i := 0; i < 5; i++ { + wg.Add(1) + go func() { + defer wg.Done() + _ = classifier.Close() // Should not panic + }() + } + wg.Wait() +} + +// TestClassifyAfterClose verifies behavior when Classify is called after Close. +func TestClassifyAfterClose(t *testing.T) { + modelPath := findModelPath(t) + if modelPath == "" { + t.Skip("Classifier model not found, skipping classify-after-close test") + } + + const poolSize = 2 + logger := zap.NewNop() + + classifier, err := NewPooledHugotClassifier(modelPath, poolSize, logger) + if err != nil { + t.Fatalf("Failed to create classifier: %v", err) + } + + // Close first + err = classifier.Close() + if err != nil { + t.Fatalf("Close failed: %v", err) + } + + // Now try to classify + ctx := context.Background() + texts := []string{"Test sentence"} + labels := []string{"positive", "negative"} + _, err = classifier.Classify(ctx, texts, labels) + if err == nil { + t.Error("Expected error when classifying after close, got nil") + } + if err != ErrClassifierClosed { + t.Errorf("Expected ErrClassifierClosed, got: %v", err) + } +} diff --git a/pkg/termite/lib/embeddings/close_race_test.go b/pkg/termite/lib/embeddings/close_race_test.go index 73039b4..6471fbc 100644 --- a/pkg/termite/lib/embeddings/close_race_test.go +++ b/pkg/termite/lib/embeddings/close_race_test.go @@ -130,85 +130,6 @@ func TestCloseWhileEmbedding(t *testing.T) { } } -// TestCloseWhileEmbeddingStress runs many iterations to increase chance -// of triggering the race condition. -func TestCloseWhileEmbeddingStress(t *testing.T) { - modelPath := findModelPath(t) - if modelPath == "" { - t.Skip("Model not found, skipping close race stress test") - } - - const iterations = 20 - const poolSize = 2 - logger := zap.NewNop() - - var panicCount atomic.Int32 - var errorCount atomic.Int32 - var successCount atomic.Int32 - - for iter := 0; iter < iterations; iter++ { - t.Run(fmt.Sprintf("iter_%d", iter), func(t *testing.T) { - embedder, err := NewPooledHugotEmbedder(modelPath, "model.onnx", poolSize, logger) - if err != nil { - t.Fatalf("Failed to create embedder: %v", err) - } - - var embedDone sync.WaitGroup - embedDone.Add(1) - - var panicked atomic.Bool - var embedError atomic.Value - - ctx := context.Background() - - go func() { - defer embedDone.Done() - defer func() { - if r := recover(); r != nil { - panicked.Store(true) - panicCount.Add(1) - } - }() - - contents := make([][]ai.ContentPart, 20) - for i := 0; i < len(contents); i++ { - contents[i] = []ai.ContentPart{ - ai.TextContent{Text: fmt.Sprintf("stress test sentence %d for iteration %d", i, iter)}, - } - } - - _, err := embedder.Embed(ctx, contents) - if err != nil { - embedError.Store(err) - errorCount.Add(1) - } else { - successCount.Add(1) - } - }() - - // Variable delay to hit different points in the inference - delay := time.Duration(iter%10) * time.Millisecond - time.Sleep(delay) - - // Close while embed is (probably) running - _ = embedder.Close() - - embedDone.Wait() - - if panicked.Load() { - t.Errorf("Iteration %d: PANIC detected", iter) - } - }) - } - - t.Logf("Summary: %d panics, %d errors, %d successes out of %d iterations", - panicCount.Load(), errorCount.Load(), successCount.Load(), iterations) - - if panicCount.Load() > 0 { - t.Errorf("BUG CONFIRMED: %d panics detected - Close() race is dangerous", panicCount.Load()) - } -} - // TestMultipleCloseIsSafe verifies that calling Close() multiple times // doesn't cause issues (tests assumption A5: session destroy idempotence). func TestMultipleCloseIsSafe(t *testing.T) { diff --git a/pkg/termite/lib/generation/close_race_test.go b/pkg/termite/lib/generation/close_race_test.go new file mode 100644 index 0000000..6318084 --- /dev/null +++ b/pkg/termite/lib/generation/close_race_test.go @@ -0,0 +1,166 @@ +// Copyright 2025 Antfly, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//go:build onnx && ORT + +package generation + +import ( + "context" + "os" + "path/filepath" + "sync" + "testing" + "time" + + "go.uber.org/zap" +) + +// findModelPath searches for a generator model in common locations. +func findModelPath(t *testing.T) string { + t.Helper() + + // Check common model locations + homeDir, _ := os.UserHomeDir() + paths := []string{ + filepath.Join(homeDir, ".termite", "models", "generators", "google", "gemma-2-2b-it"), + "../../../../testdata/generators/tiny-random-gemma-3", + } + + for _, p := range paths { + // Check for genai_config.json which is required for ONNX RT GenAI + if _, err := os.Stat(filepath.Join(p, "genai_config.json")); err == nil { + t.Logf("Found model at %s", p) + return p + } + } + + return "" +} + +// TestCloseWhileGenerating tests the race condition between Close() and Generate(). +// With the fix in place, the -race detector should not find any races. +func TestCloseWhileGenerating(t *testing.T) { + modelPath := findModelPath(t) + if modelPath == "" { + t.Skip("Generator model not found, skipping close race test") + } + + const poolSize = 2 + logger := zap.NewNop() + + generator, err := NewPooledHugotGenerator(modelPath, poolSize, logger) + if err != nil { + t.Fatalf("Failed to create generator: %v", err) + } + + ctx := context.Background() + messages := []Message{ + {Role: "user", Content: "Write a short poem about clouds."}, + } + opts := GenerateOptions{ + MaxTokens: 50, + } + + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + // This runs concurrently with Close() - the -race detector catches unsync access + _, _ = generator.Generate(ctx, messages, opts) + }() + + time.Sleep(10 * time.Millisecond) // Let inference begin + _ = generator.Close() + wg.Wait() +} + +// TestMultipleCloseIsSafe verifies that calling Close() multiple times +// doesn't cause issues (protected by sync.Once). +func TestMultipleCloseIsSafe(t *testing.T) { + modelPath := findModelPath(t) + if modelPath == "" { + t.Skip("Generator model not found, skipping multiple close test") + } + + const poolSize = 2 + logger := zap.NewNop() + + generator, err := NewPooledHugotGenerator(modelPath, poolSize, logger) + if err != nil { + t.Fatalf("Failed to create generator: %v", err) + } + + // Do one successful generate first + ctx := context.Background() + messages := []Message{ + {Role: "user", Content: "Hello"}, + } + opts := GenerateOptions{ + MaxTokens: 5, + } + _, err = generator.Generate(ctx, messages, opts) + if err != nil { + t.Fatalf("Initial generate failed: %v", err) + } + + // Now close multiple times concurrently + var wg sync.WaitGroup + for i := 0; i < 5; i++ { + wg.Add(1) + go func() { + defer wg.Done() + _ = generator.Close() // Should not panic + }() + } + wg.Wait() +} + +// TestGenerateAfterClose verifies behavior when Generate is called after Close. +func TestGenerateAfterClose(t *testing.T) { + modelPath := findModelPath(t) + if modelPath == "" { + t.Skip("Generator model not found, skipping generate-after-close test") + } + + const poolSize = 2 + logger := zap.NewNop() + + generator, err := NewPooledHugotGenerator(modelPath, poolSize, logger) + if err != nil { + t.Fatalf("Failed to create generator: %v", err) + } + + // Close first + err = generator.Close() + if err != nil { + t.Fatalf("Close failed: %v", err) + } + + // Now try to generate + ctx := context.Background() + messages := []Message{ + {Role: "user", Content: "Hello"}, + } + opts := GenerateOptions{ + MaxTokens: 5, + } + _, err = generator.Generate(ctx, messages, opts) + if err == nil { + t.Error("Expected error when generating after close, got nil") + } + if err != ErrGeneratorClosed { + t.Errorf("Expected ErrGeneratorClosed, got: %v", err) + } +} diff --git a/pkg/termite/lib/generation/hugot.go b/pkg/termite/lib/generation/hugot.go index 1595a30..82cd08b 100644 --- a/pkg/termite/lib/generation/hugot.go +++ b/pkg/termite/lib/generation/hugot.go @@ -22,6 +22,7 @@ import ( "os" "path/filepath" "runtime" + "sync" "sync/atomic" "github.com/antflydb/termite/pkg/termite/lib/hugot" @@ -511,8 +512,17 @@ type PooledHugotGenerator struct { imageToken string // image placeholder token from model's special_tokens_map.json toolParser ToolParser // tool call parser (from genai_config.json) toolCallFormat string // the tool call format name (e.g., "functiongemma") + + // Synchronization for safe Close() behavior + closed atomic.Bool // Prevents new operations after Close() + wg sync.WaitGroup // Waits for in-flight operations to complete + closeOnce sync.Once // Ensures Close() runs exactly once + closeErr error // Stores error from Close() } +// ErrGeneratorClosed is returned when Generate is called on a closed generator. +var ErrGeneratorClosed = errors.New("generator is closed") + // NewPooledHugotGenerator creates a new pooled generator using the Hugot runtime. // poolSize determines how many concurrent requests can be processed (0 = auto-detect from CPU count). // Note: For generative models, hugot uses genai_config.json to determine model files. @@ -647,6 +657,20 @@ func (p *PooledHugotGenerator) convertMessages(messages []Message) []backends.Me // Thread-safe: uses semaphore to limit concurrent pipeline access. // Uses the streaming pipeline internally and collects all tokens into the response. func (p *PooledHugotGenerator) Generate(ctx context.Context, messages []Message, opts GenerateOptions) (*GenerateResult, error) { + // Check if closed before starting + if p.closed.Load() { + return nil, ErrGeneratorClosed + } + + // Track this in-flight operation so Close() waits for us + p.wg.Add(1) + defer p.wg.Done() + + // Double-check after registration (handles race with Close()) + if p.closed.Load() { + return nil, ErrGeneratorClosed + } + if len(messages) == 0 { return nil, errors.New("messages are required") } @@ -715,12 +739,28 @@ func (p *PooledHugotGenerator) Generate(ctx context.Context, messages []Message, // GenerateStream produces tokens one at a time via channels. // Thread-safe: uses semaphore to limit concurrent pipeline access. func (p *PooledHugotGenerator) GenerateStream(ctx context.Context, messages []Message, opts GenerateOptions) (<-chan TokenDelta, <-chan error, error) { + // Check if closed before starting + if p.closed.Load() { + return nil, nil, ErrGeneratorClosed + } + + // Note: wg.Done() is called in the goroutine when streaming completes + p.wg.Add(1) + + // Double-check after registration (handles race with Close()) + if p.closed.Load() { + p.wg.Done() + return nil, nil, ErrGeneratorClosed + } + if len(messages) == 0 { + p.wg.Done() return nil, nil, errors.New("messages are required") } // Acquire semaphore slot (blocks if all pipelines busy) if err := p.sem.Acquire(ctx, 1); err != nil { + p.wg.Done() return nil, nil, fmt.Errorf("acquiring pipeline slot: %w", err) } @@ -737,6 +777,7 @@ func (p *PooledHugotGenerator) GenerateStream(ctx context.Context, messages []Me output, err := pipeline.RunMessages(ctx, [][]backends.Message{hugotMessages}) if err != nil { p.sem.Release(1) + p.wg.Done() p.logger.Error("Streaming pipeline generation failed", zap.Int("pipelineIndex", idx), zap.Error(err)) @@ -748,7 +789,8 @@ func (p *PooledHugotGenerator) GenerateStream(ctx context.Context, messages []Me errChan := make(chan error, 1) go func() { - defer p.sem.Release(1) // Release semaphore when done streaming + defer p.wg.Done() // Signal completion to WaitGroup + defer p.sem.Release(1) // Release semaphore when done streaming defer close(tokenChan) defer close(errChan) @@ -777,12 +819,23 @@ func (p *PooledHugotGenerator) GenerateStream(ctx context.Context, messages []Me // Close releases resources. // Only destroys the session if it was created by this generator (not shared). +// Thread-safe: waits for in-flight operations to complete before destroying. +// Safe to call multiple times (only the first call takes effect). func (p *PooledHugotGenerator) Close() error { - if p.session != nil && !p.sessionShared { - p.logger.Info("Destroying Hugot session (owned by this pooled generator)") - return p.session.Destroy() - } else if p.sessionShared { - p.logger.Debug("Skipping session destruction (shared session)") - } - return nil + p.closeOnce.Do(func() { + // Set closed flag to prevent new operations + p.closed.Store(true) + + // Wait for all in-flight operations to complete + p.wg.Wait() + + // Now safe to destroy the session + if p.session != nil && !p.sessionShared { + p.logger.Info("Destroying Hugot session (owned by this pooled generator)") + p.closeErr = p.session.Destroy() + } else if p.sessionShared { + p.logger.Debug("Skipping session destruction (shared session)") + } + }) + return p.closeErr } diff --git a/pkg/termite/lib/ner/hugot.go b/pkg/termite/lib/ner/hugot.go index 6449fab..bf9926c 100644 --- a/pkg/termite/lib/ner/hugot.go +++ b/pkg/termite/lib/ner/hugot.go @@ -19,6 +19,7 @@ import ( "errors" "fmt" "runtime" + "sync" "sync/atomic" "github.com/antflydb/termite/pkg/termite/lib/hugot" @@ -316,8 +317,17 @@ type PooledHugotNER struct { logger *zap.Logger sessionShared bool poolSize int + + // Synchronization for safe Close() behavior + closed atomic.Bool // Prevents new operations after Close() + wg sync.WaitGroup // Waits for in-flight operations to complete + closeOnce sync.Once // Ensures Close() runs exactly once + closeErr error // Stores error from Close() } +// ErrNERClosed is returned when Recognize is called on a closed NER model. +var ErrNERClosed = errors.New("NER model is closed") + // NewPooledHugotNER creates a new pooled NER model using the Hugot ONNX runtime. // poolSize determines how many concurrent requests can be processed (0 = auto-detect from CPU count). // onnxFilename specifies which ONNX file to load (e.g., "model.onnx", "model_i8.onnx"). @@ -516,6 +526,20 @@ func NewPooledHugotNERWithSessionManager(modelPath string, onnxFilename string, // Recognize extracts named entities from the given texts. // Thread-safe: uses semaphore to limit concurrent pipeline access. func (p *PooledHugotNER) Recognize(ctx context.Context, texts []string) ([][]Entity, error) { + // Check if closed before starting + if p.closed.Load() { + return nil, ErrNERClosed + } + + // Track this in-flight operation so Close() waits for us + p.wg.Add(1) + defer p.wg.Done() + + // Double-check after registration (handles race with Close()) + if p.closed.Load() { + return nil, ErrNERClosed + } + if len(texts) == 0 { return nil, nil } @@ -603,14 +627,25 @@ func (p *PooledHugotNER) parseEntities(text string, pipelineEntities []pipelines // Close releases resources. // Only destroys the session if it was created by this NER model (not shared). +// Thread-safe: waits for in-flight operations to complete before destroying. +// Safe to call multiple times (only the first call takes effect). func (p *PooledHugotNER) Close() error { - if p.session != nil && !p.sessionShared { - p.logger.Info("Destroying Hugot session (owned by this pooled NER model)") - return p.session.Destroy() - } else if p.sessionShared { - p.logger.Debug("Skipping session destruction (shared session)") - } - return nil + p.closeOnce.Do(func() { + // Set closed flag to prevent new operations + p.closed.Store(true) + + // Wait for all in-flight operations to complete + p.wg.Wait() + + // Now safe to destroy the session + if p.session != nil && !p.sessionShared { + p.logger.Info("Destroying Hugot session (owned by this pooled NER model)") + p.closeErr = p.session.Destroy() + } else if p.sessionShared { + p.logger.Debug("Skipping session destruction (shared session)") + } + }) + return p.closeErr } // countEntities returns the total number of entities across all texts. diff --git a/pkg/termite/lib/ner/hugot_test.go b/pkg/termite/lib/ner/hugot_test.go new file mode 100644 index 0000000..a25da90 --- /dev/null +++ b/pkg/termite/lib/ner/hugot_test.go @@ -0,0 +1,154 @@ +// Copyright 2025 Antfly, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//go:build onnx && ORT + +package ner + +import ( + "context" + "os" + "path/filepath" + "sync" + "testing" + "time" + + "go.uber.org/zap" +) + +// findModelPath searches for a NER model in common locations. +func findModelPath(t *testing.T) string { + t.Helper() + + // Check common model locations + homeDir, _ := os.UserHomeDir() + paths := []string{ + filepath.Join(homeDir, ".termite", "models", "ner", "dslim", "bert-base-NER"), + filepath.Join(homeDir, ".termite", "models", "ner", "dbmdz", "bert-large-cased-finetuned-conll03-english"), + "../../../../testdata/ner/bert-base-NER", + } + + for _, p := range paths { + if _, err := os.Stat(filepath.Join(p, "model.onnx")); err == nil { + t.Logf("Found model at %s", p) + return p + } + } + + return "" +} + +// TestCloseWhileRecognizing tests the race condition between Close() and Recognize(). +// With the fix in place, the -race detector should not find any races. +func TestCloseWhileRecognizing(t *testing.T) { + modelPath := findModelPath(t) + if modelPath == "" { + t.Skip("NER model not found, skipping close race test") + } + + const poolSize = 2 + logger := zap.NewNop() + + recognizer, err := NewPooledHugotNER(modelPath, "model.onnx", poolSize, logger) + if err != nil { + t.Fatalf("Failed to create NER recognizer: %v", err) + } + + ctx := context.Background() + texts := []string{ + "John Smith works at Google in Mountain View, California.", + "Apple Inc. was founded by Steve Jobs in Cupertino.", + } + + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + // This runs concurrently with Close() - the -race detector catches unsync access + _, _ = recognizer.Recognize(ctx, texts) + }() + + time.Sleep(10 * time.Millisecond) // Let inference begin + _ = recognizer.Close() + wg.Wait() +} + +// TestMultipleCloseIsSafe verifies that calling Close() multiple times +// doesn't cause issues (protected by sync.Once). +func TestMultipleCloseIsSafe(t *testing.T) { + modelPath := findModelPath(t) + if modelPath == "" { + t.Skip("NER model not found, skipping multiple close test") + } + + const poolSize = 2 + logger := zap.NewNop() + + recognizer, err := NewPooledHugotNER(modelPath, "model.onnx", poolSize, logger) + if err != nil { + t.Fatalf("Failed to create NER recognizer: %v", err) + } + + // Do one successful recognize first + ctx := context.Background() + texts := []string{"John works at Google."} + _, err = recognizer.Recognize(ctx, texts) + if err != nil { + t.Fatalf("Initial recognize failed: %v", err) + } + + // Now close multiple times concurrently + var wg sync.WaitGroup + for i := 0; i < 5; i++ { + wg.Add(1) + go func() { + defer wg.Done() + _ = recognizer.Close() // Should not panic + }() + } + wg.Wait() +} + +// TestRecognizeAfterClose verifies behavior when Recognize is called after Close. +func TestRecognizeAfterClose(t *testing.T) { + modelPath := findModelPath(t) + if modelPath == "" { + t.Skip("NER model not found, skipping recognize-after-close test") + } + + const poolSize = 2 + logger := zap.NewNop() + + recognizer, err := NewPooledHugotNER(modelPath, "model.onnx", poolSize, logger) + if err != nil { + t.Fatalf("Failed to create NER recognizer: %v", err) + } + + // Close first + err = recognizer.Close() + if err != nil { + t.Fatalf("Close failed: %v", err) + } + + // Now try to recognize + ctx := context.Background() + texts := []string{"John works at Google."} + _, err = recognizer.Recognize(ctx, texts) + if err == nil { + t.Error("Expected error when recognizing after close, got nil") + } + if err != ErrNERClosed { + t.Errorf("Expected ErrNERClosed, got: %v", err) + } +} diff --git a/pkg/termite/lib/reranking/hugot.go b/pkg/termite/lib/reranking/hugot.go index e4ac572..3fd148b 100644 --- a/pkg/termite/lib/reranking/hugot.go +++ b/pkg/termite/lib/reranking/hugot.go @@ -19,6 +19,7 @@ import ( "errors" "fmt" "runtime" + "sync" "sync/atomic" "github.com/antflydb/antfly-go/libaf/reranking" @@ -70,8 +71,17 @@ type PooledHugotReranker struct { logger *zap.Logger sessionShared bool poolSize int + + // Synchronization for safe Close() behavior + closed atomic.Bool // Prevents new operations after Close() + wg sync.WaitGroup // Waits for in-flight operations to complete + closeOnce sync.Once // Ensures Close() runs exactly once + closeErr error // Stores error from Close() } +// ErrRerankerClosed is returned when Rerank is called on a closed reranker. +var ErrRerankerClosed = errors.New("reranker is closed") + // NewPooledHugotReranker creates a new pooled reranker using the Hugot ONNX runtime. // poolSize determines how many concurrent requests can be processed (0 = auto-detect from CPU count). // onnxFilename specifies which ONNX file to load (e.g., "model.onnx", "model_f16.onnx", "model_i8.onnx"). @@ -203,6 +213,20 @@ func newPooledHugotRerankerInternal(modelPath string, onnxFilename string, poolS // Rerank scores pre-rendered prompts based on relevance to the query. // Thread-safe: uses semaphore to limit concurrent pipeline access. func (p *PooledHugotReranker) Rerank(ctx context.Context, query string, prompts []string) ([]float32, error) { + // Check if closed before starting + if p.closed.Load() { + return nil, ErrRerankerClosed + } + + // Track this in-flight operation so Close() waits for us + p.wg.Add(1) + defer p.wg.Done() + + // Double-check after registration (handles race with Close()) + if p.closed.Load() { + return nil, ErrRerankerClosed + } + if len(prompts) == 0 { return []float32{}, nil } @@ -251,12 +275,23 @@ func (p *PooledHugotReranker) Rerank(ctx context.Context, query string, prompts // Close releases resources. // Only destroys the session if it was created by this reranker (not shared). +// Thread-safe: waits for in-flight operations to complete before destroying. +// Safe to call multiple times (only the first call takes effect). func (p *PooledHugotReranker) Close() error { - if p.session != nil && !p.sessionShared { - p.logger.Info("Destroying Hugot session (owned by this pooled reranker)") - return p.session.Destroy() - } else if p.sessionShared { - p.logger.Debug("Skipping session destruction (shared session)") - } - return nil + p.closeOnce.Do(func() { + // Set closed flag to prevent new operations + p.closed.Store(true) + + // Wait for all in-flight operations to complete + p.wg.Wait() + + // Now safe to destroy the session + if p.session != nil && !p.sessionShared { + p.logger.Info("Destroying Hugot session (owned by this pooled reranker)") + p.closeErr = p.session.Destroy() + } else if p.sessionShared { + p.logger.Debug("Skipping session destruction (shared session)") + } + }) + return p.closeErr } diff --git a/pkg/termite/lib/reranking/hugot_test.go b/pkg/termite/lib/reranking/hugot_test.go new file mode 100644 index 0000000..328b6ec --- /dev/null +++ b/pkg/termite/lib/reranking/hugot_test.go @@ -0,0 +1,172 @@ +// Copyright 2025 Antfly, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//go:build onnx && ORT + +package reranking + +import ( + "context" + "os" + "path/filepath" + "sync" + "testing" + "time" + + "go.uber.org/zap" +) + +// findModelPath searches for a reranker model in common locations. +func findModelPath(t *testing.T) string { + t.Helper() + + // Check common model locations + homeDir, _ := os.UserHomeDir() + paths := []string{ + filepath.Join(homeDir, ".termite", "models", "rerankers", "mixedbread-ai", "mxbai-rerank-base-v1"), + filepath.Join(homeDir, ".termite", "models", "rerankers", "BAAI", "bge-reranker-base"), + "../../../../testdata/rerankers/mxbai-rerank-base-v1", + } + + for _, p := range paths { + // Check for model_i8.onnx first (smaller), then model.onnx + for _, modelFile := range []string{"model_i8.onnx", "model.onnx"} { + if _, err := os.Stat(filepath.Join(p, modelFile)); err == nil { + t.Logf("Found model at %s with %s", p, modelFile) + return p + } + } + } + + return "" +} + +// getOnnxFilename returns the ONNX filename to use based on what's available. +func getOnnxFilename(modelPath string) string { + if _, err := os.Stat(filepath.Join(modelPath, "model_i8.onnx")); err == nil { + return "model_i8.onnx" + } + return "model.onnx" +} + +// TestCloseWhileReranking tests the race condition between Close() and Rerank(). +// With the fix in place, the -race detector should not find any races. +func TestCloseWhileReranking(t *testing.T) { + modelPath := findModelPath(t) + if modelPath == "" { + t.Skip("Reranker model not found, skipping close race test") + } + + const poolSize = 2 + logger := zap.NewNop() + onnxFile := getOnnxFilename(modelPath) + + reranker, err := NewPooledHugotReranker(modelPath, onnxFile, poolSize, logger) + if err != nil { + t.Fatalf("Failed to create reranker: %v", err) + } + + ctx := context.Background() + query := "What is machine learning?" + prompts := []string{ + "Machine learning is a subset of artificial intelligence.", + "The weather today is sunny and warm.", + "Deep learning uses neural networks with many layers.", + } + + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + // This runs concurrently with Close() - the -race detector catches unsync access + _, _ = reranker.Rerank(ctx, query, prompts) + }() + + time.Sleep(10 * time.Millisecond) // Let inference begin + _ = reranker.Close() + wg.Wait() +} + +// TestMultipleCloseIsSafe verifies that calling Close() multiple times +// doesn't cause issues (protected by sync.Once). +func TestMultipleCloseIsSafe(t *testing.T) { + modelPath := findModelPath(t) + if modelPath == "" { + t.Skip("Reranker model not found, skipping multiple close test") + } + + const poolSize = 2 + logger := zap.NewNop() + onnxFile := getOnnxFilename(modelPath) + + reranker, err := NewPooledHugotReranker(modelPath, onnxFile, poolSize, logger) + if err != nil { + t.Fatalf("Failed to create reranker: %v", err) + } + + // Do one successful rerank first + ctx := context.Background() + query := "test query" + prompts := []string{"test prompt"} + _, err = reranker.Rerank(ctx, query, prompts) + if err != nil { + t.Fatalf("Initial rerank failed: %v", err) + } + + // Now close multiple times concurrently + var wg sync.WaitGroup + for i := 0; i < 5; i++ { + wg.Add(1) + go func() { + defer wg.Done() + _ = reranker.Close() // Should not panic + }() + } + wg.Wait() +} + +// TestRerankAfterClose verifies behavior when Rerank is called after Close. +func TestRerankAfterClose(t *testing.T) { + modelPath := findModelPath(t) + if modelPath == "" { + t.Skip("Reranker model not found, skipping rerank-after-close test") + } + + const poolSize = 2 + logger := zap.NewNop() + onnxFile := getOnnxFilename(modelPath) + + reranker, err := NewPooledHugotReranker(modelPath, onnxFile, poolSize, logger) + if err != nil { + t.Fatalf("Failed to create reranker: %v", err) + } + + // Close first + err = reranker.Close() + if err != nil { + t.Fatalf("Close failed: %v", err) + } + + // Now try to rerank + ctx := context.Background() + query := "test query" + prompts := []string{"test prompt"} + _, err = reranker.Rerank(ctx, query, prompts) + if err == nil { + t.Error("Expected error when reranking after close, got nil") + } + if err != ErrRerankerClosed { + t.Errorf("Expected ErrRerankerClosed, got: %v", err) + } +}