Skip to content

fix(types): Resolve 'tfx.' prefixed KFP v2/Vertex AI artifact types to native Python classes#7869

Open
vkarampudi wants to merge 1 commit into
tensorflow:masterfrom
vkarampudi:feature/fix-artifact-type-resolution
Open

fix(types): Resolve 'tfx.' prefixed KFP v2/Vertex AI artifact types to native Python classes#7869
vkarampudi wants to merge 1 commit into
tensorflow:masterfrom
vkarampudi:feature/fix-artifact-type-resolution

Conversation

@vkarampudi

Copy link
Copy Markdown
Collaborator

Description

This PR resolves a major type resolution bug (Issue #7849) where standard TFX artifacts produced under Kubeflow Pipelines (KFP) v2 or Google Cloud Vertex AI Pipelines fail to resolve to their native Python classes.

The Problem

Under KFP v2 / Vertex AI (which uses IR-based pipeline execution), TFX artifacts are registered in MLMD using their YAML schema title which is prefixed with tfx. (e.g. tfx.Examples, tfx.Model).

When a downstream Python step queries MLMD, TFX's get_artifact_type_class() utility attempts to match the MLMD type name against native Python classes. Because native classes have TYPE_NAME attributes without the prefix (e.g. TYPE_NAME = "Examples"), the exact match check fails:

    if (artifact_type.name == candidate_type.name and ...):

Consequently, TFX prints a warning ("Could not find matching artifact class for type...") and generates an ephemeral generic Artifact subclass. This is highly brittle and breaks any downstream code performing class checks (e.g. isinstance(artifact, standard_artifacts.Examples)) or using subclass-specific properties/methods.

The Solution

We normalize artifact_type.name by stripping the tfx. prefix before checking for a match with native TFX classes:

    artifact_type_name = artifact_type.name
    if artifact_type_name.startswith('tfx.'):
      artifact_type_name = artifact_type_name[4:]

Testing

  • Added a new unit test in artifact_utils_test.py (testArtifactTypeRoundTrip) that constructs a mock KFP-prefixed tfx.Examples artifact type and asserts that it is correctly resolved to standard_artifacts.Examples.
  • All artifact_utils_test tests pass successfully in the tfx-py310 environment.

@nithyak0204 nithyak0204 self-assigned this Jun 24, 2026
@nithyak0204 nithyak0204 requested a review from nikelite June 24, 2026 08:30
@nithyak0204 nithyak0204 removed the request for review from nikelite June 24, 2026 09:01
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants