ImportExampleGen on Kubeflow not working
After building tfx image on latest master, I just discovered using ImportExampleGen on Kubeflow results in:
```
usage: container_entrypoint.py [-h] --exec_properties EXEC_PROPERTIES
                               --outputs OUTPUTS --executor_class_path
                               EXECUTOR_CLASS_PATH
                               {CsvExampleGen,BigQueryExampleGen,StatisticsGen,SchemaGen,ExampleValidator,Transform,Trainer,Evaluator,ModelValidator,Pusher}
                               ...
container_entrypoint.py: error: argument executor: invalid choice: 'ImportExampleGen' (choose from 'CsvExampleGen', 'BigQueryExampleGen', 'StatisticsGen', 'SchemaGen', 'ExampleValidator', 'Transform', 'Trainer', 'Evaluator', 'ModelValidator', 'Pusher')
```
diff --git a/tfx/orchestration/kubeflow/container_entrypoint.py b/tfx/orchestration/kubeflow/container_entrypoint.py
index 4d41e5ee..635c1cc3 100644
--- a/tfx/orchestration/kubeflow/container_entrypoint.py
+++ b/tfx/orchestration/kubeflow/container_entrypoint.py
@@ -51,6 +51,10 @@ def main():
   subparser.add_argument('--input_base', type=str, required=True)
   subparser.set_defaults(executor=wrappers.CsvExampleGenWrapper)
 
+  subparser = subparsers.add_parser('ImportExampleGen')
+  subparser.add_argument('--input_base', type=str, required=True)
+  subparser.set_defaults(executor=wrappers.ImportExampleGenWrapper)
+
   subparser = subparsers.add_parser('BigQueryExampleGen')
   subparser.set_defaults(executor=wrappers.BigQueryExampleGenWrapper)
 
diff --git a/tfx/orchestration/kubeflow/executor_wrappers.py b/tfx/orchestration/kubeflow/executor_wrappers.py
index ffad9492..c4f8589d 100644
--- a/tfx/orchestration/kubeflow/executor_wrappers.py
+++ b/tfx/orchestration/kubeflow/executor_wrappers.py
@@ -158,6 +158,29 @@ class CsvExampleGenWrapper(KubeflowExecutorWrapper):
       input_artifact.span = span
 
 
+class ImportExampleGenWrapper(KubeflowExecutorWrapper):
+  """Wrapper for ImportExampleGen component."""
+
+  def __init__(self, args: argparse.Namespace):
+    super(ImportExampleGenWrapper, self).__init__(
+        executor_class_path=args.executor_class_path,
+        name='ImportExampleGen',
+        input_dict={
+            'input_base': parse_tfx_type(args.input_base),
+        },
+        outputs=args.outputs,
+        exec_properties=json.loads(args.exec_properties),
+    )
+    self._set_input_artifact_span()
+    self._set_outputs()
+
+  def _set_input_artifact_span(self):
+    for input_artifact in self._input_dict['input_base']:
+      matched = re.match(r'span_([0-9]+)', input_artifact.uri)
+      span = matched.group(1) if matched else 1
+      input_artifact.span = span
+
+
 class BigQueryExampleGenWrapper(KubeflowExecutorWrapper):
   """Wrapper for BigQueryExampleGen component."""
 
['tfx/orchestration/kubeflow/container_entrypoint.py', 'tfx/orchestration/kubeflow/executor_wrappers.py']