Repository: tensorflow/tfx

Issue:
Code formatting can affect the semantic behavior of tfx pipelines (DagRunner not being registered)
https://github.com/tensorflow/tfx/blob/bfbc9587c568c1247f6a6cfb59db161ed3a4970c/tfx/tools/cli/handler/base_handler.py#L106-L117

The referenced code is where the issue happens.

As can be seen, it checks for existence of a `KubeflowDagRunner(...)` in the contents, but that's a naive assumption since the opening and closing parenthesis can be in separate lines and in that case the regex won't match.

### Steps to reproduce the issue:
The following piece of code _works just fine_.
```python
 kf_config = KubeflowDagRunnerConfig(
        kubeflow_metadata_config=kf_metadata,
        tfx_image=os.environ.get("KUBEFLOW_TFX_IMAGE", None),
 )

KubeflowDagRunner(config=kf_config)
```

The following piece of code though, _fails_ with **_kubeflow runner not found in dsl._**
```python
KubeflowDagRunner(
    config=KubeflowDagRunnerConfig(
           kubeflow_metadata_config=kf_metadata,
           tfx_image=os.environ.get("KUBEFLOW_TFX_IMAGE", None),
     )
)
```

Output of the `tfx pipeline create` cli command is as follows:
```
CLI
Updating pipeline
Detected Kubeflow.
Use --engine flag if you intend to use a different orchestrator.
Reading build spec from build.yaml
Use skaffold to build the container image.
/usr/local/bin/skaffold
kubeflow runner not found in dsl.
New container image is built. Target image is available in the build spec file.
```




Discussion:
Yeap this is problematic. Sorry for your inconvenience.

Resolving patch:
diff --git a/tfx/tools/cli/handler/base_handler.py b/tfx/tools/cli/handler/base_handler.py
--- a/tfx/tools/cli/handler/base_handler.py
+++ b/tfx/tools/cli/handler/base_handler.py
@@ -21,7 +21,6 @@
 import abc
 import json
 import os
-import re
 import subprocess
 import sys
 import tempfile
@@ -107,13 +106,12 @@ def _check_dsl_runner(self) -> None:
     engine_flag = self.flags_dict[labels.ENGINE_FLAG]
     with open(self.flags_dict[labels.PIPELINE_DSL_PATH], 'r') as f:
       dsl_contents = f.read()
-      regexes = {
-          labels.AIRFLOW_ENGINE: r'AirflowDagRunner\(.*\)',
-          labels.KUBEFLOW_ENGINE: r'KubeflowDagRunner\(.*\)',
-          labels.BEAM_ENGINE: r'BeamDagRunner\(.*\)'
+      runner_names = {
+          labels.AIRFLOW_ENGINE: 'AirflowDagRunner',
+          labels.KUBEFLOW_ENGINE: 'KubeflowDagRunner',
+          labels.BEAM_ENGINE: 'BeamDagRunner',
       }
-      match = re.search(regexes[engine_flag], dsl_contents)
-      if not match:
+      if runner_names[engine_flag] not in dsl_contents:
         sys.exit('{} runner not found in dsl.'.format(engine_flag))
 
   def _extract_pipeline_args(self) -> Dict[Text, Any]:


Test patch:
diff --git a/tfx/tools/cli/handler/base_handler_test.py b/tfx/tools/cli/handler/base_handler_test.py
--- a/tfx/tools/cli/handler/base_handler_test.py
+++ b/tfx/tools/cli/handler/base_handler_test.py
@@ -92,7 +92,13 @@ def testCheckPipelineDslPathInvalid(self):
     self.assertEqual(str(err.exception), 'Invalid pipeline path: {}'
                      .format(flags_dict[labels.PIPELINE_DSL_PATH]))
 
-  def testCheckDslRunnerWrongEngine(self):
+  def testCheckDslRunner(self):
+    flags_dict = {labels.ENGINE_FLAG: self.engine,
+                  labels.PIPELINE_DSL_PATH: self.pipeline_path}
+    handler = FakeHandler(flags_dict)
+    handler._check_dsl_runner()
+
+  def testCheckDslRunner_WrongEngine(self):
     flags_dict = {labels.ENGINE_FLAG: 'kubeflow',
                   labels.PIPELINE_DSL_PATH: self.pipeline_path}
     handler = FakeHandler(flags_dict)