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: