code: def wrap_fn(
      fn,
      cls,
      method_name,
      trace,
      stack_level=1,
      pre_fn=None,
      post_fn=None,
      is_class_method=None,
  ):
      """Wrap a function so that its execution can be traced and its args and return values modified."""
      class_name = cls.__qualname__
  
      @functools.wraps(fn)
      def fn_with_tracing(self, *args: Any, **kwargs: Any):
          if class_name not in trace:
              trace[class_name] = {}
  
          self_id = id(self)
          stack = inspect.stack()
          frame = stack[stack_level]
          frame_id = id(frame)
          stack_len = len(stack) - 1
  
          if self_id not in trace[class_name]:
              trace[class_name][self_id] = {}
  
          if method_name not in trace[class_name][self_id]:
              trace[class_name][self_id][method_name] = {}
  
          if frame_id not in trace[class_name][self_id][method_name]:
              trace[class_name][self_id][method_name][frame_id] = {}
  
          trace_entry = trace[class_name][self_id][method_name][
              frame_id
          ]
  
          if pre_fn:
              # If a pre_fn is specified, it can both record information
              # in a trace, as well as return modified args and kwargs
              # that will be provided to the actual fn being wrappped
              pre_trace, args, kwargs = pre_fn(self, *args, **kwargs)
              trace_entry["pre"] = pre_trace
  
          # We record the invocation and the calling location in the trace
          trace_entry["frame"] = {
              "filename": frame.filename,
              "lineno": frame.lineno,
              "function": frame.function,
              "depth": stack_len,
          }
  
          # we cache the dfeault parameters used during the function call
          trace_entry["default_args"] = get_default_args(fn)
  
          # we cache also the parameters used during the function call
          trace_entry["call_args"] = kwargs
  
          trace_entry["call"] = {"start": time.time_ns()}
  
          ret = (
              fn(self, *args, **kwargs)
              if not is_class_method
              else fn(*args, **kwargs)
          )
  
          trace_entry["call"]["end"] = time.time_ns()
  
          if post_fn:
              # If a post_fn is specified, it can both record information
              # in a trace, as well as modify the value returned from fn
              post_trace, ret = post_fn(self, ret)
              trace_entry["post"] = post_trace
  
          return ret
  
      return fn_with_tracing
  
file_path: pytorch-lightning/src/lightning/app/utilities/tracer.py
imports: import functools
  import inspect
  import runpy
  import sys
  import time
  from pathlib import Path
  from typing import Any, Optional
label_desc: Use image node_id 2 for calling a global function with example usage: wrap_fn(fn, cls, method_name, trace, stack_level, pre_fn, post_fn, is_class_method) and returns: fn_with_tracing, ret