frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

Revert "[Pytorch] Consolidate Strobelight compile time profiler between OSS and fbcode (#135953)"

Browse Source 
This reverts commit b8637503c036abb898f6b880b325aeffe6f09c03.

Reverted https://github.com/pytorch/pytorch/pull/135953 on behalf of https://github.com/kollasb due to Broke internal module factory compatibility, revert from Phabricator failed ([comment](https://github.com/pytorch/pytorch/pull/135953#issuecomment-2351381777))
This commit is contained in:
PyTorch MergeBot
2024-09-15 05:32:38 +00:00
parent 31e42a45dd
commit 357b7fb579
2 changed files with 14 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
torch/_strobelight/examples/compile_time_profile_example.py
View File

@ -1,10 +1,11 @@
# mypy: allow-untyped-defs
import torch
from torch._utils_internal import enable_compiletime_strobelight
from torch._strobelight.compile_time_profiler import StrobelightCompileTimeProfiler
if __name__ == "__main__":
enable_compiletime_strobelight()
# You can pass TORCH_COMPILE_STROBELIGHT=True instead.
StrobelightCompileTimeProfiler.enable()
def fn(x, y, z):
return x * y + z

20
torch/_utils_internal.py
View File

@ -12,6 +12,17 @@ from torch._strobelight.compile_time_profiler import StrobelightCompileTimeProfi
log = logging.getLogger(__name__)
if os.environ.get("TORCH_COMPILE_STROBELIGHT", False):
import shutil
if not shutil.which("strobeclient"):
log.info(
"TORCH_COMPILE_STROBELIGHT is true, but seems like you are not on a FB machine."
)
else:
log.info("Strobelight profiler is enabled via environment variable")
StrobelightCompileTimeProfiler.enable()
# this arbitrary-looking assortment of functionality is provided here
# to have a central place for overrideable behavior. The motivating
# use is the FB build environment, where this source file is replaced
@ -64,15 +75,6 @@ def throw_abstract_impl_not_imported_error(opname, module, context):
)
def enable_compiletime_strobelight():
StrobelightCompileTimeProfiler.enable()
if os.environ.get("TORCH_COMPILE_STROBELIGHT", False):
log.info("Strobelight profiler is enabled via environment variable")
enable_compiletime_strobelight()
# NB! This treats "skip" kwarg specially!!
def compile_time_strobelight_meta(phase_name):
def compile_time_strobelight_meta_inner(function):