def test_equal_objects(self): self.assertEqual( CounterName('counter_name', 'stage_name', 'step_name'), CounterName('counter_name', 'stage_name', 'step_name')) self.assertNotEqual( CounterName('counter_name', 'stage_name', 'step_name'), CounterName('counter_name', 'stage_name', 'step_nam')) # Testing objects with an IOTarget. self.assertEqual( CounterName('counter_name', 'stage_name', 'step_name', io_target=counters.side_input_id(1, 's9')), CounterName('counter_name', 'stage_name', 'step_name', io_target=counters.side_input_id(1, 's9'))) self.assertNotEqual( CounterName('counter_name', 'stage_name', 'step_name', io_target=counters.side_input_id(1, 's')), CounterName('counter_name', 'stage_name', 'step_name', io_target=counters.side_input_id(1, 's9')))
def test_equal_objects(self): self.assertEqual(CounterName('counter_name', 'stage_name', 'step_name'), CounterName('counter_name', 'stage_name', 'step_name')) self.assertNotEqual(CounterName('counter_name', 'stage_name', 'step_name'), CounterName('counter_name', 'stage_name', 'step_nam')) # Testing objects with an IOTarget. self.assertEqual(CounterName('counter_name', 'stage_name', 'step_name', io_target=counters.side_input_id(1, 's9')), CounterName('counter_name', 'stage_name', 'step_name', io_target=counters.side_input_id(1, 's9'))) self.assertNotEqual(CounterName('counter_name', 'stage_name', 'step_name', io_target=counters.side_input_id(1, 's')), CounterName('counter_name', 'stage_name', 'step_name', io_target=counters.side_input_id(1, 's9')))
def _update_counters_for_requesting_step(self, step_name): side_input_id = counters.side_input_id(step_name, self.input_index) self.scoped_state = self._state_sampler.scoped_state( self.declaring_step, 'read-sideinput', io_target=side_input_id) self.bytes_read_counter = self._counter_factory.get_counter( CounterName('read-sideinput-byte-count', step_name=self.declaring_step, io_target=side_input_id), Counter.SUM)
def test_basic_counters(self): counter_factory = CounterFactory() sampler = statesampler.StateSampler('stage1', counter_factory) sampler.start() with sampler.scoped_state('step1', 'stateA'): counter = opcounters.SideInputReadCounter(counter_factory, sampler, declaring_step='step1', input_index=1) with sampler.scoped_state('step2', 'stateB'): with counter: counter.add_bytes_read(10) counter.update_current_step() sampler.stop() sampler.commit_counters() actual_counter_names = set( [c.name for c in counter_factory.get_counters()]) expected_counter_names = set([ # Counter names for STEP 1 counters.CounterName('read-sideinput-msecs', stage_name='stage1', step_name='step1', io_target=counters.side_input_id('step1', 1)), counters.CounterName('read-sideinput-byte-count', step_name='step1', io_target=counters.side_input_id('step1', 1)), # Counter names for STEP 2 counters.CounterName('read-sideinput-msecs', stage_name='stage1', step_name='step1', io_target=counters.side_input_id('step2', 1)), counters.CounterName('read-sideinput-byte-count', step_name='step1', io_target=counters.side_input_id('step2', 1)), ]) self.assertTrue( actual_counter_names.issuperset(expected_counter_names))
def _update_counters_for_requesting_step(self, step_name): side_input_id = counters.side_input_id(step_name, self.input_index) self.scoped_state = self._state_sampler.scoped_state( self.declaring_step, 'read-sideinput', io_target=side_input_id) self.bytes_read_counter = self._counter_factory.get_counter( CounterName( 'read-sideinput-byte-count', step_name=self.declaring_step, io_target=side_input_id), Counter.SUM)
def test_basic_counters(self): counter_factory = CounterFactory() sampler = statesampler.StateSampler('stage1', counter_factory) sampler.start() with sampler.scoped_state('step1', 'stateA'): counter = opcounters.SideInputReadCounter(counter_factory, sampler, declaring_step='step1', input_index=1) with sampler.scoped_state('step2', 'stateB'): with counter: counter.add_bytes_read(10) counter.update_current_step() sampler.stop() sampler.commit_counters() actual_counter_names = set([c.name for c in counter_factory.get_counters()]) expected_counter_names = set([ # Counter names for STEP 1 counters.CounterName('read-sideinput-msecs', stage_name='stage1', step_name='step1', io_target=counters.side_input_id('step1', 1)), counters.CounterName('read-sideinput-byte-count', step_name='step1', io_target=counters.side_input_id('step1', 1)), # Counter names for STEP 2 counters.CounterName('read-sideinput-msecs', stage_name='stage1', step_name='step1', io_target=counters.side_input_id('step2', 1)), counters.CounterName('read-sideinput-byte-count', step_name='step1', io_target=counters.side_input_id('step2', 1)), ]) self.assertTrue(actual_counter_names.issuperset(expected_counter_names))
def update_current_step(self): """Update the current running step. Due to the fusion optimization, user code may choose to emit the data structure that holds side inputs (Iterable, Dict, or others). This call updates the current step, to attribute the data consumption to the step that is responsible for actual consumption. CounterName uses the io_target field for information pertinent to the consumption of side inputs. """ current_state = self._state_sampler.current_state() operation_name = current_state.name.step_name self.scoped_state = self._state_sampler.scoped_state( self.declaring_step, 'read-sideinput', io_target=counters.side_input_id(operation_name, self.input_index)) self.bytes_read_counter = self._counter_factory.get_counter( CounterName( 'read-sideinput-byte-count', step_name=self.declaring_step, io_target=counters.side_input_id(operation_name, self.input_index)), Counter.SUM)