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

@@ -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

.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

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)

packaging/build_conda.sh

@@ -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