| # |
| # 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. |
| # |
| |
| """Tests for apache_beam.runners.worker.sdk_worker.""" |
| |
| # pytype: skip-file |
| |
| import contextlib |
| import logging |
| import unittest |
| from collections import namedtuple |
| |
| import grpc |
| import hamcrest as hc |
| import mock |
| |
| from apache_beam.coders import FastPrimitivesCoder |
| from apache_beam.coders import VarIntCoder |
| from apache_beam.internal.metrics.metric import Metrics as InternalMetrics |
| from apache_beam.metrics import monitoring_infos |
| from apache_beam.metrics.execution import MetricsEnvironment |
| from apache_beam.portability.api import beam_fn_api_pb2 |
| from apache_beam.portability.api import beam_fn_api_pb2_grpc |
| from apache_beam.portability.api import beam_runner_api_pb2 |
| from apache_beam.portability.api import metrics_pb2 |
| from apache_beam.runners.worker import sdk_worker |
| from apache_beam.runners.worker import statecache |
| from apache_beam.runners.worker.sdk_worker import BundleProcessorCache |
| from apache_beam.runners.worker.sdk_worker import GlobalCachingStateHandler |
| from apache_beam.runners.worker.sdk_worker import SdkWorker |
| from apache_beam.utils import thread_pool_executor |
| |
| _LOGGER = logging.getLogger(__name__) |
| |
| |
| class BeamFnControlServicer(beam_fn_api_pb2_grpc.BeamFnControlServicer): |
| def __init__(self, requests, raise_errors=True): |
| self.requests = requests |
| self.instruction_ids = set(r.instruction_id for r in requests) |
| self.responses = {} |
| self.raise_errors = raise_errors |
| |
| def Control(self, response_iterator, context): |
| for request in self.requests: |
| _LOGGER.info("Sending request %s", request) |
| yield request |
| for response in response_iterator: |
| _LOGGER.info("Got response %s", response) |
| if response.instruction_id != -1: |
| assert response.instruction_id in self.instruction_ids |
| assert response.instruction_id not in self.responses |
| self.responses[response.instruction_id] = response |
| if self.raise_errors and response.error: |
| raise RuntimeError(response.error) |
| elif len(self.responses) == len(self.requests): |
| _LOGGER.info("All %s instructions finished.", len(self.requests)) |
| return |
| raise RuntimeError( |
| "Missing responses: %s" % |
| (self.instruction_ids - set(self.responses.keys()))) |
| |
| |
| class SdkWorkerTest(unittest.TestCase): |
| def _get_process_bundles(self, prefix, size): |
| return [ |
| beam_fn_api_pb2.ProcessBundleDescriptor( |
| id=str(str(prefix) + "-" + str(ix)), |
| transforms={ |
| str(ix): beam_runner_api_pb2.PTransform(unique_name=str(ix)) |
| }) for ix in range(size) |
| ] |
| |
| def _check_fn_registration_multi_request(self, *args): |
| """Check the function registration calls to the sdk_harness. |
| |
| Args: |
| tuple of request_count, number of process_bundles per request and workers |
| counts to process the request. |
| """ |
| for (request_count, process_bundles_per_request) in args: |
| requests = [] |
| process_bundle_descriptors = [] |
| |
| for i in range(request_count): |
| pbd = self._get_process_bundles(i, process_bundles_per_request) |
| process_bundle_descriptors.extend(pbd) |
| requests.append( |
| beam_fn_api_pb2.InstructionRequest( |
| instruction_id=str(i), |
| register=beam_fn_api_pb2.RegisterRequest( |
| process_bundle_descriptor=process_bundle_descriptors))) |
| |
| test_controller = BeamFnControlServicer(requests) |
| |
| server = grpc.server(thread_pool_executor.shared_unbounded_instance()) |
| beam_fn_api_pb2_grpc.add_BeamFnControlServicer_to_server( |
| test_controller, server) |
| test_port = server.add_insecure_port("[::]:0") |
| server.start() |
| |
| harness = sdk_worker.SdkHarness( |
| "localhost:%s" % test_port, state_cache_size=100) |
| harness.run() |
| |
| self.assertEqual( |
| harness._bundle_processor_cache.fns, |
| {item.id: item |
| for item in process_bundle_descriptors}) |
| |
| def test_fn_registration(self): |
| self._check_fn_registration_multi_request((1, 4), (4, 4)) |
| |
| def test_inactive_bundle_processor_returns_empty_progress_response(self): |
| bundle_processor = mock.MagicMock() |
| bundle_processor_cache = BundleProcessorCache(None, None, None, {}) |
| bundle_processor_cache.activate('instruction_id') |
| worker = SdkWorker(bundle_processor_cache) |
| split_request = beam_fn_api_pb2.InstructionRequest( |
| instruction_id='progress_instruction_id', |
| process_bundle_progress=beam_fn_api_pb2.ProcessBundleProgressRequest( |
| instruction_id='instruction_id')) |
| self.assertEqual( |
| worker.do_instruction(split_request), |
| beam_fn_api_pb2.InstructionResponse( |
| instruction_id='progress_instruction_id', |
| process_bundle_progress=beam_fn_api_pb2. |
| ProcessBundleProgressResponse())) |
| |
| # Add a mock bundle processor as if it was running before it's released |
| bundle_processor_cache.active_bundle_processors['instruction_id'] = ( |
| 'descriptor_id', bundle_processor) |
| bundle_processor_cache.release('instruction_id') |
| self.assertEqual( |
| worker.do_instruction(split_request), |
| beam_fn_api_pb2.InstructionResponse( |
| instruction_id='progress_instruction_id', |
| process_bundle_progress=beam_fn_api_pb2. |
| ProcessBundleProgressResponse())) |
| |
| def test_failed_bundle_processor_returns_failed_progress_response(self): |
| bundle_processor = mock.MagicMock() |
| bundle_processor_cache = BundleProcessorCache(None, None, None, {}) |
| bundle_processor_cache.activate('instruction_id') |
| worker = SdkWorker(bundle_processor_cache) |
| |
| # Add a mock bundle processor as if it was running before it's discarded |
| bundle_processor_cache.active_bundle_processors['instruction_id'] = ( |
| 'descriptor_id', bundle_processor) |
| bundle_processor_cache.discard('instruction_id') |
| split_request = beam_fn_api_pb2.InstructionRequest( |
| instruction_id='progress_instruction_id', |
| process_bundle_progress=beam_fn_api_pb2.ProcessBundleProgressRequest( |
| instruction_id='instruction_id')) |
| hc.assert_that( |
| worker.do_instruction(split_request).error, |
| hc.contains_string( |
| 'Bundle processing associated with instruction_id has failed')) |
| |
| def test_inactive_bundle_processor_returns_empty_split_response(self): |
| bundle_processor = mock.MagicMock() |
| bundle_processor_cache = BundleProcessorCache(None, None, None, {}) |
| bundle_processor_cache.activate('instruction_id') |
| worker = SdkWorker(bundle_processor_cache) |
| split_request = beam_fn_api_pb2.InstructionRequest( |
| instruction_id='split_instruction_id', |
| process_bundle_split=beam_fn_api_pb2.ProcessBundleSplitRequest( |
| instruction_id='instruction_id')) |
| self.assertEqual( |
| worker.do_instruction(split_request), |
| beam_fn_api_pb2.InstructionResponse( |
| instruction_id='split_instruction_id', |
| process_bundle_split=beam_fn_api_pb2.ProcessBundleSplitResponse())) |
| |
| # Add a mock bundle processor as if it was running before it's released |
| bundle_processor_cache.active_bundle_processors['instruction_id'] = ( |
| 'descriptor_id', bundle_processor) |
| bundle_processor_cache.release('instruction_id') |
| self.assertEqual( |
| worker.do_instruction(split_request), |
| beam_fn_api_pb2.InstructionResponse( |
| instruction_id='split_instruction_id', |
| process_bundle_split=beam_fn_api_pb2.ProcessBundleSplitResponse())) |
| |
| def get_responses(self, instruction_requests, data_sampler=None): |
| """Evaluates and returns {id: InstructionResponse} for the requests.""" |
| test_controller = BeamFnControlServicer(instruction_requests) |
| |
| server = grpc.server(thread_pool_executor.shared_unbounded_instance()) |
| beam_fn_api_pb2_grpc.add_BeamFnControlServicer_to_server( |
| test_controller, server) |
| test_port = server.add_insecure_port("[::]:0") |
| server.start() |
| |
| harness = sdk_worker.SdkHarness( |
| "localhost:%s" % test_port, |
| state_cache_size=100, |
| data_sampler=data_sampler) |
| harness.run() |
| return test_controller.responses |
| |
| def test_harness_monitoring_infos_and_metadata(self): |
| # Clear the process wide metric container. |
| MetricsEnvironment.process_wide_container().metrics = {} |
| # Create a process_wide metric. |
| urn = 'my.custom.urn' |
| labels = {'key': 'value'} |
| InternalMetrics.counter(urn=urn, labels=labels, process_wide=True).inc(10) |
| |
| harness_monitoring_infos_request = beam_fn_api_pb2.InstructionRequest( |
| instruction_id="monitoring_infos", |
| harness_monitoring_infos=beam_fn_api_pb2.HarnessMonitoringInfosRequest( |
| )) |
| |
| responses = self.get_responses([harness_monitoring_infos_request]) |
| |
| expected_monitoring_info = monitoring_infos.int64_counter( |
| urn, 10, labels=labels) |
| monitoring_data = ( |
| responses['monitoring_infos'].harness_monitoring_infos.monitoring_data) |
| |
| # Request the full MonitoringInfo metadata for the returned short_ids. |
| short_ids = list(monitoring_data.keys()) |
| monitoring_infos_metadata_request = beam_fn_api_pb2.InstructionRequest( |
| instruction_id="monitoring_infos_metadata", |
| monitoring_infos=beam_fn_api_pb2.MonitoringInfosMetadataRequest( |
| monitoring_info_id=short_ids)) |
| |
| responses = self.get_responses([monitoring_infos_metadata_request]) |
| |
| # Request the full MonitoringInfo metadata to be returned now. |
| expected_monitoring_info.ClearField("payload") |
| |
| # Verify that one of the returned monitoring infos is our expected |
| # monitoring info. |
| short_id_to_mi = ( |
| responses['monitoring_infos_metadata'].monitoring_infos.monitoring_info) |
| found = False |
| for mi in short_id_to_mi.values(): |
| # Clear the timestamp before comparing |
| mi.ClearField("start_time") |
| if mi == expected_monitoring_info: |
| found = True |
| self.assertTrue(found, str(responses['monitoring_infos_metadata'])) |
| |
| def test_failed_bundle_processor_returns_failed_split_response(self): |
| bundle_processor = mock.MagicMock() |
| bundle_processor_cache = BundleProcessorCache(None, None, None, {}) |
| bundle_processor_cache.activate('instruction_id') |
| worker = SdkWorker(bundle_processor_cache) |
| |
| # Add a mock bundle processor as if it was running before it's discarded |
| bundle_processor_cache.active_bundle_processors['instruction_id'] = ( |
| 'descriptor_id', bundle_processor) |
| bundle_processor_cache.discard('instruction_id') |
| split_request = beam_fn_api_pb2.InstructionRequest( |
| instruction_id='split_instruction_id', |
| process_bundle_split=beam_fn_api_pb2.ProcessBundleSplitRequest( |
| instruction_id='instruction_id')) |
| hc.assert_that( |
| worker.do_instruction(split_request).error, |
| hc.contains_string( |
| 'Bundle processing associated with instruction_id has failed')) |
| |
| def test_data_sampling_response(self): |
| # Create a data sampler with some fake sampled data. This data will be seen |
| # in the sample response. |
| coder = FastPrimitivesCoder() |
| |
| class FakeDataSampler: |
| def samples(self, pcollection_ids): |
| return beam_fn_api_pb2.SampleDataResponse( |
| element_samples={ |
| 'pcoll_id_1': beam_fn_api_pb2.SampleDataResponse.ElementList( |
| elements=[ |
| beam_fn_api_pb2.SampledElement( |
| element=coder.encode_nested('a')) |
| ]), |
| 'pcoll_id_2': beam_fn_api_pb2.SampleDataResponse.ElementList( |
| elements=[ |
| beam_fn_api_pb2.SampledElement( |
| element=coder.encode_nested('b')) |
| ]) |
| }) |
| |
| def stop(self): |
| pass |
| |
| data_sampler = FakeDataSampler() |
| |
| # Create and send the fake reponse. The SdkHarness should query the |
| # DataSampler and fill out the sample response. |
| sample_request = beam_fn_api_pb2.InstructionRequest( |
| instruction_id='sample_request', |
| sample_data=beam_fn_api_pb2.SampleDataRequest( |
| pcollection_ids=['pcoll_id_1', 'pcoll_id_2'])) |
| responses = self.get_responses([sample_request], data_sampler) |
| self.assertEqual(len(responses), 1) |
| |
| # Get and assert the correct response. |
| response = responses['sample_request'] |
| expected_response = beam_fn_api_pb2.InstructionResponse( |
| instruction_id='sample_request', |
| sample_data=beam_fn_api_pb2.SampleDataResponse( |
| element_samples={ |
| 'pcoll_id_1': beam_fn_api_pb2.SampleDataResponse.ElementList( |
| elements=[ |
| beam_fn_api_pb2.SampledElement( |
| element=coder.encode_nested('a')) |
| ]), |
| 'pcoll_id_2': beam_fn_api_pb2.SampleDataResponse.ElementList( |
| elements=[ |
| beam_fn_api_pb2.SampledElement( |
| element=coder.encode_nested('b')) |
| ]) |
| })) |
| |
| self.assertEqual(response, expected_response) |
| |
| |
| class CachingStateHandlerTest(unittest.TestCase): |
| def test_caching(self): |
| |
| coder = VarIntCoder() |
| coder_impl = coder.get_impl() |
| |
| class FakeUnderlyingState(object): |
| """Simply returns an incremented counter as the state "value." |
| """ |
| def set_counter(self, n): |
| self._counter = n |
| |
| def get_raw(self, *args): |
| self._counter += 1 |
| return coder.encode(self._counter), None |
| |
| @contextlib.contextmanager |
| def process_instruction_id(self, bundle_id): |
| yield |
| |
| underlying_state = FakeUnderlyingState() |
| state_cache = statecache.StateCache(100 << 20) |
| caching_state_hander = GlobalCachingStateHandler( |
| state_cache, underlying_state) |
| |
| state1 = beam_fn_api_pb2.StateKey( |
| bag_user_state=beam_fn_api_pb2.StateKey.BagUserState( |
| user_state_id='state1')) |
| state2 = beam_fn_api_pb2.StateKey( |
| bag_user_state=beam_fn_api_pb2.StateKey.BagUserState( |
| user_state_id='state2')) |
| side1 = beam_fn_api_pb2.StateKey( |
| multimap_side_input=beam_fn_api_pb2.StateKey.MultimapSideInput( |
| transform_id='transform', side_input_id='side1')) |
| side2 = beam_fn_api_pb2.StateKey( |
| iterable_side_input=beam_fn_api_pb2.StateKey.IterableSideInput( |
| transform_id='transform', side_input_id='side2')) |
| |
| state_token1 = beam_fn_api_pb2.ProcessBundleRequest.CacheToken( |
| token=b'state_token1', |
| user_state=beam_fn_api_pb2.ProcessBundleRequest.CacheToken.UserState()) |
| state_token2 = beam_fn_api_pb2.ProcessBundleRequest.CacheToken( |
| token=b'state_token2', |
| user_state=beam_fn_api_pb2.ProcessBundleRequest.CacheToken.UserState()) |
| side1_token1 = beam_fn_api_pb2.ProcessBundleRequest.CacheToken( |
| token=b'side1_token1', |
| side_input=beam_fn_api_pb2.ProcessBundleRequest.CacheToken.SideInput( |
| transform_id='transform', side_input_id='side1')) |
| side1_token2 = beam_fn_api_pb2.ProcessBundleRequest.CacheToken( |
| token=b'side1_token2', |
| side_input=beam_fn_api_pb2.ProcessBundleRequest.CacheToken.SideInput( |
| transform_id='transform', side_input_id='side1')) |
| |
| def get_as_list(key): |
| return list(caching_state_hander.blocking_get(key, coder_impl)) |
| |
| underlying_state.set_counter(100) |
| with caching_state_hander.process_instruction_id('bundle1', []): |
| self.assertEqual(get_as_list(state1), [101]) # uncached |
| self.assertEqual(get_as_list(state2), [102]) # uncached |
| self.assertEqual(get_as_list(state1), [101]) # cached on bundle |
| self.assertEqual(get_as_list(side1), [103]) # uncached |
| self.assertEqual(get_as_list(side2), [104]) # uncached |
| |
| underlying_state.set_counter(200) |
| with caching_state_hander.process_instruction_id( |
| 'bundle2', [state_token1, side1_token1]): |
| self.assertEqual(get_as_list(state1), [201]) # uncached |
| self.assertEqual(get_as_list(state2), [202]) # uncached |
| self.assertEqual(get_as_list(state1), [201]) # cached on state token1 |
| self.assertEqual(get_as_list(side1), [203]) # uncached |
| self.assertEqual(get_as_list(side1), [203]) # cached on side1_token1 |
| self.assertEqual(get_as_list(side2), [204]) # uncached |
| self.assertEqual(get_as_list(side2), [204]) # cached on bundle |
| |
| underlying_state.set_counter(300) |
| with caching_state_hander.process_instruction_id( |
| 'bundle3', [state_token1, side1_token1]): |
| self.assertEqual(get_as_list(state1), [201]) # cached on state token1 |
| self.assertEqual(get_as_list(state2), [202]) # cached on state token1 |
| self.assertEqual(get_as_list(state1), [201]) # cached on state token1 |
| self.assertEqual(get_as_list(side1), [203]) # cached on side1_token1 |
| self.assertEqual(get_as_list(side1), [203]) # cached on side1_token1 |
| self.assertEqual(get_as_list(side2), [301]) # uncached |
| self.assertEqual(get_as_list(side2), [301]) # cached on bundle |
| |
| underlying_state.set_counter(400) |
| with caching_state_hander.process_instruction_id( |
| 'bundle4', [state_token2, side1_token1]): |
| self.assertEqual(get_as_list(state1), [401]) # uncached |
| self.assertEqual(get_as_list(state2), [402]) # uncached |
| self.assertEqual(get_as_list(state1), [401]) # cached on state token2 |
| self.assertEqual(get_as_list(side1), [203]) # cached on side1_token1 |
| self.assertEqual(get_as_list(side1), [203]) # cached on side1_token1 |
| self.assertEqual(get_as_list(side2), [403]) # uncached |
| self.assertEqual(get_as_list(side2), [403]) # cached on bundle |
| |
| underlying_state.set_counter(500) |
| with caching_state_hander.process_instruction_id( |
| 'bundle5', [state_token2, side1_token2]): |
| self.assertEqual(get_as_list(state1), [401]) # cached on state token2 |
| self.assertEqual(get_as_list(state2), [402]) # cached on state token2 |
| self.assertEqual(get_as_list(state1), [401]) # cached on state token2 |
| self.assertEqual(get_as_list(side1), [501]) # uncached |
| self.assertEqual(get_as_list(side1), [501]) # cached on side1_token2 |
| self.assertEqual(get_as_list(side2), [502]) # uncached |
| self.assertEqual(get_as_list(side2), [502]) # cached on bundle |
| |
| class UnderlyingStateHandler(object): |
| """Simply returns an incremented counter as the state "value." |
| """ |
| def __init__(self): |
| self._encoded_values = [] |
| self._continuations = False |
| |
| def set_value(self, value, coder): |
| self._encoded_values = [coder.encode(value)] |
| |
| def set_values(self, values, coder): |
| self._encoded_values = [coder.encode(value) for value in values] |
| |
| def set_continuations(self, continuations): |
| self._continuations = continuations |
| |
| def get_raw(self, _state_key, continuation_token=None): |
| if self._continuations and len(self._encoded_values) > 0: |
| if not continuation_token: |
| continuation_token = '0' |
| idx = int(continuation_token) |
| next_token = str(idx + |
| 1) if idx + 1 < len(self._encoded_values) else None |
| return self._encoded_values[idx], next_token |
| else: |
| return b''.join(self._encoded_values), None |
| |
| def append_raw(self, _key, bytes): |
| self._encoded_values.append(bytes) |
| |
| def clear(self, *args): |
| self._encoded_values = [] |
| |
| @contextlib.contextmanager |
| def process_instruction_id(self, bundle_id): |
| yield |
| |
| def test_append_clear_with_preexisting_state(self): |
| state = beam_fn_api_pb2.StateKey( |
| bag_user_state=beam_fn_api_pb2.StateKey.BagUserState( |
| user_state_id='state1')) |
| |
| cache_token = beam_fn_api_pb2.ProcessBundleRequest.CacheToken( |
| token=b'state_token1', |
| user_state=beam_fn_api_pb2.ProcessBundleRequest.CacheToken.UserState()) |
| |
| coder = VarIntCoder() |
| |
| underlying_state_handler = self.UnderlyingStateHandler() |
| state_cache = statecache.StateCache(100 << 20) |
| handler = GlobalCachingStateHandler(state_cache, underlying_state_handler) |
| |
| def get(): |
| return handler.blocking_get(state, coder.get_impl()) |
| |
| def append(iterable): |
| handler.extend(state, coder.get_impl(), iterable) |
| |
| def clear(): |
| handler.clear(state) |
| |
| # Initialize state |
| underlying_state_handler.set_value(42, coder) |
| with handler.process_instruction_id('bundle', [cache_token]): |
| # Append without reading beforehand |
| append([43]) |
| self.assertEqual(get(), [42, 43]) |
| clear() |
| self.assertEqual(get(), []) |
| append([44, 45]) |
| self.assertEqual(get(), [44, 45]) |
| append((46, 47)) |
| self.assertEqual(get(), [44, 45, 46, 47]) |
| clear() |
| append(range(1000)) |
| self.assertEqual(get(), list(range(1000))) |
| |
| def test_continuation_token(self): |
| underlying_state_handler = self.UnderlyingStateHandler() |
| state_cache = statecache.StateCache(100 << 20) |
| handler = GlobalCachingStateHandler(state_cache, underlying_state_handler) |
| |
| coder = VarIntCoder() |
| |
| state = beam_fn_api_pb2.StateKey( |
| bag_user_state=beam_fn_api_pb2.StateKey.BagUserState( |
| user_state_id='state1')) |
| |
| cache_token = beam_fn_api_pb2.ProcessBundleRequest.CacheToken( |
| token=b'state_token1', |
| user_state=beam_fn_api_pb2.ProcessBundleRequest.CacheToken.UserState()) |
| |
| def get(materialize=True): |
| result = handler.blocking_get(state, coder.get_impl()) |
| return list(result) if materialize else result |
| |
| def get_type(): |
| return type(get(materialize=False)) |
| |
| def append(*values): |
| handler.extend(state, coder.get_impl(), values) |
| |
| def clear(): |
| handler.clear(state) |
| |
| underlying_state_handler.set_continuations(True) |
| underlying_state_handler.set_values([45, 46, 47], coder) |
| with handler.process_instruction_id('bundle', [cache_token]): |
| self.assertEqual( |
| get_type(), GlobalCachingStateHandler.ContinuationIterable) |
| self.assertEqual(get(), [45, 46, 47]) |
| append(48, 49) |
| self.assertEqual( |
| get_type(), GlobalCachingStateHandler.ContinuationIterable) |
| self.assertEqual(get(), [45, 46, 47, 48, 49]) |
| clear() |
| self.assertEqual(get_type(), list) |
| self.assertEqual(get(), []) |
| append(1) |
| self.assertEqual(get(), [1]) |
| append(2, 3) |
| self.assertEqual(get(), [1, 2, 3]) |
| clear() |
| for i in range(1000): |
| append(i) |
| self.assertEqual(get_type(), list) |
| self.assertEqual(get(), [i for i in range(1000)]) |
| |
| |
| class ShortIdCacheTest(unittest.TestCase): |
| def testShortIdAssignment(self): |
| TestCase = namedtuple('TestCase', ['expected_short_id', 'info']) |
| test_cases = [ |
| TestCase(*args) for args in [ |
| ( |
| "1", |
| metrics_pb2.MonitoringInfo( |
| urn="beam:metric:user:distribution_int64:v1", |
| type="beam:metrics:distribution_int64:v1")), |
| ( |
| "2", |
| metrics_pb2.MonitoringInfo( |
| urn="beam:metric:element_count:v1", |
| type="beam:metrics:sum_int64:v1")), |
| ( |
| "3", |
| metrics_pb2.MonitoringInfo( |
| urn="beam:metric:ptransform_progress:completed:v1", |
| type="beam:metrics:progress:v1")), |
| ( |
| "4", |
| metrics_pb2.MonitoringInfo( |
| urn="beam:metric:user:distribution_double:v1", |
| type="beam:metrics:distribution_double:v1")), |
| ( |
| "5", |
| metrics_pb2.MonitoringInfo( |
| urn="TestingSentinelUrn", type="TestingSentinelType")), |
| ( |
| "6", |
| metrics_pb2.MonitoringInfo( |
| urn= |
| "beam:metric:pardo_execution_time:finish_bundle_msecs:v1", |
| type="beam:metrics:sum_int64:v1")), |
| # This case and the next one validates that different labels |
| # with the same urn are in fact assigned different short ids. |
| ( |
| "7", |
| metrics_pb2.MonitoringInfo( |
| urn="beam:metric:user:sum_int64:v1", |
| type="beam:metrics:sum_int64:v1", |
| labels={ |
| "PTRANSFORM": "myT", |
| "NAMESPACE": "harness", |
| "NAME": "metricNumber7" |
| })), |
| ( |
| "8", |
| metrics_pb2.MonitoringInfo( |
| urn="beam:metric:user:sum_int64:v1", |
| type="beam:metrics:sum_int64:v1", |
| labels={ |
| "PTRANSFORM": "myT", |
| "NAMESPACE": "harness", |
| "NAME": "metricNumber8" |
| })), |
| ( |
| "9", |
| metrics_pb2.MonitoringInfo( |
| urn="beam:metric:user:top_n_double:v1", |
| type="beam:metrics:top_n_double:v1", |
| labels={ |
| "PTRANSFORM": "myT", |
| "NAMESPACE": "harness", |
| "NAME": "metricNumber7" |
| })), |
| ( |
| "a", |
| metrics_pb2.MonitoringInfo( |
| urn="beam:metric:element_count:v1", |
| type="beam:metrics:sum_int64:v1", |
| labels={"PCOLLECTION": "myPCol"})), |
| # validate payload is ignored for shortId assignment |
| ( |
| "3", |
| metrics_pb2.MonitoringInfo( |
| urn="beam:metric:ptransform_progress:completed:v1", |
| type="beam:metrics:progress:v1", |
| payload=b"this is ignored!")) |
| ] |
| ] |
| |
| cache = sdk_worker.ShortIdCache() |
| |
| for case in test_cases: |
| self.assertEqual( |
| case.expected_short_id, |
| cache.get_short_id(case.info), |
| "Got incorrect short id for monitoring info:\n%s" % case.info) |
| |
| # Retrieve all of the monitoring infos by short id, and verify that the |
| # metadata (everything but the payload) matches the originals |
| actual_recovered_infos = cache.get_infos( |
| case.expected_short_id for case in test_cases).values() |
| for recoveredInfo, case in zip(actual_recovered_infos, test_cases): |
| self.assertEqual( |
| monitoringInfoMetadata(case.info), |
| monitoringInfoMetadata(recoveredInfo)) |
| |
| # Retrieve short ids one more time in reverse |
| for case in reversed(test_cases): |
| self.assertEqual( |
| case.expected_short_id, |
| cache.get_short_id(case.info), |
| "Got incorrect short id on second retrieval for monitoring info:\n%s" |
| % case.info) |
| |
| |
| def monitoringInfoMetadata(info): |
| return { |
| descriptor.name: value |
| for descriptor, |
| value in info.ListFields() if not descriptor.name == "payload" |
| } |
| |
| |
| if __name__ == "__main__": |
| logging.getLogger().setLevel(logging.INFO) |
| unittest.main() |
