feat(sdk+api) Implement Semaphores and Mutexes Pipeline Configuration

Signed-off-by: ddalvi <[email protected]>

Author: DharmitD

api/v2alpha1/go/pipelinespec/pipeline_spec.pb.go

@@ -1103,5 +1103,9 @@ message PlatformDeploymentConfig {
 
 // Spec for pipeline-level config options. See PipelineConfig DSL class.
 message PipelineConfig {
-  // TODO add pipeline-level configs
+  // Name of the semaphore key to control pipeline concurrency
+  string semaphore_key = 1;
+
+  // Name of the mutex to ensure mutual exclusion
+  string mutex_name = 2;
 }

api/v2alpha1/pipeline_spec.proto

@@ -2011,11 +2011,13 @@ def write_pipeline_spec_to_file(
 
 def _merge_pipeline_config(pipelineConfig: pipeline_config.PipelineConfig,
                            platformSpec: pipeline_spec_pb2.PlatformSpec):
-    # TODO: add pipeline config options (ttl, semaphore, etc.) to the dict
-    # json_format.ParseDict(
-    #     {'pipelineConfig': {
-    #         '<some pipeline config option>': pipelineConfig.<get that value>,
-    #     }}, platformSpec.platforms['kubernetes'])
+    json_format.ParseDict(
+        {
+            'pipelineConfig': {
+                'semaphoreName': pipelineConfig.semaphore_key,
+                'mutexName': pipelineConfig.mutex_name,
+            }
+        }, platformSpec.platforms['kubernetes'])
 
     return platformSpec
 

sdk/python/kfp/compiler/pipeline_spec_builder.py

@@ -18,6 +18,21 @@ class PipelineConfig:
     """PipelineConfig contains pipeline-level config options."""
 
     def __init__(self):
-        pass
+        self.semaphore_key = None
+        self.mutex_name = None
 
-    # TODO add pipeline level configs
+    def set_semaphore_key(self, semaphore_key: str):
+        """Set the name of the semaphore to control pipeline concurrency.
+
+        Args:
+            semaphore_key (str): Name of the semaphore.
+        """
+        self.semaphore_key = semaphore_key.strip()
+
+    def set_mutex_name(self, mutex_name: str):
+        """Set the name of the mutex to ensure mutual exclusion.
+
+        Args:
+            mutex_name (str): Name of the mutex.
+        """
+        self.mutex_name = mutex_name.strip()