feat: ability to run tests in batches

Author: f-allian

causal_testing/__main__.py

@@ -31,7 +31,13 @@ def main() -> None:
 
     # Load and run tests
     framework.load_tests()
-    results = framework.run_tests(silent=args.silent)
+
+    if args.batch_size > 0:
+        logging.info(f"Running tests in batches of size {args.batch_size}")
+        results = framework.run_tests_in_batches(batch_size=args.batch_size, silent=args.silent)
+    else:
+        logging.info("Running tests in regular mode")
+        results = framework.run_tests(silent=args.silent)
 
     # Save results
     framework.save_results(results)

causal_testing/main.py

@@ -9,6 +9,7 @@
 from tqdm import tqdm
 
 import pandas as pd
+import numpy as np
 
 from .specification.causal_dag import CausalDAG
 from .specification.scenario import Scenario
@@ -22,7 +23,6 @@
 from .estimation.logistic_regression_estimator import LogisticRegressionEstimator
 
 logger = logging.getLogger(__name__)
-logger.setLevel(logging.ERROR)
 
 
 @dataclass
@@ -307,6 +307,62 @@ def create_causal_test(self, test: dict, base_test: BaseTestCase) -> CausalTestC
             estimator=estimator,
         )
 
+    def run_tests_in_batches(self, batch_size=100, silent=False) -> List[CausalTestResult]:
+        """
+        Run tests in batches to reduce memory usage.
+
+        :param batch_size: Number of tests to run in each batch
+        :param silent: Whether to suppress errors
+        :return: List of all test results
+        :raises: ValueError if no tests are loaded
+        """
+        logger.info("Running causal tests in batches...")
+
+        if not self.test_cases:
+            raise ValueError("No tests loaded. Call load_tests() first.")
+
+        num_tests = len(self.test_cases)
+        num_batches = int(np.ceil(num_tests / batch_size))
+
+        logger.info(f"Processing {num_tests} tests in {num_batches} batches of up to {batch_size} tests each")
+        all_results = []
+        with tqdm(total=num_tests, desc="Overall progress", mininterval=0.1) as overall_pbar:
+            # Process each batch
+            for batch_idx in range(num_batches):
+                start_idx = batch_idx * batch_size
+                end_idx = min(start_idx + batch_size, num_tests)
+                current_batch_size = end_idx - start_idx
+
+                logger.info(f"Processing batch {batch_idx + 1} of {num_batches} (tests {start_idx} to {end_idx - 1})")
+
+                # Get current batch of tests
+                current_batch = self.test_cases[start_idx:end_idx]
+
+                # Process the current batch
+                batch_results = []
+                for test_case in current_batch:
+                    try:
+                        result = test_case.execute_test()
+                        batch_results.append(result)
+                    except Exception as e:
+                        if not silent:
+                            logger.error(f"Error running test: {str(e)}")
+                            raise
+                        result = CausalTestResult(
+                            estimator=test_case.estimator,
+                            test_value=TestValue("Error", str(e)),
+                        )
+                        batch_results.append(result)
+
+                    overall_pbar.update(1)
+
+                all_results.extend(batch_results)
+
+                logger.info(f"Completed batch {batch_idx + 1} of {num_batches} ({current_batch_size} tests)")
+
+        logger.info(f"Completed processing all {len(all_results)} tests in {num_batches} batches")
+        return all_results
+
     def run_tests(self, silent=False) -> List[CausalTestResult]:
         """
         Run all test cases and return their results.
@@ -418,5 +474,11 @@ def parse_args(args: Optional[Sequence[str]] = None) -> argparse.Namespace:
         help="Do not crash on error. If set to true, errors are recorded as test results.",
         default=False,
     )
+    parser.add_argument(
+        "--batch-size",
+        type=int,
+        default=0,
+        help="Run tests in batches of the specified size (default: 0, which means no batching)",
+    )
 
     return parser.parse_args(args)