def test_get_ir_value(self):
self.assertDictEqual(
json_format.MessageToDict(pipeline_spec_pb2.Value(int_value=42)),
json_format.MessageToDict(dsl_utils.get_value(42)))
self.assertDictEqual(
json_format.MessageToDict(pipeline_spec_pb2.Value(double_value=12.2)),
json_format.MessageToDict(dsl_utils.get_value(12.2)))
self.assertDictEqual(
json_format.MessageToDict(
pipeline_spec_pb2.Value(string_value='hello world')),
json_format.MessageToDict(dsl_utils.get_value('hello world')))
with self.assertRaisesRegex(TypeError, 'Got unexpected type'):
dsl_utils.get_value(_DummyClass())
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_value(value: Union[str, int, float]) -> pipeline_spec_pb2.Value:
"""Gets pipeline value proto from Python value."""
result = pipeline_spec_pb2.Value()
if isinstance(value, str):
result.string_value = value
elif isinstance(value, int):
result.int_value = value
elif isinstance(value, float):
result.double_value = value
else:
raise TypeError(
'Got unexpected type %s for value %s. Currently only support str, int '
'and float.' % (type(value), value))
return result
def _get_value(value: Union[int, float, str]) -> pipeline_spec_pb2.Value:
assert value is not None, 'None values should be filterd out.'
result = pipeline_spec_pb2.Value()
if isinstance(value, int):
result.int_value = value
elif isinstance(value, float):
result.double_value = value
elif isinstance(value, str):
result.string_value = value
else:
raise TypeError('Got unknown type of value: {}'.format(value))
return result
def _get_pipeline_value(
value: Union[int, float, str]) -> Optional[pipeline_spec_pb2.Value]:
"""Converts Python primitive value to pipeline value pb."""
if value is None:
return None
result = pipeline_spec_pb2.Value()
if isinstance(value, int):
result.int_value = value
elif isinstance(value, float):
result.double_value = value
elif isinstance(value, str):
result.string_value = value
else:
raise TypeError('Got unknown type of value: {}'.format(value))
return result
def get_kubeflow_value(
tfx_value: Union[int, float, str]) -> 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):
result.string_value = tfx_value
else:
raise TypeError('Got unknown type of value: {}'.format(tfx_value))
return result
def _get_value(
value: Optional[Union[int, float, str]]
) -> Optional[pipeline_spec_pb2.Value]:
if value is None:
return None
result = pipeline_spec_pb2.Value()
if isinstance(value, int):
result.int_value = value
elif isinstance(value, float):
result.double_value = value
elif isinstance(value, str):
result.string_value = value
else:
raise TypeError('Got unknown type of value: {}'.format(value))
return result
def _get_value(
param: _pipeline_param.PipelineParam) -> pipeline_spec_pb2.Value:
assert param.value is not None, 'None values should be filterd out.'
result = pipeline_spec_pb2.Value()
# TODO(chensun): remove defaulting to 'String' for None param_type once we
# fix importer behavior.
param_type = type_utils.get_parameter_type(param.param_type
or 'String')
if param_type == pipeline_spec_pb2.PrimitiveType.INT:
result.int_value = int(param.value)
elif param_type == pipeline_spec_pb2.PrimitiveType.DOUBLE:
result.double_value = float(param.value)
elif param_type == pipeline_spec_pb2.PrimitiveType.STRING:
result.string_value = str(param.value)
else:
# For every other type, defaults to 'String'.
# TODO(chensun): remove this default behavior once we migrate from
# `pipeline_spec_pb2.Value` to `protobuf.Value`.
result.string_value = str(param.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:
