From e7ae1d46e5e54881a261fdabf34209e413c2ec81 Mon Sep 17 00:00:00 2001
From: qihqi <hanq@google.com>
Date: Thu, 27 Jul 2023 17:03:19 -0700
Subject: [PATCH] Add a script for running stablehlo tests. (#5360)

* Add kokoro presubmit for stablehlo tests
---
 .kokoro/build.sh                         | 27 +--------------------
 .kokoro/build_and_run_stablehlo_tests.sh | 31 ++++++++++++++++++++++++
 .kokoro/common.sh                        | 29 ++++++++++++++++++++++
 .kokoro/presubmit.cfg                    |  7 ++++++
 4 files changed, 68 insertions(+), 26 deletions(-)
 create mode 100644 .kokoro/build_and_run_stablehlo_tests.sh
 create mode 100644 .kokoro/presubmit.cfg

diff --git a/.kokoro/build.sh b/.kokoro/build.sh
index 83199e875447..8ba3a46767a8 100755
--- a/.kokoro/build.sh
+++ b/.kokoro/build.sh
@@ -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 \
diff --git a/.kokoro/build_and_run_stablehlo_tests.sh b/.kokoro/build_and_run_stablehlo_tests.sh
new file mode 100644
index 000000000000..0c39bd3a330e
--- /dev/null
+++ b/.kokoro/build_and_run_stablehlo_tests.sh
@@ -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
+
+
+
diff --git a/.kokoro/common.sh b/.kokoro/common.sh
index 59942511c587..fd04317efc23 100644
--- a/.kokoro/common.sh
+++ b/.kokoro/common.sh
@@ -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
 
diff --git a/.kokoro/presubmit.cfg b/.kokoro/presubmit.cfg
new file mode 100644
index 000000000000..98b285b3c6e3
--- /dev/null
+++ b/.kokoro/presubmit.cfg
@@ -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