Skip to content

Commit

Permalink
Fix package variant versioning (pytorch#2045)
Browse files Browse the repository at this point in the history 
Summary:
- Fix package variant versioning for ROCm Nova builds

Pull Request resolved: pytorch#2045

Reviewed By: sryap

Differential Revision: D49716498

Pulled By: q10

fbshipit-source-id: ff2e9f98c2acfe79fee2fb7af703e8c21587164d
  • Loading branch information
@q10 @facebook-github-bot
q10 authored and facebook-github-bot committed Sep 28, 2023
1 parent e4ed378 commit 1595e2a
Show file tree
Hide file tree
Showing 3 changed files with 126 additions and 69 deletions.
9 changes: 7 additions & 2 deletions .github/scripts/fbgemm_gpu_build.bash
View file
Original file line number Diff line number Diff line change
Expand Up @@ -58,7 +58,9 @@ prepare_fbgemm_gpu_build () {
__configure_fbgemm_gpu_build_cpu () {
# Update the package name and build args depending on if CUDA is specified
echo "[BUILD] Setting CPU-only build args ..."
build_args=(--cpu_only)
build_args=(
--package_variant=cpu
)
}

__configure_fbgemm_gpu_build_rocm () {
Expand Down Expand Up @@ -88,7 +90,9 @@ __configure_fbgemm_gpu_build_rocm () {
print_exec conda env config vars set ${env_prefix} PYTORCH_ROCM_ARCH="${arch_list}"

echo "[BUILD] Setting ROCm build args ..."
build_args=()
build_args=(
--package_variant=rocm
)
}

__configure_fbgemm_gpu_build_cuda () {
Expand Down Expand Up @@ -133,6 +137,7 @@ __configure_fbgemm_gpu_build_cuda () {
# shellcheck disable=SC2155,SC2086
local nvml_lib_path=$(conda run --no-capture-output ${env_prefix} printenv NVML_LIB_PATH)
build_args=(
--package_variant=cuda
--nvml_lib_path="${nvml_lib_path}"
-DTORCH_CUDA_ARCH_LIST="'${arch_list}'"
)
Expand Down
9 changes: 6 additions & 3 deletions fbgemm_gpu/docs/BuildInstructions.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -354,13 +354,15 @@ unset TORCH_CUDA_ARCH_LIST
# Build the wheel artifact only
python setup.py bdist_wheel \
--package_name="${package_name}" \
--package_variant=cuda \
--python-tag="${python_tag}" \
--plat-name="manylinux1_${ARCH}" \
--nvml_lib_path=${NVML_LIB_PATH} \
-DTORCH_CUDA_ARCH_LIST="${cuda_arch_list}"

# Build and install the library into the Conda environment
python setup.py install \
--package_variant=cuda \
--nvml_lib_path=${NVML_LIB_PATH} \
-DTORCH_CUDA_ARCH_LIST="${cuda_arch_list}"
```
Expand All @@ -384,11 +386,12 @@ package_name=fbgemm_gpu_rocm
# Build the wheel artifact only
python setup.py bdist_wheel \
--package_name="${package_name}" \
--package_variant=rocm \
--python-tag="${python_tag}" \
--plat-name="manylinux1_${ARCH}"

# Build and install the library into the Conda environment
python setup.py install develop
python setup.py install develop --package_variant=rocm
```

### CPU-Only Build
Expand All @@ -406,12 +409,12 @@ package_name=fbgemm_gpu_cpu
# Build the wheel artifact only
python setup.py bdist_wheel \
--package_name="${package_name}" \
--package_variant=cpu \
--python-tag="${python_tag}" \
--plat-name="manylinux1_${ARCH}" \
--cpu_only

# Build and install the library into the Conda environment
python setup.py install --cpu_only
python setup.py install --package_variant=cpu
```

### Post-Build Checks (For Developers)
Expand Down
177 changes: 113 additions & 64 deletions fbgemm_gpu/setup.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,47 +24,20 @@
from tabulate import tabulate


def generate_package_version(package_name: str, version_variant: str):
print("[SETUP.PY] Generating the package version ...")

if "nightly" in package_name:
# Use date stamp for nightly versions
print("[SETUP.PY] Package is for NIGHTLY; using timestamp for the versioning")
today = date.today()
version = f"{today.year}.{today.month}.{today.day}"

elif "test" in package_name:
# Use date stamp for nightly versions
print("[SETUP.PY] Package is for TEST: using random number for the versioning")
version = (f"0.0.{random.randint(0, 1000)}",)

else:
# Use git tag / branch / commit info to generate a PEP-440-compliant version string
print("[SETUP.PY] Package is for RELEASE: using git info for the versioning")
print(
f"[SETUP.PY] TAG: {gitversion.get_tag()}, BRANCH: {gitversion.get_branch()}, SHA: {gitversion.get_sha()}"
)
# Remove the local version identifier, if any (e.g. 0.4.0rc0.post0+git.6a63116c.dirty => 0.4.0rc0.post0)
# Then remove post0 (keep postN for N > 0) (e.g. 0.4.0rc0.post0 => 0.4.0rc0)
version = re.sub(".post0$", "", gitversion.version_from_git().split("+")[0])
version = str(version) + version_variant
print(f"[SETUP.PY] Setting the package version: {version}")
return version


def parse_args(argv: List[str]) -> argparse.Namespace:
parser = argparse.ArgumentParser(description="fbgemm_gpu setup")
parser.add_argument(
"--cpu_only",
dest="cpu_only",
action="store_true",
help="build for cpu-only (no GPU support)",
"--package_variant",
type=str,
choices=["cpu", "cuda", "rocm"],
default="cpu",
help="The FBGEMM_GPU variant to build.",
)
parser.add_argument(
"--package_name",
type=str,
default="fbgemm_gpu",
help="the name of this output wheel",
help="The candidate name of the output wheel.",
)
parser.add_argument(
"--nvml_lib_path",
Expand Down Expand Up @@ -171,7 +144,7 @@ def _get_cxx11_abi():
os.environ["CMAKE_BUILD_PARALLEL_LEVEL"] = str(os.cpu_count() // 2)

cmake_args = [f"-DCMAKE_PREFIX_PATH={torch_root}", _get_cxx11_abi()]
if args.cpu_only:
if args.package_variant == "cpu":
cmake_args.append("-DFBGEMM_CPU_ONLY=ON")
if args.nvml_lib_path:
cmake_args.append(f"-DNVML_LIB_PATH={args.nvml_lib_path}")
Expand All @@ -181,6 +154,100 @@ def _get_cxx11_abi():
class FbgemmGpuInstaller(PipInstall):
"""FBGEMM_GPU PIP Installer"""

@classmethod
def extract_package_name(cls, args_package_name: str) -> str:
package_name: str = ""

if "BUILD_FROM_NOVA" in os.environ:
nova_flag = os.getenv("BUILD_FROM_NOVA")
print(f"[SETUP.PY] BUILD_FROM_NOVA={nova_flag}")

# The package name is the same for all build variants in Nova
package_name = "fbgemm_gpu"

if str(nova_flag) != "0":
# Skip build clean and build wheel steps in Nova workflow since
# they are done in pre-script
print("[SETUP.PY] Build from Nova detected... exiting.")
sys.exit(0)

else:
package_name = args_package_name

print(f"[SETUP.PY] Extracted the package name: '{package_name}'")
return package_name

@classmethod
def extract_variant_version(cls, variant: str) -> str:
variant_version: str = ""

if variant == "cpu":
variant_version = "+cpu"
elif variant == "cuda":
set_cuda_environment_variables()
if torch.version.cuda is not None:
cuda_version = torch.version.cuda.split(".")
variant_version = f"+cu{cuda_version[0]}{cuda_version[1]}"
else:
sys.exit(
"[SETUP.PY] Installed PyTorch variant is not CUDA; cannot determine the CUDA version!"
)
elif variant == "rocm":
if torch.version.hip is not None:
rocm_version = torch.version.hip.split(".")
variant_version = f"+rocm{rocm_version[0]}.{rocm_version[1]}"
else:
sys.exit(
"[SETUP.PY] Installed PyTorch variant is not ROCm; cannot determine the ROCm version!"
)
else:
sys.exit(
f"[SETUP.PY] Unrecognized build variant variant '{variant}'; cannot proceed with FBGEMM_GPU build!"
)

if "BUILD_FROM_NOVA" not in os.environ:
# If not building from Nova, use the fbgemm_gpu-<variant>
# PyPI does not accept version+xx in the name convention.
print("[SETUP.PY] Not building FBGEMM_GPU from Nova.")
variant_version = ""

print(f"[SETUP.PY] Extracted the package variant+version: '{variant_version}'")
return variant_version

@classmethod
def generate_package_version(cls, package_name: str, variant_version: str):
print("[SETUP.PY] Generating the package version ...")

if "nightly" in package_name:
# Use date stamp for nightly versions
print(
"[SETUP.PY] Package is for NIGHTLY; using timestamp for the versioning"
)
today = date.today()
version = f"{today.year}.{today.month}.{today.day}"

elif "test" in package_name and "BUILD_FROM_NOVA" not in os.environ:
# Use random numbering for test versions
print(
"[SETUP.PY] Package is for TEST: using random number for the versioning"
)
version = (f"0.0.{random.randint(0, 1000)}",)

else:
# Use git tag / branch / commit info to generate a PEP-440-compliant version string
print(
"[SETUP.PY] Package is for RELEASE: using git info for the versioning"
)
print(
f"[SETUP.PY] TAG: {gitversion.get_tag()}, BRANCH: {gitversion.get_branch()}, SHA: {gitversion.get_sha()}"
)
# Remove the local version identifier, if any (e.g. 0.4.0rc0.post0+git.6a63116c.dirty => 0.4.0rc0.post0)
# Then remove post0 (keep postN for N > 0) (e.g. 0.4.0rc0.post0 => 0.4.0rc0)
version = re.sub(".post0$", "", gitversion.version_from_git().split("+")[0])
version = str(version) + variant_version
print(f"[SETUP.PY] Setting the full package version string: {version}")
return version

@classmethod
def generate_version_file(cls, package_version: str) -> None:
with open("fbgemm_gpu/_fbgemm_gpu_version.py", "w") as file:
Expand Down Expand Up @@ -262,41 +329,23 @@ def run(self):
def main(argv: List[str]) -> None:
# Handle command line args before passing to main setup() method.
args, unknown = parse_args(argv)
print("Parsed Arguments: ", args)
print(f"[SETUP.PY] Parsed Arguments: {args}")
if len(unknown) != 0 and (len(unknown) != 1 or unknown[0] != "clean"):
print("Unknown Arguments: ", unknown)
print(f"[SETUP.PY] Unknown Arguments: {unknown}")

if args.cpu_only:
version_variant = "+cpu"
else:
set_cuda_environment_variables()
if torch.version.cuda is not None:
cuda_version = torch.version.cuda.split(".")
version_variant = "+cu" + str(cuda_version[0]) + str(cuda_version[1])
else:
# rocm or other gpus - to be specified if we offcially support them
version_variant = ""

# Skip Nova build steps since it will be done in pre-script
if "BUILD_FROM_NOVA" in os.environ:
build_from_nova = os.getenv("BUILD_FROM_NOVA")
print("build_from_nova", build_from_nova)
# Package name is the same for all variants in Nova
package_name = "fbgemm_gpu"
if str(build_from_nova) != "0":
# Skip build clean and build wheel steps in Nova workflow since they are done in pre-script
print("Build from Nova detected... exiting")
sys.exit(0)
else:
# If not building from Nova, use the fbgemm_gpu-<variant>
# PyPi does not accept version+xx in the name convention.
version_variant = ""
package_name = args.package_name
# Repair command line args for setup.
sys.argv = [sys.argv[0]] + unknown

# Determine the package version
package_version = generate_package_version(args.package_name, version_variant)
# Extract the package name
package_name = FbgemmGpuInstaller.extract_package_name(args.package_name)

# Extract the variant version, e.g. cpu, cu121, rocm5.6
variant_version = FbgemmGpuInstaller.extract_variant_version(args.package_variant)

# Generate the full package version string
package_version = FbgemmGpuInstaller.generate_package_version(
package_name, variant_version
)

# Generate the version file
FbgemmGpuInstaller.generate_version_file(package_version)
Expand Down

0 comments on commit 1595e2a

Please sign in to comment.