Implementation:Tensorflow Serving Graph Rewriter

Knowledge Sources	Tensorflow_Serving
Domains	Model Serving, Graph Optimization
Last Updated	2026-02-13 00:00 GMT

Overview

Provides a global registry for a single graph rewrite function that is applied to all MetaGraphDefs immediately after loading, before session creation.

Description

The GraphRewriter class implements a singleton-based global registration pattern for a graph rewriting function. The registered function receives a mutable tensorflow::MetaGraphDef pointer and can modify the graph definition before any session or runtime is created from it.

Key design aspects:

• Singleton pattern: GetGlobal() returns a static singleton instance, ensuring a single global graph rewriter per process.
• Thread safety: All access to the rewriter function is protected by an absl::Mutex with ABSL_GUARDED_BY annotations.
• Single registration: Set() returns an AlreadyExists error if a rewriter has already been registered, preventing accidental overwrites.
• IsRegistered(): Allows callers to check if a rewriter is available before attempting to use it.
• ResetForTesting(): Sets the rewriter back to nullptr for test isolation.

Two convenience inline functions provide a simpler interface:

• SetGraphRewriter: Registers a graph rewrite function globally (can only be called once per process).
• ResetGraphRewriterForTesting: Resets the rewriter (for testing only).

The rewriter is called by TfrtSavedModelFactory during model loading, between MetaGraphDef reading and TFRT SavedModel creation. This EXPERIMENTAL API may change or be removed.

Usage

Use this module to register a custom graph transformation that should be applied to all loaded models. Common use cases include custom optimizations, op replacements, or graph modifications needed for specific deployment targets. Register the rewriter at process startup before any models are loaded.

Code Reference

Source Location

• Repository: Tensorflow_Serving
• File: tensorflow_serving/session_bundle/graph_rewriter.h (lines 1-91)

Signature

class GraphRewriter {
 public:
  static GraphRewriter& GetGlobal();
  Status Set(std::function<Status(tensorflow::MetaGraphDef*)>&& rewriter);
  std::function<Status(tensorflow::MetaGraphDef*)>& Get();
  bool IsRegistered();
  Status ResetForTesting();
};

inline Status SetGraphRewriter(
    std::function<Status(tensorflow::MetaGraphDef*)>&& rewriter);
inline Status ResetGraphRewriterForTesting();

Import

#include "tensorflow_serving/session_bundle/graph_rewriter.h"

I/O Contract

Inputs

Name	Type	Required	Description
rewriter	std::function<Status(MetaGraphDef*)>	Yes	Function that modifies a MetaGraphDef in place

Outputs

Name	Type	Description
Set() return	Status	OK on success; AlreadyExists if a rewriter was already registered
Get() return	std::function<Status(MetaGraphDef*)>&	Reference to the registered rewriter function (may be nullptr)
IsRegistered()	bool	True if a rewriter has been registered

Usage Examples

Registering a Graph Rewriter

// At process startup
Status status = SetGraphRewriter(
    [](tensorflow::MetaGraphDef* meta_graph_def) -> Status {
      // Custom graph modifications
      return OkStatus();
    });

// During model loading (done automatically by TfrtSavedModelFactory)
if (auto& rewriter = GraphRewriter::GetGlobal(); rewriter.IsRegistered()) {
  TF_RETURN_IF_ERROR(rewriter.Get()(&meta_graph_def));
}

Related Pages

• Principle:Tensorflow_Serving_Graph_Rewriting