Repository: tensorflow/tfx

Issue #584: Pipeline fails with tf.keras.Model subclass
Created: 2019-09-05T20:00:11Z
Running a pipeline throws a KeyError: user_module while running 'Run[Transform]' when using a custom tf.Estimator whose model_fn references a tf.keras.Model subclass. This bug was first tracked by #500 before I figured out what the root cause was. Replacing the subclass with a direct call to tf.keras.Sequential with the same layers works just fine, however, this is only a stopgap solution since more complicated models might require subclassing or the keras functional API (which I will test).

Duplicate issue #500: User module not found although it's properly specified
Created: 2019-08-18T18:54:49Z
My Transform component fails with the following error. Not sure if I'm doing something wrong or if it's a bug...
Local TFX build: ef9b28b
TFX image used in pipeline: 0.14.0.dev20190818 from https://hub.docker.com/r/tensorflow/tfx/tags
Kubeflow: v0.5.0
Running on GKE
EDIT: The TFX image used is the release image on TFX DockerHub tagged 0.14.0.dev20190818, which should be in sync with ef9b28b. I'm not using DataflowRunner in my pipeline. The pipeline is a custom one, it's not the taxi example. I've confirmed that my module file is properly stored in GCP and that the pipeline can access it — the Transform component is successfully calling my module file's preprocessing_fn (confirmed by print statements), so the KeyError: 'user_module' error seems to be happening in the Transform component but downstream of preprocessing_fn.
Traceback (most recent call last):
  File "/tfx-src/tfx/orchestration/kubeflow/container_entrypoint.py", line 106, in <module>
    main()
  File "/tfx-src/tfx/orchestration/kubeflow/container_entrypoint.py", line 102, in main
    executor.run()
  File "/opt/venv/lib/python3.5/site-packages/tfx/orchestration/kubeflow/executor_wrappers.py", line 123, in run
    self._exec_properties)
  File "/opt/venv/lib/python3.5/site-packages/tfx/components/transform/executor.py", line 285, in Do
    self.Transform(label_inputs, label_outputs, status_file)
  File "/opt/venv/lib/python3.5/site-packages/tfx/components/transform/executor.py", line 752, in Transform
    compute_statistics, materialize_output_paths)
  File "/opt/venv/lib/python3.5/site-packages/tfx/components/transform/executor.py", line 1042, in _RunBeamImpl
    raw_examples_file_format, output_path))
  File "/opt/venv/lib/python3.5/site-packages/apache_beam/pipeline.py", line 426, in __exit__
    self.run().wait_until_finish()
  File "/opt/venv/lib/python3.5/site-packages/apache_beam/pipeline.py", line 404, in run
    self.to_runner_api(use_fake_coders=True),
  File "/opt/venv/lib/python3.5/site-packages/apache_beam/pipeline.py", line 656, in to_runner_api
    root_transform_id = context.transforms.get_id(self._root_transform())
  File "/opt/venv/lib/python3.5/site-packages/apache_beam/runners/pipeline_context.py", line 84, in get_id
    self._id_to_proto[id] = obj.to_runner_api(self._pipeline_context)
  File "/opt/venv/lib/python3.5/site-packages/apache_beam/pipeline.py", line 877, in to_runner_api
    for part in self.parts],
  File "/opt/venv/lib/python3.5/site-packages/apache_beam/pipeline.py", line 877, in <listcomp>
    for part in self.parts],
  File "/opt/venv/lib/python3.5/site-packages/apache_beam/runners/pipeline_context.py", line 84, in get_id
    self._id_to_proto[id] = obj.to_runner_api(self._pipeline_context)
  File "/opt/venv/lib/python3.5/site-packages/apache_beam/pipeline.py", line 875, in to_runner_api
    spec=transform_to_runner_api(self.transform, context),
  File "/opt/venv/lib/python3.5/site-packages/apache_beam/pipeline.py", line 870, in transform_to_runner_api
    return transform.to_runner_api(context, has_parts=bool(self.parts))
  File "/opt/venv/lib/python3.5/site-packages/apache_beam/transforms/ptransform.py", line 591, in to_runner_api
    urn, typed_param = self.to_runner_api_pickled(context)
  File "/opt/venv/lib/python3.5/site-packages/apache_beam/transforms/ptransform.py", line 624, in to_runner_api_pickled
    pickler.dumps(self))
  File "/opt/venv/lib/python3.5/site-packages/apache_beam/internal/pickler.py", line 230, in dumps
    s = dill.dumps(o)
  File "/opt/venv/lib/python3.5/site-packages/dill/_dill.py", line 294, in dumps
    dump(obj, file, protocol, byref, fmode, recurse)#, strictio)
  File "/opt/venv/lib/python3.5/site-packages/dill/_dill.py", line 287, in dump
    pik.dump(obj)
  File "/usr/lib/python3.5/pickle.py", line 408, in dump
    self.save(obj)
  File "/usr/lib/python3.5/pickle.py", line 520, in save
    self.save_reduce(obj=obj, *rv)
  File "/usr/lib/python3.5/pickle.py", line 627, in save_reduce
    save(state)
  File "/usr/lib/python3.5/pickle.py", line 475, in save
    f(self, obj) # Call unbound method with explicit self
  File "/opt/venv/lib/python3.5/site-packages/apache_beam/internal/pickler.py", line 198, in new_save_module_dict
    return old_save_module_dict(pickler, obj)
  File "/opt/venv/lib/python3.5/site-packages/dill/_dill.py", line 902, in save_module_dict
    StockPickler.save_dict(pickler, obj)
  File "/usr/lib/python3.5/pickle.py", line 814, in save_dict
    self._batch_setitems(obj.items())
  File "/usr/lib/python3.5/pickle.py", line 840, in _batch_setitems
    save(v)
  File "/usr/lib/python3.5/pickle.py", line 475, in save
    f(self, obj) # Call unbound method with explicit self
  File "/usr/lib/python3.5/pickle.py", line 744, in save_tuple
    save(element)
  File "/usr/lib/python3.5/pickle.py", line 475, in save
    f(self, obj) # Call unbound method with explicit self
  File "/opt/venv/lib/python3.5/site-packages/dill/_dill.py", line 1386, in save_function
    obj.__dict__), obj=obj)
  File "/usr/lib/python3.5/pickle.py", line 603, in save_reduce
    save(args)
  File "/usr/lib/python3.5/pickle.py", line 475, in save
    f(self, obj) # Call unbound method with explicit self
  File "/usr/lib/python3.5/pickle.py", line 744, in save_tuple
    save(element)
  File "/usr/lib/python3.5/pickle.py", line 475, in save
    f(self, obj) # Call unbound method with explicit self
  File "/opt/venv/lib/python3.5/site-packages/apache_beam/internal/pickler.py", line 198, in new_save_module_dict
    return old_save_module_dict(pickler, obj)
  File "/opt/venv/lib/python3.5/site-packages/dill/_dill.py", line 902, in save_module_dict
    StockPickler.save_dict(pickler, obj)
  File "/usr/lib/python3.5/pickle.py", line 814, in save_dict
    self._batch_setitems(obj.items())
  File "/usr/lib/python3.5/pickle.py", line 840, in _batch_setitems
    save(v)
  File "/usr/lib/python3.5/pickle.py", line 475, in save
    f(self, obj) # Call unbound method with explicit self
  File "/opt/venv/lib/python3.5/site-packages/apache_beam/internal/pickler.py", line 101, in wrapper
    if _is_nested_class(obj) and obj.__module__ != '__main__':
  File "/opt/venv/lib/python3.5/site-packages/apache_beam/internal/pickler.py", line 60, in _is_nested_class
    and cls.__name__ not in sys.modules[cls.__module__].__dict__)
KeyError: 'user_module'

Closing comment (2019-09-05T20:02:59Z): It appears that this bug is due to using a tf.keras.Model subclass in the pipeline. This is now tracked by #584. Closing.