Repository: tensorflow/tfx

Issue:
TFX taxi template KubeflowDagRunner Transform throws "ModuleNotFoundError"
I built a TFX pipeline based on the taxi template. All components run okay when using the BeamDagRunner. When using the KubeflowDagRunner all but Transform run okay. I build the pipeline on Kubeflow notebook, upload the packaged pipeline via UI, or run via KFP.Client().create_run_from_pipeline_package directly. 

The Kubeflow Pipelines UI shows all the components but Transform run okay. I see the respective DataFlow jobs in the DataFlow UI. Transform fails with the error below and I do not even see a respective DataFlow job in the DataFlow UI.

INFO:absl:Nonempty beam arg setup_file already includes dependency
WARNING:tensorflow:From /tfx-src/tfx/components/transform/executor.py:511: Schema (from tensorflow_transform.tf_metadata.dataset_schema) is deprecated and will be removed in a future version.
Instructions for updating:
Schema is a deprecated, use schema_utils.schema_from_feature_spec to create a `Schema`
Traceback (most recent call last):
  File "/tfx-src/tfx/orchestration/kubeflow/container_entrypoint.py", line 382, in <module>
    main()
  File "/tfx-src/tfx/orchestration/kubeflow/container_entrypoint.py", line 375, in main
    execution_info = launcher.launch()
  File "/tfx-src/tfx/orchestration/launcher/base_component_launcher.py", line 205, in launch
    execution_decision.exec_properties)
  File "/tfx-src/tfx/orchestration/launcher/in_process_component_launcher.py", line 67, in _run_executor
    executor.Do(input_dict, output_dict, exec_properties)
  File "/tfx-src/tfx/components/transform/executor.py", line 391, in Do
    self.Transform(label_inputs, label_outputs, status_file)
  File "/tfx-src/tfx/components/transform/executor.py", line 833, in Transform
    preprocessing_fn = self._GetPreprocessingFn(inputs, outputs)
  File "/tfx-src/tfx/components/transform/executor.py", line 775, in _GetPreprocessingFn
    preprocessing_fn_path_split[-1])
  File "/tfx-src/tfx/utils/import_utils.py", line 77, in import_func_from_module
    user_module = importlib.import_module(module_path)
  File "/opt/venv/lib/python3.6/importlib/__init__.py", line 126, in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
  File "<frozen importlib._bootstrap>", line 994, in _gcd_import
  File "<frozen importlib._bootstrap>", line 971, in _find_and_load
  File "<frozen importlib._bootstrap>", line 953, in _find_and_load_unlocked
ModuleNotFoundError: No module named 'preprocessing'

BEAM_PIPELINE_ARGS = [
    '--project=' + GCP_PROJECT_ID,
    '--runner=DataflowRunner',
    # '--experiments=shuffle_mode=auto',
    '--temp_location=' + os.path.join('gs://', GCS_BUCKET_NAME, 'tmp'),
    '--setup_file=./setup.py',
    #'--requirements_file=./requirements.txt',
    '--region=' + GCP_REGION,
]

Here's the setup.py where I attempted to include the preprocessing module explicitly to no avail

import setuptools

setuptools.setup(
    name='markdowns-tfx-pipeline',
    version='0.2',
    py_modules=['configs', 'features', 'hparams', 'model', 'pipeline', 'preprocessing'],
    install_requires=['tfx'],
    packages=setuptools.find_packages(),
)


