tf.data.experimental.distributed_save

Initiates the process of saving a dataset to disk using tf.data service.

tf.data.experimental.distributed_save(
    dataset: tf.data.Dataset,
    path: str,
    data_service_address: str,
    compression: str = 'AUTO'
) -> Optional[ops.OperationType]

The op uses tf.data service (https://www.tensorflow.org/api_docs/python/tf/data/experimental/service) to write a dataset snapshot. Returns immediately after submitting the request. Does not wait for the snapshot to be finished. Requires that the tf.data service run a fixed number of worker replicas.

To load the snapshot, users may optionally pass wait=True to tf.data.Dataset.load so it can read snapshots as they are being written.

Example usage:

import os
import tempfile
 
# Runs tf.data service.
tempdir = tempfile.gettempdir()
dispatcher = tf.data.experimental.service.DispatchServer(
    tf.data.experimental.service.DispatcherConfig(
        fault_tolerant_mode=True,
        work_dir=os.path.join(tempdir, "work_dir")))
dispatcher_address = dispatcher.target.split("://")[1]
worker = tf.data.experimental.service.WorkerServer(
    tf.data.experimental.service.WorkerConfig(
        dispatcher_address=dispatcher_address))
 
# Writes dataset snapshot.
path = os.path.join(tempdir, "dataset_snapshot")
dataset = tf.data.Dataset.range(1)
tf.data.experimental.distributed_save(dataset, path, dispatcher_address)
 
# Loads a dataset snapshot.
loaded_dataset = tf.data.Dataset.load(path, wait=True)
for elem in loaded_dataset:
  print(elem)
tf.Tensor(0, shape=(), dtype=int64)

dataset The tf.data.Dataset to save.
path The directory path to save the dataset. Requires that:

  • The directory does not exist and will create the directory.
  • The file system supports atomic move (rename).
data_service_address tf.data service dispatcher address.
compression (Optional.) Whether and how to compress the dataset snapshot. If "AUTO", the tf.data runtime decides which algorithm to use. If "GZIP" or "SNAPPY", that specific algorithm is used. If None, the dataset snapshot is not compressed.

An operation which when executed performs the distributed save.

ValueError If dispatcher_address is invalid.
tf.errors.AlreadyExistsError If the snapshot has already started or has finished.
tf.errors.FailedPreconditionError If the file system does not support atomic move (rename).
tf.errors.InvalidArgumentError If tf.data service is not running in the fault tolerant mode.