1 | #pragma once |
2 | |
3 | #include <c10/util/Exception.h> |
4 | |
5 | /// This file provides some simple utilities for detecting common deadlocks in |
6 | /// PyTorch. For now, we focus exclusively on detecting Python GIL deadlocks, |
7 | /// as the GIL is a wide ranging lock that is taken out in many situations. |
8 | /// The basic strategy is before performing an operation that may block, you |
9 | /// can use TORCH_ASSERT_NO_GIL_WITHOUT_PYTHON_DEP() to assert that the GIL is |
10 | /// not held. This macro is to be used in contexts where no static dependency |
11 | /// on Python is available (we will handle indirecting a virtual call for you). |
12 | /// |
13 | /// If the GIL is held by a torchdeploy interpreter, we always report false. |
14 | /// If you are in a context where Python bindings are available, it's better |
15 | /// to directly assert on PyGILState_Check (as it avoids a vcall and also |
16 | /// works correctly with torchdeploy.) |
17 | |
18 | namespace c10 { |
19 | |
20 | #define TORCH_ASSERT_NO_GIL_WITHOUT_PYTHON_DEP() \ |
21 | TORCH_INTERNAL_ASSERT( \ |
22 | !c10::impl::check_python_gil(), \ |
23 | "Holding GIL before a blocking operation! Please release the GIL before blocking, or see https://github.com/pytorch/pytorch/issues/56297 for how to release the GIL for destructors of objects") |
24 | |
25 | namespace impl { |
26 | |
27 | C10_API bool check_python_gil(); |
28 | |
29 | struct C10_API PythonGILHooks { |
30 | virtual ~PythonGILHooks() = default; |
31 | // Returns true if we hold the GIL. If not linked against Python we |
32 | // always return false. |
33 | virtual bool check_python_gil() const = 0; |
34 | }; |
35 | |
36 | C10_API void SetPythonGILHooks(PythonGILHooks* factory); |
37 | |
38 | // DO NOT call this registerer from a torch deploy instance! You will clobber |
39 | // other registrations |
40 | struct C10_API PythonGILHooksRegisterer { |
41 | explicit PythonGILHooksRegisterer(PythonGILHooks* factory) { |
42 | SetPythonGILHooks(factory); |
43 | } |
44 | ~PythonGILHooksRegisterer() { |
45 | SetPythonGILHooks(nullptr); |
46 | } |
47 | }; |
48 | |
49 | } // namespace impl |
50 | } // namespace c10 |
51 | |