[DistAutograd x JIT] Capture global state, dist autograd current context id, before thread switching triggered by JIT future.wait() (#36395)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/36395
titled
Test Plan:
# Unit tests
```
buck test mode/dev-nosan //caffe2/test/distributed/rpc/jit:dist_autograd_fork -- test_restore_context_after_swtich_to_jit_thread
buck build mode/dev-nosan //caffe2/test/distributed/rpc/jit:dist_autograd_fork
buck-out/gen/caffe2/test/distributed/rpc/jit/dist_autograd_fork\#binary.par -r test_restore_context_after_swtich_to_jit_thread
```
```
buck test mode/dev-nosan //caffe2/test/distributed/rpc:dist_autograd_fork -- test_backward_simple_script_call
buck build mode/dev-nosan //caffe2/test/distributed/rpc:dist_autograd_fork
buck-out/gen/caffe2/test/distributed/rpc/dist_autograd_fork\#binary.par -r test_backward_simple_script_call
```
Differential Revision: D7857991
fbshipit-source-id: 168e0e3846a50ea92d4f9450a30ccc6c13e2fcec
