Getting Tensorflow Extended (TFX) to work with apple silicon (m1, etc.)

If you’re using Tensorflow and trying to get a model into production, the officially sanctioned toolset for doing so is Tensorflow Extended (TFX). Unfortunately, if you’re on apple silicon as of 28/11/2023, there is no support for tfx (open issue here). If you try and install tfx in your mac environment (1.14.0 being the latest), you will be told there are no installation candidates or version, with varying reasons depending on your set up. From wrong python versions (if you’re not on python 3.9 for ml-metadata) to unavailable packages for your architecture (tfx-bsl which requires x86_64)

UPDATE: See this post for how to get TFX master working w/TF 2.15.0 or this post for TFX 1.14.0

If you’re willing to take a performance hit of running tfx in docker with rosetta2, you can get around this by running a devcontainer in vscode with micromamba. Here is the magic incantation Dockerfile required, referenced from your devcontainer.json:

FROM --platform=linux/amd64 ghcr.io/mamba-org/micromamba

USER root

RUN apt-get update && apt-get install -y --no-install-recommends \
build-essential \
&& rm -rf /var/lib/apt/lists/*

USER $MAMBA_USER
COPY --chown=$MAMBA_USER:$MAMBA_USER env.yaml /tmp/env.yaml
RUN micromamba install -y -n base -f /tmp/env.yaml && \
micromamba clean --all --yes

and the env.yaml:

name: base
channels:
- conda-forge
dependencies:
- tensorflow=2.13.1
- python=3.10
- pip:
- tfx==1.14.0
- jsonschema==4.17.3

Note the use of platform=linux/amd64 in the Dockerfile to get compatible pypi distributions. Big credit to this issue which allowed me to find compatible versions of jsonschema to use with tfx 1.14.0 .