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__":
