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: