From 3ab5da5801eb560e4698a7823830ecc49df6624a Mon Sep 17 00:00:00 2001 From: Madhur Karampudi <142544288+vkarampudi@users.noreply.github.com> Date: Mon, 13 May 2024 10:58:17 -0700 Subject: [PATCH] [Fix] Support YAML files in kubeflow v2 dag runner (#6808) * Support pipeline spec 2.1 with backward compatibility PiperOrigin-RevId: 631251911 * Support YAML files in kubeflow v2 dag runner PiperOrigin-RevId: 631970915 * Update version.py for 1.15.1 Release * Automated rollback of commit 9d86d1a1261909983503448072610dbeb5766498 PiperOrigin-RevId: 632488503 * Support YAML files in kubeflow v2 dag runner PiperOrigin-RevId: 633095531 --------- Co-authored-by: tfx-team --- RELEASE.md | 3 ++- tfx/orchestration/kubeflow/v2/kubeflow_v2_dag_runner.py | 4 ++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/RELEASE.md b/RELEASE.md index c17b4ae6e4..9bb07f6d49 100644 --- a/RELEASE.md +++ b/RELEASE.md @@ -19,4 +19,5 @@ | -- | -- | -- | -- | | `kfp-pipeline-spec` | `kfp-pipeline-spec>=0.1.10,<0.2` | `>0.1.13,<0.2` | | -## Documentation Updates \ No newline at end of file +## Documentation Updates + diff --git a/tfx/orchestration/kubeflow/v2/kubeflow_v2_dag_runner.py b/tfx/orchestration/kubeflow/v2/kubeflow_v2_dag_runner.py index 1302c881a9..ed610b7eef 100644 --- a/tfx/orchestration/kubeflow/v2/kubeflow_v2_dag_runner.py +++ b/tfx/orchestration/kubeflow/v2/kubeflow_v2_dag_runner.py @@ -89,12 +89,12 @@ def _write_pipeline_spec_to_file( yaml_comments = _extract_comments_from_pipeline_spec( pipeline_spec_dict, pipeline_description ) - with open(package_path, 'w') as yaml_file: + with fileio.open(package_path, 'w') as yaml_file: yaml_file.write(yaml_comments) documents = [pipeline_spec_dict] yaml.dump_all(documents, yaml_file, sort_keys=True) else: - with open(package_path, 'w') as json_file: + with fileio.open(package_path, 'w') as json_file json.dump(pipeline_job_dict, json_file, sort_keys=True)