Использование API набора данных TensorFlow приводит к зависанию процесса в деструкторе сеанса

Резюме:

Мы используем API TensorFlow Datasets. Более конкретно, мы используем tf.data.Dataset.from_generator для создания набора данных на основе функции генератора.

Когда Python приходит к мусору, собираем наш объект tf.Session его деструктор делает вызов в TensorFlow для удаления сеанса (tf_session.TF_DeleteSession). Этот вызов зависает, потому что он пытается выполнить функцию tf.py_func, но не может получить блокировку глобального интерпретатора Python. Функция, которую пытается выполнить, оказывается функцией "финализировать" из нашего набора данных.

Подробнее:

Когда наш объект tf.Session - это мусор, собранный в Python, его деструктор (метод __del__) висит бесконечно. Проблема заключается в том, что этот вызов в BaseSession:

tf_session.TF_DeleteSession(self._session)

Запуск lldb показывает следующую трассировку стека:

* thread #1, queue = 'com.apple.main-thread', stop reason = signal SIGSTOP
  * frame #0: 0x0000000101855e7e libsystem_kernel.dylib'__psynch_cvwait + 10
    frame #1: 0x000000010188d662 libsystem_pthread.dylib'_pthread_cond_wait + 732
    frame #2: 0x00000001019b6cb0 libc++.1.dylib'std::__1::condition_variable::wait(std::__1::unique_lock<std::__1::mutex>&) + 18
    frame #3: 0x000000011279a63b libtensorflow_framework.so'nsync::nsync_mu_semaphore_p_with_deadline(nsync::nsync_semaphore_s_*, timespec) + 283
    frame #4: 0x0000000112796eb7 libtensorflow_framework.so'nsync::nsync_cv_wait_with_deadline_generic(nsync::nsync_cv_s_*, void*, void (*)(void*), void (*)(void*), timespec, nsync::nsync_note_s_*) + 423
    frame #5: 0x0000000112797621 libtensorflow_framework.so'nsync::nsync_cv_wait(nsync::nsync_cv_s_*, nsync::nsync_mu_s_*) + 49
    frame #6: 0x00000001090810e3 _pywrap_tensorflow_internal.so'tensorflow::Notification::WaitForNotification() + 67
    frame #7: 0x0000000109d4d809 _pywrap_tensorflow_internal.so'tensorflow::CapturedFunction::RunInstantiated(std::__1::vector<tensorflow::Tensor, std::__1::allocator<tensorflow::Tensor> > const&, std::__1::vector<tensorflow::Tensor, std::__1::allocator<tensorflow::Tensor> >*) + 649
    frame #8: 0x0000000109cffa21 _pywrap_tensorflow_internal.so'tensorflow::(anonymous namespace)::GeneratorDatasetOp::Dataset::Iterator::~Iterator() + 97
    frame #9: 0x0000000109cffb8e _pywrap_tensorflow_internal.so'tensorflow::(anonymous namespace)::GeneratorDatasetOp::Dataset::Iterator::~Iterator() + 14
    frame #10: 0x0000000109cfd669 _pywrap_tensorflow_internal.so'tensorflow::(anonymous namespace)::FlatMapDatasetOp::Dataset::Iterator::~Iterator() + 105
    frame #11: 0x0000000109cfd6de _pywrap_tensorflow_internal.so'tensorflow::(anonymous namespace)::FlatMapDatasetOp::Dataset::Iterator::~Iterator() + 14
    frame #12: 0x00000001019e98fd libc++.1.dylib'std::__1::__shared_weak_count::__release_shared() + 43
    frame #13: 0x0000000109d0a579 _pywrap_tensorflow_internal.so'tensorflow::(anonymous namespace)::IteratorResource::~IteratorResource() + 169
    frame #14: 0x0000000109d0a5fe _pywrap_tensorflow_internal.so'tensorflow::(anonymous namespace)::IteratorResource::~IteratorResource() + 14
    frame #15: 0x000000011226db4d libtensorflow_framework.so'tensorflow::ResourceMgr::DoDelete(std::__1::basic_string<char, std::__1::char_traits<char>, std::__1::allocator<char> > const&, unsigned long long, std::__1::basic_string<char, std::__1::char_traits<char>, std::__1::allocator<char> > const&, std::__1::basic_string<char, std::__1::char_traits<char>, std::__1::allocator<char> > const&) + 301
    frame #16: 0x000000011226dd50 libtensorflow_framework.so'tensorflow::ResourceMgr::DoDelete(std::__1::basic_string<char, std::__1::char_traits<char>, std::__1::allocator<char> > const&, std::__1::type_index, std::__1::basic_string<char, std::__1::char_traits<char>, std::__1::allocator<char> > const&) + 192
    frame #17: 0x0000000109d0c558 _pywrap_tensorflow_internal.so'tensorflow::(anonymous namespace)::OneShotIteratorOp::~OneShotIteratorOp() + 104
    frame #18: 0x0000000109d0c71e _pywrap_tensorflow_internal.so'tensorflow::(anonymous namespace)::OneShotIteratorOp::~OneShotIteratorOp() + 14
    frame #19: 0x00000001122670ff libtensorflow_framework.so'tensorflow::OpSegment::Item::~Item() + 63
    frame #20: 0x0000000112267ffd libtensorflow_framework.so'tensorflow::OpSegment::RemoveHold(std::__1::basic_string<char, std::__1::char_traits<char>, std::__1::allocator<char> > const&) + 205
    frame #21: 0x000000010b880b42 _pywrap_tensorflow_internal.so'tensorflow::DirectSession::~DirectSession() + 546
    frame #22: 0x000000010b88108e _pywrap_tensorflow_internal.so'tensorflow::DirectSession::~DirectSession() + 14
    frame #23: 0x000000010935dfd3 _pywrap_tensorflow_internal.so'TF_DeleteSession + 931
    frame #24: 0x0000000109006e5a _pywrap_tensorflow_internal.so'_wrap_TF_DeleteSession(_object*, _object*) + 122
    frame #25: 0x00000001007bb688 Python'_PyCFunction_FastCallDict + 568
    frame #26: 0x00000001008443e4 Python'call_function + 612
    frame #27: 0x0000000100849d84 Python'_PyEval_EvalFrameDefault + 21892
    frame #28: 0x00000001008447cc Python'_PyFunction_FastCallDict + 828
    frame #29: 0x000000010075f984 Python'_PyObject_FastCallDict + 356
    frame #30: 0x000000010075faa0 Python'_PyObject_Call_Prepend + 208
    frame #31: 0x000000010075f8d4 Python'_PyObject_FastCallDict + 180
    frame #32: 0x00000001007d6579 Python'slot_tp_finalize + 121
    frame #33: 0x000000010089b18a Python'collect + 1418
    frame #34: 0x000000010089b8c3 Python'_PyGC_CollectIfEnabled + 99
    frame #35: 0x000000010087af57 Python'Py_FinalizeEx + 119
    frame #36: 0x000000010087b0e0 Python'Py_Exit + 16
    frame #37: 0x000000010087ef4c Python'handle_system_exit + 252
    frame #38: 0x000000010087f1a5 Python'PyErr_PrintEx + 437
    frame #39: 0x0000000100880a1d Python'PyRun_SimpleStringFlags + 125
    frame #40: 0x00000001008992a4 Python'Py_Main + 1812
    frame #41: 0x0000000100000dfe Python
    frame #42: 0x0000000100000c34 Python

Похоже, что деструктор сеанса ждет завершения операции op. Преступник, похоже, PyFuncOp, который не проходит мимо этой строки:

py_threadstate = PyGILState_Ensure();

Таким образом, похоже, что этот op пытается получить GIL, но не может. Мое предположение заключается в том, что этот py_func является "финализацией" для набора данных (из _GeneratorDataset).

Мое предположение заключается в том, что когда Python вызывает tf_session.TF_DeleteSession(self._session) что GIL должен быть выпущен, и поэтому PyFuncOp должен затем получить его снова. В самом деле, когда я пишу изолированный тест, чтобы попытаться воспроизвести это, я не вижу этой проблемы, и GIL приобретается успешно.

Я бы поднял ошибку в TensorFlow для этого, но мне не удалось воспроизвести ее в небольшом тестовом примере (это происходит только в нашей "реальной" системе, хотя она полностью воспроизводима там).

Может ли кто-нибудь помочь мне понять, что здесь происходит?