Make a custom rule for LiteRT WORKSPACE to enable usage of both http_archive or local_repository

LiteRT-PiperOrigin-RevId: 775364074

Author: ecalubaquib

.bazelrc

@@ -528,3 +528,12 @@ common --noenable_bzlmod
 # TODO: Migrate for https://github.com/bazelbuild/bazel/issues/7260
 common --noincompatible_enable_cc_toolchain_resolution
 common --noincompatible_enable_android_toolchain_resolution
+
+# Flags for tensorflow source rules to use local_repository rather than http_archive.
+# NOTE: due to bazel caching, when switching from http_archive to local_repository,
+# do a bazel clean --expunge --async before building or use the
+# flag --override_repository=org_tensorflow=$TF_LOCAL_SOURCE_PATH
+build:use_local_tf --action_env=USE_LOCAL_TF=true
+build:use_local_tf --action_env=TF_LOCAL_SOURCE_PATH
+test:use_local_tf --action_env=USE_LOCAL_TF=true
+test:use_local_tf --action_env=TF_LOCAL_SOURCE_PATH

WORKSPACE

@@ -26,11 +26,16 @@ http_archive(
     url = "https://github.com/bazelbuild/rules_java/releases/download/5.3.5/rules_java-5.3.5.tar.gz",
 )
 
-# TensorFlow
-# TODO(ecalubaquib): Revert this to http_archive once tflite is fully separated from tensorflow.
-local_repository(
+# Load the custom repository rule to select either a local TensorFlow source or a remote http_archive.
+load("//litert:tensorflow_source_rules.bzl", "tensorflow_source_repo")
+
+tensorflow_source_repo(
     name = "org_tensorflow",
-    path = "third_party/tensorflow",
+    sha256 = "",
+    strip_prefix = "",
+    urls = [
+        "https://github.com/tensorflow/tensorflow/archive/master.tar.gz",
+    ],
 )
 
 # Initialize the TensorFlow repository and all dependencies.

ci/build_android_package.sh

@@ -177,7 +177,8 @@ BUILD_FLAGS=("-c" "opt" \
     "--define=android_dexmerger_tool=d8_dexmerger" \
     "--define=android_incremental_dexing_tool=d8_dexbuilder" \
     "--repo_env=HERMETIC_PYTHON_VERSION=3.11" \
-    "--show_timestamps")
+    "--show_timestamps"
+    "--config=use_local_tf")
 
 # Merge extra config flags from the environment
 BUILD_FLAGS+=(${BAZEL_CONFIG_FLAGS})

ci/build_maven_with_docker.sh

@@ -65,6 +65,7 @@ else
   cp .tf_configure.bazelrc /root_dir
 
   cd /root_dir
+  export TF_LOCAL_SOURCE_PATH="/root_dir/third_party/tensorflow"
   bash /script_dir/build_android_package.sh
 
   # Bundle the Maven package

ci/build_pip_package_with_bazel.sh

@@ -16,6 +16,7 @@
 set -ex
 
 # Run this script under the root directory.
+export TF_LOCAL_SOURCE_PATH=${TF_LOCAL_SOURCE_PATH:"$(pwd)/third_party/tensorflow"}
 
 ARCH="$(uname -m)"
 TENSORFLOW_TARGET=${TENSORFLOW_TARGET:-$1}
@@ -30,25 +31,28 @@ case "${TENSORFLOW_TARGET}" in
       --copt=-march=armv7-a --copt=-mfpu=neon-vfpv4
       --copt=-O3 --copt=-fno-tree-pre --copt=-fpermissive
       --define tensorflow_mkldnn_contraction_kernel=0
-      --define=raspberry_pi_with_neon=true"
+      --define=raspberry_pi_with_neon=true
+      --config=use_local_tf"
     ;;
   rpi0)
     BAZEL_FLAGS="--config=elinux_armhf
       --copt=-march=armv6 -mfpu=vfp -mfloat-abi=hard
-      --copt=-O3 --copt=-fno-tree-pre --copt=-fpermissive
+      --copt=-O3 --copt=-fno-tree-pre --copt=-fpermissivec
       --define tensorflow_mkldnn_contraction_kernel=0
-      --define=raspberry_pi_with_neon=true"
+      --define=raspberry_pi_with_neon=true
+      --config=use_local_tf"
     ;;
   aarch64)
     BAZEL_FLAGS="--config=release_arm64_linux
       --define tensorflow_mkldnn_contraction_kernel=0
