def value_converter( tfx_value: Any) -> Optional[pipeline_pb2.ValueOrRuntimeParameter]: """Converts TFX/MLMD values into Kubeflow pipeline ValueOrRuntimeParameter.""" if tfx_value is None: return None result = pipeline_pb2.ValueOrRuntimeParameter() if isinstance(tfx_value, (int, float, str, Text)): result.constant_value.CopyFrom(get_kubeflow_value(tfx_value)) elif isinstance(tfx_value, (Dict, List)): result.constant_value.CopyFrom( pipeline_pb2.Value(string_value=json.dumps(tfx_value))) elif isinstance(tfx_value, data_types.RuntimeParameter): # Attach the runtime parameter to the context. parameter_utils.attach_parameter(tfx_value) result.runtime_parameter = tfx_value.name elif isinstance(tfx_value, metadata_store_pb2.Value): if tfx_value.WhichOneof('value') == 'int_value': result.constant_value.CopyFrom( pipeline_pb2.Value(int_value=tfx_value.int_value)) elif tfx_value.WhichOneof('value') == 'double_value': result.constant_value.CopyFrom( pipeline_pb2.Value(double_value=tfx_value.double_value)) elif tfx_value.WhichOneof('value') == 'string_value': result.constant_value.CopyFrom( pipeline_pb2.Value(string_value=tfx_value.string_value)) elif isinstance(tfx_value, message.Message): result.constant_value.CopyFrom( pipeline_pb2.Value(string_value=json_format.MessageToJson( message=tfx_value, sort_keys=True))) else: # By default will attempt to encode the object using json_utils.dumps. result.constant_value.CopyFrom( pipeline_pb2.Value(string_value=json_utils.dumps(tfx_value))) return result
def get_kubeflow_value( mlmd_value: metadata_store_pb2.Value) -> pipeline_pb2.Value: result = pipeline_pb2.Value() if not mlmd_value.HasField('value'): return result if mlmd_value.WhichOneof('value') == 'int_value': result.int_value = mlmd_value.int_value elif mlmd_value.WhichOneof('value') == 'double_value': result.double_value = mlmd_value.double_value elif mlmd_value.WhichOneof('value') == 'string_value': result.string_value = mlmd_value.string_value else: raise TypeError('Get unknown type of value: {}'.format(mlmd_value)) return result
def get_kubeflow_value( tfx_value: Union[int, float, str, Text]) -> Optional[pipeline_pb2.Value]: """Converts TFX/MLMD values into Kubeflow pipeline Value proto message.""" if tfx_value is None: return None result = pipeline_pb2.Value() if isinstance(tfx_value, int): result.int_value = tfx_value elif isinstance(tfx_value, float): result.double_value = tfx_value elif isinstance(tfx_value, (str, Text)): result.string_value = tfx_value else: raise TypeError('Got unknown type of value: {}'.format(tfx_value)) return result
from google.protobuf import message _TEST_TWO_STEP_PIPELINE_NAME = 'two-step-pipeline' _TEST_FULL_PIPELINE_NAME = 'full-taxi-pipeline' _TEST_PIPELINE_ROOT = 'path/to/my/root' _TEST_INPUT_DATA = 'path/to/my/data' _TEST_MODULE_FILE_LOCATION = 'path/to/my/module_utils.py' TEST_RUNTIME_CONFIG = pipeline_pb2.PipelineJob.RuntimeConfig( gcs_output_directory=_TEST_PIPELINE_ROOT, parameters={ 'string_param': pipeline_pb2.Value(string_value='test-string'), 'int_param': pipeline_pb2.Value(int_value=42), 'float_param': pipeline_pb2.Value(double_value=3.14) }) _POLLING_INTERVAL_IN_SECONDS = 60 _MAX_JOB_EXECUTION_TIME_IN_SECONDS = 2400 _KUBEFLOW_SUCCEEDED_STATE = 'SUCCEEDED' _KUBEFLOW_RUNNING_STATES = frozenset(('PENDING', 'RUNNING')) # TODO(b/158245564): Reevaluate whether to keep this test helper function def two_step_pipeline() -> tfx_pipeline.Pipeline: