[ci] refactor push_ray_image shared functions to image_tags_lib (#60793)

There is significant overlap with the next feature of
push_anyscale_image, so this refactor reduces on a lot of duplicated
work.

Signed-off-by: andrew <andrew@anyscale.com>

Author: andrew-anyscale

ci/ray_ci/automation/BUILD.bazel

@@ -114,6 +114,31 @@ py_library(
     ],
 )
 
+py_library(
+    name = "image_tags_lib",
+    srcs = ["image_tags_lib.py"],
+    visibility = ["//ci/ray_ci/automation:__subpackages__"],
+    deps = [
+        ":crane_lib",
+        "//ci/ray_ci:ray_ci_lib",
+    ],
+)
+
+py_test(
+    name = "test_image_tags_lib",
+    size = "small",
+    srcs = ["test_image_tags_lib.py"],
+    exec_compatible_with = ["//bazel:py3"],
+    tags = [
+        "ci_unit",
+        "team:ci",
+    ],
+    deps = [
+        ":image_tags_lib",
+        ci_require("pytest"),
+    ],
+)
+
 py_library(
     name = "test_utils",
     srcs = ["test_utils.py"],
@@ -322,7 +347,7 @@ py_binary(
     srcs = ["push_ray_image.py"],
     exec_compatible_with = ["//bazel:py3"],
     deps = [
-        ":crane_lib",
+        ":image_tags_lib",
         "//ci/ray_ci:ray_ci_lib",
         "//release:ray_release",
         ci_require("click"),

ci/ray_ci/automation/image_tags_lib.py

@@ -0,0 +1,137 @@
+"""
+Shared utilities for generating Docker image tags for Ray and Anyscale images.
+"""
+
+import logging
+from typing import List
+
+from ci.ray_ci.automation.crane_lib import (
+    CraneError,
+    call_crane_copy,
+    call_crane_manifest,
+)
+from ci.ray_ci.configs import DEFAULT_PYTHON_TAG_VERSION
+from ci.ray_ci.docker_container import GPU_PLATFORM
+
+logger = logging.getLogger(__name__)
+
+
+class ImageTagsError(Exception):
+    """Error raised when image tag operations fail."""
+
+
+def format_platform_tag(platform: str) -> str:
+    """
+    Format platform as -cpu or shortened CUDA version.
+
+    Examples:
+        cpu -> -cpu
+        cu12.1.1-cudnn8 -> -cu121
+        cu12.3.2-cudnn9 -> -cu123
+    """
+    if platform == "cpu":
+        return "-cpu"
+    # cu12.3.2-cudnn9 -> -cu123
+    platform_base = platform.split("-", 1)[0]
+    parts = platform_base.split(".")
+    if len(parts) < 2:
+        raise ImageTagsError(f"Unrecognized GPU platform format: {platform}")
+    return f"-{parts[0]}{parts[1]}"
+
+
+def format_python_tag(python_version: str) -> str:
+    """
+    Format python version as -py310 (no dots, with hyphen prefix).
+
+    Examples:
+        3.10 -> -py310
+        3.11 -> -py311
+    """
+    return f"-py{python_version.replace('.', '')}"
+
+
+def get_python_suffixes(python_version: str) -> List[str]:
+    """
+    Get python version suffixes (includes empty for default version).
+
+    For DEFAULT_PYTHON_TAG_VERSION (3.10), returns both the explicit
+    suffix and an empty suffix for backward compatibility.
+
+    Examples:
+        3.10 -> ["-py310", ""]
+        3.11 -> ["-py311"]
+    """
+    suffixes = [format_python_tag(python_version)]
+    if python_version == DEFAULT_PYTHON_TAG_VERSION:
+        suffixes.append("")
+    return suffixes
+
+
+def get_platform_suffixes(platform: str, image_type: str) -> List[str]:
+    """
+    Get platform suffixes (includes aliases like -gpu for GPU_PLATFORM).
+
+    Handles the following cases:
+    - CPU with ray/ray-extra: adds empty suffix alias
+    - GPU_PLATFORM: adds -gpu alias
+    - GPU_PLATFORM with ray-ml/ray-ml-extra: adds empty suffix alias
+
+    Args:
+        platform: The platform string (e.g., "cpu", "cu12.1.1-cudnn8")
+        image_type: The image type (e.g., "ray", "ray-ml", "ray-extra")
+
+    Returns:
+        List of platform suffixes to use for tagging.
+    """
+    platform_tag = format_platform_tag(platform)
+    suffixes = [platform_tag]
+
+    if platform == "cpu":
+        # no tag is alias to cpu for ray image
+        if image_type in ("ray", "ray-extra"):
+            suffixes.append("")
+    elif platform == GPU_PLATFORM:
+        # gpu is alias to GPU_PLATFORM
+        suffixes.append("-gpu")
+        # no tag is alias to gpu for ray-ml image
+        if image_type in ("ray-ml", "ray-ml-extra"):
+            suffixes.append("")
+
+    return suffixes
+
+
+def get_variation_suffix(image_type: str) -> str:
+    """
+    Get variation suffix for -extra image types.
+
+    Examples:
+        ray -> ""
+        ray-extra -> "-extra"
+        ray-ml-extra -> "-extra"
+    """
+    if image_type in ("ray-extra", "ray-ml-extra", "ray-llm-extra"):
+        return "-extra"
+    return ""
+
+
+def image_exists(tag: str) -> bool:
+    """Check if a container image manifest exists using crane."""
+    try:
+        call_crane_manifest(tag)
+        return True
+    except CraneError:
+        return False
+
+
+def copy_image(source: str, destination: str, dry_run: bool = False) -> None:
+    """Copy a container image from source to destination using crane."""
+    if dry_run:
+        logger.info(f"DRY RUN: Would copy {source} -> {destination}")
+        return
+
+    logger.info(f"Copying {source} -> {destination}")
+    try:
+        call_crane_copy(source, destination)
+        logger.info(f"Successfully copied to {destination}")
+    except CraneError as e:
+        raise ImageTagsError(f"Crane copy failed: {e}")

ci/ray_ci/automation/push_ray_image.py

@@ -5,22 +5,24 @@
 
 import click
 
-from ci.ray_ci.automation.crane_lib import (
-    CraneError,
-    call_crane_copy,
-    call_crane_manifest,
+from ci.ray_ci.automation.image_tags_lib import (
+    ImageTagsError,
+    copy_image,
+    format_platform_tag,
+    get_platform_suffixes,
+    get_python_suffixes,
+    get_variation_suffix,
+    image_exists,
 )
 from ci.ray_ci.configs import (
     ARCHITECTURE,
     DEFAULT_ARCHITECTURE,
-    DEFAULT_PYTHON_TAG_VERSION,
     PYTHON_VERSIONS,
 )
 from ci.ray_ci.docker_container import (
     ARCHITECTURES_RAY,
     ARCHITECTURES_RAY_LLM,
     ARCHITECTURES_RAY_ML,
-    GPU_PLATFORM,
     PLATFORMS_RAY,
     PLATFORMS_RAY_LLM,
     PLATFORMS_RAY_ML,
@@ -48,14 +50,10 @@ class PushRayImageError(Exception):
     """Error raised when pushing ray images fails."""
 
 
+# Re-export for backward compatibility with tests
 def compact_cuda_suffix(platform: str) -> str:
     """Convert a CUDA platform string to compact suffix (e.g. cu12.1.1-cudnn8 -> -cu121)."""
-    platform_base = platform.split("-", 1)[0]
-    parts = platform_base.split(".")
-    if len(parts) < 2:
-        raise PushRayImageError(f"Unrecognized GPU platform format: {platform}")
-
-    return f"-{parts[0]}{parts[1]}"
+    return format_platform_tag(platform)
 
 
 class RayImagePushContext:
@@ -188,62 +186,28 @@ def _versions(self) -> List[str]:
 
     def _variation_suffix(self) -> str:
         """Get -extra suffix for extra image types."""
-        if self.ray_type in {
-            RayType.RAY_EXTRA,
-            RayType.RAY_ML_EXTRA,
-            RayType.RAY_LLM_EXTRA,
-        }:
-            return "-extra"
-        return ""
+        return get_variation_suffix(self.ray_type.value)
 
     def _python_suffixes(self) -> List[str]:
         """Get python version suffixes (includes empty for default version)."""
-        suffixes = [f"-py{self.python_version.replace('.', '')}"]
-        if self.python_version == DEFAULT_PYTHON_TAG_VERSION:
-            suffixes.append("")
-        return suffixes
+        return get_python_suffixes(self.python_version)
 
     def _platform_suffixes(self) -> List[str]:
         """Get platform suffixes (includes aliases like -gpu for GPU_PLATFORM)."""
-        if self.platform == "cpu":
-            suffixes = ["-cpu"]
-            # no tag is alias to cpu for ray image
-            if self.ray_type in {RayType.RAY, RayType.RAY_EXTRA}:
-                suffixes.append("")
-            return suffixes
-
-        suffixes = [compact_cuda_suffix(self.platform)]
-        if self.platform == GPU_PLATFORM:
-            # gpu is alias to GPU_PLATFORM value for ray image
-            suffixes.append("-gpu")
-            # no tag is alias to gpu for ray-ml image
-            if self.ray_type in {RayType.RAY_ML, RayType.RAY_ML_EXTRA}:
-                suffixes.append("")
-
-        return suffixes
+        return get_platform_suffixes(self.platform, self.ray_type.value)
 
 
 def _image_exists(tag: str) -> bool:
     """Check if a container image manifest exists using crane."""
-    try:
-        call_crane_manifest(tag)
-        return True
-    except CraneError:
-        return False
+    return image_exists(tag)
 
 
 def _copy_image(reference: str, destination: str, dry_run: bool = False) -> None:
     """Copy a container image from source to destination using crane."""
-    if dry_run:
-        logger.info(f"DRY RUN: Would copy {reference} -> {destination}")
-        return
-
-    logger.info(f"Copying {reference} -> {destination}")
     try:
-        call_crane_copy(reference, destination)
-        logger.info(f"Successfully copied to {destination}")
-    except CraneError as e:
-        raise PushRayImageError(f"Crane copy failed: {e}")
+        copy_image(reference, destination, dry_run)
+    except ImageTagsError as e:
+        raise PushRayImageError(str(e))
 
 
 def _should_upload(pipeline_id: str, branch: str, rayci_schedule: str) -> bool: