Amazon SageMaker Fast File Mode. Methods for Streaming Training Data... | by Chaim Rand | Towards Data Science

Last year we published a blog post in which we surveyed different methods for streaming training data stored in Amazon S3 into an Amazon SageMaker training session. We highlighted some of the strengths and weaknesses of the different options and examined their ability to address some specific needs such as:
dynamic boosting, which requires changing how the data is sampled mid-training,the freedom to choose data formats other than the TFRecord format, andrandom access to data records.
dynamic boosting, which requires changing how the data is sampled mid-training,
the freedom to choose data formats other than the TFRecord format, and
random access to data records.
Last month AWS announced a new mechanism called Fast File Mode (FFM) for training in Amazon SageMaker with data from Amazon S3. In this post we will review some of the properties of this new method and compare it to some of the alternative methods we reviewed in our previous post.
As in our previous post, the scenario we address is one in which the size of our dataset is so large that it is either impossible or inefficient to download it in its entirety to a local disk before training. Instead, we stream the data directly from S3 into the training loop. We will focus on training in TensorFlow, specifically version 2.6, although most of the observations we will make are just as relevant to other training frameworks, as well.
To program your Amazon SageMaker training session to stream data from S3 using Fast File Mode you need to set the input_mode to FastFile and enter your dataset S3 path(s) to the fit function, as shown in the code block below:
from sagemaker.tensorflow import TensorFlowestimator=TensorFlow(entry_point='main.py',                     role=<AWS IAM role>,                     py_version='py37',                     framework_version='2.6.0', #tf version                     instance_count=1,                     instance_type='ml.g4dn.xlarge',                     input_mode='FastFile'                    )train_data="s3://sagemaker-path-to-train-data"test_data="s3://sagemaker-path-to-test-data"estimator.fit({'train':train_data, 'test':test_data})
Here we have defined two channels named train and test. The training job will start up with two environment variables SM_CHANNEL_TRAIN and SM_CHANNEL_TEST corresponding to the two data channels. Each will be set to a local path that appears to the training application as a POSIX file system mount of all of the data within programmed paths in S3. Thus, in order to access a file at s3://sagemaker-path-to-train-data/relative-path, rather than downloading the file (e.g. using boto3) the training application would simply open:path=os.path.join(os.environ["SM_CHANNEL_TRAIN"],<relative-path>)For example, if your files are stored in TFRecord format you could use TFRecordDataset to create a dataset as in the code block below:
paths=[<list of relative filepaths>]local_root=os.environ["SM_CHANNEL_TRAIN"]filepaths=[os.path.join(local_root,path) for path in paths]ds=tf.data.TFRecordDataset(               filepaths,                num_parallel_reads=tf.data.experimental.AUTOTUNE)
In our previous blog post we covered several options for streaming data from S3 into an Amazon SageMaker training session. In this section we highlight some of the advantages of FFM compared to other methods, with particular focus on the Amazon SageMaker Pipe Mode alternative. For more details on Pipe Mode check out this previous post.
As we have discussed in the past, the number of data streams that is supported by Pipe Mode is limited to 20 (as of the time of this writing). Using multiple data streams can solve certain needs such as data partitioning (where we partition data according to class) or stream duplication (e.g. streaming the same dataset with unique shuffling to multiple processes when performing data distributed training). As detailed in this blog post, the restriction on the number of channel streams can pose a challenge that may require creativity to overcome. With FFM, multiple datasets can be defined using a single data channel configuration. All you need to do is to configure the channel’s base S3 path to be the root of all the datasets.
When you program a job to use Pipe Mode you configure a fixed set of files to be streamed through the pipe and may optionally delegate shuffling and/or sharding to the pipe mode manager. Once the training starts you cannot change the contents of the pipe or its configuration. Relative to pipe mode, FFM offers far greater control:
Shuffling — you can apply your own data shuffling policy as well as capture the randomly generated parameters.
Sharding — you can define your own sharding policy.
Updating dataset contents during training — you can change the mixture of the data while the model is training.
File chunk support — FFM supports reading a portion of a file without pulling the entire file. This may have implications on your choice of file format for your data in S3 as some formats rely on file chunking. Note that while file chunking could be used (assuming an appropriate file format) to implement random access to data records, this may have severe implications on performance. However, you could implement some level of random access by analyzing the minimum acceptable chunk size for your machine learning project.
In the code block below we demonstrate how a custom generator can be created that controls (and records) shuffling, sharding, and changing the data mixture mid-training.
import randomdef generator():    paths=... # optionally shard filepaths    for e in range(num_epochs):        seed=random.randrange(100)        print(f"starting epoch {e}. Random seed is {seed}")        random.seed(seed)        random.shuffle(paths)        # optionally modify contents based on current metrics                   paths=...         for p in paths:            yield pfilepaths_ds=tf.data.Dataset.from_generator(               generator,               output_types=tf.string,               output_shapes=())ds=tf.data.TFRecordDataset(               filepaths_ds,                num_parallel_reads=tf.data.experimental.AUTOTUNE)
See here for an example of how to employ this scheme to implement dynamic boosting with two classes.
While Pipe Mode can be configured to stream any binary data, and thus any file format, the high level Amazon SageMaker SDK provides support only for the TFRecord file format, via the PipeModeDataset. Using any other file format with pipe mode would require a bit of effort. Contrary to pipe mode, FFM does not favor any file format over the other.In a previous post, we surveyed several different file format options for training in TensorFlow. Using FFM you are free to use any one of them.
Keep in mind that there are best practices for storing files in the cloud. In particular, these practices dictate optimal file sizes. We typically store our data in files with sizes that are a few hundred MB.Another consideration is the number of files. According to the FFM documentation, the training start up time may be impacted by the total number of files.
As demonstrated above, when using FFM we can use the same standard TensorFlow dataset classes that we would use in a local environment. This is contrary to pipe mode that requires a special custom dataset, PIPEmodeDataset. This has several implications:
Easier to program — fewer modifications are required to adapt your code from a local training session to a cloud based training session.
Easier to debug — the similarity to a local run makes it easier to debug issues.
API support — the use of standardized TensorFlow datasets means automatic support for all TensorFlow dataset utility APIs.
In this section we provide a performance comparison between FFM, Pipe Mode, and direct S3 access on a dataset comprised of TFRecord files. The metrics we have compared are: 1. the average data throughput, measured in the number of records per second, and 2. the average CPU utilization. Measuring the CPU utilization is important in order to assess the degree to which the input mode might contribute to potential CPU bottlenecks. The code we have used to perform the test can be found in the code block below:
count = 0begin = time.time()stime = beginfor x in ds:  count = count + 1  if count % 1000 == 0:    etime = time.time()    print(f"step: {count} step/sec: {float(1000)/(etime-stime)}")    stime = time.time()  if count==100000:    breakprint(f"average step/sec: {float(100000)/(time.time()-stime)}")
Note that we isolate our analysis to the data input portion of the pipeline by iterating over the dataset without actually feeding the data into the training pipeline. The tests were performed on an ml.c5.xlarge instance type. The results are shared in the chart below. In the test we have run, FFM has the best (lowest) step time with Pipe Mode coming in at a close second. At the same time, note the increase in CPU utilization. These results are provided as one example of the type of comparative metrics you might see. These metrics may be highly dependent on your model architecture and the results you get may be very different. For example, if your step time is currently solely determined by your GPU activity, you may be wholly indifferent to the data input mode.
The recently announced Amazon SageMaker Fast File Mode provides a new method for efficient streaming of training data directly into an Amazon SageMaker training session. It offers the potential performance of Pipe Mode combined with the conveniences and flexibility of a local dataset. Importantly, every machine learning project is unique; performance behaviors observed in one model may be wholly different than those observed in another. It is critical that you evaluate the performance impact of adopting FFM before converting your project and recommended that you code your solution in a way that makes switching between streaming options easy. Best of luck!