def test_source_split(self):
source = RangeSource(0, 100)
expected_splits = list(source.split(30))
worker = sdk_harness.SdkWorker(
None, data_plane.GrpcClientDataChannelFactory())
worker.register(
beam_fn_api_pb2.RegisterRequest(
process_bundle_descriptor=[beam_fn_api_pb2.ProcessBundleDescriptor(
primitive_transform=[beam_fn_api_pb2.PrimitiveTransform(
function_spec=sdk_harness.serialize_and_pack_py_fn(
SourceBundle(1.0, source, None, None),
sdk_harness.PYTHON_SOURCE_URN,
id="src"))])]))
split_response = worker.initial_source_split(
beam_fn_api_pb2.InitialSourceSplitRequest(
desired_bundle_size_bytes=30,
source_reference="src"))
self.assertEqual(
expected_splits,
[sdk_harness.unpack_and_deserialize_py_fn(s.source)
for s in split_response.splits])
self.assertEqual(
[s.weight for s in expected_splits],
[s.relative_size for s in split_response.splits])
def test_source_split(self):
source = RangeSource(0, 100)
expected_splits = list(source.split(30))
worker = sdk_harness.SdkWorker(
None, data_plane.GrpcClientDataChannelFactory())
worker.register(
beam_fn_api_pb2.RegisterRequest(process_bundle_descriptor=[
beam_fn_api_pb2.ProcessBundleDescriptor(primitive_transform=[
beam_fn_api_pb2.PrimitiveTransform(
function_spec=sdk_harness.serialize_and_pack_py_fn(
SourceBundle(1.0, source, None, None),
sdk_harness.PYTHON_SOURCE_URN,
id="src"))
])
]))
split_response = worker.initial_source_split(
beam_fn_api_pb2.InitialSourceSplitRequest(
desired_bundle_size_bytes=30, source_reference="src"))
self.assertEqual(expected_splits, [
sdk_harness.unpack_and_deserialize_py_fn(s.source)
for s in split_response.splits
])
self.assertEqual([s.weight for s in expected_splits],
[s.relative_size for s in split_response.splits])
def test_source_split_via_instruction(self):
source = RangeSource(0, 100)
expected_splits = list(source.split(30))
test_controller = BeamFnControlServicer([
beam_fn_api_pb2.InstructionRequest(
instruction_id="register_request",
register=beam_fn_api_pb2.RegisterRequest(
process_bundle_descriptor=[
beam_fn_api_pb2.ProcessBundleDescriptor(
primitive_transform=[
beam_fn_api_pb2.PrimitiveTransform(
function_spec=sdk_harness.
serialize_and_pack_py_fn(
SourceBundle(1.0, source, None, None),
sdk_harness.PYTHON_SOURCE_URN,
id="src"))
])
])),
beam_fn_api_pb2.InstructionRequest(
instruction_id="split_request",
initial_source_split=beam_fn_api_pb2.InitialSourceSplitRequest(
desired_bundle_size_bytes=30, source_reference="src"))
])
server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
beam_fn_api_pb2.add_BeamFnControlServicer_to_server(
test_controller, server)
test_port = server.add_insecure_port("[::]:0")
server.start()
channel = grpc.insecure_channel("localhost:%s" % test_port)
harness = sdk_harness.SdkHarness(channel)
harness.run()
split_response = test_controller.responses[
"split_request"].initial_source_split
self.assertEqual(expected_splits, [
sdk_harness.unpack_and_deserialize_py_fn(s.source)
for s in split_response.splits
])
self.assertEqual([s.weight for s in expected_splits],
[s.relative_size for s in split_response.splits])
def test_source_split_via_instruction(self):
source = RangeSource(0, 100)
expected_splits = list(source.split(30))
test_controller = BeamFnControlServicer([
beam_fn_api_pb2.InstructionRequest(
instruction_id="register_request",
register=beam_fn_api_pb2.RegisterRequest(
process_bundle_descriptor=[
beam_fn_api_pb2.ProcessBundleDescriptor(
primitive_transform=[beam_fn_api_pb2.PrimitiveTransform(
function_spec=sdk_harness.serialize_and_pack_py_fn(
SourceBundle(1.0, source, None, None),
sdk_harness.PYTHON_SOURCE_URN,
id="src"))])])),
beam_fn_api_pb2.InstructionRequest(
instruction_id="split_request",
initial_source_split=beam_fn_api_pb2.InitialSourceSplitRequest(
desired_bundle_size_bytes=30,
source_reference="src"))
])
server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
beam_fn_api_pb2.add_BeamFnControlServicer_to_server(test_controller, server)
test_port = server.add_insecure_port("[::]:0")
server.start()
channel = grpc.insecure_channel("localhost:%s" % test_port)
harness = sdk_harness.SdkHarness(channel)
harness.run()
split_response = test_controller.responses[
"split_request"].initial_source_split
self.assertEqual(
expected_splits,
[sdk_harness.unpack_and_deserialize_py_fn(s.source)
for s in split_response.splits])
self.assertEqual(
[s.weight for s in expected_splits],
[s.relative_size for s in split_response.splits])
