Repository: pytorch
Issue: torch.compile+cudagraphs asserts in multithreaded context
Labels: triaged,module: multithreading,oncall: pt2,module: inductor

```
import threading
import torch

def foo(x, y):
    a = torch.sin(x)
    b = torch.cos(y)
    return a + b

opt_foo1 = torch.compile(foo, mode="max-autotune")

threads = []
for _ in range(32):
    threads.append(
        threading.Thread(
            target=opt_foo1,
            args=(torch.rand(2, 3).to("cuda"), torch.rand(2, 3).to("cuda")),
        )
    )

for t in threads:
    t.start()

for t in threads:
    t.join()
```

Stack trace looks like:
```
Exception in thread Thread-2 (foo):
Traceback (most recent call last):
  File "/usr/local/fbcode/platform010/lib/python3.10/threading.py", line 1016, in _bootstrap_inner
    self.run()
  File "/usr/local/fbcode/platform010/lib/python3.10/threading.py", line 953, in run
    self._target(*self._args, **self._kwargs)
  File "/data/users/suo/fbsource/buck-out/v2/gen/fbcode/da91e1664cbd9051/scripts/suo/__foo__/foo-inplace#link-tree/torch/_dynamo/eval_frame.py", line 410, in _fn
    return fn(*args, **kwargs)
  File "/data/users/suo/fbsource/buck-out/v2/gen/fbcode/da91e1664cbd9051/scripts/suo/__foo__/foo-inplace#link-tree/scripts/suo/foo.py", line 6, in foo
    def foo(x, y):
  File "/data/users/suo/fbsource/buck-out/v2/gen/fbcode/da91e1664cbd9051/scripts/suo/__foo__/foo-inplace#link-tree/torch/_dynamo/eval_frame.py", line 410, in _fn
    return fn(*args, **kwargs)
  File "/data/users/suo/fbsource/buck-out/v2/gen/fbcode/da91e1664cbd9051/scripts/suo/__foo__/foo-inplace#link-tree/torch/_dynamo/external_utils.py", line 36, in inner
    return fn(*args, **kwargs)
  File "/data/users/suo/fbsource/buck-out/v2/gen/fbcode/da91e1664cbd9051/scripts/suo/__foo__/foo-inplace#link-tree/torch/_functorch/aot_autograd.py", line 917, in forward
    return compiled_fn(full_args)
  File "/data/users/suo/fbsource/buck-out/v2/gen/fbcode/da91e1664cbd9051/scripts/suo/__foo__/foo-inplace#link-tree/torch/_functorch/_aot_autograd/utils.py", line 89, in g
    return f(*args)
  File "/data/users/suo/fbsource/buck-out/v2/gen/fbcode/da91e1664cbd9051/scripts/suo/__foo__/foo-inplace#link-tree/torch/_functorch/_aot_autograd/runtime_wrappers.py", line 101, in runtime_wrapper
    all_outs = call_func_at_runtime_with_args(
  File "/data/users/suo/fbsource/buck-out/v2/gen/fbcode/da91e1664cbd9051/scripts/suo/__foo__/foo-inplace#link-tree/torch/_functorch/_aot_autograd/utils.py", line 113, in call_func_at_runtime_with_args
    out = normalize_as_list(f(args))
  File "/data/users/suo/fbsource/buck-out/v2/gen/fbcode/da91e1664cbd9051/scripts/suo/__foo__/foo-inplace#link-tree/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py", line 181, in rng_functionalization_wrapper
    return compiled_fw(args)
  File "/data/users/suo/fbsource/buck-out/v2/gen/fbcode/da91e1664cbd9051/scripts/suo/__foo__/foo-inplace#link-tree/torch/_inductor/codecache.py", line 939, in __call__
    return self.current_callable(inputs)
  File "/data/users/suo/fbsource/buck-out/v2/gen/fbcode/da91e1664cbd9051/scripts/suo/__foo__/foo-inplace#link-tree/torch/_inductor/compile_fx.py", line 902, in run
    return compiled_fn(new_inputs)
  File "/data/users/suo/fbsource/buck-out/v2/gen/fbcode/da91e1664cbd9051/scripts/suo/__foo__/foo-inplace#link-tree/torch/_inductor/cudagraph_trees.py", line 382, in deferred_cudagraphify
    fn, out = cudagraphify(model, inputs, new_static_input_idxs, *args, **kwargs)
  File "/data/users/suo/fbsource/buck-out/v2/gen/fbcode/da91e1664cbd9051/scripts/suo/__foo__/foo-inplace#link-tree/torch/_inductor/cudagraph_trees.py", line 402, in cudagraphify
    manager = get_container(device_index).get_tree_manager()
  File "/data/users/suo/fbsource/buck-out/v2/gen/fbcode/da91e1664cbd9051/scripts/suo/__foo__/foo-inplace#link-tree/torch/_inductor/cudagraph_trees.py", line 338, in get_container
    container_dict = get_obj(local, "tree_manager_containers")
  File "/data/users/suo/fbsource/buck-out/v2/gen/fbcode/da91e1664cbd9051/scripts/suo/__foo__/foo-inplace#link-tree/torch/_inductor/cudagraph_trees.py", line 333, in get_obj
    assert torch._C._is_key_in_tls(attr_name)
AssertionError
```

Context: it would be useful in inference contexts to be able to run multiple threads per GPU, to improve GPU utilization in the presence of small-batch inference requests. 

As far as I know we don't have any thread-safety guarantees in PT2, so this issue is not really about this particular stack trace but rather to track the general question of using torch.compile in a multithreaded context.

cc @ezyang @gchanan @zou3519 @kadeng @msaroufim @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @yf225 @chenyang78 @muchulee8 @ColinPeppler @amjames @desertfire @bdhirsh @anijain2305 @peterbell10 @aakhundov