Discussion:
Hi @yantriks-edi-bice 
If I understand it correctly the reason is perhaps in the container entrypoint running on KFP, the setup_file will be overriden by the code [here](https://github.com/tensorflow/tfx/blob/0fb096ff605f9b45f29a9f93621c57c682cac4ac/tfx/orchestration/kubeflow/container_entrypoint.py#L118).

Perhaps we can try fixing this by detecting the collision before constructing the beam pipeline args? WDYT @neuromage ?
Yes, that looks like a bug. @numerology do you mind sending out a fix for ensuring user setup.py takes precedence over the default for Beam?
Yes @numerology that does look like the reason. I could see from the logs that the TFX setup.py was being used and all things TFX were working as expected but anything from my code (and I believe the Transform component is the first one which refers to anything from my code) was failing. 

Now the code you pointed to does give you access to the provided arguments and could easily check if user supplied a setup.py file however it cannot simply override the TFX setup.py one, right? Isn't that required as well, or would it work if user supplies tfx as a requirement in user setup.py file like I do above?
> Yes, that looks like a bug. @numerology do you mind sending out a fix for ensuring user setup.py takes precedence over the default for Beam?

Sure

> Isn't that required as well, or would it work if user supplies tfx as a requirement in user setup.py file like I do above?

I believe you'll need to declare tfx as a requirement in your setup.py to make sure TFX library works correctly in Dataflow worker as well.
I applied the patch to my TFX container entry point in my Kubeflow notebook but not seeing the effect. In fact, I'm not sure if my notebook local tfx package is being called at all as I added some other logging info and not seeing it. Dataflow worker may be using an image with tfx in it or installing a clean (not my patched) one but the orchestration part I was expecting to be done locally on the Kubeflow notebook container - is that true?
> I applied the patch to my TFX container entry point in my Kubeflow notebook but not seeing the effect. 

Did you rebuild and push an container image, and use that in the KubeflowDagRunner? Thanks
No, I did not but I do see now how I can pass a custom container to KubeflowDagRunner. How is the default tfx container image put together, and what's the best way for me to build a patched one while waiting for this to be merged and an image released?
@yantriks-edi-bice 

The Dockerfile of the default TFX container released by TFX team can be found at [1], which is built based on a base image (containing some heavy and common dependencies) [2]. We have a helper script for building image at [3].

TFX official release usually comes at biweek cadence, but we have a nightly release at [3], for example, `docker pull tensorflow/tfx:0.22.0.dev20200428` will give you the nightly release image built on April 28.

[1] https://github.com/tensorflow/tfx/blob/master/tfx/tools/docker/Dockerfile
[2] https://github.com/tensorflow/tfx/blob/master/tfx/tools/docker/base/Dockerfile
[3] https://github.com/tensorflow/tfx/blob/master/tfx/tools/docker/build_docker_image.sh
[4] https://hub.docker.com/r/tensorflow/tfx/tags
Thanks @numerology, I just came across that script [3] now as well. Will give it a try.

Resolving patch:
diff --git a/tfx/orchestration/kubeflow/container_entrypoint.py b/tfx/orchestration/kubeflow/container_entrypoint.py
--- a/tfx/orchestration/kubeflow/container_entrypoint.py
+++ b/tfx/orchestration/kubeflow/container_entrypoint.py
@@ -115,28 +115,6 @@ def _get_grpc_metadata_connection_config(
   return connection_config
 
 
-def _make_beam_pipeline_args(json_beam_pipeline_args: Text) -> List[Text]:
-  """Constructs beam_pipeline_args for ComponentLauncher.
-
-  Args:
-    json_beam_pipeline_args: JSON serialized list of beam pipeline args.
-
-  Returns:
-    List containing `beam_pipeline_args`.
-  """
-  beam_pipeline_args = json.loads(json_beam_pipeline_args)
-
-  # Ensure beam pipelines args has a setup.py file so we can use
-  # DataflowRunner.
-  module_dir = os.environ['TFX_SRC_DIR']
-  setup_file = os.path.join(module_dir, 'setup.py')
-  absl.logging.info('Using setup_file \'%s\' to capture TFX dependencies',
-                    setup_file)
-  beam_pipeline_args.append('--setup_file={}'.format(setup_file))
-
-  return beam_pipeline_args
-
-
 def _sanitize_underscore(name: Text) -> Text:
   """Sanitize the underscore in pythonic name for markdown visualization."""
   if name:
@@ -350,7 +328,7 @@ def main():
       _get_metadata_connection_config(kubeflow_metadata_config))
   driver_args = data_types.DriverArgs(enable_cache=args.enable_cache)
 
-  beam_pipeline_args = _make_beam_pipeline_args(args.beam_pipeline_args)
+  beam_pipeline_args = json.loads(args.beam_pipeline_args)
 
   additional_pipeline_args = json.loads(args.additional_pipeline_args)
 


Test patch: