Merge pull request PSLmodels#1049 from jdebacker/dask-performance-benchmarks

Merging

Author: rickecon

.github/workflows/build_and_test.yml

@@ -50,7 +50,7 @@ jobs:
         shell: bash -l {0}
         working-directory: ./
         run: |
-          python -m pytest -m "not local" --cov=./ --cov-report=xml
+          python -m pytest -m "not local and not benchmark" --cov=./ --cov-report=xml
       - name: Upload coverage to Codecov
         if: matrix.os == 'ubuntu-latest' && contains(github.repository, 'PSLmodels/OG-Core')
         uses: codecov/codecov-action@v4

CHANGELOG.md

@@ -6,6 +6,37 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 
+## [0.14.8] - 2025-08-26 12:00:00
+
+### Added
+
+- Adds a complete benchmark suite for measuring and optimizing Dask performance in OG-Core, with particular focus on Windows performance issues.
+- New and updated files:
+    - tests/test_dask_benchmarks.py: Mock benchmark tests with synthetic workloads
+    - tests/test_real_txfunc_benchmarks.py: Real-world tax function benchmarks
+    - tests/run_benchmarks.py: Automated benchmark runner with reporting
+    - tests/BENCHMARK_README.md: Comprehensive documentation and usage guide
+    - pytest.ini: Updated with benchmark test markers
+- Key features:
+    - Platform-specific optimization tests (Windows, macOS, Linux)
+    - Memory usage and compute time benchmarking
+    - Baseline establishment and performance regression detection
+    - Comparison of different Dask schedulers and client configurations
+    - Real tax function estimation performance measurement
+    - Automated identification of optimal Dask settings per platform
+- Benefits:
+    - Establishes performance baselines before optimization work
+    - Identifies Windows-specific Dask performance bottlenecks
+    - Provides automated regression detection for future changes
+    - Enables data-driven optimization decisions
+    - Supports continuous performance monitoring
+- Usage:
+    - `python tests/run_benchmarks.py  # Run all benchmarks`
+    - `python tests/run_benchmarks.py --quick  # Quick benchmarks only`
+    - `python tests/run_benchmarks.py --save-baseline  # Save performance baseline`
+    - `python tests/run_benchmarks.py --compare-baseline # Compare against baseline`
+- 🤖 Generated with help from Claude Code
+
 ## [0.14.7] - 2025-08-21 17:00:00
 
 ### Added
@@ -409,6 +440,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Any earlier versions of OG-USA can be found in the [`OG-Core`](https://github.com/PSLmodels/OG-Core) repository [release history](https://github.com/PSLmodels/OG-Core/releases) from [v.0.6.4](https://github.com/PSLmodels/OG-Core/releases/tag/v0.6.4) (Jul. 20, 2021) or earlier.
 
 
+[0.14.8]: https://github.com/PSLmodels/OG-Core/compare/v0.14.7...v0.14.8
 [0.14.7]: https://github.com/PSLmodels/OG-Core/compare/v0.14.6...v0.14.7
 [0.14.6]: https://github.com/PSLmodels/OG-Core/compare/v0.14.5...v0.14.6
 [0.14.5]: https://github.com/PSLmodels/OG-Core/compare/v0.14.4...v0.14.5

ogcore/__init__.py

@@ -20,4 +20,4 @@
 from ogcore.txfunc import *
 from ogcore.utils import *
 
-__version__ = "0.14.7"
+__version__ = "0.14.8"

ogcore/household.py

@@ -7,6 +7,14 @@
 # Packages
 import numpy as np
 from ogcore import tax, utils
+import logging
+
+# Configure logging
+VERBOSE = False
+log_level = logging.INFO if VERBOSE else logging.WARNING
+logging.basicConfig(
+    level=log_level, format="%(message)s"  # Only show the message itself
+)
 
 """
 ------------------------------------------------------------------------
@@ -764,28 +772,30 @@ def constraint_checker_SS(bssmat, nssmat, cssmat, ltilde):
         Warnings: if constraints are violated, warnings printed
 
     """
-    print("Checking constraints on capital, labor, and consumption.")
+    logging.info("Checking constraints on capital, labor, and consumption.")
 
     if (bssmat < 0).any():
-        print("\tWARNING: There is negative capital stock")
+        logging.info("\tWARNING: There is negative capital stock")
     flag2 = False
     if (nssmat < 0).any():
-        print(
+        logging.info(
             "\tWARNING: Labor supply violates nonnegativity ", "constraints."
         )
         flag2 = True
     if (nssmat > ltilde).any():
-        print("\tWARNING: Labor supply violates the ltilde constraint.")
+        logging.info("\tWARNING: Labor supply violates the ltilde constraint.")
         flag2 = True
     if flag2 is False:
-        print(
+        logging.info(
             "\tThere were no violations of the constraints on labor",
             " supply.",
         )
     if (cssmat < 0).any():
-        print("\tWARNING: Consumption violates nonnegativity", " constraints.")
+        logging.info(
+            "\tWARNING: Consumption violates nonnegativity", " constraints."
+        )
     else:
-        print(
+        logging.info(
             "\tThere were no violations of the constraints on", " consumption."
         )
 
@@ -810,22 +820,22 @@ def constraint_checker_TPI(b_dist, n_dist, c_dist, t, ltilde):
 
     """
     if (b_dist <= 0).any():
-        print(
+        logging.info(
             "\tWARNING: Aggregate capital is less than or equal to ",
             "zero in period %.f." % t,
         )
     if (n_dist < 0).any():
-        print(
+        logging.info(
             "\tWARNING: Labor supply violates nonnegativity",
             " constraints in period %.f." % t,
         )
     if (n_dist > ltilde).any():
-        print(
+        logging.info(
             "\tWARNING: Labor suppy violates the ltilde constraint",
             " in period %.f." % t,
         )
     if (c_dist < 0).any():
-        print(
+        logging.info(
             "\tWARNING: Consumption violates nonnegativity",
             " constraints in period %.f." % t,
         )

pytest.ini

@@ -5,3 +5,10 @@ testpaths =
     ./tests
 markers =
     local: marks tests that run locally and not on GH Actions (mostly due to run time)
+    benchmark: marks tests that measure performance and memory usage
+    distributed: marks tests that use distributed Dask clients
+    memory: marks tests focused on memory usage measurement
+    performance: marks tests focused on compute time measurement
+    slow: marks tests that take longer to run
+    real: marks tests using real OG-Core tax function code
+    platform: marks tests for platform-specific optimization

setup.py

@@ -5,7 +5,7 @@
 
 setuptools.setup(
     name="ogcore",
-    version="0.14.7",
+    version="0.14.8",
     author="Jason DeBacker and Richard W. Evans",
     license="CC0 1.0 Universal (CC0 1.0) Public Domain Dedication",
     description="A general equilibrium overlapping generations model for fiscal policy analysis",