1 /* Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 16 #ifndef TENSORFLOW_COMPILER_XLA_PYTHON_PYTHON_REF_MANAGER_H_ 17 #define TENSORFLOW_COMPILER_XLA_PYTHON_PYTHON_REF_MANAGER_H_ 18 19 #include <atomic> 20 #include <deque> 21 22 #include "absl/base/thread_annotations.h" 23 #include "absl/container/inlined_vector.h" 24 #include "absl/synchronization/mutex.h" 25 #include "absl/types/span.h" 26 #include "pybind11/pybind11.h" 27 28 namespace xla { 29 30 // Class that manages destruction of Python objects. 31 // 32 // We must not destroy Python objects without holding the GIL. However, we 33 // frequently want to hold references to Python objects for the duration of 34 // an asynchronous transfer on a Stream, and release our reference when the 35 // transfer completes. 36 // 37 // This class holds references to Python objects outside a GIL scope, that can 38 // be collected later when the GIL is held by calling CollectGarbage(). 39 class PythonRefManager { 40 public: 41 PythonRefManager() = default; 42 43 // Holds references to a set of pybind11::objects, adding the references to 44 // the PythonRefManager on destruction. 45 class ManagedPyObjects { 46 public: 47 ManagedPyObjects() = default; 48 ManagedPyObjects(PythonRefManager* manager, 49 absl::Span<pybind11::object> objects); 50 51 ~ManagedPyObjects(); 52 53 ManagedPyObjects(const ManagedPyObjects& other) = delete; 54 ManagedPyObjects(ManagedPyObjects&& other) = default; 55 ManagedPyObjects& operator=(const ManagedPyObjects& other) = delete; 56 ManagedPyObjects& operator=(ManagedPyObjects&& other) = default; 57 58 private: 59 PythonRefManager* manager_ = nullptr; 60 absl::InlinedVector<pybind11::object, 1> objects_; 61 }; 62 63 // Creates a managed std::shared_ptr to an object. When the shared_ptr is 64 // destroyed, the reference to 'object' will be added to python_garbage_, 65 // and collected next time CollectGarbage() is called. 66 std::shared_ptr<ManagedPyObjects> ManageReference(pybind11::object object); 67 std::shared_ptr<ManagedPyObjects> ManageReferences( 68 absl::Span<pybind11::object> objects); 69 70 // Adds garbage objects to the manager. 71 void AddGarbage(absl::Span<pybind11::object> garbage); 72 void AddGarbage(absl::Span<std::pair<PyCodeObject*, int> const> garbage); 73 74 // Releases the contents of python_garbage_. Requires that the GIL is held. 75 // The client calls this method during API entry points where the GIL is held 76 // to free any garbage that has accumulated. 77 void CollectGarbage(); 78 79 // Cheaper version of CollectGarbage() with relaxed consistency and frequency. 80 // The purpose of this function is to amortize lock acquisition costs over 81 // a larger number of API calls. MaybeCollectGarbage()82 void MaybeCollectGarbage() { 83 if (garbage_count_.load(std::memory_order_relaxed) >= 100) { 84 CollectGarbage(); 85 } 86 } 87 88 private: 89 absl::Mutex mu_; 90 std::deque<pybind11::object> python_garbage_ ABSL_GUARDED_BY(mu_); 91 92 // Writes to garbage_count_ are protected by mu_, reads are not protected. 93 std::atomic<int> garbage_count_{0}; 94 }; 95 96 // A global PythonRefManager. Unless `CollectGarbage()` is called before 97 // shutdown, this container will hold on to Python objects and thus cause a 98 // leak. This behavior is similar to `tensorflow::ClearDecRefCache()`. 99 PythonRefManager* GlobalPyRefManager(); 100 101 } // namespace xla 102 103 #endif // TENSORFLOW_COMPILER_XLA_PYTHON_PYTHON_REF_MANAGER_H_ 104