Flink runner not splitting tasks when parallelism is turned on in BEAM python pipeline - apache-flink

Question

I have a beam pipeline written in python that when deployed to a flink runner doesn't make use of the parallelism correctly.
There is unbounded data coming in through a kafka connector and I want the data to be read when split in parallel.
My understanding is that it should split up the tasks but as shown in the image one parallelism is used and all the other 5 sub tasks finished instantly leaving the one running to do all the work.
The pipeline settings are:
options = PipelineOptions([
"--runner=PortableRunner",
"--sdk_worker_parallelism=3",
"--artifact_endpoint=localhost:8098",
"--job_endpoint=localhost:8099",
"--environment_type=EXTERNAL",
"--environment_config=localhost:50000",
"--checkpointing_interval=30000",
])
options._all_options['parallelism'] = 3
Is this a missing config on the Flink runner or something that can be configured in the BEAM pipeline?
The full pipeline:
import apache_beam as beam
from apache_beam.options.pipeline_options import PipelineOptions
options = PipelineOptions([
"--runner=PortableRunner",
"--sdk_worker_parallelism=3",
"--artifact_endpoint=localhost:8098",
"--job_endpoint=localhost:8099",
"--environment_type=EXTERNAL",
"--environment_config=localhost:50000",
"--checkpointing_interval=30000",
])
options._all_options['parallelism'] = 3
class CountProvider(beam.RestrictionProvider):
def __init__(self, initial_split_size=5):
self._initial_split_size = initial_split_size
self.OffsetRestrictionTracker = None
def imports(self):
if self.OffsetRestrictionTracker is not None: return
from apache_beam.io.restriction_trackers import OffsetRestrictionTracker, OffsetRange
self.OffsetRestrictionTracker = OffsetRestrictionTracker
self.OffsetRange = OffsetRange
def initial_restriction(self, element):
self.imports()
return self.OffsetRange(0, 10)
def create_tracker(self, restriction):
self.imports()
return self.OffsetRestrictionTracker(restriction)
def restriction_size(self, element, restriction):
return restriction.size()*100_000
def split(self, element, restriction):
self.imports()
if restriction.start + 1 >= restriction.stop:
yield self.OffsetRange(restriction.start, restriction.stop)
else:
last_val = restriction.start
for i in range(1, self._initial_split_size):
next_stop = i * (restriction.start + restriction.stop) // self._initial_split_size
yield self.OffsetRange(last_val, next_stop)
last_val = next_stop
yield self.OffsetRange(last_val, restriction.stop)
class CountFn(beam.DoFn):
def setup(self):
print("setup")
def process(self, element, tracker=beam.DoFn.RestrictionParam(CountProvider())):
res = tracker.current_restriction()
print(f"Current Restriction {res.start}, {res.stop}")
for i in range(res.start, res.stop):
if not tracker.try_claim(i):
return
for j in range(10_000):
yield i, j
def get_initial_restriction(self, filename):
return (0, 10)
def teardown(self):
print("Teardown")
p = beam.Pipeline(options=options)
out = (p | f'Create' >> beam.Create([tuple()])
| f'Gen Data' >> beam.ParDo(CountFn())
| beam.Map(print)
)
result = p.run()
result.wait_until_finish()