From c773830b88897f8b22ac023146f6ab7aa4231b4e Mon Sep 17 00:00:00 2001 From: Jeremy Reizenstein Date: Thu, 15 Dec 2022 03:22:07 -0800 Subject: [PATCH] fix and pythonify conda build (#1394) Summary: Pull Request resolved: https://github.com/facebookresearch/pytorch3d/pull/1394 The bash logic for building conda packages became fiddly to edit. We need to switch cuda-toolkit to pytorch-cuda when PyTorch>=1.12 which was going to be a pain, so here I rewrite the code in python and do it. Reviewed By: shapovalov Differential Revision: D42036406 fbshipit-source-id: 8bb80c2f7545477182b23fc97c8514dcafcee176 --- .circleci/config.in.yml | 4 +- .circleci/config.yml | 4 +- packaging/build_conda.py | 134 +++++++++++++++++++++++++++++++++++++++ packaging/build_conda.sh | 34 ---------- 4 files changed, 138 insertions(+), 38 deletions(-) create mode 100644 packaging/build_conda.py delete mode 100755 packaging/build_conda.sh diff --git a/.circleci/config.in.yml b/.circleci/config.in.yml index 909aeaf1..3186bcbe 100644 --- a/.circleci/config.in.yml +++ b/.circleci/config.in.yml @@ -117,7 +117,7 @@ jobs: - run: name: build no_output_timeout: 20m - command: MAX_JOBS=15 TEST_FLAG=--no-test packaging/build_conda.sh + command: MAX_JOBS=15 TEST_FLAG=--no-test python3 packaging/build_conda.py - store_artifacts: path: /opt/conda/conda-bld/linux-64 - persist_to_workspace: @@ -154,7 +154,7 @@ jobs: export JUST_TESTRUN=1 VARS_TO_PASS="-e PYTHON_VERSION -e BUILD_VERSION -e PYTORCH_VERSION -e CU_VERSION -e JUST_TESTRUN" - docker run --gpus all --ipc=host -v $(pwd):/remote -w /remote ${VARS_TO_PASS} ${TESTRUN_DOCKER_IMAGE} ./packaging/build_conda.sh + docker run --gpus all --ipc=host -v $(pwd):/remote -w /remote ${VARS_TO_PASS} ${TESTRUN_DOCKER_IMAGE} python3 ./packaging/build_conda.py binary_macos_wheel: <<: *binary_common diff --git a/.circleci/config.yml b/.circleci/config.yml index b4c014fb..3990661b 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -117,7 +117,7 @@ jobs: - run: name: build no_output_timeout: 20m - command: MAX_JOBS=15 TEST_FLAG=--no-test packaging/build_conda.sh + command: MAX_JOBS=15 TEST_FLAG=--no-test python3 packaging/build_conda.py - store_artifacts: path: /opt/conda/conda-bld/linux-64 - persist_to_workspace: @@ -154,7 +154,7 @@ jobs: export JUST_TESTRUN=1 VARS_TO_PASS="-e PYTHON_VERSION -e BUILD_VERSION -e PYTORCH_VERSION -e CU_VERSION -e JUST_TESTRUN" - docker run --gpus all --ipc=host -v $(pwd):/remote -w /remote ${VARS_TO_PASS} ${TESTRUN_DOCKER_IMAGE} ./packaging/build_conda.sh + docker run --gpus all --ipc=host -v $(pwd):/remote -w /remote ${VARS_TO_PASS} ${TESTRUN_DOCKER_IMAGE} python3 ./packaging/build_conda.py binary_macos_wheel: <<: *binary_common diff --git a/packaging/build_conda.py b/packaging/build_conda.py new file mode 100644 index 00000000..8be886d2 --- /dev/null +++ b/packaging/build_conda.py @@ -0,0 +1,134 @@ +# 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. + +import os.path +import runpy +import subprocess +from typing import List + +# required env vars: +# CU_VERSION: E.g. cu112 +# JUST_TESTRUN: 1 to not set nvcc flags +# PYTORCH_VERSION: e.g. 1.12.0 +# PYTHON_VERSION: e.g. 3.9 + +# should be run from pytorch3d root + +CU_VERSION = os.environ["CU_VERSION"] +PYTORCH_VERSION = os.environ["PYTORCH_VERSION"] +pytorch_major_minor = tuple(int(i) for i in PYTORCH_VERSION.split(".")[:2]) +source_root_dir = os.environ["PWD"] + + +def version_constraint(version): + """ + Given version "11.3" returns " >=11.3,<11.4" + """ + last_part = version.rindex(".") + 1 + upper = version[:last_part] + str(1 + int(version[last_part:])) + return f" >={version},<{upper}" + + +def get_cuda_major_minor(): + if CU_VERSION == "cpu": + raise ValueError("fn only for cuda builds") + if len(CU_VERSION) != 5 or CU_VERSION[:2] != "cu": + raise ValueError(f"Bad CU_VERSION {CU_VERSION}") + major = CU_VERSION[2:4] + minor = CU_VERSION[4] + return major, minor + + +def setup_cuda(): + if CU_VERSION == "cpu": + return + major, minor = get_cuda_major_minor() + os.environ["CUDA_HOME"] = f"/usr/local/cuda-{major}.{minor}/" + os.environ["FORCE_CUDA"] = "1" + + basic_nvcc_flags = ( + "-gencode=arch=compute_35,code=sm_35 " + "-gencode=arch=compute_50,code=sm_50 " + "-gencode=arch=compute_60,code=sm_60 " + "-gencode=arch=compute_70,code=sm_70 " + "-gencode=arch=compute_75,code=sm_75 " + "-gencode=arch=compute_50,code=compute_50" + ) + if CU_VERSION == "cu102": + nvcc_flags = basic_nvcc_flags + elif CU_VERSION == "cu110": + nvcc_flags = "-gencode=arch=compute_80,code=sm_80 " + basic_nvcc_flags + else: + nvcc_flags = ( + "-gencode=arch=compute_80,code=sm_80 " + + "-gencode=arch=compute_86,code=sm_86 " + + basic_nvcc_flags + ) + + if os.environ.get("JUST_TESTRUN", "0") != "1": + os.environ["NVCC_FLAGS"] = nvcc_flags + + +def setup_conda_pytorch_constraint() -> List[str]: + pytorch_constraint = f"- pytorch=={PYTORCH_VERSION}" + os.environ["CONDA_PYTORCH_CONSTRAINT"] = pytorch_constraint + os.environ["CONDA_PYTORCH_BUILD_CONSTRAINT"] = pytorch_constraint + os.environ["PYTORCH_VERSION_NODOT"] = PYTORCH_VERSION.replace(".", "") + + if pytorch_major_minor < (1, 13): + return ["-c", "pytorch"] + else: + return ["-c", "pytorch", "-c", "nvidia"] + + +def setup_conda_cudatoolkit_constraint(): + if CU_VERSION == "cpu": + os.environ["CONDA_CPUONLY_FEATURE"] = "- cpuonly" + os.environ["CONDA_CUDATOOLKIT_CONSTRAINT"] = "" + return + os.environ["CONDA_CPUONLY_FEATURE"] = "" + + if CU_VERSION in ("cu102", "cu110"): + os.environ["CONDA_CUB_CONSTRAINT"] = "- nvidiacub" + else: + os.environ["CONDA_CUB_CONSTRAINT"] = "" + + major, minor = get_cuda_major_minor() + version_clause = version_constraint(f"{major}.{minor}") + if pytorch_major_minor < (1, 13): + toolkit = f"- cudatoolkit {version_clause}" + else: + toolkit = f"- pytorch-cuda {version_clause}" + os.environ["CONDA_CUDATOOLKIT_CONSTRAINT"] = toolkit + + +def do_build(start_args: List[str]): + args = start_args.copy() + + test_flag = os.environ.get("TEST_FLAG") + if test_flag is not None: + args.append(test_flag) + + args.extend(["-c", "bottler", "-c", "fvcore", "-c", "iopath", "-c", "conda-forge"]) + args.append("--no-anaconda-upload") + args.extend(["--python", os.environ["PYTHON_VERSION"]]) + args.append("packaging/pytorch3d") + print(args) + subprocess.check_call(args) + + +if __name__ == "__main__": + args = ["conda", "build"] + setup_cuda() + + init_path = source_root_dir + "/pytorch3d/__init__.py" + build_version = runpy.run_path(init_path)["__version__"] + os.environ["BUILD_VERSION"] = build_version + + os.environ["SOURCE_ROOT_DIR"] = source_root_dir + args += setup_conda_pytorch_constraint() + setup_conda_cudatoolkit_constraint() + do_build(args) diff --git a/packaging/build_conda.sh b/packaging/build_conda.sh deleted file mode 100755 index c07572c2..00000000 --- a/packaging/build_conda.sh +++ /dev/null @@ -1,34 +0,0 @@ -#!/bin/bash -# 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. - -set -ex - -script_dir="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )" -. "$script_dir/pkg_helpers.bash" - -VERSION=$(python -c "exec(open('${script_dir}/../pytorch3d/__init__.py').read()); print(__version__)") - -# Prevent dev tag in the version string. -export BUILD_VERSION=$VERSION - -export BUILD_TYPE=conda -setup_env "$VERSION" -export SOURCE_ROOT_DIR="$PWD" -setup_conda_pytorch_constraint -setup_conda_cudatoolkit_constraint -setup_visual_studio_constraint - -if [[ "$JUST_TESTRUN" == "1" ]] -then - # We are not building for other users, we - # are only trying to see if the tests pass. - # So save time by only building for our own GPU. - unset NVCC_FLAGS -fi - -# shellcheck disable=SC2086 -conda build $CONDA_CHANNEL_FLAGS ${TEST_FLAG:-} -c bottler -c fvcore -c iopath -c conda-forge --no-anaconda-upload --python "$PYTHON_VERSION" packaging/pytorch3d