def test_apply_physical_no_filter(self): metrics = DirectMetrics() metrics.update_physical( object(), MetricUpdates( counters={ MetricKey('step1', self.name1): 5, MetricKey('step1', self.name3): 8 })) metrics.update_physical( object(), MetricUpdates( counters={ MetricKey('step2', self.name1): 7, MetricKey('step1', self.name3): 4 })) results = metrics.query() hc.assert_that( results['counters'], hc.contains_inanyorder(*[ MetricResult(MetricKey('step1', self.name1), 0, 5), MetricResult(MetricKey('step1', self.name3), 0, 12), MetricResult(MetricKey('step2', self.name1), 0, 7) ])) metrics.commit_physical(object(), MetricUpdates()) results = metrics.query() hc.assert_that( results['counters'], hc.contains_inanyorder(*[ MetricResult(MetricKey('step1', self.name1), 0, 5), MetricResult(MetricKey('step1', self.name3), 0, 12), MetricResult(MetricKey('step2', self.name1), 0, 7) ]))
def __init__( self, pipeline_options, bundle_factory, # type: BundleFactory root_transforms, value_to_consumers, step_names, views, # type: Iterable[pvalue.AsSideInput] clock): self.pipeline_options = pipeline_options self._bundle_factory = bundle_factory self._root_transforms = root_transforms self._value_to_consumers = value_to_consumers self._step_names = step_names self.views = views self._pcollection_to_views = collections.defaultdict( list) # type: DefaultDict[pvalue.PValue, List[pvalue.AsSideInput]] for view in views: self._pcollection_to_views[view.pvalue].append(view) self._transform_keyed_states = self._initialize_keyed_states( root_transforms, value_to_consumers) self._side_inputs_container = _SideInputsContainer(views) self._watermark_manager = WatermarkManager( clock, root_transforms, value_to_consumers, self._transform_keyed_states) self._pending_unblocked_tasks = [ ] # type: List[Tuple[TransformExecutor, Timestamp]] self._counter_factory = counters.CounterFactory() self._metrics = DirectMetrics() self._lock = threading.Lock() self.shutdown_requested = False
def test_apply_physical_no_filter(self): metrics = DirectMetrics() metrics.update_physical(object(), MetricUpdates( counters={MetricKey('step1', self.name1): 5, MetricKey('step1', self.name3): 8})) metrics.update_physical(object(), MetricUpdates( counters={MetricKey('step2', self.name1): 7, MetricKey('step1', self.name3): 4})) results = metrics.query() hc.assert_that(results['counters'], hc.contains_inanyorder(*[ MetricResult(MetricKey('step1', self.name1), 0, 5), MetricResult(MetricKey('step1', self.name3), 0, 12), MetricResult(MetricKey('step2', self.name1), 0, 7)])) metrics.commit_physical(object(), MetricUpdates()) results = metrics.query() hc.assert_that(results['counters'], hc.contains_inanyorder(*[ MetricResult(MetricKey('step1', self.name1), 0, 5), MetricResult(MetricKey('step1', self.name3), 0, 12), MetricResult(MetricKey('step2', self.name1), 0, 7)]))
def __init__(self, pipeline_options, bundle_factory, root_transforms, value_to_consumers, step_names, views): self.pipeline_options = pipeline_options self._bundle_factory = bundle_factory self._root_transforms = root_transforms self._value_to_consumers = value_to_consumers self._step_names = step_names self.views = views # AppliedPTransform -> Evaluator specific state objects self._application_state_interals = {} self._watermark_manager = WatermarkManager(Clock(), root_transforms, value_to_consumers) self._side_inputs_container = _SideInputsContainer(views) self._pending_unblocked_tasks = [] self._counter_factory = counters.CounterFactory() self._cache = None self._metrics = DirectMetrics() self._lock = threading.Lock()
def test_commit_logical_no_filter(self): metrics = DirectMetrics() metrics.commit_logical( self.bundle1, MetricUpdates( counters={MetricKey('step1', self.name1): 5, MetricKey('step1', self.name2): 8}, distributions={ MetricKey('step1', self.name1): DistributionData(8, 2, 3, 5)})) metrics.commit_logical( self.bundle1, MetricUpdates( counters={MetricKey('step2', self.name1): 7, MetricKey('step1', self.name2): 4}, distributions={ MetricKey('step1', self.name1): DistributionData(4, 1, 4, 4)})) results = metrics.query() hc.assert_that( results['counters'], hc.contains_inanyorder(*[ MetricResult(MetricKey('step1', self.name2), 12, 0), MetricResult(MetricKey('step2', self.name1), 7, 0), MetricResult(MetricKey('step1', self.name1), 5, 0)])) hc.assert_that( results['distributions'], hc.contains_inanyorder( MetricResult(MetricKey('step1', self.name1), DistributionResult( DistributionData(12, 3, 3, 5)), DistributionResult( DistributionData(0, 0, None, None)))))
def __init__(self, pipeline_options, bundle_factory, root_transforms, value_to_consumers, step_names, views, clock): self.pipeline_options = pipeline_options self._bundle_factory = bundle_factory self._root_transforms = root_transforms self._value_to_consumers = value_to_consumers self._step_names = step_names self.views = views self._pcollection_to_views = collections.defaultdict(list) for view in views: self._pcollection_to_views[view.pvalue].append(view) self._transform_keyed_states = self._initialize_keyed_states( root_transforms, value_to_consumers) self._side_inputs_container = _SideInputsContainer(views) self._watermark_manager = WatermarkManager( clock, root_transforms, value_to_consumers, self._transform_keyed_states) self._pending_unblocked_tasks = [] self._counter_factory = counters.CounterFactory() self._metrics = DirectMetrics() self._lock = threading.Lock()
def test_combiner_functions(self): metrics = DirectMetrics() counter = metrics._counters['anykey'] counter.commit_logical(self.bundle1, 5) self.assertEqual(counter.extract_committed(), 5) with self.assertRaises(TypeError): counter.commit_logical(self.bundle1, None) distribution = metrics._distributions['anykey'] distribution.commit_logical(self.bundle1, DistributionData(4, 1, 4, 4)) self.assertEqual(distribution.extract_committed(), DistributionResult(DistributionData(4, 1, 4, 4))) with self.assertRaises(AttributeError): distribution.commit_logical(self.bundle1, None)
def __init__(self, pipeline_options, bundle_factory, root_transforms, value_to_consumers, step_names, views): self.pipeline_options = pipeline_options self._bundle_factory = bundle_factory self._root_transforms = root_transforms self._value_to_consumers = value_to_consumers self._step_names = step_names self.views = views # AppliedPTransform -> Evaluator specific state objects self._application_state_interals = {} self._watermark_manager = WatermarkManager( Clock(), root_transforms, value_to_consumers) self._side_inputs_container = _SideInputsContainer(views) self._pending_unblocked_tasks = [] self._counter_factory = counters.CounterFactory() self._cache = None self._metrics = DirectMetrics() self._lock = threading.Lock()
def __init__(self, pipeline_options, bundle_factory, root_transforms, value_to_consumers, step_names, views, clock): self.pipeline_options = pipeline_options self._bundle_factory = bundle_factory self._root_transforms = root_transforms self._value_to_consumers = value_to_consumers self._step_names = step_names self.views = views self._pcollection_to_views = collections.defaultdict(list) for view in views: self._pcollection_to_views[view.pvalue].append(view) self._transform_keyed_states = self._initialize_keyed_states( root_transforms, value_to_consumers) self._side_inputs_container = _SideInputsContainer(views) self._watermark_manager = WatermarkManager( clock, root_transforms, value_to_consumers, self._transform_keyed_states) self._pending_unblocked_tasks = [] self._counter_factory = counters.CounterFactory() self._metrics = DirectMetrics() self._lock = threading.Lock()
def test_commit_logical_no_filter(self): metrics = DirectMetrics() metrics.commit_logical( self.bundle1, MetricUpdates( counters={ MetricKey('step1', self.name1): 5, MetricKey('step1', self.name2): 8 }, distributions={ MetricKey('step1', self.name1): DistributionData(8, 2, 3, 5) })) metrics.commit_logical( self.bundle1, MetricUpdates( counters={ MetricKey('step2', self.name1): 7, MetricKey('step1', self.name2): 4 }, distributions={ MetricKey('step1', self.name1): DistributionData(4, 1, 4, 4) })) results = metrics.query() hc.assert_that( results['counters'], hc.contains_inanyorder( *[ MetricResult(MetricKey('step1', self.name2), 12, 0), MetricResult(MetricKey('step2', self.name1), 7, 0), MetricResult(MetricKey('step1', self.name1), 5, 0) ])) hc.assert_that( results['distributions'], hc.contains_inanyorder( MetricResult( MetricKey('step1', self.name1), DistributionResult(DistributionData(12, 3, 3, 5)), DistributionResult(DistributionData(0, 0, None, None)))))
class EvaluationContext(object): """Evaluation context with the global state information of the pipeline. The evaluation context for a specific pipeline being executed by the DirectRunner. Contains state shared within the execution across all transforms. EvaluationContext contains shared state for an execution of the DirectRunner that can be used while evaluating a PTransform. This consists of views into underlying state and watermark implementations, access to read and write side inputs, and constructing counter sets and execution contexts. This includes executing callbacks asynchronously when state changes to the appropriate point (e.g. when a side input is requested and known to be empty). EvaluationContext also handles results by committing finalizing bundles based on the current global state and updating the global state appropriately. This includes updating the per-(step,key) state, updating global watermarks, and executing any callbacks that can be executed. """ def __init__(self, pipeline_options, bundle_factory, root_transforms, value_to_consumers, step_names, views, clock): self.pipeline_options = pipeline_options self._bundle_factory = bundle_factory self._root_transforms = root_transforms self._value_to_consumers = value_to_consumers self._step_names = step_names self.views = views self._pcollection_to_views = collections.defaultdict(list) for view in views: self._pcollection_to_views[view.pvalue].append(view) self._transform_keyed_states = self._initialize_keyed_states( root_transforms, value_to_consumers) self._watermark_manager = WatermarkManager( clock, root_transforms, value_to_consumers, self._transform_keyed_states) self._side_inputs_container = _SideInputsContainer(views) self._pending_unblocked_tasks = [] self._counter_factory = counters.CounterFactory() self._cache = None self._metrics = DirectMetrics() self._lock = threading.Lock() def _initialize_keyed_states(self, root_transforms, value_to_consumers): transform_keyed_states = {} for transform in root_transforms: transform_keyed_states[transform] = {} for consumers in value_to_consumers.values(): for consumer in consumers: transform_keyed_states[consumer] = {} return transform_keyed_states def use_pvalue_cache(self, cache): assert not self._cache self._cache = cache def metrics(self): # TODO. Should this be made a @property? return self._metrics @property def has_cache(self): return self._cache is not None def append_to_cache(self, applied_ptransform, tag, elements): with self._lock: assert self._cache self._cache.append(applied_ptransform, tag, elements) def is_root_transform(self, applied_ptransform): return applied_ptransform in self._root_transforms def handle_result( self, completed_bundle, completed_timers, result): """Handle the provided result produced after evaluating the input bundle. Handle the provided TransformResult, produced after evaluating the provided committed bundle (potentially None, if the result of a root PTransform). The result is the output of running the transform contained in the TransformResult on the contents of the provided bundle. Args: completed_bundle: the bundle that was processed to produce the result. completed_timers: the timers that were delivered to produce the completed_bundle. result: the TransformResult of evaluating the input bundle Returns: the committed bundles contained within the handled result. """ with self._lock: committed_bundles, unprocessed_bundles = self._commit_bundles( result.uncommitted_output_bundles, result.unprocessed_bundles) self._watermark_manager.update_watermarks( completed_bundle, result.transform, completed_timers, committed_bundles, unprocessed_bundles, result.keyed_watermark_holds) self._metrics.commit_logical(completed_bundle, result.logical_metric_updates) # If the result is for a view, update side inputs container. if (result.uncommitted_output_bundles and result.uncommitted_output_bundles[0].pcollection in self._pcollection_to_views): for view in self._pcollection_to_views[ result.uncommitted_output_bundles[0].pcollection]: for committed_bundle in committed_bundles: # side_input must be materialized. self._side_inputs_container.add_values( view, committed_bundle.get_elements_iterable(make_copy=True)) if (self.get_execution_context(result.transform) .watermarks.input_watermark == WatermarkManager.WATERMARK_POS_INF): self._pending_unblocked_tasks.extend( self._side_inputs_container.finalize_value_and_get_tasks(view)) if result.counters: for counter in result.counters: merged_counter = self._counter_factory.get_counter( counter.name, counter.combine_fn) merged_counter.accumulator.merge([counter.accumulator]) # Commit partial GBK states existing_keyed_state = self._transform_keyed_states[result.transform] for k, v in result.partial_keyed_state.iteritems(): existing_keyed_state[k] = v return committed_bundles def get_aggregator_values(self, aggregator_or_name): return self._counter_factory.get_aggregator_values(aggregator_or_name) def schedule_pending_unblocked_tasks(self, executor_service): if self._pending_unblocked_tasks: with self._lock: for task in self._pending_unblocked_tasks: executor_service.submit(task) self._pending_unblocked_tasks = [] def _commit_bundles(self, uncommitted_bundles, unprocessed_bundles): """Commits bundles and returns a immutable set of committed bundles.""" for in_progress_bundle in uncommitted_bundles: producing_applied_ptransform = in_progress_bundle.pcollection.producer watermarks = self._watermark_manager.get_watermarks( producing_applied_ptransform) in_progress_bundle.commit(watermarks.synchronized_processing_output_time) for unprocessed_bundle in unprocessed_bundles: unprocessed_bundle.commit(None) return tuple(uncommitted_bundles), tuple(unprocessed_bundles) def get_execution_context(self, applied_ptransform): return _ExecutionContext( self._watermark_manager.get_watermarks(applied_ptransform), self._transform_keyed_states[applied_ptransform]) def create_bundle(self, output_pcollection): """Create an uncommitted bundle for the specified PCollection.""" return self._bundle_factory.create_bundle(output_pcollection) def create_empty_committed_bundle(self, output_pcollection): """Create empty bundle useful for triggering evaluation.""" return self._bundle_factory.create_empty_committed_bundle( output_pcollection) def extract_all_timers(self): return self._watermark_manager.extract_all_timers() def is_done(self, transform=None): """Checks completion of a step or the pipeline. Args: transform: AppliedPTransform to check for completion. Returns: True if the step will not produce additional output. If transform is None returns true if all steps are done. """ if transform: return self._is_transform_done(transform) for applied_ptransform in self._step_names: if not self._is_transform_done(applied_ptransform): return False return True def _is_transform_done(self, transform): tw = self._watermark_manager.get_watermarks(transform) return tw.output_watermark == WatermarkManager.WATERMARK_POS_INF def get_value_or_schedule_after_output(self, side_input, task): assert isinstance(task, TransformExecutor) return self._side_inputs_container.get_value_or_schedule_after_output( side_input, task)
class EvaluationContext(object): """Evaluation context with the global state information of the pipeline. The evaluation context for a specific pipeline being executed by the DirectRunner. Contains state shared within the execution across all transforms. EvaluationContext contains shared state for an execution of the DirectRunner that can be used while evaluating a PTransform. This consists of views into underlying state and watermark implementations, access to read and write side inputs, and constructing counter sets and execution contexts. This includes executing callbacks asynchronously when state changes to the appropriate point (e.g. when a side input is requested and known to be empty). EvaluationContext also handles results by committing finalizing bundles based on the current global state and updating the global state appropriately. This includes updating the per-(step,key) state, updating global watermarks, and executing any callbacks that can be executed. """ def __init__(self, pipeline_options, bundle_factory, root_transforms, value_to_consumers, step_names, views, clock): self.pipeline_options = pipeline_options self._bundle_factory = bundle_factory self._root_transforms = root_transforms self._value_to_consumers = value_to_consumers self._step_names = step_names self.views = views self._pcollection_to_views = collections.defaultdict(list) for view in views: self._pcollection_to_views[view.pvalue].append(view) self._transform_keyed_states = self._initialize_keyed_states( root_transforms, value_to_consumers) self._side_inputs_container = _SideInputsContainer(views) self._watermark_manager = WatermarkManager( clock, root_transforms, value_to_consumers, self._transform_keyed_states) self._pending_unblocked_tasks = [] self._counter_factory = counters.CounterFactory() self._metrics = DirectMetrics() self._lock = threading.Lock() def _initialize_keyed_states(self, root_transforms, value_to_consumers): """Initialize user state dicts. These dicts track user state per-key, per-transform and per-window. """ transform_keyed_states = {} for transform in root_transforms: transform_keyed_states[transform] = {} for consumers in value_to_consumers.values(): for consumer in consumers: transform_keyed_states[consumer] = {} return transform_keyed_states def metrics(self): # TODO. Should this be made a @property? return self._metrics def is_root_transform(self, applied_ptransform): return applied_ptransform in self._root_transforms def handle_result(self, completed_bundle, completed_timers, result): """Handle the provided result produced after evaluating the input bundle. Handle the provided TransformResult, produced after evaluating the provided committed bundle (potentially None, if the result of a root PTransform). The result is the output of running the transform contained in the TransformResult on the contents of the provided bundle. Args: completed_bundle: the bundle that was processed to produce the result. completed_timers: the timers that were delivered to produce the completed_bundle. result: the ``TransformResult`` of evaluating the input bundle Returns: the committed bundles contained within the handled result. """ with self._lock: committed_bundles, unprocessed_bundles = self._commit_bundles( result.uncommitted_output_bundles, result.unprocessed_bundles) self._metrics.commit_logical(completed_bundle, result.logical_metric_updates) # If the result is for a view, update side inputs container. self._update_side_inputs_container(committed_bundles, result) # Tasks generated from unblocked side inputs as the watermark progresses. tasks = self._watermark_manager.update_watermarks( completed_bundle, result.transform, completed_timers, committed_bundles, unprocessed_bundles, result.keyed_watermark_holds, self._side_inputs_container) self._pending_unblocked_tasks.extend(tasks) if result.counters: for counter in result.counters: merged_counter = self._counter_factory.get_counter( counter.name, counter.combine_fn) merged_counter.accumulator.merge([counter.accumulator]) # Commit partial GBK states existing_keyed_state = self._transform_keyed_states[ result.transform] for k, v in result.partial_keyed_state.items(): existing_keyed_state[k] = v return committed_bundles def _update_side_inputs_container(self, committed_bundles, result): """Update the side inputs container if we are outputting into a side input. Look at the result, and if it's outputing into a PCollection that we have registered as a PCollectionView, we add the result to the PCollectionView. """ if (result.uncommitted_output_bundles and result.uncommitted_output_bundles[0].pcollection in self._pcollection_to_views): for view in self._pcollection_to_views[ result.uncommitted_output_bundles[0].pcollection]: for committed_bundle in committed_bundles: # side_input must be materialized. self._side_inputs_container.add_values( view, committed_bundle.get_elements_iterable(make_copy=True)) def get_aggregator_values(self, aggregator_or_name): return self._counter_factory.get_aggregator_values(aggregator_or_name) def schedule_pending_unblocked_tasks(self, executor_service): if self._pending_unblocked_tasks: with self._lock: for task in self._pending_unblocked_tasks: executor_service.submit(task) self._pending_unblocked_tasks = [] def _commit_bundles(self, uncommitted_bundles, unprocessed_bundles): """Commits bundles and returns a immutable set of committed bundles.""" for in_progress_bundle in uncommitted_bundles: producing_applied_ptransform = in_progress_bundle.pcollection.producer watermarks = self._watermark_manager.get_watermarks( producing_applied_ptransform) in_progress_bundle.commit( watermarks.synchronized_processing_output_time) for unprocessed_bundle in unprocessed_bundles: unprocessed_bundle.commit(None) return tuple(uncommitted_bundles), tuple(unprocessed_bundles) def get_execution_context(self, applied_ptransform): return _ExecutionContext( self._watermark_manager.get_watermarks(applied_ptransform), self._transform_keyed_states[applied_ptransform]) def create_bundle(self, output_pcollection): """Create an uncommitted bundle for the specified PCollection.""" return self._bundle_factory.create_bundle(output_pcollection) def create_empty_committed_bundle(self, output_pcollection): """Create empty bundle useful for triggering evaluation.""" return self._bundle_factory.create_empty_committed_bundle( output_pcollection) def extract_all_timers(self): return self._watermark_manager.extract_all_timers() def is_done(self, transform=None): """Checks completion of a step or the pipeline. Args: transform: AppliedPTransform to check for completion. Returns: True if the step will not produce additional output. If transform is None returns true if all steps are done. """ if transform: return self._is_transform_done(transform) for applied_ptransform in self._step_names: if not self._is_transform_done(applied_ptransform): return False return True def _is_transform_done(self, transform): tw = self._watermark_manager.get_watermarks(transform) return tw.output_watermark == WatermarkManager.WATERMARK_POS_INF def get_value_or_block_until_ready(self, side_input, task, block_until): assert isinstance(task, TransformExecutor) return self._side_inputs_container.get_value_or_block_until_ready( side_input, task, block_until)
def test_apply_physical_logical(self): metrics = DirectMetrics() dist_zero = DistributionData(0, 0, None, None) metrics.update_physical( object(), MetricUpdates(counters={ MetricKey('step1', self.name1): 7, MetricKey('step1', self.name2): 5, MetricKey('step2', self.name1): 1 }, distributions={ MetricKey('step1', self.name1): DistributionData(3, 1, 3, 3), MetricKey('step2', self.name3): DistributionData(8, 2, 4, 4) })) results = metrics.query() hc.assert_that( results['counters'], hc.contains_inanyorder(*[ MetricResult(MetricKey('step1', self.name1), 0, 7), MetricResult(MetricKey('step1', self.name2), 0, 5), MetricResult(MetricKey('step2', self.name1), 0, 1) ])) hc.assert_that( results['distributions'], hc.contains_inanyorder(*[ MetricResult(MetricKey('step1', self.name1), DistributionResult(dist_zero), DistributionResult(DistributionData(3, 1, 3, 3))), MetricResult(MetricKey('step2', self.name3), DistributionResult(dist_zero), DistributionResult(DistributionData(8, 2, 4, 4))) ])) metrics.commit_physical( object(), MetricUpdates(counters={ MetricKey('step1', self.name1): -3, MetricKey('step2', self.name1): -5 }, distributions={ MetricKey('step1', self.name1): DistributionData(8, 4, 1, 5), MetricKey('step2', self.name2): DistributionData(8, 8, 1, 1) })) results = metrics.query() hc.assert_that( results['counters'], hc.contains_inanyorder(*[ MetricResult(MetricKey('step1', self.name1), 0, 4), MetricResult(MetricKey('step1', self.name2), 0, 5), MetricResult(MetricKey('step2', self.name1), 0, -4) ])) hc.assert_that( results['distributions'], hc.contains_inanyorder(*[ MetricResult(MetricKey('step1', self.name1), DistributionResult(dist_zero), DistributionResult(DistributionData(11, 5, 1, 5))), MetricResult(MetricKey('step2', self.name3), DistributionResult(dist_zero), DistributionResult(DistributionData(8, 2, 4, 4))), MetricResult(MetricKey('step2', self.name2), DistributionResult(dist_zero), DistributionResult(DistributionData(8, 8, 1, 1))) ])) metrics.commit_logical( object(), MetricUpdates(counters={ MetricKey('step1', self.name1): 3, MetricKey('step1', self.name2): 5, MetricKey('step2', self.name1): -3 }, distributions={ MetricKey('step1', self.name1): DistributionData(11, 5, 1, 5), MetricKey('step2', self.name2): DistributionData(8, 8, 1, 1), MetricKey('step2', self.name3): DistributionData(4, 1, 4, 4) })) results = metrics.query() hc.assert_that( results['counters'], hc.contains_inanyorder(*[ MetricResult(MetricKey('step1', self.name1), 3, 4), MetricResult(MetricKey('step1', self.name2), 5, 5), MetricResult(MetricKey('step2', self.name1), -3, -4) ])) hc.assert_that( results['distributions'], hc.contains_inanyorder(*[ MetricResult(MetricKey('step1', self.name1), DistributionResult(DistributionData(11, 5, 1, 5)), DistributionResult(DistributionData(11, 5, 1, 5))), MetricResult(MetricKey('step2', self.name3), DistributionResult(DistributionData(4, 1, 4, 4)), DistributionResult(DistributionData(8, 2, 4, 4))), MetricResult(MetricKey('step2', self.name2), DistributionResult(DistributionData(8, 8, 1, 1)), DistributionResult(DistributionData(8, 8, 1, 1))) ]))
class EvaluationContext(object): """Evaluation context with the global state information of the pipeline. The evaluation context for a specific pipeline being executed by the DirectRunner. Contains state shared within the execution across all transforms. EvaluationContext contains shared state for an execution of the DirectRunner that can be used while evaluating a PTransform. This consists of views into underlying state and watermark implementations, access to read and write PCollectionViews, and constructing counter sets and execution contexts. This includes executing callbacks asynchronously when state changes to the appropriate point (e.g. when a PCollectionView is requested and known to be empty). EvaluationContext also handles results by committing finalizing bundles based on the current global state and updating the global state appropriately. This includes updating the per-(step,key) state, updating global watermarks, and executing any callbacks that can be executed. """ def __init__(self, pipeline_options, bundle_factory, root_transforms, value_to_consumers, step_names, views): self.pipeline_options = pipeline_options self._bundle_factory = bundle_factory self._root_transforms = root_transforms self._value_to_consumers = value_to_consumers self._step_names = step_names self.views = views # AppliedPTransform -> Evaluator specific state objects self._application_state_interals = {} self._watermark_manager = WatermarkManager(Clock(), root_transforms, value_to_consumers) self._side_inputs_container = _SideInputsContainer(views) self._pending_unblocked_tasks = [] self._counter_factory = counters.CounterFactory() self._cache = None self._metrics = DirectMetrics() self._lock = threading.Lock() def use_pvalue_cache(self, cache): assert not self._cache self._cache = cache def metrics(self): # TODO. Should this be made a @property? return self._metrics @property def has_cache(self): return self._cache is not None def append_to_cache(self, applied_ptransform, tag, elements): with self._lock: assert self._cache self._cache.append(applied_ptransform, tag, elements) def is_root_transform(self, applied_ptransform): return applied_ptransform in self._root_transforms def handle_result(self, completed_bundle, completed_timers, result): """Handle the provided result produced after evaluating the input bundle. Handle the provided TransformResult, produced after evaluating the provided committed bundle (potentially None, if the result of a root PTransform). The result is the output of running the transform contained in the TransformResult on the contents of the provided bundle. Args: completed_bundle: the bundle that was processed to produce the result. completed_timers: the timers that were delivered to produce the completed_bundle. result: the TransformResult of evaluating the input bundle Returns: the committed bundles contained within the handled result. """ with self._lock: committed_bundles = self._commit_bundles(result.output_bundles) self._watermark_manager.update_watermarks(completed_bundle, result.transform, completed_timers, committed_bundles, result.watermark_hold) self._metrics.commit_logical(completed_bundle, result.logical_metric_updates()) # If the result is for a view, update side inputs container. if (result.output_bundles and result.output_bundles[0].pcollection in self.views): if committed_bundles: assert len(committed_bundles) == 1 # side_input must be materialized. side_input_result = committed_bundles[ 0].get_elements_iterable(make_copy=True) else: side_input_result = [] tasks = self._side_inputs_container.set_value_and_get_callables( result.output_bundles[0].pcollection, side_input_result) self._pending_unblocked_tasks.extend(tasks) if result.counters: for counter in result.counters: merged_counter = self._counter_factory.get_counter( counter.name, counter.combine_fn) merged_counter.accumulator.merge([counter.accumulator]) self._application_state_interals[result.transform] = result.state return committed_bundles def get_aggregator_values(self, aggregator_or_name): return self._counter_factory.get_aggregator_values(aggregator_or_name) def schedule_pending_unblocked_tasks(self, executor_service): if self._pending_unblocked_tasks: with self._lock: for task in self._pending_unblocked_tasks: executor_service.submit(task) self._pending_unblocked_tasks = [] def _commit_bundles(self, uncommitted_bundles): """Commits bundles and returns a immutable set of committed bundles.""" for in_progress_bundle in uncommitted_bundles: producing_applied_ptransform = in_progress_bundle.pcollection.producer watermarks = self._watermark_manager.get_watermarks( producing_applied_ptransform) in_progress_bundle.commit( watermarks.synchronized_processing_output_time) return tuple(uncommitted_bundles) def get_execution_context(self, applied_ptransform): return _ExecutionContext( self._watermark_manager.get_watermarks(applied_ptransform), self._application_state_interals.get(applied_ptransform)) def create_bundle(self, output_pcollection): """Create an uncommitted bundle for the specified PCollection.""" return self._bundle_factory.create_bundle(output_pcollection) def create_empty_committed_bundle(self, output_pcollection): """Create empty bundle useful for triggering evaluation.""" return self._bundle_factory.create_empty_committed_bundle( output_pcollection) def extract_fired_timers(self): return self._watermark_manager.extract_fired_timers() def is_done(self, transform=None): """Checks completion of a step or the pipeline. Args: transform: AppliedPTransform to check for completion. Returns: True if the step will not produce additional output. If transform is None returns true if all steps are done. """ if transform: return self._is_transform_done(transform) else: for applied_ptransform in self._step_names: if not self._is_transform_done(applied_ptransform): return False return True def _is_transform_done(self, transform): tw = self._watermark_manager.get_watermarks(transform) return tw.output_watermark == WatermarkManager.WATERMARK_POS_INF def get_value_or_schedule_after_output(self, pcollection_view, task): assert isinstance(task, TransformExecutor) return self._side_inputs_container.get_value_or_schedule_after_output( pcollection_view, task)
def test_apply_physical_logical(self): metrics = DirectMetrics() dist_zero = DistributionData(0, 0, None, None) metrics.update_physical( object(), MetricUpdates( counters={MetricKey('step1', self.name1): 7, MetricKey('step1', self.name2): 5, MetricKey('step2', self.name1): 1}, distributions={MetricKey('step1', self.name1): DistributionData(3, 1, 3, 3), MetricKey('step2', self.name3): DistributionData(8, 2, 4, 4)})) results = metrics.query() hc.assert_that(results['counters'], hc.contains_inanyorder(*[ MetricResult(MetricKey('step1', self.name1), 0, 7), MetricResult(MetricKey('step1', self.name2), 0, 5), MetricResult(MetricKey('step2', self.name1), 0, 1)])) hc.assert_that(results['distributions'], hc.contains_inanyorder(*[ MetricResult( MetricKey('step1', self.name1), DistributionResult(dist_zero), DistributionResult(DistributionData(3, 1, 3, 3))), MetricResult( MetricKey('step2', self.name3), DistributionResult(dist_zero), DistributionResult(DistributionData(8, 2, 4, 4)))])) metrics.commit_physical( object(), MetricUpdates( counters={MetricKey('step1', self.name1): -3, MetricKey('step2', self.name1): -5}, distributions={MetricKey('step1', self.name1): DistributionData(8, 4, 1, 5), MetricKey('step2', self.name2): DistributionData(8, 8, 1, 1)})) results = metrics.query() hc.assert_that(results['counters'], hc.contains_inanyorder(*[ MetricResult(MetricKey('step1', self.name1), 0, 4), MetricResult(MetricKey('step1', self.name2), 0, 5), MetricResult(MetricKey('step2', self.name1), 0, -4)])) hc.assert_that(results['distributions'], hc.contains_inanyorder(*[ MetricResult( MetricKey('step1', self.name1), DistributionResult(dist_zero), DistributionResult(DistributionData(11, 5, 1, 5))), MetricResult( MetricKey('step2', self.name3), DistributionResult(dist_zero), DistributionResult(DistributionData(8, 2, 4, 4))), MetricResult( MetricKey('step2', self.name2), DistributionResult(dist_zero), DistributionResult(DistributionData(8, 8, 1, 1)))])) metrics.commit_logical( object(), MetricUpdates( counters={MetricKey('step1', self.name1): 3, MetricKey('step1', self.name2): 5, MetricKey('step2', self.name1): -3}, distributions={MetricKey('step1', self.name1): DistributionData(11, 5, 1, 5), MetricKey('step2', self.name2): DistributionData(8, 8, 1, 1), MetricKey('step2', self.name3): DistributionData(4, 1, 4, 4)})) results = metrics.query() hc.assert_that(results['counters'], hc.contains_inanyorder(*[ MetricResult(MetricKey('step1', self.name1), 3, 4), MetricResult(MetricKey('step1', self.name2), 5, 5), MetricResult(MetricKey('step2', self.name1), -3, -4)])) hc.assert_that(results['distributions'], hc.contains_inanyorder(*[ MetricResult( MetricKey('step1', self.name1), DistributionResult(DistributionData(11, 5, 1, 5)), DistributionResult(DistributionData(11, 5, 1, 5))), MetricResult( MetricKey('step2', self.name3), DistributionResult(DistributionData(4, 1, 4, 4)), DistributionResult(DistributionData(8, 2, 4, 4))), MetricResult( MetricKey('step2', self.name2), DistributionResult(DistributionData(8, 8, 1, 1)), DistributionResult(DistributionData(8, 8, 1, 1)))]))