Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fixed #1066 Permanently Remove Recompile #1067

Closed
wants to merge 1 commit into from
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 0 additions & 35 deletions stumpy/cache.py
Original file line number Diff line number Diff line change
Expand Up @@ -156,40 +156,6 @@ def _get_cache():
return [f.name for f in pathlib.Path(numba_cache_dir).glob("*nb*") if f.is_file()]


def _recompile():
"""
Recompile all njit functions

Parameters
----------
None

Returns
-------
None

Notes
-----
If the `numba` cache is enabled, this results in saving (and/or overwriting)
the cached numba functions to disk.
"""
for module_name, func_name in get_njit_funcs():
module = importlib.import_module(f".{module_name}", package="stumpy")
func = getattr(module, func_name)
try:
func.recompile()
except AttributeError as e:
if (
numba.config.DISABLE_JIT
and str(e) == "'function' object has no attribute 'recompile'"
):
pass
else: # pragma: no cover
raise

return


def _save():
"""
Save all njit functions
Expand All @@ -203,7 +169,6 @@ def _save():
None
"""
_enable()
_recompile()
Copy link
Collaborator

@NimaSarajpoor NimaSarajpoor Jan 27, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Removing _recompile here means that _save only does cache._enable(), which does nothing if a njit function is already compiled. The reason that tests/test_cache.py::test_cache_save_after_clear is still passing is because both ref_cache and comp_cache are now empty.

def test_cache_save_after_clear():
T = np.random.rand(10)
m = 3
stump(T, m)
cache.save()
ref_cache = cache._get_cache()
cache.clear()
# testing cache._clear()
assert len(cache._get_cache()) == 0
cache.save()
comp_cache = cache._get_cache()
# testing cache._save() after cache._clear()
assert sorted(ref_cache) == sorted(comp_cache)

To have a better test, we should move up the first call of cache.save(), and put it just before calling the function stumpy.stump. This way we know that we have some files in ref. Then, the test in this version should fail and expose the issue to us.

However, when I made such change, I noticed that the tests are NOT failing still. The returned cache files is still an empty list. There is a deeper issue here. According to Numba's document:

If not defined, Numba picks the cache directory in the following order:
In-tree cache. Put the cache next to the corresponding source file under a pycache directory following how .pyc files are stored.

So, I turned on NUMBA_CACHE_DIR and noticed that it writes to the local git directory "stumpy", in the path ./stumpy/__pycache__, however, cache._get_cache() reads the files from the following path, which is where stumpy is installed.

.../miniconda3/envs/py310/lib/python3.10/site-packages/stumpy/__pycache__

and that's why it returns 0 files.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So, I turned on NUMBA_CACHE_DIR and noticed that it writes to the local git directory "stumpy", in the path ./stumpy/pycache, however, cache._get_cache() reads the files from the following path, which is where stumpy is installed.

This is unexpected behavior. I will look into this

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@seanlaw

So, I turned on NUMBA_CACHE_DIR and noticed that ...

Correction: "I turned on NUMBA_DEBUG_CACHE"


return

Expand Down
11 changes: 7 additions & 4 deletions stumpy/fastmath.py
Original file line number Diff line number Diff line change
Expand Up @@ -54,12 +54,15 @@ def _set(module_name, func_name, flag):
module = importlib.import_module(f".{module_name}", package="stumpy")
func = getattr(module, func_name)
try:
Copy link
Collaborator

@NimaSarajpoor NimaSarajpoor Jan 27, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is there any particular advantage in using try-except block? If not, what do you think about using a simple if-block?

if numba.config.DISABLE_JIT:  # pragma no cover
        # update fastmath if in JIT mode
        py_func = func.py_func  # Copy raw Python function (independent of `njit`)
        njit_signature = func.targetoptions.copy()  # Copy the `njit` arguments
        njit_signature.pop("nopython", None)  # Pop redundant `nopython` declaration
        njit_signature["fastmath"] = flag  # Apply new `fastmath` flag
        func = njit(py_func, **njit_signature)  # Assign `njit` function with new target
 
# Monkey-patch [njit] func into global space
 setattr(module, func_name, func)  

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is there any particular advantage in using try-except block? If not, what do you think about using a simple if-block?

If I think about "normal" usage, njit is ALWAYS "on". So it feels "wrong" to check numba.config.DISABLE_JIT as njit should be assumed to always be there. Attempting to access the fastmath attributes via try at least assumes that njit is always "on", the default behavior. Note that I would only consider try/except when the function is called infrequently/rarely.

I will think about this more.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If I think about "normal" usage, njit is ALWAYS "on". So it feels "wrong" to check numba.config.DISABLE_JIT as njit should be assumed to always be there.

Makes sense!! Thanks for helping me understand your view! The if-block might seem cleaner but it creates branching here which is not the case when we think about the regular scenario.

func.targetoptions["fastmath"] = flag
func.recompile()
py_func = func.py_func # Copy raw Python function (independent of `njit`)
njit_signature = func.targetoptions.copy() # Copy the `njit` arguments
njit_signature.pop("nopython", None) # Pop redundant `nopython` declaration
njit_signature["fastmath"] = flag # Apply new `fastmath` flag
func = njit(py_func, **njit_signature) # Assign `njit` function with new target
setattr(module, func_name, func) # Monkey-patch `njit` into global space
except AttributeError as e:
if numba.config.DISABLE_JIT and (
str(e) == "'function' object has no attribute 'targetoptions'"
or str(e) == "'function' object has no attribute 'recompile'"
str(e) == "'function' object has no attribute 'py_func'"
):
pass
else: # pragma: no cover
Expand Down
4 changes: 1 addition & 3 deletions tests/test_precision.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@
from numba import cuda

import stumpy
from stumpy import cache, config, core, fastmath
from stumpy import config, core, fastmath

try:
from numba.errors import NumbaPerformanceWarning
Expand Down Expand Up @@ -156,7 +156,6 @@ def test_snippets():
fastmath._set(
"core", "_calculate_squared_distance", {"nsz", "arcp", "contract", "afn"}
)
cache._recompile()
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

To remove cache._recompile(), we need to find an alternative approach for updating all relevant callers of the function core._calculate_squared_distance whose fastmath flag was updated in previous line.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Right. I did not appreciate this and did not capture this in my original code. I will explore more


(
cmp_snippets,
Expand Down Expand Up @@ -187,7 +186,6 @@ def test_snippets():
if not numba.config.DISABLE_JIT: # pragma: no cover
# Revert fastmath flag back to their default values
fastmath._reset("core", "_calculate_squared_distance")
cache._recompile()


@pytest.mark.filterwarnings("ignore", category=NumbaPerformanceWarning)
Expand Down
Loading