Introducing JitPlugin (#56708)

Summary:
This PR is step 1 to covering JIT'd methods and functions. Step 2 (using it in CI) is here: https://github.com/pytorch/pytorch/issues/56310.

1. This PR introduces a package `coverage_plugins` that hosts JITPlugin.
2. We also bring in a `.coveragerc` file that is used in CI to omit the files we don't want to report on (e.g., temporary directories or test or utils.)

**Disclaimer: This PR does NOT use the plug-in. Nothing should change as a result.**

Pull Request resolved: https://github.com/pytorch/pytorch/pull/56708

Test Plan:
CI. Coverage should not go down.

If you're interested in testing this plug-in locally, you should:
`pip install -e tools/coverage_plugins_package` from the root directory.
Add the following lines to `.coveragerc` under `[run]`
```
plugins =
    coverage_plugins.jit_plugin
```
And then try:
`coverage run test/test_jit.py TestAsync.test_async_script_no_script_mod`

You should see `.coverage.jit` show up at the end. You can then run `coverage combine --append` and `coverage debug data` to see that some files in `torch/jit` are covered.

Reviewed By: samestep

Differential Revision: D27945570

Pulled By: janeyx99

fbshipit-source-id: 78732940fcb498d5ec37d4075c4e7e08e96a8d55

.coveragerc

@@ -0,0 +1,13 @@
+[run]
+omit =
+    */tmp*
+    */Temp/*
+    */usr/local/lib*
+    *test/*
+
+[report]
+omit =
+    */tmp*
+    */Temp/*
+    */usr/local/lib*
+    *test/*

.gitignore

@@ -55,6 +55,7 @@ test/generated_type_hints_smoketest.py
 test/htmlcov
 test/cpp_extensions/install/
 third_party/build/
+tools/coverage_plugins_package/pip-wheel-metadata/
 tools/shared/_utils_internal.py
 tools/fast_nvcc/wrap_nvcc.sh
 tools/fast_nvcc/tmp/

.jenkins/pytorch/test.sh

@@ -25,6 +25,7 @@ fi
 
 if [[ "$BUILD_ENVIRONMENT" == *coverage* ]]; then
   export PYTORCH_COLLECT_COVERAGE=1
+  export COVERAGE_RCFILE="$PWD/.coveragerc" # coverage config file needed for plug-ins and settings to work
 fi
 
 if [[ "$BUILD_ENVIRONMENT" == *cuda* ]]; then

.jenkins/pytorch/win-test-helpers/setup_pytorch_env.bat

@@ -39,7 +39,12 @@ if %errorlevel% neq 0 ( exit /b %errorlevel% )
 popd
 
 :: The version is fixed to avoid flakiness: https://github.com/pytorch/pytorch/issues/31136
-pip install "ninja==1.10.0.post1" future "hypothesis==4.53.2" "librosa>=0.6.2" psutil pillow unittest-xml-reporting pytest coverage
+pip install "ninja==1.10.0.post1" future "hypothesis==4.53.2" "librosa>=0.6.2" psutil pillow unittest-xml-reporting pytest
+
+:: TODO: All sharded configs run coverage. We should change that to be only one config, but right now we will just
+:: install coverage everywhere. Tracked: https://github.com/pytorch/pytorch/issues/56264
+python -mpip install coverage
+
 if %errorlevel% neq 0 ( exit /b %errorlevel% )
 
 set DISTUTILS_USE_SDK=1

.jenkins/pytorch/win-test.sh

@@ -71,6 +71,7 @@ run_tests() {
         "$SCRIPT_HELPERS_DIR"/test_libtorch.bat
     else
         export PYTORCH_COLLECT_COVERAGE=1
+        export COVERAGE_RCFILE="$PWD/.coveragerc" # coverage config file needed for plug-ins and settings to work
         if [[ "${JOB_BASE_NAME}" == *-test1 ]]; then
             "$SCRIPT_HELPERS_DIR"/test_python_first_shard.bat "$DETERMINE_FROM"
             "$SCRIPT_HELPERS_DIR"/test_libtorch.bat

tools/coverage_plugins_package/README.md

@@ -0,0 +1,3 @@
+# What is this?
+
+This folder hosts a minimal package for coverage plug-ins. Currently, the only plug-in is a JIT plug-in that helps coverage mark functions and methods passed through `torch.jit.script` and `torch.jit.script_method` as covered code.

tools/coverage_plugins_package/pyproject.toml

@@ -0,0 +1,6 @@
+[build-system]
+requires = [
+    "setuptools>=42",
+    "wheel"
+]
+build-backend = "setuptools.build_meta"

tools/coverage_plugins_package/setup.py

@@ -0,0 +1,26 @@
+import setuptools
+
+with open("README.md", "r", encoding="utf-8") as fh:
+    long_description = fh.read()
+
+setuptools.setup(
+    name="coverage-plugins",
+    version="0.0.1",
+    author='PyTorch Team',
+    author_email='packages@pytorch.org',
+    description="plug-in to coverage for PyTorch JIT",
+    long_description=long_description,
+    long_description_content_type="text/markdown",
+    url="https://github.com/pytorch/pytorch",
+    project_urls={
+        "Bug Tracker": "https://github.com/pytorch/pytorch/issues",
+    },
+    classifiers=[
+        "Programming Language :: Python :: 3",
+        "License :: OSI Approved :: MIT License",
+        "Operating System :: OS Independent",
+    ],
+    package_dir={"": "src"},
+    packages=setuptools.find_packages(where="src"),
+    python_requires=">=3.6",
+)

tools/coverage_plugins_package/src/coverage_plugins/jit_plugin.py

@@ -0,0 +1,47 @@
+'''
+This coverage plug-in attempts to cover JIT'd functions and methods that were previously missed in code coverage. Any
+function and method that was passed through/decorated with torch.jit.script or torch.jit.script_method should now be
+marked covered when coverage is run with this plug-in.
+
+DISCLAIMER: note that this will mark the entire JIT'd function/method as covered without seeking proof that the
+compiled code has been executed. This means that even if the code chunk is merely compiled and not run, it will get
+marked as covered.
+'''
+
+from coverage import CoveragePlugin, CoverageData
+from inspect import ismodule, isclass, ismethod, isfunction, iscode, getsourcefile, getsourcelines
+
+# All coverage stats resulting from this plug-in will be in a separate .coverage file that should be merged later with
+# `coverage combine`. The convention seems to be .coverage.dotted.suffix based on the following link:
+# https://coverage.readthedocs.io/en/coverage-5.5/cmd.html#combining-data-files-coverage-combine
+cov_data = CoverageData(basename='.coverage.jit')
+
+
+def is_not_builtin_class(obj):
+    return isclass(obj) and not type(obj).__module__ == 'builtins'
+
+
+class JitPlugin(CoveragePlugin):
+    '''
+    dynamic_context is an overridden function that gives us access to every frame run during the coverage process. We
+    look for when the function being run is `should_drop`, as all functions that get passed into `should_drop` will be
+    compiled and thus should be marked as covered.
+    '''
+    def dynamic_context(self, frame):
+        if frame.f_code.co_name == 'should_drop':
+            obj = frame.f_locals['fn']
+            # The many conditions in the if statement below are based on the accepted arguments to getsourcefile. Based
+            # on its documentation (https://docs.python.org/3/library/inspect.html#inspect.getsourcefile), the argument
+            # must be a module, class, method, function, traceback, frame, or code object AND it cannot be a built-in
+            # module, class, or function.
+            # Currently, we DO NOT include tracebacks or frames as they should not be JIT'd, and we have not checked for
+            # built-in modules or functions as those do not seem to be JIT'd either.
+            if is_not_builtin_class(obj) or ismodule(obj) or ismethod(obj) or isfunction(obj) or iscode(obj):
+                filename = getsourcefile(obj)
+                sourcelines, starting_lineno = getsourcelines(obj)
+                line_data = {filename: range(starting_lineno, starting_lineno + len(sourcelines))}
+                cov_data.add_lines(line_data)
+        super().dynamic_context(frame)
+
+def coverage_init(reg, options):
+    reg.add_dynamic_context(JitPlugin())