Repository: numba/numba

Issue:
Numba 0.54.0: unexpected assignment behavior within `objmode` context.
<!--

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:

-->

- [Y] 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).
- [Y] I have included a self contained code sample to reproduce the problem.

<!--

Please include details of the bug here, including, if applicable, what you
expected to happen!

-->

```
import numba
from numba import objmode


class Test(object):
    a = None

@numba.jit(nopython=True)
def test():
    with objmode():
        Test.a = 100
        print(Test.a)
    
    with objmode():
        print(Test.a)


test()
```

The above code runs as expected if Numba <= 0.53.1. The expected outputs are:
```
100
100
```
However, after update Numba to 0.54.0, the outputs are:
```
100
None
```
, which is not consistent with pure python.




Discussion:
Thanks for the report, I can reproduce.

`git bisect` suggests:

```
The first bad commit could be any of:
```
b9a7fc148ad298e60b0b1197ef75eb04c4109b2c
e13d80c9af2238ae2d510b5669272debce7c3ef3
033452bc2b16e06121a3fa78dda592eb17fba32d
58c4e3282607199ad3e260919c3a4849b8088153
34db404262b50c1b90ef12108da4d4b2fa092900
```
We cannot bisect more!
```
Which are all #6977.

CC @sklam 
I found out that `cloudpickle` is serializing `Test` when it is not importable. A class (e.g. `Test`) is usually importable when it is defined as a global of a module. However, the failing case is when the class is defined in the `__main__` module. The cloudpickle code is rejecting to import from `__main__` even though it appears to be possible.

Since `Test` is serialized allow with the code in `objmode()`, it is creating a fresh `Test` class everytime.
The problem is pretty much down to this if-statement in cloudpickle: https://github.com/cloudpipe/cloudpickle/blob/343da119685f622da2d1658ef7b3e2516a01817f/cloudpickle/cloudpickle.py#L300-L301
Why does cloudpickle reject handling `__main__`? Commenting it out appears to work for this issue.
> Why does cloudpickle reject handling __main__? Commenting it out appears to work for this issue.

answer: https://github.com/cloudpipe/cloudpickle/issues/131

Resolving patch:
diff --git a/numba/cloudpickle/cloudpickle.py b/numba/cloudpickle/cloudpickle.py
--- a/numba/cloudpickle/cloudpickle.py
+++ b/numba/cloudpickle/cloudpickle.py
@@ -93,6 +93,7 @@ def g():
 _DYNAMIC_CLASS_TRACKER_BY_CLASS = weakref.WeakKeyDictionary()
 _DYNAMIC_CLASS_TRACKER_BY_ID = weakref.WeakValueDictionary()
 _DYNAMIC_CLASS_TRACKER_LOCK = threading.Lock()
+_DYNAMIC_CLASS_TRACKER_REUSING = weakref.WeakSet()
 
 PYPY = platform.python_implementation() == "PyPy"
 
@@ -117,9 +118,14 @@ def _get_or_create_tracker_id(class_def):
 def _lookup_class_or_track(class_tracker_id, class_def):
     if class_tracker_id is not None:
         with _DYNAMIC_CLASS_TRACKER_LOCK:
+            orig_class_def = class_def
             class_def = _DYNAMIC_CLASS_TRACKER_BY_ID.setdefault(
                 class_tracker_id, class_def)
             _DYNAMIC_CLASS_TRACKER_BY_CLASS[class_def] = class_tracker_id
+            # Check if we are reusing a previous class_def
+            if orig_class_def is not class_def:
+                # Remember the class_def is being reused
+                _DYNAMIC_CLASS_TRACKER_REUSING.add(class_def)
     return class_def
 
 
diff --git a/numba/cloudpickle/cloudpickle_fast.py b/numba/cloudpickle/cloudpickle_fast.py
--- a/numba/cloudpickle/cloudpickle_fast.py
+++ b/numba/cloudpickle/cloudpickle_fast.py
@@ -36,6 +36,7 @@
     parametrized_type_hint_getinitargs, _create_parametrized_type_hint,
     builtin_code_type,
     _make_dict_keys, _make_dict_values, _make_dict_items,
+    _DYNAMIC_CLASS_TRACKER_REUSING,
 )
 
 
@@ -460,6 +461,10 @@ def _function_setstate(obj, state):
 
 
 def _class_setstate(obj, state):
+    # Check if class is being reused and needs bypass setstate logic.
+    if obj in _DYNAMIC_CLASS_TRACKER_REUSING:
+        return obj
+
     state, slotstate = state
     registry = None
     for attrname, attr in state.items():


Test patch:
diff --git a/numba/tests/cloudpickle_main_class.py b/numba/tests/cloudpickle_main_class.py
new file mode 100644
--- /dev/null
+++ b/numba/tests/cloudpickle_main_class.py
@@ -0,0 +1,6 @@
+# Expected to run this module as __main__
+
+
+# Cloudpickle will think this is a dynamic class when this module is __main__
+class Klass:
+    classvar = None
diff --git a/numba/tests/test_serialize.py b/numba/tests/test_serialize.py
--- a/numba/tests/test_serialize.py
+++ b/numba/tests/test_serialize.py
@@ -1,14 +1,18 @@
 import contextlib
 import gc
 import pickle
+import runpy
 import subprocess
 import sys
+import unittest
+from multiprocessing import get_context
 
+import numba
 from numba.core.errors import TypingError
-from numba.tests.support import TestCase, tag
-from .serialize_usecases import *
-import unittest
+from numba.tests.support import TestCase
 from numba.core.target_extension import resolve_dispatcher_from_str
+from numba.cloudpickle import dumps, loads
+from .serialize_usecases import *
 
 
 class TestDispatcherPickling(TestCase):
@@ -205,5 +209,117 @@ def test_numba_unpickle(self):
         self.assertIs(got1, got2)
 
 
+
+class TestCloudPickleIssues(TestCase):
+    """This test case includes issues specific to the cloudpickle implementation.
+    """
+    _numba_parallel_test_ = False
+
+    def test_dynamic_class_reset_on_unpickle(self):
+        # a dynamic class
+        class Klass:
+            classvar = None
+
+        def mutator():
+            Klass.classvar = 100
+
+        def check():
+            self.assertEqual(Klass.classvar, 100)
+
+        saved = dumps(Klass)
+        mutator()
+        check()
+        loads(saved)
+        # Without the patch, each `loads(saved)` will reset `Klass.classvar`
+        check()
+        loads(saved)
+        check()
+
+    @unittest.skipIf(__name__ == "__main__",
+                     "Test cannot run as when module is __main__")
+    def test_main_class_reset_on_unpickle(self):
+        mp = get_context('spawn')
+        proc = mp.Process(target=check_main_class_reset_on_unpickle)
+        proc.start()
+        proc.join(timeout=10)
+        self.assertEqual(proc.exitcode, 0)
+
+    def test_dynamic_class_reset_on_unpickle_new_proc(self):
+        # a dynamic class
+        class Klass:
+            classvar = None
+
+        # serialize Klass in this process
+        saved = dumps(Klass)
+
+        # Check the reset problem in a new process
+        mp = get_context('spawn')
+        proc = mp.Process(target=check_unpickle_dyn_class_new_proc, args=(saved,))
+        proc.start()
+        proc.join(timeout=10)
+        self.assertEqual(proc.exitcode, 0)
+
+    def test_dynamic_class_issue_7356(self):
+        cfunc = numba.njit(issue_7356)
+        self.assertEqual(cfunc(), (100, 100))
+
+
+class DynClass(object):
+    # For testing issue #7356
+    a = None
+
+
+def issue_7356():
+    with numba.objmode(before="intp"):
+        DynClass.a = 100
+        before = DynClass.a
+    with numba.objmode(after="intp"):
+        after = DynClass.a
+    return before, after
+
+
+def check_main_class_reset_on_unpickle():
+    # Load module and get its global dictionary
+    glbs = runpy.run_module(
+        "numba.tests.cloudpickle_main_class",
+        run_name="__main__",
+    )
+    # Get the Klass and check it is from __main__
+    Klass = glbs['Klass']
+    assert Klass.__module__ == "__main__"
+    assert Klass.classvar != 100
+    saved = dumps(Klass)
+    # mutate
+    Klass.classvar = 100
+    # check
+    _check_dyn_class(Klass, saved)
+
+
+def check_unpickle_dyn_class_new_proc(saved):
+    Klass = loads(saved)
+    assert Klass.classvar != 100
+    # mutate
+    Klass.classvar = 100
+    # check
+    _check_dyn_class(Klass, saved)
+
+
+def _check_dyn_class(Klass, saved):
+    def check():
+        if Klass.classvar != 100:
+            raise AssertionError("Check failed. Klass reset.")
+
+    check()
+    loaded = loads(saved)
+    if loaded is not Klass:
+        raise AssertionError("Expected reuse")
+    # Without the patch, each `loads(saved)` will reset `Klass.classvar`
+    check()
+    loaded = loads(saved)
+    if loaded is not Klass:
+        raise AssertionError("Expected reuse")
+    check()
+
+
 if __name__ == '__main__':
     unittest.main()