function import from `import_utils.import_func_from_source` not working in dataflow job

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

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: 0.27
• Python version: 3.6
• Python dependencies (from pip freeze output): cherry-pick a few relevant ones apache-beam==2.27.0 dill==0.3.1.1

Describe the current behavior We have a custom component that uses import_utils.import_func_from_source from tfx.utils to import a function in a user module for a dataflow job. The code is roughly looks as follows:

predicate_fn = import_utils.import_func_from_source
with self._make_beam_pipeline() as pipeline:
     pipeline | beam.PTransform(predicate_fn)

After we bump to TFX 0.27, this custom component stops working and we see the following error in our dataflow job:

021-05-17 19:27:52.594 EDTError message from worker: Traceback (most recent call last): File 
"/usr/local/lib/python3.6/site-packages/apache_beam/internal/pickler.py", line 290, in loads return dill.loads(s) File 
"/usr/local/lib/python3.6/site-packages/dill/_dill.py", line 275, in loads return load(file, ignore, **kwds) File 
"/usr/local/lib/python3.6/site-packages/dill/_dill.py", line 270, in load return Unpickler(file, ignore=ignore, **kwds).load() File 
"/usr/local/lib/python3.6/site-packages/dill/_dill.py", line 472, in load obj = StockUnpickler.load(self) File 
"/usr/local/lib/python3.6/site-packages/dill/_dill.py", line 462, in find_class return StockUnpickler.find_class(self, module, 
name) ModuleNotFoundError: No module named 'user_module_0' During handling of the above exception, another 
exception occurred: Traceback (most recent call last): File "/usr/local/lib/python3.6/site-
packages/dataflow_worker/batchworker.py", line 649, in do_work work_executor.execute() File "/usr/local/lib/python3.6/site-
packages/dataflow_worker/executor.py", line 179, in execute op.start() File "apache_beam/runners/worker/operations.py", 
line 710, in apache_beam.runners.worker.operations.DoOperation.start File "apache_beam/runners/worker/operations.py", 
line 712, in apache_beam.runners.worker.operations.DoOperation.start File "apache_beam/runners/worker/operations.py",
 line 713, in apache_beam.runners.worker.operations.DoOperation.start File "apache_beam/runners/worker/operations.py", 
line 311, in apache_beam.runners.worker.operations.Operation.start File "apache_beam/runners/worker/operations.py", line 
317, in apache_beam.runners.worker.operations.Operation.start File "apache_beam/runners/worker/operations.py", line 659,
in apache_beam.runners.worker.operations.DoOperation.setup File "apache_beam/runners/worker/operations.py", line 660, 
in apache_beam.runners.worker.operations.DoOperation.setup File "apache_beam/runners/worker/operations.py", line 292, 
in apache_beam.runners.worker.operations.Operation.setup File "apache_beam/runners/worker/operations.py", line 306, in 
apache_beam.runners.worker.operations.Operation.setup File "apache_beam/runners/worker/operations.py", line 799, in 
apache_beam.runners.worker.operations.DoOperation._get_runtime_performance_hints File "/usr/local/lib/python3.6/site-
packages/apache_beam/internal/pickler.py", line 294, in loads return dill.loads(s) File "/usr/local/lib/python3.6/site-
packages/dill/_dill.py", line 275, in loads return load(file, ignore, **kwds) File "/usr/local/lib/python3.6/site-
packages/dill/_dill.py", line 270, in load return Unpickler(file, ignore=ignore, **kwds).load() File "/usr/local/lib/python3.6/site-
packages/dill/_dill.py", line 472, in load obj = StockUnpickler.load(self) File "/usr/local/lib/python3.6/site-
packages/dill/_dill.py", line 462, in find_class return StockUnpickler.find_class(self, module, name) ModuleNotFoundError: 
No module named 'user_module_0'

Describe the expected behavior The imported function should work in dataflow job

Name of your Organization (Optional) Spotify

[entrpn]

Update:

If someone runs into this issue with Transform component, the fix for me was to set force_tf_compat_v1=True in Transform object.

---

Same issue here with tfx 1.0.0 and dataflow when using Transforms component in a vertexai pipeline.

Vertex settings:

Container: gcr.io/tfx-oss-public/tfx:1.0.0

Dataflow settings:

tfx_executor : tfx-components-transform-executor-executor tfx_py_version : 3-7 tfx_version : 1-0-0 SDK version : Apache Beam Python 3.7 SDK 2.31.0

[charlesccychen]

Hi @axelborja, we are having some more extensive discussion of this issue in https://github.com/tensorflow/tfx/issues/3761. Can we get a better sense of your environment? Which runner are you using, and are you disabling the user code wheel packaging logic because you’ve encountered issues?

Re @daikeshi: this may be a related but different problem from what @axelborja mentions. We discussed this briefly before, but it would be great to understand the type of user code avoids this in 0.22.2 but broke recently, since I don’t recall us doing any special handling for this before.