_utils_internal.py 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367
  1. # mypy: allow-untyped-defs
  2. import functools
  3. import logging
  4. import os
  5. import sys
  6. import tempfile
  7. import typing_extensions
  8. from typing import Any, Callable, Optional, TypeVar
  9. from typing_extensions import ParamSpec
  10. import torch
  11. from torch._strobelight.compile_time_profiler import StrobelightCompileTimeProfiler
  12. _T = TypeVar("_T")
  13. _P = ParamSpec("_P")
  14. log = logging.getLogger(__name__)
  15. if os.environ.get("TORCH_COMPILE_STROBELIGHT", False):
  16. import shutil
  17. if not shutil.which("strobeclient"):
  18. log.info(
  19. "TORCH_COMPILE_STROBELIGHT is true, but seems like you are not on a FB machine."
  20. )
  21. else:
  22. log.info("Strobelight profiler is enabled via environment variable")
  23. StrobelightCompileTimeProfiler.enable()
  24. # this arbitrary-looking assortment of functionality is provided here
  25. # to have a central place for overridable behavior. The motivating
  26. # use is the FB build environment, where this source file is replaced
  27. # by an equivalent.
  28. if os.path.basename(os.path.dirname(__file__)) == "shared":
  29. torch_parent = os.path.dirname(os.path.dirname(os.path.dirname(__file__)))
  30. else:
  31. torch_parent = os.path.dirname(os.path.dirname(__file__))
  32. def get_file_path(*path_components: str) -> str:
  33. return os.path.join(torch_parent, *path_components)
  34. def get_file_path_2(*path_components: str) -> str:
  35. return os.path.join(*path_components)
  36. def get_writable_path(path: str) -> str:
  37. if os.access(path, os.W_OK):
  38. return path
  39. return tempfile.mkdtemp(suffix=os.path.basename(path))
  40. def prepare_multiprocessing_environment(path: str) -> None:
  41. pass
  42. def resolve_library_path(path: str) -> str:
  43. return os.path.realpath(path)
  44. def throw_abstract_impl_not_imported_error(opname, module, context):
  45. if module in sys.modules:
  46. raise NotImplementedError(
  47. f"{opname}: We could not find the fake impl for this operator. "
  48. )
  49. else:
  50. raise NotImplementedError(
  51. f"{opname}: We could not find the fake impl for this operator. "
  52. f"The operator specified that you may need to import the '{module}' "
  53. f"Python module to load the fake impl. {context}"
  54. )
  55. # NB! This treats "skip" kwarg specially!!
  56. def compile_time_strobelight_meta(
  57. phase_name: str,
  58. ) -> Callable[[Callable[_P, _T]], Callable[_P, _T]]:
  59. def compile_time_strobelight_meta_inner(
  60. function: Callable[_P, _T],
  61. ) -> Callable[_P, _T]:
  62. @functools.wraps(function)
  63. def wrapper_function(*args: _P.args, **kwargs: _P.kwargs) -> _T:
  64. if "skip" in kwargs and isinstance(skip := kwargs["skip"], int):
  65. kwargs["skip"] = skip + 1
  66. # This is not needed but we have it here to avoid having profile_compile_time
  67. # in stack traces when profiling is not enabled.
  68. if not StrobelightCompileTimeProfiler.enabled:
  69. return function(*args, **kwargs)
  70. return StrobelightCompileTimeProfiler.profile_compile_time(
  71. function, phase_name, *args, **kwargs
  72. )
  73. return wrapper_function
  74. return compile_time_strobelight_meta_inner
  75. # Meta only, see
  76. # https://www.internalfb.com/intern/wiki/ML_Workflow_Observability/User_Guides/Adding_instrumentation_to_your_code/
  77. #
  78. # This will cause an event to get logged to Scuba via the signposts API. You
  79. # can view samples on the API at https://fburl.com/scuba/workflow_signpost/zh9wmpqs
  80. # we log to subsystem "torch", and the category and name you provide here.
  81. # Each of the arguments translate into a Scuba column. We're still figuring
  82. # out local conventions in PyTorch, but category should be something like
  83. # "dynamo" or "inductor", and name should be a specific string describing what
  84. # kind of event happened.
  85. #
  86. # Killswitch is at
  87. # https://www.internalfb.com/intern/justknobs/?name=pytorch%2Fsignpost#event
  88. def signpost_event(category: str, name: str, parameters: dict[str, Any]):
  89. log.info("%s %s: %r", category, name, parameters)
  90. def add_mlhub_insight(category: str, insight: str, insight_description: str):
  91. pass
  92. def log_compilation_event(metrics):
  93. log.info("%s", metrics)
  94. def upload_graph(graph):
  95. pass
  96. def set_pytorch_distributed_envs_from_justknobs():
  97. pass
  98. def log_export_usage(**kwargs):
  99. pass
  100. def log_draft_export_usage(**kwargs):
  101. pass
  102. def log_trace_structured_event(*args, **kwargs) -> None:
  103. pass
  104. def log_cache_bypass(*args, **kwargs) -> None:
  105. pass
  106. def log_torchscript_usage(api: str, **kwargs):
  107. _ = api
  108. return
  109. def check_if_torch_exportable():
  110. return False
  111. def export_training_ir_rollout_check() -> bool:
  112. return True
  113. def full_aoti_runtime_assert() -> bool:
  114. return True
  115. def log_torch_jit_trace_exportability(
  116. api: str,
  117. type_of_export: str,
  118. export_outcome: str,
  119. result: str,
  120. ):
  121. _, _, _, _ = api, type_of_export, export_outcome, result
  122. return
  123. def justknobs_check(name: str, default: bool = True) -> bool:
  124. """
  125. This function can be used to killswitch functionality in FB prod,
  126. where you can toggle this value to False in JK without having to
  127. do a code push. In OSS, we always have everything turned on all
  128. the time, because downstream users can simply choose to not update
  129. PyTorch. (If more fine-grained enable/disable is needed, we could
  130. potentially have a map we lookup name in to toggle behavior. But
  131. the point is that it's all tied to source code in OSS, since there's
  132. no live server to query.)
  133. This is the bare minimum functionality I needed to do some killswitches.
  134. We have a more detailed plan at
  135. https://docs.google.com/document/d/1Ukerh9_42SeGh89J-tGtecpHBPwGlkQ043pddkKb3PU/edit
  136. In particular, in some circumstances it may be necessary to read in
  137. a knob once at process start, and then use it consistently for the
  138. rest of the process. Future functionality will codify these patterns
  139. into a better high level API.
  140. WARNING: Do NOT call this function at module import time, JK is not
  141. fork safe and you will break anyone who forks the process and then
  142. hits JK again.
  143. """
  144. return default
  145. def justknobs_getval_int(name: str) -> int:
  146. """
  147. Read warning on justknobs_check
  148. """
  149. return 0
  150. def is_fb_unit_test() -> bool:
  151. return False
  152. @functools.cache
  153. def max_clock_rate():
  154. """
  155. unit: MHz
  156. """
  157. if not torch.version.hip:
  158. from triton.testing import nvsmi
  159. return nvsmi(["clocks.max.sm"])[0]
  160. else:
  161. # Manually set max-clock speeds on ROCm until equivalent nvmsi
  162. # functionality in triton.testing or via pyamdsmi enablement. Required
  163. # for test_snode_runtime unit tests.
  164. gcn_arch = str(torch.cuda.get_device_properties(0).gcnArchName.split(":", 1)[0])
  165. if "gfx94" in gcn_arch:
  166. return 1700
  167. elif "gfx90a" in gcn_arch:
  168. return 1700
  169. elif "gfx908" in gcn_arch:
  170. return 1502
  171. elif "gfx12" in gcn_arch:
  172. return 1700
  173. elif "gfx11" in gcn_arch:
  174. return 1700
  175. elif "gfx103" in gcn_arch:
  176. return 1967
  177. elif "gfx101" in gcn_arch:
  178. return 1144
  179. elif "gfx95" in gcn_arch:
  180. return 1700 # TODO: placeholder, get actual value
  181. else:
  182. return 1100
  183. def get_mast_job_name_version() -> Optional[tuple[str, int]]:
  184. return None
  185. TEST_MASTER_ADDR = "127.0.0.1"
  186. TEST_MASTER_PORT = 29500
  187. # USE_GLOBAL_DEPS controls whether __init__.py tries to load
  188. # libtorch_global_deps, see Note [Global dependencies]
  189. USE_GLOBAL_DEPS = True
  190. # USE_RTLD_GLOBAL_WITH_LIBTORCH controls whether __init__.py tries to load
  191. # _C.so with RTLD_GLOBAL during the call to dlopen.
  192. USE_RTLD_GLOBAL_WITH_LIBTORCH = False
  193. # If an op was defined in C++ and extended from Python using the
  194. # torch.library.register_fake, returns if we require that there be a
  195. # m.set_python_module("mylib.ops") call from C++ that associates
  196. # the C++ op with a python module.
  197. REQUIRES_SET_PYTHON_MODULE = False
  198. def maybe_upload_prof_stats_to_manifold(profile_path: str) -> Optional[str]:
  199. print("Uploading profile stats (fb-only otherwise no-op)")
  200. return None
  201. def log_chromium_event_internal(
  202. event: dict[str, Any],
  203. stack: list[str],
  204. logger_uuid: str,
  205. start_time_ns: int,
  206. ):
  207. return None
  208. def record_chromium_event_internal(
  209. event: dict[str, Any],
  210. ):
  211. return None
  212. def profiler_allow_cudagraph_cupti_lazy_reinit_cuda12():
  213. return True
  214. def deprecated():
  215. """
  216. When we deprecate a function that might still be in use, we make it internal
  217. by adding a leading underscore. This decorator is used with a private function,
  218. and creates a public alias without the leading underscore, but has a deprecation
  219. warning. This tells users "THIS FUNCTION IS DEPRECATED, please use something else"
  220. without breaking them, however, if they still really really want to use the
  221. deprecated function without the warning, they can do so by using the internal
  222. function name.
  223. """
  224. def decorator(func: Callable[_P, _T]) -> Callable[_P, _T]:
  225. # Validate naming convention – single leading underscore, not dunder
  226. if not (func.__name__.startswith("_")):
  227. raise ValueError(
  228. "@deprecate must decorate a function whose name "
  229. "starts with a single leading underscore (e.g. '_foo') as the api should be considered internal for deprecation."
  230. )
  231. public_name = func.__name__[1:] # drop exactly one leading underscore
  232. module = sys.modules[func.__module__]
  233. # Don't clobber an existing symbol accidentally.
  234. if hasattr(module, public_name):
  235. raise RuntimeError(
  236. f"Cannot create alias '{public_name}' -> symbol already exists in {module.__name__}. \
  237. Please rename it or consult a pytorch developer on what to do"
  238. )
  239. warning_msg = f"{func.__name__[1:]} is DEPRECATED, please consider using an alternative API(s). "
  240. # public deprecated alias
  241. alias = typing_extensions.deprecated(
  242. warning_msg, category=UserWarning, stacklevel=1
  243. )(func)
  244. alias.__name__ = public_name
  245. # Adjust qualname if nested inside a class or another function
  246. if "." in func.__qualname__:
  247. alias.__qualname__ = func.__qualname__.rsplit(".", 1)[0] + "." + public_name
  248. else:
  249. alias.__qualname__ = public_name
  250. setattr(module, public_name, alias)
  251. return func
  252. return decorator
  253. def get_default_numa_options():
  254. """
  255. When using elastic agent, if no numa options are provided, we will use these
  256. as the default.
  257. For external use cases, we return None, i.e. no numa binding. If you would like
  258. to use torch's automatic numa binding capabilities, you should provide
  259. NumaOptions to your launch config directly or use the numa binding option
  260. available in torchrun.
  261. Must return None or NumaOptions, but not specifying to avoid circular import.
  262. """
  263. return None
  264. def log_triton_builds(fail: Optional[str]):
  265. pass
  266. def find_compile_subproc_binary() -> Optional[str]:
  267. """
  268. Allows overriding the binary used for subprocesses
  269. """
  270. return None