Repository: pytorch
Issue: [Dynamo] support logging via reorderable_logging_functions
Labels: triaged,oncall: pt2,module: dynamo,dynamo-logging

### 🐛 Describe the bug

As https://github.com/pytorch/pytorch/pull/116106 implemented, some logging functions like print and warnings.warn can be reordered to print at the end without introducing any graph breaks. Can it be extended to support methods in logging? Such as logger.info. 

### Error logs

```
Traceback (most recent call last):
  File "dyn_logging.py", line 22, in <module>
    main()
  File "dyn_logging.py", line 18, in main
    fn(a)
  File "/Users/chenxiny/miniforge3/envs/torch-metal/lib/python3.8/contextlib.py", line 75, in inner
    return func(*args, **kwds)
  File "/Users/chenxiny/miniforge3/envs/torch-metal/lib/python3.8/site-packages/torch/_dynamo/eval_frame.py", line 462, in _fn
    return fn(*args, **kwargs)
  File "/Users/chenxiny/miniforge3/envs/torch-metal/lib/python3.8/site-packages/torch/_dynamo/convert_frame.py", line 1182, in __call__
    return self._torchdynamo_orig_callable(
  File "/Users/chenxiny/miniforge3/envs/torch-metal/lib/python3.8/site-packages/torch/_dynamo/convert_frame.py", line 506, in __call__
    return _compile(
  File "/Users/chenxiny/miniforge3/envs/torch-metal/lib/python3.8/contextlib.py", line 75, in inner
    return func(*args, **kwds)
  File "/Users/chenxiny/miniforge3/envs/torch-metal/lib/python3.8/site-packages/torch/_dynamo/convert_frame.py", line 862, in _compile
    guarded_code = compile_inner(code, one_graph, hooks, transform)
  File "/Users/chenxiny/miniforge3/envs/torch-metal/lib/python3.8/site-packages/torch/_dynamo/utils.py", line 272, in time_wrapper
    r = func(*args, **kwargs)
  File "/Users/chenxiny/miniforge3/envs/torch-metal/lib/python3.8/site-packages/torch/_utils_internal.py", line 85, in wrapper_function
    return StrobelightCompileTimeProfiler.profile_compile_time(
  File "/Users/chenxiny/miniforge3/envs/torch-metal/lib/python3.8/site-packages/torch/_strobelight/compile_time_profiler.py", line 129, in profile_compile_time
    return func(*args, **kwargs)
  File "/Users/chenxiny/miniforge3/envs/torch-metal/lib/python3.8/site-packages/torch/_dynamo/convert_frame.py", line 680, in compile_inner
    out_code = transform_code_object(code, transform)
  File "/Users/chenxiny/miniforge3/envs/torch-metal/lib/python3.8/site-packages/torch/_dynamo/bytecode_transformation.py", line 1322, in transform_code_object
    transformations(instructions, code_options)
  File "/Users/chenxiny/miniforge3/envs/torch-metal/lib/python3.8/site-packages/torch/_dynamo/convert_frame.py", line 201, in _fn
    return fn(*args, **kwargs)
  File "/Users/chenxiny/miniforge3/envs/torch-metal/lib/python3.8/site-packages/torch/_dynamo/convert_frame.py", line 622, in transform
    tracer.run()
  File "/Users/chenxiny/miniforge3/envs/torch-metal/lib/python3.8/site-packages/torch/_dynamo/symbolic_convert.py", line 2657, in run
    super().run()
  File "/Users/chenxiny/miniforge3/envs/torch-metal/lib/python3.8/site-packages/torch/_dynamo/symbolic_convert.py", line 925, in run
    while self.step():
  File "/Users/chenxiny/miniforge3/envs/torch-metal/lib/python3.8/site-packages/torch/_dynamo/symbolic_convert.py", line 837, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "/Users/chenxiny/miniforge3/envs/torch-metal/lib/python3.8/site-packages/torch/_dynamo/symbolic_convert.py", line 529, in wrapper
    return inner_fn(self, inst)
  File "/Users/chenxiny/miniforge3/envs/torch-metal/lib/python3.8/site-packages/torch/_dynamo/symbolic_convert.py", line 1532, in CALL_FUNCTION
    self.call_function(fn, args, {})
  File "/Users/chenxiny/miniforge3/envs/torch-metal/lib/python3.8/site-packages/torch/_dynamo/symbolic_convert.py", line 775, in call_function
    self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
  File "/Users/chenxiny/miniforge3/envs/torch-metal/lib/python3.8/site-packages/torch/_dynamo/variables/misc.py", line 757, in call_function
    return self.obj.call_method(tx, self.name, args, kwargs)
  File "/Users/chenxiny/miniforge3/envs/torch-metal/lib/python3.8/site-packages/torch/_dynamo/variables/misc.py", line 1202, in call_method
    unimplemented("Logger not supported for non-export cases")
  File "/Users/chenxiny/miniforge3/envs/torch-metal/lib/python3.8/site-packages/torch/_dynamo/exc.py", line 271, in unimplemented
    raise Unsupported(msg, case_name=case_name)
torch._dynamo.exc.Unsupported: Logger not supported for non-export cases

from user code:
   File "dyn_logging.py", line 11, in fn
    logger.info("hi")

Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information


You can suppress this exception and fall back to eager by setting:
    import torch._dynamo
    torch._dynamo.config.suppress_errors = True
```

### Minified repro

```python
import torch
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
# using fullgraph to raise an exception in place of graph breaks
@torch._dynamo.config.patch(reorderable_logging_functions={logger.info})
@torch.compile(backend='eager', fullgraph=True) 
def fn(x):
    y = x + 1
    logger.info("hi")
    y = x * 2
    return y


def main():
    a = torch.tensor([1.,]).requires_grad_(True)
    fn(a)


if __name__ == "__main__":
    main()
```

### Versions

```
PyTorch version: 2.5.0.dev20240729
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 14.3.1 (arm64)
GCC version: Could not collect
Clang version: 15.0.0 (clang-1500.3.9.4)
CMake version: version 3.30.0
Libc version: N/A

Python version: 3.8.18 (default, Sep 11 2023, 08:17:16)  [Clang 14.0.6 ] (64-bit runtime)
Python platform: macOS-14.3.1-arm64-arm-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Apple M1 Pro

Versions of relevant libraries:
[pip3] numpy==1.24.3
[pip3] onnx==1.16.1
[pip3] onnx2torch==1.5.14
[pip3] optree==0.11.0
[pip3] pytorch-lightning==1.7.7
[pip3] torch==2.5.0a0+git1f961ad
[pip3] torchaudio==2.4.0.dev20240729
[pip3] torchmetrics==0.10.0
[pip3] torchvision==0.12.0
[conda] numpy                     1.22.3                   pypi_0    pypi
[conda] numpy-base                1.24.3           py38h90707a3_0
[conda] onnx2torch                1.5.14                   pypi_0    pypi
[conda] optree                    0.11.0                   pypi_0    pypi
[conda] pytorch-lightning         1.7.7                    pypi_0    pypi
[conda] torch                     1.12.0.dev20220518          pypi_0    pypi
[conda] torchaudio                2.4.0.dev20240729          pypi_0    pypi
[conda] torchmetrics              0.10.0                   pypi_0    pypi
[conda] torchvision               0.12.0                   pypi_0    pypi
```

cc @ezyang @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @amjames