Merge remote-tracking branch 'kubeflow/master' into rebase-2.4-master

Author: hbelmiro

backend/metadata_writer/src/metadata_helpers.py

@@ -12,10 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import functools
 import json
 import os
 import sys
-import ml_metadata
 from time import sleep
 from ml_metadata.proto import metadata_store_pb2
 from ml_metadata.metadata_store import metadata_store
@@ -169,7 +169,6 @@ def create_context_with_type(
     return context
 
 
-import functools
 @functools.lru_cache(maxsize=128)
 def get_context_by_name(
     store,

manifests/kustomize/third-party/metacontroller/base/cluster-role-binding.yaml

@@ -5,7 +5,7 @@ metadata:
 roleRef:
   apiGroup: rbac.authorization.k8s.io
   kind: ClusterRole
-  name: cluster-admin
+  name: kubeflow-metacontroller
 subjects:
 - kind: ServiceAccount
   name: meta-controller-service

manifests/kustomize/third-party/metacontroller/base/cluster-role.yaml

@@ -0,0 +1,26 @@
+apiVersion: rbac.authorization.k8s.io/v1
+kind: ClusterRole
+metadata:
+  name: kubeflow-metacontroller
+rules:
+- apiGroups: [""]
+  resources: ["namespaces"]
+  verbs: ["get", "list", "watch"]
+- apiGroups: [""]
+  resources: ["secrets", "configmaps"]
+  verbs: ["get", "list", "watch", "create", "update", "patch", "delete"]
+- apiGroups: ["apps"]
+  resources: ["deployments"]
+  verbs: ["get", "list", "watch", "create", "update", "patch", "delete"]
+- apiGroups: [""]
+  resources: ["services"]
+  verbs: ["get", "list", "watch", "create", "update", "patch", "delete"]
+- apiGroups: ["networking.istio.io"]
+  resources: ["destinationrules"]
+  verbs: ["get", "list", "watch", "create", "update", "patch", "delete"]
+- apiGroups: ["security.istio.io"]
+  resources: ["authorizationpolicies"]
+  verbs: ["get", "list", "watch", "create", "update", "patch", "delete"]
+- apiGroups: ["metacontroller.k8s.io"]
+  resources: ["compositecontrollers", "controllerrevisions", "decoratorcontrollers"]
+  verbs: ["get", "list", "watch"]

manifests/kustomize/third-party/metacontroller/base/kustomization.yaml

@@ -2,13 +2,13 @@ apiVersion: kustomize.config.k8s.io/v1beta1
 kind: Kustomization
 namespace: kubeflow
 resources:
+- cluster-role.yaml
 - cluster-role-binding.yaml
 - crd.yaml
 - service-account.yaml
 - stateful-set.yaml
 commonLabels:
   kustomize.component: metacontroller
-
 # Update metacontroller CRD:
 # Copy the upstream file to crd.yaml in this folder.
-# Upstream file: https://github.com/metacontroller/metacontroller/blob/master/manifests/production/metacontroller-crds-v1.yaml
+# Upstream file: https://github.com/metacontroller/metacontroller/blob/master/manifests/production/metacontroller-crds-v1.yaml

manifests/kustomize/third-party/metacontroller/base/stateful-set.yaml

@@ -41,7 +41,7 @@ spec:
             privileged: false
             allowPrivilegeEscalation: false
           name: metacontroller
-          image: 'ghcr.io/metacontroller/metacontroller:v4.11.21'
+          image: 'ghcr.io/metacontroller/metacontroller:v2.6.1'
       serviceAccountName: meta-controller-service
   # Workaround for https://github.com/kubernetes-sigs/kustomize/issues/677
   volumeClaimTemplates: []

sdk/python/kfp/client/client.py

@@ -686,6 +686,7 @@ def run_pipeline(
         version_id: Optional[str] = None,
         pipeline_root: Optional[str] = None,
         enable_caching: Optional[bool] = None,
+        cache_key: Optional[str] = None,
         service_account: Optional[str] = None,
     ) -> kfp_server_api.V2beta1Run:
         """Runs a specified pipeline.
@@ -709,6 +710,8 @@ def run_pipeline(
                 is ``True`` for all tasks by default. If set, the
                 setting applies to all tasks in the pipeline (overrides the
                 compile time settings).
+            cache_key (optional): Customized cache key for this task.
+                If set, the cache_key will be used as the key for the task's cache.
             service_account: Specifies which Kubernetes service
                 account to use for this run.
 
@@ -721,6 +724,7 @@ def run_pipeline(
             pipeline_id=pipeline_id,
             version_id=version_id,
             enable_caching=enable_caching,
+            cache_key=cache_key,
             pipeline_root=pipeline_root,
         )
 
@@ -806,6 +810,7 @@ def create_recurring_run(
         enabled: bool = True,
         pipeline_root: Optional[str] = None,
         enable_caching: Optional[bool] = None,
+        cache_key: Optional[str] = None,
         service_account: Optional[str] = None,
     ) -> kfp_server_api.V2beta1RecurringRun:
         """Creates a recurring run.
@@ -850,6 +855,8 @@ def create_recurring_run(
                 different caching options for individual tasks. If set, the
                 setting applies to all tasks in the pipeline (overrides the
                 compile time settings).
+            cache_key (optional): Customized cache key for this task.
+                If set, the cache_key will be used as the key for the task's cache.
             service_account: Specifies which Kubernetes service
                 account this recurring run uses.
         Returns:
@@ -862,6 +869,7 @@ def create_recurring_run(
             pipeline_id=pipeline_id,
             version_id=version_id,
             enable_caching=enable_caching,
+            cache_key=cache_key,
             pipeline_root=pipeline_root,
         )
 
@@ -908,6 +916,7 @@ def _create_job_config(
         pipeline_id: Optional[str],
         version_id: Optional[str],
         enable_caching: Optional[bool],
+        cache_key: Optional[str],
         pipeline_root: Optional[str],
     ) -> _JobConfig:
         """Creates a JobConfig with spec and resource_references.
@@ -928,6 +937,8 @@ def _create_job_config(
                 different caching options for individual tasks. If set, the
                 setting applies to all tasks in the pipeline (overrides the
                 compile time settings).
+            cache_key (optional): Customized cache key for this task.
+                If set, the cache_key will be used as the key for the task's cache.
             pipeline_root: Root path of the pipeline outputs.
 
         Returns:
@@ -956,7 +967,7 @@ def _create_job_config(
             # settings.
             if enable_caching is not None:
                 _override_caching_options(pipeline_doc.pipeline_spec,
-                                          enable_caching)
+                                          enable_caching, cache_key)
             pipeline_spec = pipeline_doc.to_dict()
 
         pipeline_version_reference = None
@@ -983,6 +994,7 @@ def create_run_from_pipeline_func(
         namespace: Optional[str] = None,
         pipeline_root: Optional[str] = None,
         enable_caching: Optional[bool] = None,
+        cache_key: Optional[str] = None,
         service_account: Optional[str] = None,
         experiment_id: Optional[str] = None,
     ) -> RunPipelineResult:
@@ -1004,6 +1016,8 @@ def create_run_from_pipeline_func(
                 different caching options for individual tasks. If set, the
                 setting applies to all tasks in the pipeline (overrides the
                 compile time settings).
+            cache_key (optional): Customized cache key for this task.
+                If set, the cache_key will be used as the key for the task's cache.
             service_account: Specifies which Kubernetes service
                 account to use for this run.
             experiment_id: ID of the experiment to add the run to. You cannot specify both experiment_id and experiment_name.
@@ -1032,6 +1046,7 @@ def create_run_from_pipeline_func(
                 namespace=namespace,
                 pipeline_root=pipeline_root,
                 enable_caching=enable_caching,
+                cache_key=cache_key,
                 service_account=service_account,
             )
 
@@ -1044,6 +1059,7 @@ def create_run_from_pipeline_package(
         namespace: Optional[str] = None,
         pipeline_root: Optional[str] = None,
         enable_caching: Optional[bool] = None,
+        cache_key: Optional[str] = None,
         service_account: Optional[str] = None,
         experiment_id: Optional[str] = None,
     ) -> RunPipelineResult:
@@ -1065,6 +1081,8 @@ def create_run_from_pipeline_package(
                 different caching options for individual tasks. If set, the
                 setting applies to all tasks in the pipeline (overrides the
                 compile time settings).
+            cache_key (optional): Customized cache key for this task.
+                If set, the cache_key will be used as the key for the task's cache.
             service_account: Specifies which Kubernetes service
                 account to use for this run.
             experiment_id: ID of the experiment to add the run to. You cannot specify both experiment_id and experiment_name.
@@ -1105,6 +1123,7 @@ def create_run_from_pipeline_package(
             params=arguments,
             pipeline_root=pipeline_root,
             enable_caching=enable_caching,
+            cache_key=cache_key,
             service_account=service_account,
         )
         return RunPipelineResult(self, run_info)
@@ -1681,12 +1700,16 @@ def _safe_load_yaml(stream: TextIO) -> _PipelineDoc:
 def _override_caching_options(
     pipeline_spec: pipeline_spec_pb2.PipelineSpec,
     enable_caching: bool,
+    cache_key: Optional[str] = None,
 ) -> None:
     """Overrides caching options.
 
     Args:
         pipeline_spec: The PipelineSpec object to update in-place.
         enable_caching: Overrides options, one of True, False.
+        cache_key: Overrides cache_key, default None, no-op.
     """
     for _, task_spec in pipeline_spec.root.dag.tasks.items():
         task_spec.caching_options.enable_cache = enable_caching
+        if cache_key:
+            task_spec.caching_options.cache_key = cache_key

sdk/python/kfp/client/client_test.py

@@ -94,6 +94,52 @@ def pipeline_with_two_component(text: str = 'hi there'):
                                 ['hello-word']['cachingOptions']['enableCache'])
                 self.assertTrue(pipeline_obj['root']['dag']['tasks']['to-lower']
                                 ['cachingOptions']['enableCache'])
+                self.assertFalse('cacheKey' in pipeline_obj['root']['dag']
+                                 ['tasks']['hello-word']['cachingOptions'])
+                self.assertFalse('cacheKey' in pipeline_obj['root']['dag']
+                                 ['tasks']['to-lower']['cachingOptions'])
+
+    def test_override_cache_key(self):
+
+        @component
+        def hello_word(text: str) -> str:
+            return text
+
+        @component
+        def to_lower(text: str) -> str:
+            return text.lower()
+
+        @pipeline(
+            name='sample two-step pipeline',
+            description='a minimal two-step pipeline')
+        def pipeline_with_two_component(text: str = 'hi there'):
+            component_1 = hello_word(text=text).set_caching_options(True)
+            component_2 = to_lower(
+                text=component_1.output).set_caching_options(False)
+
+        with tempfile.TemporaryDirectory() as tempdir:
+            temp_filepath = os.path.join(tempdir, 'hello_world_pipeline.yaml')
+            Compiler().compile(
+                pipeline_func=pipeline_with_two_component,
+                package_path=temp_filepath)
+
+            with open(temp_filepath, 'r') as f:
+                pipeline_obj = yaml.safe_load(f)
+                pipeline_spec = json_format.ParseDict(
+                    pipeline_obj, pipeline_spec_pb2.PipelineSpec())
+                client._override_caching_options(
+                    pipeline_spec, True, cache_key='OVERRIDE_KEY')
+                pipeline_obj = json_format.MessageToDict(pipeline_spec)
+                self.assertTrue(pipeline_obj['root']['dag']['tasks']
+                                ['hello-word']['cachingOptions']['enableCache'])
+                self.assertEqual(
+                    pipeline_obj['root']['dag']['tasks']['hello-word']
+                    ['cachingOptions']['cacheKey'], 'OVERRIDE_KEY')
+                self.assertTrue(pipeline_obj['root']['dag']['tasks']['to-lower']
+                                ['cachingOptions']['enableCache'])
+                self.assertEqual(
+                    pipeline_obj['root']['dag']['tasks']['to-lower']
+                    ['cachingOptions']['cacheKey'], 'OVERRIDE_KEY')
 
 
 class TestExtractPipelineYAML(parameterized.TestCase):

sdk/python/kfp/compiler/compiler_test.py

@@ -1001,6 +1001,32 @@ def my_pipeline():
 
             self.assertTrue(caching_options['enableCache'])
 
+    def test_compile_pipeline_with_cache_key(self):
+        """Test pipeline compilation with cache key."""
+
+        @dsl.component
+        def my_component():
+            pass
+
+        @dsl.pipeline(name='tiny-pipeline')
+        def my_pipeline():
+            my_task = my_component()
+            my_task.set_caching_options(True, cache_key='MY_KEY')
+
+        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_spec = yaml.safe_load(f)
+
+            task_spec = pipeline_spec['root']['dag']['tasks']['my-component']
+            caching_options = task_spec['cachingOptions']
+
+            self.assertTrue(caching_options['enableCache'])
+            self.assertEqual(caching_options['cacheKey'], 'MY_KEY')
+
     def test_compile_pipeline_with_caching_disabled(self):
         """Test pipeline compilation with caching disabled."""
 

sdk/python/kfp/compiler/pipeline_spec_builder.py

@@ -124,6 +124,9 @@ def build_task_spec_for_task(
         utils.sanitize_component_name(task.name))
     pipeline_task_spec.caching_options.enable_cache = (
         task._task_spec.enable_caching)
+    if task._task_spec.cache_key:
+        pipeline_task_spec.caching_options.cache_key = (
+            task._task_spec.cache_key)
 
     if task._task_spec.retry_policy is not None:
         pipeline_task_spec.retry_policy.CopyFrom(

sdk/python/kfp/dsl/pipeline_task.py

@@ -99,6 +99,7 @@ def __init__(
         args: Dict[str, Any],
         execute_locally: bool = False,
         execution_caching_default: bool = True,
+        execution_cache_key: Optional[str] = None,
     ) -> None:
         """Initilizes a PipelineTask instance."""
         # import within __init__ to avoid circular import
@@ -131,7 +132,8 @@ def __init__(
             inputs=dict(args.items()),
             dependent_tasks=[],
             component_ref=component_spec.name,
-            enable_caching=execution_caching_default)
+            enable_caching=execution_caching_default,
+            cache_key=execution_cache_key)
         self._run_after: List[str] = []
 
         self.importer_spec = None
@@ -301,16 +303,20 @@ def _extract_container_spec_and_convert_placeholders(
         return container_spec
 
     @block_if_final()
-    def set_caching_options(self, enable_caching: bool) -> 'PipelineTask':
+    def set_caching_options(self,
+                            enable_caching: bool,
+                            cache_key: Optional[str] = None) -> 'PipelineTask':
         """Sets caching options for the task.
 
         Args:
             enable_caching: Whether to enable caching.
+            cache_key: Customized cache key for this task. Default None.
 
         Returns:
             Self return to allow chained setting calls.
         """
         self._task_spec.enable_caching = enable_caching
+        self._task_spec.cache_key = cache_key
         return self
 
     def _ensure_container_spec_exists(self) -> None: