使用pyspark 的udf进行tensorflow 模型的预测报错 _pickle.PicklingError: Could not serialize object:

训练了一个bert分类模型，想通过pyspark 调用,一开始将tf模型的加载，预测以及转化成udf和pyspark调用写到了一个文件中，遇到了如下错误：

具体报错： _pickle.PicklingError: Could not serialize object: TypeError: can't pickle _thread.RLock objects

Traceback (most recent call last):
  File "/spark-3.0/python/lib/pyspark.zip/pyspark/serializers.py", line 468, in dumps
    return cloudpickle.dumps(obj, pickle_protocol)
  File "/spark-3.0/python/lib/pyspark.zip/pyspark/cloudpickle.py", line 1097, in dumps
    cp.dump(obj)
  File "/spark-3.0/python/lib/pyspark.zip/pyspark/cloudpickle.py", line 357, in dump
    return Pickler.dump(self, obj)
  File "/tmp/py3_env.zip/notebook/lib/python3.6/pickle.py", line 409, in dump
    self.save(obj)
  File "/tmp/py3_env.zip/notebook/lib/python3.6/pickle.py", line 476, in save
    f(self, obj) # Call unbound method with explicit self
  File "/tmp/py3_env.zip/notebook/lib/python3.6/pickle.py", line 736, in save_tuple
    save(element)
  File "/tmp/py3_env.zip/notebook/lib/python3.6/pickle.py", line 476, in save
    f(self, obj) # Call unbound method with explicit self
  File "/spark-3.0/python/lib/pyspark.zip/pyspark/cloudpickle.py", line 496, in save_function
    self.save_function_tuple(obj)
  File "/spark-3.0/python/lib/pyspark.zip/pyspark/cloudpickle.py", line 730, in save_function_tuple
    save(state)
  File "/tmp/py3_env.zip/notebook/lib/python3.6/pickle.py", line 476, in save
    f(self, obj) # Call unbound method with explicit self
  File "/tmp/py3_env.zip/notebook/lib/python3.6/pickle.py", line 821, in save_dict
    self._batch_setitems(obj.items())
  File "/tmp/py3_env.zip/notebook/lib/python3.6/pickle.py", line 847, in _batch_setitems
    save(v)
  File "/tmp/py3_env.zip/notebook/lib/python3.6/pickle.py", line 476, in save
    f(self, obj) # Call unbound method with explicit self
  File "/tmp/py3_env.zip/notebook/lib/python3.6/pickle.py", line 821, in save_dict
    self._batch_setitems(obj.items())
  File "/tmp/py3_env.zip/notebook/lib/python3.6/pickle.py", line 847, in _batch_setitems
    save(v)
  File "/tmp/py3_env.zip/notebook/lib/python3.6/pickle.py", line 521, in save
    self.save_reduce(obj=obj, *rv)
  File "/tmp/py3_env.zip/notebook/lib/python3.6/pickle.py", line 634, in save_reduce
    save(state)
  File "/tmp/py3_env.zip/notebook/lib/python3.6/pickle.py", line 476, in save
    f(self, obj) # Call unbound method with explicit self
  File "/tmp/py3_env.zip/notebook/lib/python3.6/pickle.py", line 821, in save_dict
    self._batch_setitems(obj.items())
  File "/tmp/py3_env.zip/notebook/lib/python3.6/pickle.py", line 847, in _batch_setitems
    save(v)
  File "/tmp/py3_env.zip/notebook/lib/python3.6/pickle.py", line 521, in save
    self.save_reduce(obj=obj, *rv)
  File "/tmp/py3_env.zip/notebook/lib/python3.6/pickle.py", line 634, in save_reduce
    save(state)
  File "/tmp/py3_env.zip/notebook/lib/python3.6/pickle.py", line 476, in save
    f(self, obj) # Call unbound method with explicit self
  File "/tmp/py3_env.zip/notebook/lib/python3.6/pickle.py", line 821, in save_dict
    self._batch_setitems(obj.items())
  File "/tmp/py3_env.zip/notebook/lib/python3.6/pickle.py", line 847, in _batch_setitems
    save(v)
  File "/tmp/py3_env.zip/notebook/lib/python3.6/pickle.py", line 496, in save
    rv = reduce(self.proto)
TypeError: can't pickle _thread.RLock objects
Traceback (most recent call last):
  File "/spark-3.0/python/lib/pyspark.zip/pyspark/serializers.py", line 468, in dumps
  File "/spark-3.0/python/lib/pyspark.zip/pyspark/cloudpickle.py", line 1097, in dumps
  File "/spark-3.0/python/lib/pyspark.zip/pyspark/cloudpickle.py", line 357, in dump
  File "/tmp/py3_env.zip/notebook/lib/python3.6/pickle.py", line 409, in dump
    self.save(obj)
  File "/tmp/py3_env.zip/notebook/lib/python3.6/pickle.py", line 476, in save
    f(self, obj) # Call unbound method with explicit self
  File "/tmp/py3_env.zip/notebook/lib/python3.6/pickle.py", line 736, in save_tuple
    save(element)
  File "/tmp/py3_env.zip/notebook/lib/python3.6/pickle.py", line 476, in save
    f(self, obj) # Call unbound method with explicit self
  File "/spark-3.0/python/lib/pyspark.zip/pyspark/cloudpickle.py", line 496, in save_function
  File "/spark-3.0/python/lib/pyspark.zip/pyspark/cloudpickle.py", line 730, in save_function_tuple
  File "/tmp/py3_env.zip/notebook/lib/python3.6/pickle.py", line 476, in save
    f(self, obj) # Call unbound method with explicit self
  File "/tmp/py3_env.zip/notebook/lib/python3.6/pickle.py", line 821, in save_dict
    self._batch_setitems(obj.items())
  File "/tmp/py3_env.zip/notebook/lib/python3.6/pickle.py", line 847, in _batch_setitems
    save(v)
  File "/tmp/py3_env.zip/notebook/lib/python3.6/pickle.py", line 476, in save
    f(self, obj) # Call unbound method with explicit self
  File "/tmp/py3_env.zip/notebook/lib/python3.6/pickle.py", line 821, in save_dict
    self._batch_setitems(obj.items())
  File "/tmp/py3_env.zip/notebook/lib/python3.6/pickle.py", line 847, in _batch_setitems
    save(v)
  File "/tmp/py3_env.zip/notebook/lib/python3.6/pickle.py", line 521, in save
    self.save_reduce(obj=obj, *rv)
  File "/tmp/py3_env.zip/notebook/lib/python3.6/pickle.py", line 634, in save_reduce
    save(state)
  File "/tmp/py3_env.zip/notebook/lib/python3.6/pickle.py", line 476, in save
    f(self, obj) # Call unbound method with explicit self
  File "/tmp/py3_env.zip/notebook/lib/python3.6/pickle.py", line 821, in save_dict
    self._batch_setitems(obj.items())
  File "/tmp/py3_env.zip/notebook/lib/python3.6/pickle.py", line 847, in _batch_setitems
    save(v)
  File "/tmp/py3_env.zip/notebook/lib/python3.6/pickle.py", line 521, in save
    self.save_reduce(obj=obj, *rv)
  File "/tmp/py3_env.zip/notebook/lib/python3.6/pickle.py", line 634, in save_reduce
    save(state)
  File "/tmp/py3_env.zip/notebook/lib/python3.6/pickle.py", line 476, in save
    f(self, obj) # Call unbound method with explicit self
  File "/tmp/py3_env.zip/notebook/lib/python3.6/pickle.py", line 821, in save_dict
    self._batch_setitems(obj.items())
  File "/tmp/py3_env.zip/notebook/lib/python3.6/pickle.py", line 847, in _batch_setitems
    save(v)
  File "/tmp/py3_env.zip/notebook/lib/python3.6/pickle.py", line 496, in save
    rv = reduce(self.proto)
TypeError: can't pickle _thread.RLock objects

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/tmp/tmp/spark-2949189d-9a2e-41fa-b392-4d53e3e9f9d0/bert_predict_new.py", line 75, in <module>
    cmt = df.withColumn("title_vec", model_prediction_udf(F.col("product_name")))
  File "/spark-3.0/python/lib/pyspark.zip/pyspark/sql/udf.py", line 197, in wrapper
  File "/spark-3.0/python/lib/pyspark.zip/pyspark/sql/udf.py", line 175, in __call__
  File "/spark-3.0/python/lib/pyspark.zip/pyspark/sql/udf.py", line 159, in _judf
  File "/spark-3.0/python/lib/pyspark.zip/pyspark/sql/udf.py", line 168, in _create_judf
  File "/spark-3.0/python/lib/pyspark.zip/pyspark/sql/udf.py", line 34, in _wrap_function
  File "/spark-3.0/python/lib/pyspark.zip/pyspark/rdd.py", line 2503, in _prepare_for_python_RDD
  File "/spark-3.0/python/lib/pyspark.zip/pyspark/serializers.py", line 478, in dumps
_pickle.PicklingError: Could not serialize object: TypeError: can't pickle _thread.RLock objects

这个原因可能在于tf 的对象以及代码逻辑不支持序列化，或者需要专门的序列化操作，，，为了解决问题，将tf的代码逻辑放到一个新的文件当中，暴露一个预测接口，，将pyspark代码放到一个文件中，然后在pyspark代码中引入该接口，将其转成udf即可。

from bert_inference import model_pred
model_pred_udf = F.udf(model_pred, T.FloatType())

之后就可以正常运行了