def testExecutableSpecSerialization(self):
python_executable_spec = text_format.Parse(
"""
class_path: 'path_to_my_class'
extra_flags: '--flag=my_flag'
""", executable_spec_pb2.PythonClassExecutableSpec())
python_serialized = python_execution_binary_utils.serialize_executable_spec(
python_executable_spec)
python_rehydrated = python_execution_binary_utils.deserialize_executable_spec(
python_serialized)
self.assertProtoEquals(python_rehydrated, python_executable_spec)
beam_executable_spec = text_format.Parse(
"""
python_executor_spec {
class_path: 'path_to_my_class'
extra_flags: '--flag1=1'
}
beam_pipeline_args: '--arg=my_beam_pipeline_arg'
""", executable_spec_pb2.BeamExecutableSpec())
beam_serialized = python_execution_binary_utils.serialize_executable_spec(
beam_executable_spec)
beam_rehydrated = python_execution_binary_utils.deserialize_executable_spec(
beam_serialized, with_beam=True)
self.assertProtoEquals(beam_rehydrated, beam_executable_spec)
def encode(
self,
component_spec: Optional[types.ComponentSpec] = None) -> message.Message:
result = executable_spec_pb2.BeamExecutableSpec()
result.python_executor_spec.CopyFrom(
super().encode(component_spec=component_spec))
result.beam_pipeline_args.extend(self.beam_pipeline_args)
return result
def testRunExecutorWithBeamPipelineArgs(self):
executor_spec = text_format.Parse(
"""
python_executor_spec: {
class_path: "tfx.orchestration.portable.beam_executor_operator_test.ValidateBeamPipelineArgsExecutor"
}
beam_pipeline_args: "--runner=DirectRunner"
""", executable_spec_pb2.BeamExecutableSpec())
operator = beam_executor_operator.BeamExecutorOperator(executor_spec)
executor_output_uri = os.path.join(self.tmp_dir, 'executor_output')
operator.run_executor(
data_types.ExecutionInfo(
execution_id=1,
input_dict={},
output_dict={},
exec_properties={},
execution_output_uri=executor_output_uri))
def testRunExecutorWithBeamPipelineArgs(self):
executor_spec = text_format.Parse(
"""
python_executor_spec: {
class_path: "tfx.orchestration.portable.beam_executor_operator_test.ValidateBeamPipelineArgsExecutor"
}
beam_pipeline_args: "--runner=DirectRunner"
""", executable_spec_pb2.BeamExecutableSpec())
operator = beam_executor_operator.BeamExecutorOperator(executor_spec)
pipeline_node = pipeline_pb2.PipelineNode(
node_info={'id': 'MyBeamNode'})
pipeline_info = pipeline_pb2.PipelineInfo(id='MyPipeline')
executor_output_uri = os.path.join(self.tmp_dir, 'executor_output')
executor_output = operator.run_executor(
data_types.ExecutionInfo(
execution_id=1,
input_dict={'input_key': [standard_artifacts.Examples()]},
output_dict={'output_key': [standard_artifacts.Model()]},
exec_properties={},
execution_output_uri=executor_output_uri,
pipeline_node=pipeline_node,
pipeline_info=pipeline_info,
pipeline_run_id=99))
self.assertProtoPartiallyEquals(
"""
output_artifacts {
key: "output_key"
value {
artifacts {
custom_properties {
key: "name"
value {
string_value: "MyPipeline.MyBeamNode.my_model"
}
}
}
}
}""", executor_output)