SDK/Client - Supporting pipeline packages with multiple files

kubeflow · Apr 23, 2019 · b4d9678 · b4d9678
1 parent 4e41128
commit b4d9678
Showing 1 changed file with 12 additions and 4 deletions.
diff --git a/sdk/python/kfp/_client.py b/sdk/python/kfp/_client.py
@@ -174,8 +174,12 @@ def _extract_pipeline_yaml(self, package_file):
         if len(all_yaml_files) == 0:
           raise ValueError('Invalid package. Missing pipeline yaml file in the package.')
 
-        if len(all_yaml_files) > 1:
-          raise ValueError('Invalid package. Multiple yaml files in the package.')
+        pipeline_file = all_yaml_files[0]
+        if 'pipeline.yaml' in all_yaml_files:
+          pipeline_file = 'pipeline.yaml'
+        else:
+          if len(all_yaml_files) > 1:
+            raise ValueError('Invalid package. Multiple yaml files in the package and no pipeline.yaml file.')
 
         with tar.extractfile(all_yaml_files[0]) as f:
           return yaml.load(f)
@@ -186,8 +190,12 @@ def _extract_pipeline_yaml(self, package_file):
         if len(all_yaml_files) == 0:
           raise ValueError('Invalid package. Missing pipeline yaml file in the package.')
 
-        if len(all_yaml_files) > 1:
-          raise ValueError('Invalid package. Multiple yaml files in the package.')
+        pipeline_file = all_yaml_files[0]
+        if 'pipeline.yaml' in all_yaml_files:
+          pipeline_file = 'pipeline.yaml'
+        else:
+          if len(all_yaml_files) > 1:
+            raise ValueError('Invalid package. Missing pipeline yaml file in the package and no pipeline.yaml file.')
 
         filename = zip.extract(all_yaml_files[0])
         with open(filename, 'r') as f: