From 34b2c69eab6e0f43d6c06f335795b4ce21bf2688 Mon Sep 17 00:00:00 2001
From: Qianli Scott Zhu <scottzhu@google.com>
Date: Tue, 2 May 2023 10:40:28 -0700
Subject: [PATCH] Cherrypick the release script fix for RC.

---
 pip_build.py | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/pip_build.py b/pip_build.py
index 708f1dc75d5..2ef43bce205 100644
--- a/pip_build.py
+++ b/pip_build.py
@@ -338,6 +338,7 @@ def build_pip_package(
     src_directory,
     dist_directory,
     is_nightly=False,
+    rc=None,
 ):
     # Build Keras with Bazel to get the protobuf .py files
     os.chdir(keras_root_directory)
@@ -383,6 +384,8 @@ def build_pip_package(
     if is_nightly:
         date = datetime.datetime.now()
         version += f".dev{date.strftime('%Y%m%d%H')}"
+    elif rc:
+        version += rc
     with open(os.path.join(package_directory, "__init__.py")) as f:
         init_contents = f.read()
     with open(os.path.join(package_directory, "__init__.py"), "w") as f:
@@ -455,8 +458,14 @@ def test_wheel(wheel_path, expected_version, requirements_path):
         action="store_true",
         help="Whether this is for the `keras-nightly` package.",
     )
+    parser.add_argument(
+        "--RC",
+        type=str,
+        help="Whether this is for the release candidate.",
+    )
     args = parser.parse_args()
     is_nightly = args.nightly
+    rc = args.RC
 
     build_directory = os.path.join(tempfile.gettempdir(), TMP_BUILD_DIRNAME)
     keras_root_directory = pathlib.Path(__file__).parent.resolve()
@@ -472,6 +481,7 @@ def test_wheel(wheel_path, expected_version, requirements_path):
             f"package_directory={package_directory}\n"
             f"src_directory={src_directory}\n"
             f"is_nightly={is_nightly}"
+            f"rc={rc}"
         )
     if os.path.exists(build_directory):
         raise ValueError(f"Directory already exists: {build_directory}")
@@ -487,6 +497,7 @@ def test_wheel(wheel_path, expected_version, requirements_path):
             src_directory,
             dist_directory,
             is_nightly,
+            rc,
         )
         wheel_filename = [f for f in saved_filenames if f.endswith(".whl")][0]
         if VERBOSE:
@@ -504,4 +515,4 @@ def test_wheel(wheel_path, expected_version, requirements_path):
         # Clean up: remove the build directory (no longer needed)
         if VERBOSE:
             print(f"Deleting temp build directory at {build_directory}...")
-        shutil.rmtree(build_directory)
+        shutil.rmtree(build_directory)
\ No newline at end of file