Python cannot pickle JIT class under multiprocessing

Question

I was trying to use multiprocessing and set a jit class as one of the arguments of the functions so that there would be a new jit class for each process. But it shows type error below.

TypeError: cannot pickle 'class1' object

Below is the python coding used.

spec = {}
spec['list_f'] = nb.types.List(nb.float64) #nb.int32
spec['array1d'] = nb.float32[:] 
spec['array2d'] = nb.float32[:,:] #nb.types.Array(dtype=nb.float32, ndim=1, layout="C")
spec['int_Yr'] = nb.int32
spec['list_i'] = nb.types.List(nb.int32)

Var_Idx = np.asarray(l_info.Var_Idx, dtype='int')
@jitclass(spec)
class class1:
    def __init__(self, list_f, array2d, array1d):
        self.list_f = list_f
        self.array2d = array2d
        self.array1d = array1d
        #self.list_i = Var_Idx
        self.int_Yr = 1 #Var_Idx[Liab_Data.Mat_Yr.value] #int(self.list_f[l_info.Var_Idx[Liab_Data.Mat_Yr.value]])
        self.method_print(self.int_Yr)
        
    def method_print(self, args):
        print(args)

#@njit
def function_class(class1, args):
    pid = os.getpid()
    print("PID:"+str(pid))
    class1.method_print(class1.int_Yr)
    
if __name__ == '__main__':
    list_f = [1.0, 1.0, 1.0, 1.0, 1.0, 100.0, 0.0, 0.0, 0.0, 99.0, 1.0, 9.8, 3.0, 2.0]
    
    #pdb.set_trace()
    array2d = np.ones((2, 1), dtype='f')
    array1d = np.zeros(1, dtype='f')
    
    
    class1 = class1(list_f, array2d, array1d)
    
    worker_count = os.cpu_count()
    queue = Queue()
    
    worker_pool = []
    idx_worker = 0
    for _ in range(worker_count):
        idx_worker += 1
        p = Process(target=function_class,  name='worker %d' % (idx_worker), args=(class1,idx_worker))
        p.start()
        worker_pool.append(p)
        
    for p in worker_pool:
        p.join()

score 0 · Answer 1 · answered Feb 07 '23 at 08:37

Jitclass cannot be passed to subprocess you could create the class in the sub process.

first pass the args to create the class

p = multiprocessing.Process(target=function_class, name='worker %d' % (idx_worker), args=(list_f, array2d, array1d, idx_worker))

then create the class in the function_class

def function_class(list_f, array2d, array1d, args):

    class2 = class1(list_f, array2d, array1d)
    pid = os.getpid()
    print("PID:" + str(pid), args)
    class2.method_print(class2.int_Yr)

BTW, I don't think calling a variable the same way of your class is a good idea

class1 = class1(list_f, array2d, array1d)

Python cannot pickle JIT class under multiprocessing

1 Answers1