Merge pull request CEMeNT-PSAAP#161 from jpmorgan98/main

Adding Caching and a Debug Runtime Option

.github/workflows/python-publish.yml

@@ -6,6 +6,8 @@
 # separate terms of service, privacy policy, and support
 # documentation.
 
+# assumming this is being done in a repo with trusted publishing permissions in pypi
+
 name: Upload Python Package
 
 on:
@@ -17,10 +19,10 @@ permissions:
 
 jobs:
   deploy:
-
     runs-on: ubuntu-latest
-
-    steps:
+    permissions:
+        id-token: write
+    steps: 
     - uses: actions/checkout@v3
     - name: Set up Python 3.11
       uses: actions/setup-python@v3
@@ -38,8 +40,9 @@ jobs:
         bash .github/workflows/patch.sh
     - name: Build package
       run: python -m build
-    - name: Publish package
-      uses: pypa/gh-action-pypi-publish@27b31702a0e7fc50959f5ad993c78deac1bdfc29
+    - name: Publish package distributions to PyPI
+      uses: pypa/gh-action-pypi-publish@release/v1
+
       #with:
       #  user: __token__
       #  password: ${{ secrets.PYPI_API_TOKEN }}

docs/paper.md

@@ -107,7 +107,7 @@ We have implemented novel algorithms using this scheme and, in some verification
 It uses the Numba compiler for Python to compile compute kernels to a desired hardware target, including support for graphics processing units (GPUs) [@lam_numba_2015].
 `MC/DC` uses `mpi4py` for distributed-memory parallelism [@mpi4py_2021] and has run at the scale of tens of thousands of processors [@variansyah_mc23_mcdc].
 These acceleration and abstraction techniques allow `MC/DC` developers to remain in a pure Python development environment without needing to support compiled or domain-specific languages.
-This has allowed `MC/DC` to grow from its instantiation less than two years ago into a codebase that supports full performant neutron transport and investigation of novel transport algorithms, with development mostly from relative novices.
+This has allowed `MC/DC` to grow from its initialization less than two years ago into a codebase that supports full performant neutron transport and investigation of novel transport algorithms, with development mostly from relative novices.
 
 Many of the traditionally developed neutron-transport codes are export-controlled (i.e., are not open source and difficult to access) and notoriously difficult to install, use, and develop in.
 Because `MC/DC` is an open-source and easily installable Python package (with a `pip`-installable distribution), it is ideal for use in an academic environment for both research and education.

docs/source/conf.py

@@ -42,6 +42,7 @@
     "sphinx.ext.autosummary",
     "sphinx_toolbox.github",
     "sphinx_toolbox.sidebar_links",
+    "sphinx.ext.autosectionlabel",
 ]
 autosummary_generate = True
 

docs/source/contribution.rst

@@ -30,6 +30,64 @@ It's pretty easy to do this locally, just run,
 
 in the top level MC/DC directory and all necessary changes will be automatically made for you.
 
+---------
+Debugging
+---------
+
+MCDC includes options to debug the Numba JIT code.
+It does this by toggling Numba options using the numba.config submodule.
+This will result in less performant code and longer compile times but will allow for better error messages from Numba and other packages.
+`See Numba documentation of a list of all possible debug and compiler options. <https://numba.readthedocs.io/en/stable/reference/envvars.html#debugging>`_
+The most useful set of debug options for MC/DC can be enabled with
+
+.. code-block:: python3
+
+    python input.py --mode=numba_debug
+
+Which will toggle the following debug and compiler options in Numba:
+
+* ``DISABLE_JIT=False`` turns on the jitter
+* ``NUMBA_OPT=0`` Forces the compilers to form un-optimized code (other options for this are ``1``, ``2``, and ``3`` with ``3`` being the most optimized). This option might need to be changed if errors only result from more optimization.
+* ``DEBUG=False`` turns on all debugging options. This is still disabled in ``mcdc numba_debug`` as it will print ALOT of info on your terminal screen
+* ``NUMBA_FULL_TRACEBACKS=1`` allows errors from sub-packages to be printed (i.e. Numpy)
+* ``NUMBA_BOUNDSCHECK=1`` numba will check vectors for bounds errors. If this is disabled it bound errors will result in a ``seg_fault``. This in consort with the previous option allows for the exact location of a bound error to be printed from Numpy subroutines
+* ``NUMBA_DEBUG_NRT=1`` enables the `Numba run time (NRT) statistics counter <https://numba.readthedocs.io/en/stable/developer/numba-runtime.html>`_ This helps with debugging memory leaks.
+* ``NUMBA_DEBUG_TYPEINFER= 1`` print out debugging information about type inferences that numba might need to make if a function is ill-defined
+* ``NUMBA_ENABLE_PROFILING=1`` enables profiler use
+* ``NUMBA_DUMP_CFG=1`` prints out a control flow diagram
+
+If extra debug options or alteration to these options are required they can be toggled and passed under the ``mode==numba_debug`` option tree near the top of ``mcdc/main.py``.
+
+-------
+Caching
+-------
+
+MC/DC is a just-in-time (JIT) compiled code.
+This is sometimes disadvantageous, especially for users who might run many versions of the same simulation with slightly different parameters.
+As the JIT compilation scheme will only compile functions that are actually used in a given simulation, it is not a grantee that any one function will be compiled.
+
+Developers should be very cautious about using caching features.
+Numba has a few documented errors around caching.
+The most critical of which is that functions in other files that are called by cached functions will not force a recompile, even if there are changes in those sub-functions.
+In this case caching should be disabled.
+
+In MC/DC the outer most loop functions (in ``mcdc/loop.py``) are called to be cached.
+This is done with a option on the jit flag above the individual function declarations like
+
+.. code-block:: python3
+
+    nb.njit(cache=True)
+    def loop_fixed_source(mcdc):
+        # Loop over
+        ...
+
+To disable caching toggle these jit flags from ``True`` to ``False``.
+Alteratively a developer could delete the ``__pycache__`` directory or other cache directory which is system dependent (`see more about clearing the numba cache <https://numba.readthedocs.io/en/stable/developer/caching.html>`_)
+
+
+At some point MC/DC will enable `Numba's Ahead of Time compilation abilities <https://numba.readthedocs.io/en/stable/user/pycc.html>`_. But the core development team is holding off until scheduled `upgrades to AOT functionality in Numba are implemented <https://numba.readthedocs.io/en/stable/reference/deprecation.html#deprecation-numba-pycc>`_.
+However if absolutely required by users numba does allow for some `cache sharing <https://numba.readthedocs.io/en/stable/developer/caching.html>`_.
+
 -------
 Testing
 -------

docs/source/user.rst

@@ -203,6 +203,23 @@ Numba Mode
 
     python input.py --mode=numba
 
+When running in Numba mode a significant amount of time is taken compiling Python functions to performant binaries.
+Only functions used in a specific simulation will be compiled.
+These binaries will be cached meaning that in subsequent runs of the same simulation the compilation step can be avoided.
+The cache can be used as an effective ahead of time compilation scheme where binaries can be compiled once and shared between machines.
+For more information on caching see :ref:`Caching` and `Numba Caching <https://numba.readthedocs.io/en/stable/developer/caching.html>`_.
+
+MC/DC also has the ability to run Numba in a debugging mode.
+This will result in less performant code and longer compile times but will allow for better error messages from Numba and other packages.
+
+.. code-block:: python3
+
+    python input.py --mode=numba_debug
+
+
+For more information on the exact behavior of this option see :ref:`Debugging`
+
+
 Using MPI
 ^^^^^^^^^
 
@@ -216,6 +233,7 @@ Below, ``--mode`` can equal python or numba.
     srun python input.py --mode=<python/numba>
 
 
+
 Postprocessing Results
 ----------------------
 

mcdc/loop.py

@@ -23,8 +23,13 @@
 # Fixed-source loop
 # =========================================================================
 
+# about caching:
+#     it is enabled as a default at the jit call level
+#     to effectivly disable cache, delete the cache folder (often located in /MCDC/mcdc/__pycache__)
+#     see more about cacheing here https://numba.readthedocs.io/en/stable/developer/caching.html
 
-@njit
+
+@njit(cache=True)
 def loop_fixed_source(mcdc):
     # Loop over batches
     for idx_batch in range(mcdc["setting"]["N_batch"]):
@@ -89,7 +94,7 @@ def loop_fixed_source(mcdc):
 # =========================================================================
 
 
-@njit
+@njit(cache=True)
 def loop_eigenvalue(mcdc):
     # Loop over power iteration cycles
     for idx_cycle in range(mcdc["setting"]["N_cycle"]):
@@ -128,7 +133,7 @@ def loop_eigenvalue(mcdc):
 # =============================================================================
 
 
-@njit
+@njit(cache=True)
 def loop_source(seed, mcdc):
     # Progress bar indicator
     N_prog = 0
@@ -210,7 +215,7 @@ def loop_source(seed, mcdc):
 # =========================================================================
 
 
-@njit
+@njit(cache=True)
 def loop_particle(P, mcdc):
     # Particle tracker
     if mcdc["setting"]["track_particle"]:
@@ -303,7 +308,7 @@ def loop_particle(P, mcdc):
 # =============================================================================
 
 
-@njit
+@njit(cache=True)
 def loop_iqmc(mcdc):
     # function calls from specified solvers
     iqmc = mcdc["technique"]["iqmc"]
@@ -320,7 +325,7 @@ def loop_iqmc(mcdc):
             gmres(mcdc)
 
 
-@njit
+@njit(cache=True)
 def source_iteration(mcdc):
     simulation_end = False
     iqmc = mcdc["technique"]["iqmc"]
@@ -358,7 +363,7 @@ def source_iteration(mcdc):
         total_source_old = iqmc["total_source"].copy()
 
 
-@njit
+@njit(cache=True)
 def gmres(mcdc):
     """
     GMRES solver.
@@ -513,7 +518,7 @@ def gmres(mcdc):
             return
 
 
-@njit
+@njit(cache=True)
 def power_iteration(mcdc):
     simulation_end = False
     iqmc = mcdc["technique"]["iqmc"]
@@ -559,7 +564,7 @@ def power_iteration(mcdc):
                 print_iqmc_eigenvalue_exit_code(mcdc)
 
 
-@njit
+@njit(cache=True)
 def davidson(mcdc):
     """
     The generalized Davidson method is a Krylov subspace method for solving
@@ -669,7 +674,7 @@ def davidson(mcdc):
 # =============================================================================
 
 
-@njit
+@njit(cache=True)
 def loop_source_precursor(seed, mcdc):
     # TODO: censussed neutrons seeding is still not reproducible
 

mcdc/main.py

@@ -1,25 +1,74 @@
-import argparse, os
+import argparse, os, sys
 import numba as nb
 
 # Parse command-line arguments
 #   TODO: Will be inside run() once Python/Numba adapter is integrated
 parser = argparse.ArgumentParser(description="MC/DC: Monte Carlo Dynamic Code")
 parser.add_argument(
-    "--mode", type=str, help="Run mode", choices=["python", "numba"], default="python"
+    "--mode",
+    type=str,
+    help="Run mode",
+    choices=["python", "numba", "numba_debug"],
+    default="python",
 )
 parser.add_argument("--N_particle", type=int, help="Number of particles")
 parser.add_argument("--output", type=str, help="Output file name")
 parser.add_argument("--progress_bar", default=True, action="store_true")
 parser.add_argument("--no-progress_bar", dest="progress_bar", action="store_false")
 args, unargs = parser.parse_known_args()
 
+from mcdc.print_ import (
+    print_banner,
+    print_msg,
+    print_runtime,
+    print_header_eigenvalue,
+    print_warning,
+)
+
 # Set mode
 #   TODO: Will be inside run() once Python/Numba adapter is integrated
 mode = args.mode
 if mode == "python":
     nb.config.DISABLE_JIT = True
 elif mode == "numba":
     nb.config.DISABLE_JIT = False
