test_caching.py

import sys
import os.path
import re
import subprocess

import numpy as np

from numba.tests.support import capture_cache_log
from numba.tests.test_caching import BaseCacheTest
from numba.core import config
import unittest


class UfuncCacheTest(BaseCacheTest):
    """
    Since the cache stats is not exposed by ufunc, we test by looking at the
    cache debug log.
    """
    _numba_parallel_test_ = False

    here = os.path.dirname(__file__)
    usecases_file = os.path.join(here, "cache_usecases.py")
    modname = "ufunc_caching_test_fodder"

    regex_data_saved = re.compile(r'\[cache\] data saved to')
    regex_index_saved = re.compile(r'\[cache\] index saved to')

    regex_data_loaded = re.compile(r'\[cache\] data loaded from')
    regex_index_loaded = re.compile(r'\[cache\] index loaded from')

    def check_cache_saved(self, cachelog, count):
        """
        Check number of cache-save were issued
        """
        data_saved = self.regex_data_saved.findall(cachelog)
        index_saved = self.regex_index_saved.findall(cachelog)
        self.assertEqual(len(data_saved), count)
        self.assertEqual(len(index_saved), count)

    def check_cache_loaded(self, cachelog, count):
        """
        Check number of cache-load were issued
        """
        data_loaded = self.regex_data_loaded.findall(cachelog)
        index_loaded = self.regex_index_loaded.findall(cachelog)
        self.assertEqual(len(data_loaded), count)
        self.assertEqual(len(index_loaded), count)

    def check_ufunc_cache(self, usecase_name, n_overloads, **kwargs):
        """
        Check number of cache load/save.
        There should be one per overloaded version.
        """
        mod = self.import_module()
        usecase = getattr(mod, usecase_name)
        # New cache entry saved
        with capture_cache_log() as out:
            new_ufunc = usecase(**kwargs)
        cachelog = out.getvalue()
        self.check_cache_saved(cachelog, count=n_overloads)

        # Use cached version
        with capture_cache_log() as out:
            cached_ufunc = usecase(**kwargs)
        cachelog = out.getvalue()
        self.check_cache_loaded(cachelog, count=n_overloads)

        return new_ufunc, cached_ufunc


class TestUfuncCacheTest(UfuncCacheTest):

    def test_direct_ufunc_cache(self, **kwargs):
        new_ufunc, cached_ufunc = self.check_ufunc_cache(
            "direct_ufunc_cache_usecase", n_overloads=2, **kwargs)
        # Test the cached and original versions
        inp = np.random.random(10).astype(np.float64)
        np.testing.assert_equal(new_ufunc(inp), cached_ufunc(inp))
        inp = np.arange(10, dtype=np.intp)
        np.testing.assert_equal(new_ufunc(inp), cached_ufunc(inp))

    def test_direct_ufunc_cache_objmode(self):
        self.test_direct_ufunc_cache(forceobj=True)

    def test_direct_ufunc_cache_parallel(self):
        self.test_direct_ufunc_cache(target='parallel')

    def test_indirect_ufunc_cache(self, **kwargs):
        new_ufunc, cached_ufunc = self.check_ufunc_cache(
            "indirect_ufunc_cache_usecase", n_overloads=3, **kwargs)
        # Test the cached and original versions
        inp = np.random.random(10).astype(np.float64)
        np.testing.assert_equal(new_ufunc(inp), cached_ufunc(inp))
        inp = np.arange(10, dtype=np.intp)
        np.testing.assert_equal(new_ufunc(inp), cached_ufunc(inp))

    def test_indirect_ufunc_cache_parallel(self):
        self.test_indirect_ufunc_cache(target='parallel')


class TestDUfuncCacheTest(UfuncCacheTest):
    # Note: DUFunc doesn't support parallel target yet

    def check_dufunc_usecase(self, usecase_name):
        mod = self.import_module()
        usecase = getattr(mod, usecase_name)
        # Create dufunc
        with capture_cache_log() as out:
            ufunc = usecase()
        self.check_cache_saved(out.getvalue(), count=0)
        # Compile & cache
        with capture_cache_log() as out:
            ufunc(np.arange(10))
        self.check_cache_saved(out.getvalue(), count=1)
        self.check_cache_loaded(out.getvalue(), count=0)
        # Use cached
        with capture_cache_log() as out:
            ufunc = usecase()
            ufunc(np.arange(10))
        self.check_cache_loaded(out.getvalue(), count=1)

    def test_direct_dufunc_cache(self):
        # We don't test for objmode because DUfunc don't support it.
        self.check_dufunc_usecase('direct_dufunc_cache_usecase')

    def test_indirect_dufunc_cache(self):
        self.check_dufunc_usecase('indirect_dufunc_cache_usecase')


def _fix_raw_path(rstr):
    if config.IS_WIN32:
        rstr = rstr.replace(r'/', r'\\\\')
    return rstr


class TestGUfuncCacheTest(UfuncCacheTest):

    def test_filename_prefix(self):
        mod = self.import_module()
        usecase = getattr(mod, "direct_gufunc_cache_usecase")
        with capture_cache_log() as out:
            usecase()
        cachelog = out.getvalue()
        # find number filename with "guf-" prefix
        fmt1 = _fix_raw_path(r'/__pycache__/guf-{}')
        prefixed = re.findall(fmt1.format(self.modname), cachelog)
        fmt2 = _fix_raw_path(r'/__pycache__/{}')
        normal = re.findall(fmt2.format(self.modname), cachelog)
        # expecting 2 overloads
        self.assertGreater(len(normal), 2)
        # expecting equal number of wrappers and overloads cache entries
        self.assertEqual(len(normal), len(prefixed))

    def test_direct_gufunc_cache(self, **kwargs):
        # 2 cache entry for the 2 overloads
        # and 2 cache entry for the gufunc wrapper
        new_ufunc, cached_ufunc = self.check_ufunc_cache(
            "direct_gufunc_cache_usecase", n_overloads=2 + 2, **kwargs)
        # Test the cached and original versions
        inp = np.random.random(10).astype(np.float64)
        np.testing.assert_equal(new_ufunc(inp), cached_ufunc(inp))
        inp = np.arange(10, dtype=np.intp)
        np.testing.assert_equal(new_ufunc(inp), cached_ufunc(inp))

    def test_direct_gufunc_cache_objmode(self):
        self.test_direct_gufunc_cache(forceobj=True)

    def test_direct_gufunc_cache_parallel(self):
        self.test_direct_gufunc_cache(target='parallel')

    def test_indirect_gufunc_cache(self, **kwargs):
        # 3 cache entry for the 3 overloads
        # and no cache entry for the gufunc wrapper
        new_ufunc, cached_ufunc = self.check_ufunc_cache(
            "indirect_gufunc_cache_usecase", n_overloads=3, **kwargs)
        # Test the cached and original versions
        inp = np.random.random(10).astype(np.float64)
        np.testing.assert_equal(new_ufunc(inp), cached_ufunc(inp))
        inp = np.arange(10, dtype=np.intp)
        np.testing.assert_equal(new_ufunc(inp), cached_ufunc(inp))

    def test_indirect_gufunc_cache_parallel(self, **kwargs):
        self.test_indirect_gufunc_cache(target='parallel')


class TestCacheSpecificIssue(UfuncCacheTest):

    def run_in_separate_process(self, runcode):
        # Based on the same name util function in test_dispatcher but modified
        # to allow user to define what to run.
        code = """if 1:
            import sys

            sys.path.insert(0, %(tempdir)r)
            mod = __import__(%(modname)r)
            mod.%(runcode)s
            """ % dict(tempdir=self.tempdir, modname=self.modname,
                       runcode=runcode)

        popen = subprocess.Popen([sys.executable, "-c", code],
                                 stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        out, err = popen.communicate()
        if popen.returncode != 0:
            raise AssertionError("process failed with code %s: stderr follows"
                                 "\n%s\n" % (popen.returncode, err.decode()))

    #
    # The following test issue #2198 that loading cached (g)ufunc first
    # bypasses some target context initialization.
    #

    def test_first_load_cached_ufunc(self):
        # ensure function is cached
        self.run_in_separate_process('direct_ufunc_cache_usecase()')
        # use the cached function
        # this will fail if the target context is not init'ed
        self.run_in_separate_process('direct_ufunc_cache_usecase()')

    def test_first_load_cached_gufunc(self):
        # ensure function is cached
        self.run_in_separate_process('direct_gufunc_cache_usecase()')
        # use the cached function
        # this will fail out if the target context is not init'ed
        self.run_in_separate_process('direct_gufunc_cache_usecase()')


if __name__ == '__main__':
    unittest.main()