File size: 8,629 Bytes
a64dea9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 | """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())
|