| """Make torch._dynamo's handler registration idempotent. |
| |
| torch/_dynamo/variables/torch.py builds a local `handlers` dict and registers |
| functions into it from a nested closure, raising |
| AssertionError("Handler already registered for {fn}") when a fn turns up twice. |
| The closure is recreated per call, so it cannot be reached by patching module |
| attributes - the source file itself has to change. |
| |
| The assert runs during dynamo tracing, which compressed-tensors triggers from |
| round_to_quantized_type_args while compressing the NVFP4 experts, and it fires |
| with fn == torch._C._cuda_exchangeDevice. Turning the assert into a plain |
| overwrite makes the registration idempotent. |
| |
| install() must run before torch._dynamo is imported, so dwct/__init__.py calls |
| it first and nothing here imports torch (the path comes from find_spec). |
| |
| The module also scans already-imported torch/compressed_tensors/transformers |
| modules for any module-level callable carrying the same assert and arms a |
| meta-path hook so later imports get scanned too. DWCT_PATCH_FORCE=0 disables it. |
| """ |
| from __future__ import annotations |
|
|
| import importlib |
| import importlib.abc |
| import importlib.util |
| import inspect |
| import os |
| import pathlib |
| import re |
| import sys |
|
|
| _installed = False |
| _MARK = "Handler already registered" |
| _MARKER = "dwct: idempotent registration" |
| _REL = "torch/_dynamo/variables/torch.py" |
| _ROOTS = ("torch", "compressed_tensors", "transformers", "spaces") |
|
|
| _BLOCK = re.compile( |
| r"^(?P<ind>[ \t]*)if fn in handlers:\n" |
| r"[ \t]*raise AssertionError\(.*Handler already registered.*\)\n" |
| r"(?P=ind)handlers\[fn\] = handler\n", |
| re.MULTILINE, |
| ) |
| _LINE = re.compile( |
| r"^(?P<ind>[ \t]*)raise AssertionError\(.*Handler already registered.*\)[ \t]*$", |
| re.MULTILINE, |
| ) |
|
|
| _wrapped = {} |
| _seen = set() |
|
|
|
|
| def _torch_file(): |
| spec = None |
| try: |
| spec = importlib.util.find_spec("torch") |
| except Exception: |
| spec = None |
| roots = [] |
| if spec is not None and getattr(spec, "submodule_search_locations", None): |
| roots.extend(list(spec.submodule_search_locations)) |
| roots.extend(sys.path) |
| for root in roots: |
| if not root: |
| continue |
| p = pathlib.Path(root) / "_dynamo" / "variables" / "torch.py" |
| if p.is_file(): |
| return p |
| return None |
|
|
|
|
| def _source_patch(): |
| p = _torch_file() |
| if p is None: |
| print("[patchfix] %s not found" % _REL) |
| return |
| try: |
| src = p.read_text() |
| except Exception as exc: |
| print("[patchfix] read %s failed: %r" % (p, exc)) |
| return |
| if _MARKER in src: |
| print("[patchfix] source already patched: %s" % p) |
| return |
|
|
| m = _BLOCK.search(src) |
| if m: |
| repl = "%shandlers[fn] = handler # %s\n" % (m.group("ind"), _MARKER) |
| src = src[: m.start()] + repl + src[m.end():] |
| how = "block" |
| else: |
| m = _LINE.search(src) |
| if not m: |
| print("[patchfix] assert text not found in %s" % p) |
| return |
| repl = "%spass # %s" % (m.group("ind"), _MARKER) |
| src = src[: m.start()] + repl + src[m.end():] |
| how = "line" |
|
|
| try: |
| p.write_text(src) |
| except Exception as exc: |
| print("[patchfix] write %s failed: %r" % (p, exc)) |
| return |
|
|
| for pyc in (p.parent / "__pycache__").glob("torch.cpython-*.pyc"): |
| try: |
| pyc.unlink() |
| except Exception: |
| pass |
| importlib.invalidate_caches() |
| print("[patchfix] source patched (%s): %s" % (how, p)) |
| if "torch._dynamo.variables.torch" in sys.modules: |
| print("[patchfix] note: module already imported; edit lands next run") |
|
|
|
|
| def _tolerant(fn): |
| def wrapper(*args, **kwargs): |
| try: |
| return fn(*args, **kwargs) |
| except AssertionError as exc: |
| if _MARK in str(exc): |
| print("[patchfix] repeat registration ignored: %s" % exc) |
| return None |
| raise |
|
|
| wrapper.__name__ = getattr(fn, "__name__", "wrapper") |
| wrapper.__qualname__ = getattr(fn, "__qualname__", wrapper.__name__) |
| wrapper._dwct_tolerant = True |
| return wrapper |
|
|
|
|
| def _unwrap(obj): |
| if isinstance(obj, (staticmethod, classmethod)): |
| return obj.__func__ |
| return obj |
|
|
|
|
| def _rewrap(orig, new): |
| if isinstance(orig, staticmethod): |
| return staticmethod(new) |
| if isinstance(orig, classmethod): |
| return classmethod(new) |
| return new |
|
|
|
|
| def _is_guard(obj): |
| if not inspect.isfunction(obj): |
| return False |
| if id(obj) in _seen: |
| return False |
| _seen.add(id(obj)) |
| try: |
| return _MARK in inspect.getsource(obj) |
| except Exception: |
| return False |
|
|
|
|
| def _register(orig, label): |
| if id(orig) in _wrapped: |
| return False |
| _wrapped[id(orig)] = _tolerant(orig) |
| where = "" |
| try: |
| where = " (%s:%d)" % ( |
| inspect.getsourcefile(orig), |
| inspect.getsourcelines(orig)[1], |
| ) |
| except Exception: |
| pass |
| print("[patchfix] guard %s%s" % (label, where)) |
| return True |
|
|
|
|
| def _scan_module(mod): |
| name = getattr(mod, "__name__", "") |
| found = 0 |
| try: |
| items = list(vars(mod).items()) |
| except Exception: |
| return 0 |
| for attr, obj in items: |
| if getattr(obj, "__module__", None) == name and _is_guard(obj): |
| if _register(obj, "%s.%s" % (name, attr)): |
| found += 1 |
| elif inspect.isclass(obj) and getattr(obj, "__module__", None) == name: |
| for mname, meth in list(vars(obj).items()): |
| raw = _unwrap(meth) |
| if _is_guard(raw): |
| if _register(raw, "%s.%s.%s" % (name, attr, mname)): |
| found += 1 |
| return found |
|
|
|
|
| def _repoint(): |
| n = 0 |
| for mod in list(sys.modules.values()): |
| if mod is None: |
| continue |
| try: |
| items = list(vars(mod).items()) |
| except Exception: |
| continue |
| for attr, obj in items: |
| w = _wrapped.get(id(obj)) |
| if w is not None: |
| try: |
| setattr(mod, attr, w) |
| n += 1 |
| except Exception: |
| pass |
| continue |
| if inspect.isclass(obj): |
| for mname, meth in list(vars(obj).items()): |
| w = _wrapped.get(id(_unwrap(meth))) |
| if w is not None: |
| try: |
| setattr(obj, mname, _rewrap(meth, w)) |
| n += 1 |
| except Exception: |
| pass |
| return n |
|
|
|
|
| class _Scanner(importlib.abc.MetaPathFinder): |
| def find_spec(self, fullname, path=None, target=None): |
| if fullname.split(".")[0] not in _ROOTS: |
| return None |
| spec = None |
| for finder in list(sys.meta_path): |
| if finder is self: |
| continue |
| try: |
| spec = finder.find_spec(fullname, path, target) |
| except Exception: |
| spec = None |
| if spec is not None: |
| break |
| if spec is None or getattr(spec, "loader", None) is None: |
| return None |
| loader = spec.loader |
| if getattr(loader, "_dwct_scanned", False): |
| return spec |
| try: |
| orig = loader.exec_module |
| except Exception: |
| return spec |
|
|
| def exec_module(module, _orig=orig): |
| _orig(module) |
| try: |
| if _scan_module(module): |
| _repoint() |
| except Exception as exc: |
| print( |
| "[patchfix] scan %s failed: %r" |
| % (getattr(module, "__name__", "?"), exc) |
| ) |
|
|
| try: |
| loader.exec_module = exec_module |
| loader._dwct_scanned = True |
| except Exception: |
| pass |
| return spec |
|
|
|
|
| def install(): |
| global _installed |
| if _installed or os.environ.get("DWCT_PATCH_FORCE", "1") == "0": |
| return |
| _installed = True |
|
|
| _source_patch() |
|
|
| hits = 0 |
| for mod in list(sys.modules.values()): |
| if mod is None: |
| continue |
| if getattr(mod, "__name__", "").split(".")[0] not in _ROOTS: |
| continue |
| try: |
| hits += _scan_module(mod) |
| except Exception: |
| pass |
| if hits: |
| print("[patchfix] %d module-level guard(s), %d ref(s) re-pointed" |
| % (hits, _repoint())) |
| else: |
| print("[patchfix] no module-level guard; import hook armed") |
|
|
| sys.meta_path.insert(0, _Scanner()) |
|
|
