Added package test script.

PiperOrigin-RevId: 592057630
Change-Id: I9e47f76baabc5d3ab9cc1060c7da4945b214b72b

Author: ssjhv

.bazelrc

@@ -1,2 +1,9 @@
 # This flag is required for CUDA repo that @org_tensorflow depends on.
 common --experimental_repo_remote_exec
+
+build:manylinux_2_17_x86_64 --host_crosstool_top="@ubuntu20.04-gcc9_manylinux2014-cuda11.2-cudnn8.1-tensorrt7.2_config_cuda//crosstool:toolchain"
+build:manylinux_2_17_x86_64 --crosstool_top="@ubuntu20.04-gcc9_manylinux2014-cuda11.2-cudnn8.1-tensorrt7.2_config_cuda//crosstool:toolchain"
+build:manylinux_2_17_x86_64 --extra_toolchains="@ubuntu20.04-gcc9_manylinux2014-cuda11.2-cudnn8.1-tensorrt7.2_config_cuda//crosstool:toolchain-linux-x86_64"
+build:manylinux_2_17_x86_64 --extra_execution_platforms="@ubuntu20.04-gcc9_manylinux2014-cuda11.2-cudnn8.1-tensorrt7.2_config_platform//:platform"
+build:manylinux_2_17_x86_64 --host_platform="@ubuntu20.04-gcc9_manylinux2014-cuda11.2-cudnn8.1-tensorrt7.2_config_platform//:platform"
+build:manylinux_2_17_x86_64 --platforms="@ubuntu20.04-gcc9_manylinux2014-cuda11.2-cudnn8.1-tensorrt7.2_config_platform//:platform"

BUILD

@@ -45,7 +45,7 @@ py_binary(
 
 py_binary(
     name = "build_pip_pkg",
-    srcs = ["build_pip_pkg.py"],
+    srcs = ["tools/build_pip_pkg.py"],
     data = [
         "LICENSE",
         "README.md",

build_pip_pkg.py renamed to tools/build_pip_pkg.py

@@ -15,30 +15,36 @@
 # ==============================================================================
 
 # This script must run at the workspace root directory.
-# In addition, `tensorflow_compression/cc:libtensorflow_compression.so` should
-# have been built using bazel.
 
 set -ex  # Fail if any command fails, echo commands.
 
-OUTPUT_DIR="${1-/tmp}"
-WHEEL_VERSION="${2-0.dev0}"
+# Script configuration --------------------------------------------------------
+OUTPUT_DIR="${1-/tmp/tensorflow_compression}"
+WHEEL_VERSION=${2-0.dev0}
 
-PKGDIR="$(mktemp -d)"
-trap 'rm -r -- "${PKGDIR}"' EXIT
+# Optionally exported environment variables.
+: ${BAZEL_OPT:=}
+# -----------------------------------------------------------------------------
 
-cp LICENSE README.md MANIFEST.in requirements.txt build_pip_pkg.py "${PKGDIR}"
+python -m pip install -U pip setuptools wheel
+python -m pip install -r requirements.txt
+bazel build ${BAZEL_OPT} -c opt --copt=-mavx tensorflow_compression/cc:libtensorflow_compression.so
 
-mkdir -p "${PKGDIR}/tensorflow_compression/cc"
+SRCDIR="$(mktemp -d)"
+trap 'rm -r -- "${SRCDIR}"' EXIT
+
+cp LICENSE README.md MANIFEST.in requirements.txt "${SRCDIR}"
+
+mkdir -p "${SRCDIR}/tensorflow_compression/cc"
 cp "$(bazel info -c opt bazel-genfiles)/tensorflow_compression/cc/libtensorflow_compression.so" \
-  "${PKGDIR}/tensorflow_compression/cc"
+  "${SRCDIR}/tensorflow_compression/cc"
 
 
-copy_file()
-{
-  FILENAME="${1#./}"
-  DEST="${PKGDIR%/}/$(dirname "${FILENAME}")"
-  mkdir -p "${DEST}"
-  cp "${FILENAME}" "${DEST}"
+copy_file() {
+  local FILENAME="${1#./}"
+  local DST="${SRCDIR%/}/$(dirname "${FILENAME}")"
+  mkdir -p "${DST}"
+  cp "${FILENAME}" "${DST}"
 }
 
 copy_file "tensorflow_compression/__init__.py"
@@ -48,6 +54,4 @@ copy_file "tensorflow_compression/all_tests.py"
 find tensorflow_compression/python -name "*.py" \
   | while read filename; do copy_file "${filename}"; done
 
-pushd "${PKGDIR}"
-python build_pip_pkg.py . "${OUTPUT_DIR}" "${WHEEL_VERSION}"
-popd
+python tools/build_pip_pkg.py "${SRCDIR}" "${OUTPUT_DIR}" "${WHEEL_VERSION}"

build_pip_pkg.sh renamed to tools/build_pip_pkg.sh

@@ -0,0 +1,36 @@
+#!/usr/bin/env bash
+# Copyright 2023 Google LLC. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+set -ex  # Fail if any command fails, echo commands.
+
+WHEEL="${1}"
+
+# `import tensorflow_compression` in the bazel root directory produces cryptic
+# error messages, because Python ends up looking for .so files under the
+# subdirectories in the src repo instead of Python module libraries. Changing
+# the current directory helps avoid running tests inside the bazel root
+# direcotory by accident.
+pushd /tmp
+
+python -m pip install -U pip setuptools wheel
+python -m pip install "${WHEEL}"
+python -m pip list -v
+
+# Logs elements of tfc namespace and runs unit tests.
+python -c "import tensorflow_compression as tfc; print('\n'.join(sorted(dir(tfc))))"
+python -m tensorflow_compression.all_tests
+
+popd  # /tmp