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)