Changed run_tests_in_batches. Still takes up all the memory

Author: jmafoster1

causal_testing/__main__.py

@@ -1,6 +1,10 @@
 """This module contains the main entrypoint functionality to the Causal Testing Framework."""
 
 import logging
+import tempfile
+import json
+import os
+
 from .main import setup_logging, parse_args, CausalTestingPaths, CausalTestingFramework
 
 
@@ -34,13 +38,27 @@ def main() -> None:
 
     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)
+        with tempfile.TemporaryDirectory() as tmpdir:
+            output_files = []
+            for i, results in enumerate(framework.run_tests_in_batches(batch_size=args.batch_size, silent=args.silent)):
+                temp_file_path = os.path.join(tmpdir, f"output_{i}.json")
+                framework.save_results(results, temp_file_path)
+                output_files.append(temp_file_path)
+                del results
+
+            # Now stitch the results together from the temporary files
+            all_results = []
+            for file_path in output_files:
+                with open(file_path, "r") as f:
+                    all_results.extend(json.load(f))
+
+            # Save the final stitched results to your desired location
+            with open(args.output, "w") as f:
+                json.dump(all_results, f, indent=4)
     else:
         logging.info("Running tests in regular mode")
         results = framework.run_tests(silent=args.silent)
-
-    # Save results
-    framework.save_results(results)
+        framework.save_results(results)
 
     logging.info("Causal testing completed successfully.")
 

causal_testing/main.py

@@ -8,6 +8,7 @@
 from typing import Dict, Any, Optional, List, Union, Sequence
 from tqdm import tqdm
 
+
 import pandas as pd
 import numpy as np
 
@@ -344,7 +345,6 @@ def run_tests_in_batches(self, batch_size: int = 100, silent: bool = False) -> L
         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 progress:
             # Process each batch
             for batch_idx in range(num_batches):
@@ -360,26 +360,23 @@ def run_tests_in_batches(self, batch_size: int = 100, silent: bool = False) -> L
                 batch_results = []
                 for test_case in current_batch:
                     try:
-                        result = test_case.execute_test()
-                        batch_results.append(result)
-                    except (TypeError, AttributeError) as e:
+                        batch_results.append(test_case.execute_test())
+                    # pylint: disable=broad-exception-caught
+                    except Exception as e:
                         if not silent:
                             logger.error(f"Type or attribute error in test: {str(e)}")
                             raise
-                        result = CausalTestResult(
-                            estimator=test_case.estimator,
-                            test_value=TestValue("Error", str(e)),
+                        batch_results.append(
+                            CausalTestResult(
+                                estimator=test_case.estimator,
+                                test_value=TestValue("Error", str(e)),
+                            )
                         )
-                        batch_results.append(result)
 
                     progress.update(1)
 
-                all_results.extend(batch_results)
-
-                logger.info(f"Completed batch {batch_idx + 1} of {num_batches}")
-
-        logger.info(f"Completed processing all {len(all_results)} tests in {num_batches} batches")
-        return all_results
+                yield batch_results
+        logger.info(f"Completed processing in {num_batches} batches")
 
     def run_tests(self, silent=False) -> List[CausalTestResult]:
         """
@@ -399,7 +396,6 @@ def run_tests(self, silent=False) -> List[CausalTestResult]:
             try:
                 result = test_case.execute_test()
                 results.append(result)
-                logger.info(f"Test completed: {test_case}")
             # pylint: disable=broad-exception-caught
             except Exception as e:
                 if not silent:
@@ -414,9 +410,11 @@ def run_tests(self, silent=False) -> List[CausalTestResult]:
 
         return results
 
-    def save_results(self, results: List[CausalTestResult]) -> None:
+    def save_results(self, results: List[CausalTestResult], output_path: str = None) -> None:
         """Save test results to JSON file in the expected format."""
-        logger.info(f"Saving results to {self.paths.output_path}")
+        if output_path is None:
+            output_path = self.paths.output_path
+        logger.info(f"Saving results to {output_path}")
 
         # Load original test configs to preserve test metadata
         with open(self.paths.test_config_path, "r", encoding="utf-8") as f:
@@ -460,7 +458,7 @@ def save_results(self, results: List[CausalTestResult]) -> None:
             json_results.append(output)
 
         # Save to file
-        with open(self.paths.output_path, "w", encoding="utf-8") as f:
+        with open(output_path, "w", encoding="utf-8") as f:
             json.dump(json_results, f, indent=2)
 
         logger.info("Results saved successfully")