Add "model provider" pattern to litert device scripts. This is used to specify scripts that generate data dependencies at runtime vs specifying them directly. Wire this up with ats to support models downloaded from the public bucket.

LukeBoyer · copybara-github · commit cf52fe3e3039 · 2025-10-29T19:01:23.000-07:00
LiteRT-PiperOrigin-RevId: 825802418
diff --git a/litert/ats/BUILD b/litert/ats/BUILD
@@ -17,6 +17,7 @@ load("@rules_cc//cc:cc_test.bzl", "cc_test")
 load("//litert/ats:ats.bzl", "litert_define_ats")
 load("//litert/build_common:litert_build_defs.bzl", "litert_test")
 load("//litert/integration_test:litert_device.bzl", "litert_device_test")
+load("//litert/integration_test:litert_device_script.bzl", "make_download_model_provider")
 
 package(
     # copybara:uncomment default_applicable_licenses = ["@org_tensorflow//tensorflow:license"],
@@ -328,6 +329,11 @@ cc_library(
 
 # PRE-CONFIGURED CTS SUITES ########################################################################
 
+make_download_model_provider(
+    name = "ats_models_provider",
+    url = "https://storage.googleapis.com/litert/ats_models.tar.gz",
+)
+
 litert_define_ats(
     name = "sample_cpu_ats",
     backend = "cpu",
@@ -340,8 +346,9 @@ litert_define_ats(
     backend = "example",
     compile_only_suffix = "_aot",
     do_register = [
-        "sub.*f32",
-        "mul.*f32",
+        # "sub.*f32",
+        # "mul.*f32",
+        ".*ExtraModel.*",
     ],
     extra_flags = ["--limit=1"],
     jit_suffix = "",
diff --git a/litert/ats/ats.bzl b/litert/ats/ats.bzl
@@ -89,7 +89,8 @@ def litert_define_ats(
     if compile_aot_and_run_suffix:
         fail("Compile aot and run on device is not supported yet.")
 
-    init_run_args = []
+    # TODO: Unify local workdir paths for scripting.
+    init_run_args = ["--extra_models={}".format("/data/local/tmp/runfiles/user/tmp/litert_extras")]
     if is_npu_backend(backend):
         init_run_args += [
             "--dispatch_dir=\"{}\"".format(dispatch_device_rlocation(backend)),
@@ -114,6 +115,7 @@ def litert_define_ats(
             local_suffix = "",
             exec_args = run_args,
             backend_id = backend,
+            model_providers = ["//litert/integration_test:ats_models_provider"],
         )
 
     compile_args = _make_ats_args(
@@ -135,4 +137,5 @@ def litert_define_ats(
             exec_args = compile_args,
             build_for_host = True,
             build_for_device = False,
+            model_providers = ["//litert/integration_test:ats_models_provider"],
         )
diff --git a/litert/ats/ats_aot.sh b/litert/ats/ats_aot.sh
@@ -16,7 +16,8 @@
 
 source "${0%.*}_lib.sh" || exit 1
 
-models_out="/tmp/ats_models"
+# TODO: Unify workdirs with other scripts.
+readonly models_out="/tmp/litert_extras/ats"
 readonly exec_args=("${@:1}")
 
 compile_bin=""
@@ -29,6 +30,7 @@ compiler_libs=()
 
 function setup_context() {
   mkdir -p "$models_out"
+  rm -rf "$models_out"/*
 
   local in_flags=$1
   for a in ${in_flags[@]}; do
@@ -67,6 +69,15 @@ function setup_context() {
       link_path=$(dirname ${lib}):${link_path}
     fi
   done
+
+  local input_models=($(get_provided_models))
+  if [[ $? -ne 0 ]]; then
+    fatal "Failed to get provided models."
+  fi
+
+  if [[ -n "${input_models[*]}" ]]; then
+    compile_args+=("--extra_models=$(str_join "," ${input_models[@]})")
+  fi
 }
 
 function print_args() {
diff --git a/litert/integration_test/BUILD b/litert/integration_test/BUILD
@@ -19,7 +19,7 @@ load("@rules_cc//cc:cc_test.bzl", "cc_test")
 load("//litert/build_common:special_rule.bzl", "litert_android_linkopts")
 load("//litert/integration_test:litert_device.bzl", "litert_device_exec", "litert_device_test", "litert_integration_test")
 load("//litert/integration_test:litert_device_common.bzl", "device_rlocation", "get_libs")
-load("//litert/integration_test:litert_device_script.bzl", "litert_device_script")
+load("//litert/integration_test:litert_device_script.bzl", "litert_device_script", "make_download_model_provider")
 # copybara:uncomment load("//litert/integration_test/google:litert_device_guitar.bzl", "litert_cpu_mh_guitar_test", "litert_mediatek_mh_guitar_test", "litert_pixel_9_mh_guitar_test", "litert_qualcomm_mh_guitar_test")
 
 package(
@@ -427,6 +427,17 @@ cc_binary(
     ],
 )
 
+sh_binary(
+    name = "dummy_model_provider",
+    srcs = ["dummy_model_provider.sh"],
+)
+
+sh_library(
+    name = "device_script_common",
+    srcs = ["device_script_common.sh"],
+    visibility = ["//litert:litert_public"],
+)
+
 litert_device_script(
     name = "check_script_device",
     testonly = True,
@@ -441,11 +452,40 @@ litert_device_script(
     exec_args = [
         "--check_device",
     ],
+    model_providers = [
+        ":dummy_model_provider",
+    ],
+    script = "device_script_test.sh",
+)
+
+litert_device_script(
+    name = "check_script_host",
+    testonly = True,
+    backend_id = "example",
+    bin = ":dummy_binary",
+    build_for_device = False,
+    build_for_host = True,
+    data = [
+        "//litert/test:testdata/mobilenet_v2_1.0_224.tflite",
+        "//litert/test:testdata/simple_add_op_qc_v75_precompiled.tflite",
+    ],
+    exec_args = [
+        "--check_host",
+    ],
+    model_providers = [
+        ":dummy_model_provider",
+    ],
     script = "device_script_test.sh",
 )
 
+make_download_model_provider(
+    name = "ats_models_provider",
+    url = "https://storage.googleapis.com/litert/ats_models.tar.gz",
+)
+
 litert_device_exec(
     name = "exec_for_testing_device",
+    testonly = True,
     backend_id = "example",
     data = [
         "//litert/test:testdata/mobilenet_v2_1.0_224.tflite",
@@ -459,6 +499,9 @@ litert_device_exec(
         "--expected_libs_on_ld={}".format(",".join([device_rlocation(lib) for lib in get_libs("example")])),
     ],
     local_suffix = "",
+    model_providers = [
+        ":ats_models_provider",
+    ],
     remote_suffix = None,
     target = ":dummy_binary",
 )
diff --git a/litert/integration_test/device_script_common.sh b/litert/integration_test/device_script_common.sh
@@ -0,0 +1,52 @@
+# Copyright 2025 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+#!/bin/bash
+
+reset_color='\033[0m'
+host_color='\033[34m'
+hightlight_color='\033[36m'
+error_color='\033[31m'
+
+# Print a message in the canonical color.
+function print() {
+  echo -e "${host_color}${1}${reset_color}"
+}
+
+# Print a message in the canonical hightlight color.
+function print_hightlight() {
+  echo -e "${hightlight_color}${1}${reset_color}"
+}
+
+# Print a file or host file device file pair.
+function print_file() {
+  if [[ "$#" -ne 2 ]]; then
+    echo -e "    ${1}"
+  else
+    echo -e "    ${1} => ${2}"
+  fi
+}
+
+# Print message and exit.
+function fatal() {
+  echo -e "${error_color}ERROR: ${reset_color}${1}"
+  exit 1
+}
+
+# Join with delim.
+function str_join() {
+  local IFS=$1
+  shift
+  echo "$*"
+}
diff --git a/litert/integration_test/device_script_template.sh b/litert/integration_test/device_script_template.sh
@@ -1,5 +1,3 @@
-#!/bin/bash
-
 # Copyright 2025 Google LLC.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -19,10 +17,7 @@
 # Shell library for working with data and executable files from bzl between host
 # and device. Meant to be templated via litert_device_script.bzl.
 
-reset_color='\033[0m'
-host_color='\033[34m'
-hightlight_color='\033[36m'
-error_color='\033[31m'
+source third_party/odml/litert/litert/integration_test/device_script_common.sh || exit 1
 
 # Root of runfiles on the device.
 device_runfiles_root="/data/local/tmp/runfiles"
@@ -169,28 +164,17 @@ function find_device_runtime_lib() {
   done
 }
 
-# Print a message in the canonical color.
-function print() {
-  echo -e "${host_color}${1}${reset_color}"
-}
-
-# Print a message in the canonical hightlight color.
-function print_hightlight() {
-  echo -e "${hightlight_color}${1}${reset_color}"
-}
-
-# Print a file or host file device file pair.
-function print_file() {
-  if [[ "$#" -ne 2 ]]; then
-    echo -e "    ${1}"
+# Call any/all model provider scripts built with this tool. The return code
+# of this need to be checked by callers.
+function get_provided_models() {
+  local model_providers=@@model_providers@@
+  if [[ "$model_providers" == "@@"*"@@" ]]; then
+    echo ""
   else
-    echo -e "    ${1} => ${2}"
+    for provider in ${model_providers[@]}; do
+      $provider
+    done
   fi
 }
 
-# Print message and exit.
-function fatal() {
-  echo -e "${error_color}ERROR: ${reset_color}${1}"
-  exit 1
-}
 
diff --git a/litert/integration_test/device_script_test.sh b/litert/integration_test/device_script_test.sh
@@ -72,11 +72,19 @@ function check_len() {
   if (( ${#array[@]} == ${len} )); then
     echo "${tag} array len OK"
   else
-    echo "${tag}: array len NOT OK"
+    echo "${tag}: array len NOT OK, expected ${len}, got ${#array[@]}"
     exit 1
   fi
 }
 
+provided_models=($(get_provided_models))
+if [ $? -ne 0 ]; then
+    echo "Failed to get provided models."
+    exit 1
+else
+    check_len "provided_models" 2 ${provided_models[*]}
+fi
+
 check_len "data_files" 2 $(data_files)
 for file in $(data_files); do
   check_file "$file" "data_files" ".tflite"
@@ -105,3 +113,6 @@ if [[ -n "$check_device" ]]; then
   check_file "$(find_device_dispatch)" "find_device_dispatch" "Example.so"
   check_file "$(find_device_runtime_lib)" "find_device_runtime_lib" "libLiteRtRuntimeCApi.so"
 fi
+
+
+# echo "$(get_provided_models)"
diff --git a/litert/integration_test/download_model_provider.sh b/litert/integration_test/download_model_provider.sh
@@ -0,0 +1,46 @@
+# Copyright 2025 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+#!/bin/bash
+
+source third_party/odml/litert/litert/integration_test/device_script_common.sh || exit 1
+
+# TODO: Unify workdirs with other scripts.
+readonly work_dir="/tmp/litert_extras"
+mkdir -p "${work_dir}"
+
+readonly url=@@url@@
+
+if [[ "$url" == "@@"*"@@" ]]; then
+  fatal "No url templated into the script."
+elif [[ -z "${url}" ]]; then
+  fatal "Url is empty."
+fi
+
+readonly target_file="${work_dir}/$(basename ${url})"
+
+if [[ ${target_file} != *".tar.gz" ]]; then
+  fatal "Target file is not a .tar.gz: ${target_file}"
+fi
+
+wget -p -O ${target_file} ${url}
+if [[ $? -ne 0 ]]; then
+  fatal "Failed to download model from ${url}."
+fi
+
+models=($(tar -xhvf ${target_file} -C ${work_dir}))
+rm -f ${target_file}
+for model in ${models[@]}; do
+  echo "${work_dir}/${model}"
+done
diff --git a/litert/integration_test/dummy_model_provider.sh b/litert/integration_test/dummy_model_provider.sh
@@ -0,0 +1,18 @@
+# Copyright 2025 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+#!/bin/bash
+
+echo "${work_dir}/dummy_model1.tflite"
+echo "${work_dir}/dummy_model2.tflite"
diff --git a/litert/integration_test/litert_device.bzl b/litert/integration_test/litert_device.bzl
@@ -133,6 +133,7 @@ def litert_device_exec(
         exec_args = [],
         remote_suffix = "",
         local_suffix = "_adb",
+        model_providers = [],
         testonly = True):
     """
     Macro to execute a binary target on a device through adb.
@@ -159,12 +160,14 @@ def litert_device_exec(
         bin = target,
         script = "//litert/integration_test:mobile_install.sh",
         exec_args = exec_args,
+        model_providers = model_providers,
         testonly = testonly,
         backend_id = backend_id,
     )
 
     # Copybara comment doesn't work right if it is inside an if statement (breaks formatting).
     if remote_suffix != None:
+        # Note model providers are not compatible with mobile harness.
         _litert_mh_exec(
             name = name + remote_suffix,
             target = target,
diff --git a/litert/integration_test/litert_device_script.bzl b/litert/integration_test/litert_device_script.bzl
diff --git a/litert/integration_test/mobile_install.sh b/litert/integration_test/mobile_install.sh
