Getting Tensorflow Extended (TFX) master to work with apple silicon natively

Dec 19, 2023 · 662 words · 4 minute read

TFX is not compatible with Apple Silicon yet, though there are a few pull requests in flight to make this happen. In my previous post I ran through how to get it working with TFX version 1.14.0. This post covers how to get TFX master working against Tensorflow 2.15.0, using nightly components. Very similar, with an additional TFX build and an added twist for tensorflow-serving-api.

You will need to clone or build patched master versions of

  1. google/ml-metadata [Pull Request] [Branch]
  2. tensorflow/tfx-bsl [Pull Request] [Branch],
  3. tensorflow/data-validation [Pull Request] [Branch]
  4. tensorflow/serving [Branch]
  5. tensorflow/tfx [Branch]

Installing them to a venv or conda/mamba environment should then allow you to install tfx==1.14.0. All the usual caveats of installing hand-rolled versions of libraries apply. Here be dragons!

Building and installing

Pre-requisites:

  1. Xcode >=15 (Xcode Command Line tools being at 15 is not enough, you really want Xcode)
  2. Cmake and Bazelisk (homebrew install works)
  3. Python 3.9
  4. Docker for Mac

I’ll use micromamba but venv/conda like environments should work too. I’ve tried this with an M1 Pro on Sonoma 14.2 and Xcode 15.0, Python 3.9 and micromamba but YMMV

Steps:

  1. Create and activate your environment 
    micromamba create -n tfx-nightly -c conda-forge python=3.9
micromamba activate tfx-nightly
  2. Pin the version of Bazel to 5.3.2 (from ml-metadata’s version) 
    export USE_BAZEL_VERSION=5.3.2
  3. Prepare environment for building nightly tfx components 
    export TFX_DEPENDENCY_SELECTOR=NIGHTLY
  4. Clone and build the required projects:

    1. google/ml-metadata

      git clone https://github.com/nicholasjng/ml-metadata.git
cd ml-metadata
git checkout m1fix
python setup.py bdist_wheel
pip install dist/ml_metadata-1.15.0.dev0-cp39-cp39-macosx_11_0_arm64.whl

    2. tensorflow/tfx-bsl

      git clone https://github.com/tangm/tfx-bsl.git
cd tfx-bsl
git checkout 48-Allow-compilation-on-m1-macs
pip install numpy # (per `tfx-bsl` source building instructions)
python setup.py bdist_wheel
pip install --extra-index-url https://pypi-nightly.tensorflow.org/simple dist/tfx_bsl-1.15.0.dev0-cp39-cp39-macosx_11_0_arm64.whl
      • We add the additional tf-nightly pypi repository as tfx-bsl has a dependency on tensorflow/metadata 1.15.0-dev0. We could build it ourselves as well, without any patches from tensorflow/metadata:master

    3. tensorflow/data-validation

      git clone https://github.com/tangm/data-validation.git
cd data-validation
git checkout 205-allow-apple-silicon
python setup.py bdist_wheel
pip install dist/tensorflow_data_validation-1.15.0.dev0-cp39-cp39-macosx_11_0_arm64.whl

    4. tensorflow/serving

      # prepare the output dir
mkdir serving-dist
cd serving-dist
docker run -it -v $(pwd):/output tensorflow/serving:nightly-devel bash

# in the docker container:
cd ~
git clone https://github.com/tangm/serving
git checkout m1-local-build
export USE_BAZEL_VERSION=6.4.0
bazel build --color=yes --curses=yes --verbose_failures --output_filter=DONT_MATCH_ANYTHING --config=release tensorflow_serving/tools/pip_package:build_pip_package
bazel-bin/tensorflow_serving/tools/pip_package/build_pip_package /output 
exit

# back in the real world
pip install tensorflow_serving_api-2.15.0-py2.py3-none-any.whl

      There’s alot to unpack here.

      • We need to build the python package for tensorflow-serving-api, which is heavily tied to the tensorflow version. The latest version available is 2.14.1 , which is a problem since the tfx nightly packages depend on the latest tensorflow, 2.15.0.
      • The changes on my branch are loosely based on the release branch commits for tensorflow-serving 2.14.x
      • You can find the instructions for how to build the python package in the provided Dockerfile for development environments. I also tried building it by running the bazel build command locally and using the provided tools/run_in_docker.sh but ran into strange issues like
      ERROR: <builtin>: BazelWorkspaceStatusAction stable-status.txt failed: Failed to determine workspace status: Process exited with status 127
/bin/sh: /proc/self/cwd/tools/gen_status_stamp.sh: No such file or directory
Target //tensorflow_serving/tools/pip_package:build_pip_package up-to-date:
bazel-bin/tensorflow_serving/tools/pip_package/build_pip_package

      So I’ve just isolated the commands required, running it in the provided dev environment.

    5. And finally install tfx!

      git clone https://github.com/tangm/tfx.git
cd tfx
git checkout m1-local-build
# as per package_build/README.md
package_build/initialize.sh
python package_build/ml-pipelines-sdk/setup.py bdist_wheel
python package_build/tfx/setup.py bdist_wheel

pip install --extra-index-url https://pypi-nightly.tensorflow.org/simple dist/tfx-1.15.0.dev0-py3-none-any.whl dist/ml_pipelines_sdk-1.15.0.dev0-py3-none-any.whl jsonschema==4.17.3
      • The branch has a change to allow the version of TF to be specified, otherwise a version conflict will occur. I have no idea how to build TFX dependencies (with constraint tensorflow>=2.15.0,<3) in conjunction with the tfx repo constraint tensorflow>=2.13.0,<2.14.
      • Note the pin for jsonschema to aid in version resolution.

Sanity Testing

python -c "from tfx import version ; print('TFX version: {}'.format(version.__version__))"

Should show 1.15.0.dev

Using the penguin template:

export PIPELINE_NAME=sanity_check
export PROJECT_DIR=$PWD/$PIPELINE_NAME
tfx template copy \
  --pipeline_name="${PIPELINE_NAME}" \
  --destination_path="${PROJECT_DIR}" \
  --model=penguin
cd sanity_check
tfx pipeline create --engine=local --pipeline_path=local_runner.py
tfx run create --engine=local --pipeline_name="${PIPELINE_NAME}"

If you update the pipeline/pipeline.py file to uncomment other components like in the tutorials, remember to update the pipeline before running it again

tfx pipeline update --engine=local --pipeline_path=local_runner.py
tfx run create --engine=local --pipeline_name="${PIPELINE_NAME}"

And we are done!