Repository: CoffeaTeam/coffea

Issue:
not able to load LumiMask
**Describe the bug**
Cannot load LumiMask from a textfile.

**To Reproduce**
Steps to reproduce the behavior:
```
from coffea.lumi_tools import LumiMask
# load lumimask from e.g. this file https://github.com/cmantill/boostedhiggs/blob/main/boostedhiggs/data/Cert_294927-306462_13TeV_UL2017_Collisions17_GoldenJSON.txt
lumiMask = LumiMask("Cert_294927-306462_13TeV_UL2017_Collisions17_GoldenJSON.txt")
# call lumimask
lumiMask
```

Error:
```
/opt/conda/lib/python3.8/site-packages/coffea/processor/executor.py in run_uproot_job(fileset, treename, processor_instance, executor, executor_args, pre_executor, pre_args, chunksize, maxchunks, metadata_cache, dynamic_chunksize, dynamic_chunksize_targets)
   1619     else:
   1620         pi_to_send = lz4f.compress(
-> 1621             cloudpickle.dumps(processor_instance), compression_level=pi_compression
   1622         )
   1623     closure = partial(

/opt/conda/lib/python3.8/site-packages/cloudpickle/cloudpickle_fast.py in dumps(obj, protocol, buffer_callback)
     71                 file, protocol=protocol, buffer_callback=buffer_callback
     72             )
---> 73             cp.dump(obj)
     74             return file.getvalue()
     75 

/opt/conda/lib/python3.8/site-packages/cloudpickle/cloudpickle_fast.py in dump(self, obj)
    561     def dump(self, obj):
    562         try:
--> 563             return Pickler.dump(self, obj)
    564         except RuntimeError as e:
    565             if "recursion" in e.args[0]:

TypeError: cannot pickle '_nrt_python._MemInfo' object
```
**Expected behavior**
I would expect to be able to load the lumiMask when called as ```lumiMask(events.run, events.luminosityBlock)```.

Discussion:

Patch:
diff --git a/coffea/lumi_tools/lumi_tools.py b/coffea/lumi_tools/lumi_tools.py
index 932eab9f..e4d92fe2 100644
--- a/coffea/lumi_tools/lumi_tools.py
+++ b/coffea/lumi_tools/lumi_tools.py
@@ -34,28 +34,6 @@ class LumiData(object):
                 ],  # not sure what lumi:0 means, appears to be always zero (DAQ off before beam dump?)
             },
         )
-        self.index = Dict.empty(
-            key_type=types.Tuple([types.uint32, types.uint32]), value_type=types.float64
-        )
-        self.build_lumi_table()
-
-    def build_lumi_table(self):
-        """Build index for numba-compiled functions
-
-        This needs to be executed upon unpickling, it should be part of
-        a custom deserialize function.
-        """
-        runs = self._lumidata[:, 0].astype("u4")
-        lumis = self._lumidata[:, 1].astype("u4")
-        LumiData._build_lumi_table_kernel(runs, lumis, self._lumidata, self.index)
-
-    @staticmethod
-    @numba.njit(parallel=False, fastmath=False)
-    def _build_lumi_table_kernel(runs, lumis, lumidata, index):
-        for i in range(len(runs)):
-            run = runs[i]
-            lumi = lumis[i]
-            index[(run, lumi)] = float(lumidata[i, 2])
 
     def get_lumi(self, runlumis):
         """Calculate integrated lumi
@@ -66,12 +44,28 @@ class LumiData(object):
                 A 2d numpy array of ``[[run,lumi], [run,lumi], ...]`` or `LumiList` object
                 of the lumiSections to integrate over.
         """
+        self.index = Dict.empty(
+            key_type=types.Tuple([types.uint32, types.uint32]), value_type=types.float64
+        )
+        runs = self._lumidata[:, 0].astype("u4")
+        lumis = self._lumidata[:, 1].astype("u4")
+        # fill self.index
+        LumiData._build_lumi_table_kernel(runs, lumis, self._lumidata, self.index)
+
         if isinstance(runlumis, LumiList):
             runlumis = runlumis.array
         tot_lumi = np.zeros((1,), dtype=np.float64)
         LumiData._get_lumi_kernel(runlumis[:, 0], runlumis[:, 1], self.index, tot_lumi)
         return tot_lumi[0]
 
