Add cmake rule iree_fetch_artifact (#10516)
diff --git a/CMakeLists.txt b/CMakeLists.txt index 5df8be9..416578f 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt
@@ -369,6 +369,7 @@ include(iree_microbenchmark_suite) include(iree_hal_cts_test_suite) include(iree_static_linker_test) +include(iree_fetch_artifact) set(CMAKE_POSITION_INDEPENDENT_CODE TRUE)
diff --git a/benchmarks/generated_benchmark_suites.cmake b/benchmarks/generated_benchmark_suites.cmake index abf9456..a67e32f 100644 --- a/benchmarks/generated_benchmark_suites.cmake +++ b/benchmarks/generated_benchmark_suites.cmake
@@ -16,19 +16,14 @@ # Below is generated by build_tools/benchmarks/suites/cmake_rule_generator.py # ################################################################################ # Fetch the model from "https://storage.googleapis.com/iree-model-artifacts/mobilenet_v2_1.0_224.tflite" -add_custom_command( - OUTPUT "${_MODEL_ARTIFACTS_DIR}/7d45f8e5-bb5e-48d0-928d-8f125104578f_mobilenet_v2.tflite" - COMMAND - "${Python3_EXECUTABLE}" "${IREE_ROOT_DIR}/build_tools/scripts/download_file.py" - "https://storage.googleapis.com/iree-model-artifacts/mobilenet_v2_1.0_224.tflite" -o "${_MODEL_ARTIFACTS_DIR}/7d45f8e5-bb5e-48d0-928d-8f125104578f_mobilenet_v2.tflite" - DEPENDS - "${IREE_ROOT_DIR}/build_tools/scripts/download_file.py" - COMMENT "Downloading https://storage.googleapis.com/iree-model-artifacts/mobilenet_v2_1.0_224.tflite" -) -add_custom_target( - "${_PACKAGE_NAME}_model-7d45f8e5-bb5e-48d0-928d-8f125104578f" - DEPENDS +iree_fetch_artifact( + NAME + "model-7d45f8e5-bb5e-48d0-928d-8f125104578f" + SOURCE_URL + "https://storage.googleapis.com/iree-model-artifacts/mobilenet_v2_1.0_224.tflite" + OUTPUT "${_MODEL_ARTIFACTS_DIR}/7d45f8e5-bb5e-48d0-928d-8f125104578f_mobilenet_v2.tflite" + UNPACK ) # Import the TFLite model "${_MODEL_ARTIFACTS_DIR}/7d45f8e5-bb5e-48d0-928d-8f125104578f_mobilenet_v2.tflite"
diff --git a/build_tools/cmake/iree_benchmark_suite.cmake b/build_tools/cmake/iree_benchmark_suite.cmake index ee4e013..3795bd7 100644 --- a/build_tools/cmake/iree_benchmark_suite.cmake +++ b/build_tools/cmake/iree_benchmark_suite.cmake
@@ -223,26 +223,23 @@ # Update the source file to the downloaded-to place. string(REPLACE "/" ";" _SOURCE_URL_SEGMENTS "${_SOURCE_URL}") list(POP_BACK _SOURCE_URL_SEGMENTS _LAST_URL_SEGMENT) - set(_DOWNLOAD_TARGET "${_PACKAGE_NAME}_iree-download-benchmark-source-${_LAST_URL_SEGMENT}") + set(_DOWNLOAD_TARGET_NAME "iree-download-benchmark-source-${_LAST_URL_SEGMENT}") # Strip off gzip/tar suffix if present (downloader unpacks if necessary) string(REGEX REPLACE "(\.gz)|(\.tar\.gz)$" "" _SOURCE_FILE_BASENAME "${_LAST_URL_SEGMENT}") set(_MODULE_SOURCE "${_ROOT_ARTIFACTS_DIR}/${_SOURCE_FILE_BASENAME}") - if(NOT TARGET "${_DOWNLOAD_TARGET}") - add_custom_command( - OUTPUT "${_MODULE_SOURCE}" - COMMAND - "${Python3_EXECUTABLE}" "${IREE_ROOT_DIR}/build_tools/scripts/download_file.py" - "${_SOURCE_URL}" -o "${_MODULE_SOURCE}" - DEPENDS - "${IREE_ROOT_DIR}/build_tools/scripts/download_file.py" - COMMENT "Downloading ${_SOURCE_URL}" - ) - add_custom_target("${_DOWNLOAD_TARGET}" - DEPENDS "${_MODULE_SOURCE}" + if(NOT TARGET "${_PACKAGE_NAME}_${_DOWNLOAD_TARGET_NAME}") + iree_fetch_artifact( + NAME + "${_DOWNLOAD_TARGET_NAME}" + SOURCE_URL + "${_SOURCE_URL}" + OUTPUT + "${_MODULE_SOURCE}" + UNPACK ) endif() - set(_MODULE_SOURCE_TARGET "${_DOWNLOAD_TARGET}") + set(_MODULE_SOURCE_TARGET "${_PACKAGE_NAME}_${_DOWNLOAD_TARGET_NAME}") endif() # If the source is a TFLite file, import it.
diff --git a/build_tools/cmake/iree_fetch_artifact.cmake b/build_tools/cmake/iree_fetch_artifact.cmake new file mode 100644 index 0000000..b282b7d --- /dev/null +++ b/build_tools/cmake/iree_fetch_artifact.cmake
@@ -0,0 +1,56 @@ +# Copyright 2022 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 + +# iree_fetch_artifact() +# +# Download file from URL. NEVER Use this rule to download from untrusted +# sources, it doesn't unpack the file safely. +# +# Parameters: +# NAME: Name of target (see Note). +# SOURCE_URL: Source URL to donwload the file. +# OUTPUT: Path to the output file or directory to unpack. +# UNPACK: When added, it will try to unpack the archive if supported. +# +# Note: +# By default, it will create a target named ${_PACKAGE_NAME}_${_RULE_NAME}. +function(iree_fetch_artifact) + cmake_parse_arguments( + _RULE + "UNPACK" + "NAME;SOURCE_URL;OUTPUT" + "" + ${ARGN} + ) + + set(_ARGS "${IREE_ROOT_DIR}/build_tools/scripts/download_file.py") + list(APPEND _ARGS "${_RULE_SOURCE_URL}") + list(APPEND _ARGS "-o") + list(APPEND _ARGS "${_RULE_OUTPUT}") + + if(_RULE_UNPACK) + list(APPEND _ARGS "--unpack") + endif() + + # TODO: CMake built-in file command can replace the python script. But python + # script also provides streaming unpack (doesn't use double space when + # unpacking). Need to evaluate if we want to replace. + add_custom_command( + OUTPUT "${_RULE_OUTPUT}" + COMMAND + "${Python3_EXECUTABLE}" + ${_ARGS} + DEPENDS + "${IREE_ROOT_DIR}/build_tools/scripts/download_file.py" + COMMENT "Downloading ${_RULE_SOURCE_URL}" + ) + + iree_package_name(_PACKAGE_NAME) + add_custom_target("${_PACKAGE_NAME}_${_RULE_NAME}" + DEPENDS + "${_RULE_OUTPUT}" + ) +endfunction()
diff --git a/build_tools/python/e2e_test_framework/iree_download_artifact_template.cmake b/build_tools/python/e2e_test_framework/iree_download_artifact_template.cmake index 5dc8577..3c74089 100644 --- a/build_tools/python/e2e_test_framework/iree_download_artifact_template.cmake +++ b/build_tools/python/e2e_test_framework/iree_download_artifact_template.cmake
@@ -1,15 +1,10 @@ # Fetch the model from "$__SOURCE_URL" -add_custom_command( - OUTPUT "$__OUTPUT_PATH" - COMMAND - "$${Python3_EXECUTABLE}" "$${IREE_ROOT_DIR}/build_tools/scripts/download_file.py" - "$__SOURCE_URL" -o "$__OUTPUT_PATH" - DEPENDS - "$${IREE_ROOT_DIR}/build_tools/scripts/download_file.py" - COMMENT "Downloading $__SOURCE_URL" -) -add_custom_target( - "$${_PACKAGE_NAME}_$__TARGET_NAME" - DEPENDS +iree_fetch_artifact( + NAME + "$__TARGET_NAME" + SOURCE_URL + "$__SOURCE_URL" + OUTPUT "$__OUTPUT_PATH" + UNPACK )
diff --git a/build_tools/scripts/download_file.py b/build_tools/scripts/download_file.py index 492ad1c..7e7a689 100755 --- a/build_tools/scripts/download_file.py +++ b/build_tools/scripts/download_file.py
@@ -34,6 +34,10 @@ required=True, metavar="<output-file>", help="Output file path") + parser.add_argument("--unpack", + action='store_true', + default=False, + help="Unpack the downloaded file if it's an archive.") return parser.parse_args() @@ -50,24 +54,26 @@ f"Failed to download file with status {response.status} {response.msg}" ) - if args.source_url.endswith(".tar.gz"): - # Open tar.gz in the streaming mode. - with tarfile.open(fileobj=response, mode="r|*") as tar_file: - if os.path.exists(args.output): - shutil.rmtree(args.output) - os.makedirs(args.output) - tar_file.extractall(args.output) + if args.unpack: + if args.source_url.endswith(".tar.gz"): + # Open tar.gz in the streaming mode. + with tarfile.open(fileobj=response, mode="r|*") as tar_file: + if os.path.exists(args.output): + shutil.rmtree(args.output) + os.makedirs(args.output) + tar_file.extractall(args.output) + return + elif args.source_url.endswith(".gz"): + # Open gzip from a file-like object, which will be in the streaming mode. + with gzip.open(filename=response, mode="rb") as input_file: + with open(args.output, "wb") as output_file: + shutil.copyfileobj(input_file, output_file) + return - elif args.source_url.endswith(".gz"): - # Open gzip from a file-like object, which will be in the streaming mode. - with gzip.open(filename=response, mode="rb") as input_file: - with open(args.output, "wb") as output_file: - shutil.copyfileobj(input_file, output_file) - - else: - with open(args.output, "wb") as output_file: - # Streaming copy. - shutil.copyfileobj(response, output_file) + # Fallback to download the file only. + with open(args.output, "wb") as output_file: + # Streaming copy. + shutil.copyfileobj(response, output_file) if __name__ == "__main__":