build_tools/docker/get_image_name.py

#!/usr/bin/env python3

# Copyright 2023 The IREE Authors
#
# Licensed under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""Command line tool to get the fully qualified image name given short name.

Syntax:
  ./build_tools/docker/get_image_name.py {short_name}

Where {short_name} is the last name component of an image in prod_digests.txt
(i.e. "base", "nvidia", etc).

This is used both in tree and out of tree to get a image name and current
version without adding fully referencing sha256 hashes, etc.
"""

from pathlib import Path
import sys


def find_image_by_name(img_name):
    this_dir = Path(__file__).resolve().parent

    with open(this_dir / "prod_digests.txt", "rt") as f:
        for line in f.readlines():
            line = line.strip()
            if line.startswith(f"gcr.io/iree-oss/{img_name}@"):
                return line
        else:
            raise ValueError(
                f"ERROR: Image name {img_name} not found in prod_digests.txt"
            )


if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("ERROR: Expected image short name", file=sys.stderr)
        sys.exit(1)
    short_name = sys.argv[1]
    image_name = find_image_by_name(short_name)
    print(image_name)