[RFC] Emit better Telemetry in PyTorch
Labels: feature, module: logging, triaged, oncall: pt2

### 🚀 The feature, motivation and pitch

**Summary**

The existing PyTorch 2.0 logging needs to be enhanced to emit better Telemetry. For this purpose, this document proposes the following changes:

1. Emit function parameters along with the function name using [_log_api_usage_once](https://github.com/pytorch/pytorch/blob/7b47cd0a6c206ef1de1ad74881942352089ddc72/torch/csrc/Module.cpp#L1420)
2. Control dumping dynamo compile metrics using compile_times via an environment variable PYTORCH_DYNAMO_COMPILE_TIMES_PERIOD.

**Motivation**

Enhancing the logging will highlight the features actively used by the Pytorch Users. Controlling the dumping of dynamo compile metrics using an environment variable will enable the feature without modifying the User training/inference script.

**Proposed Implementation**

**PyTorch API level logging**

PyTorch uses Google’s logging library ([glog](https://github.com/google/glog#severity-levels)) for logging in C++. The environment variable [PYTORCH_API_USAGE_STDERR](https://github.com/pytorch/pytorch/blob/main/c10/util/Logging.cpp#L93)
can be used to log what APIs are used. This corresponds to the API  [_log_api_usage_once](https://github.com/pytorch/pytorch/blob/7b47cd0a6c206ef1de1ad74881942352089ddc72/torch/csrc/Module.cpp#L1420) in python and C10_LOG_API_USAGE_ONCE in C++.  The current implementation only provides insight into which APIs were used, leaving out their arguments. 

We would like to capture the parameters of APIs that we want to log without compromising the users’s IP and privacy. We will modify the below APIs logging to include their parameters. We will only include the arguments of type bool, enum and constrained strings.


```
def format_api(api_name, param_dict):
    # format_string will be of the format <API><seperator:-><param_seperator:,>
    format_string = os.environ['PYTORCH_LOG_API_FORMAT']
    if format_string is None:
        format_string = <API><seprator: ><param_seperator:,>
    separator_start = "<seperator:"
    separator_end = ">"
    param_separator_start = "<param_seperator:"
    param_separator_end = ">"

    #default values
    separator = " " 
    param_separator = ","
    
    separator_index_start = format_string.find(separator_start)
    separator_index_end = format_string.find(separator_end, separator_index_start)
    param_separator_index_start = format_string.find(param_separator_start)
    param_separator_index_end = format_string.find(param_separator_end, param_separator_index_start)

    if separator_index_start != -1 and separator_index_end != -1:
        separator = format_string[separator_index_start + len(separator_start):separator_index_end]
    
    if param_separator_index_start != -1 and param_separator_index_end != -1:
        param_separator = format_string[param_separator_index_start + len(param_separator_start):param_separator_index_end]
    
    formatted_api = format_string.replace("<API>", api_name)
    formatted_api = formatted_api.replace(f"{separator_start}{separator}{separator_end}", separator)
    formatted_api = formatted_api.replace(f"{param_separator_start}{param_separator}{param_separator_end}", param_separator)

    param_string = param_separator.join([f"{key}{param_separator}{value}" for key, value in param_dict.items()])

    formatted_api = f"{formatted_api}{separator}{param_string}"
    return formatted_api
```


We can use a custom format_api function that takes API name and parameters dict and provide the formatted  string as output. The format string can be set using an environment variable **PYTORCH_LOG_API_FORMAT**.



**[torch.compile](https://github.com/pytorch/pytorch/blob/7b47cd0a6c206ef1de1ad74881942352089ddc72/torch/__init__.py#L1574)**
```
param_dict = {}

if(isinstance(backend, str)) {
  param_dict[“backend”] = backend
}
param_dict[“fullgraph”] = str(fullgraph)
param_dict[“dynamic”] = str(dynamic)
param_dict[“mode”] = str(mode)
param_dict[“disable”] = str(disable)
C._log_api_usage_once(format_api(torch.compile, param_dict))
param_dict = {sync_module_states:str(sync_module_states), forward_prefetch:str(forward_prefetch)}
torch._C._log_api_usage_once(format_api(torch.distributed.fully_shard, param_dict))
```

**[torch._dynamo.optimize](https://github.com/pytorch/pytorch/blob/7b47cd0a6c206ef1de1ad74881942352089ddc72/torch/_dynamo/eval_frame.py#LL532C35-L532C57)**

```
param_dict = {nopython:str(nopython), disable:str(disable), dynamic:str(dynamic)}
torch._C._log_api_usage_once(format_api(torch._dynamo.optimize, param_dict)
```

**[torch._dynamo.export](https://github.com/pytorch/pytorch/blob/7b47cd0a6c206ef1de1ad74881942352089ddc72/torch/_dynamo/eval_frame.py#L832)**
```
param_dict = {aten_graph:str(aten_graph), pre_autograd:str(pre_autograd), tracing_mode:tracing_mode}, assume_static_by_default:str(assume_static_by_default), functionalize:str(functionalize)}
torch._C._log_api_usage_once(format_api(torch._dynamo.export, param_dict))
```



**Compilation runtime**

TorchDynamo uses a function that can be used as a decorator to [capture](https://github.com/pytorch/pytorch/blob/main/torch/_dynamo/utils.py#L171) the breakdown of compilation time in terms of graph capture and backend compilation times. However, this breakdown is only made available with an explicit call to a [summary API](https://github.com/pytorch/pytorch/blob/main/torch/_dynamo/utils.py#L201) at the end of training/serving. As a result, this information is not available without customer intervention. 

The API can be automatically invoked after training/inference. This can be achieved by registering these APIs using [atexit](https://docs.python.org/3/library/atexit.html) exit handler in the torch._dynamo.__init__.py.

`atexit.register(lambda: compile_times(str, True))`

We can also dump this compile_times info periodically using a [Timeloop](https://pypi.org/project/timeloop/) library. We can periodically invoke compile_times API with a defined interval set using  an environment variable **PYTORCH_DYNAMO_COMPILE_TIMES_PERIOD** in seconds. This will help live hosted endpoints where the script never exits. The default value will be 0 and gets updated based on user provided value.


```
#Pseudo code
import os
import time
from timeloop import Timeloop
from datetime import timedelta
import torch._dynamo.utils

tl = Timeloop()

period = os.environ['PYTORCH_PRINT_DYNAMO_COMPILE_TIMES_PERIOD']
if not isinstance(period, int):
    period = 0;
@tl.job(interval=timedelta(seconds=period))
def check_and_invoke_compile_times():
    torch._dynamo.utils.compile_times()
    
if __name__ == "__main__":
    if period > 0:
      tl.start(block=True)

```
To prevent redundant dumping of previously dumped metrics in the previous compile_times invocation, it is necessary to maintain an index indicating the point until which the compile times of specific APIs have already been logged. This can be achieved by enhancing the existing compilation_metrics dictionary to include a sub-dictionary for each API. The sub-dictionary will contain the values under the key 'values' and the new index that is pending logging under the key 'log_index'.
By implementing this modification, we can effectively track the progress of logging the compile times for individual APIs and avoid re-dumping already logged metrics.

We will modify the compile_times API as follows:

```
#Pseudo code
def dynamo_timed(original_function=None, phase_name=None):
    def dynamo_timed_inner(func):
        @wraps(func)
        def time_wrapper(*args, **kwargs):
            key = func.__qualname__
            if key not in compilation_metrics:
                compilation_metrics[key] = {values:[], index:0}
            with torch.profiler.record_function(f"{key} (dynamo_timed)"):
                t0 = time.time()
                r = func(*args, **kwargs)
                time_spent = time.time() - t0
            # print(f"Dynamo timer: key={key}, latency={latency:.2f} sec")
            compilation_metrics[key]["values"].append(time_spent)
            if phase_name:
                frame_key = str(curr_frame)
                if frame_key not in frame_phase_timing:
                    frame_phase_timing[frame_key] = {}
                assert (
                    phase_name not in frame_phase_timing[frame_key]
                ), f"Duplicate phase name {phase_name} for frame {frame_key}"
                frame_phase_timing[frame_key][phase_name] = time_spent
            return r

        return time_wrapper

    if original_function:
        return dynamo_timed_inner(original_function)
    return dynamo_timed_inner


def compile_times(repr="str", aggregate=False):
    """
    Get metrics about torchdynamo frontend/backend compilation times.

    Accumulates information from functions tagged with `@dynamo_timed`.

    repr='str' returns a printable string for user interaction, and 'csv'
    returns headers, rows which can be logged for output

    aggregate causes values from multiple compilations (e.g. split graphs)
    to be accumulated into one value.  If false, expect more than one value
    per metric.
    """

    def fmt_fn(values, item_fn=lambda x: x):
        if aggregate:
            return item_fn(sum(values))
        return ", ".join(map(item_fn, values))

    if repr == "str":
        rows = [
            (k, fmt_fn(compilation_metrics[k]["values"] if aggregate else compilation_metrics[k]["values"][compilation_metrics[k]["index"]:], item_fn=lambda x: f"{x:.4f}"))
            for k in compilation_metrics
        ]
        out = "TorchDynamo compilation metrics:\n"
        out += tabulate(rows, headers=("Function", "Runtimes (s)"))
        return out
    elif repr == "csv":
        values = [
            fmt_fn(v, item_fn=lambda x: f"{x:.6f}")
            for v in compilation_metrics.values()
        ]
        headers = list(compilation_metrics.keys())
        return headers, values
```

**Metrics** 

This feature will enable PyTorch developers to understand the % usage of different PyTorch features. 



**Drawbacks**

This will not break any existing features.

**Alternatives**

NA

### Alternatives

NA

### Additional context

NA

cc @ezyang @msaroufim @wconstab @bdhirsh @anijain2305