Tfx: StatisticsGen fails with OutOfMemoryError error while using Flink, external beam worker pool and large amount of data

[ferryvg]

Hi there!

Flink 1.10.3
Beam 2.27.0
TFX (and subpackages) 0.27.0
MiniO (S3) as storage

I'm develop some TFX pipeline: ImportExampleGen, StatisticsGen, etc.

I have ~5gb TF Example data and 10 million data rows.

I configured Flink cluster with using almost 15gb ram, but StatisticsGen fails with OutOfMemory error.

My dataset is not production. Production datasets is more huge.

Error occures on GenerateStatistics[train]/RunStatsGenerators/GenerateSlicedStatisticsImpl/TopKUniquesStatsGenerator/CombineCountsAndWeights/GroupByKey -> [2]GenerateStatistics[train]/RunStatsGenerators/GenerateSlicedStatisticsImpl/TopKUniquesStatsGenerator/{CombineCountsAndWeights, Rearrange} -> ([2]GenerateStatistics[train]/RunStatsGenerators/GenerateSlicedStatisticsImpl/TopKUniquesStatsGenerator/{Uniques_Keys, Uniques_CountPerFeatureName} -> ToKeyedWorkItem, ToKeyedWorkItem) task

Detailed error from Flink logs:

Exception in thread "grpc-default-executor-4" java.lang.OutOfMemoryError: Java heap space
        at java.util.Arrays.copyOfRange(Arrays.java:3520)
        at org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString$ArraysByteArrayCopier.copyFrom(ByteString.java:126)
        at org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString.copyFrom(ByteString.java:362)
        at org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.ByteString.copyFrom(ByteString.java:372)
        at org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.CodedInputStream$StreamDecoder.readBytesSlowPath(CodedInputStream.java:2978)
        at org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.CodedInputStream$StreamDecoder.readBytes(CodedInputStream.java:2386)
        at org.apache.beam.model.fnexecution.v1.BeamFnApi$Elements$Data.<init>(BeamFnApi.java:29659)
        at org.apache.beam.model.fnexecution.v1.BeamFnApi$Elements$Data.<init>(BeamFnApi.java:29600)
        at org.apache.beam.model.fnexecution.v1.BeamFnApi$Elements$Data$1.parsePartialFrom(BeamFnApi.java:30539)
        at org.apache.beam.model.fnexecution.v1.BeamFnApi$Elements$Data$1.parsePartialFrom(BeamFnApi.java:30533)
        at org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.CodedInputStream$StreamDecoder.readMessage(CodedInputStream.java:2366)
        at org.apache.beam.model.fnexecution.v1.BeamFnApi$Elements.<init>(BeamFnApi.java:29463)
        at org.apache.beam.model.fnexecution.v1.BeamFnApi$Elements.<init>(BeamFnApi.java:29412)
        at org.apache.beam.model.fnexecution.v1.BeamFnApi$Elements$1.parsePartialFrom(BeamFnApi.java:32943)
        at org.apache.beam.model.fnexecution.v1.BeamFnApi$Elements$1.parsePartialFrom(BeamFnApi.java:32937)
        at org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.AbstractParser.parseFrom(AbstractParser.java:86)
        at org.apache.beam.vendor.grpc.v1p26p0.com.google.protobuf.AbstractParser.parseFrom(AbstractParser.java:48)
        at org.apache.beam.vendor.grpc.v1p26p0.io.grpc.protobuf.lite.ProtoLiteUtils$MessageMarshaller.parseFrom(ProtoLiteUtils.java:223)
        at org.apache.beam.vendor.grpc.v1p26p0.io.grpc.protobuf.lite.ProtoLiteUtils$MessageMarshaller.parse(ProtoLiteUtils.java:215)
        at org.apache.beam.vendor.grpc.v1p26p0.io.grpc.protobuf.lite.ProtoLiteUtils$MessageMarshaller.parse(ProtoLiteUtils.java:118)
        at org.apache.beam.vendor.grpc.v1p26p0.io.grpc.MethodDescriptor.parseRequest(MethodDescriptor.java:299)
        at org.apache.beam.vendor.grpc.v1p26p0.io.grpc.internal.ServerCallImpl$ServerStreamListenerImpl.messagesAvailableInternal(ServerCallImpl.java:309)
        at org.apache.beam.vendor.grpc.v1p26p0.io.grpc.internal.ServerCallImpl$ServerStreamListenerImpl.messagesAvailable(ServerCallImpl.java:292)
        at org.apache.beam.vendor.grpc.v1p26p0.io.grpc.internal.ServerImpl$JumpToApplicationThreadServerStreamListener$1MessagesAvailable.runInContext(ServerImpl.java:782)
        at org.apache.beam.vendor.grpc.v1p26p0.io.grpc.internal.ContextRunnable.run(ContextRunnable.java:37)
        at org.apache.beam.vendor.grpc.v1p26p0.io.grpc.internal.SerializingExecutor.run(SerializingExecutor.java:123)
        at java.util.concurrent.ThreadPoolExecutor.runWorker(ThreadPoolExecutor.java:1142)
        at java.util.concurrent.ThreadPoolExecutor$Worker.run(ThreadPoolExecutor.java:617)
        at java.lang.Thread.run(Thread.java:745)

How can i decrease memory usage?

Thanks!

---

Pipline options:

--runner=FlinkRunner
--flink_master=flink-cluster:8081
--streaming
--auto_balance_write_files_sharding_enabled
--faster_copy
--max_parallelism=5
--object_reuse
--environment_type=EXTERNAL
--environment_config=127.0.0.1:50000
--flink_submit_uber_jar
--parallelism=5
--sdk_worker_parallelism=5
--execution_mode_for_batch=BATCH_FORCED
--s3_access_key=valid-access-key
--s3_secret_access_key=valid-secret-key
--s3_endpoint_url=minio:9000"
--s3_verify=false

Flink memory configuration:

jobmanager.memory.process.size: 2048m
jobmanager.heap.size: 4096m
taskmanager.runtime.large-record-handler: true
taskmanager.runtime.hashjoin-bloom-filters: true
compiler.delimited-informat.max-line-samples: 100
taskmanager.memory.task.heap.size: 12gb
taskmanager.memory.managed.size: 4gb
taskmanager.memory.jvm-overhead.max: 4gb
taskmanager.memory.jvm-overhead.min: 256m
taskmanager.memory.jvm-overhead.fraction: 0.4
taskmanager.memory.task.off-heap.size: 4gb
taskmanager.memory.framework.heap.size: 512m
taskmanager.memory.framework.off-heap.size: 2gb
taskmanager.numberOfTaskSlots: 5
taskmanager.memory.network.fraction: 0.4
taskmanager.memory.network.min: 64mb
taskmanager.memory.network.max: 4gb

There is Flink task manager metrics from Grafana

As can you see, ImportExampleGen processed successfully. But, after mem idle, starts and fails StatisticsGen.

[ibzib]

You may want to try reducing sdk_worker_paralleism to 1. sdk_worker_parallelism sets the number of Beam SDK worker clients _per Flink worker node_, so setting it too high could easily overwhelm your workers.

[ferryvg]

@ibzib Yeah, i know. But, this increase processing time.

[ibzib]

@ibzib Yeah, i know. But, this increase processing time.

Not necessarily. There are two layers of parallelism here: Flink's (parallelism) and Beam's (sdk_worker_parallelism). To reiterate, I'm not suggesting you change parallelism, I'm suggesting you change sdk_worker_parallelism, which sets the number of Beam SDK worker clients _per Flink worker node_. That means the number of worker clients equals parallelism × sdk_worker_parallelism = 25. And more worker clients is not necessarily better; you only have a limited amount of CPU, so there is a point of diminishing returns somewhere.

And more generally, processing time is moot when the pipeline is failing altogether.

[arghyaganguly]

@ferryvg , Flink runner is currently not supported.Closing this issue.

[tensorflow-butler[bot]]

Are you satisfied with the resolution of your issue?
Yes
No

[ibzib]

@ferryvg , Flink runner is currently not supported.Closing this issue.

@arghyaganguly Can you please reopen this? I have worked on the Flink runner and I don't think ferryvg is trying to do anything unsupported here. There are instructions for running TFX on the Flink runner in the docs, so in that sense it should be considered supported. https://github.com/tensorflow/tfx/blob/master/tfx/examples/chicago_taxi_pipeline/README.md#chicago-taxi-flink-example

[brills]

@ferryvg Does your data contain large string/bytes features (e.g. images, videos or large blob of texts)?

On common reason that StatisticsGen could go OOM is that it will try to compute top-k for all string/bytes features that have large values. You can annotate the feature in the schema (StatsOptions.schema) with an image_domain to tell StatsticsGen to not generate top-k for a feature.

[ferryvg]

@brills
No, we does not have features with large data, but we have features with non-unique values (almost 60% non-unique of all), like mobile carrier code, country region, domain.
Currently we was refuse some built-in components like StatisticsGen and makes custom SchemaGen. But we have same issue with Evaluator in DataFlow

[brills]

@ibzib @ferryvg

Also to clarify what "support" means in a comment left by @arghyaganguly

Flink runner is currently not supported.Closing this issue.

As @ibzib correctly pointed out, the Flink runner is supported by Apache Beam, but TFX does not test its components on Flink runner and the team has no expertise on the Flink runner. TFX at this time can only provide best-effort support here (someone may be able to provide debugging hints like the comment I posted earlier, but it's difficult for us to try to reproduce and drill into your issue).

OTOH, Apache Beam may have committed support for the Flink runner.

[vaskozl]

schema (StatsOptions.schema) with an image_domain

Thank's I think I've encountered this and didn't know there was a solution!

TFX does not test its components on Flink runner and the team has no expertise

I myself have had very mixed success with Flink since release 0.21 onward, it's helped with larger datasets but it's its own can of worms (jobs hanging inexplicably sometimes). The thing is, when datasets start hitting a few gigabytes one invariably ends up hitting OOM with DirectRunner even if your dataset in theory should be able to easily fit in RAM. E.g. in this issue https://github.com/tensorflow/tfx/issues/1907 it's recommended to use Spark/Flink/Dataflow. I've also found https://github.com/tensorflow/transform/issues/143 to be a problem (atleast in 0.22 haven't confirmed in 0.27).

I think Spark is even less likely to work in general than Flink and seems to have basically no support from the beam folk. As far as I can see this makes Dataflow the only option for running datasets over a few GB large which causes people who are just starting out and not on GCP to have a hard time.

I think it would be great to test with larger datasets with DirectRunner and Flink at-least to document some expectation of what should work in production with a chunky dataset. Being able to run more reliably on large datasets without Dataflow also means one can use Jupyter notebooks more readily without crashing which I offers a lot of value.

TLDR: I would love if you guys tested the components with a huge dataset if you don't already.

[vaskozl]

@ferryvg Does your data contain large string/bytes features (e.g. images, videos or large blob of texts)?

On common reason that StatisticsGen could go OOM is that it will try to compute top-k for all string/bytes features that have large values. You can annotate the feature in the schema (StatsOptions.schema) with an image_domain to tell StatsticsGen to not generate top-k for a feature.

I'm struggling to figure out exactly how to set the stats_schema, does it need to be a copy of the examplegen schema but with the extra image_domain key set? An example would be amazing :)

Cheers for the help so far

[1025KB]

@ferryvg Does your data contain large string/bytes features (e.g. images, videos or large blob of texts)?
On common reason that StatisticsGen could go OOM is that it will try to compute top-k for all string/bytes features that have large values. You can annotate the feature in the schema (StatsOptions.schema) with an image_domain to tell StatsticsGen to not generate top-k for a feature.

I'm struggling to figure out exactly how to set the stats_schema, does it need to be a copy of the examplegen schema but with the extra image_domain key set? An example would be amazing :)

You need to copy out the schema generated from the schema gen, and then make the modification, and use an importernode (check this) to import it back to the pipeline (and from then on, stop using the SchemaGen)

Cheers for the help so far

[ibzib]

re @vaskozl

I myself have had very mixed success with Flink since release 0.21 onward, it's helped with larger datasets but it's its own can of worms (jobs hanging inexplicably sometimes). The thing is, when datasets start hitting a few gigabytes one invariably ends up hitting OOM with DirectRunner even if your dataset in theory should be able to easily fit in RAM. E.g. in this issue #1907 it's recommended to use Spark/Flink/Dataflow. I've also found tensorflow/transform#143 to be a problem (atleast in 0.22 haven't confirmed in 0.27).

Anecdotally, I've found more recent versions of Flink seem to have fewer stuckness problems, so you may want to try upgrading to to TFX 0.28 / Beam 2.28 / Flink 1.12.

I think Spark is even less likely to work in general than Flink and seems to have basically no support from the beam folk.

I'm curious if there are any specific bugs/missing features informing this?

TLDR: I would love if you guys tested the components with a huge dataset if you don't already.

I agree. Can you file this feature request as a separate issue so it gets more visibility?

[ConverJens]

@ibzib @vaskozl @ferryvg I recently tried to get the FlinkRunner working in kubernetes but ended up failing when jobs hanged in the beam worker. Would any of you mind sharing how you have deployed Flink w/ Beam and the pipeline options you are using?

[ferryvg]

@ConverJens
As finally, we abandoned dag-managers, as well as most of the components such as StatisticsGen, SchemaGen. Of the standard ones, only Trainer remained, which we modified for ourselves, and Pusher.
So, you can trashup StatisticsGen, if you not require it, and load schema manually

[vaskozl]

@ibzib @vaskozl @ferryvg I recently tried to get the FlinkRunner working in kubernetes but ended up failing when jobs hanged in the beam worker. Would any of you mind sharing how you have deployed Flink w/ Beam and the pipeline options you are using?

I also experienced consistent hanging with Flink on k8s with the worker pool image. I was able to get it working on docker with host network in 0.22 but haven’t got it working since 0.27 and k8s without hitting hanging. Thought I was alone.

[ferryvg]

@vaskozl @ConverJens you tried 10.x of Flink or newest? I was have success expiriens with 10.x

[ibzib]

@ConverJens @vaskozl can you file separate issues for your hanging issues? It'll help us investigate without cluttering up the thread chasing various unrelated bugs.

[ConverJens]

@ibzib I have filed a new issue here: https://github.com/tensorflow/tfx/issues/3435

@ferryvg I've only tested with Flink 1.10.3 since that is the newest one supported by beam version >=2.27.0.

@vaskozl I think I have a pretty similar experience, getting it to work once or twice locally but not in k8s. Maybe we can continue this discussion in the other issue I posted?

[ConverJens]

@ferryvg Are you using a Flink session cluster or another deploy setup?