Dockerfile

# Two targets of interest here, "test" and "wheel". These are invoked
# as follows (from the butterfly/tensorflow source directory):
#
# docker build --target [test|wheel] -t <IMAGE[:TAG]> .
# 
# The "test" target is used to run tests on our custom operators.
# The "wheel" target is used to build a wheel for running python tflite
# models with our custom ops.

# If you want to pull the wheel out of the built image by hand, try this:
#
# whl_name=tflite_runtime-2.9.2-cp310-cp310-linux_x86_64.whl
# whl_path=/usr/src/app/tensorflow/lite/tools/pip_package/gen/tflite_pip/python3/dist/${whl_name}
# id=$(docker create [docker-image-name])
# docker cp $id:$whl_path $whl_name
# docker rm -v $id

# --- BASEL SETUP -------------------------------------------------------------

FROM golang:1.17 as base
RUN go install github.com/bazelbuild/bazelisk@latest && \
    ln -s $(go env GOPATH)/bin/bazelisk /usr/local/bin/bazel


# --- TF ENVIRONMENT & WHEEL GENERATION ---------------------------------------

FROM python:3.10 as build
COPY --from=base /usr/local/bin/bazel /usr/local/bin/bazel
RUN apt-get update && \
    apt-get install -y default-jdk && \
    apt-get clean
ENV JAVA_HOME /usr/lib/jvm/default-java
RUN export JAVA_HOME
RUN pip install -U pip numpy==1.23.3 wheel packaging requests opt_einsum && \
    pip install -U keras_preprocessing --no-deps
WORKDIR /usr/src/app
COPY . /usr/src/app
# Defaults for all configuration inputs.
RUN yes '' | ./configure
RUN tensorflow/lite/tools/pip_package/build_pip_package_with_bazel.sh


# --- WHEEL COPY AND IMAGE TRIM -----------------------------------------------

FROM debian:bullseye-slim as wheel
WORKDIR /usr/src/app
COPY --from=build /usr/src/app/tensorflow/lite/tools/pip_package/gen/tflite_pip/python3/dist/tflite_runtime-2.9.2-cp310-cp310-linux_x86_64.whl .


# --- TEST COMMAND ARGS -------------------------------------------------------

FROM build as test
ARG FLAGS=--experimental_repo_remote_exec

# Enumeration of targets that are butterfly specific.
CMD bazel test $FLAGS //tensorflow/lite/kernels:dlops_bline_coordinates_test \
    //tensorflow/lite/kernels:connected_components_test \
    //tensorflow/lite/kernels:unique_with_counts_test \
    //tensorflow/lite/kernels:dlops_class_persisted_masks_test