Introduction

Machine Learning Metadata (MLMD) is an integral component of TensorFlow Extended (TFX), facilitating the management and orchestration of metadata in ML workflows. MLMD offers a structured and organized approach to track artifacts, executions, and lineage of the machine learning pipeline, significantly enhancing reproducibility, debugging, and interoperability. This article dives into how data is stored in MLMD within TFX, its structure, and its critical role in operationalizing machine learning models.

Understanding MLMD and Its Role in TFX

TensorFlow TFX is an end-to-end platform for deploying production-ready machine learning pipelines. In this context, MLMD acts as a central hub to track and manage all the metadata associated with these pipelines. It stores information in a metadata store, persisting data related to various components of the ML workflow, like dataset versions, trained models, evaluation metrics, and more.

Key Concepts in MLMD

• Artifacts: These are the data items produced or consumed by components in a pipeline, such as datasets, models, or any form of data that can be versioned and tracked. MLMD maintains a record of these artifacts, storing metadata like version, type, and URI location.
• Executions: Executions represent components' run instances within a pipeline. Each execution corresponds to a specific operation and records metadata about the runtime environment, including parameters used, execution time, and status. It helps in tracing the computation history.
• Contexts: Contexts provide a way to group executions and artifacts that are related in some logical way, such as all executions of a particular pipeline or all artifacts of a specific user. Contexts are useful for organizing and managing entities within a metadata store.

Metadata Storage Structure

MLMD's data model is built on connected graphs representing the relationships between different components of a machine learning pipeline. At a high-level view, MLMD distinguishes between types and instances:

1. Types: Define a schema of metadata entities, like artifact types and execution types.
2. Instances: Represent specific records or states of metadata entities.

The hierarchical structure ensures that all components are traceable through their interconnections.

MLMD CRUD Operations

MLMD provides a variety of operations to manipulate and query metadata:

• get_artifacts() : Fetches artifacts stored in the metadata store.
• put_artifacts() : Persists artifact metadata to the store.
• get_executions() : Retrieves execution records.
• put_executions() : Adds new execution entries to the store.
• get_contexts() : Retrieves context entries.

Practical Example

Consider a TFX pipeline with three components: ExampleGen , Trainer , and Evaluator . When the pipeline runs, MLMD tracks:

• Input Examples generated by ExampleGen .
• The model artifact produced by Trainer .
• Evaluation metrics output by Evaluator .

These components' interactions and outputs are orchestrated and stored in MLMD as artifacts, executions, and contexts, providing full traceability and detailed logging of each step.

Here's a code snippet illustrating how to record an artifact in MLMD: