├── .github └── workflows │ └── python-package-build-and-publish.yml ├── .gitignore ├── LICENSE ├── MANIFEST.in ├── README.md ├── build └── lib │ └── numba_progress │ └── __init__.py ├── examples ├── __init__.py ├── clock_print.py ├── example_nested_loops.py ├── example_notebook.ipynb ├── example_parallel.py ├── example_sequential.py ├── example_signature.py └── sleep.py ├── numba_progress ├── __init__.py ├── _version.py ├── numba_atomic.py └── progress.py └── pyproject.toml /.github/workflows/python-package-build-and-publish.yml: -------------------------------------------------------------------------------- 1 | 2 | name: Python package build and publish 3 | 4 | on: workflow_dispatch 5 | 6 | jobs: 7 | deploy: 8 | runs-on: ubuntu-latest 9 | steps: 10 | - uses: actions/checkout@v2 11 | - name: Set up Python 12 | uses: actions/setup-python@v1 13 | with: 14 | python-version: 3.9 15 | - name: Install dependencies 16 | run: | 17 | python -m pip install --upgrade pip 18 | pip install twine flake8 build 19 | - name: Lint with flake8 for syntax errors 20 | run: | 21 | pip install flake8 22 | flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics 23 | flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics 24 | - name: Build Python source distribution and wheel 25 | run: | 26 | python -m build --sdist --wheel --outdir dist/ . 27 | - name: Publish wheels to PyPI 28 | env: 29 | TWINE_USERNAME: ${{ secrets.PYPI_USERNAME }} 30 | TWINE_PASSWORD: ${{ secrets.PYPI_PASSWORD }} 31 | run: | 32 | twine upload dist/*.tar.gz 33 | twine upload dist/*.whl 34 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | .idea/ 2 | __pycache__/ 3 | *.egg-info/ 4 | *.so 5 | build/ 6 | dist/ 7 | _version.py -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2021 Felix Igelbrink 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /MANIFEST.in: -------------------------------------------------------------------------------- 1 | include LICENSE -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Numba-progress 2 | 3 | A progress bar implementation for numba functions using tqdm. 4 | The module provides the class `ProgressBar` that works as a wrapper around the 5 | `tqdm.tqdm` progress bar. 6 | 7 | It works by spawning a separate thread that updates the `tqdm` progress bar 8 | based on an atomic counter which can be accessed and updated in a numba nopython function. 9 | 10 | The progress bar works with parallel as well as sequential numba functions. 11 | 12 | ## Installation 13 | 14 | ### Using pip 15 | ``` 16 | pip install numba-progress 17 | ``` 18 | 19 | ### From source 20 | ``` 21 | git clone https://github.com/mortacious/numba-progress.git 22 | cd numba-progress 23 | python setup.py install 24 | ``` 25 | 26 | ## Usage 27 | 28 | ```python 29 | from numba import njit 30 | from numba_progress import ProgressBar 31 | 32 | num_iterations = 100 33 | 34 | @njit(nogil=True) 35 | def numba_function(num_iterations, progress_proxy): 36 | for i in range(num_iterations): 37 | # 38 | progress_proxy.update(1) 39 | 40 | with ProgressBar(total=num_iterations) as progress: 41 | numba_function(num_iterations, progress) 42 | ``` 43 | 44 | The `ProgressBar` also works within parallel functions out of the box: 45 | 46 | ```python 47 | from numba import njit, prange 48 | from numba_progress import ProgressBar 49 | 50 | num_iterations = 100 51 | 52 | @njit(nogil=True, parallel=True) 53 | def numba_function(num_iterations, progress_proxy): 54 | for i in prange(num_iterations): 55 | # 56 | progress_proxy.update(1) 57 | 58 | with ProgressBar(total=num_iterations) as progress: 59 | numba_function(num_iterations, progress) 60 | ``` 61 | 62 | Refer to the `examples` folder for more usage examples. -------------------------------------------------------------------------------- /build/lib/numba_progress/__init__.py: -------------------------------------------------------------------------------- 1 | from .progress import ProgressBar 2 | from ._version import __version__ 3 | -------------------------------------------------------------------------------- /examples/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mortacious/numba-progress/57e34124ffbd4ff2bf0f09c8642f09e365972a0a/examples/__init__.py -------------------------------------------------------------------------------- /examples/clock_print.py: -------------------------------------------------------------------------------- 1 | from sleep import clock, usleep 2 | 3 | import numba as nb 4 | 5 | @nb.njit(nogil=True) 6 | def numba_clock(): 7 | c1 = clock() 8 | print("c1", c1) 9 | usleep(1000000) 10 | c2 = clock() 11 | print("c2", c2) 12 | print("time", c2-c1) 13 | 14 | if __name__ == "__main__": 15 | numba_clock() -------------------------------------------------------------------------------- /examples/example_nested_loops.py: -------------------------------------------------------------------------------- 1 | from sleep import usleep 2 | import numba as nb 3 | from numba_progress import ProgressBar 4 | 5 | @nb.njit(nogil=True) 6 | def numba_sleeper(num_iterations, sleep_us, progress): 7 | for i in range(num_iterations): 8 | progress[0].update() 9 | for j in range(num_iterations): 10 | usleep(sleep_us) 11 | progress[1].update(1) 12 | # reset the second progress bar to 0 13 | progress[1].set(0) 14 | 15 | 16 | if __name__ == "__main__": 17 | num_iterations = 30 18 | sleep_time_us = 25_000 19 | with ProgressBar(total=num_iterations, ncols=80) as numba_progress1, ProgressBar(total=num_iterations, ncols=80) as numba_progress2: 20 | # note: progressbar object must be passed as a tuple (a list will not work due to different treatment in numba) 21 | numba_sleeper(num_iterations, sleep_time_us, (numba_progress1, numba_progress2)) -------------------------------------------------------------------------------- /examples/example_notebook.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "code", 5 | "execution_count": 1, 6 | "metadata": { 7 | "pycharm": { 8 | "name": "#%%\n" 9 | } 10 | }, 11 | "outputs": [], 12 | "source": [ 13 | "from numba_progress import ProgressBar\n", 14 | "from sleep import usleep\n", 15 | "import numba as nb" 16 | ] 17 | }, 18 | { 19 | "cell_type": "code", 20 | "execution_count": 2, 21 | "metadata": {}, 22 | "outputs": [], 23 | "source": [ 24 | "@nb.njit(nogil=True, parallel=True)\n", 25 | "def numba_parallel_sleeper(num_iterations, sleep_us, progress_hook):\n", 26 | " for i in nb.prange(num_iterations):\n", 27 | " usleep(sleep_us)\n", 28 | " progress_hook.update(1)" 29 | ] 30 | }, 31 | { 32 | "cell_type": "code", 33 | "execution_count": 3, 34 | "metadata": { 35 | "pycharm": { 36 | "name": "#%%\n" 37 | } 38 | }, 39 | "outputs": [ 40 | { 41 | "data": { 42 | "application/vnd.jupyter.widget-view+json": { 43 | "model_id": "85383479b280434bb1fbec9c0394ba8b", 44 | "version_major": 2, 45 | "version_minor": 0 46 | }, 47 | "text/plain": [ 48 | " 0%| | 0/240 [00:00>> is_notebook() 21 | False 22 | """ 23 | # http://stackoverflow.com/questions/34091701/determine-if-were-in-an-ipython-notebook-session 24 | if "IPython" not in sys.modules: # IPython hasn't been imported 25 | return False 26 | from IPython import get_ipython 27 | 28 | # check for `kernel` attribute on the IPython instance 29 | return getattr(get_ipython(), "kernel", None) is not None 30 | 31 | 32 | class ProgressBar(object): 33 | """ 34 | Wraps the tqdm progress bar enabling it to be updated from within a numba nopython function. 35 | It works by spawning a separate thread that updates the tqdm progress bar based on an atomic counter which can be 36 | accessed within the numba function. The progress bar works with parallel as well as sequential numba functions. 37 | 38 | Note: As this Class contains python objects not useable or convertable into numba, it will be boxed as a 39 | proxy object, that only exposes the minimum subset of functionality to update the progress bar. Attempts 40 | to return or create a ProgressBar within a numba function will result in an error. 41 | 42 | Parameters 43 | ---------- 44 | file: `io.TextIOWrapper` or `io.StringIO`, optional 45 | Specifies where to output the progress messages 46 | (default: sys.stdout). Uses `file.write(str)` and `file.flush()` 47 | methods. For encoding, see `write_bytes`. 48 | update_interval: float, optional 49 | The interval in seconds used by the internal thread to check for updates [default: 0.1]. 50 | notebook: bool, optional 51 | If set, forces or forbits the use of the notebook progress bar. By default the best progress bar will be 52 | determined automatically. 53 | dynamic_ncols: bool, optional 54 | If true, the number of columns (the width of the progress bar) is constantly adjusted. This improves the 55 | output of the notebook progress bar a lot. 56 | kwargs: dict-like, optional 57 | Addtional parameters passed to the tqdm class. See https://github.com/tqdm/tqdm for a documentation of 58 | the available parameters. Noteable exceptions are the parameters: 59 | - file is redefined above (see above) 60 | - iterable is not available because it would not make sense here 61 | - dynamic_ncols is defined above 62 | """ 63 | def __init__(self, file=None, update_interval=0.1, notebook=None, dynamic_ncols=True, **kwargs): 64 | if file is None: 65 | file = sys.stdout 66 | self._last_value = 0 67 | 68 | if notebook is None: 69 | notebook = is_notebook() 70 | 71 | if notebook: 72 | self._tqdm = tqdm_notebook(iterable=None, dynamic_ncols=dynamic_ncols, file=file, **kwargs) 73 | else: 74 | self._tqdm = tqdm(iterable=None, dynamic_ncols=dynamic_ncols, file=file, **kwargs) 75 | 76 | self.hook = np.zeros(1, dtype=np.uint64) 77 | self._updater_thread = None 78 | self._exit_event = Event() 79 | self.update_interval = update_interval 80 | self._start() 81 | 82 | def _start(self): 83 | self._timer = Thread(target=self._update_function) 84 | self._timer.start() 85 | 86 | def close(self): 87 | self._exit_event.set() 88 | self._timer.join() 89 | self._update_tqdm() # update to set the progressbar to it's final value in case the thread missed a loop 90 | self._tqdm.refresh() 91 | self._tqdm.close() 92 | 93 | @property 94 | def n(self): 95 | return self.hook[0] 96 | 97 | def set(self, n=0): 98 | atomic_xchg(self.hook, 0, n) 99 | self._update_tqdm() 100 | 101 | def update(self, n=1): 102 | atomic_add(self.hook, 0, n) 103 | self._update_tqdm() 104 | 105 | def _update_tqdm(self): 106 | value = self.hook[0] 107 | #diff = value - self._last_value 108 | #self._last_value = value 109 | self._tqdm.n = value 110 | self._tqdm.refresh() 111 | #self._tqdm.update(diff) 112 | 113 | def _update_function(self): 114 | """Background thread for updating the progress bar. 115 | """ 116 | while not self._exit_event.is_set(): 117 | self._update_tqdm() 118 | self._exit_event.wait(self.update_interval) 119 | 120 | def __enter__(self): 121 | return self 122 | 123 | def __exit__(self, exc_type, exc_val, exc_tb): 124 | self.close() 125 | 126 | 127 | # Numba Native Implementation for the ProgressBar Class 128 | 129 | class ProgressBarTypeImpl(types.Type): 130 | def __init__(self): 131 | super().__init__(name='ProgressBar') 132 | 133 | 134 | # This is the numba type representation of the ProgressBar class to be used in signatures 135 | ProgressBarType = ProgressBarTypeImpl() 136 | 137 | 138 | @typeof_impl.register(ProgressBar) 139 | def typeof_index(val, c): 140 | return ProgressBarType 141 | 142 | 143 | as_numba_type.register(ProgressBar, ProgressBarType) 144 | 145 | 146 | @register_model(ProgressBarTypeImpl) 147 | class ProgressBarModel(models.StructModel): 148 | def __init__(self, dmm, fe_type): 149 | members = [ 150 | ('hook', types.Array(types.uint64, 1, 'C')), 151 | ] 152 | models.StructModel.__init__(self, dmm, fe_type, members) 153 | 154 | 155 | # make the hook attribute accessible 156 | make_attribute_wrapper(ProgressBarTypeImpl, 'hook', 'hook') 157 | 158 | 159 | 160 | @overload_attribute(ProgressBarTypeImpl, 'n') 161 | def get_value(progress_bar): 162 | def getter(progress_bar): 163 | return progress_bar.hook[0] 164 | return getter 165 | 166 | 167 | @unbox(ProgressBarTypeImpl) 168 | def unbox_progressbar(typ, obj, c): 169 | """ 170 | Convert a ProgressBar to it's native representation (proxy object) 171 | """ 172 | hook_obj = c.pyapi.object_getattr_string(obj, 'hook') 173 | progress_bar = cgutils.create_struct_proxy(typ)(c.context, c.builder) 174 | progress_bar.hook = unbox_array(types.Array(types.uint64, 1, 'C'), hook_obj, c).value 175 | c.pyapi.decref(hook_obj) 176 | is_error = cgutils.is_not_null(c.builder, c.pyapi.err_occurred()) 177 | return NativeValue(progress_bar._getvalue(), is_error=is_error) 178 | 179 | 180 | @box(ProgressBarTypeImpl) 181 | def box_progressbar(typ, val, c): 182 | raise TypeError("Native representation of ProgressBar cannot be converted back to a python object " 183 | "as it contains internal python state.") 184 | 185 | 186 | @overload_method(ProgressBarTypeImpl, "update", jit_options={"nogil": True}) 187 | def _ol_update(self, n=1): 188 | """ 189 | Numpy implementation of the update method. 190 | """ 191 | if isinstance(self, ProgressBarTypeImpl): 192 | def _update_impl(self, n=1): 193 | atomic_add(self.hook, 0, n) 194 | return _update_impl 195 | 196 | @overload_method(ProgressBarTypeImpl, "set", jit_options={"nogil": True}) 197 | def _ol_set(self, n=0): 198 | """ 199 | Numpy implementation of the update method. 200 | """ 201 | if isinstance(self, ProgressBarTypeImpl): 202 | def _set_impl(self, n=0): 203 | atomic_xchg(self.hook, 0, n) 204 | return _set_impl 205 | 206 | 207 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [build-system] 2 | requires = ["setuptools", "setuptools_scm>=6.2"] 3 | build-backend = "setuptools.build_meta" 4 | 5 | [project] 6 | name = "numba-progress" 7 | authors = [ 8 | {name = "Felix Igelbrink", email = "felix.igelbrink@dfki.de"}, 9 | ] 10 | description = 'A tqdm-compatible progress bar implementation for numba functions' 11 | readme = "README.md" 12 | requires-python = ">= 3.9" 13 | keywords = [] 14 | license = {text = "MIT"} 15 | classifiers = [ 16 | "Programming Language :: Python :: 3.8", 17 | "Programming Language :: Python :: 3.9", 18 | "Programming Language :: Python :: 3.10", 19 | "Programming Language :: Python :: 3.11", 20 | "Programming Language :: Python :: 3.12", 21 | "License :: OSI Approved :: MIT License", 22 | "Operating System :: POSIX :: Linux", 23 | "Intended Audience :: Developers", 24 | "Topic :: Scientific/Engineering", 25 | "Topic :: Software Development", 26 | ] 27 | dependencies = [ 28 | 'numba>=0.52', 29 | 'numpy', 30 | 'tqdm' 31 | ] 32 | dynamic = ["version"] 33 | 34 | [project.urls] 35 | Homepage = "https://github.com/mortacious/numba-progress" 36 | Issues = "https://github.com/mortacious/numba-progress/issues" 37 | 38 | [tool.setuptools] 39 | license-files = ["../LICENSE"] 40 | 41 | [tool.setuptools_scm] 42 | write_to = "numba_progress/_version.py" 43 | version_scheme = "only-version" 44 | local_scheme = "no-local-version" 45 | 46 | [tool.setuptools.packages.find] 47 | include = [ 48 | 'numba_progress' 49 | ] 50 | --------------------------------------------------------------------------------