Build wheels for s3

Summary: For Linux, instead of uploading wheels to PyPI which will only work with one particular version of PyTorch and CUDA, from the next release we will store a range of built wheels on S3.

Reviewed By: nikhilaravi

Differential Revision: D26209398

fbshipit-source-id: 945a6907b78807e1eedb25007f87f90bbf59f80e

packaging/linux_wheels/README.md

@@ -0,0 +1,29 @@
+## Building Linux pip Packages
+
+1. Make sure this directory is on a filesystem which docker can
+use - e.g. not NFS. If you are using a local hard drive there is
+nothing to do here.
+
+2. You may want to `docker pull pytorch/conda-cuda:latest`.
+
+3. Run `bash go.sh` in this directory. This takes ages
+and writes packages to `inside/output`.
+
+4. You can upload the packages to s3, along with basic html files
+which enable them to be used, with `bash after.sh`.
+
+
+In particular, if you are in a jupyter/colab notebook you can
+then install using these wheels with the following series of
+commands.
+
+```
+import sys
+import torch
+version_str="".join([
+    f"py3{sys.version_info.minor}_cu",
+    torch.version.cuda.replace(".",""),
+    f"_pyt{torch.__version__[0:5:2]}"
+])
+!pip install pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/{version_str}/download.html
+```

packaging/linux_wheels/after.sh

@@ -0,0 +1,5 @@
+#!/usr/bin/bash
+# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
+set -ex
+sudo chown -R "$USER" output
+python publish.py

packaging/linux_wheels/go.sh

@@ -0,0 +1,3 @@
+#!/usr/bin/bash
+# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
+sudo docker run --rm  -v "$PWD/../../:/inside" pytorch/conda-cuda bash inside/packaging/linux_wheels/inside.sh

packaging/linux_wheels/inside.sh

@@ -0,0 +1,102 @@
+#!/bin/bash
+# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
+set -ex
+
+conda init bash
+# shellcheck source=/dev/null
+source ~/.bashrc
+
+cd /inside
+VERSION=$(python -c "exec(open('pytorch3d/__init__.py').read()); print(__version__)")
+
+export BUILD_VERSION=$VERSION
+export FORCE_CUDA=1
+
+wget --no-verbose https://github.com/NVIDIA/cub/archive/1.10.0.tar.gz
+tar xzf 1.10.0.tar.gz
+CUB_HOME=$(realpath ./cub-1.10.0)
+export CUB_HOME
+echo "CUB_HOME is now $CUB_HOME"
+
+
+PYTHON_VERSIONS="3.6 3.7 3.8 3.9"
+# the keys are pytorch versions
+declare -A CONDA_CUDA_VERSIONS=(
+#    ["1.4.0"]="cu101"
+#    ["1.5.0"]="cu101 cu102"
+#    ["1.5.1"]="cu101 cu102"
+#    ["1.6.0"]="cu101 cu102"
+    ["1.7.0"]="cu101 cu102 cu110"
+    ["1.7.1"]="cu101 cu102 cu110"
+)
+
+
+
+for python_version in $PYTHON_VERSIONS
+do
+    for pytorch_version in "${!CONDA_CUDA_VERSIONS[@]}"
+    do
+        if [[ "3.6 3.7 3.8" != *$python_version* ]] && [[ "1.4.0 1.5.0 1.5.1 1.6.0 1.7.0" == *$pytorch_version* ]]
+        then
+            #python 3.9 and later not supported by pytorch 1.7.0 and before
+            continue
+        fi
+
+        if [[ "3.9" == "$python_version" ]]
+        then
+            extra_channel="-c conda-forge"
+        else
+            extra_channel=""
+        fi
+
+        for cu_version in ${CONDA_CUDA_VERSIONS[$pytorch_version]}
+        do
+            case "$cu_version" in
+                cu110)
+                    export CUDA_HOME=/usr/local/cuda-11.0/
+                    export CUDA_TAG=11.0
+                    export 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_80,code=sm_80 -gencode=arch=compute_50,code=compute_50"
+                ;;
+                cu102)
+                    export CUDA_HOME=/usr/local/cuda-10.2/
+                    export CUDA_TAG=10.2
+                    export 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"
+                ;;
+                cu101)
+                    export CUDA_HOME=/usr/local/cuda-10.1/
+                    export CUDA_TAG=10.1
+                    export 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"
+                ;;
+                *)
+                    echo "Unrecognized cu_version=$cu_version"
+                    exit 1
+                ;;
+            esac
+            tag=py"${python_version//./}"_"${cu_version}"_pyt"${pytorch_version//./}"
+
+            outdir="/inside/packaging/linux_wheels/output/$tag"
+            if [[ -d "$outdir" ]]
+            then
+                continue
+            fi
+
+            conda create -y -n "$tag" "python=$python_version"
+            conda activate "$tag"
+            conda install -y -c pytorch $extra_channel "pytorch=$pytorch_version" "cudatoolkit=$CUDA_TAG" torchvision
+            pip install fvcore iopath
+            echo "python version" "$python_version" "pytorch version" "$pytorch_version" "cuda version" "$cu_version" "tag" "$tag"
+
+            rm -rf dist
+
+            python setup.py clean
+            python setup.py bdist_wheel
+
+            rm -rf "$outdir"
+            mkdir -p "$outdir"
+            cp dist/*whl "$outdir"
+
+            conda deactivate
+        done
+    done
+done
+echo "DONE"

packaging/linux_wheels/publish.py

@@ -0,0 +1,76 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
+import os
+import subprocess
+from pathlib import Path
+from typing import List
+
+
+dest = "s3://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/"
+
+output = Path("output")
+
+
+def fs3cmd(args, allow_failure: bool = False) -> List[str]:
+    """
+    This function returns the args for subprocess to mimic the bash command
+    fs3cmd available in the fairusers_aws module on the FAIR cluster.
+    """
+    os.environ["FAIR_CLUSTER_NAME"] = os.environ["FAIR_ENV_CLUSTER"].lower()
+    cmd_args = ["/public/apps/fairusers_aws/bin/fs3cmd"] + args
+    return cmd_args
+
+
+def fs3_exists(path) -> bool:
+    """
+    Returns True if the path exists inside dest on S3.
+    In fact, will also return True if there is a file which has the given
+    path as a prefix, but we are careful about this.
+    """
+    out = subprocess.check_output(fs3cmd(["ls", path]))
+    return len(out) != 0
+
+
+def get_html_wrappers() -> None:
+    for directory in sorted(output.iterdir()):
+        output_wrapper = directory / "download.html"
+        assert not output_wrapper.exists()
+        dest_wrapper = dest + directory.name + "/download.html"
+        if fs3_exists(dest_wrapper):
+            subprocess.check_call(fs3cmd(["get", dest_wrapper, str(output_wrapper)]))
+
+
+def write_html_wrappers() -> None:
+    html = """
+    <a href="$">$</a><br>
+    """
+
+    for directory in sorted(output.iterdir()):
+        files = list(directory.glob("*.whl"))
+        assert len(files) == 1, files
+        [wheel] = files
+
+        this_html = html.replace("$", wheel.name)
+        output_wrapper = directory / "download.html"
+        if output_wrapper.exists():
+            contents = output_wrapper.read_text()
+            if this_html not in contents:
+                with open(output_wrapper, "a") as f:
+                    f.write(this_html)
+        else:
+            output_wrapper.write_text(this_html)
+
+
+def to_aws() -> None:
+    for directory in output.iterdir():
+        for file in directory.iterdir():
+            print(file)
+            subprocess.check_call(
+                fs3cmd(["put", str(file), dest + str(file.relative_to(output))])
+            )
+
+
+if __name__ == "__main__":
+    # Uncomment this for subsequent releases.
+    # get_html_wrappers()
+    write_html_wrappers()
+    to_aws()