We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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.
```



--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `tfx/tools/cli/handler/base_handler.py`
Content:
```
1 # Lint as: python2, python3
2 # Copyright 2019 Google LLC. All Rights Reserved.
3 #
4 # Licensed under the Apache License, Version 2.0 (the "License");
5 # you may not use this file except in compliance with the License.
6 # You may obtain a copy of the License at
7 #
8 #     http://www.apache.org/licenses/LICENSE-2.0
9 #
10 # Unless required by applicable law or agreed to in writing, software
11 # distributed under the License is distributed on an "AS IS" BASIS,
12 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 # See the License for the specific language governing permissions and
14 # limitations under the License.
15 """Base handler class."""
16 
17 from __future__ import absolute_import
18 from __future__ import division
19 from __future__ import print_function
20 
21 import abc
22 import json
23 import os
24 import re
25 import subprocess
26 import sys
27 import tempfile
28 from typing import Any, Dict, List, Text
29 
30 import click
31 from six import with_metaclass
32 import tensorflow as tf
33 
34 from tfx.tools.cli import labels
35 from tfx.utils import io_utils
36 
37 
38 class BaseHandler(with_metaclass(abc.ABCMeta, object)):
39   """Base Handler for CLI.
40 
41   Attributes:
42     flags_dict: A dictionary with flags provided in a command.
43   """
44 
45   def __init__(self, flags_dict: Dict[Text, Any]):
46     self.flags_dict = flags_dict
47     self._handler_home_dir = self._get_handler_home()
48 
49   @abc.abstractmethod
50   def create_pipeline(self) -> None:
51     """Creates pipeline for the handler."""
52     pass
53 
54   @abc.abstractmethod
55   def update_pipeline(self) -> None:
56     """Updates pipeline for the handler."""
57     pass
58 
59   @abc.abstractmethod
60   def list_pipelines(self) -> None:
61     """List all the pipelines in the environment."""
62     pass
63 
64   @abc.abstractmethod
65   def delete_pipeline(self) -> None:
66     """Deletes pipeline for the handler."""
67     pass
68 
69   @abc.abstractmethod
70   def compile_pipeline(self) -> None:
71     """Compiles pipeline for the handler."""
72     pass
73 
74   @abc.abstractmethod
75   def create_run(self) -> None:
76     """Runs a pipeline for the handler."""
77     pass
78 
79   @abc.abstractmethod
80   def delete_run(self) -> None:
81     """Deletes a run."""
82     pass
83 
84   @abc.abstractmethod
85   def terminate_run(self) -> None:
86     """Stops a run."""
87     pass
88 
89   @abc.abstractmethod
90   def list_runs(self) -> None:
91     """Lists all runs of a pipeline."""
92     pass
93 
94   @abc.abstractmethod
95   def get_run(self) -> None:
96     """Checks run status."""
97     pass
98 
99   def _check_pipeline_dsl_path(self) -> None:
100     """Check if pipeline dsl path exists."""
101     pipeline_dsl_path = self.flags_dict[labels.PIPELINE_DSL_PATH]
102     if not tf.io.gfile.exists(pipeline_dsl_path):
103       sys.exit('Invalid pipeline path: {}'.format(pipeline_dsl_path))
104 
105   def _check_dsl_runner(self) -> None:
106     """Check if runner in dsl is same as engine flag."""
107     engine_flag = self.flags_dict[labels.ENGINE_FLAG]
108     with open(self.flags_dict[labels.PIPELINE_DSL_PATH], 'r') as f:
109       dsl_contents = f.read()
110       regexes = {
111           labels.AIRFLOW_ENGINE: r'AirflowDagRunner\(.*\)',
112           labels.KUBEFLOW_ENGINE: r'KubeflowDagRunner\(.*\)',
113           labels.BEAM_ENGINE: r'BeamDagRunner\(.*\)'
114       }
115       match = re.search(regexes[engine_flag], dsl_contents)
116       if not match:
117         sys.exit('{} runner not found in dsl.'.format(engine_flag))
118 
119   def _extract_pipeline_args(self) -> Dict[Text, Any]:
120     """Get pipeline args from the DSL.
121 
122     Returns:
123       Python dictionary with pipeline details extracted from DSL.
124     """
125     pipeline_dsl_path = self.flags_dict[labels.PIPELINE_DSL_PATH]
126     if os.path.isdir(pipeline_dsl_path):
127       sys.exit('Provide dsl file path.')
128 
129     # Create an environment for subprocess.
130     temp_env = os.environ.copy()
131 
132     # Create temp file to store pipeline_args from pipeline dsl.
133     temp_file = tempfile.mkstemp(prefix='cli_tmp_', suffix='_pipeline_args')[1]
134 
135     # Store temp_file path in temp_env.
136     temp_env[labels.TFX_JSON_EXPORT_PIPELINE_ARGS_PATH] = temp_file
137 
138     # Mark the SDK environment if not in a template.
139     if 'pipelines.kubeflow.org/pipeline-sdk-type' not in temp_env:
140       temp_env['pipelines.kubeflow.org/pipeline-sdk-type'] = 'tfx-cli'
141 
142     # Run dsl with mock environment to store pipeline args in temp_file.
143     self._subprocess_call([sys.executable, pipeline_dsl_path], env=temp_env)
144     if os.stat(temp_file).st_size != 0:
145       # Load pipeline_args from temp_file for TFX pipelines
146       with open(temp_file, 'r') as f:
147         pipeline_args = json.load(f)
148     else:
149       # For non-TFX pipelines, extract pipeline name from the dsl filename.
150       pipeline_args = {
151           labels.PIPELINE_NAME:
152               os.path.basename(pipeline_dsl_path).split('.')[0]
153       }
154 
155     # Delete temp file
156     io_utils.delete_dir(temp_file)
157 
158     return pipeline_args
159 
160   def _get_handler_home(self) -> Text:
161     """Sets handler home.
162 
163     Returns:
164       Path to handler home directory.
165     """
166     engine_flag = self.flags_dict[labels.ENGINE_FLAG]
167     handler_home_dir = engine_flag.upper() + '_HOME'
168     if handler_home_dir in os.environ:
169       return os.environ[handler_home_dir]
170     return os.path.join(os.environ['HOME'], engine_flag, '')
171 
172   def _subprocess_call(self,
173                        command: List[Text],
174                        env: Dict[Text, Any] = None) -> None:
175     return_code = subprocess.call(command, env=env)
176     if return_code != 0:
177       sys.exit('Error while running "{}" '.format(' '.join(command)))
178 
179   def _check_pipeline_existence(self,
180                                 pipeline_name: Text,
181                                 required: bool = True) -> None:
182     """Check if pipeline folder exists and if not, exit system.
183 
184     Args:
185       pipeline_name: Name of the pipeline.
186       required: Set it as True if pipeline needs to exist else set it to False.
187     """
188     handler_pipeline_path = os.path.join(self._handler_home_dir, pipeline_name,
189                                          '')
190     # Check if pipeline folder exists.
191     exists = tf.io.gfile.exists(handler_pipeline_path)
192     if required and not exists:
193       sys.exit('Pipeline "{}" does not exist.'.format(pipeline_name))
194     elif not required and exists:
195       sys.exit('Pipeline "{}" already exists.'.format(pipeline_name))
196 
197   def get_schema(self):
198     pipeline_name = self.flags_dict[labels.PIPELINE_NAME]
199 
200     # Check if pipeline exists.
201     self._check_pipeline_existence(pipeline_name)
202 
203     # Path to pipeline args.
204     pipeline_args_path = os.path.join(self._handler_home_dir,
205                                       self.flags_dict[labels.PIPELINE_NAME],
206                                       'pipeline_args.json')
207 
208     # Get pipeline_root.
209     with open(pipeline_args_path, 'r') as f:
210       pipeline_args = json.load(f)
211 
212     # Check if pipeline root created. If not, it means that the user has not
213     # created a run yet or the pipeline is still running for the first time.
214     pipeline_root = pipeline_args[labels.PIPELINE_ROOT]
215     if not tf.io.gfile.exists(pipeline_root):
216       sys.exit(
217           'Create a run before inferring schema. If pipeline is already running, then wait for it to successfully finish.'
218       )
219 
220     # If pipeline_root exists, then check if SchemaGen output exists.
221     components = tf.io.gfile.listdir(pipeline_root)
222     if 'SchemaGen' not in components:
223       sys.exit(
224           'Either SchemaGen component does not exist or pipeline is still running. If pipeline is running, then wait for it to successfully finish.'
225       )
226 
227     # Get the latest SchemaGen output.
228     schemagen_outputs = tf.io.gfile.listdir(
229         os.path.join(pipeline_root, 'SchemaGen', 'schema', ''))
230     latest_schema_folder = max(schemagen_outputs, key=int)
231 
232     # Copy schema to current dir.
233     latest_schema_path = os.path.join(pipeline_root, 'SchemaGen', 'schema',
234                                       latest_schema_folder, 'schema.pbtxt')
235     curr_dir_path = os.path.join(os.getcwd(), 'schema.pbtxt')
236     io_utils.copy_file(latest_schema_path, curr_dir_path, overwrite=True)
237 
238     # Print schema and path to schema
239     click.echo('Path to schema: {}'.format(curr_dir_path))
240     click.echo('*********SCHEMA FOR {}**********'.format(pipeline_name.upper()))
241     with open(curr_dir_path, 'r') as f:
242       click.echo(f.read())
243
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

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 @@
     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]:

{"golden_diff": "diff --git a/tfx/tools/cli/handler/base_handler.py b/tfx/tools/cli/handler/base_handler.py\n--- a/tfx/tools/cli/handler/base_handler.py\n+++ b/tfx/tools/cli/handler/base_handler.py\n@@ -21,7 +21,6 @@\n import abc\n import json\n import os\n-import re\n import subprocess\n import sys\n import tempfile\n@@ -107,13 +106,12 @@\n     engine_flag = self.flags_dict[labels.ENGINE_FLAG]\n     with open(self.flags_dict[labels.PIPELINE_DSL_PATH], 'r') as f:\n       dsl_contents = f.read()\n-      regexes = {\n-          labels.AIRFLOW_ENGINE: r'AirflowDagRunner\\(.*\\)',\n-          labels.KUBEFLOW_ENGINE: r'KubeflowDagRunner\\(.*\\)',\n-          labels.BEAM_ENGINE: r'BeamDagRunner\\(.*\\)'\n+      runner_names = {\n+          labels.AIRFLOW_ENGINE: 'AirflowDagRunner',\n+          labels.KUBEFLOW_ENGINE: 'KubeflowDagRunner',\n+          labels.BEAM_ENGINE: 'BeamDagRunner',\n       }\n-      match = re.search(regexes[engine_flag], dsl_contents)\n-      if not match:\n+      if runner_names[engine_flag] not in dsl_contents:\n         sys.exit('{} runner not found in dsl.'.format(engine_flag))\n \n   def _extract_pipeline_args(self) -> Dict[Text, Any]:\n", "issue": "Code formatting can affect the semantic behavior of tfx pipelines (DagRunner not being registered)\nhttps://github.com/tensorflow/tfx/blob/bfbc9587c568c1247f6a6cfb59db161ed3a4970c/tfx/tools/cli/handler/base_handler.py#L106-L117\r\n\r\nThe referenced code is where the issue happens.\r\n\r\nAs 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.\r\n\r\n### Steps to reproduce the issue:\r\nThe following piece of code _works just fine_.\r\n```python\r\n kf_config = KubeflowDagRunnerConfig(\r\n        kubeflow_metadata_config=kf_metadata,\r\n        tfx_image=os.environ.get(\"KUBEFLOW_TFX_IMAGE\", None),\r\n )\r\n\r\nKubeflowDagRunner(config=kf_config)\r\n```\r\n\r\nThe following piece of code though, _fails_ with **_kubeflow runner not found in dsl._**\r\n```python\r\nKubeflowDagRunner(\r\n    config=KubeflowDagRunnerConfig(\r\n           kubeflow_metadata_config=kf_metadata,\r\n           tfx_image=os.environ.get(\"KUBEFLOW_TFX_IMAGE\", None),\r\n     )\r\n)\r\n```\r\n\r\nOutput of the `tfx pipeline create` cli command is as follows:\r\n```\r\nCLI\r\nUpdating pipeline\r\nDetected Kubeflow.\r\nUse --engine flag if you intend to use a different orchestrator.\r\nReading build spec from build.yaml\r\nUse skaffold to build the container image.\r\n/usr/local/bin/skaffold\r\nkubeflow runner not found in dsl.\r\nNew container image is built. Target image is available in the build spec file.\r\n```\r\n\r\n\n", "before_files": [{"content": "# Lint as: python2, python3\n# Copyright 2019 Google LLC. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"Base handler class.\"\"\"\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport abc\nimport json\nimport os\nimport re\nimport subprocess\nimport sys\nimport tempfile\nfrom typing import Any, Dict, List, Text\n\nimport click\nfrom six import with_metaclass\nimport tensorflow as tf\n\nfrom tfx.tools.cli import labels\nfrom tfx.utils import io_utils\n\n\nclass BaseHandler(with_metaclass(abc.ABCMeta, object)):\n  \"\"\"Base Handler for CLI.\n\n  Attributes:\n    flags_dict: A dictionary with flags provided in a command.\n  \"\"\"\n\n  def __init__(self, flags_dict: Dict[Text, Any]):\n    self.flags_dict = flags_dict\n    self._handler_home_dir = self._get_handler_home()\n\n  @abc.abstractmethod\n  def create_pipeline(self) -> None:\n    \"\"\"Creates pipeline for the handler.\"\"\"\n    pass\n\n  @abc.abstractmethod\n  def update_pipeline(self) -> None:\n    \"\"\"Updates pipeline for the handler.\"\"\"\n    pass\n\n  @abc.abstractmethod\n  def list_pipelines(self) -> None:\n    \"\"\"List all the pipelines in the environment.\"\"\"\n    pass\n\n  @abc.abstractmethod\n  def delete_pipeline(self) -> None:\n    \"\"\"Deletes pipeline for the handler.\"\"\"\n    pass\n\n  @abc.abstractmethod\n  def compile_pipeline(self) -> None:\n    \"\"\"Compiles pipeline for the handler.\"\"\"\n    pass\n\n  @abc.abstractmethod\n  def create_run(self) -> None:\n    \"\"\"Runs a pipeline for the handler.\"\"\"\n    pass\n\n  @abc.abstractmethod\n  def delete_run(self) -> None:\n    \"\"\"Deletes a run.\"\"\"\n    pass\n\n  @abc.abstractmethod\n  def terminate_run(self) -> None:\n    \"\"\"Stops a run.\"\"\"\n    pass\n\n  @abc.abstractmethod\n  def list_runs(self) -> None:\n    \"\"\"Lists all runs of a pipeline.\"\"\"\n    pass\n\n  @abc.abstractmethod\n  def get_run(self) -> None:\n    \"\"\"Checks run status.\"\"\"\n    pass\n\n  def _check_pipeline_dsl_path(self) -> None:\n    \"\"\"Check if pipeline dsl path exists.\"\"\"\n    pipeline_dsl_path = self.flags_dict[labels.PIPELINE_DSL_PATH]\n    if not tf.io.gfile.exists(pipeline_dsl_path):\n      sys.exit('Invalid pipeline path: {}'.format(pipeline_dsl_path))\n\n  def _check_dsl_runner(self) -> None:\n    \"\"\"Check if runner in dsl is same as engine flag.\"\"\"\n    engine_flag = self.flags_dict[labels.ENGINE_FLAG]\n    with open(self.flags_dict[labels.PIPELINE_DSL_PATH], 'r') as f:\n      dsl_contents = f.read()\n      regexes = {\n          labels.AIRFLOW_ENGINE: r'AirflowDagRunner\\(.*\\)',\n          labels.KUBEFLOW_ENGINE: r'KubeflowDagRunner\\(.*\\)',\n          labels.BEAM_ENGINE: r'BeamDagRunner\\(.*\\)'\n      }\n      match = re.search(regexes[engine_flag], dsl_contents)\n      if not match:\n        sys.exit('{} runner not found in dsl.'.format(engine_flag))\n\n  def _extract_pipeline_args(self) -> Dict[Text, Any]:\n    \"\"\"Get pipeline args from the DSL.\n\n    Returns:\n      Python dictionary with pipeline details extracted from DSL.\n    \"\"\"\n    pipeline_dsl_path = self.flags_dict[labels.PIPELINE_DSL_PATH]\n    if os.path.isdir(pipeline_dsl_path):\n      sys.exit('Provide dsl file path.')\n\n    # Create an environment for subprocess.\n    temp_env = os.environ.copy()\n\n    # Create temp file to store pipeline_args from pipeline dsl.\n    temp_file = tempfile.mkstemp(prefix='cli_tmp_', suffix='_pipeline_args')[1]\n\n    # Store temp_file path in temp_env.\n    temp_env[labels.TFX_JSON_EXPORT_PIPELINE_ARGS_PATH] = temp_file\n\n    # Mark the SDK environment if not in a template.\n    if 'pipelines.kubeflow.org/pipeline-sdk-type' not in temp_env:\n      temp_env['pipelines.kubeflow.org/pipeline-sdk-type'] = 'tfx-cli'\n\n    # Run dsl with mock environment to store pipeline args in temp_file.\n    self._subprocess_call([sys.executable, pipeline_dsl_path], env=temp_env)\n    if os.stat(temp_file).st_size != 0:\n      # Load pipeline_args from temp_file for TFX pipelines\n      with open(temp_file, 'r') as f:\n        pipeline_args = json.load(f)\n    else:\n      # For non-TFX pipelines, extract pipeline name from the dsl filename.\n      pipeline_args = {\n          labels.PIPELINE_NAME:\n              os.path.basename(pipeline_dsl_path).split('.')[0]\n      }\n\n    # Delete temp file\n    io_utils.delete_dir(temp_file)\n\n    return pipeline_args\n\n  def _get_handler_home(self) -> Text:\n    \"\"\"Sets handler home.\n\n    Returns:\n      Path to handler home directory.\n    \"\"\"\n    engine_flag = self.flags_dict[labels.ENGINE_FLAG]\n    handler_home_dir = engine_flag.upper() + '_HOME'\n    if handler_home_dir in os.environ:\n      return os.environ[handler_home_dir]\n    return os.path.join(os.environ['HOME'], engine_flag, '')\n\n  def _subprocess_call(self,\n                       command: List[Text],\n                       env: Dict[Text, Any] = None) -> None:\n    return_code = subprocess.call(command, env=env)\n    if return_code != 0:\n      sys.exit('Error while running \"{}\" '.format(' '.join(command)))\n\n  def _check_pipeline_existence(self,\n                                pipeline_name: Text,\n                                required: bool = True) -> None:\n    \"\"\"Check if pipeline folder exists and if not, exit system.\n\n    Args:\n      pipeline_name: Name of the pipeline.\n      required: Set it as True if pipeline needs to exist else set it to False.\n    \"\"\"\n    handler_pipeline_path = os.path.join(self._handler_home_dir, pipeline_name,\n                                         '')\n    # Check if pipeline folder exists.\n    exists = tf.io.gfile.exists(handler_pipeline_path)\n    if required and not exists:\n      sys.exit('Pipeline \"{}\" does not exist.'.format(pipeline_name))\n    elif not required and exists:\n      sys.exit('Pipeline \"{}\" already exists.'.format(pipeline_name))\n\n  def get_schema(self):\n    pipeline_name = self.flags_dict[labels.PIPELINE_NAME]\n\n    # Check if pipeline exists.\n    self._check_pipeline_existence(pipeline_name)\n\n    # Path to pipeline args.\n    pipeline_args_path = os.path.join(self._handler_home_dir,\n                                      self.flags_dict[labels.PIPELINE_NAME],\n                                      'pipeline_args.json')\n\n    # Get pipeline_root.\n    with open(pipeline_args_path, 'r') as f:\n      pipeline_args = json.load(f)\n\n    # Check if pipeline root created. If not, it means that the user has not\n    # created a run yet or the pipeline is still running for the first time.\n    pipeline_root = pipeline_args[labels.PIPELINE_ROOT]\n    if not tf.io.gfile.exists(pipeline_root):\n      sys.exit(\n          'Create a run before inferring schema. If pipeline is already running, then wait for it to successfully finish.'\n      )\n\n    # If pipeline_root exists, then check if SchemaGen output exists.\n    components = tf.io.gfile.listdir(pipeline_root)\n    if 'SchemaGen' not in components:\n      sys.exit(\n          'Either SchemaGen component does not exist or pipeline is still running. If pipeline is running, then wait for it to successfully finish.'\n      )\n\n    # Get the latest SchemaGen output.\n    schemagen_outputs = tf.io.gfile.listdir(\n        os.path.join(pipeline_root, 'SchemaGen', 'schema', ''))\n    latest_schema_folder = max(schemagen_outputs, key=int)\n\n    # Copy schema to current dir.\n    latest_schema_path = os.path.join(pipeline_root, 'SchemaGen', 'schema',\n                                      latest_schema_folder, 'schema.pbtxt')\n    curr_dir_path = os.path.join(os.getcwd(), 'schema.pbtxt')\n    io_utils.copy_file(latest_schema_path, curr_dir_path, overwrite=True)\n\n    # Print schema and path to schema\n    click.echo('Path to schema: {}'.format(curr_dir_path))\n    click.echo('*********SCHEMA FOR {}**********'.format(pipeline_name.upper()))\n    with open(curr_dir_path, 'r') as f:\n      click.echo(f.read())\n", "path": "tfx/tools/cli/handler/base_handler.py"}], "after_files": [{"content": "# Lint as: python2, python3\n# Copyright 2019 Google LLC. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"Base handler class.\"\"\"\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport abc\nimport json\nimport os\nimport subprocess\nimport sys\nimport tempfile\nfrom typing import Any, Dict, List, Text\n\nimport click\nfrom six import with_metaclass\nimport tensorflow as tf\n\nfrom tfx.tools.cli import labels\nfrom tfx.utils import io_utils\n\n\nclass BaseHandler(with_metaclass(abc.ABCMeta, object)):\n  \"\"\"Base Handler for CLI.\n\n  Attributes:\n    flags_dict: A dictionary with flags provided in a command.\n  \"\"\"\n\n  def __init__(self, flags_dict: Dict[Text, Any]):\n    self.flags_dict = flags_dict\n    self._handler_home_dir = self._get_handler_home()\n\n  @abc.abstractmethod\n  def create_pipeline(self) -> None:\n    \"\"\"Creates pipeline for the handler.\"\"\"\n    pass\n\n  @abc.abstractmethod\n  def update_pipeline(self) -> None:\n    \"\"\"Updates pipeline for the handler.\"\"\"\n    pass\n\n  @abc.abstractmethod\n  def list_pipelines(self) -> None:\n    \"\"\"List all the pipelines in the environment.\"\"\"\n    pass\n\n  @abc.abstractmethod\n  def delete_pipeline(self) -> None:\n    \"\"\"Deletes pipeline for the handler.\"\"\"\n    pass\n\n  @abc.abstractmethod\n  def compile_pipeline(self) -> None:\n    \"\"\"Compiles pipeline for the handler.\"\"\"\n    pass\n\n  @abc.abstractmethod\n  def create_run(self) -> None:\n    \"\"\"Runs a pipeline for the handler.\"\"\"\n    pass\n\n  @abc.abstractmethod\n  def delete_run(self) -> None:\n    \"\"\"Deletes a run.\"\"\"\n    pass\n\n  @abc.abstractmethod\n  def terminate_run(self) -> None:\n    \"\"\"Stops a run.\"\"\"\n    pass\n\n  @abc.abstractmethod\n  def list_runs(self) -> None:\n    \"\"\"Lists all runs of a pipeline.\"\"\"\n    pass\n\n  @abc.abstractmethod\n  def get_run(self) -> None:\n    \"\"\"Checks run status.\"\"\"\n    pass\n\n  def _check_pipeline_dsl_path(self) -> None:\n    \"\"\"Check if pipeline dsl path exists.\"\"\"\n    pipeline_dsl_path = self.flags_dict[labels.PIPELINE_DSL_PATH]\n    if not tf.io.gfile.exists(pipeline_dsl_path):\n      sys.exit('Invalid pipeline path: {}'.format(pipeline_dsl_path))\n\n  def _check_dsl_runner(self) -> None:\n    \"\"\"Check if runner in dsl is same as engine flag.\"\"\"\n    engine_flag = self.flags_dict[labels.ENGINE_FLAG]\n    with open(self.flags_dict[labels.PIPELINE_DSL_PATH], 'r') as f:\n      dsl_contents = f.read()\n      runner_names = {\n          labels.AIRFLOW_ENGINE: 'AirflowDagRunner',\n          labels.KUBEFLOW_ENGINE: 'KubeflowDagRunner',\n          labels.BEAM_ENGINE: 'BeamDagRunner',\n      }\n      if runner_names[engine_flag] not in dsl_contents:\n        sys.exit('{} runner not found in dsl.'.format(engine_flag))\n\n  def _extract_pipeline_args(self) -> Dict[Text, Any]:\n    \"\"\"Get pipeline args from the DSL.\n\n    Returns:\n      Python dictionary with pipeline details extracted from DSL.\n    \"\"\"\n    pipeline_dsl_path = self.flags_dict[labels.PIPELINE_DSL_PATH]\n    if os.path.isdir(pipeline_dsl_path):\n      sys.exit('Provide dsl file path.')\n\n    # Create an environment for subprocess.\n    temp_env = os.environ.copy()\n\n    # Create temp file to store pipeline_args from pipeline dsl.\n    temp_file = tempfile.mkstemp(prefix='cli_tmp_', suffix='_pipeline_args')[1]\n\n    # Store temp_file path in temp_env.\n    temp_env[labels.TFX_JSON_EXPORT_PIPELINE_ARGS_PATH] = temp_file\n\n    # Mark the SDK environment if not in a template.\n    if 'pipelines.kubeflow.org/pipeline-sdk-type' not in temp_env:\n      temp_env['pipelines.kubeflow.org/pipeline-sdk-type'] = 'tfx-cli'\n\n    # Run dsl with mock environment to store pipeline args in temp_file.\n    self._subprocess_call([sys.executable, pipeline_dsl_path], env=temp_env)\n    if os.stat(temp_file).st_size != 0:\n      # Load pipeline_args from temp_file for TFX pipelines\n      with open(temp_file, 'r') as f:\n        pipeline_args = json.load(f)\n    else:\n      # For non-TFX pipelines, extract pipeline name from the dsl filename.\n      pipeline_args = {\n          labels.PIPELINE_NAME:\n              os.path.basename(pipeline_dsl_path).split('.')[0]\n      }\n\n    # Delete temp file\n    io_utils.delete_dir(temp_file)\n\n    return pipeline_args\n\n  def _get_handler_home(self) -> Text:\n    \"\"\"Sets handler home.\n\n    Returns:\n      Path to handler home directory.\n    \"\"\"\n    engine_flag = self.flags_dict[labels.ENGINE_FLAG]\n    handler_home_dir = engine_flag.upper() + '_HOME'\n    if handler_home_dir in os.environ:\n      return os.environ[handler_home_dir]\n    return os.path.join(os.environ['HOME'], engine_flag, '')\n\n  def _subprocess_call(self,\n                       command: List[Text],\n                       env: Dict[Text, Any] = None) -> None:\n    return_code = subprocess.call(command, env=env)\n    if return_code != 0:\n      sys.exit('Error while running \"{}\" '.format(' '.join(command)))\n\n  def _check_pipeline_existence(self,\n                                pipeline_name: Text,\n                                required: bool = True) -> None:\n    \"\"\"Check if pipeline folder exists and if not, exit system.\n\n    Args:\n      pipeline_name: Name of the pipeline.\n      required: Set it as True if pipeline needs to exist else set it to False.\n    \"\"\"\n    handler_pipeline_path = os.path.join(self._handler_home_dir, pipeline_name,\n                                         '')\n    # Check if pipeline folder exists.\n    exists = tf.io.gfile.exists(handler_pipeline_path)\n    if required and not exists:\n      sys.exit('Pipeline \"{}\" does not exist.'.format(pipeline_name))\n    elif not required and exists:\n      sys.exit('Pipeline \"{}\" already exists.'.format(pipeline_name))\n\n  def get_schema(self):\n    pipeline_name = self.flags_dict[labels.PIPELINE_NAME]\n\n    # Check if pipeline exists.\n    self._check_pipeline_existence(pipeline_name)\n\n    # Path to pipeline args.\n    pipeline_args_path = os.path.join(self._handler_home_dir,\n                                      self.flags_dict[labels.PIPELINE_NAME],\n                                      'pipeline_args.json')\n\n    # Get pipeline_root.\n    with open(pipeline_args_path, 'r') as f:\n      pipeline_args = json.load(f)\n\n    # Check if pipeline root created. If not, it means that the user has not\n    # created a run yet or the pipeline is still running for the first time.\n    pipeline_root = pipeline_args[labels.PIPELINE_ROOT]\n    if not tf.io.gfile.exists(pipeline_root):\n      sys.exit(\n          'Create a run before inferring schema. If pipeline is already running, then wait for it to successfully finish.'\n      )\n\n    # If pipeline_root exists, then check if SchemaGen output exists.\n    components = tf.io.gfile.listdir(pipeline_root)\n    if 'SchemaGen' not in components:\n      sys.exit(\n          'Either SchemaGen component does not exist or pipeline is still running. If pipeline is running, then wait for it to successfully finish.'\n      )\n\n    # Get the latest SchemaGen output.\n    schemagen_outputs = tf.io.gfile.listdir(\n        os.path.join(pipeline_root, 'SchemaGen', 'schema', ''))\n    latest_schema_folder = max(schemagen_outputs, key=int)\n\n    # Copy schema to current dir.\n    latest_schema_path = os.path.join(pipeline_root, 'SchemaGen', 'schema',\n                                      latest_schema_folder, 'schema.pbtxt')\n    curr_dir_path = os.path.join(os.getcwd(), 'schema.pbtxt')\n    io_utils.copy_file(latest_schema_path, curr_dir_path, overwrite=True)\n\n    # Print schema and path to schema\n    click.echo('Path to schema: {}'.format(curr_dir_path))\n    click.echo('*********SCHEMA FOR {}**********'.format(pipeline_name.upper()))\n    with open(curr_dir_path, 'r') as f:\n      click.echo(f.read())\n", "path": "tfx/tools/cli/handler/base_handler.py"}]}