Fix build deps (#116)

* Fix build deps * ixup * More fixes * pypi jax * fixup * Cleanup * Fix case insensitive file system * Fixup deps * fix hashes * Fix openstack * fix * More cleanup * Cleanup * GPU results * fixes * fix * Update llama.py * fix * Fix * fix * Update pipeline.yml * Update gpu_pipeline.yml * Update gpu_pipeline.yml * Update pipeline.yml
EnzymeAD · Aug 25, 2024 · 61281ad · 61281ad
1 parent 928b66f
commit 61281ad
Show file tree

Hide file tree

Showing 26 changed files with 1,634 additions and 752 deletions.
diff --git a/.buildkite/gpu_pipeline.yml b/.buildkite/gpu_pipeline.yml
@@ -16,19 +16,15 @@ steps:
 
       mv bazel* .local/bin/bazel
       chmod +x .local/bin/bazel
-      wget https://repo.anaconda.com/miniconda/Miniconda3-py310_23.3.1-0-Linux-x86_64.sh
-      chmod +x Miniconda*.sh
-      ./Miniconda*.sh -b -p `pwd`/conda
-      rm Miniconda*.sh
 
-      python -m ensurepip --upgrade
-      python -m pip install --user numpy wheel
       mkdir -p .baztmp
-      rm -f bazel-bin/*.whl
-      HERMETIC_PYTHON_VERSION=`python -c "import sys; print('{0[0]}.{0[1]}'.format(sys.version_info))"` bazel --output_user_root=`pwd`/.baztmp build --python_path=`which python` --define=no_nccl_support=true :enzyme_ad
-      cp bazel-bin/*.whl .
-      python -m pip install --user *.whl "jax[cuda12]"
 
       echo "--- :python: Test"
 
-      HERMETIC_PYTHON_VERSION=`python -c "import sys; print('{0[0]}.{0[1]}'.format(sys.version_info))"` bazel --output_user_root=`pwd`/.baztmp test --python_path=`which python` --test_output=errors //test/...
+      HERMETIC_PYTHON_VERSION="3.12" bazel --output_user_root=`pwd`/.baztmp test --test_output=errors //test/...
+      HERMETIC_PYTHON_VERSION="3.12" bazel --output_user_root=`pwd`/.baztmp test --cache_test_results=no //test:bench_vs_xla
+      HERMETIC_PYTHON_VERSION="3.12" bazel --output_user_root=`pwd`/.baztmp test --cache_test_results=no //test:llama
+      cat bazel-out/*/testlogs/test/llama/test.log
+    artifact_paths:
+      - "bazel-out/*/testlogs/test/llama/test.log"
+      - "bazel-out/*/testlogs/test/llama/bench_vs_xla.log"
diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
@@ -1,12 +1,14 @@
 steps:
-  - name: "CI {{matrix.arch}} -- {{matrix.os}}"
+  - name: "CI {{matrix.arch}} -- {{matrix.os}} python {{matrix.python}}"
     matrix:
       setup:
         arch:
           - aarch64
           - x86_64
         os:
           - macos
+        python:
+          - "3.12"
     agents:
       queue: "juliaecosystem"
       os: "{{matrix.os}}"
@@ -30,52 +32,31 @@ steps:
       chmod +x .local/bin/md5
 
       if [ "{{matrix.os}}" == "macos" ]; then
-        if [ "{{matrix.arch}}" == "aarch64" ]; then
-          sed -i.bak 's~targets = \[.*\]~targets = \[\"AArch64\", \"AMDGPU\"]~g' WORKSPACE
-          curl -fLO https://repo.anaconda.com/miniconda/Miniconda3-py310_23.3.1-0-MacOSX-arm64.sh
-        else
-          sed -i.bak 's~targets = \[.*\]~targets = \[\"X86\", \"AMDGPU\"]~g' WORKSPACE
-          curl -fLO https://repo.anaconda.com/miniconda/Miniconda3-py310_23.3.1-0-MacOSX-{{matrix.arch}}.sh
-        fi
         curl -fLO "https://github.com/bazelbuild/bazelisk/releases/download/v1.19.0/bazelisk-darwin"
         mv bazelisk-darwin .local/bin/bazel
         chmod +x .local/bin/bazel
-        chmod +x Miniconda*.sh
-        ./Miniconda*.sh -b -p `pwd`/conda
-        rm Miniconda*.sh
       elif [ "{{matrix.os}}" == "linux" ]; then
-        if [ "{{matrix.arch}}" == "aarch64" ]; then
-          curl -fLO https://github.com/bazelbuild/bazelisk/releases/download/v1.19.0/bazelisk-linux-arm64
-        else
-          curl -fLO https://github.com/bazelbuild/bazelisk/releases/download/v1.19.0/bazelisk-linux-amd64
-        fi
         mv bazel* .local/bin/bazel
         chmod +x .local/bin/bazel
-        wget https://repo.anaconda.com/miniconda/Miniconda3-py310_23.3.1-0-Linux-{{matrix.arch}}.sh
-        chmod +x Miniconda*.sh
-        ./Miniconda*.sh -b -p `pwd`/conda
-        rm Miniconda*.sh
       else
-        wget https://repo.anaconda.com/miniconda/Miniconda3-py310_23.3.1-0-Windows-{{matrix.arch}}.exe
         if [ "{{matrix.arch}}" == "aarch64" ]; then
           wget https://github.com/bazelbuild/bazel/releases/download/6.2.1/bazel-6.2.1-windows-arm64.exe
         else
           wget https://github.com/bazelbuild/bazel/releases/download/6.2.1/bazel-6.2.1-windows-x86_64.exe
         fi
         mv bazel* .local/bin/bazel.exe
-        start /wait "" Miniconda3*.exe /InstallationType=JustMe /RegisterPython=0 /S /D=`pwd`/conda
-        rm Miniconda*.exe
       fi
-      # conda install -c conda-forge cxx-compiler -y
-      python -m ensurepip --upgrade
-      python -m pip install --user numpy wheel
       mkdir -p .baztmp
+      HERMETIC_PYTHON_VERSION={{matrix.python}} bazel --output_user_root=`pwd`/.baztmp test --test_output=errors //test/...
+      HERMETIC_PYTHON_VERSION={{matrix.python}} bazel --output_user_root=`pwd`/.baztmp test --cache_test_results=no //test:bench_vs_xla
+      HERMETIC_PYTHON_VERSION={{matrix.python}} bazel --output_user_root=`pwd`/.baztmp test --cache_test_results=no //test:llama
+      cat bazel-out/*/testlogs/test/llama/test.log
       rm -f bazel-bin/*.whl
-      HERMETIC_PYTHON_VERSION=`python -c "import sys; print('{0[0]}.{0[1]}'.format(sys.version_info))"` bazel --output_user_root=`pwd`/.baztmp build --python_path=`which python` --define=no_nccl_support=true :enzyme_ad
+      HERMETIC_PYTHON_VERSION={{matrix.python}} bazel --output_user_root=`pwd`/.baztmp build :wheel
       cp bazel-bin/*.whl .
-      python -m pip install --user *.whl "jax[cpu]"
-      HERMETIC_PYTHON_VERSION=`python -c "import sys; print('{0[0]}.{0[1]}'.format(sys.version_info))"` bazel --output_user_root=`pwd`/.baztmp test --python_path=`which python` --test_output=errors //test/...
     artifact_paths:
       - "*.whl"
+      - "bazel-out/*/testlogs/test/llama/test.log"
+      - "bazel-out/*/testlogs/test/llama/bench_vs_xla.log"
 
     timeout_in_minutes: 180
diff --git a/.buildkite/secure_pipeline.yml b/.buildkite/secure_pipeline.yml
@@ -1,5 +1,5 @@
 steps:
-  - name: "Tag {{matrix.arch}} -- {{matrix.os}}"
+  - name: "Tag {{matrix.arch}} -- {{matrix.os}} python {{matrix.python}}"
     matrix:
       setup:
         arch:
@@ -30,19 +30,9 @@ steps:
       chmod +x .local/bin/md5
 
       if [ "{{matrix.os}}" == "macos" ]; then
-        if [ "{{matrix.arch}}" == "aarch64" ]; then
-          sed -i.bak 's~targets = \[.*\]~targets = \[\"AArch64\", \"AMDGPU\"]~g' WORKSPACE
-          curl -fLO https://repo.anaconda.com/miniconda/Miniconda3-py3`echo {{matrix.python}} | cut -c 3-`_24.7.1-0-MacOSX-arm64.sh
-        else
-          sed -i.bak 's~targets = \[.*\]~targets = \[\"X86\", \"AMDGPU\"]~g' WORKSPACE
-          curl -fLO https://repo.anaconda.com/miniconda/Miniconda3-py3`echo {{matrix.python}} | cut -c 3-`_24.7.1-0-MacOSX-{{matrix.arch}}.sh
-        fi
         curl -fLO "https://github.com/bazelbuild/bazelisk/releases/download/v1.19.0/bazelisk-darwin"
         mv bazelisk-darwin .local/bin/bazel
         chmod +x .local/bin/bazel
-        chmod +x Miniconda*.sh
-        ./Miniconda*.sh -b -p `pwd`/conda
-        rm Miniconda*.sh
       elif [ "{{matrix.os}}" == "linux" ]; then      
         if [ "{{matrix.arch}}" == "aarch64" ]; then
           curl -fLO https://github.com/bazelbuild/bazelisk/releases/download/v1.19.0/bazelisk-linux-arm64
@@ -51,25 +41,20 @@ steps:
         fi
         mv bazel* .local/bin/bazel
         chmod +x .local/bin/bazel
-        wget https://repo.anaconda.com/miniconda/Miniconda3-py3`echo {{matrix.python}} | cut -c 3-`_24.7.1-0-Linux-{{matrix.arch}}.sh
-        chmod +x Miniconda*.sh
-        ./Miniconda*.sh -b -p `pwd`/conda
-        rm Miniconda*.sh
       else
-        wget https://repo.anaconda.com/miniconda/Miniconda3-py3`echo {{matrix.python}} | cut -c 3-`_24.7.1-0-Windows-{{matrix.arch}}.exe
         if [ "{{matrix.arch}}" == "aarch64" ]; then
           wget https://github.com/bazelbuild/bazel/releases/download/6.2.1/bazel-6.2.1-windows-arm64.exe
         else
           wget https://github.com/bazelbuild/bazel/releases/download/6.2.1/bazel-6.2.1-windows-x86_64.exe
         fi
         mv bazel* .local/bin/bazel.exe
       fi
-      python -m ensurepip --upgrade
-      python -m pip install --user numpy wheel
       mkdir baztmp
       export TAG=`echo $BUILDKITE_TAG | cut -c2-`
       sed -i.bak "s~version = \"[0-9.]*\"~version = \"\$TAG\"~g" BUILD
-      HERMETIC_PYTHON_VERSION={{matrix.python}} bazel --output_user_root=`pwd`/baztmp build --define=no_nccl_support=true :enzyme_ad
+      HERMETIC_PYTHON_VERSION={{matrix.python}} bazel --output_user_root=`pwd`/baztmp build @llvm-project//llvm:FileCheck
+      rm bazel-bin/*.whl
+      HERMETIC_PYTHON_VERSION={{matrix.python}} bazel --output_user_root=`pwd`/baztmp build :wheel
       cp bazel-bin/*.whl .
       python -m pip install *.whl
       python -m pip install --user twine

diff --git a/.buildkite/secure_pipeline.yml.signature b/.buildkite/secure_pipeline.yml.signature
@@ -1 +1 @@
-Salted__Q�"���=x��:�W�ԇ.�<+���9�	E�g�h�d;�h?}��e�^Vr�?n�T�Wdˤ��N�.T�A��Y=q#V:��M7�
+Salted__�F�	�܌ɈLY�o�j�vQ��Z��4�*w��5�5W��y#����WY΀ۇ��Ѝ�86��9����Ř;��I�Q�]PvkU'�b�IJ
diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
@@ -9,22 +9,23 @@ on:
 
 jobs:
   build:
-    name: Build ${{ matrix.os }}
+    name: Build ${{ matrix.os }} python ${{ matrix.python }}
     runs-on: ${{ matrix.os }}
     strategy:
       fail-fast: false
       matrix:
         os: [openstack22]
+        python: ["3.12"]
     timeout-minutes: 500
     steps:
       - name: add llvm
         run: |
             if [ "$(expr substr $(uname -s) 1 5)" == "Linux" ]; then
               sudo apt-get update
-              sudo apt-get install -y git gcc g++ python3 python3-dev python3-pip
-              sudo python3 -m pip install --upgrade lit
+              sudo apt-get install -y git gcc g++ #python3 python3-dev python3-pip
+              #sudo python3 -m pip install --upgrade lit
+              #- run: python3 -m pip install --user numpy
             fi
-      - run: python3 -m pip install --user numpy
       - uses: actions/checkout@v3
         with:
           submodules: recursive
@@ -42,20 +43,15 @@ jobs:
           repository-cache: true
           bazelisk-version: 1.x
 
-      - run: |
-          HERMETIC_PYTHON_VERSION=`python3 -c "import sys; print('{0[0]}.{0[1]}'.format(sys.version_info))"` bazel build @llvm-project//llvm:FileCheck
-          sudo rm bazel-bin/*.whl || echo
-          HERMETIC_PYTHON_VERSION=`python3 -c "import sys; print('{0[0]}.{0[1]}'.format(sys.version_info))"` bazel build :enzyme_ad
-
-      - run: cp bazel-bin/*.whl .
-
       - name: test
         run: |
-          ls -all .
-          ls -all bazel-bin
-          python3 -m pip uninstall enzyme-ad -y || echo
-          python3 -m pip install --user --force-reinstall "jax[cpu]" *.whl
-          HERMETIC_PYTHON_VERSION=`python3 -c "import sys; print('{0[0]}.{0[1]}'.format(sys.version_info))"` bazel test --test_output=errors ...
+          HERMETIC_PYTHON_VERSION=${{ matrix.python }} bazel test --test_output=errors ...
+          
+      - name: Build Wheel
+        run: |
+          sudo rm bazel-bin/*.whl || echo
+          HERMETIC_PYTHON_VERSION=${{ matrix.python }} bazel build :wheel
+          cp bazel-bin/*.whl .
 
       - name: Upload Build
         uses: actions/upload-artifact@v3

diff --git a/.github/workflows/tag.yml b/.github/workflows/tag.yml
@@ -6,7 +6,7 @@ on:
 
 jobs:
   build:
-    name: Build ${{ matrix.os }} python ${{ matrix.python }}
+    name: Tag ${{ matrix.os }} python ${{ matrix.python }}
     runs-on: ${{ matrix.os }}
     strategy:
       fail-fast: false
@@ -45,7 +45,7 @@ jobs:
       - run: |
           HERMETIC_PYTHON_VERSION=${{ matrix.python }} bazel build @llvm-project//llvm:FileCheck
           sudo rm bazel-bin/*.whl || echo
-          HERMETIC_PYTHON_VERSION=${{ matrix.python }} bazel build :enzyme_ad
+          HERMETIC_PYTHON_VERSION=${{ matrix.python }} bazel build :wheel
       - env:
           TWINE_USERNAME: __token__
           TWINE_PASSWORD: ${{ secrets.PYPI_TOKEN }}

diff --git a/BUILD b/BUILD
@@ -66,8 +66,22 @@ cc_binary(
     ],
 )
 
-py_wheel(
+py_library(
     name = "enzyme_ad",
+    visibility = ["//visibility:public"],
+    deps = [
+        "@pypi_jax//:pkg",
+        "@pypi_absl_py//:pkg",
+    ],
+    imports=["src"],
+    data = [
+        "//:enzyme_jax_data",
+        "//src/enzyme_ad/jax:enzyme_jax_internal",
+    ]
+)
+
+py_wheel(
+    name = "wheel",
     author = "Enzyme Authors",
     author_email = "[email protected], [email protected]",
     distribution = "enzyme_ad",

diff --git a/README.md b/README.md
@@ -44,13 +44,13 @@ pip install enzyme-ad
 
 ## Building from source
 
-Requirements: `bazel-6.2.1`, `clang++`, `python`, `python-virtualenv`,
+Requirements: `bazel-6.5`, `clang++`, `python`, `python-virtualenv`,
 `python3-dev`.
 
 Build our extension with:
 ```sh
 # Will create a whl in bazel-bin/enzyme_ad-VERSION-SYSTEM.whl
-bazel build :enzyme_ad
+bazel build :wheel
 ```
 
 Finally, install the built library with:
@@ -61,6 +61,12 @@ Note that you cannot run code from the root of the git directory. For instance,
 
 ## Running the test
 
+To run tests, you can simply execute the following bazel commands (this does not require building or installing the wheel).
+```sh
+bazel test //test/...
+```
+
+Alternatively, if you have installed the wheel, you can manually invoke the tests as follows
 ```sh
 cd test && python test.py
 ```
diff --git a/WORKSPACE b/WORKSPACE
@@ -41,12 +41,13 @@ python_init_rules()
 load("@xla//third_party/py:python_init_repositories.bzl", "python_init_repositories")
 python_init_repositories(
     requirements = {
-        "3.9": "//build:requirements_lock_3_9.txt",
-        "3.10": "//build:requirements_lock_3_10.txt",
-        "3.11": "//build:requirements_lock_3_11.txt",
-        "3.12": "//build:requirements_lock_3_12.txt",
-        "3.13": "//build:requirements_lock_3_13.txt",
+        "3.10": "//builddeps:requirements_lock_3_10.txt",
+        "3.11": "//builddeps:requirements_lock_3_11.txt",
+        "3.12": "//builddeps:requirements_lock_3_12.txt",
     },
+    local_wheel_inclusion_list = [
+        "enzyme_ad*",
+    ]
 )
 
 load("@xla//third_party/py:python_init_toolchains.bzl", "python_init_toolchains")
@@ -55,16 +56,8 @@ python_init_toolchains()
 load("@xla//third_party/py:python_init_pip.bzl", "python_init_pip")
 python_init_pip()
 
-load("@xla//third_party/py:python_init_rules.bzl", "python_init_rules")
-python_init_rules()
-
-load("@rules_python//python:repositories.bzl", "py_repositories")
-
-py_repositories()
-
-load("@rules_python//python/pip_install:repositories.bzl", "pip_install_dependencies")
-
-pip_install_dependencies()
+load("@pypi//:requirements.bzl", "install_deps")
+install_deps()
 
 http_archive(
     name = "enzyme",

diff --git a/builddeps/BUILD b/builddeps/BUILD
@@ -0,0 +1,17 @@
+licenses(["notice"])
+
+load("@python//:defs.bzl", "compile_pip_requirements")
+load("@python_version_repo//:py_version.bzl", "REQUIREMENTS")
+
+compile_pip_requirements(
+    name = "requirements",
+    extra_args = [
+        "--allow-unsafe",
+        "--build-isolation",
+        "--rebuild",
+    ],
+    requirements_in = "requirements.in",
+    requirements_txt = REQUIREMENTS,
+    generate_hashes = True,
+    data = ["test-requirements.txt"]
+)
diff --git a/builddeps/requirements.in b/builddeps/requirements.in
@@ -0,0 +1,8 @@
+#
+# test deps
+#
+-r test-requirements.txt
+
+jax >= 0.4.21
+jaxlib >= 0.4.21
+absl_py >= 2.0.0
Original file line number	Diff line number	Diff line change
		@@ -1 +1 @@
		Salted__Q�"��=x��:�W�ԇ.�<+��9� E�g�h�d;�h?}��e�^Vr�?n�T�Wdˤ��N�.T�A��Y=q#V:��M7�
		Salted__�F� �܌ɈLY�o�j�vQ��Z��4�*w��5�5W��y#��WY΀ۇ��Ѝ�86��9��Ř;��I�Q�]PvkU'�b�IJ