fix(parallel): cache factory modules and probe count once to avoid double calls (#26)

Signed-off-by: assagman <ahmetsercansagman@gmail.com>

Author: assagman

internal/module/parallel.go

@@ -68,6 +68,13 @@ type Parallel struct {
 	// Runtime state for logging/context (captured once)
 	moduleInfoOnce sync.Once
 	moduleInfo     parallelModuleInfo
+
+	mu sync.Mutex
+
+	// Factory cache to avoid double factory calls during probing and execution
+	factoryCache   []core.Module
+	factoryCount   int
+	factoryCounted bool
 }
 
 // getParallelModuleInfo extracts module type and LM model information for logging
@@ -214,12 +221,21 @@ func NewParallel(module core.Module) *Parallel {
 		batchKey:       "_batch",
 		repeat:         1,
 		verbose:        false,
+		factoryCache:   nil,
+		factoryCount:   0,
+		factoryCounted: false,
 	}
 }
 
 // NewParallelWithFactory creates a Parallel module with a factory function.
 // The factory is called for each task with the task index.
-// This is the recommended approach for stateful modules.
+//
+// When inputs don't specify a batch (_batch or slices), Parallel probes the
+// factory to determine module count by calling factory(i) until it returns nil.
+// The factory MUST return nil for out-of-bounds indices.
+//
+// In factory mode, WithRepeat(n) is ignored as the factory count determines
+// the number of parallel tasks.
 func NewParallelWithFactory(factory func(i int) core.Module) *Parallel {
 	return &Parallel{
 		factory:        factory,
@@ -231,11 +247,17 @@ func NewParallelWithFactory(factory func(i int) core.Module) *Parallel {
 		batchKey:       "_batch",
 		repeat:         1,
 		verbose:        false,
+		factoryCache:   nil,
+		factoryCount:   0,
+		factoryCounted: false,
 	}
 }
 
 // NewParallelWithInstances creates a Parallel module with pre-created instances.
-// Each task will use instances[i % len(instances)].
+// Each task uses instances[i % len(instances)].
+//
+// By default, repeat is set to len(instances), so a single input runs through
+// all instances. Override with WithRepeat() if using batch inputs.
 func NewParallelWithInstances(instances []core.Module) *Parallel {
 	if len(instances) == 0 {
 		panic("NewParallelWithInstances: instances slice cannot be empty")
@@ -248,8 +270,11 @@ func NewParallelWithInstances(instances []core.Module) *Parallel {
 		returnAll:      true,
 		onlySuccessful: true,
 		batchKey:       "_batch",
-		repeat:         1,
+		repeat:         len(instances), // Auto-set to run all instances with same input
 		verbose:        false,
+		factoryCache:   nil,
+		factoryCount:   0,
+		factoryCounted: false,
 	}
 }
 
@@ -367,7 +392,7 @@ func (p *Parallel) Forward(ctx context.Context, inputs map[string]any) (*core.Pr
 	}()
 
 	// Expand inputs into batch
-	batch, err := p.expandInputs(inputs)
+	batch, err := p.expandInputs(ctx, inputs)
 	if err != nil {
 		predErr = fmt.Errorf("failed to expand inputs: %w", err)
 		return nil, predErr
@@ -416,6 +441,14 @@ func (p *Parallel) Forward(ctx context.Context, inputs map[string]any) (*core.Pr
 	// Module getter
 	getModule := func(i int) core.Module {
 		if p.factory != nil {
+			// Check cache first (probed during expandInputs)
+			p.mu.Lock()
+			if i < len(p.factoryCache) && p.factoryCache[i] != nil {
+				mod := p.factoryCache[i]
+				p.mu.Unlock()
+				return mod
+			}
+			p.mu.Unlock()
 			return p.factory(i)
 		}
 		if len(p.instances) > 0 {
@@ -712,7 +745,7 @@ func (p *Parallel) Forward(ctx context.Context, inputs map[string]any) (*core.Pr
 }
 
 // expandInputs converts inputs into a slice of input maps
-func (p *Parallel) expandInputs(inputs map[string]any) ([]map[string]any, error) {
+func (p *Parallel) expandInputs(ctx context.Context, inputs map[string]any) ([]map[string]any, error) {
 	// Check for explicit batch
 	if batchVal, ok := inputs[p.batchKey]; ok {
 		batch, ok := batchVal.([]map[string]any)
@@ -760,7 +793,26 @@ func (p *Parallel) expandInputs(inputs map[string]any) ([]map[string]any, error)
 		return batch, nil
 	}
 
-	// No batch, no slices - repeat if configured
+	// No batch, no slices - check for factory-based module count
+	if p.factory != nil {
+		// Log warning if WithRepeat was also used
+		if p.repeat > 1 {
+			logging.GetLogger().Warn(ctx, "Parallel: WithRepeat(n) ignored because factory is present. Factory count takes precedence.", map[string]any{
+				"repeat_ignored": p.repeat,
+			})
+		}
+		// Probe factory to determine module count
+		n := p.probeFactoryCount()
+		batch := make([]map[string]any, n)
+		for i := 0; i < n; i++ {
+			taskInputs := make(map[string]any)
+			maps.Copy(taskInputs, inputs)
+			batch[i] = taskInputs
+		}
+		return batch, nil
+	}
+
+	// Repeat if configured
 	if p.repeat > 1 {
 		batch := make([]map[string]any, p.repeat)
 		for i := 0; i < p.repeat; i++ {
@@ -776,6 +828,39 @@ func (p *Parallel) expandInputs(inputs map[string]any) ([]map[string]any, error)
 	return []map[string]any{inputs}, nil
 }
 
+// probeFactoryCount determines how many modules the factory can create
+// by calling factory(i) until it returns nil. Returns at least 1.
+func (p *Parallel) probeFactoryCount() int {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	if p.factoryCounted {
+		return p.factoryCount
+	}
+	if p.factory == nil {
+		p.factoryCount = 1
+		p.factoryCounted = true
+		return 1
+	}
+
+	const maxProbe = 1000 // Safety limit
+	var cache []core.Module
+	for i := 0; i < maxProbe; i++ {
+		mod := p.factory(i)
+		if mod == nil {
+			p.factoryCount = i
+			p.factoryCache = cache
+			p.factoryCounted = true
+			return p.factoryCount
+		}
+		cache = append(cache, mod)
+	}
+
+	p.factoryCount = maxProbe
+	p.factoryCache = cache
+	p.factoryCounted = true
+	return maxProbe
+}
+
 // summarizeLatencies calculates min/max/avg/p50 from latencies
 func summarizeLatencies(latencies []time.Duration) struct {
 	MinMs int64 `json:"min_ms"`
@@ -857,6 +942,9 @@ func (p *Parallel) Clone() core.Module {
 		batchKey:       p.batchKey,
 		repeat:         p.repeat,
 		verbose:        p.verbose,
+		factoryCache:   nil,
+		factoryCount:   0,
+		factoryCounted: false,
 	}
 
 	// Only clone module if it exists

internal/module/parallel_test.go

@@ -1271,6 +1271,7 @@ type capturingLogger struct {
 	mu     sync.Mutex
 	infos  []logEntry
 	debugs []logEntry
+	warns  []logEntry
 }
 
 type logEntry struct {
@@ -1300,7 +1301,13 @@ func (c *capturingLogger) Debug(ctx context.Context, message string, fields map[
 }
 
 func (c *capturingLogger) Warn(ctx context.Context, message string, fields map[string]any) {
-	// Not used in these tests
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.warns = append(c.warns, logEntry{
+		level:   "warn",
+		message: message,
+		fields:  fields,
+	})
 }
 
 func (c *capturingLogger) Error(ctx context.Context, message string, fields map[string]any) {
@@ -1333,6 +1340,7 @@ func (c *capturingLogger) reset() {
 	defer c.mu.Unlock()
 	c.infos = c.infos[:0]
 	c.debugs = c.debugs[:0]
+	c.warns = c.warns[:0]
 }
 
 // mockLMWithName is a MockLM that returns a specific model name
@@ -2381,3 +2389,83 @@ func TestParallelWithFactory_ExactFactoryCalls(t *testing.T) {
 		t.Errorf("Expected 3 factory calls, got %d", calls)
 	}
 }
+
+func TestParallelFactoryCaching(t *testing.T) {
+	t.Parallel()
+	sig := core.NewSignature("Test").
+		AddOutput("value", core.FieldTypeString, "Value")
+
+	var factoryCalls atomic.Int32
+
+	factory := func(i int) core.Module {
+		factoryCalls.Add(1)
+		if i >= 3 {
+			return nil
+		}
+		lm := &MockLM{
+			GenerateFunc: func(ctx context.Context, messages []core.Message, options *core.GenerateOptions) (*core.GenerateResult, error) {
+				return &core.GenerateResult{
+					Content: "[[ ## value ## ]]\nsuccess",
+					Usage:   core.Usage{TotalTokens: 1, Cost: 0.001},
+				}, nil
+			},
+		}
+		return NewPredict(sig, lm)
+	}
+
+	p := NewParallelWithFactory(factory).WithMaxWorkers(2)
+
+	_, err := p.Forward(context.Background(), map[string]any{"some_input": "val"})
+	if err != nil {
+		t.Fatalf("Forward failed: %v", err)
+	}
+
+	if factoryCalls.Load() != 4 {
+		t.Errorf("Expected exactly 4 factory calls (3 successes + 1 nil probe), got %d", factoryCalls.Load())
+	}
+}
+
+func TestParallelWithRepeatWarning(t *testing.T) {
+	capturingLog := &capturingLogger{}
+	originalLogger := logging.GetLogger()
+	defer logging.SetLogger(originalLogger)
+	logging.SetLogger(capturingLog)
+
+	sig := core.NewSignature("Test")
+	factory := func(i int) core.Module {
+		if i >= 1 {
+			return nil
+		}
+		return NewPredict(sig, &MockLM{
+			GenerateFunc: func(ctx context.Context, messages []core.Message, options *core.GenerateOptions) (*core.GenerateResult, error) {
+				return &core.GenerateResult{
+					Content: "[[ ## result ## ]]\nok",
+					Usage:   core.Usage{},
+				}, nil
+			},
+		})
+	}
+
+	p := NewParallelWithFactory(factory).WithRepeat(5)
+
+	_, err := p.Forward(context.Background(), map[string]any{"dummy": "val"})
+	if err != nil {
+		t.Fatalf("Forward failed: %v", err)
+	}
+
+	capturingLog.mu.Lock()
+	defer capturingLog.mu.Unlock()
+
+	found := false
+	expectedMsg := "Parallel: WithRepeat(n) ignored because factory is present. Factory count takes precedence."
+	for _, w := range capturingLog.warns {
+		if w.message == expectedMsg {
+			found = true
+			break
+		}
+	}
+
+	if !found {
+		t.Errorf("Expected warning message %q not found in captured logs", expectedMsg)
+	}
+}