From 9f693e35a2a86b90af07ee9665f05c508bc1ba76 Mon Sep 17 00:00:00 2001
From: Amit Patankar <amitpatankar@google.com>
Date: Tue, 14 Apr 2020 10:54:46 -0700
Subject: [PATCH] Add a binary size check for CPU binaries as well.

PiperOrigin-RevId: 306472255
Change-Id: Id1c760d984311a1037c71b167d1ccbb110bff9ce
---
 .../builds/nightly_release_smoke_test.sh      | 23 +++++++++++++++----
 1 file changed, 18 insertions(+), 5 deletions(-)

diff --git a/tensorflow/tools/ci_build/builds/nightly_release_smoke_test.sh b/tensorflow/tools/ci_build/builds/nightly_release_smoke_test.sh
index 93a1888571e..f0184788a96 100644
--- a/tensorflow/tools/ci_build/builds/nightly_release_smoke_test.sh
+++ b/tensorflow/tools/ci_build/builds/nightly_release_smoke_test.sh
@@ -18,7 +18,8 @@
 set -e
 set -x
 
-MAX_WHL_SIZE=550M
+CPU_MAX_WHL_SIZE=190M
+GPU_MAX_WHL_SIZE=510M
 
 function run_smoke_test() {
   VENV_TMP_DIR=$(mktemp -d)
@@ -82,10 +83,22 @@ function test_tf_imports() {
 }
 
 function test_tf_whl_size() {
-  if [[ $(find $WHL_NAME -type f -size +${MAX_WHL_SIZE}) ]]; then
-    echo "The whl size has exceeded 550MB. To keep within pypi's CDN
-distribution limit, we must not exceed that threshold."
-    return 1
+  # We do not need a separate check for MacOS regular binaries.
+  # We check for the `_cpu` string in the whl file name.
+  if [[ "$WHL_NAME" == *"_cpu"* ]]; then
+    # Check CPU whl size.
+    if [[ $(find $WHL_NAME -type f -size +${CPU_MAX_WHL_SIZE}) ]]; then
+      echo "The CPU whl size has exceeded ${CPU_MAX_WHL_SIZE}MB. To keep
+within pypi's CDN distribution limit, we must not exceed that threshold."
+      return 1
+    fi
+  else
+    # Check GPU whl size.
+    if [[ $(find $WHL_NAME -type f -size +${GPU_MAX_WHL_SIZE}) ]]; then
+      echo "The GPU whl size has exceeded ${GPU_MAX_WHL_SIZE}MB. To keep
+within pypi's CDN distribution limit, we must not exceed that threshold."
+      return 1
+    fi
   fi
 }