Move BUILD_TF/TORCH_OPS settings to github actions file

Enable users to build wheel with tfops locally with docker container.

.github/workflows/ubuntu-wheel.yml

@@ -19,6 +19,9 @@ concurrency:
 
 env:
   GCE_CLI_GHA_VERSION: '416.0.0'      # Fixed to avoid dependency on API changes
+  BUILD_CUDA_MODULE: 'ON'
+  BUILD_PYTORCH_OPS: 'ON'
+  BUILD_TENSORFLOW_OPS: 'OFF'   # Turn ON when cxx11_abi is same for TF and PyTorch
 
 jobs:
   build-wheel:

docker/Dockerfile.wheel

@@ -8,14 +8,18 @@ ARG CCACHE_TAR_NAME
 ARG CMAKE_VERSION
 ARG CCACHE_VERSION
 ARG PYTHON_VERSION
+ARG BUILD_TENSORFLOW_OPS
+ARG BUILD_PYTORCH_OPS
 
 # Forward all ARG to ENV
 # ci_utils.sh requires these environment variables
 ENV DEVELOPER_BUILD=${DEVELOPER_BUILD}
 ENV CCACHE_TAR_NAME=${CCACHE_TAR_NAME}
 ENV CMAKE_VERSION=${CMAKE_VERSION}
-ARG CCACHE_VERSION=${CCACHE_VERSION}
+ENV CCACHE_VERSION=${CCACHE_VERSION}
 ENV PYTHON_VERSION=${PYTHON_VERSION}
+ENV BUILD_PYTORCH_OPS=${BUILD_PYTORCH_OPS}
+ENV BUILD_TENSORFLOW_OPS=${BUILD_TENSORFLOW_OPS}
 
 # Prevent interactive inputs when installing packages
 ENV DEBIAN_FRONTEND=noninteractive
@@ -123,11 +127,8 @@ COPY . /root/Open3D
 WORKDIR /root/Open3D
 
 # Build python wheel
- # TODO: re-enable tensorflow support, off due to due to cxx11_abi issue with PyTorch
 RUN export NPROC=$(nproc) \
  && export BUILD_SHARED_LIBS=OFF \
- && export BUILD_TENSORFLOW_OPS=OFF \
- && export BUILD_PYTORCH_OPS=ON \
  && source /root/Open3D/util/ci_utils.sh \
  && build_pip_package build_azure_kinect build_jupyter
 

docker/docker_build.sh

@@ -193,6 +193,8 @@ cuda_wheel_build() {
     fi
     echo "[cuda_wheel_build()] PYTHON_VERSION: ${PYTHON_VERSION}"
     echo "[cuda_wheel_build()] DEVELOPER_BUILD: ${DEVELOPER_BUILD}"
+    echo "[cuda_wheel_build()] BUILD_TENSORFLOW_OPS=${BUILD_TENSORFLOW_OPS:?'env var must be set.'}"
+    echo "[cuda_wheel_build()] BUILD_PYTORCH_OPS=${BUILD_PYTORCH_OPS:?'env var must be set.'}"
 
     pushd "${HOST_OPEN3D_ROOT}"
     docker build \
@@ -203,6 +205,8 @@ cuda_wheel_build() {
         --build-arg CMAKE_VERSION="${CMAKE_VERSION}" \
         --build-arg CCACHE_VERSION="${CCACHE_VERSION}" \
         --build-arg PYTHON_VERSION="${PYTHON_VERSION}" \
+        --build-arg BUILD_TENSORFLOW_OPS="${BUILD_TENSORFLOW_OPS}" \
+        --build-arg BUILD_PYTORCH_OPS="${BUILD_PYTORCH_OPS}" \
         -t open3d-ci:wheel \
         -f docker/Dockerfile.wheel .
     popd