#

# Licensed to the Apache Software Foundation (ASF) under one or more

# contributor license agreements. See the NOTICE file distributed with

# this work for additional information regarding copyright ownership.

# The ASF licenses this file to You under the Apache License, Version 2.0

# (the "License"); you may not use this file except in compliance with

# the License. You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

#

"""A set of utilities to write pipelines for performance tests.

This module offers a way to create pipelines using synthetic sources and steps.

Exact shape of the pipeline and the behaviour of sources and steps can be

controlled through arguments. Please see function 'parse_args()' for more

details about the arguments.

Shape of the pipeline is primarily controlled through two arguments. Argument

'steps' can be used to define a list of steps as a JSON string. Argument

'barrier' describes how these steps are separated from each other. Argument

'barrier' can be use to build a pipeline as a series of steps or a tree of

steps with a fanin or a fanout of size 2.

Other arguments describe what gets generated by synthetic sources that produce

data for the pipeline.

"""

# pytype: skip-file

import argparse

import json

import logging

import math

import os

import time

from random import Random

from typing import Optional

from typing import Tuple

import apache_beam as beam

from apache_beam import pvalue

from apache_beam import typehints

from apache_beam.io import WriteToText

from apache_beam.io import iobase

from apache_beam.io import range_trackers

from apache_beam.io import restriction_trackers

from apache_beam.io.restriction_trackers import OffsetRange

from apache_beam.io.restriction_trackers import OffsetRestrictionTracker

from apache_beam.options.pipeline_options import PipelineOptions

from apache_beam.options.pipeline_options import SetupOptions

from apache_beam.testing.test_pipeline import TestPipeline

from apache_beam.transforms import userstate

from apache_beam.transforms.core import RestrictionProvider

try:

import numpy as np

except ImportError:

np = None

try:

from .fast_test_utils import LCGenerator

except ImportError:

# Cythonized lib not available

from .test_utils import LCGenerator

class _Random(Random):

"""A subclass of `random.Random` from the Python Standard Library that

provides a method returning random bytes of arbitrary length.

"""

# `numpy.random.RandomState` does not provide `random()` method, we keep this

# for compatibility reasons.

random_sample = Random.random

def get_generator(seed: Optional[int] = None, algorithm: Optional[str] = None):

if algorithm is None or algorithm == 'builtin':

return _Random(seed)

elif algorithm == 'lcg':

generator = LCGenerator()

if seed is not None:

generator.seed(seed)

return generator

else:

raise ValueError(

'Unknown algorithm %s. Supported algorithms are "builtin" or "lcg".',

algorithm)

def parse_byte_size(s):

suffixes = 'BKMGTP'

if s[-1] in suffixes:

return int(float(s[:-1]) * 1024**suffixes.index(s[-1]))

return int(s)

def div_round_up(a, b):

"""Return ceil(a/b)."""

return int(math.ceil(float(a) / b))

def rotate_key(element):

"""Returns a new key-value pair of the same size but with a different key."""

(key, value) = element

return key[-1:] + key[:-1], value

def initial_splitting_zipf(

start_position,

stop_position,

desired_num_bundles,

distribution_parameter,

num_total_records=None):

"""Split the given range (defined by start_position, stop_position) into

desired_num_bundles using zipf with the given distribution_parameter.

"""

if not num_total_records:

num_total_records = stop_position - start_position

samples = np.random.zipf(distribution_parameter, desired_num_bundles)

total = sum(samples)

relative_bundle_sizes = [(float(sample) / total) for sample in samples]

bundle_ranges = []

start = start_position

index = 0

while start < stop_position:

if index == desired_num_bundles - 1:

bundle_ranges.append((start, stop_position))

break

stop = start + int(num_total_records * relative_bundle_sizes[index])

bundle_ranges.append((start, stop))

start = stop

index += 1

return bundle_ranges

class SyntheticStep(beam.DoFn):

"""A DoFn of which behavior can be controlled through prespecified parameters.

"""

def __init__(

self,

per_element_delay_sec=0,

per_bundle_delay_sec=0,

output_records_per_input_record=1,

output_filter_ratio=0):

if per_element_delay_sec and per_element_delay_sec < 1e-3:

raise ValueError(

'Per element sleep time must be at least 1e-3. '

'Received: %r',

per_element_delay_sec)

self._per_element_delay_sec = per_element_delay_sec

self._per_bundle_delay_sec = per_bundle_delay_sec

self._output_records_per_input_record = output_records_per_input_record

self._output_filter_ratio = output_filter_ratio

def start_bundle(self):

self._start_time = time.time()

def finish_bundle(self):

# The target is for the enclosing stage to take as close to as possible

# the given number of seconds, so we only sleep enough to make up for

# overheads not incurred elsewhere.

to_sleep = self._per_bundle_delay_sec - (time.time() - self._start_time)

# Ignoring sub-millisecond sleep times.

if to_sleep >= 1e-3:

time.sleep(to_sleep)

def process(self, element):

if self._per_element_delay_sec >= 1e-3:

time.sleep(self._per_element_delay_sec)

filter_element = False

if self._output_filter_ratio > 0:

if np.random.random() < self._output_filter_ratio:

filter_element = True

if not filter_element:

for _ in range(self._output_records_per_input_record):

yield element

class NonLiquidShardingOffsetRangeTracker(OffsetRestrictionTracker):

"""An OffsetRangeTracker that doesn't allow splitting. """

def try_split(self, split_offset):

pass # Don't split.

def checkpoint(self):

pass # Don't split.

class SyntheticSDFStepRestrictionProvider(RestrictionProvider):

"""A `RestrictionProvider` for SyntheticSDFStep.

An initial_restriction and split that operate on num_records and ignores

source description (element). Splits into initial_splitting_num_bundles.

Returns size_estimate_override as restriction size, if set. Otherwise uses

element size.

If initial_splitting_uneven_chunks, produces uneven chunks.

"""

def __init__(

self,

num_records,

initial_splitting_num_bundles,

initial_splitting_uneven_chunks,

disable_liquid_sharding,

size_estimate_override):

self._num_records = num_records

self._initial_splitting_num_bundles = initial_splitting_num_bundles

self._initial_splitting_uneven_chunks = initial_splitting_uneven_chunks

self._disable_liquid_sharding = disable_liquid_sharding

self._size_estimate_override = size_estimate_override

def initial_restriction(self, element):

return OffsetRange(0, self._num_records)

def create_tracker(self, restriction):

if self._disable_liquid_sharding:

return NonLiquidShardingOffsetRangeTracker(restriction)

else:

return OffsetRestrictionTracker(restriction)

def split(self, element, restriction):

elems = restriction.size()

if (self._initial_splitting_uneven_chunks and

self._initial_splitting_num_bundles > 1 and elems > 1):

bundle_ranges = initial_splitting_zipf(

restriction.start,

restriction.stop,

self._initial_splitting_num_bundles,

3.0)

for start, stop in bundle_ranges:

yield OffsetRange(start, stop)

else:

offsets_per_split = max(1, (elems // self._initial_splitting_num_bundles))

for split in restriction.split(offsets_per_split, offsets_per_split // 2):

yield split

def restriction_size(self, element, restriction):

if self._size_estimate_override is not None:

return self._size_estimate_override

element_size = len(element) if isinstance(element, str) else 1

return restriction.size() * element_size

def get_synthetic_sdf_step(

per_element_delay_sec=0,

per_bundle_delay_sec=0,

output_records_per_input_record=1,

output_filter_ratio=0,

initial_splitting_num_bundles=8,

initial_splitting_uneven_chunks=False,

disable_liquid_sharding=False,

size_estimate_override=None,

):

"""A function which returns a SyntheticSDFStep with given parameters. """

class SyntheticSDFStep(beam.DoFn):

"""A SplittableDoFn of which behavior can be controlled through prespecified

parameters.

"""

def __init__(

self,

per_element_delay_sec_arg,

per_bundle_delay_sec_arg,

output_filter_ratio_arg,

output_records_per_input_record_arg):

if per_element_delay_sec_arg:

per_element_delay_sec_arg = (

per_element_delay_sec_arg // output_records_per_input_record_arg)

if per_element_delay_sec_arg < 1e-3:

raise ValueError(

'Per element sleep time must be at least 1e-3 after being '

'divided among output elements.')

self._per_element_delay_sec = per_element_delay_sec_arg

self._per_bundle_delay_sec = per_bundle_delay_sec_arg

self._output_filter_ratio = output_filter_ratio_arg

def start_bundle(self):

self._start_time = time.time()

def finish_bundle(self):

# The target is for the enclosing stage to take as close to as possible

# the given number of seconds, so we only sleep enough to make up for

# overheads not incurred elsewhere.

to_sleep = self._per_bundle_delay_sec - (time.time() - self._start_time)

# Ignoring sub-millisecond sleep times.

if to_sleep >= 1e-3:

time.sleep(to_sleep)

def process(

self,

element,

restriction_tracker=beam.DoFn.RestrictionParam(

SyntheticSDFStepRestrictionProvider(

output_records_per_input_record,

initial_splitting_num_bundles,

initial_splitting_uneven_chunks,

disable_liquid_sharding,

size_estimate_override))):

filter_element = False

if self._output_filter_ratio > 0:

if np.random.random() < self._output_filter_ratio:

filter_element = True

current_restriction = restriction_tracker.current_restriction()

for cur in range(current_restriction.start, current_restriction.stop):

if not restriction_tracker.try_claim(cur):

return

if self._per_element_delay_sec:

time.sleep(self._per_element_delay_sec)

if not filter_element:

yield element

cur += 1

return SyntheticSDFStep(

per_element_delay_sec,

per_bundle_delay_sec,

output_filter_ratio,

output_records_per_input_record)

class SyntheticSource(iobase.BoundedSource):

"""A custom source of a specified size.

"""

def __init__(self, input_spec):

"""Initiates a synthetic source.

Args:

input_spec: Input specification of the source. See corresponding option in

function 'parse_args()' below for more details.

Raises:

ValueError: if input parameters are invalid.

"""

def maybe_parse_byte_size(s):

return parse_byte_size(s) if isinstance(s, str) else int(s)

self._num_records = input_spec['numRecords']

self._key_size = maybe_parse_byte_size(input_spec.get('keySizeBytes', 1))

self._hot_key_fraction = input_spec.get('hotKeyFraction', 0)

self._num_hot_keys = input_spec.get('numHotKeys', 0)

self._value_size = maybe_parse_byte_size(

input_spec.get('valueSizeBytes', 1))

self._total_size = self.element_size * self._num_records

self._initial_splitting = (

input_spec['bundleSizeDistribution']['type']

if 'bundleSizeDistribution' in input_spec else 'const')

if self._initial_splitting != 'const' and self._initial_splitting != 'zipf':

raise ValueError(

'Only const and zipf distributions are supported for determining '

'sizes of bundles produced by initial splitting. Received: %s',

self._initial_splitting)

self._initial_splitting_num_bundles = (

input_spec['forceNumInitialBundles']

if 'forceNumInitialBundles' in input_spec else 0)

if self._initial_splitting == 'zipf':

self._initial_splitting_distribution_parameter = (

input_spec['bundleSizeDistribution']['param'])

if self._initial_splitting_distribution_parameter < 1:

raise ValueError(

'Parameter for a Zipf distribution must be larger than 1. '

'Received %r.',

self._initial_splitting_distribution_parameter)

else:

self._initial_splitting_distribution_parameter = 0

self._dynamic_splitting = (

'none' if (

'splitPointFrequencyRecords' in input_spec and

input_spec['splitPointFrequencyRecords'] == 0) else 'perfect')

if 'delayDistribution' in input_spec:

if input_spec['delayDistribution']['type'] != 'const':

raise ValueError(

'SyntheticSource currently only supports delay '

'distributions of type \'const\'. Received %s.',

input_spec['delayDistribution']['type'])

self._sleep_per_input_record_sec = (

float(input_spec['delayDistribution']['const']) / 1000)

if (self._sleep_per_input_record_sec and

self._sleep_per_input_record_sec < 1e-3):

raise ValueError(

'Sleep time per input record must be at least 1e-3.'

' Received: %r',

self._sleep_per_input_record_sec)

else:

self._sleep_per_input_record_sec = 0

# algorithm of the generator

self.gen_algo = input_spec.get('algorithm', None)

if self.gen_algo not in (None, 'builtin', 'lcg'):

raise ValueError(

'Unknown algorithm for input_spec: %s. Supported '

'algorithms are "builtin" and "lcg".',

self.gen_algo)

@property

def element_size(self):

return self._key_size + self._value_size

def estimate_size(self):

return self._total_size

def split(self, desired_bundle_size, start_position=0, stop_position=None):

# Performs initial splitting of SyntheticSource.

#

# Exact sizes and distribution of initial splits generated here depends on

# the input specification of the SyntheticSource.

if stop_position is None:

stop_position = self._num_records

if self._initial_splitting == 'zipf':

desired_num_bundles = self._initial_splitting_num_bundles or math.ceil(

float(self.estimate_size()) / desired_bundle_size)

bundle_ranges = initial_splitting_zipf(

start_position,

stop_position,

desired_num_bundles,

self._initial_splitting_distribution_parameter,

self._num_records)

else:

if self._initial_splitting_num_bundles:

bundle_size_in_elements = max(

1, int(self._num_records / self._initial_splitting_num_bundles))

else:

bundle_size_in_elements = (

max(

div_round_up(desired_bundle_size, self.element_size),

int(math.floor(math.sqrt(self._num_records)))))

bundle_ranges = []

for start in range(start_position, stop_position,

bundle_size_in_elements):

stop = min(start + bundle_size_in_elements, stop_position)

bundle_ranges.append((start, stop))

for start, stop in bundle_ranges:

yield iobase.SourceBundle(stop - start, self, start, stop)

def get_range_tracker(self, start_position, stop_position):

if start_position is None:

start_position = 0

if stop_position is None:

stop_position = self._num_records

tracker = range_trackers.OffsetRangeTracker(start_position, stop_position)

if self._dynamic_splitting == 'none':

tracker = range_trackers.UnsplittableRangeTracker(tracker)

return tracker

def _gen_kv_pair(self, generator, index):

generator.seed(index)

rand = generator.random_sample()

# Determines whether to generate hot key or not.

if rand < self._hot_key_fraction:

# Generate hot key.

# An integer is randomly selected from the range [0, numHotKeys-1]

# with equal probability.

generator_hot = get_generator(

seed=index % self._num_hot_keys, algorithm=self.gen_algo)

bytes_ = generator_hot.randbytes(self._key_size), generator.randbytes(

self._value_size)

else:

bytes_ = generator.randbytes(self.element_size)

bytes_ = bytes_[:self._key_size], bytes_[self._key_size:]

return bytes_

def read(self, range_tracker):

index = range_tracker.start_position()

generator = get_generator(algorithm=self.gen_algo)

while range_tracker.try_claim(index):

time.sleep(self._sleep_per_input_record_sec)

yield self._gen_kv_pair(generator, index)

index += 1

def default_output_coder(self):

return beam.coders.TupleCoder(

[beam.coders.BytesCoder(), beam.coders.BytesCoder()])

class SyntheticSDFSourceRestrictionProvider(RestrictionProvider):

"""A `RestrictionProvider` for SyntheticSDFAsSource.

In initial_restriction(element) and split(element), element means source

description.

A typical element is like:

{

'key_size': 1,

'value_size': 1,

'initial_splitting_num_bundles': 8,

'initial_splitting_desired_bundle_size': 2,

'sleep_per_input_record_sec': 0,

'initial_splitting' : 'const'

}

"""

def initial_restriction(self, element):

return OffsetRange(0, element['num_records'])

def create_tracker(self, restriction):

return restriction_trackers.OffsetRestrictionTracker(restriction)

def split(self, element, restriction):

bundle_ranges = []

start_position = restriction.start

stop_position = restriction.stop

element_size = element['key_size'] + element['value_size']

estimate_size = element_size * element['num_records']

if element['initial_splitting'] == 'zipf':

desired_num_bundles = (

element['initial_splitting_num_bundles'] or div_round_up(

estimate_size, element['initial_splitting_desired_bundle_size']))

samples = np.random.zipf(

element['initial_splitting_distribution_parameter'],

desired_num_bundles)

total = sum(samples)

relative_bundle_sizes = [(float(sample) / total) for sample in samples]

start = start_position

index = 0

while start < stop_position:

if index == desired_num_bundles - 1:

bundle_ranges.append(OffsetRange(start, stop_position))

break

stop = start + int(

element['num_records'] * relative_bundle_sizes[index])

bundle_ranges.append(OffsetRange(start, stop))

start = stop

index += 1

else:

if element['initial_splitting_num_bundles']:

bundle_size_in_elements = max(

1,

int(

element['num_records'] /

element['initial_splitting_num_bundles']))

else:

bundle_size_in_elements = (

max(

div_round_up(

element['initial_splitting_desired_bundle_size'],

element_size),

int(math.floor(math.sqrt(element['num_records'])))))

for start in range(start_position, stop_position,

bundle_size_in_elements):

stop = min(start + bundle_size_in_elements, stop_position)

bundle_ranges.append(OffsetRange(start, stop))

return bundle_ranges

def restriction_size(self, element, restriction):

return (element['key_size'] + element['value_size']) * restriction.size()

class SyntheticSDFAsSource(beam.DoFn):

"""A SDF that generates records like a source.

This SDF accepts a PCollection of record-based source description.

A typical description is like:

{

'key_size': 1,

'value_size': 1,

'initial_splitting_num_bundles': 8,

'initial_splitting_desired_bundle_size': 2,

'sleep_per_input_record_sec': 0,

'initial_splitting' : 'const'

}

A simple pipeline taking this SDF as a source is like:

p

| beam.Create([description1, description2,...])

| beam.ParDo(SyntheticSDFAsSource())

NOTE:

The SDF.process() will have different param content between defining a DoFn

and runtime.

When defining an SDF.process, the restriction_tracker should be a

`RestrictionProvider`.

During runtime, the DoFnRunner.process_with_sized_restriction() will feed

a 'RestrictionTracker' based on a restriction to SDF.process().

"""

def process(

self,

element,

restriction_tracker=beam.DoFn.RestrictionParam(

SyntheticSDFSourceRestrictionProvider())):

cur = restriction_tracker.current_restriction().start

while restriction_tracker.try_claim(cur):

r = get_generator(algorithm=element.get('algorithm', None), seed=cur)

time.sleep(element['sleep_per_input_record_sec'])

yield (

r.randbytes(element['key_size']), r.randbytes(element['value_size']))

cur += 1

class ShuffleBarrier(beam.PTransform):

def expand(self, pc):

return (

pc

| beam.Map(rotate_key)

| beam.GroupByKey()

| 'Ungroup' >> beam.FlatMap(lambda elm: [(elm[0], v) for v in elm[1]]))

class SideInputBarrier(beam.PTransform):

def expand(self, pc):

return (

pc

| beam.Map(rotate_key)

| beam.Map(

lambda elem, ignored: elem,

beam.pvalue.AsIter(pc | beam.FlatMap(lambda elem: None))))

def merge_using_gbk(name, pc1, pc2):

"""Merges two given PCollections using a CoGroupByKey."""

pc1_with_key = pc1 | (name + 'AttachKey1') >> beam.Map(lambda x: (x, x))

pc2_with_key = pc2 | (name + 'AttachKey2') >> beam.Map(lambda x: (x, x))

grouped = ({

'pc1': pc1_with_key, 'pc2': pc2_with_key

} | (name + 'Group') >> beam.CoGroupByKey())

return (

grouped | (name + 'DeDup') >> beam.Map(lambda elm: elm[0])

) # Ignoring values

def merge_using_side_input(name, pc1, pc2):

"""Merges two given PCollections using side inputs."""

def join_fn(val, _): # Ignoring side input

return val

return pc1 | name >> beam.core.Map(join_fn, beam.pvalue.AsIter(pc2))

def expand_using_gbk(name, pc):

"""Expands a given PCollection into two copies using GroupByKey."""

ret = []

ret.append((pc | ('%s.a' % name) >> ShuffleBarrier()))

ret.append((pc | ('%s.b' % name) >> ShuffleBarrier()))

return ret

def expand_using_second_output(name, pc):

"""Expands a given PCollection into two copies using side outputs."""

class ExpandFn(beam.DoFn):

def process(self, element):

yield beam.pvalue.TaggedOutput('second_out', element)

yield element

pc1, pc2 = (pc | name >> beam.ParDo(

ExpandFn()).with_outputs('second_out', main='main_out'))

return [pc1, pc2]

def _parse_steps(json_str):

"""Converts the JSON step description into Python objects.

See property 'steps' for more details about the JSON step description.

Args:

json_str: a JSON string that describes the steps.

Returns:

Information about steps as a list of dictionaries. Each dictionary may have

following properties.

(1) per_element_delay - amount of delay for each element in seconds.

(2) per_bundle_delay - minimum amount of delay for a given step in seconds.

(3) output_records_per_input_record - number of output elements generated

for each input element to a step.

(4) output_filter_ratio - the probability at which a step may filter out a

given element by not producing any output for that element.

(5) splittable - if the step should be splittable.

(6) initial_splitting_num_bundles - number of bundles initial split if step

is splittable.

(7) initial_splitting_uneven_chunks - if the bundles should be

unevenly-sized

(8) disable_liquid_sharding - if liquid sharding should be disabled

(9) size_estimate_override - the size estimate or None to use default

"""

all_steps = []

json_data = json.loads(json_str)

for val in json_data:

steps = {}

steps['per_element_delay'] = ((float(val['per_element_delay_msec']) / 1000)

if 'per_element_delay_msec' in val else 0)

steps['per_bundle_delay'] = (

float(val['per_bundle_delay_sec'])

if 'per_bundle_delay_sec' in val else 0)

steps['output_records_per_input_record'] = (

int(val['output_records_per_input_record'])

if 'output_records_per_input_record' in val else 1)

steps['output_filter_ratio'] = (

float(val['output_filter_ratio'])

if 'output_filter_ratio' in val else 0)

steps['splittable'] = (

bool(val['splittable']) if 'splittable' in val else False)

steps['initial_splitting_num_bundles'] = (

int(val['initial_splitting_num_bundles'])

if 'initial_splitting_num_bundles' in val else 8)

steps['initial_splitting_uneven_chunks'] = (

bool(val['initial_splitting_uneven_chunks'])

if 'initial_splitting_uneven_chunks' in val else False)

steps['disable_liquid_sharding'] = (

bool(val['disable_liquid_sharding'])

if 'disable_liquid_sharding' in val else False)

steps['size_estimate_override'] = (

int(val['size_estimate_override'])

if 'size_estimate_override' in val else None)

all_steps.append(steps)

return all_steps

def parse_args(args):

"""Parses a given set of arguments.

Args:

args: set of arguments to be passed.

Returns:

a tuple where first item gives the set of arguments defined and parsed

within this method and second item gives the set of unknown arguments.

"""

parser = argparse.ArgumentParser()

parser.add_argument(

'--steps',

dest='steps',

type=_parse_steps,

help='A JSON string that gives a list where each entry of the list is '

'configuration information for a step. Configuration for each step '

'consists of '

'(1) A float "per_bundle_delay_sec" (in seconds). Defaults to 0.'

'(2) A float "per_element_delay_msec" (in milli seconds). '

' Defaults to 0.'

'(3) An integer "output_records_per_input_record". Defaults to 1.'

'(4) A float "output_filter_ratio" in the range [0, 1] . '

' Defaults to 0.'

'(5) A bool "splittable" that defaults to false.'

'(6) An integer "initial_splitting_num_bundles". Defaults to 8.')

parser.add_argument(

'--input',

dest='input',

type=json.loads,

help='A JSON string that describes the properties of the SyntheticSource '

'used by the pipeline. Configuration is similar to Java '

'SyntheticBoundedInput.'

'Currently supports following properties. '

'(1) An integer "numRecords". '

'(2) An integer "keySize". '

'(3) An integer "valueSize". '

'(4) A tuple "bundleSizeDistribution" with following values. '

' A string "type". Allowed values are "const" and "zipf". '

' An float "param". Only used if "type"=="zipf". Must be '

' larger than 1. '

'(5) An integer "forceNumInitialBundles". '

'(6) An integer "splitPointFrequencyRecords". '

'(7) A tuple "delayDistribution" with following values. '

' A string "type". Only allowed value is "const". '

' An integer "const". '

'(8) A string "algorithm". Allowed values are "builtin" for Python '

' builtin random generator, and "lcg" for the linear congruential '

' generator equivalent to Java (java.util.Random).')

parser.add_argument(

'--barrier',

dest='barrier',

default='shuffle',

choices=[

'shuffle',

'side-input',

'expand-gbk',

'expand-second-output',

'merge-gbk',

'merge-side-input'

],

help='Whether to use shuffle as the barrier '

'(as opposed to side inputs).')

parser.add_argument(

'--output',

dest='output',

default='',

help='Destination to write output.')

return parser.parse_known_args(args)

def run(argv=None, save_main_session=True):

"""Runs the workflow."""

known_args, pipeline_args = parse_args(argv)

pipeline_options = PipelineOptions(pipeline_args)

pipeline_options.view_as(SetupOptions).save_main_session = save_main_session

input_info = known_args.input

with TestPipeline(options=pipeline_options) as p:

source = SyntheticSource(input_info)

# pylint: disable=expression-not-assigned

barrier = known_args.barrier

pc_list = []

num_roots = 2**(len(known_args.steps) - 1) if (

barrier == 'merge-gbk' or barrier == 'merge-side-input') else 1

for read_no in range(num_roots):

pc_list.append((p | ('Read %d' % read_no) >> beam.io.Read(source)))

for step_no, steps in enumerate(known_args.steps):

if step_no != 0:

new_pc_list = []

for pc_no, pc in enumerate(pc_list):

if barrier == 'shuffle':

new_pc_list.append(

(pc | ('shuffle %d.%d' % (step_no, pc_no)) >> ShuffleBarrier()))

elif barrier == 'side-input':

new_pc_list.append((

pc | ('side-input %d.%d' %

(step_no, pc_no)) >> SideInputBarrier()))

elif barrier == 'expand-gbk':

new_pc_list.extend(

expand_using_gbk(('expand-gbk %d.%d' % (step_no, pc_no)), pc))

elif barrier == 'expand-second-output':

new_pc_list.extend(

expand_using_second_output(

('expand-second-output %d.%d' % (step_no, pc_no)), pc))

elif barrier == 'merge-gbk':

if pc_no % 2 == 0:

new_pc_list.append(

merge_using_gbk(('merge-gbk %d.%d' % (step_no, pc_no)),

pc,

pc_list[pc_no + 1]))

else:

continue

elif barrier == 'merge-side-input':

if pc_no % 2 == 0:

new_pc_list.append(

merge_using_side_input(

('merge-side-input %d.%d' % (step_no, pc_no)),

pc,

pc_list[pc_no + 1]))

else:

continue

pc_list = new_pc_list

new_pc_list = []

for pc_no, pc in enumerate(pc_list):

if steps['splittable']:

step = get_synthetic_sdf_step(

per_element_delay_sec=steps['per_element_delay'],

per_bundle_delay_sec=steps['per_bundle_delay'],

output_records_per_input_record=steps[

'output_records_per_input_record'],

output_filter_ratio=steps['output_filter_ratio'],

initial_splitting_num_bundles=steps[

'initial_splitting_num_bundles'],

initial_splitting_uneven_chunks=steps[

'initial_splitting_uneven_chunks'],

disable_liquid_sharding=steps['disable_liquid_sharding'],

size_estimate_override=steps['size_estimate_override'])

else:

step = SyntheticStep(

per_element_delay_sec=steps['per_element_delay'],

per_bundle_delay_sec=steps['per_bundle_delay'],

output_records_per_input_record=steps[

'output_records_per_input_record'],

output_filter_ratio=steps['output_filter_ratio'])

new_pc = pc | 'SyntheticStep %d.%d' % (step_no,

pc_no) >> beam.ParDo(step)

new_pc_list.append(new_pc)

pc_list = new_pc_list

if known_args.output:

# If an output location is provided we format and write output.

if len(pc_list) == 1:

(

pc_list[0]

| 'FormatOutput' >> beam.Map(lambda elm: (elm[0] + elm[1]))

| 'WriteOutput' >> WriteToText(known_args.output))

logging.info('Pipeline run completed.')

if __name__ == '__main__':

logging.getLogger().setLevel(logging.INFO)

run()

class StatefulLoadGenerator(beam.PTransform):

"""A PTransform for generating random data using Timers API."""

def __init__(self, input_options, num_keys=100):

self.num_records = input_options['num_records']

self.key_size = input_options['key_size']

self.value_size = input_options['value_size']

self.num_keys = num_keys

@typehints.with_output_types(Tuple[bytes, bytes])

class GenerateKeys(beam.DoFn):

def __init__(self, num_keys, key_size):

self.num_keys = num_keys

self.key_size = key_size

def process(self, impulse):

for _ in range(self.num_keys):

key = os.urandom(self.key_size)

yield key, b''

class GenerateLoad(beam.DoFn):

state_spec = userstate.CombiningValueStateSpec(

'bundles_remaining', combine_fn=sum)

timer_spec = userstate.TimerSpec('timer', userstate.TimeDomain.WATERMARK)

def __init__(self, num_records_per_key, value_size, bundle_size=1000):

self.num_records_per_key = num_records_per_key

self.payload = os.urandom(value_size)

self.bundle_size = bundle_size

self.time_fn = time.time

def process(

self,

_element,

records_remaining=beam.DoFn.StateParam(state_spec),

timer=beam.DoFn.TimerParam(timer_spec)):

records_remaining.add(self.num_records_per_key)

timer.set(0)

@userstate.on_timer(timer_spec)

def process_timer(

self,

key=beam.DoFn.KeyParam,

records_remaining=beam.DoFn.StateParam(state_spec),

timer=beam.DoFn.TimerParam(timer_spec)):

cur_bundle_size = min(self.bundle_size, records_remaining.read())

for _ in range(cur_bundle_size):

records_remaining.add(-1)

yield key, self.payload

if records_remaining.read() > 0:

timer.set(0)

def expand(self, pbegin):

assert isinstance(pbegin, pvalue.PBegin), (

'Input to transform must be a PBegin but found %s' % pbegin)

return (

pbegin

| 'Impulse' >> beam.Impulse()

| 'GenerateKeys' >> beam.ParDo(

StatefulLoadGenerator.GenerateKeys(self.num_keys, self.key_size))

| 'GenerateLoad' >> beam.ParDo(

StatefulLoadGenerator.GenerateLoad(

self.num_records // self.num_keys, self.value_size)))