mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
c480a479b1
This is semantics changing as if you are dealing with multiple code objects which have exactly the same filename/firstlineno/name, but are distinct objects, and need non-aliasing automatic dynamic state. Otherwise, this should be equivalent (modulo lifetime). I want to do this because when I do PGO I can't index on code object identity, need a stable identifier. Signed-off-by: Edward Z. Yang <ezyang@meta.com> Pull Request resolved: https://github.com/pytorch/pytorch/pull/138740 Approved by: https://github.com/bobrenjc93 ghstack dependencies: #138693, #138717
62 lines
1.7 KiB
Python
62 lines
1.7 KiB
Python
import itertools
|
|
import logging
|
|
from typing import Any, Callable, List
|
|
|
|
from torch.hub import _Faketqdm, tqdm
|
|
|
|
|
|
# Disable progress bar by default, not in dynamo config because otherwise get a circular import
|
|
disable_progress = True
|
|
|
|
|
|
# Return all loggers that torchdynamo/torchinductor is responsible for
|
|
def get_loggers() -> List[logging.Logger]:
|
|
return [
|
|
logging.getLogger("torch.fx.experimental.symbolic_shapes"),
|
|
logging.getLogger("torch._dynamo"),
|
|
logging.getLogger("torch._inductor"),
|
|
]
|
|
|
|
|
|
# Creates a logging function that logs a message with a step # prepended.
|
|
# get_step_logger should be lazily called (i.e. at runtime, not at module-load time)
|
|
# so that step numbers are initialized properly. e.g.:
|
|
|
|
# @functools.lru_cache(None)
|
|
# def _step_logger():
|
|
# return get_step_logger(logging.getLogger(...))
|
|
|
|
# def fn():
|
|
# _step_logger()(logging.INFO, "msg")
|
|
|
|
_step_counter = itertools.count(1)
|
|
|
|
# Update num_steps if more phases are added: Dynamo, AOT, Backend
|
|
# This is very inductor centric
|
|
# _inductor.utils.has_triton() gives a circular import error here
|
|
|
|
if not disable_progress:
|
|
try:
|
|
import triton # noqa: F401
|
|
|
|
num_steps = 3
|
|
except ImportError:
|
|
num_steps = 2
|
|
pbar = tqdm(total=num_steps, desc="torch.compile()", delay=0)
|
|
|
|
|
|
def get_step_logger(logger: logging.Logger) -> Callable[..., None]:
|
|
if not disable_progress:
|
|
pbar.update(1)
|
|
if not isinstance(pbar, _Faketqdm):
|
|
pbar.set_postfix_str(f"{logger.name}")
|
|
|
|
step = next(_step_counter)
|
|
|
|
def log(level: int, msg: str, **kwargs: Any) -> None:
|
|
if "stacklevel" not in kwargs:
|
|
kwargs["stacklevel"] = 2
|
|
logger.log(level, "Step %s: %s", step, msg, **kwargs)
|
|
|
|
return log
|