Skip to content

Commit

Permalink
Add test to verify setting of SemaphoreKey and MutexName fields in KF…
Browse files Browse the repository at this point in the history
…P DSL

Signed-off-by: ddalvi <[email protected]>
  • Loading branch information
DharmitD committed Nov 14, 2024
1 parent c4620a9 commit 78adc95
Showing 1 changed file with 36 additions and 0 deletions.
36 changes: 36 additions & 0 deletions sdk/python/kfp/compiler/compiler_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -3847,6 +3847,42 @@ def outer():
foo_platform_set_bar_feature(task, 12)


class TestPipelineSemaphoreMutex(unittest.TestCase):

def test_pipeline_with_semaphore_and_mutex(self):
from kfp import compiler
from kfp import dsl
from kfp.dsl.pipeline_config import PipelineConfig

config = PipelineConfig()
config.set_semaphore_key("semaphore")
config.set_mutex_name("mutex")

@dsl.pipeline(pipeline_config=config)
def my_pipeline():
task = comp()

with tempfile.TemporaryDirectory() as tempdir:
output_yaml = os.path.join(tempdir, 'pipeline.yaml')
compiler.Compiler().compile(
pipeline_func=my_pipeline, package_path=output_yaml)

with open(output_yaml, 'r') as f:
pipeline_docs = list(yaml.safe_load_all(f))

pipeline_spec = None
for doc in pipeline_docs:
if 'platforms' in doc:
pipeline_spec = doc
break

if pipeline_spec:
kubernetes_spec = pipeline_spec['platforms']['kubernetes'][
'pipelineConfig']
assert kubernetes_spec['semaphoreKey'] == "semaphore"
assert kubernetes_spec['mutexName'] == "mutex"


class ExtractInputOutputDescription(unittest.TestCase):

def test_no_descriptions(self):
Expand Down

0 comments on commit 78adc95

Please sign in to comment.