File tree Expand file tree Collapse file tree 1 file changed +36
-0
lines changed Expand file tree Collapse file tree 1 file changed +36
-0
lines changed Original file line number Diff line number Diff line change @@ -3847,6 +3847,42 @@ def outer():
3847
3847
foo_platform_set_bar_feature (task , 12 )
3848
3848
3849
3849
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
+
3850
3886
class ExtractInputOutputDescription (unittest .TestCase ):
3851
3887
3852
3888
def test_no_descriptions (self ):
You can’t perform that action at this time.
0 commit comments