Skip to content

Commit

Permalink
Add a script for running stablehlo tests. (#5360)
Browse files Browse the repository at this point in the history
* Add kokoro presubmit for stablehlo tests
  • Loading branch information
qihqi authored and will-cromar committed Sep 14, 2023
1 parent c49daaf commit e7ae1d4
Show file tree
Hide file tree
Showing 4 changed files with 68 additions and 26 deletions.
27 changes: 1 addition & 26 deletions .kokoro/build.sh
Original file line number Diff line number Diff line change
Expand Up @@ -12,32 +12,7 @@ source ${XLA_DIR}/.kokoro/common.sh

TORCHVISION_COMMIT="$(cat $PYTORCH_DIR/.github/ci_commit_pins/vision.txt)"

apt-get clean && apt-get update
apt-get upgrade -y
apt-get install --fix-missing -y python3-pip git curl libopenblas-dev vim jq \
apt-transport-https ca-certificates procps openssl sudo wget libssl-dev libc6-dbg

curl -LO https://github.com/bazelbuild/bazelisk/releases/download/v1.11.0/bazelisk-linux-amd64
mv bazelisk-linux-amd64 /usr/local/bin/bazel
chmod +x /usr/local/bin/bazel

pip install mkl mkl-include setuptools typing_extensions cmake requests
sudo ln -s /usr/local/lib/libmkl_intel_lp64.so.2 /usr/local/lib/libmkl_intel_lp64.so.1
sudo ln -s /usr/local/lib/libmkl_intel_thread.so.2 /usr/local/lib/libmkl_intel_thread.so.1
sudo ln -s /usr/local/lib/libmkl_core.so.2 /usr/local/lib/libmkl_core.so.1
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib/

echo "deb [signed-by=/usr/share/keyrings/cloud.google.gpg] https://packages.cloud.google.com/apt cloud-sdk main" >> /etc/apt/sources.list.d/google-cloud-sdk.list
curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key --keyring /usr/share/keyrings/cloud.google.gpg add -

# TODO(yeounoh) fix `GoogleCredentials` import error
apt-get update
apt-get -y install google-cloud-cli
pip install --upgrade google-api-python-client
pip install --upgrade oauth2client
pip install --upgrade google-cloud-storage
pip install lark-parser
pip install cloud-tpu-client
install_environments

pip install --user https://storage.googleapis.com/tpu-pytorch/wheels/tpuvm/torch-nightly-cp38-cp38-linux_x86_64.whl \
https://storage.googleapis.com/tpu-pytorch/wheels/tpuvm/torchvision-nightly-cp38-cp38-linux_x86_64.whl \
Expand Down
31 changes: 31 additions & 0 deletions .kokoro/build_and_run_stablehlo_tests.sh
Original file line number Diff line number Diff line change
@@ -0,0 +1,31 @@
#!/bin/bash

set -ex

DEBIAN_FRONTEND=noninteractive

PYTORCH_DIR="${KOKORO_ARTIFACTS_DIR}/github/pytorch"
XLA_DIR=$PYTORCH_DIR/xla
git clone --quiet https://github.com/pytorch/pytorch.git "$PYTORCH_DIR"
cp -r "${KOKORO_ARTIFACTS_DIR}/github/xla" "$XLA_DIR"
source ${XLA_DIR}/.kokoro/common.sh

TORCHVISION_COMMIT="$(cat $PYTORCH_DIR/.github/ci_commit_pins/vision.txt)"

install_environments

set +e
# build pytorch
pushd $PYTORCH_DIR
pip install --user .
popd
build_torch_xla $XLA_DIR
pushd $XLA_DIR
LOGFILE=/tmp/pytorch_py_test.log
TEST_DIR=test/stablehlo
for FILE in $(ls $TEST_DIR); do
python $TEST_DIR/$FILE | tee $LOGFILE
done



29 changes: 29 additions & 0 deletions .kokoro/common.sh
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,35 @@ function checkout_torch_pin_if_available() {
git submodule update --init --recursive
}

function install_environments() {
apt-get clean && apt-get update
apt-get upgrade -y
apt-get install --fix-missing -y python3-pip git curl libopenblas-dev vim jq \
apt-transport-https ca-certificates procps openssl sudo wget libssl-dev libc6-dbg

curl -LO https://github.com/bazelbuild/bazelisk/releases/download/v1.11.0/bazelisk-linux-amd64
mv bazelisk-linux-amd64 /usr/local/bin/bazel
chmod +x /usr/local/bin/bazel

pip install mkl mkl-include setuptools typing_extensions cmake requests
sudo ln -sf /usr/local/lib/libmkl_intel_lp64.so.2 /usr/local/lib/libmkl_intel_lp64.so.1
sudo ln -sf /usr/local/lib/libmkl_intel_thread.so.2 /usr/local/lib/libmkl_intel_thread.so.1
sudo ln -sf /usr/local/lib/libmkl_core.so.2 /usr/local/lib/libmkl_core.so.1
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib/

echo "deb [signed-by=/usr/share/keyrings/cloud.google.gpg] https://packages.cloud.google.com/apt cloud-sdk main" >> /etc/apt/sources.list.d/google-cloud-sdk.list
curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key --keyring /usr/share/keyrings/cloud.google.gpg add -

# TODO(yeounoh) fix `GoogleCredentials` import error
apt-get update
apt-get -y install google-cloud-cli
pip install --upgrade google-api-python-client
pip install --upgrade oauth2client
pip install --upgrade google-cloud-storage
pip install lark-parser
pip install cloud-tpu-client
}

function install_deps_pytorch_xla() {
XLA_DIR=$1

Expand Down
7 changes: 7 additions & 0 deletions .kokoro/presubmit.cfg
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@

# -*- protobuffer -*-
# proto-file: google3/devtools/kokoro/config/proto/build.proto
# proto-message: BuildConfig

build_file: "xla/.kokoro/build_and_run_stablehlo_tests.sh"
timeout_mins: 360

0 comments on commit e7ae1d4

Please sign in to comment.