def test_user_pipeline_intact_after_deducing_pipeline_fragment(self, cell):
with cell: # Cell 1
p = beam.Pipeline(ir.InteractiveRunner())
# Watch the pipeline `p` immediately without calling locals().
ib.watch({'p': p})
with cell: # Cell 2
# pylint: disable=range-builtin-not-iterating
init = p | 'Init' >> beam.Create(range(10))
with cell: # Cell 3
square = init | 'Square' >> beam.Map(lambda x: x * x)
with cell: # Cell 4
cube = init | 'Cube' >> beam.Map(lambda x: x**3)
# Watch every PCollection has been defined so far in local scope without
# calling locals().
ib.watch({'init': init, 'square': square, 'cube': cube})
user_pipeline_proto_before_deducing_fragment = p.to_runner_api(
return_context=False, use_fake_coders=True)
_ = pf.PipelineFragment([square]).deduce_fragment()
user_pipeline_proto_after_deducing_fragment = p.to_runner_api(
return_context=False, use_fake_coders=True)
assert_pipeline_proto_equal(
self,
user_pipeline_proto_before_deducing_fragment,
user_pipeline_proto_after_deducing_fragment)
def test_instrument_example_unbounded_pipeline_to_read_cache(self):
"""Tests that the instrumenter works for a single unbounded source.
"""
# Create the pipeline that will be instrumented.
p_original = beam.Pipeline(interactive_runner.InteractiveRunner())
ie.current_env().set_cache_manager(StreamingCache(cache_dir=None),
p_original)
source_1 = p_original | 'source1' >> beam.io.ReadFromPubSub(
subscription='projects/fake-project/subscriptions/fake_sub')
# pylint: disable=possibly-unused-variable
pcoll_1 = source_1 | 'square1' >> beam.Map(lambda x: x * x)
# Mock as if cacheable PCollections are cached.
ib.watch(locals())
for name, pcoll in locals().items():
if not isinstance(pcoll, beam.pvalue.PCollection):
continue
cache_key = self.cache_key_of(name, pcoll)
self._mock_write_cache(p_original, [b''], cache_key)
# Instrument the original pipeline to create the pipeline the user will see.
instrumenter = instr.build_pipeline_instrument(p_original)
actual_pipeline = beam.Pipeline.from_runner_api(
proto=instrumenter.instrumented_pipeline_proto(),
runner=interactive_runner.InteractiveRunner(),
options=None)
# Now, build the expected pipeline which replaces the unbounded source with
# a TestStream.
source_1_cache_key = self.cache_key_of('source_1', source_1)
p_expected = beam.Pipeline()
test_stream = (p_expected
| TestStream(output_tags=[source_1_cache_key]))
# pylint: disable=expression-not-assigned
test_stream[source_1_cache_key] | 'square1' >> beam.Map(
lambda x: x * x)
# Test that the TestStream is outputting to the correct PCollection.
class TestStreamVisitor(PipelineVisitor):
def __init__(self):
self.output_tags = set()
def enter_composite_transform(self, transform_node):
self.visit_transform(transform_node)
def visit_transform(self, transform_node):
transform = transform_node.transform
if isinstance(transform, TestStream):
self.output_tags = transform.output_tags
v = TestStreamVisitor()
actual_pipeline.visit(v)
expected_output_tags = set([source_1_cache_key])
actual_output_tags = v.output_tags
self.assertSetEqual(expected_output_tags, actual_output_tags)
# Test that the pipeline is as expected.
assert_pipeline_proto_equal(self, p_expected.to_runner_api(),
instrumenter.instrumented_pipeline_proto())
def test_background_caching_pipeline_proto(self):
p = beam.Pipeline(interactive_runner.InteractiveRunner())
# Test that the two ReadFromPubSub are correctly cut out.
a = p | 'ReadUnboundedSourceA' >> beam.io.ReadFromPubSub(
subscription='projects/fake-project/subscriptions/fake_sub')
b = p | 'ReadUnboundedSourceB' >> beam.io.ReadFromPubSub(
subscription='projects/fake-project/subscriptions/fake_sub')
# Add some extra PTransform afterwards to make sure that only the unbounded
# sources remain.
c = (a, b) | beam.CoGroupByKey()
_ = c | beam.Map(lambda x: x)
ib.watch(locals())
instrumenter = instr.pin(p)
actual_pipeline = instrumenter.background_caching_pipeline_proto()
# Now recreate the expected pipeline, which should only have the unbounded
# sources.
p = beam.Pipeline(interactive_runner.InteractiveRunner())
a = p | 'ReadUnboundedSourceA' >> beam.io.ReadFromPubSub(
subscription='projects/fake-project/subscriptions/fake_sub')
_ = a | 'a' >> cache.WriteCache(ie.current_env().cache_manager(), '')
b = p | 'ReadUnboundedSourceB' >> beam.io.ReadFromPubSub(
subscription='projects/fake-project/subscriptions/fake_sub')
_ = b | 'b' >> cache.WriteCache(ie.current_env().cache_manager(), '')
expected_pipeline = p.to_runner_api(return_context=False,
use_fake_coders=True)
assert_pipeline_proto_equal(self, expected_pipeline, actual_pipeline)
def test_read_cache(self, mocked_get_cache_manager):
p = beam.Pipeline()
pcoll = p | beam.Create([1, 2, 3])
consumer_transform = beam.Map(lambda x: x * x)
_ = pcoll | consumer_transform
ib.watch(locals())
# Create the cache in memory.
cache_manager = InMemoryCache()
mocked_get_cache_manager.return_value = cache_manager
aug_p = ap.AugmentedPipeline(p)
key = repr(aug_p._cacheables[pcoll].to_key())
cache_manager.write('test', 'full', key)
# Capture the applied transform of the consumer_transform.
pcoll_id = aug_p._context.pcollections.get_id(pcoll)
consumer_transform_id = None
pipeline_proto = p.to_runner_api()
for (transform_id,
transform) in pipeline_proto.components.transforms.items():
if pcoll_id in transform.inputs.values():
consumer_transform_id = transform_id
break
self.assertIsNotNone(consumer_transform_id)
# Read cache on the pipeline proto.
_, cache_id = read_cache.ReadCache(
pipeline_proto, aug_p._context, aug_p._cache_manager,
aug_p._cacheables[pcoll]).read_cache()
actual_pipeline = pipeline_proto
# Read cache directly on the pipeline instance.
label = '{}{}'.format('_cache_', key)
transform = read_cache._ReadCacheTransform(aug_p._cache_manager, key,
label)
p | 'source' + label >> transform
expected_pipeline = p.to_runner_api()
# This rougly checks the equivalence between two protos, not detailed
# wiring in sub transforms under top level transforms.
assert_pipeline_proto_equal(self, expected_pipeline, actual_pipeline)
# Check if the actual_pipeline uses cache as input of the
# consumer_transform instead of the original pcoll from source.
inputs = actual_pipeline.components.transforms[
consumer_transform_id].inputs
self.assertIn(cache_id, inputs.values())
self.assertNotIn(pcoll_id, inputs.values())
def test_write_cache(self, mocked_get_cache_manager):
p = beam.Pipeline()
pcoll = p | beam.Create([1, 2, 3])
ib.watch(locals())
cache_manager = InMemoryCache()
mocked_get_cache_manager.return_value = cache_manager
aug_p = ap.AugmentedPipeline(p)
key = repr(aug_p._cacheables[pcoll].to_key())
pipeline_proto = p.to_runner_api()
# Write cache on the pipeline proto.
write_cache.WriteCache(pipeline_proto, aug_p._context,
aug_p._cache_manager,
aug_p._cacheables[pcoll]).write_cache()
actual_pipeline = pipeline_proto
# Write cache directly on the piepline instance.
label = '{}{}'.format('_cache_', key)
transform = write_cache._WriteCacheTransform(aug_p._cache_manager, key,
label)
_ = pcoll | 'sink' + label >> transform
expected_pipeline = p.to_runner_api()
assert_pipeline_proto_equal(self, expected_pipeline, actual_pipeline)
# Check if the actual_pipeline uses pcoll as an input of a write transform.
pcoll_id = aug_p._context.pcollections.get_id(pcoll)
write_transform_id = None
for transform_id, transform in \
actual_pipeline.components.transforms.items():
if pcoll_id in transform.inputs.values():
write_transform_id = transform_id
break
self.assertIsNotNone(write_transform_id)
self.assertIn(
'sink', actual_pipeline.components.transforms[write_transform_id].
unique_name)
def test_instrument_example_unbounded_pipeline_to_read_cache_not_cached(
self):
"""Tests that the instrumenter works when the PCollection is not cached.
"""
# Create the pipeline that will be instrumented.
from apache_beam.options.pipeline_options import StandardOptions
options = StandardOptions(streaming=True)
p_original_read_cache = beam.Pipeline(
interactive_runner.InteractiveRunner(), options)
ie.current_env().set_cache_manager(StreamingCache(cache_dir=None),
p_original_read_cache)
source_1 = p_original_read_cache | 'source1' >> beam.io.ReadFromPubSub(
subscription='projects/fake-project/subscriptions/fake_sub')
# pylint: disable=possibly-unused-variable
pcoll_1 = source_1 | 'square1' >> beam.Map(lambda x: x * x)
# Watch but do not cache the PCollections.
ib.watch(locals())
# This should be noop.
utils.watch_sources(p_original_read_cache)
# Instrument the original pipeline to create the pipeline the user will see.
p_copy = beam.Pipeline.from_runner_api(
p_original_read_cache.to_runner_api(),
runner=interactive_runner.InteractiveRunner(),
options=options)
ie.current_env().add_derived_pipeline(p_original_read_cache, p_copy)
instrumenter = instr.build_pipeline_instrument(p_copy)
actual_pipeline = beam.Pipeline.from_runner_api(
proto=instrumenter.instrumented_pipeline_proto(),
runner=interactive_runner.InteractiveRunner(),
options=options)
# Now, build the expected pipeline which replaces the unbounded source with
# a TestStream.
source_1_cache_key = self.cache_key_of('source_1', source_1)
p_expected = beam.Pipeline()
ie.current_env().set_cache_manager(StreamingCache(cache_dir=None),
p_expected)
test_stream = (p_expected
| TestStream(output_tags=[source_1_cache_key]))
# pylint: disable=expression-not-assigned
(test_stream[source_1_cache_key]
| 'square1' >> beam.Map(lambda x: x * x)
| 'reify' >> beam.Map(lambda _: _)
| cache.WriteCache(ie.current_env().get_cache_manager(p_expected),
'unused'))
# Test that the TestStream is outputting to the correct PCollection.
class TestStreamVisitor(PipelineVisitor):
def __init__(self):
self.output_tags = set()
def enter_composite_transform(self, transform_node):
self.visit_transform(transform_node)
def visit_transform(self, transform_node):
transform = transform_node.transform
if isinstance(transform, TestStream):
self.output_tags = transform.output_tags
v = TestStreamVisitor()
actual_pipeline.visit(v)
expected_output_tags = set([source_1_cache_key])
actual_output_tags = v.output_tags
self.assertSetEqual(expected_output_tags, actual_output_tags)
# Test that the pipeline is as expected.
assert_pipeline_proto_equal(self, p_expected.to_runner_api(),
instrumenter.instrumented_pipeline_proto())
def test_instrument_mixed_streaming_batch(self):
"""Tests caching for both batch and streaming sources in the same pipeline.
This ensures that cached bounded and unbounded sources are read from the
TestStream.
"""
# Create the pipeline that will be instrumented.
from apache_beam.options.pipeline_options import StandardOptions
options = StandardOptions(streaming=True)
p_original = beam.Pipeline(interactive_runner.InteractiveRunner(),
options)
streaming_cache_manager = StreamingCache(cache_dir=None)
ie.current_env().set_cache_manager(streaming_cache_manager, p_original)
source_1 = p_original | 'source1' >> beam.io.ReadFromPubSub(
subscription='projects/fake-project/subscriptions/fake_sub')
source_2 = p_original | 'source2' >> beam.Create([1, 2, 3, 4, 5])
# pylint: disable=possibly-unused-variable
pcoll_1 = ((source_1, source_2)
| beam.Flatten()
| 'square1' >> beam.Map(lambda x: x * x))
# Watch but do not cache the PCollections.
ib.watch(locals())
# This should be noop.
utils.watch_sources(p_original)
self._mock_write_cache(p_original, [],
self.cache_key_of('source_2', source_2))
ie.current_env().mark_pcollection_computed([source_2])
# Instrument the original pipeline to create the pipeline the user will see.
p_copy = beam.Pipeline.from_runner_api(
p_original.to_runner_api(),
runner=interactive_runner.InteractiveRunner(),
options=options)
ie.current_env().add_derived_pipeline(p_original, p_copy)
instrumenter = instr.build_pipeline_instrument(p_copy)
actual_pipeline = beam.Pipeline.from_runner_api(
proto=instrumenter.instrumented_pipeline_proto(),
runner=interactive_runner.InteractiveRunner(),
options=options)
# Now, build the expected pipeline which replaces the unbounded source with
# a TestStream.
source_1_cache_key = self.cache_key_of('source_1', source_1)
source_2_cache_key = self.cache_key_of('source_2', source_2)
p_expected = beam.Pipeline()
ie.current_env().set_cache_manager(streaming_cache_manager, p_expected)
test_stream = (
p_expected
| TestStream(output_tags=[source_1_cache_key, source_2_cache_key]))
# pylint: disable=expression-not-assigned
((test_stream[self.cache_key_of('source_1', source_1)],
test_stream[self.cache_key_of('source_2', source_2)])
| beam.Flatten()
| 'square1' >> beam.Map(lambda x: x * x)
| 'reify' >> beam.Map(lambda _: _)
| cache.WriteCache(ie.current_env().get_cache_manager(p_expected),
'unused'))
# Test that the TestStream is outputting to the correct PCollection.
class TestStreamVisitor(PipelineVisitor):
def __init__(self):
self.output_tags = set()
def enter_composite_transform(self, transform_node):
self.visit_transform(transform_node)
def visit_transform(self, transform_node):
transform = transform_node.transform
if isinstance(transform, TestStream):
self.output_tags = transform.output_tags
v = TestStreamVisitor()
actual_pipeline.visit(v)
expected_output_tags = set([source_1_cache_key, source_2_cache_key])
actual_output_tags = v.output_tags
self.assertSetEqual(expected_output_tags, actual_output_tags)
# Test that the pipeline is as expected.
assert_pipeline_proto_equal(self, p_expected.to_runner_api(),
instrumenter.instrumented_pipeline_proto())
def test_able_to_cache_intermediate_unbounded_source_pcollection(self):
"""Tests being able to cache an intermediate source PCollection.
In the following pipeline, the source doesn't have a reference and so is
not automatically cached in the watch() command. This tests that this case
is taken care of.
"""
# Create the pipeline that will be instrumented.
from apache_beam.options.pipeline_options import StandardOptions
options = StandardOptions(streaming=True)
streaming_cache_manager = StreamingCache(cache_dir=None)
p_original_cache_source = beam.Pipeline(
interactive_runner.InteractiveRunner(), options)
ie.current_env().set_cache_manager(streaming_cache_manager,
p_original_cache_source)
# pylint: disable=possibly-unused-variable
source_1 = (
p_original_cache_source
| 'source1' >> beam.io.ReadFromPubSub(
subscription='projects/fake-project/subscriptions/fake_sub')
| beam.Map(lambda e: e))
# Watch but do not cache the PCollections.
ib.watch(locals())
# Make sure that sources without a user reference are still cached.
utils.watch_sources(p_original_cache_source)
intermediate_source_pcoll = None
for watching in ie.current_env().watching():
watching = list(watching)
for var, watchable in watching:
if 'synthetic' in var:
intermediate_source_pcoll = watchable
break
# Instrument the original pipeline to create the pipeline the user will see.
p_copy = beam.Pipeline.from_runner_api(
p_original_cache_source.to_runner_api(),
runner=interactive_runner.InteractiveRunner(),
options=options)
ie.current_env().add_derived_pipeline(p_original_cache_source, p_copy)
instrumenter = instr.build_pipeline_instrument(p_copy)
actual_pipeline = beam.Pipeline.from_runner_api(
proto=instrumenter.instrumented_pipeline_proto(),
runner=interactive_runner.InteractiveRunner(),
options=options)
ie.current_env().add_derived_pipeline(p_original_cache_source,
actual_pipeline)
# Now, build the expected pipeline which replaces the unbounded source with
# a TestStream.
intermediate_source_pcoll_cache_key = \
self.cache_key_of('synthetic_var_' + str(id(intermediate_source_pcoll)),
intermediate_source_pcoll)
p_expected = beam.Pipeline()
ie.current_env().set_cache_manager(streaming_cache_manager, p_expected)
test_stream = (
p_expected
| TestStream(output_tags=[intermediate_source_pcoll_cache_key]))
# pylint: disable=expression-not-assigned
(test_stream[intermediate_source_pcoll_cache_key]
| 'square1' >> beam.Map(lambda e: e)
| 'reify' >> beam.Map(lambda _: _)
| cache.WriteCache(ie.current_env().get_cache_manager(p_expected),
'unused'))
# Test that the TestStream is outputting to the correct PCollection.
class TestStreamVisitor(PipelineVisitor):
def __init__(self):
self.output_tags = set()
def enter_composite_transform(self, transform_node):
self.visit_transform(transform_node)
def visit_transform(self, transform_node):
transform = transform_node.transform
if isinstance(transform, TestStream):
self.output_tags = transform.output_tags
v = TestStreamVisitor()
actual_pipeline.visit(v)
expected_output_tags = set([intermediate_source_pcoll_cache_key])
actual_output_tags = v.output_tags
self.assertSetEqual(expected_output_tags, actual_output_tags)
# Test that the pipeline is as expected.
assert_pipeline_proto_equal(self, p_expected.to_runner_api(),
instrumenter.instrumented_pipeline_proto())
def test_instrument_example_unbounded_pipeline_direct_from_source(self):
"""Tests that the it caches PCollections from a source.
"""
# Create a new interactive environment to make the test idempotent.
ie.new_env(cache_manager=streaming_cache.StreamingCache(
cache_dir=None))
# Create the pipeline that will be instrumented.
from apache_beam.options.pipeline_options import StandardOptions
options = StandardOptions(streaming=True)
p_original = beam.Pipeline(interactive_runner.InteractiveRunner(),
options)
source_1 = p_original | 'source1' >> beam.io.ReadFromPubSub(
subscription='projects/fake-project/subscriptions/fake_sub')
# pylint: disable=possibly-unused-variable
# Watch but do not cache the PCollections.
ib.watch(locals())
def cache_key_of(name, pcoll):
return name + '_' + str(id(pcoll)) + '_' + str(id(pcoll.producer))
# Instrument the original pipeline to create the pipeline the user will see.
p_copy = beam.Pipeline.from_runner_api(
p_original.to_runner_api(),
runner=interactive_runner.InteractiveRunner(),
options=options)
instrumenter = instr.build_pipeline_instrument(p_copy)
actual_pipeline = beam.Pipeline.from_runner_api(
proto=instrumenter.instrumented_pipeline_proto(),
runner=interactive_runner.InteractiveRunner(),
options=options)
# Now, build the expected pipeline which replaces the unbounded source with
# a TestStream.
source_1_cache_key = cache_key_of('source_1', source_1)
p_expected = beam.Pipeline()
# pylint: disable=unused-variable
test_stream = (
p_expected
| TestStream(output_tags=[cache_key_of('source_1', source_1)]))
# Test that the TestStream is outputting to the correct PCollection.
class TestStreamVisitor(PipelineVisitor):
def __init__(self):
self.output_tags = set()
def enter_composite_transform(self, transform_node):
self.visit_transform(transform_node)
def visit_transform(self, transform_node):
transform = transform_node.transform
if isinstance(transform, TestStream):
self.output_tags = transform.output_tags
v = TestStreamVisitor()
actual_pipeline.visit(v)
expected_output_tags = set([source_1_cache_key])
actual_output_tags = v.output_tags
self.assertSetEqual(expected_output_tags, actual_output_tags)
# Test that the pipeline is as expected.
assert_pipeline_proto_equal(
self, p_expected.to_runner_api(use_fake_coders=True),
instrumenter.instrumented_pipeline_proto())
