← Back to Docs

PYTHON POOLING MANUAL AND TESTING

Documentation for PYTHON_POOLING_MANUAL_AND_TESTING from the Dspex repository.

Of course. Based on the extensive information in your repository, here is a comprehensive technical manual on the Python bridge pooling implementation, its design, and its integration into your three-tier testing architecture.

Technical Manual: Python Bridge Pooling

Version: 1.0
Date: 2025-07-14

1. Introduction

1.1. Purpose

This document provides a comprehensive technical overview of the DSPex Python bridge pooling system. It details the architecture, design principles, implementation, configuration, and testing strategy for managing concurrent, isolated Python processes.

1.2. Overview of the Pooling System

The Python bridge is a critical component for integrating Elixir with the Python DSPy ecosystem. The initial single-process bridge design, while functional, presented significant limitations in scalability, concurrency, and state management, leading to test flakiness and production bottlenecks.

The pooling system, built on NimblePool, addresses these challenges by replacing the single Python process with a managed pool of PoolWorker processes. This architecture provides:

  • Scalability: Handles a high volume of concurrent requests by distributing work across multiple Python runtimes.
  • Isolation: Guarantees session-level isolation, preventing state pollution and program ID conflicts between different users or tasks.
  • Resource Efficiency: Reuses warm Python processes to minimize startup overhead and manage system resources effectively.
  • Reliability: Implements health monitoring and automatic recovery for worker processes, ensuring system stability.

1.3. Target Audience

This manual is intended for:

  • Software Engineers developing and maintaining the DSPex library.
  • Site Reliability Engineers (SREs) deploying and monitoring the system in production.
  • Quality Assurance (QA) Engineers designing and executing tests for the system.

2. Architecture and Design

2.1. High-Level Architecture

The pooling system is managed by a dedicated supervision tree, ensuring fault tolerance and proper lifecycle management. It is designed to be a drop-in replacement for the single bridge, abstracted via the DSPex.Adapters.PythonPool adapter.

Application Layer
       │
       ▼
┌──────────────────────┐
│ DSPex.Adapters.      │
│   PythonPool         │  <-- Elixir Interface
└──────────────────────┘
       │
       ▼
┌──────────────────────┐
│ PoolSupervisor       │  <-- Top-Level Supervisor
└──────────────────────┘
       │
       ├─►┌──────────────────────────┐
       │  │ SessionPool (GenServer)  │  <-- Pool & Session Manager
       │  └──────────────────────────┘
       │             │
       │             ▼
       │  ┌──────────────────────────┐
       │  │      NimblePool          │
       │  └──────────────────────────┘
       │    ├───► PoolWorker 1 ◄───► Python Process 1 (dspy_bridge.py --mode pool-worker)
       │    ├───► PoolWorker 2 ◄───► Python Process 2
       │    └───► PoolWorker N ◄───► Python Process N
       │
       └─►┌──────────────────────────┐
          │ PoolMonitor (GenServer)  │  <-- Health & Metrics Monitor
          └──────────────────────────┘

2.2. Key Components

2.2.1. DSPex.PythonBridge.PoolSupervisor

The entry point and top-level supervisor for the entire pooling system. It starts and manages the SessionPool and PoolMonitor, ensuring that the core components are always running.

2.2.2. DSPex.PythonBridge.SessionPool

This GenServer is the heart of the pool management system. It leverages NimblePool to manage a dynamic set of PoolWorker processes. Its key responsibilities include:

  • Pool Lifecycle: Starting, stopping, and supervising the NimblePool instance.
  • Session Management: Tracking active sessions and their state.
  • Request Routing: Checking out available workers for session-specific or anonymous operations.
  • Graceful Shutdown: Ensuring all workers and sessions are terminated cleanly.

2.2.3. DSPex.PythonBridge.PoolWorker

An implementation of the NimblePool worker behaviour. Each PoolWorker process is responsible for a single, dedicated Python process.

  • Process Management: Spawns and manages a dspy_bridge.py process in pool-worker mode.
  • Communication: Manages the Port connection for message passing.
  • Session Affinity: Can be temporarily “bound” to a session ID during checkout to maintain context for a series of related operations.
  • Health & State: Tracks its own health status (:initializing, :healthy, :unhealthy) and statistics (requests handled, errors, etc.).

2.2.4. DSPex.Adapters.PythonPool

The public-facing Elixir module that implements the DSPex.Adapters.Adapter behaviour. It provides a clean, session-aware API for creating and executing programs, abstracting away the complexities of pool management. It transparently handles session creation and command execution via the SessionPool.

2.2.5. Python Bridge (dspy_bridge.py) Enhancements

The Python script is enhanced to support the pooling architecture:

  • --mode pool-worker: A command-line argument that starts the script in a mode ready to be managed by a PoolWorker.
  • Session-Namespaced Programs: When in pool-worker mode, the DSPyBridge class stores programs in a nested dictionary: self.session_programs[session_id][program_id]. This is the core mechanism for session isolation.
  • Session Cleanup: The bridge includes a cleanup_session command to purge all data associated with a specific session when it ends.
  • Stateless Operation: Each command from Elixir includes a session_id, making each call self-contained from the pool’s perspective.

2.3. Session Management and Isolation

The pooling architecture guarantees isolation through a combination of Elixir-side and Python-side logic.

  1. Session Creation: A session is implicitly created in the SessionPool the first time a session_id is used.
  2. Worker Checkout: When a request for a specific session_id arrives, SessionPool checks out an available PoolWorker. NimblePool gives preference to a worker that last served the same session (session affinity), but any available worker can be used.
  3. Command Execution: The session_id is passed with every command to the Python worker. The Python script uses this session_id as the primary key for its program registry, ensuring that program_ids from different sessions do not conflict.
  4. Worker Check-in: After the operation completes, the worker is checked back into the pool, ready to serve another request from any session.
  5. Session Termination: When SessionPool.end_session/1 is called, the pool broadcasts a cleanup command to any worker that might hold state for that session, ensuring resources are freed on the Python side.

3. Implementation Details

3.1. Elixir Implementation

The core of the Elixir implementation lies in the PoolWorker and SessionPool modules.

PoolWorker Initialization (init_worker/1) The worker starts a Python process in the correct mode and performs an initial “ping” to ensure readiness.

# in lib/dspex/python_bridge/pool_worker.ex
@impl NimblePool
def init_worker(pool_state) do
  worker_id = generate_worker_id()
  
  # ... validate environment ...
  
  # Start Python process in pool-worker mode
  port_opts = [
    :binary,
    :exit_status,
    {:packet, 4},
    {:args, [script_path, "--mode", "pool-worker", "--worker-id", worker_id]}
  ]
  port = Port.open({:spawn_executable, python_path}, port_opts)
  
  # ... initialize worker state ...

  # Send initialization ping to ensure the python process is ready
  case send_initialization_ping(worker_state) do
    {:ok, updated_state} ->
      Logger.info("Pool worker #{worker_id} started successfully")
      {:ok, updated_state, pool_state}
    {:error, reason} ->
      # ... handle error ...
  end
end

SessionPool Execution (execute_in_session/4) This function demonstrates the checkout-execute-checkin pattern using NimblePool.

# in lib/dspex/python_bridge/session_pool.ex
def execute_in_session(session_id, command, args, opts \\ []) do
  pool_timeout = Keyword.get(opts, :pool_timeout, @default_checkout_timeout)
  
  # ... track session ...

  # Execute with NimblePool
  try do
    NimblePool.checkout!(
      state.pool_name,
      {:session, session_id}, # Checkout with session context for affinity
      fn _from, worker_state ->
        # This block runs in the checked-out worker process
        case PoolWorker.send_command(worker_state, command, args, operation_timeout) do
          # ... handle response ...
        end
      end,
      pool_timeout
    )
  catch
    # ... handle timeout and other errors ...
  end
end

3.2. Python (dspy_bridge.py) Implementation

The Python script is designed to be stateless from the perspective of a single request, receiving all necessary context in the command arguments.

Command Handling with Session Awareness

