Trouble running custom TFX component on Vertex AI (Kubeflow Runner) — Executor import error

[crbl1122]

If the bug is related to a specific library below, please raise an issue in the
respective repo directly:

TensorFlow Data Validation Repo

TensorFlow Model Analysis Repo

TensorFlow Transform Repo

TensorFlow Serving Repo

System information

• Have I specified the code to reproduce the issue (Yes, No):
• Environment in which the code is executed (e.g., Local(Linux/MacOS/Windows),
  Interactive Notebook, Google Cloud, etc):
• TensorFlow version:
• TFX Version:
• Python version:
• Python dependencies (from pip freeze output):

Describe the current behavior

I’m trying to run a simple custom TFX component using the Kubeflow Runner on GCP Vertex AI.

I’ve defined the component in two ways: Using the @component decorator, and as a fully custom component (with a defined Executor and ComponentSpec).

In both cases, pipeline compilation succeeds, but when I run the pipeline, I get the following error:

ImportError: Executor class couldn’t be found in main

I also tried packaging the component into a .tar.gz file and uploading it to Cloud Storage, but the result was the same.

Environment details:

TFX version: 1.15 Runner: KubeflowDagRunner Platform: Vertex AI Pipelines

Has anyone successfully run a custom TFX component on GCP using the Kubeflow runner? If so, how did you structure or package your component so that the Executor is properly found during execution? I couldn't find any documentation or tutorial about this.

Describe the expected behavior

Successfully running a custom TFX component in Vertex AI using Kubeflow Runner

Standalone code to reproduce the issue

Providing a bare minimum test case or step(s) to reproduce the problem will
greatly help us to debug the issue. If possible, please share a link to
Colab/Jupyter/any notebook.

Name of your Organization (Optional)

Other info / logs

Include any logs or source code that would be helpful to diagnose the problem.
If including tracebacks, please include the full traceback. Large logs and files
should be attached.

[pritamdodeja]

Hello, can you share the yaml produced by the pipeline compilation step? Are you able to run your custom component locally using LocalDagRunner? Does the container executing your pipeline have access to the custom tfx component that you've created? My understanding of this process is:

1. You compile your tfx pipeline, which contains the pipeline in a yaml representation conforming to a pipeline spec. This is Kubeflow Dag Runner v2 converting the tfx pipeline to a Kubeflow representation. The pipeline spec contains all the input and output schemas, along with the values you want to pass to your component, and the entry point class.
2. When you submit your job, Kubeflow (Vertex AI pipelines as an executor of Kubeflow pipelines) will take this pipeline spec and deliver these inputs to your executor class, setting up the required resources from a kubernetes perspective, and deal with writing to ML metadata etc., it will also deal with the conversion of this job to the workflow execution backend (argo I believe in the case of Vertex AI pipelines).

So to your original question: running a custom component in Vertex AI pipelines, the question comes down to, at what point in the process is this failing? Given that architecturally, what Kubeflow is doing is quite agnostic of your component itself (this is why the component spec and all the things are so rigorously defined). It is purely transporting your component related information to the execution environment, and interacting with ML Metadata to make sure the information related to the execution is consistent with what actually happened.

[crbl1122]

Hey, when I run the component locally, everything works fine. However, let's first discuss why despite following all the tutorials that demonstrate using the @component decorator for Python function-based components, I still encounter the same error as with a fully customized component.

The pipeline compiles successfully and starts execution, just like the customized component. However, during the component’s execution, it fails midway with that error. As I mentioned above, I also tried packaging the component and installing it using packages_to_install, following the Kubeflow Pipelines style.

[pritamdodeja]

Can you expand your container.args? There you should see a --executor_class_path that should show the path of the class that's going to do the work. In the case of a custom component, it'd look something like tfx.components.my_custom_component.Executor. One key difference between the local execution and the remote execution is that in the remote execution, your pipeline is going through a kfp compilation step, whereas locally (likely) you're running this with something like a LocalDagRunner which doesn't require the conversion.

Not related to the problem you're trying to solve, but your CsvReader should produce tf.train based Examples. They already have beam based components (e.g. CsvExampleGen) that might do at least some of what you're trying to do.

[crbl1122]

I fixed the issue by configuring the correct executor path in the component’s JSON file.

[github-actions]

Are you satisfied with the resolution of your issue?
Yes
No

[pritamdodeja]

It'd be great if you can put a sample of how you did that as a comment or contribute to the documentation. Congrats on solving your issue!