xref: /aosp_15_r20/external/pytorch/tools/coverage_plugins_package/src/coverage_plugins/jit_plugin.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45) 
1"""
2This coverage plug-in attempts to cover JIT'd functions and methods that were previously missed in code coverage. Any
3function and method that was passed through/decorated with torch.jit.script or torch.jit.script_method should now be
4marked covered when coverage is run with this plug-in.
5
6DISCLAIMER: note that this will mark the entire JIT'd function/method as covered without seeking proof that the
7compiled code has been executed. This means that even if the code chunk is merely compiled and not run, it will get
8marked as covered.
9"""
10
11from inspect import (
12    getsourcefile,
13    getsourcelines,
14    isclass,
15    iscode,
16    isfunction,
17    ismethod,
18    ismodule,
19)
20from time import time
21from typing import Any
22
23from coverage import CoverageData, CoveragePlugin  # type: ignore[import]
24
25
26# All coverage stats resulting from this plug-in will be in a separate .coverage file that should be merged later with
27# `coverage combine`. The convention seems to be .coverage.dotted.suffix based on the following link:
28# https://coverage.readthedocs.io/en/coverage-5.5/cmd.html#combining-data-files-coverage-combine
29cov_data = CoverageData(basename=f".coverage.jit.{time()}")
30
31
32def is_not_builtin_class(obj: Any) -> bool:
33    return isclass(obj) and not type(obj).__module__ == "builtins"
34
35
36class JitPlugin(CoveragePlugin):  # type: ignore[misc, no-any-unimported]
37    """
38    dynamic_context is an overridden function that gives us access to every frame run during the coverage process. We
39    look for when the function being run is `should_drop`, as all functions that get passed into `should_drop` will be
40    compiled and thus should be marked as covered.
41    """
42
43    def dynamic_context(self, frame: Any) -> None:
44        if frame.f_code.co_name == "should_drop":
45            obj = frame.f_locals["fn"]
46            # The many conditions in the if statement below are based on the accepted arguments to getsourcefile. Based
47            # on its documentation (https://docs.python.org/3/library/inspect.html#inspect.getsourcefile), the argument
48            # must be a module, class, method, function, traceback, frame, or code object AND it cannot be a built-in
49            # module, class, or function.
50            # Currently, we DO NOT include tracebacks or frames as they should not be JIT'd, and we have not checked for
51            # built-in modules or functions as those do not seem to be JIT'd either.
52            if (
53                is_not_builtin_class(obj)
54                or ismodule(obj)
55                or ismethod(obj)
56                or isfunction(obj)
57                or iscode(obj)
58            ):
59                filename = getsourcefile(obj)
60                # We don't want to report for filename = None
61                if filename:
62                    # TODO: Because torch.jit._IgnoreContextManager relies on Python's `exec` method
63                    # which doesn't generate source codelines, getsourcelines(obj) fails. For now,
64                    # we just ignore the exception until we figure out a better way to
65                    # implement torch.jit._IgnoreContextManager.
66                    try:
67                        sourcelines, starting_lineno = getsourcelines(obj)
68                    except OSError:
69                        pass
70                    else:
71                        line_data = {
72                            filename: range(
73                                starting_lineno, starting_lineno + len(sourcelines)
74                            )
75                        }
76                        cov_data.add_lines(line_data)
77        super().dynamic_context(frame)
78
79
80def coverage_init(reg: Any, options: Any) -> None:
81    reg.add_dynamic_context(JitPlugin())
82