From 629dbb151f61c96a2f2dc514eab61fdcd5f89100 Mon Sep 17 00:00:00 2001
From: can <can@anyscale.com>
Date: Wed, 28 Jun 2023 01:14:34 +0000
Subject: [PATCH] Validate ray wheel

Signed-off-by: can <can@anyscale.com>
---
 release/ray_release/byod/build.py            | 49 ++++++++++++++++----
 release/ray_release/tests/test_byod_build.py |  3 ++
 2 files changed, 42 insertions(+), 10 deletions(-)

diff --git a/release/ray_release/byod/build.py b/release/ray_release/byod/build.py
index cc7780a7405e5..fe21de9b4166c 100644
--- a/release/ray_release/byod/build.py
+++ b/release/ray_release/byod/build.py
@@ -50,12 +50,7 @@ def build_anyscale_custom_byod_image(test: Test) -> None:
         stdout=sys.stderr,
         env=env,
     )
-    # push the image to ecr, the image will have a tag in this format
-    # {commit_sha}-py{version}-gpu-{custom_information_dict_hash}
-    subprocess.check_call(
-        ["docker", "push", byod_image],
-        stdout=sys.stderr,
-    )
+    _validate_and_push(byod_image)
 
 
 def build_anyscale_base_byod_images(tests: List[Test]) -> None:
@@ -134,13 +129,47 @@ def build_anyscale_base_byod_images(tests: List[Test]) -> None:
                     stdout=sys.stderr,
                     env=env,
                 )
-                subprocess.check_call(
-                    ["docker", "push", byod_image],
-                    stdout=sys.stderr,
-                )
+                _validate_and_push(byod_image)
                 built.add(ray_image)
 
 
+def _validate_and_push(byod_image: str) -> None:
+    """
+    Validates the given image and pushes it to ECR.
+    """
+    docker_ray_commit = (
+        subprocess.check_output(
+            [
+                "docker",
+                "run",
+                "-ti",
+                "--entrypoint",
+                "python",
+                byod_image,
+                "-c",
+                "import ray; print(ray.__commit__)",
+            ],
+        )
+        .decode("utf-8")
+        .strip()
+    )
+    expected_ray_commit = _get_ray_commit()
+    assert (
+        docker_ray_commit == expected_ray_commit
+    ), f"Expected ray commit {expected_ray_commit}, found {docker_ray_commit}"
+    subprocess.check_call(
+        ["docker", "push", byod_image],
+        stdout=sys.stderr,
+    )
+
+
+def _get_ray_commit() -> str:
+    return os.environ.get(
+        "COMMIT_TO_TEST",
+        os.environ["BUILDKITE_COMMIT"],
+    )
+
+
 def _download_dataplane_build_file() -> None:
     """
     Downloads the dataplane build file from S3.
diff --git a/release/ray_release/tests/test_byod_build.py b/release/ray_release/tests/test_byod_build.py
index 3cea0f1496e80..ba5601d039088 100644
--- a/release/ray_release/tests/test_byod_build.py
+++ b/release/ray_release/tests/test_byod_build.py
@@ -26,6 +26,9 @@ def _mock_check_call(
     ), patch(
         "subprocess.check_call",
         side_effect=_mock_check_call,
+    ), patch(
+        "subprocess.check_output",
+        return_value=b"abc123",
     ):
         test = Test(
             name="name",