+    nb.config.NUMBA_DEBUG_CACHE = 1
+elif mode == "numba_debug":
+    msg = "\n >> Entering numba debug mode\n >> will result in slower code and longer compile times\n >> to configure debug options see main.py"
+    print_warning(msg)
+
+    nb.config.DISABLE_JIT = False  # turns on the jitter
+    nb.config.DEBUG = False  # turns on debugging options
+    nb.config.NUMBA_FULL_TRACEBACKS = (
+        1  # enables errors from sub-packages to be printed
+    )
+    nb.config.NUMBA_BOUNDSCHECK = 1  # checks bounds errors of vectors
+    nb.config.NUMBA_COLOR_SCHEME = (
+        "dark_bg"  # prints error messages for dark background terminals
+    )
+    nb.config.NUMBA_DEBUG_NRT = 1  # Numba run time (NRT) statistics counter
+    nb.config.NUMBA_DEBUG_TYPEINFER = (
+        1  # print out debugging information about type inference.
+    )
+    nb.config.NUMBA_ENABLE_PROFILING = 1  # enables profiler use
+    nb.config.NUMBA_DUMP_CFG = 1  # prints out a control flow diagram
+    nb.config.NUMBA_OPT = 0  # forums un optimized code from compilers
+    nb.config.NUMBA_DEBUGINFO = 1  #
+    nb.config.NUMBA_EXTEND_VARIABLE_LIFETIMES = (
+        1  # allows for inspection of numba variables after end of compilation
+    )
+
+    # file="str.txt";file1="list.txt"
+    # out=sys.stdout
+    # sys.stdout=open('debug_numba_config.txt','w')
+    # help(nb.config)
+    # sys.stdout.close
+
+    # print_msg('>> Numba config exported to debug_numba_config.txt')
+
+# elif mode == "numba x86":
+#    nb.config.NUMBA_OPT = 3
+#    NUMBA_DISABLE_INTEL_SVML
 
 import h5py
 import numpy as np

mcdc/print_.py

@@ -2,7 +2,7 @@
 import numpy as np
 import sys
 from mpi4py import MPI
-
+from colorama import Fore, Back, Style
 
 master = MPI.COMM_WORLD.Get_rank() == 0
 
@@ -21,7 +21,8 @@ def print_error(msg):
 
 def print_warning(msg):
     if master:
-        print("Warning: %s\n" % msg)
+        print(Fore.RED + "Warning: %s\n" % msg)
+        print(Style.RESET_ALL)
         sys.stdout.flush()
 
 

mcdc/visualizer.py

@@ -1,3 +1,5 @@
+from mcdc.print_ import print_warning
+
 try:
     # launches visualization window
     # must be inside this loop so it doesn't launch when the visualizer is imported
@@ -7,16 +9,8 @@
     import distinctipy  # creates unlimited visually distinct colors for visualization
 
 except ImportError as e:
-    print("")
-    print(">>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>")
-    print("MC/DC visualization error:")
-    print("     Dependencies for visualization not installed")
-    print("     To install optional dependencies needed for visualization:")
-    print("         <pip install mcdc[viz]> ")
-    print("")
-    print(">>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>")
-    print("")
-
+    msg = "\n >> MC/DC visualization error: \n >> dependencies for visualization not installed \n >> install optional dependencies needed for visualization with \n >>     <pip install mcdc[viz]> (add para for mac)"
+    print_warning(msg)
 
 import tkinter as tk  # Tkinter is used to create the window for the time slider and color key
 import math
@@ -420,15 +414,8 @@ def visualize(start_time=0, end_time=0, tick_interval=1, material_colors={}):
         # must be inside this loop so it doesn't launch when the visualizer is imported
         import netgen.gui
     except ImportError as e:
-        print("")
-        print(">>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>")
-        print("MC/DC visualization error:")
-        print("     Dependencies for visualization not installed")
-        print("     To install optional dependencies needed for visualization:")
-        print("         <pip install mcdc[viz]> ")
-        print("")
-        print(">>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>")
-        print("")
+        msg = "\n >> MC/DC visualization error: \n >> dependencies for visualization not installed \n >> install optional dependencies needed for visualization with \n >>     <pip install mcdc[viz]> (add para for mac)"
+        print_warning(msg)
 
     color_key_dic = draw_Geometry(
         current_time=0,

pyproject.toml

@@ -18,7 +18,7 @@ authors = [
 ]
 description = "Monte Carlo / Dynamic Code, a pure python high performance Monte Carlo neutronics package"
 readme = "README.md"
-requires-python = ">=3.9, <=3.11.7"
+requires-python = ">=3.9, <=3.11.8"
 license = {file = "LICENSE"}
 keywords = ["Monte Carlo", "Nuclear", "GPU", "numba", "mpi4py", "neutron transport", "neutronics", "HPC"]
 classifiers = [