+    @staticmethod
+    @numba.njit(parallel=False, fastmath=False)
+    def _build_lumi_table_kernel(runs, lumis, lumidata, index):
+        for i in range(len(runs)):
+            run = runs[i]
+            lumi = lumis[i]
+            index[(run, lumi)] = float(lumidata[i, 2])
+
     @staticmethod
     @numba.njit(parallel=False, fastmath=False)
     def _get_lumi_kernel(runs, lumis, index, tot_lumi):
@@ -100,7 +94,7 @@ class LumiMask(object):
         with open(jsonfile) as fin:
             goldenjson = json.load(fin)
 
-        self._masks = Dict.empty(key_type=types.uint32, value_type=types.uint32[:])
+        self._masks = {}
 
         for run, lumilist in goldenjson.items():
             mask = np.array(lumilist, dtype=np.uint32).flatten()
@@ -123,18 +117,19 @@ class LumiMask(object):
                 An array of dtype `bool` where valid (run, lumi) tuples
                 will have their corresponding entry set ``True``.
         """
+        # fill numba typed dict
+        _masks = Dict.empty(key_type=types.uint32, value_type=types.uint32[:])
+        for k, v in self._masks.items():
+            _masks[k] = v
+
         if isinstance(runs, ak.highlevel.Array):
             runs = ak.to_numpy(runs)
         if isinstance(lumis, ak.highlevel.Array):
             lumis = ak.to_numpy(lumis)
         mask_out = np.zeros(dtype="bool", shape=runs.shape)
-        LumiMask._apply_run_lumi_mask(self._masks, runs, lumis, mask_out)
+        LumiMask._apply_run_lumi_mask_kernel(_masks, runs, lumis, mask_out)
         return mask_out
 
-    @staticmethod
-    def _apply_run_lumi_mask(masks, runs, lumis, mask_out):
-        LumiMask._apply_run_lumi_mask_kernel(masks, runs, lumis, mask_out)
-
     # This could be run in parallel, but windows does not support it
     @staticmethod
     @numba.njit(parallel=False, fastmath=True)


Test patch:
diff --git a/tests/test_lumi_tools.py b/tests/test_lumi_tools.py
index f8914ce4..7134615f 100644
--- a/tests/test_lumi_tools.py
+++ b/tests/test_lumi_tools.py
@@ -1,5 +1,7 @@
 from __future__ import print_function, division
 
+import cloudpickle
+
 from coffea.lumi_tools import LumiData, LumiMask, LumiList
 from coffea.util import numpy as np
 
@@ -10,51 +12,81 @@ def test_lumidata():
 
     lumidata = LumiData("tests/samples/lumi_small.csv")
 
-    runslumis = np.zeros((10, 2), dtype=np.uint32)
-    runslumis[:, 0] = lumidata._lumidata[0:10, 0]
-    runslumis[:, 1] = lumidata._lumidata[0:10, 1]
-    lumi = lumidata.get_lumi(runslumis)
-    diff = abs(lumi - 1.539941814)
-    print("lumi:", lumi, "diff:", diff)
-    assert diff < 1e-4
-
-    # test build_lumi_table_kernel
-    py_index = Dict.empty(
-        key_type=types.Tuple([types.uint32, types.uint32]), value_type=types.float64
-    )
-    pyruns = lumidata._lumidata[:, 0].astype("u4")
-    pylumis = lumidata._lumidata[:, 1].astype("u4")
-    LumiData._build_lumi_table_kernel.py_func(
-        pyruns, pylumis, lumidata._lumidata, py_index
-    )
-
-    assert len(py_index) == len(lumidata.index)
+    # pickle & unpickle
+    lumidata_pickle = cloudpickle.loads(cloudpickle.dumps(lumidata))
 
-    # test get_lumi_kernel
-    py_tot_lumi = np.zeros((1,), dtype=np.float64)
-    LumiData._get_lumi_kernel.py_func(
-        runslumis[:, 0], runslumis[:, 1], py_index, py_tot_lumi
-    )
+    # check same internal lumidata
+    assert np.all(lumidata._lumidata == lumidata_pickle._lumidata)
 
-    assert abs(py_tot_lumi[0] - lumi) < 1e-4
+    runslumis = np.zeros((10, 2), dtype=np.uint32)
+    results = {"lumi": {}, "index": {}}
+    for ld in lumidata, lumidata_pickle:
+        runslumis[:, 0] = ld._lumidata[0:10, 0]
+        runslumis[:, 1] = ld._lumidata[0:10, 1]
+        lumi = ld.get_lumi(runslumis)
+        results["lumi"][ld] = lumi
+        diff = abs(lumi - 1.539941814)
+        print("lumi:", lumi, "diff:", diff)
+        assert diff < 1e-4
+
+        # test build_lumi_table_kernel
+        py_index = Dict.empty(
+            key_type=types.Tuple([types.uint32, types.uint32]), value_type=types.float64
+        )
+        pyruns = ld._lumidata[:, 0].astype("u4")
+        pylumis = ld._lumidata[:, 1].astype("u4")
+        LumiData._build_lumi_table_kernel.py_func(
+            pyruns, pylumis, ld._lumidata, py_index
+        )
+
+        assert len(py_index) == len(ld.index)
+
+        # test get_lumi_kernel
+        py_tot_lumi = np.zeros((1,), dtype=np.float64)
+        LumiData._get_lumi_kernel.py_func(
+            runslumis[:, 0], runslumis[:, 1], py_index, py_tot_lumi
+        )
+
+        assert abs(py_tot_lumi[0] - lumi) < 1e-4
+
+        # store results:
+        results["lumi"][ld] = lumi
+        results["index"][ld] = ld.index
+
+    assert np.all(results["lumi"][lumidata] == results["lumi"][lumidata_pickle])
+    assert len(results["index"][lumidata]) == len(results["index"][lumidata_pickle])
 
 
 def test_lumimask():
     lumimask = LumiMask(
         "tests/samples/Cert_294927-306462_13TeV_EOY2017ReReco_Collisions17_JSON.txt"
     )
+
+    # pickle & unpickle
+    lumimask_pickle = cloudpickle.loads(cloudpickle.dumps(lumimask))
+
+    # check same mask keys
+    keys = lumimask._masks.keys()
+    assert keys == lumimask_pickle._masks.keys()
+    # check same mask values
+    assert all(np.all(lumimask._masks[k] == lumimask_pickle._masks[k]) for k in keys)
+
     runs = np.array([303825, 123], dtype=np.uint32)
     lumis = np.array([115, 123], dtype=np.uint32)
-    mask = lumimask(runs, lumis)
-    print("mask:", mask)
-    assert mask[0]
-    assert not mask[1]
 
-    # test underlying py_func
-    py_mask = np.zeros(dtype="bool", shape=runs.shape)
-    LumiMask._apply_run_lumi_mask_kernel.py_func(lumimask._masks, runs, lumis, py_mask)
+    for lm in lumimask, lumimask_pickle:
+        mask = lm(runs, lumis)
+        print("mask:", mask)
+        assert mask[0]
+        assert not mask[1]
+
+        # test underlying py_func
+        py_mask = np.zeros(dtype="bool", shape=runs.shape)
+        LumiMask._apply_run_lumi_mask_kernel.py_func(lm._masks, runs, lumis, py_mask)
+
+        assert np.all(mask == py_mask)
 
-    assert np.all(mask == py_mask)
+    assert np.all(lumimask(runs, lumis) == lumimask_pickle(runs, lumis))
 
 
 def test_lumilist():