diff --git a/fbgemm_gpu/experimental/example/example/__init__.py b/fbgemm_gpu/experimental/example/example/__init__.py index 6f63bdbfa2..b119f8ba1c 100644 --- a/fbgemm_gpu/experimental/example/example/__init__.py +++ b/fbgemm_gpu/experimental/example/example/__init__.py @@ -15,10 +15,6 @@ # pyre-ignore[21] # @manual=//deeplearning/fbgemm/fbgemm_gpu:test_utils from fbgemm_gpu import open_source - - # pyre-ignore[21] - # @manual=//deeplearning/fbgemm/fbgemm_gpu:test_utils - from fbgemm_gpu.docs.version import __version__ # noqa: F401 except Exception: open_source: bool = False diff --git a/fbgemm_gpu/experimental/gemm/triton_gemm/__init__.py b/fbgemm_gpu/experimental/gemm/triton_gemm/__init__.py index 8ba2244f14..0d158b1653 100644 --- a/fbgemm_gpu/experimental/gemm/triton_gemm/__init__.py +++ b/fbgemm_gpu/experimental/gemm/triton_gemm/__init__.py @@ -11,9 +11,5 @@ # pyre-ignore[21] # @manual=//deeplearning/fbgemm/fbgemm_gpu:test_utils from fbgemm_gpu import open_source - - # pyre-ignore[21] - # @manual=//deeplearning/fbgemm/fbgemm_gpu:test_utils - from fbgemm_gpu.docs.version import __version__ # noqa: F401 except Exception: open_source: bool = False diff --git a/fbgemm_gpu/experimental/gen_ai/bench/__init__.py b/fbgemm_gpu/experimental/gen_ai/bench/__init__.py index 4ddf254d44..216404efac 100644 --- a/fbgemm_gpu/experimental/gen_ai/bench/__init__.py +++ b/fbgemm_gpu/experimental/gen_ai/bench/__init__.py @@ -9,9 +9,5 @@ # pyre-ignore[21] # @manual=//deeplearning/fbgemm/fbgemm_gpu:test_utils from fbgemm_gpu import open_source - - # pyre-ignore[21] - # @manual=//deeplearning/fbgemm/fbgemm_gpu:test_utils - from fbgemm_gpu.docs.version import __version__ # noqa: F401 except Exception: open_source: bool = False diff --git a/fbgemm_gpu/experimental/gen_ai/gen_ai/__init__.py b/fbgemm_gpu/experimental/gen_ai/gen_ai/__init__.py index 61d27625e2..d8e23bd97e 100644 --- a/fbgemm_gpu/experimental/gen_ai/gen_ai/__init__.py +++ b/fbgemm_gpu/experimental/gen_ai/gen_ai/__init__.py @@ -15,10 +15,6 @@ # pyre-ignore[21] # @manual=//deeplearning/fbgemm/fbgemm_gpu:test_utils from fbgemm_gpu import open_source - - # pyre-ignore[21] - # @manual=//deeplearning/fbgemm/fbgemm_gpu:test_utils - from fbgemm_gpu.docs.version import __version__ # noqa: F401 except Exception: open_source: bool = False diff --git a/fbgemm_gpu/experimental/gen_ai/gen_ai/attention/cutlass_blackwell_fmha/__init__.py b/fbgemm_gpu/experimental/gen_ai/gen_ai/attention/cutlass_blackwell_fmha/__init__.py index fbfe659a22..30d6f8e2c6 100644 --- a/fbgemm_gpu/experimental/gen_ai/gen_ai/attention/cutlass_blackwell_fmha/__init__.py +++ b/fbgemm_gpu/experimental/gen_ai/gen_ai/attention/cutlass_blackwell_fmha/__init__.py @@ -10,10 +10,6 @@ # pyre-ignore[21] # @manual=//deeplearning/fbgemm/fbgemm_gpu:test_utils from fbgemm_gpu import open_source - - # pyre-ignore[21] - # @manual=//deeplearning/fbgemm/fbgemm_gpu:test_utils - from fbgemm_gpu.docs.version import __version__ # noqa: F401 except Exception: open_source: bool = False diff --git a/fbgemm_gpu/experimental/gen_ai/gen_ai/attention/cutlass_blackwell_fmha/cutlass_blackwell_fmha_interface.py b/fbgemm_gpu/experimental/gen_ai/gen_ai/attention/cutlass_blackwell_fmha/cutlass_blackwell_fmha_interface.py index ea9bdc3768..6f7c684cb5 100644 --- a/fbgemm_gpu/experimental/gen_ai/gen_ai/attention/cutlass_blackwell_fmha/cutlass_blackwell_fmha_interface.py +++ b/fbgemm_gpu/experimental/gen_ai/gen_ai/attention/cutlass_blackwell_fmha/cutlass_blackwell_fmha_interface.py @@ -12,10 +12,6 @@ # pyre-ignore[21] # @manual=//deeplearning/fbgemm/fbgemm_gpu:test_utils from fbgemm_gpu import open_source - - # pyre-ignore[21] - # @manual=//deeplearning/fbgemm/fbgemm_gpu:test_utils - from fbgemm_gpu.docs.version import __version__ # noqa: F401 except Exception: open_source: bool = False diff --git a/fbgemm_gpu/experimental/gen_ai/gen_ai/moe/__init__.py b/fbgemm_gpu/experimental/gen_ai/gen_ai/moe/__init__.py index 2ab5be4910..720de0ac76 100644 --- a/fbgemm_gpu/experimental/gen_ai/gen_ai/moe/__init__.py +++ b/fbgemm_gpu/experimental/gen_ai/gen_ai/moe/__init__.py @@ -14,10 +14,6 @@ # pyre-ignore[21] # @manual=//deeplearning/fbgemm/fbgemm_gpu:test_utils from fbgemm_gpu import open_source - - # pyre-ignore[21] - # @manual=//deeplearning/fbgemm/fbgemm_gpu:test_utils - from fbgemm_gpu.docs.version import __version__ # noqa: F401 except Exception: open_source: bool = False diff --git a/fbgemm_gpu/experimental/gen_ai/gen_ai/moe/layers.py b/fbgemm_gpu/experimental/gen_ai/gen_ai/moe/layers.py index 1497028d42..2df16741cd 100644 --- a/fbgemm_gpu/experimental/gen_ai/gen_ai/moe/layers.py +++ b/fbgemm_gpu/experimental/gen_ai/gen_ai/moe/layers.py @@ -38,10 +38,6 @@ # pyre-ignore[21] # @manual=//deeplearning/fbgemm/fbgemm_gpu:test_utils from fbgemm_gpu import open_source - - # pyre-ignore[21] - # @manual=//deeplearning/fbgemm/fbgemm_gpu:test_utils - from fbgemm_gpu.docs.version import __version__ # noqa: F401 except Exception: open_source: bool = False diff --git a/fbgemm_gpu/experimental/gen_ai/test/moe/gather_scatter_test.py b/fbgemm_gpu/experimental/gen_ai/test/moe/gather_scatter_test.py index bc697ab639..4186986bc1 100644 --- a/fbgemm_gpu/experimental/gen_ai/test/moe/gather_scatter_test.py +++ b/fbgemm_gpu/experimental/gen_ai/test/moe/gather_scatter_test.py @@ -31,9 +31,6 @@ try: # @manual=//deeplearning/fbgemm/fbgemm_gpu:test_utils from fbgemm_gpu import open_source - - # @manual=//deeplearning/fbgemm/fbgemm_gpu:test_utils - from fbgemm_gpu.docs.version import __version__ # noqa: F401 except Exception: open_source: bool = False diff --git a/fbgemm_gpu/experimental/gen_ai/test/moe/shuffling_test.py b/fbgemm_gpu/experimental/gen_ai/test/moe/shuffling_test.py index bd025b655f..0d6c040a5a 100644 --- a/fbgemm_gpu/experimental/gen_ai/test/moe/shuffling_test.py +++ b/fbgemm_gpu/experimental/gen_ai/test/moe/shuffling_test.py @@ -29,9 +29,6 @@ try: # @manual=//deeplearning/fbgemm/fbgemm_gpu:test_utils from fbgemm_gpu import open_source - - # @manual=//deeplearning/fbgemm/fbgemm_gpu:test_utils - from fbgemm_gpu.docs.version import __version__ # noqa: F401 except Exception: open_source: bool = False diff --git a/fbgemm_gpu/fbgemm_gpu/__init__.py b/fbgemm_gpu/fbgemm_gpu/__init__.py index 942d7f10b7..afb69cfa87 100644 --- a/fbgemm_gpu/fbgemm_gpu/__init__.py +++ b/fbgemm_gpu/fbgemm_gpu/__init__.py @@ -5,6 +5,7 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import json import logging import os import re @@ -26,6 +27,19 @@ } +def _load_target_info(target: str) -> dict[str, str]: + try: + filepath = os.path.join( + os.path.dirname(__file__), "docs", f"target.{target}.json.py" + ) + with open(filepath, "r") as file: + data = json.load(file) + except Exception: + data = {} + + return data + + def _load_library(filename: str, version: str, no_throw: bool = False) -> None: """Load a shared library from the given filename.""" @@ -98,13 +112,16 @@ def _load_library(filename: str, version: str, no_throw: bool = False) -> None: # Trigger the manual addition of docstrings to pybind11-generated operators import fbgemm_gpu.docs # noqa: F401, E402 + +__targets_infos__ = { + target: _load_target_info(target) for target in ["default", "genai", "hstu"] +} +__targets_infos__ = {k: v for (k, v) in __targets_infos__.items() if v} + try: - # Export the version string from the version file auto-generated by setup.py - from fbgemm_gpu.docs.version import ( # noqa: F401, E402 - __target__, - __variant__, - __version__, - ) + __target__, __info__ = next(iter(__targets_infos__.items())) + __variant__ = __info__["variant"] + __version__ = __info__["version"] except Exception: __variant__: str = "INTERNAL" __version__: str = "INTERNAL" @@ -145,18 +162,19 @@ def _load_library(filename: str, version: str, no_throw: bool = False) -> None: "genai": fbgemm_genai_libraries, } -for library in libraries_to_load.get(__target__, []): - # NOTE: In all cases, we want to throw an error if we cannot load the - # library. However, this appears to break the OSS documentation build, - # where the Python documentation doesn't show up in the generated docs. - # - # To work around this problem, we introduce a fake build variant called - # `docs` and we only throw a library load error when the variant is not - # `docs`. For more information, see: - # - # https://github.com/pytorch/FBGEMM/pull/3477 - # https://github.com/pytorch/FBGEMM/pull/3717 - _load_library(f"{library}.so", __version__, __variant__ == "docs") +for target, info in __targets_infos__.items(): + for library in libraries_to_load.get(target, []): + # NOTE: In all cases, we want to throw an error if we cannot load the + # library. However, this appears to break the OSS documentation build, + # where the Python documentation doesn't show up in the generated docs. + # + # To work around this problem, we introduce a fake build variant called + # `docs` and we only throw a library load error when the variant is not + # `docs`. For more information, see: + # + # https://github.com/pytorch/FBGEMM/pull/3477 + # https://github.com/pytorch/FBGEMM/pull/3717 + _load_library(f"{library}.so", info["version"], info["variant"] == "docs") try: # Trigger meta operator registrations diff --git a/fbgemm_gpu/setup.py b/fbgemm_gpu/setup.py index 1031297307..97600fc0bb 100644 --- a/fbgemm_gpu/setup.py +++ b/fbgemm_gpu/setup.py @@ -58,8 +58,8 @@ def from_args(cls, argv: list[str]): parser.add_argument( "--build-target", type=str, - choices=["default", "genai", "hstu"], - default="default", + choices=["none", "default", "genai", "hstu"], + default="none", help="The FBGEMM build target to build.", ) parser.add_argument( @@ -134,7 +134,9 @@ def nova_non_prebuild_step(self) -> bool: # flag. As such, we skip building in the clean and build wheel steps. return self.nova_flag() == 1 - def target(self) -> str: + def target(self) -> Optional[str]: + if self.args.build_target == "none": + return None return self.args.build_target def variant(self) -> str: @@ -145,8 +147,10 @@ def package_name(self) -> str: pkg_name: str = "fbgemm_gpu_genai" elif self.target() == "hstu": pkg_name: str = "fbgemm_gpu_hstu" - else: + elif self.target() == "default": pkg_name: str = "fbgemm_gpu" + else: + pkg_name: str = "none" if self.nova_flag() is None: # If running outside of Nova workflow context, append the channel @@ -501,27 +505,25 @@ class FbgemmGpuInstall(PipInstall): @classmethod def generate_version_file(cls, build: FbgemmGpuBuild) -> None: - with open("fbgemm_gpu/docs/version.py", "w") as file: - package_version = build.package_version() - - print( - f"[SETUP.PY] Generating version file at: {os.path.realpath(file.name)}" - ) - text = textwrap.dedent( - f""" - #!/usr/bin/env python3 - # Copyright (c) Meta Platforms, Inc. and affiliates. - # All rights reserved. - # - # This source code is licensed under the BSD-style license found in the - # LICENSE file in the root directory of this source tree. - - __version__: str = "{package_version}" - __target__: str = "{build.target()}" - __variant__: str = "{build.variant()}" - """ - ) - file.write(text) + if build.target(): + # NOTE: This is a JSON file needs to be named with *.py extension so + # that setup.py can pick it up and package it into the wheel + with open(f"fbgemm_gpu/docs/target.{build.target()}.json.py", "w") as file: + package_version = build.package_version() + + print( + f"[SETUP.PY] Generating target info file at: {os.path.realpath(file.name)}" + ) + text = textwrap.dedent( + f""" + {{ + "version": "{package_version}", + "target": "{build.target()}", + "variant": "{build.variant()}" + }} + """ + ) + file.write(text) @classmethod def description(cls) -> str: