From 9c586b1351df48a5e36461475bfb2d9ab961d749 Mon Sep 17 00:00:00 2001
From: bottler <bottler@users.noreply.github.com>
Date: Thu, 31 Oct 2024 08:41:20 -0700
Subject: [PATCH] Run tests in github action not circleci (#1896)

Summary: Pull Request resolved: https://github.com/facebookresearch/pytorch3d/pull/1896

Differential Revision: D65272512

Pulled By: bottler
---
 .github/workflows/build.yml   | 20 ++++++++++++++++++++
 packaging/build_conda.py      | 35 ++++++++++++++++++++++++++---------
 packaging/pytorch3d/meta.yaml |  9 +++++++++
 3 files changed, 55 insertions(+), 9 deletions(-)
 create mode 100644 .github/workflows/build.yml

diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
new file mode 100644
index 000000000..2b7b5d27d
--- /dev/null
+++ b/.github/workflows/build.yml
@@ -0,0 +1,20 @@
+name: facebookresearch/pytorch3d/build_and_test
+on:
+  pull_request:
+    branches:
+      - main
+jobs:
+  binary_linux_conda_cuda:
+    runs-on: 4-core-ubuntu-gpu-t4
+    env:
+      PYTHON_VERSION: "3.12"
+      BUILD_VERSION: "${{ github.run_number }}"
+      PYTORCH_VERSION: "2.4.1"
+      CU_VERSION: "cu121"
+      JUST_TESTRUN: 1
+    steps:
+    - uses: actions/checkout@v4
+    - name: Build and run tests
+      run: |-
+        conda create --name env --yes --quiet conda-build
+        conda run --no-capture-output --name env python3 ./packaging/build_conda.py --use-conda-cuda
diff --git a/packaging/build_conda.py b/packaging/build_conda.py
index 554e86c6d..cc0c9521b 100644
--- a/packaging/build_conda.py
+++ b/packaging/build_conda.py
@@ -4,10 +4,11 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+import argparse
 import os.path
 import runpy
 import subprocess
-from typing import List
+from typing import List, Tuple
 
 # required env vars:
 # CU_VERSION: E.g. cu112
@@ -23,7 +24,7 @@
 source_root_dir = os.environ["PWD"]
 
 
-def version_constraint(version):
+def version_constraint(version) -> str:
     """
     Given version "11.3" returns " >=11.3,<11.4"
     """
@@ -32,7 +33,7 @@ def version_constraint(version):
     return f" >={version},<{upper}"
 
 
-def get_cuda_major_minor():
+def get_cuda_major_minor() -> Tuple[str, str]:
     if CU_VERSION == "cpu":
         raise ValueError("fn only for cuda builds")
     if len(CU_VERSION) != 5 or CU_VERSION[:2] != "cu":
@@ -42,11 +43,10 @@ def get_cuda_major_minor():
     return major, minor
 
 
-def setup_cuda():
+def setup_cuda(use_conda_cuda: bool) -> List[str]:
     if CU_VERSION == "cpu":
-        return
+        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 = (
@@ -75,6 +75,15 @@ def setup_cuda():
 
     if os.environ.get("JUST_TESTRUN", "0") != "1":
         os.environ["NVCC_FLAGS"] = nvcc_flags
+    if use_conda_cuda:
+        os.environ["CONDA_CUDA_TOOLKIT_BUILD_CONSTRAINT1"] = "- cuda-toolkit"
+        os.environ["CONDA_CUDA_TOOLKIT_BUILD_CONSTRAINT2"] = (
+            f"- cuda-version={major}.{minor}"
+        )
+        return ["-c", f"nvidia/label/cuda-{major}.{minor}.0"]
+    else:
+        os.environ["CUDA_HOME"] = f"/usr/local/cuda-{major}.{minor}/"
+        return []
 
 
 def setup_conda_pytorch_constraint() -> List[str]:
@@ -95,7 +104,7 @@ def setup_conda_pytorch_constraint() -> List[str]:
         return ["-c", "pytorch", "-c", "nvidia"]
 
 
-def setup_conda_cudatoolkit_constraint():
+def setup_conda_cudatoolkit_constraint() -> None:
     if CU_VERSION == "cpu":
         os.environ["CONDA_CPUONLY_FEATURE"] = "- cpuonly"
         os.environ["CONDA_CUDATOOLKIT_CONSTRAINT"] = ""
@@ -116,7 +125,7 @@ def setup_conda_cudatoolkit_constraint():
     os.environ["CONDA_CUDATOOLKIT_CONSTRAINT"] = toolkit
 
 
-def do_build(start_args: List[str]):
+def do_build(start_args: List[str]) -> None:
     args = start_args.copy()
 
     test_flag = os.environ.get("TEST_FLAG")
@@ -132,8 +141,16 @@ def do_build(start_args: List[str]):
 
 
 if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Build the conda package.")
+    parser.add_argument(
+        "--use-conda-cuda",
+        action="store_true",
+        help="get cuda from conda ignoring local cuda",
+    )
+    our_args = parser.parse_args()
+
     args = ["conda", "build"]
-    setup_cuda()
+    args += setup_cuda(use_conda_cuda=our_args.use_conda_cuda)
 
     init_path = source_root_dir + "/pytorch3d/__init__.py"
     build_version = runpy.run_path(init_path)["__version__"]
diff --git a/packaging/pytorch3d/meta.yaml b/packaging/pytorch3d/meta.yaml
index 9f115ff99..a268a24b9 100644
--- a/packaging/pytorch3d/meta.yaml
+++ b/packaging/pytorch3d/meta.yaml
@@ -8,10 +8,13 @@ source:
 requirements:
   build:
     - {{ compiler('c') }} # [win]
+    {{ environ.get('CONDA_CUDA_TOOLKIT_BUILD_CONSTRAINT1', '') }}
+    {{ environ.get('CONDA_CUDA_TOOLKIT_BUILD_CONSTRAINT2', '') }}
     {{ environ.get('CONDA_CUB_CONSTRAINT') }}
 
   host:
     - python
+    - mkl =2023  # [x86_64]
     {{ environ.get('SETUPTOOLS_CONSTRAINT') }}
     {{ environ.get('CONDA_PYTORCH_BUILD_CONSTRAINT') }}
     {{ environ.get('CONDA_PYTORCH_MKL_CONSTRAINT') }}
@@ -22,12 +25,14 @@ requirements:
     - python
     - numpy >=1.11
     - torchvision >=0.5
+    - mkl =2023  # [x86_64]
     - iopath
     {{ environ.get('CONDA_PYTORCH_CONSTRAINT') }}
     {{ environ.get('CONDA_CUDATOOLKIT_CONSTRAINT') }}
 
 build:
   string: py{{py}}_{{ environ['CU_VERSION'] }}_pyt{{ environ['PYTORCH_VERSION_NODOT']}}
+  # script: LD_LIBRARY_PATH=$PREFIX/lib:$BUILD_PREFIX/lib:$LD_LIBRARY_PATH python setup.py install --single-version-externally-managed --record=record.txt # [not win]
   script: python setup.py install --single-version-externally-managed --record=record.txt # [not win]
   script_env:
     - CUDA_HOME
@@ -47,6 +52,10 @@ test:
     - imageio
     - hydra-core
     - accelerate
+    - matplotlib
+    - tabulate
+    - pandas
+    - sqlalchemy
   commands:
     #pytest .
     python -m unittest discover -v -s tests -t .