Create an environment variable to specify BYOD image tag for running …

…release tests Currently release test infra generate the BYOD image tag for a test automatically. This PR gives folks a way to curate the tag by themselves. Signed-off-by: can <[email protected]>
cread · Aug 9, 2023 · 9b747da · 9b747da
1 parent 273a2a1
commit 9b747da
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 4 deletions.
diff --git a/release/ray_release/byod/build.py b/release/ray_release/byod/build.py
@@ -203,10 +203,14 @@ def _validate_and_push(byod_image: str) -> None:
         .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}"
+    if os.environ.get("RAY_IMAGE_TAG"):
+        logger.info(f"Ray commit from image: {docker_ray_commit}")
+    else:
+        expected_ray_commit = _get_ray_commit()
+        assert (
+            docker_ray_commit == expected_ray_commit
+        ), f"Expected ray commit {expected_ray_commit}, found {docker_ray_commit}"
+    logger.info(f"Pushing image to ECR: {byod_image}")
     subprocess.check_call(
         ["docker", "push", byod_image],
         stdout=sys.stderr,

diff --git a/release/ray_release/test.py b/release/ray_release/test.py
@@ -213,6 +213,12 @@ def get_byod_base_image_tag(self) -> str:
         """
         Returns the byod image tag to use for this test.
         """
+        byod_image_tag = os.environ.get("RAY_IMAGE_TAG")
+        if byod_image_tag:
+            # Use the image tag specified in the environment variable.
+            # TODO(can): this is a temporary backdoor that should be removed
+            # once civ2 is fully rolled out.
+            return byod_image_tag
         commit = os.environ.get(
             "COMMIT_TO_TEST",
             os.environ["BUILDKITE_COMMIT"],

diff --git a/release/ray_release/tests/test_test.py b/release/ray_release/tests/test_test.py
@@ -91,6 +91,11 @@ def test_get_ray_image():
             _stub_test({"cluster": {"byod": {}}}).get_ray_image()
             == "rayproject/ray:pr-123.123456-py38-cpu"
         )
+    with mock.patch.dict(os.environ, {"RAY_IMAGE_TAG": "my_tag"}):
+        assert (
+            _stub_test({"cluster": {"byod": {}}}).get_ray_image()
+            == "rayproject/ray:my_tag"
+        )
 
 
 def test_get_anyscale_byod_image():