def testRunExecutorWithBeamPipelineArgs(self):
executor_sepc = text_format.Parse(
"""
class_path: "tfx.orchestration.portable.python_executor_operator_test.ValidateBeamPipelineArgsExecutor"
extra_flags: "--runner=DirectRunner"
""", executable_spec_pb2.PythonClassExecutableSpec())
operator = python_executor_operator.PythonExecutorOperator(executor_sepc)
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 testRunExecutor_with_InplaceUpdateExecutor(self):
executor_sepc = text_format.Parse(
"""
class_path: "tfx.orchestration.portable.python_executor_operator_test.InplaceUpdateExecutor"
""", executable_spec_pb2.PythonClassExecutableSpec())
operator = python_executor_operator.PythonExecutorOperator(
executor_sepc)
input_dict = {'input_key': [standard_artifacts.Examples()]}
output_dict = {'output_key': [standard_artifacts.Model()]}
exec_properties = {
'string': 'value',
'int': 1,
'float': 0.0,
# This should not happen on production and will be
# dropped.
'proto': execution_result_pb2.ExecutorOutput()
}
stateful_working_dir = os.path.join(self.tmp_dir,
'stateful_working_dir')
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_dict,
output_dict=output_dict,
exec_properties=exec_properties,
stateful_working_dir=stateful_working_dir,
execution_output_uri=executor_output_uri))
self.assertProtoPartiallyEquals(
"""
output_artifacts {
key: "output_key"
value {
artifacts {
custom_properties {
key: "name"
value {
string_value: "my_model"
}
}
}
}
}""", executor_output)
def testRunExecutor_with_InprocessExecutor(self):
executor_sepc = text_format.Parse(
"""
class_path: "tfx.orchestration.portable.python_executor_operator_test.InprocessExecutor"
""", executable_spec_pb2.PythonClassExecutableSpec())
operator = python_executor_operator.PythonExecutorOperator(
executor_sepc)
input_dict = {'input_key': [standard_artifacts.Examples()]}
output_dict = {'output_key': [standard_artifacts.Model()]}
exec_properties = {'key': 'value'}
executor_output = operator.run_executor(
self._get_execution_info(input_dict, output_dict, exec_properties))
self.assertProtoPartiallyEquals(
"""
output_artifacts {
key: "output_key"
value {
artifacts {
}
}
}""", executor_output)
def main(_):
flags.mark_flag_as_required('tfx_execution_info_b64')
flags.mark_flag_as_required('tfx_python_class_executor_spec_b64')
execution_info = python_executor_binary_utils.deserialize_execution_info(
FLAGS.tfx_execution_info_b64)
python_class_executor_spec = executable_spec_pb2.PythonClassExecutableSpec.FromString(
base64.b64decode(FLAGS.tfx_python_class_executor_spec_b64))
logging.info('execution_info = %s\n',
text_format.MessageToString(execution_info))
logging.info('python_class_executor_spec = %s\n',
text_format.MessageToString(python_class_executor_spec))
operator = python_executor_operator.PythonExecutorOperator(
python_class_executor_spec)
run_result = operator.run_executor(execution_info)
if run_result:
with tf.io.gfile.GFile(execution_info.executor_output_uri, 'wb') as f:
f.write(run_result.SerializeToString())
def testRunExecutor_with_InprocessExecutor(self):
executor_sepc = text_format.Parse(
"""
class_path: "tfx.orchestration.portable.python_executor_operator_test.InprocessExecutor"
""", local_deployment_config_pb2.ExecutableSpec.PythonClassExecutableSpec(
))
operator = python_executor_operator.PythonExecutorOperator(
executor_sepc)
input_dict = {'input_key': [standard_artifacts.Examples()]}
output_dict = {'output_key': [standard_artifacts.Model()]}
exec_properties = {'key': 'value'}
stateful_working_dir = os.path.join(self.tmp_dir,
'stateful_working_dir')
executor_output_uri = os.path.join(self.tmp_dir, 'executor_output')
executor_output = operator.run_executor(
base_executor_operator.ExecutionInfo(
input_dict=input_dict,
output_dict=output_dict,
exec_properties=exec_properties,
stateful_working_dir=stateful_working_dir,
executor_output_uri=executor_output_uri))
self.assertProtoPartiallyEquals(
"""
execution_properties {
key: "key"
value {
string_value: "value"
}
}
output_artifacts {
key: "output_key"
value {
artifacts {
}
}
}""", executor_output)
def _run_executor(
executable_spec: executable_spec_pb2.PythonClassExecutableSpec,
execution_info: data_types.ExecutionInfo
) -> execution_result_pb2.ExecutorOutput:
operator = python_executor_operator.PythonExecutorOperator(executable_spec)
return operator.run_executor(execution_info)
