Repository: numba/numba

Issue:
Issue with caching in Jupyterlab
<!--

Thanks for opening an issue! To help the Numba team handle your information
efficiently, please first ensure that there is no other issue present that
already describes the issue you have
(search at https://github.com/numba/numba/issues?&q=is%3Aissue).

-->

## Reporting a bug

<!--

Before submitting a bug report please ensure that you can check off these boxes:

-->

- [x] I have tried using the latest released version of Numba (most recent is
 visible in the change log (https://github.com/numba/numba/blob/master/CHANGE_LOG).
- [x] I have included a self contained code sample to reproduce the problem.
  i.e. it's possible to run as 'python bug.py'. --> Attaching an NB since that's where the issue is.

<!--

Please include details of the bug here, including, if applicable, what you
expected to happen!

-->

Having an issue with `cache=True` when used in a Jupyter Notebook. It works fine in a regular IPython shell or in python scripts.

Attaching a NB since that's where the issue is visible (it wouldn't let me upload a ipynb file, so I renamed it to a txt file; to use it you can download it and rename it as ipynb):

[numba-cache-reproducer.txt](https://github.com/numba/numba/files/7023363/numba-cache-reproducer.txt)

This is a screenshot:

<img width="1256" alt="Screen Shot 2021-08-20 at 1 54 19 PM" src="https://user-images.githubusercontent.com/14302941/130275522-778783fa-09dc-4ff5-923f-bb240d5c8c4d.png">

This is on numba=0.53.1. I noticed that 0.54 has been tagged, but I couldn't find it on conda-forge at the time of creating the issue.

Also attaching the output of `numba -s`:
[numba-s-out.txt](https://github.com/numba/numba/files/7023378/numba-s-out.txt)


I played around a little, and the following diff did _seem_ to fix it (i.e. it doesn't throw an error anymore):
[numba-nb-potential-fix-diff.txt](https://github.com/numba/numba/files/7023379/numba-nb-potential-fix-diff.txt)
but I'm not really sure if it's the right way to do it or if it works in general.


Discussion:

Resolving patch:
diff --git a/numba/core/caching.py b/numba/core/caching.py
--- a/numba/core/caching.py
+++ b/numba/core/caching.py
@@ -307,7 +307,10 @@ def get_disambiguator(self):
 
     @classmethod
     def from_function(cls, py_func, py_file):
-        if not py_file.startswith("<ipython-"):
+        if not (
+            py_file.startswith("<ipython-")
+            or os.path.basename(os.path.dirname(py_file)).startswith("ipykernel_")
+        ):
             return
         self = cls(py_func, py_file)
         try:


Test patch:
diff --git a/numba/tests/test_dispatcher.py b/numba/tests/test_dispatcher.py
--- a/numba/tests/test_dispatcher.py
+++ b/numba/tests/test_dispatcher.py
@@ -46,6 +46,11 @@
 except ImportError:
     pygments = None
 
+try:
+    import ipykernel
+except ImportError:
+    ipykernel = None
+
 _is_armv7l = platform.machine() == 'armv7l'
 
 
@@ -1468,6 +1473,63 @@ def execute_with_input():
         err = execute_with_input()
         self.assertIn("cache hits = 1", err.strip())
 
+    @unittest.skipIf((ipykernel is None) or (ipykernel.version_info[0] < 6),
+                     "requires ipykernel >= 6")
+    def test_ipykernel(self):
+        # Test caching in an IPython session using ipykernel
+
+        base_cmd = [sys.executable, '-m', 'IPython']
+        base_cmd += ['--quiet', '--quick', '--no-banner', '--colors=NoColor']
+        try:
+            ver = subprocess.check_output(base_cmd + ['--version'])
+        except subprocess.CalledProcessError as e:
+            self.skipTest("ipython not available: return code %d"
+                          % e.returncode)
+        ver = ver.strip().decode()
+        # Create test input
+        from ipykernel import compiler
+        inputfn = compiler.get_tmp_directory()
+        with open(inputfn, "w") as f:
+            f.write(r"""
+                import os
+                import sys
+
+                from numba import jit
+
+                # IPython 5 does not support multiline input if stdin isn't
+                # a tty (https://github.com/ipython/ipython/issues/9752)
+                f = jit(cache=True)(lambda: 42)
+
+                res = f()
+                # IPython writes on stdout, so use stderr instead
+                sys.stderr.write(u"cache hits = %d\n" % f.stats.cache_hits[()])
+
+                # IPython hijacks sys.exit(), bypass it
+                sys.stdout.flush()
+                sys.stderr.flush()
+                os._exit(res)
+                """)
+
+        def execute_with_input():
+            # Feed the test input as stdin, to execute it in REPL context
+            with open(inputfn, "rb") as stdin:
+                p = subprocess.Popen(base_cmd, stdin=stdin,
+                                     stdout=subprocess.PIPE,
+                                     stderr=subprocess.PIPE,
+                                     universal_newlines=True)
+                out, err = p.communicate()
+                if p.returncode != 42:
+                    self.fail("unexpected return code %d\n"
+                              "-- stdout:\n%s\n"
+                              "-- stderr:\n%s\n"
+                              % (p.returncode, out, err))
+                return err
+
+        execute_with_input()
+        # Run a second time and check caching
+        err = execute_with_input()
+        self.assertIn("cache hits = 1", err.strip())
+
 
 @skip_parfors_unsupported
 class TestSequentialParForsCache(BaseCacheUsecasesTest):