Repository: pytorch/torchdynamo

Issue:
Make torchdynamo not import third party package in `skipfiles.py`
@xuzhao9 in https://github.com/facebookresearch/torchdynamo/issues/107#issuecomment-1095681515 found that the following line makes alexnet 18% slower: 

https://github.com/jansel/torchdynamo/blob/bf90b8cdbacf35944fa8c12185b1823dc5cb90bb/torchdynamo/skipfiles.py#L123

It seems importing:  "networkx", "omegaconf", "onnx", "pandas",  and "sklearn" cause performance issues.

TorchDynamo is only importing these modules to find the filename, which is also a bit wasteful.  We should rewrite `skipfiles.py` to use [find_spec](https://docs.python.org/3/library/importlib.html#importlib.abc.PathEntryFinder.find_spec) instead, so we don't need to import unused packages.

Also, I think we can cut down the list of modules in skipfiles dramatically.  Most of those were added when TorchDynamo didn't automatically skip backends and supported much less of python, so likely many (most?) can be removed.



Discussion:
Can we skip tracing all non-pytorch source code, instead of enumerating all common third party libraries?
We could have some sort of heuristic to decide if we should analyze a given frame.  E.g. if neither args nor globals contain `torch.*` objects, just skip the frame.
For more on skipping see https://github.com/facebookresearch/torchdynamo/issues/69
Another annoying thing about this is that it makes import times significantly longer.

For example, we have a `skip_dynamo` decorator in AOTAutograd. These imports mean that whenever we import AOTAutograd, we now also need to import {tensorflow, onnx, etc.}
Reopening this issue since we are reverting #153 

Resolving patch:
diff --git a/torchdynamo/skipfiles.py b/torchdynamo/skipfiles.py
--- a/torchdynamo/skipfiles.py
+++ b/torchdynamo/skipfiles.py
@@ -32,8 +32,12 @@
 import torch
 
 
+def _strip_init_py(s):
+    return re.sub(r"__init__.py$", "", s)
+
+
 def _module_dir(m: types.ModuleType):
-    return re.sub(r"__init__.py$", "", m.__file__)
+    return _strip_init_py(m.__file__)
 
 
 SKIP_DIRS = [
@@ -79,22 +83,32 @@ def _module_dir(m: types.ModuleType):
         _weakrefset,
     )
 ]
-SKIP_DIRS_RE = None  # set in add() below
 FILENAME_ALLOWLIST = {
     torch.nn.Sequential.__init__.__code__.co_filename,
 }
+SKIP_DIRS_RE = None
 
 
-def add(module: types.ModuleType):
-    assert isinstance(module, types.ModuleType)
+def _recompile_re():
     global SKIP_DIRS_RE
-    name = module.__file__
-    if name is None:
-        return
-    SKIP_DIRS.append(_module_dir(module))
     SKIP_DIRS_RE = re.compile(f"^({'|'.join(map(re.escape, SKIP_DIRS))})")
 
 
+def add(import_name: str):
+    if isinstance(import_name, types.ModuleType):
+        return add(import_name.__name__)
+    assert isinstance(import_name, str)
+    module_spec = importlib.util.find_spec(import_name)
+    if not module_spec:
+        return
+    origin = module_spec.origin
+    if origin is None:
+        return
+    global SKIP_DIRS_RE
+    SKIP_DIRS.append(_strip_init_py(origin))
+    _recompile_re()
+
+
 def check(filename, allow_torch=False):
     """Should skip this file?"""
     if filename is None:
@@ -127,10 +141,9 @@ def check(filename, allow_torch=False):
     "tvm",
     "fx2trt_oss",
 ):
-    try:
-        add(importlib.import_module(_name))
-    except (ImportError, TypeError):
-        pass
+    add(_name)
+
+_recompile_re()
 
 
 def is_torch_inline_allowed(filename):


Test patch: