diff --git a/src/pb_stub.cc b/src/pb_stub.cc index 173a9c3e..8d87ab1f 100644 --- a/src/pb_stub.cc +++ b/src/pb_stub.cc @@ -1059,13 +1059,7 @@ Stub::~Stub() } #endif - // Ensure the interpreter is active before trying to clean up. - if (Py_IsInitialized()) { - py::gil_scoped_acquire acquire; - py::object async_event_loop_local(std::move(async_event_loop_)); - py::object background_futures_local(std::move(background_futures_)); - py::object model_instance_local(std::move(model_instance_)); - } + DestroyPythonObjects(); stub_message_queue_.reset(); parent_message_queue_.reset(); @@ -1088,6 +1082,11 @@ Stub::GetOrCreateInstance() void Stub::DestroyInstance() { + if (!stub_instance) { + return; + } + + stub_instance->DestroyPythonObjects(); stub_instance.reset(); } @@ -1503,6 +1502,20 @@ Stub::GetCUDAMemoryPoolAddress(std::unique_ptr& ipc_message) #endif } +void +Stub::DestroyPythonObjects() +{ + // Ensure the interpreter is active before trying to clean up. + if (Py_IsInitialized()) { + py::gil_scoped_acquire acquire; + py::object async_event_loop_local(std::move(async_event_loop_)); + py::object background_futures_local(std::move(background_futures_)); + py::object model_instance_local(std::move(model_instance_)); + py::object deserialize_bytes_local(std::move(deserialize_bytes_)); + py::object serialize_bytes_local(std::move(serialize_bytes_)); + } +} + void Stub::ProcessBLSResponseDecoupled(std::unique_ptr& ipc_message) { diff --git a/src/pb_stub.h b/src/pb_stub.h index e39014f8..b2a97178 100644 --- a/src/pb_stub.h +++ b/src/pb_stub.h @@ -267,6 +267,11 @@ class Stub { /// Get the CUDA memory pool address from the parent process. void GetCUDAMemoryPoolAddress(std::unique_ptr& ipc_message); + /// Cleans up Python objects and must be called before the destructor. + /// This prevents problems that occur when Python object destructors + /// call Stub::GetOrCreateInstance. + void DestroyPythonObjects(); + /// Calls the user's is_ready() Python method and returns its response /// when handling model readiness check requests. void ProcessUserModelReadinessRequest(