Repository: pytorch
Issue: Race condition on shutdown involving PThreadPool and autograd
Labels: triaged,module: multithreading,module: sanitizers

### 🐛 Describe the bug

I don't have an OSS repro. The internal test that (only very rarely) triggers the race is https://www.internalfb.com/intern/test/281475084563302?ref_report_id=0

Here is the TSAN error (I verified all the backtraces are benign):

```
WARNING: ThreadSanitizer: data race (pid=11205)
  Write of size 8 at 0x00000340fef0 by main thread:
    #0 std::unique_ptr<caffe2::PThreadPool, std::default_delete<caffe2::PThreadPool> >::~unique_ptr() fbcode/third-party-buck/platform010/build/libgcc/include/c++/trunk/bits/unique_ptr.h:362 (lite_trainer_test+0x2b9a1d1)
    #1 cxa_at_exit_wrapper(void*) <null> (lite_trainer_test+0x3219eb0)

  Previous read of size 8 at 0x00000340fef0 by thread T1:
    #0 std::__uniq_ptr_impl<caffe2::PThreadPool, std::default_delete<caffe2::PThreadPool> >::_M_ptr() const fbcode/third-party-buck/platform010/build/libgcc/include/c++/trunk/bits/unique_ptr.h:173 (lite_trainer_test+0x2b9ba61)
    #1 std::unique_ptr<caffe2::PThreadPool, std::default_delete<caffe2::PThreadPool> >::get() const fbcode/third-party-buck/platform010/build/libgcc/include/c++/trunk/bits/unique_ptr.h:422 (lite_trainer_test+0x2b9a385)
    #2 caffe2::pthreadpool() xplat/caffe2/caffe2/utils/threadpool/pthreadpool-cpp.cc:105 (lite_trainer_test+0x2b99fe2)
    #3 at::init_num_threads() xplat/caffe2/aten/src/ATen/ParallelNative.cpp:207 (lite_trainer_test+0x8961b7)
    #4 torch::autograd::Engine::thread_init(int, std::shared_ptr<torch::autograd::ReadyQueue> const&, bool) xplat/caffe2/torch/csrc/autograd/engine.cpp:353 (lite_trainer_test+0x2c5ac25)
    #5 __invoke_impl<void, void (torch::autograd::Engine::*)(int, const std::shared_ptr<torch::autograd::ReadyQueue> &, bool), torch::autograd::Engine *, signed char, std::shared_ptr<torch::autograd::ReadyQueue>, bool> fbcode/third-party-buck/platform010/build/libgcc/include/c++/trunk/bits/invoke.h:74 (lite_trainer_test+0x2c6c629)
    #6 __invoke<void (torch::autograd::Engine::*)(int, const std::shared_ptr<torch::autograd::ReadyQueue> &, bool), torch::autograd::Engine *, signed char, std::shared_ptr<torch::autograd::ReadyQueue>, bool> fbcode/third-party-buck/platform010/build/libgcc/include/c++/trunk/bits/invoke.h:96 (lite_trainer_test+0x2c6c629)
    #7 void std::thread::_Invoker<std::tuple<void (torch::autograd::Engine::*)(int, std::shared_ptr<torch::autograd::ReadyQueue> const&, bool), torch::autograd::Engine*, signed char, std::shared_ptr<torch::autograd::ReadyQueue>, bool> >::_M_invoke<0ul, 1ul, 2ul, 3ul, 4ul>(std::_Index_tuple<0ul, 1ul, 2ul, 3ul, 4ul>) fbcode/third-party-buck/platform010/build/libgcc/include/c++/trunk/bits/std_thread.h:253 (lite_trainer_test+0x2c6c629)
    #8 operator() fbcode/third-party-buck/platform010/build/libgcc/include/c++/trunk/bits/std_thread.h:260 (lite_trainer_test+0x2c6c441)
    #9 std::thread::_State_impl<std::thread::_Invoker<std::tuple<void (torch::autograd::Engine::*)(int, std::shared_ptr<torch::autograd::ReadyQueue> const&, bool), torch::autograd::Engine*, signed char, std::shared_ptr<torch::autograd::ReadyQueue>, bool> > >::_M_run() fbcode/third-party-buck/platform010/build/libgcc/include/c++/trunk/bits/std_thread.h:211 (lite_trainer_test+0x2c6c441)
    #10 execute_native_thread_routine /home/engshare/third-party2/libgcc/11.x/src/gcc-11.x/x86_64-facebook-linux/libstdc++-v3/src/c++11/../../../.././libstdc++-v3/src/c++11/thread.cc:82:18 (libstdc++.so.6+0xdf4e4)

  Location is global 'caffe2::pthreadpool()::threadpool' of size 8 at 0x00000340fef0 (lite_trainer_test+0x00000340fef0)

  Thread T1 (tid=11291, running) created by main thread at:
    #0 pthread_create <null> (lite_trainer_test+0x3258e8d)
    #1 __gthread_create /home/engshare/third-party2/libgcc/11.x/src/gcc-11.x/x86_64-facebook-linux/libstdc++-v3/include/x86_64-facebook-linux/bits/gthr-default.h:663:35 (libstdc++.so.6+0xdf80e)
    #2 std::thread::_M_start_thread(std::unique_ptr<std::thread::_State, std::default_delete<std::thread::_State> >, void (*)()) /home/engshare/third-party2/libgcc/11.x/src/gcc-11.x/x86_64-facebook-linux/libstdc++-v3/src/c++11/../../../.././libstdc++-v3/src/c++11/thread.cc:147:37 (libstdc++.so.6+0xdf80e)
    #3 torch::autograd::Engine::start_device_threads() xplat/caffe2/torch/csrc/autograd/engine.cpp:1483 (lite_trainer_test+0x2c60efe)
    #4 __invoke_impl<void, void (torch::autograd::Engine::*const &)(), torch::autograd::Engine *> fbcode/third-party-buck/platform010/build/libgcc/include/c++/trunk/bits/invoke.h:74 (lite_trainer_test+0x2c6b8d6)
    #5 __invoke<void (torch::autograd::Engine::*const &)(), torch::autograd::Engine *> fbcode/third-party-buck/platform010/build/libgcc/include/c++/trunk/bits/invoke.h:96 (lite_trainer_test+0x2c6b8d6)
    #6 operator()<torch::autograd::Engine *> fbcode/third-party-buck/platform010/build/libgcc/include/c++/trunk/functional:131 (lite_trainer_test+0x2c6b8d6)
    #7 invoke<void (torch::autograd::Engine::*&)(), torch::autograd::Engine *> fbsource/c10/util/C++17.h:195 (lite_trainer_test+0x2c6b8d6)
    #8 void c10::once_flag::call_once_slow<void (torch::autograd::Engine::*)(), torch::autograd::Engine*>(void (torch::autograd::Engine::*&&)(), torch::autograd::Engine*&&) fbsource/c10/util/CallOnce.h:50 (lite_trainer_test+0x2c6b8d6)
    #9 call_once<c10::once_flag, void (torch::autograd::Engine::*)(), torch::autograd::Engine *> fbsource/c10/util/CallOnce.h:20 (lite_trainer_test+0x2c60b96)
    #10 torch::autograd::Engine::initialize_device_threads_pool() xplat/caffe2/torch/csrc/autograd/engine.cpp:1277 (lite_trainer_test+0x2c60b96)
    #11 torch::autograd::Engine::execute_with_graph_task(std::shared_ptr<torch::autograd::GraphTask> const&, std::shared_ptr<torch::autograd::Node>, torch::autograd::InputBuffer&&) xplat/caffe2/torch/csrc/autograd/engine.cpp:1285 (lite_trainer_test+0x2c61011)
    #12 torch::autograd::Engine::execute(std::vector<torch::autograd::Edge, std::allocator<torch::autograd::Edge> > const&, std::vector<at::Tensor, std::allocator<at::Tensor> > const&, bool, bool, bool, std::vector<torch::autograd::Edge, std::allocator<torch::autograd::Edge> > const&) xplat/caffe2/torch/csrc/autograd/engine.cpp:1257 (lite_trainer_test+0x2c63515)
    #13 torch::autograd::run_backward(std::vector<at::Tensor, std::allocator<at::Tensor> > const&, std::vector<at::Tensor, std::allocator<at::Tensor> > const&, bool, bool, std::vector<at::Tensor, std::allocator<at::Tensor> > const&, bool, bool) xplat/caffe2/torch/csrc/autograd/autograd.cpp:139 (lite_trainer_test+0x2c48a46)
    #14 torch::autograd::backward(std::vector<at::Tensor, std::allocator<at::Tensor> > const&, std::vector<at::Tensor, std::allocator<at::Tensor> > const&, c10::optional<bool>, bool, std::vector<at::Tensor, std::allocator<at::Tensor> > const&) xplat/caffe2/torch/csrc/autograd/autograd.cpp:171 (lite_trainer_test+0x2c47f62)
    #15 torch::autograd::VariableHooks::_backward(at::Tensor const&, c10::ArrayRef<at::Tensor>, c10::optional<at::Tensor> const&, c10::optional<bool>, bool) const xplat/caffe2/torch/csrc/autograd/variable.cpp:564 (lite_trainer_test+0x2c88c02)
    #16 at::Tensor::_backward(c10::ArrayRef<at::Tensor>, c10::optional<at::Tensor> const&, c10::optional<bool>, bool) const xplat/caffe2/aten/src/ATen/core/Tensor.cpp:140 (lite_trainer_test+0x8ba331)
    #17 at::Tensor::backward(at::Tensor const&, c10::optional<bool>, bool, c10::optional<c10::ArrayRef<at::Tensor> >) const fbsource/ATen/core/TensorBody.h:445 (lite_trainer_test+0x721f3c)
    #18 torch::jit::LiteTrainer_LinearRegression_Test::TestBody() xplat/caffe2/fb/lite_trainer/lite_trainer_test.cpp:71 (lite_trainer_test+0x714c95)
    #19 void testing::internal::HandleExceptionsInMethodIfSupported<testing::Test, void>(testing::Test*, void (testing::Test::*)(), char const*) third-party/googletest/1.11.0/googletest/googletest/src/gtest.cc:2665 (lite_trainer_test+0x31b3eed)
    #20 testing::Test::Run() third-party/googletest/1.11.0/googletest/googletest/src/gtest.cc:2682 (lite_trainer_test+0x31b3b13)
    #21 testing::TestInfo::Run() third-party/googletest/1.11.0/googletest/googletest/src/gtest.cc:2861 (lite_trainer_test+0x31b50b6)
    #22 testing::TestSuite::Run() third-party/googletest/1.11.0/googletest/googletest/src/gtest.cc:3015 (lite_trainer_test+0x31b6c42)
    #23 testing::internal::UnitTestImpl::RunAllTests() third-party/googletest/1.11.0/googletest/googletest/src/gtest.cc:5917 (lite_trainer_test+0x31cc067)
    #24 bool testing::internal::HandleExceptionsInMethodIfSupported<testing::internal::UnitTestImpl, bool>(testing::internal::UnitTestImpl*, bool (testing::internal::UnitTestImpl::*)(), char const*) third-party/googletest/1.11.0/googletest/googletest/src/gtest.cc:2665 (lite_trainer_test+0x31cb8f4)
    #25 testing::UnitTest::Run() third-party/googletest/1.11.0/googletest/googletest/src/gtest.cc:5500 (lite_trainer_test+0x31cb34e)
    #26 RUN_ALL_TESTS() fbsource/gtest/gtest.h:2490 (lite_trainer_test+0x31a7a0a)
    #27 main third-party/googletest/1.11.0/googletest/googletest/src/gtest_main.cc:52 (lite_trainer_test+0x31a79b2)
```

### Versions

main