Petastorm Documentation

The SparkDatasetConverter API allows you to convert Spark DataFrames into torch.utils.data.DataLoader objects for PyTorch training.

Workflow:

1. Set the PARENT_CACHE_DIR_URL_CONF in the Spark configuration to define the cache location.
2. Use make_spark_converter(df) to create a converter for your training and/or testing DataFrames.
3. Use the converter.make_torch_dataloader() context manager to obtain a torch.utils.data.DataLoader.
4. Call converter.delete() on each converter to remove the cached Parquet files.

from petastorm.spark import SparkDatasetConverter, make_spark_converter

# 1. Specify a cache directory
spark.conf.set(SparkDatasetConverter.PARENT_CACHE_DIR_URL_CONF, 'hdfs:/...')

df_train, df_test = ... 
model = Net()

# 2. Create converters
converter_train = make_spark_converter(df_train)
converter_test = make_spark_converter(df_test)

# 3. Use context managers for PyTorch DataLoaders
with converter_train.make_torch_dataloader() as dataloader_train:
    # dataloader_train is a torch.utils.data.DataLoader
    train(model, dataloader_train, ...)

with converter_test.make_torch_dataloader() as dataloader_test:
    test(model, dataloader_test, ...)

# 4. Clean up
converter_train.delete()
converter_test.delete()

To use Petastorm with PyTorch, use the petastorm.pytorch.DataLoader adapter. This allows you to supply custom PyTorch collating functions and transforms via a TransformSpec.

PyTorch DataLoader variants:

• petastorm.pytorch.DataLoader: The standard adapter for integrating with PyTorch training loops.
• petastorm.pytorch.BatchedDataLoader: Optimized for very large batch sizes. It buffers using Torch tensors (CPU or CUDA) and provides significantly higher throughput. It does not support Decimal or string types.
• petastorm.pytorch.InMemBatchedDataLoader: Designed for datasets that fit in system memory. It reads the dataset once and caches it in memory to avoid repeated I/O across multiple epochs.

import torch
from petastorm.pytorch import DataLoader

torch.manual_seed(1)
device = torch.device('cpu')
model = Net().to(device)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.5)

def _transform_row(mnist_row):
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])
    return (transform(mnist_row['image']), mnist_row['digit'])

transform = TransformSpec(_transform_row, removed_fields=['idx'])

with DataLoader(make_reader('file:///localpath/mnist/train', num_epochs=10, 
                            transform_spec=transform, seed=1, shuffle_rows=True), 
                batch_size=64) as train_loader:
    train(model, device, train_loader, 10, optimizer, 1)

with DataLoader(make_reader('file:///localpath/mnist/test', num_epochs=10, 
                            transform_spec=transform), 
                batch_size=1000) as test_loader:
    test(model, device, test_loader)

Petastorm datasets are stored in Apache Parquet format with additional high-level schema information for multidimensional arrays. Dataset generation is typically performed using PySpark.

Key Components:

• Unischema: Defines the dataset schema. It can render types into Spark StructType, TensorFlow tf.DType, and numpy.dtype. Each field requires a type, shape, codec instance, and a nullability flag.
• Codecs: Supports extensible data codecs like ScalarCodec, CompressedImageCodec (e.g., 'png', 'jpeg'), and NdarrayCodec.
• materialize_dataset: A context manager used to wrap the Spark dataset generation. It handles setting up Spark environment variables and writing Petastorm-specific metadata at the end of the process.
• dict_to_spark_row: Converts a Python dictionary into a pyspark.Row while ensuring compliance with the Unischema (checking shape, type, and nullability).

Workflow:

1. Define a Unischema.
2. Create a row generator function that returns a dictionary.
3. Use materialize_dataset as a context manager.
4. Parallelize your data, map it using dict_to_spark_row, and write it to Parquet using Spark.

import numpy as np
from pyspark.sql import SparkSession
from pyspark.sql.types import IntegerType

from petastorm.codecs import ScalarCodec, CompressedImageCodec, NdarrayCodec
from petastorm.etl.dataset_metadata import materialize_dataset
from petastorm.unischema import dict_to_spark_row, Unischema, UnischemaField

# The schema defines how the dataset schema looks like
HelloWorldSchema = Unischema('HelloWorldSchema', [
    UnischemaField('id', np.int32, (), ScalarCodec(IntegerType()), False),
    UnischemaField('image1', np.uint8, (128, 256, 3), CompressedImageCodec('png'), False),
    UnischemaField('array_4d', np.uint8, (None, 128, 30, None), NdarrayCodec(), False),
])


def row_generator(x):
    """Returns a single entry in the generated dataset. Return a bunch of random values as an example.""
    return {'id': x,
            'image1': np.random.randint(0, 255, dtype=np.uint8, size=(128, 256, 3)),
            'array_4d': np.random.randint(0, 255, dtype=np.uint8, size=(4, 128, 30, 3))}


def generate_petastorm_dataset(output_url='file:///tmp/hello_world_dataset'):
    rowgroup_size_mb = 256

    spark = SparkSession.builder.config('spark.driver.memory', '2g').master('local[2]').getOrCreate()
    sc = spark.sparkContext

    # Wrap dataset materialization portion. Will take care of setting up spark environment variables as
    # well as save petastorm specific metadata
    rows_count = 10
    with materialize_dataset(spark, output_url, HelloWorldSchema, rowgroup_size_mb):

        rows_rdd = sc.parallelize(range(rows_count))\
            .map(row_generator)\
            .map(lambda x: dict_to_spark_row(HelloWorldSchema, x))

        spark.createDataFrame(rows_rdd, HelloWorldSchema.as_spark_schema()) \
            .coalesce(10) \
            .write \
            .mode('overwrite') \
            .parquet(output_url)

In Petastorm 0.5.0 and later, you should use make_reader instead of instantiating the Reader class directly. make_reader simplifies the configuration of reader pools and sharding.

Key changes when moving from Reader to make_reader:

• Reader Pool: Instead of passing pool objects like ThreadPool(), use the reader_pool_type string argument ('thread', 'process', or 'dummy') and specify the number of workers via workers_count.
• Sharding: The arguments training_partition and num_training_partitions have been renamed to cur_shard and shard_count respectively.
• Shuffling: The shuffle and shuffle_options arguments are replaced by shuffle_row_groups (boolean) and shuffle_row_drop_partitions (integer).

from petastorm import make_reader
reader = make_reader(dataset_url,
                     reader_pool_type='thread',
                     workers_count=5,
                     cur_shard=1, shard_count=5,
                     shuffle_row_groups=False)

Use the petastorm-throughput.py command line tool to measure the sample throughput of a petastorm.reader.Reader for a specific dataset. This tool helps identify optimal performance by testing different parallelism configurations.

To run a basic benchmark, provide the dataset URI (e.g., file:///path/to/dataset) as an argument.

Tuning Throughput

To find the optimal throughput for your system, vary these parameters:

• -w: The number of workers used to load and decode data (can be threads or processes).
• -p: Determines whether parallelism is thread-based or process-based.
• -m: The number of warmup reads to execute before measurement to reach a steady state.
• -n: The number of measurement reads to perform.

Example Workflow

1. Generate a sample dataset:

   python examples/hello_world/generate_hello_world_dataset.py

2. Run the throughput benchmark with specific warmup and measurement counts:

   petastorm-throughput.py file:///tmp/hello_world_dataset -m 1000 -n 5000

$ petastorm-throughput.py file:///tmp/hello_world_dataset -m 1000 -n 5000

If you encounter the warning You are using a deprecated metadata version. Please run petastorm-generate-metadata.py on spark to update., you should regenerate the dataset metadata to use the new structure. This process is fast (seconds) and reuses existing schema information.

Command: Run the following command from your terminal:

petastorm-generate-metadata.py --dataset_url hdfs://namenode-host:port/path/to/dataset

Important Considerations:

• Connectivity: If running locally, PySpark might struggle with Hadoop configuration files (core-site.xml, etc.). You may need to provide the namenode hostname directly instead of a nameservice.
• Permissions: You must have write permissions to the dataset directory and the _common_metadata/_metadata files. If running on remote storage like HDFS, ensure your user has appropriate access or run the script from within a Spark environment that shares the original runtime permissions.