Skip to content

Commit 68e0ef5

Browse files
author
Felix Igelbrink
committed
initial commit
1 parent 52ae95b commit 68e0ef5

13 files changed

+474
-27
lines changed

LICENSE

+17-25
Original file line numberDiff line numberDiff line change
@@ -1,29 +1,21 @@
1-
BSD 3-Clause License
1+
MIT License
22

3-
Copyright (c) 2021, Felix Igelbrink
4-
All rights reserved.
3+
Copyright (c) 2021 Felix Igelbrink
54

6-
Redistribution and use in source and binary forms, with or without
7-
modification, are permitted provided that the following conditions are met:
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:
811

9-
1. Redistributions of source code must retain the above copyright notice, this
10-
list of conditions and the following disclaimer.
12+
The above copyright notice and this permission notice shall be included in all
13+
copies or substantial portions of the Software.
1114

12-
2. Redistributions in binary form must reproduce the above copyright notice,
13-
this list of conditions and the following disclaimer in the documentation
14-
and/or other materials provided with the distribution.
15-
16-
3. Neither the name of the copyright holder nor the names of its
17-
contributors may be used to endorse or promote products derived from
18-
this software without specific prior written permission.
19-
20-
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
21-
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
22-
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
23-
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
24-
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
25-
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
26-
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
27-
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
28-
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
29-
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
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.

MANIFEST.in

+1
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
include LICENSE

README.md

