Skip to content

Conversation

@jhamon
Copy link
Collaborator

@jhamon jhamon commented Nov 7, 2025

Summary

This PR implements a custom pytest plugin for test sharding, allowing tests to be automatically distributed across multiple CI jobs for parallel execution. This replaces the previous manual directory-based test splitting approach with a more flexible, hash-based distribution system. Additionally, integration tests have been reorganized into top-level folders that group tests by the client type and setup requirements needed.

This should bring down the total CI runtime to ~8 minutes or less.

Changes

Core Implementation

  • New pytest plugin (tests/pytest_shard.py):

    • Implements pytest_addoption hook to add --splits and --group command-line options
    • Implements pytest_collection_modifyitems hook to filter tests based on shard assignment
    • Uses hash-based distribution (MD5 hash of test node ID) for deterministic test assignment
    • Supports environment variables PYTEST_SPLITS and PYTEST_GROUP as alternatives to command-line options
    • Includes validation for shard parameters with helpful error messages
  • Plugin registration (tests/conftest.py):

    • Registers the plugin globally so it's available for all test runs
    • Plugin is automatically loaded when running pytest

CI Workflow Updates

  • Updated .github/workflows/testing-integration.yaml:

    • Replaced manual directory-based test splitting with automatic sharding
    • rest_sync tests: Now uses 8 shards (was manually split by directory)
    • rest_asyncio tests: Now uses 5 shards (was manually split by directory)
    • grpc tests: No sharding (runs all tests in single job, including tests/integration/rest_sync/db/data with USE_GRPC='true')
  • Updated .github/actions/run-integration-test/action.yaml:

    • Added pytest_splits and pytest_group input parameters
    • Updated test execution to pass sharding arguments when provided

Test Reorganization

  • Integration tests reorganized by client type (tests/integration/):
    • rest_sync/: Tests using the synchronous REST client (Pinecone())
      • Uses standard Index() objects for database operations
      • Supports optional GRPC mode via USE_GRPC='true' environment variable
      • Contains subdirectories for db/ (control and data operations), inference/, and admin/ tests
    • rest_asyncio/: Tests using the asynchronous REST client (Pinecone().IndexAsyncio())
      • Uses async fixtures and IndexAsyncio() objects
      • Requires pytest-asyncio for async test execution
      • Contains subdirectories for db/ (control and data operations) and inference/ tests
    • grpc/: Tests using the GRPC client (PineconeGRPC())
      • Uses PineconeGRPC() client and GRPCIndex objects
      • Contains db/data/ tests for GRPC-specific functionality
    • This organization makes it clear which client type each test requires and simplifies fixture setup

Bug Fixes

  • Fixed race condition in test cleanup (tests/integration/rest_sync/db/control/pod/conftest.py):
    • Added NotFoundException handling in attempt_delete_index function
    • Prevents teardown errors when index is deleted between has_index check and describe_index call

Testing

  • Unit tests (tests/unit/test_pytest_shard.py):
    • Tests for hash-based distribution logic
    • Tests for validation and error handling
    • Tests for deterministic shard assignment
    • Tests for edge cases (single shard, environment variables, etc.)
    • Tests gracefully handle testdir limitations (plugin loading in isolated environments)

Documentation

  • Updated docs/maintainers/testing-guide.md:
    • Added "Test Sharding" section with usage examples
    • Documented command-line options and environment variables
    • Explained how sharding works and its use in CI
    • Documented actual shard counts used in CI workflows
    • Fixed broken link to testing-integration.yaml

Benefits

  1. Automatic test distribution: Tests are automatically distributed across shards using a deterministic hash algorithm, eliminating the need to manually maintain directory-based splits
  2. Better load balancing: Hash-based distribution ensures more even test distribution across shards compared to directory-based splitting
  3. Easier maintenance: No need to manually update CI workflows when test files are added, removed, or reorganized
  4. Flexibility: Shard counts can be easily adjusted in CI workflows without code changes
  5. Deterministic: Same test always goes to the same shard, making debugging easier
  6. Clear test organization: Tests are grouped by client type, making it immediately clear which setup and fixtures are needed for each test
  7. Simplified fixture management: Each client type has its own conftest.py with appropriate fixtures, reducing complexity and potential conflicts

Usage

Command-line

pytest tests/integration/rest_sync --splits=8 --group=1

Environment variables

export PYTEST_SPLITS=8
export PYTEST_GROUP=1
pytest tests/integration/rest_sync

Testing

  • Plugin works correctly in real pytest environment
  • CI workflows updated and ready for use

Notes

  • The plugin is automatically available when running pytest (no installation needed)
  • Shard counts in CI can be adjusted based on test suite size and CI capacity

@jhamon jhamon changed the title Integration test reorg Integration test reorg and build sharding Nov 14, 2025
@jhamon jhamon marked this pull request as ready for review November 14, 2025 18:32
@jhamon jhamon merged commit f44ca91 into release-candidate/2025-10 Nov 14, 2025
60 checks passed
@jhamon jhamon deleted the jhamon/test-reorg branch November 14, 2025 18:32
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants