Skip to content

Commit

Permalink
SDK/Compiler - Preventing pipeline entrypoint template name from clas…
Browse files Browse the repository at this point in the history
…hing with other template names

Case exhibiting the problem:
```
def add(a, b):
    ...
@dsl.pipeline(name="add')
def some_name():
    add(...)
```
  • Loading branch information
Ark-kun committed Nov 25, 2019
1 parent c74a9d8 commit 5f6b165
Show file tree
Hide file tree
Showing 2 changed files with 31 additions and 3 deletions.
23 changes: 20 additions & 3 deletions sdk/python/kfp/compiler/compiler.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
from deprecated import deprecated
import inspect
import tarfile
import uuid
import zipfile
from typing import Callable, Set, List, Text, Dict, Tuple, Any, Union, Optional

Expand Down Expand Up @@ -619,9 +620,13 @@ def _create_pipeline_workflow(self, args, pipeline, op_transformers=None, pipeli
param['value'] = str(arg.value)
input_params.append(param)

# Making the pipeline group name unique to prevent name clashes with templates
pipeline_group = pipeline.groups[0]
temp_pipeline_group_name = uuid.uuid4().hex
pipeline_group.name = temp_pipeline_group_name

# Templates
templates = self._create_dag_templates(pipeline, op_transformers)
templates.sort(key=lambda x: x['name'])

# Exit Handler
exit_handler = None
Expand All @@ -632,12 +637,24 @@ def _create_pipeline_workflow(self, args, pipeline, op_transformers=None, pipeli

# The whole pipeline workflow
pipeline_name = pipeline.name or 'Pipeline'

# Workaround for pipeline name clashing with container template names
# TODO: Make sure template names cannot clash at all (container, DAG, workflow)
template_map = {template['name'].lower(): template for template in templates}
from ..components._naming import _make_name_unique_by_adding_index
pipeline_template_name = _make_name_unique_by_adding_index(pipeline_name, template_map, '-')

# Restoring the name of the pipeline template
pipeline_template = template_map[temp_pipeline_group_name]
pipeline_template['name'] = pipeline_template_name

templates.sort(key=lambda x: x['name'])
workflow = {
'apiVersion': 'argoproj.io/v1alpha1',
'kind': 'Workflow',
'metadata': {'generateName': pipeline_name + '-'},
'metadata': {'generateName': pipeline_template_name + '-'},
'spec': {
'entrypoint': pipeline_name,
'entrypoint': pipeline_template_name,
'templates': templates,
'arguments': {'parameters': input_params},
'serviceAccountName': 'pipeline-runner'
Expand Down
11 changes: 11 additions & 0 deletions sdk/python/tests/compiler/compiler_tests.py
Original file line number Diff line number Diff line change
Expand Up @@ -784,3 +784,14 @@ def test_py_input_artifact_raw_value(self):
"""Test pipeline input_artifact_raw_value."""
self._test_py_compile_yaml('input_artifact_raw_value')

def test_pipeline_name_same_as_task_name(self):
def some_name():
dsl.ContainerOp(
name='some_name',
image='alpine:latest',
)

workflow_dict = compiler.Compiler()._compile(some_name)
template_names = set(template['name'] for template in workflow_dict['spec']['templates'])
self.assertGreater(len(template_names), 1)
self.assertEqual(template_names, {'some-name', 'some-name-2'})

0 comments on commit 5f6b165

Please sign in to comment.