ML Metadata: recording what your pipeline actually did

A metadata store, and the question it exists to answer

ML Metadata, usually shortened to MLMD, is a library for recording and retrieving metadata associated with ML developer and data scientist workflows. The primary language in the repository is C++, and it is Apache-2.0 licensed, with the homepage pointing at the TFX guide.

What it stores is the structure of a machine learning workflow: which datasets, which models, which metrics, and how they relate. The description says it records and retrieves metadata associated with those workflows, and the project documentation refers to a getting started guide for background and instructions.

The important thing about the scope is what the library is not. It does not run your training, it does not read your data, and it does not instrument your code. It is the place where the facts about a workflow live once you have decided to record them, and the value of it is entirely in the vocabulary it gives you: a dataset is a type of artifact, a model is a type of artifact, an execution is a type of event, and the relationships between them are first class records rather than strings you put in a table.

That is the difference from the usual homegrown solution. A notebook that writes a JSON file next to its output, or a spreadsheet of run parameters, records the same facts in a form that cannot answer a question like which datasets were used to train the model that is currently serving traffic. A store with a schema can.

Python package, C++ core, and an optional server binary

The install instructions show the shape of the project, and the first sign is that there are three things you can install.

The recommended way is the PyPI package:

pip install ml-metadata

And then you import the relevant packages:

from ml_metadata import metadata_store
from ml_metadata.proto import metadata_store_pb2

The second import is the interesting one, because the name says what it is. The proto suffix means the generated Protocol Buffers types, which is the wire format of the store, and it is the part you touch when you are constructing or interpreting artifacts rather than calling convenience helpers.

The third thing is a server. The source build has an optional step to build the gRPC server binary:

bazel build -c opt --define grpc_no_ares=true  //ml_metadata/metadata_store:metadata_store_server

That target name, metadata_store_server, is the whole architecture in one line. The store is a service with a client and a server, and the Python package is a client for it. The grpc_no_ares define is there because the c-ares DNS resolver adds a dependency that the build does not want, which is the kind of detail you only find in a build flag.

The repository layout confirms the split: ml_metadata/ holds the package and the C++ code including the docker build files, and there is a separate docs/ directory with a mkdocs.yml, so the documentation is built with MkDocs and lives beside the code.

Three ways to install, and one of them is the one Google uses

There are three installation paths and the documentation is clear about which one it prefers for which purpose.

PyPI is the recommended way to install ML Metadata, and it is what a user should normally do. Nightly packages are also published, at pypi-nightly.tensorflow.org on Google Cloud, and they are installed with the extra index URL:

pip install --extra-index-url https://pypi-nightly.tensorflow.org/simple ml-metadata

The nightly index is the TensorFlow one, which tells you that this project shares its release infrastructure with the rest of that stack, and a nightly is what you install when you need a fix that has not reached PyPI.

Then there is Docker, and the documentation makes a stronger claim about it: this is the recommended way to build ML Metadata under Linux, and it is continuously tested at Google. The procedure is to install docker and docker-compose, and then at the project root:

DOCKER_SERVICE=manylinux-python${PY_VERSION}
sudo docker compose build ${DOCKER_SERVICE}
sudo docker compose run ${DOCKER_SERVICE}

where PY_VERSION is one of 310, 311, 312 or 313. A wheel is produced under dist/ and installed with:

pip install dist/*.whl

The manylinux naming is the reason for the docker route rather than a native build. The compose file in the repository defines a service per Python version, each based on a manylinux2010 image and each with the project root mounted at /build, which is the WORKDIR of the image. The comment in that file explains the repetition: extensions are not supported until Python 3.4 in that base, hence the repeated boilerplate per version.

Building from source means Bazel, cmake and a version decision

The source installation is five steps and the second one is the one people get wrong.

The prerequisites are Bazel and cmake, each with a link to its own installation instructions, and then the repository is cloned:

git clone https://github.com/google/ml-metadata
cd ml-metadata

The note on that step is important. These instructions install the latest master branch of ML Metadata, and if you want a specific branch, such as a release branch, you pass -b <branchname> to the git clone command. Since MLMD may be backwards incompatible before version 1.0, the version you clone is the version you inherit.

The pip package is then built with Bazel, driven through setup.py:

python setup.py bdist_wheel

and the generated .whl file is in the dist subdirectory, installed with pip install dist/*.whl. The optional server build is the bazel command shown earlier.

The setup.py file is where the C++ half of the project shows through. It defines a custom build command that overrides distutils build so that the Bazel build runs first, before any of the sub commands, because the build command is also invoked from bdist_wheel and from install, and therefore from pip install . and python setup.py install. The custom class also sets the build base to build_mlmd_tmp, and the file carries a compatibility shim for Python 3.11 and later, where distutils was removed from the standard library, with a comment that setuptools is imported first to avoid legacy distutils behaviour.

Platforms, and the one that is on the way out

MLMD is built and tested on 64-bit operating systems, and the list has one line that reads differently from the others.

macOS 10.14.6 (Mojave) or later. Ubuntu 20.04 or later. And then Windows 10 or later, marked DEPRECATED, with a pointer: for a Windows-compatible library, please refer to MLMD 1.14.0 or earlier versions.

That is a real fork in the road for a Windows user. The current line does not support Windows, and the escape hatch is to pin an older release, which means pinning an older Python, an older dependency set and missing whatever has been added since. The Python badge in the readme lists 3.10, 3.11, 3.12 and 3.13, and the docker compose services cover 310, 311 and 312, which is a slightly different range from the badge.

The licence and the versioning note complete the picture. The project is Apache-2.0, and the readme carries a warning that ML Metadata may be backwards incompatible before version 1.0, which is a line worth taking literally even now, because the 1.x releases are the ones the current documentation describes.

The repository also carries a test configuration that is worth knowing about. The pytest settings put testpaths at ml_metadata/metadata_store/ and the python_files pattern at *_test.py, so the test suite is run against the store specifically and the naming convention is not the pytest default.

Where ML Metadata is the wrong tool

The first limit is that nothing captures the metadata for you. The library is a store and a vocabulary, so a pipeline has to be modified to write to it, and a pipeline that was not written to record its artifacts produces an empty store that looks authoritative. A tool that logs everything for you, and a tool that stores what you tell it, are different products, and only one of them is this.

The second is the C++ and Bazel build. A team whose pipeline tooling is pure Python now has a project that needs Bazel, cmake, a manylinux container and, for the server, a gRPC build, in order to use a Python client. The PyPI path avoids all of that, but anyone who needs to build from source, patch the C++ core, or produce a custom distribution takes on the whole toolchain.

The third is that the schema is the hard part and the documentation does not design it for you. The readme points at a getting started guide for background and instructions, and the visible readme does not describe the type system, the artifact properties, the event structure or the lineage model. Any team evaluating this is signing up to reading the guide and the generated API documentation before it can judge whether the vocabulary fits, and the model has to fit, because the value of a metadata store collapses if the types do not match how your organisation actually describes its work.

The fourth is the platform situation. Windows is deprecated, and the fallback is an old release. A Windows shop is choosing to run an outdated version of a schema-driven store, and schema-driven stores do not age well, because the metadata of a project changes as the project changes and an old vocabulary will not describe new work.

The fifth is the scale question. The store is a gRPC service, which is a good design for a shared system of record that several services query, and a heavier one than a local file. A team with three notebooks and a shared drive does not need a service, and a team with hundreds of pipelines and a governance requirement probably does.

Editorial conclusion

ML Metadata is worth adopting for a team that has to answer where a model came from, because it gives the lineage a home, because the store is a gRPC service that other tools can query rather than a log file, and because the artifact type system makes a dataset, a model and a metric first class records instead of strings in a spreadsheet. It is the wrong tool if you want the metadata captured for you, because nothing here instruments your pipeline, and the wrong tool if you only need to look up one run. Before you adopt it, check the three things the project states about itself: that ML Metadata may be backwards incompatible before version 1.0, so a pinned version matters; that the supported platforms are macOS 10.14.6 or later, Ubuntu 20.04 or later, and Windows 10 or later is deprecated with MLMD 1.14.0 or earlier recommended for a Windows-compatible library; and that the build needs Bazel and cmake, which is a heavier source-build than a pure Python project would suggest.

Frequently asked questions

Official sources

1. google/ml-metadata on GitHub
2. License: Apache-2.0
3. Project website
4. README
5. Releases