support loading libraries when installing multiple FBGEMM targets

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
 

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

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

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
 

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
 

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
 

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
 

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
 

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
 

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
 

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

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: