Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 71e67ba

Browse files
pytorchmergebotxuhancn
pytorchmergebot
authored and
xuhancn
committed
Revert "Implement thread-local config with contextvars (#173568)"
This reverts commit c904bdc. Reverted #173568 on behalf of https://github.com/georgehong due to tests failing with torch._dynamo hit config.recompile_limit (8) ([comment](#173568 (comment)))
1 parent 65254e4 commit 71e67ba

10 files changed

Lines changed: 34 additions & 237 deletions

File tree

‎benchmarks/dynamo/ci_expected_accuracy/inductor_torchbench_training.csv‎

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -170,10 +170,6 @@ soft_actor_critic,pass,6
170170

171171

172172

173-
speech_transformer,pass,16
174-
175-
176-
177173
squeezenet1_1,pass,6
178174

179175

‎test/dynamo/test_misc.py‎

Lines changed: 0 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,6 @@
4444
from torch import Tensor
4545
from torch._C import FileCheck
4646
from torch._dynamo import allow_in_graph
47-
from torch._dynamo.comptime import comptime
4847
from torch._dynamo.eval_frame import _debug_get_cache_entry_list
4948
from torch._dynamo.exc import Unsupported
5049
from torch._dynamo.source import ConstantSource, GetItemSource, LocalSource
@@ -7121,49 +7120,6 @@ def test_recompile(foo, *, exp_frame_count):
71217120
torch._dynamo.reset()
71227121
test_recompile(foo_graph_break, exp_frame_count=2)
71237122

7124-
def test_multithread_compile_dynamic(self):
7125-
def f(x):
7126-
comptime.assert_static(x.shape[0])
7127-
return x * x
7128-
7129-
def _do_test(func):
7130-
success = True
7131-
7132-
def run(offset):
7133-
for i in range(20):
7134-
print(func(torch.randn(i * 2 + offset)))
7135-
7136-
t1 = threading.Thread(target=run, args=[0])
7137-
t2 = threading.Thread(target=run, args=[1])
7138-
7139-
def exc_hook(x):
7140-
nonlocal success
7141-
success = False
7142-
7143-
try:
7144-
threading.excepthook = exc_hook
7145-
t1.start()
7146-
t2.start()
7147-
7148-
t1.join()
7149-
t2.join()
7150-
finally:
7151-
threading.excepthook = threading.__excepthook__
7152-
self.assertTrue(success)
7153-
7154-
_do_test(torch.compile(f, backend="eager", dynamic=False))
7155-
torch._dynamo.reset()
7156-
7157-
f_opt = torch.compile(f, backend="eager")
7158-
7159-
def g(x):
7160-
with torch._dynamo.config.patch(
7161-
automatic_dynamic_shapes=False, assume_static_by_default=True
7162-
):
7163-
f_opt(x)
7164-
7165-
_do_test(g)
7166-
71677123
def test_backend_match_guard_multi_threads(self):
71687124
x = torch.randn([3, 4])
71697125

‎test/inductor/test_config.py‎

Lines changed: 1 addition & 64 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,7 @@
77
from torch._inductor import config
88
from torch._inductor.pattern_matcher import PatternMatcherPass
99
from torch._inductor.test_case import run_tests, TestCase
10-
from torch.testing._internal.inductor_utils import GPU_TYPE, HAS_CPU, HAS_TRITON
11-
from torch.testing._internal.triton_utils import requires_gpu
10+
from torch.testing._internal.inductor_utils import HAS_CPU, HAS_TRITON
1211

1312

1413
def dummy_fn(x):
@@ -328,68 +327,6 @@ def fn2(x, y):
328327
self.assertEqual(counters["inductor"]["fxgraph_cache_hit"], 0)
329328
self.assertEqual(counters["inductor"]["fxgraph_lookup_write_file"], 0)
330329

331-
@requires_gpu
332-
@torch._inductor.config.patch(fx_graph_cache=False)
333-
def test_config_read_in_backwards(self):
334-
@torch.compile
335-
def f(x, y):
336-
z = x @ y
337-
return z.sin().sum()
338-
339-
called = False
340-
341-
def my_pass(graph):
342-
nonlocal called
343-
called = True
344-
345-
x, y = (
346-
torch.randn(3, 3, device=GPU_TYPE, requires_grad=True),
347-
torch.randn(3, 3, device=GPU_TYPE),
348-
)
349-
z = f(x, y)
350-
z.backward()
351-
self.assertFalse(called)
352-
torch._dynamo.reset()
353-
z = f(x, y)
354-
with torch._inductor.config.patch(post_grad_custom_pre_pass=my_pass):
355-
z.backward()
356-
357-
self.assertTrue(called)
358-
359-
called = False
360-
torch._dynamo.reset()
361-
z = f(x, y)
362-
with torch._inductor.config.patch(post_grad_custom_pre_pass=my_pass):
363-
torch.autograd.grad(z, x)
364-
self.assertTrue(called)
365-
366-
@torch._inductor.config.patch(fx_graph_cache=False)
367-
def test_config_read_in_grad_fn(self):
368-
@torch.compile
369-
def f(x, y):
370-
z = x @ y
371-
return z.sin().sum()
372-
373-
called = False
374-
375-
def my_pass(graph):
376-
nonlocal called
377-
called = True
378-
379-
x, y = (
380-
torch.randn(3, 3, requires_grad=True),
381-
torch.randn(3, 3),
382-
)
383-
384-
with torch._inductor.config.patch(post_grad_custom_pre_pass=my_pass):
385-
z = f(x, y)
386-
self.assertTrue(called)
387-
388-
# Make sure the context gets cleared after forward pass
389-
called = False
390-
z.grad_fn.apply(torch.tensor(0))
391-
self.assertFalse(called)
392-
393330

394331
if __name__ == "__main__":
395332
run_tests()

‎test/inductor/test_torchinductor.py‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
# Owner(s): ["module: inductor"]
22
# ruff: noqa: F841
33
import contextlib
4-
import contextvars
54
import copy
65
import dataclasses
76
import functools
@@ -5414,9 +5413,8 @@ def run_weights_sharing_model(m, inp):
54145413
threads = []
54155414
compiled_m = torch.compile(model)
54165415
for _ in range(1, numb_instance + 1):
5417-
ctx = contextvars.copy_context()
54185416
thread = threading.Thread(
5419-
target=ctx.run, args=(run_weights_sharing_model, compiled_m, inp)
5417+
target=run_weights_sharing_model, args=(compiled_m, inp)
54205418
)
54215419
threads.append(thread)
54225420
thread.start()

‎test/test_utils_config_module.py‎

Lines changed: 5 additions & 86 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ class TestConfigModule(TestCase):
2323
def tearDown(self):
2424
# Config changes get persisted between test cases
2525
for k in config._config:
26-
config._config[k].user_override.set(_UNSET_SENTINEL)
26+
config._config[k].user_override = _UNSET_SENTINEL
2727
config._hash_digest = None
2828
# Reset deprecation warning flags
2929
for k in config._config:
@@ -86,7 +86,7 @@ def test_none_override_semantics(self):
8686
config.e_bool = None
8787
self.assertIsNone(config.e_bool)
8888
for k in config._config:
89-
config._config[k].user_override.set(_UNSET_SENTINEL)
89+
config._config[k].user_override = _UNSET_SENTINEL
9090

9191
def test_reference_semantics(self):
9292
config.e_list.append(2)
@@ -157,42 +157,6 @@ def test_save_config(self):
157157
self.assertTrue(config.e_bool)
158158
self.assertFalse(config.e_ignored)
159159

160-
def test_save_config_with_patch(self):
161-
self.assertTrue(config.e_bool)
162-
with config.patch(e_bool=False):
163-
p = config.save_config()
164-
self.assertDictEqual(
165-
pickle.loads(p),
166-
{
167-
"_cache_config_ignore_prefix": ["magic_cache_config"],
168-
"e_bool": False,
169-
"e_dict": {1: 2},
170-
"e_float": 1.0,
171-
"e_int": 1,
172-
"e_list": [1],
173-
"e_none": None,
174-
"e_set": {1},
175-
"e_string": "string",
176-
"e_tuple": (1,),
177-
"nested.e_bool": True,
178-
"_e_ignored": True,
179-
"e_compile_ignored": True,
180-
"magic_cache_config_ignored": True,
181-
"_save_config_ignore": ["e_ignored"],
182-
"e_config": True,
183-
"e_jk": True,
184-
"e_jk_false": False,
185-
"e_env_default": True,
186-
"e_env_default_FALSE": False,
187-
"e_env_default_str": "1234",
188-
"e_env_default_str_empty": "",
189-
"e_env_force": True,
190-
"e_optional": True,
191-
"e_deprecated": True,
192-
"e_not_deprecated": False,
193-
},
194-
)
195-
196160
def test_save_config_portable(self):
197161
with warnings.catch_warnings():
198162
warnings.filterwarnings("ignore", category=FutureWarning)
@@ -444,7 +408,7 @@ def test_bad_jk_type(self):
444408
AssertionError,
445409
msg="AssertionError: justknobs only support booleans, thisisnotvalid is not a boolean",
446410
):
447-
_ConfigEntry(Config(default="bad", justknob="fake_knob"), "test")
411+
_ConfigEntry(Config(default="bad", justknob="fake_knob"))
448412

449413
def test_alias(self):
450414
self.assertFalse(config2.e_aliasing_bool)
@@ -464,15 +428,13 @@ def test_reference_is_default(self):
464428
def test_invalid_config_int(self):
465429
with self.assertRaises(AssertionError):
466430
_ConfigEntry(
467-
Config(default=2, env_name_default="FAKE_DISABLE", value_type=int),
468-
"test",
431+
Config(default=2, env_name_default="FAKE_DISABLE", value_type=int)
469432
)
470433

471434
def test_invalid_config_float(self):
472435
with self.assertRaises(AssertionError):
473436
_ConfigEntry(
474-
Config(default=2, env_name_force="FAKE_DISABLE", value_type=float),
475-
"test",
437+
Config(default=2, env_name_force="FAKE_DISABLE", value_type=float)
476438
)
477439

478440
def test_deprecated_config(self):
@@ -560,49 +522,6 @@ def fn(x):
560522
f"Unexpected config deprecation warnings: {[str(x.message) for x in deprecation_warnings]}",
561523
)
562524

563-
def test_patch_then_global(self):
564-
self.assertTrue(config.e_bool)
565-
with config.patch(e_bool=False):
566-
self.assertFalse(config.e_bool)
567-
568-
config.e_bool = False
569-
self.assertFalse(config.e_bool)
570-
571-
def test_is_default_patch(self):
572-
self.assertTrue(config.e_bool)
573-
with config.patch(e_bool=False):
574-
self.assertFalse(config._is_default("e_bool"))
575-
576-
def test_dict_patch(self):
577-
self.assertTrue(config.e_bool)
578-
with config.patch(e_bool=False):
579-
d = config._get_dict()
580-
self.assertFalse(d["e_bool"])
581-
582-
def test_set_in_patch(self):
583-
self.assertEqual(config.e_int, 1)
584-
with config.patch(e_int=2):
585-
self.assertEqual(config.e_int, 2)
586-
config.e_int = 3
587-
self.assertEqual(config.e_int, 3)
588-
# Exiting the patch resets e_int to 1, losing the fact that we explicitly set it to 4 - see ConfigModule._do_setattr
589-
self.assertEqual(config.e_int, 1)
590-
591-
def test_nested_patch(self):
592-
self.assertEqual(config.e_int, 1)
593-
self.assertEqual(config.e_string, "string")
594-
with config.patch(e_int=2, e_string="inner"):
595-
self.assertEqual(config.e_int, 2)
596-
self.assertEqual(config.e_string, "inner")
597-
with config.patch(e_int=3):
598-
self.assertEqual(config.e_int, 3)
599-
self.assertEqual(config.e_string, "inner")
600-
config.e_int = 4
601-
self.assertEqual(config.e_int, 4)
602-
self.assertEqual(config.e_int, 2)
603-
self.assertEqual(config.e_int, 1)
604-
self.assertEqual(config.e_string, "string")
605-
606525

607526
if __name__ == "__main__":
608527
run_tests()

‎torch/_functorch/_aot_autograd/runtime_wrappers.py‎

Lines changed: 0 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2842,11 +2842,6 @@ def impl_fn(double_ctx: Any = None) -> Any:
28422842
out,
28432843
)
28442844

2845-
if (
2846-
torch._C._is_key_in_tls("context")
2847-
and (config_ctx := torch._C._get_obj_in_tls("context")) is not None
2848-
):
2849-
impl_fn = functools.partial(config_ctx.run, impl_fn)
28502845
needs_grad = torch.is_grad_enabled() and any(
28512846
t.requires_grad for t in all_args if isinstance(t, torch.Tensor)
28522847
)

‎torch/_inductor/autotune_process.py‎

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@
22
from __future__ import annotations
33

44
import atexit
5-
import contextvars
65
import ctypes
76
import dataclasses
87
import functools
@@ -1335,9 +1334,6 @@ def get_instance(cls) -> PrecompileThreadPool:
13351334
return cls._instance
13361335

13371336
def submit(self, fn, *args, **kwargs):
1338-
ctx = contextvars.copy_context()
1339-
# Need to copy context so workers have access to the correct config settings
1340-
fn = functools.partial(ctx.run, fn)
13411337
return self._executor.submit(fn, *args, **kwargs)
13421338

13431339
def _shutdown(self, wait: bool = False):

‎torch/autograd/graph.py‎

Lines changed: 0 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
11
import abc
22
import contextlib
3-
import contextvars
43
import functools
54
import logging
65
import threading
@@ -874,15 +873,10 @@ def _engine_run_backward(
874873
attach_logging_hooks = log.getEffectiveLevel() <= logging.DEBUG
875874
if attach_logging_hooks:
876875
unregister_hooks = _register_logging_hooks_on_whole_graph(t_outputs)
877-
878-
# Need to save the context so compiler config will be visible in device threads
879-
torch._C._stash_obj_in_tls("context", contextvars.copy_context())
880-
881876
try:
882877
return Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass
883878
t_outputs, *args, **kwargs
884879
) # Calls into the C++ engine to run the backward pass
885880
finally:
886881
if attach_logging_hooks:
887882
unregister_hooks() # type: ignore[possibly-undefined]
888-
torch._C._stash_obj_in_tls("context", None)

‎torch/csrc/Module.cpp‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2778,7 +2778,6 @@ Call this whenever a new thread is created in order to propagate values from
27782778

27792779
py_module.def(
27802780
"_stash_obj_in_tls", [](const std::string& key, py::handle arg) {
2781-
Py_INCREF(arg.ptr());
27822781
at::impl::ThreadLocalPythonObjects::get_state().set(
27832782
key,
27842783
std::make_shared<c10::SafePyObject>(arg.ptr(), getPyInterpreter()));

0 commit comments

Comments
 (0)