def testGetMetadataValueTypeFailed(self): tfx_value = pipeline_pb2.Value() text_format.Parse( """ runtime_parameter { name: 'rp' }""", tfx_value) with self.assertRaisesRegex(RuntimeError, 'Expecting field_value but got'): common_utils.get_metadata_value_type(tfx_value)
def prepare_execution( metadata_handler: metadata.Metadata, execution_type: metadata_store_pb2.ExecutionType, state: metadata_store_pb2.Execution.State, exec_properties: Optional[Mapping[Text, types.Property]] = None, ) -> metadata_store_pb2.Execution: """Creates an execution proto based on the information provided. Args: metadata_handler: A handler to access MLMD store. execution_type: A metadata_pb2.ExecutionType message describing the type of the execution. state: The state of the execution. exec_properties: Execution properties that need to be attached. Returns: A metadata_store_pb2.Execution message. """ execution = metadata_store_pb2.Execution() execution.last_known_state = state execution.type_id = common_utils.register_type_if_not_exist( metadata_handler, execution_type).id exec_properties = exec_properties or {} # For every execution property, put it in execution.properties if its key is # in execution type schema. Otherwise, put it in execution.custom_properties. for k, v in exec_properties.items(): if (execution_type.properties.get(k) == common_utils.get_metadata_value_type(v)): common_utils.set_metadata_value(execution.properties[k], v) else: common_utils.set_metadata_value(execution.custom_properties[k], v) logging.debug('Prepared EXECUTION:\n %s', execution) return execution
def testGetMetadataValueType(self): tfx_value = pipeline_pb2.Value() text_format.Parse( """ field_value { int_value: 1 }""", tfx_value) self.assertEqual(common_utils.get_metadata_value_type(tfx_value), metadata_store_pb2.INT)
def _generate_context_proto( metadata_handler: metadata.Metadata, context_spec: pipeline_pb2.ContextSpec) -> metadata_store_pb2.Context: """Generates metadata_pb2.Context based on the ContextSpec message. Args: metadata_handler: A handler to access MLMD store. context_spec: A pipeline_pb2.ContextSpec message that instructs registering of a context. Returns: A metadata_store_pb2.Context message. Raises: RuntimeError: When actual property type does not match provided metadata type schema. """ context_type = common_utils.register_type_if_not_exist( metadata_handler, context_spec.type) context_name = common_utils.get_value(context_spec.name) assert isinstance(context_name, Text), 'context name should be string.' context = metadata_store_pb2.Context(type_id=context_type.id, name=context_name) for k, v in context_spec.properties.items(): if k in context_type.properties: actual_property_type = common_utils.get_metadata_value_type(v) if context_type.properties.get(k) == actual_property_type: common_utils.set_metadata_value(context.properties[k], v) else: raise RuntimeError( 'Property type %s different from provided metadata type property type %s for key %s' % (actual_property_type, context_type.properties.get(k), k)) else: common_utils.set_metadata_value(context.custom_properties[k], v) return context
def testGetMetadataValueTypePrimitiveValue(self): self.assertEqual(common_utils.get_metadata_value_type(1), metadata_store_pb2.INT)