Skip to content

Commit 4f8cd17

Browse files
committed
Add test to verify setting of SemaphoreKey and MutexName fields in DSL
Signed-off-by: ddalvi <ddalvi@redhat.com>
1 parent 6af6032 commit 4f8cd17

File tree

1 file changed

+36
-0
lines changed

1 file changed

+36
-0
lines changed

sdk/python/kfp/compiler/compiler_test.py

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3847,6 +3847,42 @@ def outer():
38473847
foo_platform_set_bar_feature(task, 12)
38483848

38493849

3850+
class TestPipelineSemaphoreMutex(unittest.TestCase):
3851+
3852+
def test_pipeline_with_semaphore_and_mutex(self):
3853+
from kfp import compiler
3854+
from kfp import dsl
3855+
from kfp.dsl.pipeline_config import PipelineConfig
3856+
3857+
config = PipelineConfig()
3858+
config.set_semaphore_key('semaphore')
3859+
config.set_mutex_name('mutex')
3860+
3861+
@dsl.pipeline(pipeline_config=config)
3862+
def my_pipeline():
3863+
task = comp()
3864+
3865+
with tempfile.TemporaryDirectory() as tempdir:
3866+
output_yaml = os.path.join(tempdir, 'pipeline.yaml')
3867+
compiler.Compiler().compile(
3868+
pipeline_func=my_pipeline, package_path=output_yaml)
3869+
3870+
with open(output_yaml, 'r') as f:
3871+
pipeline_docs = list(yaml.safe_load_all(f))
3872+
3873+
pipeline_spec = None
3874+
for doc in pipeline_docs:
3875+
if 'platforms' in doc:
3876+
pipeline_spec = doc
3877+
break
3878+
3879+
if pipeline_spec:
3880+
kubernetes_spec = pipeline_spec['platforms']['kubernetes'][
3881+
'pipelineConfig']
3882+
assert kubernetes_spec['semaphoreKey'] == 'semaphore'
3883+
assert kubernetes_spec['mutexName'] == 'mutex'
3884+
3885+
38503886
class ExtractInputOutputDescription(unittest.TestCase):
38513887

38523888
def test_no_descriptions(self):

0 commit comments

Comments
 (0)