Memory leak in TFX Transform when using Apache Beam runner

[wsuchy]

There is a dataset https://www.kaggle.com/generall/oneshotwikilinks (~6GB uncompressed) that has been transformed into parquet format:

df = pd.read_csv("shuffled_dedup_entities.tsv", sep='\t', names=['concept', 'left','mention','right'])
df.to_parquet("/tmp/wikilinks.parquet")

The code shown here: https://gist.github.com/wsuchy/90fcb59d7a97e377634e096b030abfff will consume more than 120GB (which is a server max) of memory even though its only task is to read from parquet, transform, and store as TFRecords.

Is there a problem with the code, or rather TFX / Apache Beam?

[rmothukuru]

@wsuchy,
Can you please share Dummy Data so that we can try reproducing the Memory Leak Issue with your code. Thanks!

[wsuchy]

@rmothukuru I've provided you the link to Kaggle dataset and steps to create parquet file from it. If you still need the file then please provide a place where I could upload it (~3.6GB compressed)

[wsuchy]

System config:

tensorflow-transform==0.14.0
tensorflow==1.14.0
tensorflow-metadata==0.14.0
sys.version: 3.7.3 (default, Mar 27 2019, 22:11:17) [GCC 7.3.0]

[DmitriyBaykov]

Hi, all. I'm facing a similar issue. Due to my investigations, the reason is beam.io.WriteToTFRecord function. @wsuchy Did you manage to fix this or somehow avoid it? I would be appreciated for any help. My config is:
tensorflow 2.1.0 tensorflow-metadata 0.21.2 tensorflow-transform 0.21.2

[DmitriyBaykov]

I was able to avoid a memory leak. As I've mentioned previously the issue is with beam.io.WriteToTFRecord, so I reimplemented it with beam.ParDo function. Here's some code in case if someone faced same issue.
Replace beam.io.WriteToTFRecord in pipeline with this line of code:
_ = (train_dataset | 'WriteTrainDataset' >> beam.ParDo(WriteToTFRecord('save_path')))

Where WriteToTFRecord function is class:

class WriteToTFRecord(beam.DoFn):
     def __init__(self, save_path):
        self.save_path = save_path

     def start_bundle(self):
        self.writer = tf.io.TFRecordWriter(self.save_path)

    def process(self, transformed_dict):
        self.writer.write(serialize_dict(transformed_dict))

    def finish_bundle(self):
        self.writer.close()

    def serialize_dict(tensors_dict):
        example_string = {}
        for key in tensors_dict:
            example_string[key] = create_feature(tensors_dict[key])
        return tf.train.Example(features=tf.train.Features(feature=example_string)).SerializeToString()
    
    def create_feature(value):
        if type(value) == list or type(value) == np.ndarray:
            if type(value[0]) == str:
                return tf.train.Feature(bytes_list=tf.train.BytesList(value=[val.encode('utf-8') for val in value]))
            elif type(value[0]) == np.int64:
                return tf.train.Feature(int64_list=tf.train.Int64List(value=value))
            elif type(value[0]) == np.float32 or type(value[0]) == np.float64:
                return tf.train.Feature(float_list=tf.train.FloatList(value=value))
            elif type(value[0]) == np.ndarray:
                return create_feature(np.concatenate(value, axis=0))
            elif  type(value[0]) == bytes:
                return tf.train.Feature(bytes_list=tf.train.BytesList(value=value))
        elif type(value) == str:
            return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value.encode('utf-8')]))
        elif type(value) == np.int64 or type(value) == int:
            return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))
        elif type(value) == np.float32 or type(value) == np.float64 or type(value) == float:
            return tf.train.Feature(float_list=tf.train.FloatList(value=[value]))

You can remove/add some ifs in create_feature func.
This code works for me very fast with low memory usage.

[vaskozl]

I'v also found that beam.io.WriteToTFRecord was using progressively more memory as the database was being written the record artifacts. This was only really obvious after removing reshuffle from the ExampleGen as per tensorflow/tfx#1907

Thank you to @DmitriyBaykov for the suggestions, I'm now also using tf.io.TFRecordWriter instead and memory usage stays low even for large datasets.

[vaskozl]

If someone has free time, could you look into this and tensorflow/tfx#1907 when running with DirectRunner?

I'm currently patching this on every release, as otherwise the containers OOM if using a dataset with anything bigger than a few gigabytes (even though the container has 32GB) of memory.