+62-2
Original file line numberDiff line numberDiff line change
@@ -1,2 +1,62 @@
1-
# numba_progress
2-
Using progress bars in numba
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 within 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+
#<DO CUSTOM WORK HERE>
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+
#<DO CUSTOM WORK HERE>
56+
progress_proxy.update(1)
57+
58+
with ProgressBar(total=num_iterations) as progress:
59+
numba_function(num_iterations, progress)
60+
```
61+
62+
See also the `examples` folder for more usage examples.

examples/__init__.py

Whitespace-only changes.

examples/example_parallel.py

+18
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
from sleep import usleep
2+
import numba as nb
3+
from numba_progress import ProgressBar
4+
5+
6+
@nb.njit(nogil=True, parallel=True)
7+
def numba_parallel_sleeper(num_iterations, sleep_us, progress_hook):
8+
for i in nb.prange(num_iterations):
9+
usleep(sleep_us)
10+
progress_hook.update(1)
11+
12+
13+
if __name__ == "__main__":
14+
num_iterations = 30*8
15+
sleep_time_us = 250_000
16+
with ProgressBar(total=num_iterations, ncols=80) as numba_progress:
17+
numba_parallel_sleeper(num_iterations, sleep_time_us, numba_progress)
18+

examples/example_sequential.py

+17
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
from sleep import usleep
2+
import numba as nb
3+
from numba_progress import ProgressBar
4+
5+
6+
@nb.njit(nogil=True)
7+
def numba_sleeper(num_iterations, sleep_us, progress_hook):
8+
for i in range(num_iterations):
9+
usleep(sleep_us)
10+
progress_hook.update(1)
11+
12+
13+
if __name__ == "__main__":
14+
num_iterations = 30
15+
sleep_time_us = 250_000
16+
with ProgressBar(total=num_iterations, ncols=80) as numba_progress:
17+
numba_sleeper(num_iterations, sleep_time_us, numba_progress)

examples/sleep.py

+15
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
import numba as nb
2+
3+
4+
@nb.generated_jit(nogil=True, nopython=True)
5+
def usleep(usec):
6+
# c usleep function in numba
7+
import ctypes
8+
libc = ctypes.CDLL('libc.so.6')
9+
libc.usleep.argtypes = (ctypes.c_uint,)
10+
func_usleep = libc.usleep
11+
12+
def usleep_impl(usec):
13+
func_usleep(usec)
14+
15+
return usleep_impl

numba_progress/__init__.py

+2
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
from .progress import ProgressBar
2+
from ._version import __version__

numba_progress/_version.py

+1
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
__version__ = "0.0.1"

numba_progress/numba_atomic.py

+150
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,150 @@
1+
# source from https://github.com/KatanaGraph/katana/blob/master/python/katana/numba_support/numpy_atomic.py
2+
3+
# The 3-Clause BSD License
4+
#
5+
# Copyright 2018 The University of Texas at Austin
6+
#
7+
# Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
8+
#
9+
# 1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
10+
#
11+
# 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.
12+
#
13+
# 3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission.
14+
#
15+
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
16+
17+
18+
from numba import types
19+
from numba.core import cgutils
20+
from numba.core.typing.arraydecl import get_array_index_type
21+
from numba.extending import lower_builtin, type_callable
22+
from numba.np.arrayobj import basic_indexing, make_array, normalize_indices
23+
24+
__all__ = ["atomic_add", "atomic_sub", "atomic_max", "atomic_min"]
25+
26+
27+
def atomic_rmw(context, builder, op, arrayty, val, ptr):
28+
assert arrayty.aligned # We probably have to have aligned arrays.
29+
dataval = context.get_value_as_data(builder, arrayty.dtype, val)
30+
return builder.atomic_rmw(op, ptr, dataval, "monotonic")
31+
32+
33+
def declare_atomic_array_op(iop, uop, fop):
34+
def decorator(func):
35+
@type_callable(func)
36+
def func_type(context):
37+
def typer(ary, idx, val):
38+
out = get_array_index_type(ary, idx)
39+
if out is not None:
40+
res = out.result
41+
if context.can_convert(val, res):
42+
return res
43+
return None
44+
45+
return typer
46+
47+
_ = func_type
48+
49+
@lower_builtin(func, types.Buffer, types.Any, types.Any)
50+
def func_impl(context, builder, sig, args):
51+
"""
52+
array[a] = scalar_or_array
53+
array[a,..,b] = scalar_or_array
54+
"""
55+
aryty, idxty, valty = sig.args
56+
ary, idx, val = args
57+
58+
if isinstance(idxty, types.BaseTuple):
59+
index_types = idxty.types
60+
indices = cgutils.unpack_tuple(builder, idx, count=len(idxty))
61+
else:
62+
index_types = (idxty,)
63+
indices = (idx,)
64+
65+
ary = make_array(aryty)(context, builder, ary)
66+
67+
# First try basic indexing to see if a single array location is denoted.
68+
index_types, indices = normalize_indices(context, builder, index_types, indices)
69+
dataptr, shapes, _strides = basic_indexing(
70+
context, builder, aryty, ary, index_types, indices, boundscheck=context.enable_boundscheck,
71+
)
72+
if shapes:
73+
raise NotImplementedError("Complex shapes are not supported")
74+
75+
# Store source value the given location
76+
val = context.cast(builder, val, valty, aryty.dtype)
77+
op = None
78+
if isinstance(aryty.dtype, types.Integer) and aryty.dtype.signed:
79+
op = iop
80+
elif isinstance(aryty.dtype, types.Integer) and not aryty.dtype.signed:
81+
op = uop
82+
elif isinstance(aryty.dtype, types.Float):
83+
op = fop
84+
if op is None:
85+
raise TypeError("Atomic operation not supported on " + str(aryty))
86+
return atomic_rmw(context, builder, op, aryty, val, dataptr)
87+
88+
_ = func_impl
89+
90+
return func
91+
92+
return decorator
93+
94+
95+
@declare_atomic_array_op("add", "add", "fadd")
96+
def atomic_add(ary, i, v):
97+
"""
98+
Atomically, perform `ary[i] += v` and return the previous value of `ary[i]`.
99+
100+
i must be a simple index for a single element of ary. Broadcasting and vector operations are not supported.
101+
102+
This should be used from numba compiled code.
103+
"""
104+
orig = ary[i]
105+
ary[i] += v
106+
return orig
107+
108+
109+
@declare_atomic_array_op("sub", "sub", "fsub")
110+
def atomic_sub(ary, i, v):
111+
"""
112+
Atomically, perform `ary[i] -= v` and return the previous value of `ary[i]`.
113+
114+
i must be a simple index for a single element of ary. Broadcasting and vector operations are not supported.
115+
116+
This should be used from numba compiled code.
117+
"""
118+
orig = ary[i]
119+
ary[i] -= v
120+
return orig
121+
122+
123+
@declare_atomic_array_op("max", "umax", None)
124+
def atomic_max(ary, i, v):
125+
"""
126+
Atomically, perform `ary[i] = max(ary[i], v)` and return the previous value of `ary[i]`.
127+
This operation does not support floating-point values.
128+
129+
i must be a simple index for a single element of ary. Broadcasting and vector operations are not supported.
130+
131+
This should be used from numba compiled code.
132+
"""
133+
orig = ary[i]
134+
ary[i] = max(ary[i], v)
135+
return orig
136+
137+
138+
@declare_atomic_array_op("min", "umin", None)
139+
def atomic_min(ary, i, v):
140+
"""
141+
Atomically, perform `ary[i] = min(ary[i], v)` and return the previous value of `ary[i]`.
142+
This operation does not support floating-point values.
143+
144+
i must be a simple index for a single element of ary. Broadcasting and vector operations are not supported.
145+
146+
This should be used from numba compiled code.
147+
"""
148+
orig = ary[i]
149+
ary[i] = min(ary[i], v)
150+
return orig

0 commit comments

Comments
 (0)