def _parse_raw_artifact()

in tfx/orchestration/kubeflow/v2/container/kubeflow_v2_entrypoint_utils.py [0:0]

• 42 lines of code
• 18 McCabe index (conditional complexity)

def _parse_raw_artifact(
    artifact_pb: pipeline_pb2.RuntimeArtifact,
    name_from_id: MutableMapping[int, str]) -> artifact.Artifact:
  """Parses RuntimeArtifact proto message without artifact_type."""
  # This parser can only reserve what's inside the RuntimeArtifact pb message.

  # Recovers the type information from artifact type schema.
  # TODO(b/170261670): Replace this workaround by a more resilient
  # implementation. Currently custom artifact type can hardly be supported.
  assert artifact_pb.type, 'RuntimeArtifact is expected to have a type.'

  # 1. Import the artifact class from preloaded TFX library.
  type_path = _retrieve_class_path(artifact_pb.type)
  artifact_cls = import_utils.import_class_by_path(type_path)

  # 2. Copy properties and custom properties to the MLMD artifact pb.
  mlmd_artifact = metadata_store_pb2.Artifact()
  # TODO(b/135056715): Change to a unified getter/setter of Artifact type
  # once it's ready.
  if artifact_pb.name:
    # TODO(b/169583143): Remove this workaround when TFX migrates to use
    # str-typed id/name to identify artifacts.
    # Convert and populate the MLMD artifact ID.
    mlmd_artifact.id = _get_hashed_id(artifact_pb.name, name_from_id)

  mlmd_artifact.uri = artifact_pb.uri
  for k, v in artifact_pb.properties.items():
    mlmd_artifact.properties[k].CopyFrom(compiler_utils.get_mlmd_value(v))

  for k, v in artifact_pb.custom_properties.items():
    mlmd_artifact.custom_properties[k].CopyFrom(
        compiler_utils.get_mlmd_value(v))

  # Translate metadata items into properties and custom properties.
  mlmd_artifact_type = artifact_cls().artifact_type
  metadata_dict = json_format.MessageToDict(artifact_pb.metadata)
  for k, v in metadata_dict.items():
    if k in mlmd_artifact_type.properties:
      property_type = mlmd_artifact_type.properties[k]
      if property_type == metadata_store_pb2.INT and isinstance(v, float):
        mlmd_artifact.properties[k].int_value = int(v)
        continue
      elif property_type == metadata_store_pb2.DOUBLE and isinstance(v, float):
        mlmd_artifact.properties[k].double_value = v
        continue
      elif property_type == metadata_store_pb2.STRING and isinstance(v, str):
        mlmd_artifact.properties[k].string_value = v
        continue
      elif property_type == metadata_store_pb2.STRUCT:
        mlmd_artifact.properties[k].struct_value.CopyFrom(
            artifact._encode_struct_value(v))  # pylint: disable=protected-access
        continue
      # We fell through, which means the property doesn't actually fit the
      # schema. Therefore, we treat it as a custom property.

    # First, we drop the custom property prefix if we had to drop it because
    # of a property name conflict.
    if k.startswith(artifact.CUSTOM_PROPERTIES_PREFIX):
      stripped_k = k[len(artifact.CUSTOM_PROPERTIES_PREFIX):]
      if stripped_k in mlmd_artifact_type.properties:
        k = stripped_k
    mlmd_artifact.custom_properties[k].struct_value.CopyFrom(
        artifact._encode_struct_value(v))  # pylint: disable=protected-access

  # 3. Instantiate the artifact Python object.
  result = artifact_cls()
  result.set_mlmd_artifact(mlmd_artifact)

  return result