Introduces PipelineConfig class with TTL configuration

Signed-off-by: Ricardo M. Oliveira <[email protected]>

api/v2alpha1/pipeline_spec.proto

@@ -1080,7 +1080,7 @@ message PlatformSpec {
   // Platform key to full platform config
   map<string, SinglePlatformSpec> platforms = 1;
 
-  optional int32 ttlStrategysecondsAfterCompletion = 2;
+
 }
 
 message SinglePlatformSpec {
@@ -1093,8 +1093,14 @@ message SinglePlatformSpec {
   // Arbitrary configuration, which will be defined by the platform
   // protos/libraries.
   google.protobuf.Struct config = 3;
+
+  PipelineConfig pipelineConfig = 4;
 }
 
+message PipelineConfig {
+  // Set TTL at pipeline-level
+  optional int32 pipelineTtl = 1;
+}
 
 message PlatformDeploymentConfig {
   // Map of executor label to executor-level config

backend/src/apiserver/template/v2_template.go

@@ -23,6 +23,7 @@ import (
 	"regexp"
 	"strings"
 
+	"github.com/golang/glog"
 	structpb "github.com/golang/protobuf/ptypes/struct"
 	"github.com/kubeflow/pipelines/api/v2alpha1/go/pipelinespec"
 	"github.com/kubeflow/pipelines/backend/src/apiserver/model"
@@ -77,9 +78,17 @@ func (t *V2Spec) ScheduledWorkflow(modelJob *model.Job) (*scheduledworkflow.Sche
 		}
 	}
 
+	var pipeline_options argocompiler.Options
+	if t.platformSpec.PipelineConfig.Ttl != nil {
+		glog.Info("Found pipeline config")
+		pipeline_options = argocompiler.Options{
+			TtlSeconds: *t.platformSpec.PipelineConfig.Ttl,
+		}
+	}
+
 	var obj interface{}
 	if util.CurrentExecutionType() == util.ArgoWorkflow {
-		obj, err = argocompiler.Compile(job, kubernetesSpec, nil)
+		obj, err = argocompiler.Compile(job, kubernetesSpec, &pipeline_options)
 	} else if util.CurrentExecutionType() == util.TektonPipelineRun {
 		obj, err = tektoncompiler.Compile(job, kubernetesSpec, &tektoncompiler.Options{LauncherImage: Launcher})
 	}
@@ -300,9 +309,17 @@ func (t *V2Spec) RunWorkflow(modelRun *model.Run, options RunWorkflowOptions) (u
 		}
 	}
 
+	var pipeline_options *argocompiler.Options
+	if t.platformSpec.PipelineConfig.Ttl != nil {
+		glog.Info("Found pipeline config")
+		pipeline_options = &argocompiler.Options{
+			TtlSeconds: *t.platformSpec.PipelineConfig.Ttl,
+		}
+	}
+
 	var obj interface{}
 	if util.CurrentExecutionType() == util.ArgoWorkflow {
-		obj, err = argocompiler.Compile(job, kubernetesSpec, nil)
+		obj, err = argocompiler.Compile(job, kubernetesSpec, pipeline_options)
 	} else if util.CurrentExecutionType() == util.TektonPipelineRun {
 		obj, err = tektoncompiler.Compile(job, kubernetesSpec, nil)
 	}
@@ -344,7 +361,7 @@ func IsPlatformSpecWithKubernetesConfig(template []byte) bool {
 		return false
 	}
 	_, ok := platformSpec.Platforms["kubernetes"]
-	return ok
+	return ok || platformSpec.PipelineConfig != nil
 }
 
 func (t *V2Spec) validatePipelineJobInputs(job *pipelinespec.PipelineJob) error {

backend/src/v2/compiler/argocompiler/argo.go

@@ -40,8 +40,13 @@ type Options struct {
 	// optional
 	PipelineRoot string
 	// TODO(Bobgy): add an option -- dev mode, ImagePullPolicy should only be Always in dev mode.
+	TtlSeconds int32
 }
 
+const (
+	pipeline_default_ttlSeconds = int32(30)
+)
+
 func Compile(jobArg *pipelinespec.PipelineJob, kubernetesSpecArg *pipelinespec.SinglePlatformSpec, opts *Options) (*wfapi.Workflow, error) {
 	// clone jobArg, because we don't want to change it
 	jobMsg := proto.Clone(jobArg)
@@ -86,6 +91,11 @@ func Compile(jobArg *pipelinespec.PipelineJob, kubernetesSpecArg *pipelinespec.S
 		}
 	}
 
+	pipeline_ttlseconds := pipeline_default_ttlSeconds
+	if &opts.TtlSeconds != nil {
+		pipeline_ttlseconds = opts.TtlSeconds
+	}
+
 	// initialization
 	wf := &wfapi.Workflow{
 		TypeMeta: k8smeta.TypeMeta{
@@ -117,6 +127,9 @@ func Compile(jobArg *pipelinespec.PipelineJob, kubernetesSpecArg *pipelinespec.S
 			},
 			ServiceAccountName: "pipeline-runner",
 			Entrypoint:         tmplEntrypoint,
+			TTLStrategy: &wfapi.TTLStrategy{
+				SecondsAfterCompletion: &pipeline_ttlseconds,
+			},
 		},
 	}
 	c := &workflowCompiler{

sdk/python/kfp/compiler/pipeline_spec_builder.py

@@ -28,6 +28,7 @@
 from kfp.dsl import component_factory
 from kfp.dsl import for_loop
 from kfp.dsl import pipeline_channel
+from kfp.dsl import pipeline_config
 from kfp.dsl import pipeline_context
 from kfp.dsl import pipeline_task
 from kfp.dsl import placeholders
@@ -1879,7 +1880,7 @@ def create_pipeline_spec(
     pipeline: pipeline_context.Pipeline,
     component_spec: structures.ComponentSpec,
     pipeline_outputs: Optional[Any] = None,
-    ttl_strategy_seconds_after_completion: Optional[int] = None,
+    pipeline_config: Optional[pipeline_config.PipelineConfig] = None,
 ) -> Tuple[pipeline_spec_pb2.PipelineSpec, pipeline_spec_pb2.PlatformSpec]:
     """Creates a pipeline spec object.
 
@@ -1947,9 +1948,10 @@ def create_pipeline_spec(
         group_name_to_group=group_name_to_group,
         condition_channels=condition_channels,
     )
-    import pdb; pdb.set_trace()
+
     platform_spec = pipeline_spec_pb2.PlatformSpec()
-    platform_spec.ttlStrategysecondsAfterCompletion = ttl_strategy_seconds_after_completion
+    if pipeline_config is not None:
+        _merge_pipeline_config(pipelineConfig=pipeline_config, platformSpec=platform_spec)
     for group in all_groups:
         build_spec_by_group(
             pipeline_spec=pipeline_spec,
@@ -2063,7 +2065,16 @@ def write_pipeline_spec_to_file(
     else:
         raise ValueError(
             f'The output path {package_path} should end with ".yaml".')
+
+def _merge_pipeline_config(pipelineConfig: pipeline_config.PipelineConfig, platformSpec: pipeline_spec_pb2.PlatformSpec):
+    pipeline_config_json = json_format.ParseDict(
+        {
+            "pipelineConfig": {
+                "pipelineTtl": pipelineConfig.get_ttl(),
+            }
+        }, platformSpec.platforms.get('kubernetes', {}))
 
+    return platformSpec
 
 def extract_comments_from_pipeline_spec(pipeline_spec: dict,
                                         pipeline_description: str) -> str:

sdk/python/kfp/dsl/component_decorator.py

@@ -16,7 +16,7 @@
 from typing import Callable, List, Optional
 import warnings
 
-from kfp.dsl import component_factory
+from kfp.dsl import component_factory, pipeline_config
 
 
 def component(func: Optional[Callable] = None,
@@ -28,7 +28,7 @@ def component(func: Optional[Callable] = None,
               output_component_file: Optional[str] = None,
               install_kfp_package: bool = True,
               kfp_package_path: Optional[str] = None,
-              ttl_active_seconds: Optional[int] = None):
+              pipeline_config: Optional[pipeline_config.PipelineConfig] = None):
     """Decorator for Python-function based components.
 
     A KFP component can either be a lightweight component or a containerized
@@ -116,7 +116,7 @@ def pipeline():
             output_component_file=output_component_file,
             install_kfp_package=install_kfp_package,
             kfp_package_path=kfp_package_path,
-            ttl_active_seconds=ttl_active_seconds)
+            pipeline_config=pipeline_config)
 
     return component_factory.create_component_from_func(
         func,
@@ -127,4 +127,4 @@ def pipeline():
         output_component_file=output_component_file,
         install_kfp_package=install_kfp_package,
         kfp_package_path=kfp_package_path,
-        ttl_active_seconds=ttl_active_seconds)
+        pipeline_config=pipeline_config)

sdk/python/kfp/dsl/component_factory.py

@@ -27,6 +27,7 @@
 from kfp.dsl import container_component_artifact_channel
 from kfp.dsl import container_component_class
 from kfp.dsl import graph_component
+from kfp.dsl import pipeline_config
 from kfp.dsl import placeholders
 from kfp.dsl import python_component
 from kfp.dsl import structures
@@ -534,14 +535,14 @@ def create_component_from_func(
     install_kfp_package: bool = True,
     kfp_package_path: Optional[str] = None,
     pip_trusted_hosts: Optional[List[str]] = None,
-    ttl_active_seconds: Optional[int] = None,
+    pipeline_config: Optional[pipeline_config.PipelineConfig] = None,
 ) -> python_component.PythonComponent:
     """Implementation for the @component decorator.
 
     The decorator is defined under component_decorator.py. See the
     decorator for the canonical documentation for this function.
     """
-
+    
     packages_to_install_command = _get_packages_to_install_command(
         install_kfp_package=install_kfp_package,
         target_image=target_image,
@@ -678,7 +679,7 @@ def create_graph_component_from_func(
     name: Optional[str] = None,
     description: Optional[str] = None,
     display_name: Optional[str] = None,
-    ttl_strategy_seconds_after_completion: Optional[int] = None,
+    pipeline_config: Optional[pipeline_config.PipelineConfig] = None,
 ) -> graph_component.GraphComponent:
     """Implementation for the @pipeline decorator.
 
@@ -695,7 +696,7 @@ def create_graph_component_from_func(
         component_spec=component_spec,
         pipeline_func=func,
         display_name=display_name,
-        ttl_strategy_seconds_after_completion=ttl_strategy_seconds_after_completion,
+        pipeline_config=pipeline_config,
     )
 
 

sdk/python/kfp/dsl/graph_component.py

@@ -20,6 +20,7 @@
 from kfp.compiler import pipeline_spec_builder as builder
 from kfp.dsl import base_component
 from kfp.dsl import pipeline_channel
+from kfp.dsl import pipeline_config
 from kfp.dsl import pipeline_context
 from kfp.dsl import structures
 from kfp.pipeline_spec import pipeline_spec_pb2
@@ -37,7 +38,7 @@ def __init__(
         component_spec: structures.ComponentSpec,
         pipeline_func: Callable,
         display_name: Optional[str] = None,
-        ttl_strategy_seconds_after_completion: Optional[int] = None,
+        pipeline_config: Optional[pipeline_config.PipelineConfig] = None,
     ):
         super().__init__(component_spec=component_spec)
         self.pipeline_func = pipeline_func
@@ -70,7 +71,7 @@ def __init__(
             pipeline=dsl_pipeline,
             component_spec=self.component_spec,
             pipeline_outputs=pipeline_outputs,
-            ttl_strategy_seconds_after_completion=ttl_strategy_seconds_after_completion,
+            pipeline_config=pipeline_config,
         )
 
         pipeline_root = getattr(pipeline_func, 'pipeline_root', None)

sdk/python/kfp/dsl/pipeline_config.py

@@ -0,0 +1,28 @@
+# Copyright 2024 The Kubeflow Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Pipeline-level config options."""
+
+
+class PipelineConfig:
+    """PipelineConfig contains pipeline-level config options."""
+
+    def __init__(self):
+        pass
+
+    # TODO add pipeline level configs
+    def set_ttl(self, ttl: int):
+        self.__ttl = ttl
+
+    def get_ttl(self) -> int:
+        return int(self.__ttl)

sdk/python/kfp/dsl/pipeline_context.py

@@ -18,6 +18,7 @@
 from typing import Callable, Optional
 
 from kfp.dsl import component_factory
+from kfp.dsl import pipeline_config
 from kfp.dsl import pipeline_task
 from kfp.dsl import tasks_group
 from kfp.dsl import utils
@@ -29,7 +30,7 @@ def pipeline(func: Optional[Callable] = None,
              description: Optional[str] = None,
              pipeline_root: Optional[str] = None,
              display_name: Optional[str] = None,
-             ttl_active_seconds: Optional[int] = None) -> Callable:
+             pipeline_config: Optional[pipeline_config.PipelineConfig] = None) -> Callable:
     """Decorator used to construct a pipeline.
 
     Example
@@ -51,7 +52,7 @@ def my_pipeline(a: str, b: int):
         pipeline_root: The root directory from which to read input and output
             parameters and artifacts.
         display_name: A human-readable name for the pipeline.
-        ttl_active_seconds: Delete completed workflows after a set time.
+        pipeline_config: Pipeline-level configuration.
     """
     if func is None:
         return functools.partial(
@@ -60,7 +61,7 @@ def my_pipeline(a: str, b: int):
             description=description,
             pipeline_root=pipeline_root,
             display_name=display_name,
-            ttl_active_seconds=ttl_active_seconds,
+            pipeline_config=pipeline_config,
         )
 
     if pipeline_root:
@@ -71,7 +72,7 @@ def my_pipeline(a: str, b: int):
         name=name,
         description=description,
         display_name=display_name,
-        ttl_strategy_seconds_after_completion=ttl_active_seconds,
+        pipeline_config=pipeline_config,
     )