-      --copt=-O3"
+      --copt=-O3
+      --config=use_local_tf"
     ;;
   native)
-    BAZEL_FLAGS="--copt=-O3 --copt=-march=native"
+    BAZEL_FLAGS="--copt=-O3 --copt=-march=native --config=use_local_tf"
     ;;
   *)
-    BAZEL_FLAGS="--copt=-O3"
+    BAZEL_FLAGS="--copt=-O3 --config=use_local_tf"
     ;;
 esac
 

ci/build_pip_package_with_docker.sh

@@ -48,6 +48,7 @@ if [ ! -d /root_dir ]; then
 else
   export CI_BUILD_PYTHON="python${DOCKER_PYTHON_VERSION}"
   export HERMETIC_PYTHON_VERSION="${DOCKER_PYTHON_VERSION}"
+  export TF_LOCAL_SOURCE_PATH="/root_dir/third_party/tensorflow"
 
   # Running inside docker container
   cd /third_party_tensorflow

ci/run_bazel_test.sh

@@ -24,6 +24,7 @@ TEST_LANG_FILTERS="${TEST_LANG_FILTERS:-cc,py}"
 BUILD_FLAGS=(
     "--config=bulk_test_cpu"
     "--config=disable_tf_lite_py"
+    "--config=use_local_tf"
     "--test_lang_filters=${TEST_LANG_FILTERS}"
     "--keep_going"
     "--repo_env=USE_PYWRAP_RULES=True"

ci/run_bazel_test_with_docker.sh

@@ -66,6 +66,7 @@ else
   cp .tf_configure.bazelrc /root_dir
 
   export HERMETIC_PYTHON_VERSION=${DOCKER_PYTHON_VERSION}
+  export TF_LOCAL_SOURCE_PATH="/root_dir/third_party/tensorflow"
 
   cd /root_dir
   bash /script_dir/run_bazel_test.sh

litert/tensorflow_source_rules.bzl

@@ -0,0 +1,51 @@
+# This file is used to define a custom repository rule for TensorFlow submodule used by LiteRT.
+#
+# The function below allows us to select between a local TensorFlow source from local_repository
+# or a remote http_archive based on the 'USE_LOCAL_TF' environment variable.
+#
+# To use this rule, you must set the 'USE_LOCAL_TF' environment variable to
+# 'true' and the 'TF_LOCAL_SOURCE_PATH' environment variable to the absolute path of
+# the tensorflow source directory. We need to pass the absolute path because we cannot use
+# ctx.path() on a relative path for lower bazel versions.
+
+"""
+Implementation function for custom TensorFlow source repository rule.
+This rule is used to select between a local TensorFlow source from local_repository
+or a remote http_archive based on the 'USE_LOCAL_TF' environment variable.
+"""
+
+def _tensorflow_source_repo_impl(ctx):
+    use_local_tf = ctx.os.environ.get("USE_LOCAL_TF", "false") == "true"
+
+    if use_local_tf:
+        # TF_LOCAL_SOURCE_PATH must be set to the absolute path to the tensorflow source directory.
+        TF_LOCAL_SOURCE_PATH_ENV = ctx.os.environ.get("TF_LOCAL_SOURCE_PATH", "")
+        if not TF_LOCAL_SOURCE_PATH_ENV:
+            fail("""ERROR: USE_LOCAL_TF is true, but TF_LOCAL_SOURCE_PATH environment variable
+                 is not set with the absolute path to TensorFlow source.""")
+
+        local_path_str = TF_LOCAL_SOURCE_PATH_ENV  # Get the path from the environment variable
+        resolved_local_path = ctx.path(local_path_str)
+
+        for f in resolved_local_path.readdir():
+            ctx.symlink(f, f.basename)
+    else:
+        ctx.download_and_extract(
+            url = ctx.attr.urls[0],
+            sha256 = ctx.attr.sha256,
+            stripPrefix = ctx.attr.strip_prefix,
+        )
+
+tensorflow_source_repo = repository_rule(
+    implementation = _tensorflow_source_repo_impl,
+    local = False,
+    attrs = {
+        "sha256": attr.string(mandatory = False),
+        "strip_prefix": attr.string(mandatory = False),
+        "urls": attr.string_list(mandatory = True),
+    },
+    doc = """
+    A custom repository rule to select between a local TensorFlow source or a remote http_archive
+    based on the'USE_LOCAL_TF' environment variable and TF_LOCAL_SOURCE_PATH flag.
+    """,
+)