# in priv/python/dspy_bridge.py
class DSPyBridge:
    def __init__(self, mode="standalone", worker_id=None):
        self.mode = mode
        self.worker_id = worker_id
        # In standalone mode, programs are stored globally
        self.programs: Dict[str, Any] = {}
        # In pool-worker mode, programs are namespaced by session
        self.session_programs: Dict[str, Dict[str, Any]] = {}
        # ... other initializations

    def create_program(self, args: Dict[str, Any]) -> Dict[str, Any]:
        program_id = args.get('id')
        
        # Handle session-based storage in pool-worker mode
        if self.mode == "pool-worker":
            session_id = args.get('session_id')
            if not session_id:
                raise ValueError("session_id is required in pool-worker mode")
            
            # Initialize session if needed
            if session_id not in self.session_programs:
                self.session_programs[session_id] = {}
            
            program_registry = self.session_programs[session_id]
        else:
            program_registry = self.programs

        if program_id in program_registry:
            raise ValueError(f"Program with ID '{program_id}' already exists.")

        # ... create program and store it in program_registry ...

4. Configuration

The pooling system is highly configurable via config/config.exs or environment-specific files.

4.1. Enabling Pooling

To enable the pooling architecture, set the :python_bridge_pool_mode configuration flag. This is typically done in your environment configs (e.g., prod.exs).

# In config/prod.exs
import Config

# Enable pool mode for the Python bridge
config :dspex, :python_bridge_pool_mode, true

4.2. Core Pool Settings

The primary configuration is for the PoolSupervisor.

# in config/pool_config.exs or your environment config
config :dspex, DSPex.PythonBridge.PoolSupervisor,
  # Number of Python worker processes
  pool_size: System.schedulers_online() * 2,
  
  # Maximum additional workers created under load
  max_overflow: System.schedulers_online(),
  
  # Max time to wait for a worker (ms)
  checkout_timeout: 5_000,
  
  # How often to perform health checks (ms)
  health_check_interval: 30_000,
  
  # Whether to start workers lazily on first use
  lazy: false

4.3. Environment-Specific Tuning

The system is designed for different configurations per environment.

# In config/dev.exs
config :dspex, DSPex.PythonBridge.PoolSupervisor,
  pool_size: 2,
  lazy: true

# In config/test.exs
config :dspex, DSPex.PythonBridge.PoolSupervisor,
  pool_size: 4,
  checkout_timeout: 10_000

# In config/prod.exs
config :dspex, DSPex.PythonBridge.PoolSupervisor,
  pool_size: System.schedulers_online() * 3,
  lazy: false

5. Testing Strategy

The pooling implementation is thoroughly tested across our Three-Tier Testing Architecture.

5.1. The Three-Tier Testing Architecture

Our testing philosophy is layered to provide a balance between speed and confidence:

  • Layer 1: Fast Unit Tests (mix test.fast)

    • Scope: Pure Elixir logic, isolated modules.
    • Mocks: All external dependencies (like Python processes) are mocked.
    • Speed: Milliseconds.
    • Goal: Rapid feedback for developers during TDD.

  • Layer 2: Protocol Validation (mix test.protocol)

    • Scope: Communication protocol between Elixir and Python.
    • Mocks: A mock server (BridgeMockServer) simulates the Python side, validating the wire protocol without requiring a full Python/DSPy environment.
    • Speed: Sub-second.
    • Goal: Ensure data serialization and command handling are correct.

  • Layer 3: Full Integration (mix test.integration)

    • Scope: The complete end-to-end system.
    • Mocks: No mocks. Uses real Python processes, a real DSPy installation, and can make real calls to LLM APIs (if configured).
    • Speed: Seconds per test.
    • Goal: Verify the entire system works together as expected.

5.2. How Pooling Factors In

Pooling tests are designed to cover all three layers:

  • Layer 1 (Unit Tests):

    • PoolWorkerUnitTest and SessionPoolMockTest test the internal logic of the components in isolation.
    • The Python process is mocked using DSPex.Test.MockPort or by using the test process PID itself, allowing for fast, deterministic testing of state transitions and logic without any I/O.
    # in test/dspex/python_bridge/pool_worker_unit_test.exs
test "session checkout binds to session", %{worker: worker} do
  checkout_type = {:session, "user_123"}
  from = {self(), make_ref()}

  # Simulates the NimblePool checkout callback
  {:ok, _, updated_state, _} = PoolWorkerHelpers.simulate_checkout(
    checkout_type, from, worker, %{}
  )

  assert updated_state.current_session == "user_123"
  assert updated_state.stats.checkouts == 1
end

  • Layer 2 (Protocol Tests):

    • Tests for the PythonPool adapter would use the BridgeMock adapter to ensure that session information is correctly serialized and passed according to the wire protocol. This confirms the pooling layer correctly communicates session context.

  • Layer 3 (Integration Tests):

    • PoolWorkerIntegrationTest starts a real Python process for a worker and verifies the full lifecycle.
    • SessionPoolTest and SessionPoolConcurrencyTest start a full NimblePool of real Python workers to test session management, concurrency, and isolation under load.
    # in test/dspex/python_bridge/session_pool_test.exs
@moduletag :layer_3
test "handles concurrent sessions without interference", %{pool: pool} do
  sessions = for i <- 1..10, do: "concurrent_#{i}"

  tasks = Enum.map(sessions, fn session_id ->
    Task.async(fn ->
      # Each session creates and executes its own program
      # The pooling system ensures these run in parallel on different workers
      # without interfering with each other.
      {:ok, program_id} = SessionPool.execute_in_session(
        pool, session_id, :create_program, %{id: "prog_#{session_id}"}
      )
      # ...
    end)
  end)

  results = Task.await_many(tasks, 30_000)
  assert Enum.all?(results, &(&1 == :ok))
end

5.3. Advanced Testing: Chaos and Stress

The pool_tests.yml.disabled workflow file outlines a strategy for advanced testing to ensure production readiness:

  • Stress Tests: Run nightly to simulate high-load conditions over extended periods, testing the pool’s stability and overflow handling.
  • Chaos Tests: Inject faults into the system, such as randomly killing PoolWorker processes, to verify that the PoolSupervisor correctly recovers and maintains system health.

5.4. CI/CD Strategy

The CI/CD pipeline is configured to run tests based on the layer, providing fast feedback for most changes while ensuring full system validation before merging to main.

  • On Pull Request: Run Layer 1 and Layer 2 tests.
  • On Merge to develop/main: Run all three layers.
  • Nightly: Run Layer 3 tests plus Stress and Chaos tests.

6. Usage and Operations

6.1. Basic Usage

The PythonPool adapter abstracts away the complexity. Developers interact with a session-aware API.

# 1. Create a session-bound adapter instance
adapter = DSPex.Adapters.PythonPool.session_adapter("user_123_session_abc")

# 2. Use the adapter to create and execute programs
config = %{id: "my_program", signature: MySignature}
{:ok, program_id} = adapter.create_program(config)

inputs = %{question: "What is Elixir?"}
{:ok, result} = adapter.execute_program(program_id, inputs, %{})

# The session is automatically managed by the adapter and pool.
# No need to manually end the session unless resources need immediate cleanup.
DSPex.Adapters.PythonPool.end_session("user_123_session_abc")

6.2. Monitoring and Telemetry

The pooling system emits telemetry events for key operations, which can be hooked into for monitoring and logging.

Available Events:

  • [:dspex, :python_bridge, :pool, :checkout]
  • [:dspex, :python_bridge, :pool, :checkin]
  • [:dspex, :python_bridge, :pool, :worker_restarted]
  • [:dspex, :python_bridge, :pool, :timeout]

The PoolMonitor continuously checks the health of the workers and the overall pool, exposing metrics like:

  • active_sessions
  • pool_size & utilization
  • healthy_workers & unhealthy_workers
  • total_commands & total_errors

6.3. Troubleshooting

  • "Pool timeout - all workers busy": The pool is saturated. Increase pool_size or max_overflow, or investigate long-running Python operations.
  • "Session not found": Ensure the session_id is consistent across calls and that the session has not been explicitly ended or timed out.
  • Frequent Worker Restarts: Check the logs for PoolWorker errors. This often points to issues in the Python environment (e.g., package incompatibilities, memory leaks in the Python script). Use PoolSupervisor.get_stats() to monitor restart counts.

7. Conclusion

The NimblePool-based architecture provides a robust, scalable, and reliable foundation for integrating Elixir with Python DSPy. It solves the critical challenges of concurrency and state isolation posed by the single-process model. By leveraging a layered testing strategy and providing comprehensive configuration and monitoring capabilities, the pooling system is designed for both high-performance development and production-grade stability.