Repository: kubeflow/pipelines

Issue:
[pH] v2 sample test - enable kaniko caching
https://cloud.google.com/build/docs/kaniko-cache

this will help improve local testing speed
https://github.com/kubeflow/pipelines/blob/master/v2/test/components/kaniko.yaml


Discussion:

Resolving patch:
diff --git a/samples/v2/hello_world.py b/samples/v2/hello_world.py
--- a/samples/v2/hello_world.py
+++ b/samples/v2/hello_world.py
@@ -24,7 +24,7 @@ def hello_world(text: str):
 
 
 @dsl.pipeline(name='hello-world', description='A simple intro pipeline')
-def pipeline_parameter_to_consumer(text: str = 'hi there'):
+def pipeline_hello_world(text: str = 'hi there'):
     '''Pipeline that passes small pipeline parameter string to consumer op'''
 
     consume_task = hello_world(
@@ -35,6 +35,6 @@ def pipeline_parameter_to_consumer(text: str = 'hi there'):
 if __name__ == "__main__":
     # execute only if run as a script
     compiler.Compiler().compile(
-        pipeline_func=pipeline_parameter_to_consumer,
+        pipeline_func=pipeline_hello_world,
         package_path='hello_world_pipeline.json'
     )


Test patch:
diff --git a/samples/test/config.yaml b/samples/test/config.yaml
--- a/samples/test/config.yaml
+++ b/samples/test/config.yaml
@@ -22,6 +22,8 @@
 # The field `path` corresponds to the test's python module path
 # e.g. if folder path is `samples/test/fail_test.py`, then module path is
 # `samples.test.fail_test`.
+- name: hello_world
+  path: samples.v2.hello_world_test
 - name: condition
   path: samples.core.condition.condition_test
 - name: nested_condition
@@ -75,14 +77,10 @@
   path: samples.test.lightweight_python_functions_v2_with_outputs_test
 - name: lightweight_python_functions_v2_pipeline
   path: samples.test.lightweight_python_functions_v2_pipeline_test
-# TODO(v2-compatible): support pipeline: pipeline_with_ontology.py
-# TODO(v2-compatible): support two_step_pipeline_with_importer.py
-# TODO(v2-compatible): support IR placeholder
 - name: placeholder_concat
   path: samples.test.placeholder_concat_test
 - name: metrics_visualization_v2
   path: samples.test.metrics_visualization_v2_test
-# TODO(Bobgy): enable this test when https://github.com/kubeflow/pipelines/pull/5424
-# is deployed to test infra.
+# TODO(Bobgy): This is currently passing, should it fail?
 # - name: fail_parameter_value_missing
 #   path: samples.test.fail_parameter_value_missing_test
diff --git a/samples/test/util.py b/samples/test/util.py
--- a/samples/test/util.py
+++ b/samples/test/util.py
@@ -23,6 +23,7 @@
 from google.protobuf.json_format import MessageToDict
 
 import kfp
+import kfp.v2.compiler
 import kfp_server_api
 from ml_metadata import metadata_store
 from ml_metadata.proto import metadata_store_pb2
@@ -118,6 +119,8 @@ def main(
         host: Optional[str] = None,
         external_host: Optional[str] = None,
         launcher_image: Optional[str] = None,
+        launcher_v2_image: Optional[str] = None,
+        driver_image: Optional[str] = None,
         experiment: str = 'v2_sample_test_samples',
         metadata_service_host: Optional[str] = None,
         metadata_service_port: int = 8080,
@@ -133,6 +136,10 @@ def main(
         :type output_directory: str, optional
         :param launcher_image: override launcher image, only used in V2_COMPATIBLE mode
         :type launcher_image: URI, optional
+        :param launcher_v2_image: override launcher v2 image, only used in V2_ENGINE mode
+        :type launcher_v2_image: URI, optional
+        :param driver_image: override driver image, only used in V2_ENGINE mode
+        :type driver_image: URI, optional
         :param experiment: experiment the run is added to, defaults to 'v2_sample_test_samples'
         :type experiment: str, optional
         :param metadata_service_host: host for metadata grpc service, defaults to METADATA_GRPC_SERVICE_HOST or 'metadata-grpc-service'
@@ -167,23 +174,32 @@ def run_pipeline(
         ) -> kfp_server_api.ApiRunDetail:
             arguments = arguments or {}
             extra_arguments = {}
-            if mode != kfp.dsl.PipelineExecutionMode.V1_LEGACY:
+            # TODO(Bobgy): support overriding pipeline root for v2_engine
+            if mode == kfp.dsl.PipelineExecutionMode.V2_COMPATIBLE:
                 extra_arguments = {
                     kfp.dsl.ROOT_PARAMETER_NAME: output_directory
                 }
 
             def _create_run():
-                return client.create_run_from_pipeline_func(
-                    pipeline_func,
-                    mode=mode,
-                    arguments={
-                        **extra_arguments,
-                        **arguments,
-                    },
-                    launcher_image=launcher_image,
-                    experiment_name=experiment,
-                    enable_caching=enable_caching,
-                )
+                if mode == kfp.dsl.PipelineExecutionMode.V2_ENGINE:
+                    return run_v2_pipeline(
+                        client=client,
+                        fn=pipeline_func,
+                        driver_image=driver_image,
+                        launcher_v2_image=launcher_v2_image
+                    )
+                else:
+                    return client.create_run_from_pipeline_func(
+                        pipeline_func,
+                        mode=mode,
+                        arguments={
+                            **extra_arguments,
+                            **arguments,
+                        },
+                        launcher_image=launcher_image,
+                        experiment_name=experiment,
+                        enable_caching=enable_caching,
+                    )
 
             run_result = _retry_with_backoff(fn=_create_run)
             print("Run details page URL:")
@@ -224,6 +240,28 @@ def _create_run():
     fire.Fire(main)
 
 
+def run_v2_pipeline(
+    client: kfp.Client, fn: Callable, driver_image: str, launcher_v2_image: str
+):
+    import tempfile
+    import subprocess
+    pipeline_spec = tempfile.mktemp(suffix='.json')
+    kfp.v2.compiler.Compiler().compile(
+        pipeline_func=fn, package_path=pipeline_spec
+    )
+    argo_workflow_spec = tempfile.mktemp(suffix='.yaml')
+    with open(argo_workflow_spec, 'w') as f:
+        args = [
+            'kfp-v2-compiler', '--spec', pipeline_spec, '--driver',
+            driver_image, '--launcher', launcher_v2_image
+        ]
+        # call v2 backend compiler CLI to compile pipeline spec to argo workflow
+        subprocess.check_call(args, stdout=f)
+    return client.create_run_from_pipeline_package(
+        pipeline_file=argo_workflow_spec, arguments={}
+    )
+
+
 def simplify_proto_struct(data: dict) -> dict:
     res = {}
     for key, value in data.items():
diff --git a/samples/v2/hello_world_test.py b/samples/v2/hello_world_test.py
new file mode 100644
--- /dev/null
+++ b/samples/v2/hello_world_test.py
@@ -0,0 +1,43 @@
+# Copyright 2021 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.
+"""Hello world v2 engine pipeline."""
+
+from __future__ import annotations
+import unittest
+from pprint import pprint
+
+import kfp
+import kfp_server_api
+
+from .hello_world import pipeline_hello_world
+from ..test.util import run_pipeline_func, TestCase, KfpMlmdClient
+
+
+def verify(run: kfp_server_api.ApiRun, mlmd_connection_config, **kwargs):
+    t = unittest.TestCase()
+    t.maxDiff = None  # we always want to see full diff
+    t.assertEqual(run.status, 'Succeeded')
+    client = KfpMlmdClient(mlmd_connection_config=mlmd_connection_config)
+    tasks = client.get_tasks(run_id=run.id)
+    pprint(tasks)
+
+
+if __name__ == '__main__':
+    run_pipeline_func([
+        TestCase(
+            pipeline_func=pipeline_hello_world,
+            verify_func=verify,
+            mode=kfp.dsl.PipelineExecutionMode.V2_ENGINE,
+        ),
+    ])
diff --git a/v2/compiler/argo_test.go b/v2/compiler/argo_test.go
--- a/v2/compiler/argo_test.go
+++ b/v2/compiler/argo_test.go
@@ -16,10 +16,13 @@ func Test_argo_compiler(t *testing.T) {
 	if err != nil {
 		t.Error(err)
 	}
+
 	expectedText := `
   apiVersion: argoproj.io/v1alpha1
   kind: Workflow
   metadata:
+    annotations:
+      pipelines.kubeflow.org/v2_pipeline: "true"
     creationTimestamp: null
     generateName: hello-world-
   spec:
@@ -53,7 +56,7 @@ func Test_argo_compiler(t *testing.T) {
         - --executor_input_path
         - '{{outputs.parameters.executor-input.path}}'
         command:
-        - /bin/kfp/driver
+        - /bin/kfp-driver
         image: gcr.io/ml-pipeline/kfp-driver:latest
         name: ""
         resources: {}
@@ -222,7 +225,7 @@ func Test_argo_compiler(t *testing.T) {
         - --context_id_path
         - '{{outputs.parameters.context-id.path}}'
         command:
-        - /bin/kfp/driver
+        - /bin/kfp-driver
         image: gcr.io/ml-pipeline/kfp-driver:latest
         name: ""
         resources: {}
diff --git a/v2/test/Dockerfile b/v2/test/Dockerfile
--- a/v2/test/Dockerfile
+++ b/v2/test/Dockerfile
@@ -12,10 +12,34 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+### Build backend compiler binary ###
+# Pin golang to 1.15 because of https://github.com/golang/go/issues/44129.
+# Remove 1.15 pin when the issue gets resolved
+FROM golang:1.15-alpine as builder
+
+# Install tools
+RUN apk add --no-cache make bash
+
+# Download go modules in a separate step for better docker caching.
+WORKDIR /build
+# api package is a dependency of v2/go.mod
+COPY api/go.mod api/go.sum api/
+COPY v2/go.mod v2/go.sum v2/
+RUN cd v2 && go mod download
+# Build the application
+COPY api api
+COPY v2 v2
+WORKDIR /build/v2
+RUN make build/compiler
+
 FROM python:3.7-slim
 
 WORKDIR /workdir
-COPY . /workdir
+COPY v2/test/requirements.txt v2/test/
+COPY sdk/python sdk/python
 # relative path in requirements.txt are relative to workdir, so we need to
 # cd to that folder first
 RUN cd v2/test && pip3 install --no-cache-dir -r requirements.txt
+COPY --from=builder /build/v2/build/compiler /bin/kfp-v2-compiler
+# copy all other code
+COPY . .
diff --git a/v2/test/components/kaniko.yaml b/v2/test/components/kaniko.yaml
--- a/v2/test/components/kaniko.yaml
+++ b/v2/test/components/kaniko.yaml
@@ -28,7 +28,9 @@ inputs:
     Input context artifact, it should be a folder.
     At least one of context_uri and context_artifact should be specified.
 - {name: context_sub_path, type: Path, default: ""}
-- {name: destination, type: URI, description: destination should not contain tag}
+- {name: destination, type: String, description: destination should not contain tag}
+- {name: cache, type: String, description: enable caching one of true/false, default: "true"}
+- {name: cache_ttl, type: String, description: 'cache TTL, format like XXh', default: 24h}
 outputs:
 - {name: digest, type: URI, description: Image URI with full digest}
 metadata:
@@ -37,7 +39,11 @@ metadata:
 implementation:
   container:
     # The debug image contains busybox, while the normal images do not.
-    image: gcr.io/kaniko-project/executor:v1.6.0-debug
+    # Pinned to 1.3.0, because unusually high memory usage is observed in later
+    # versions, causing OOM. 1.6.0 is known to break, we may try later releases
+    # when they are released.
+    # https://github.com/GoogleContainerTools/kaniko/issues/1680
+    image: gcr.io/kaniko-project/executor:v1.3.0-debug
     command:
     - sh
     - -exc
@@ -48,25 +54,33 @@ implementation:
         context_sub_path="$3"
         destination="$4"
         digest_output_path="$5"
+        cache="$6"
+        cache_ttl="$7"
         context=""
         if [[ "${context_uri}" != "" ]]; then
           context="${context_uri}"
         else
           context="dir://${context_artifact_path}"
         fi
-        mkdir -p "$(dirname "${digest_output_path}")";
+        digest_file="$(mktemp)"
         # Reference: https://github.com/GoogleContainerTools/kaniko
-        # snapshotMode added because: https://github.com/GoogleContainerTools/kaniko/issues/1333
+        # snapshotMode and use-new-run added because: https://github.com/GoogleContainerTools/kaniko/issues/1333
         /kaniko/executor \
             --dockerfile "${dockerfile}" \
             --context "${context}" \
             ${context_sub_path:+ --context-sub-path "${context_sub_path}"} \
             --destination "${destination}" \
-            --snapshotMode redo \
-            --image-name-with-digest-file "${digest_output_path}"
+            --digest-file "${digest_file}" \
+            --cache="${cache}" \
+            --cache-ttl="${cache_ttl}"
+        mkdir -p "$(dirname "${digest_output_path}")"
+        digest="$(cat $digest_file)"
+        echo "${destination}@${digest}" > "${digest_output_path}"
     - inputValue: dockerfile
     - inputValue: context_uri
     - inputPath: context_artifact
     - inputValue: context_sub_path
     - inputValue: destination
     - outputPath: digest
+    - inputValue: cache
+    - inputValue: cache_ttl
diff --git a/v2/test/components/run_sample.yaml b/v2/test/components/run_sample.yaml
--- a/v2/test/components/run_sample.yaml
+++ b/v2/test/components/run_sample.yaml
@@ -20,6 +20,8 @@ inputs:
 - {name: Host, type: URI, default: "http://ml-pipeline:8888"}
 - {name: External Host, type: URI}
 - {name: Launcher Image, type: URI, default: "gcr.io/ml-pipeline/kfp-launcher:latest"}
+- {name: Launcher v2 Image, type: URI, default: "gcr.io/ml-pipeline/kfp-launcher-v2:latest"}
+- {name: Driver Image, type: URI, default: "gcr.io/ml-pipeline/kfp-driver:latest"}
 implementation:
   container:
     image: python:3.7-alpine
@@ -35,3 +37,7 @@ implementation:
     - inputValue: External Host
     - --launcher_image
     - inputValue: Launcher Image
+    - --launcher_v2_image
+    - inputValue: Launcher v2 Image
+    - --driver_image
+    - inputValue: Driver Image
diff --git a/v2/test/sample_test.py b/v2/test/sample_test.py
--- a/v2/test/sample_test.py
+++ b/v2/test/sample_test.py
@@ -15,6 +15,7 @@
 # %%
 import yaml
 import os
+import kfp
 
 REPO_ROOT = os.path.join('..', '..')
 SAMPLES_CONFIG_PATH = os.path.join(REPO_ROOT, 'samples', 'test', 'config.yaml')
@@ -22,13 +23,6 @@
 with open(SAMPLES_CONFIG_PATH, 'r') as stream:
     SAMPLES_CONFIG = yaml.safe_load(stream)
 
-import kfp
-import kfp.components as comp
-import json
-from typing import Optional
-
-MINUTE = 60  # seconds
-
 download_gcs_tgz = kfp.components.load_component_from_file(
     'components/download_gcs_tgz.yaml'
 )
@@ -37,43 +31,80 @@
 )
 kaniko = kfp.components.load_component_from_file('components/kaniko.yaml')
 
+PIPELINE_TIME_OUT = 40 * 60  # 40 minutes
+
 
 @kfp.dsl.pipeline(name='v2 sample test')
 def v2_sample_test(
     context: 'URI' = 'gs://your-bucket/path/to/context.tar.gz',
-    launcher_destination: 'URI' = 'gcr.io/gongyuan-pipeline-test/kfp-launcher',
-    gcs_root: 'URI' = 'gs://gongyuan-test/v2',
-    samples_destination: 'URI' = 'gcr.io/gongyuan-pipeline-test/v2-sample-test',
+    gcs_root: 'URI' = 'gs://ml-pipeline-test/v2',
+    image_registry: 'URI' = 'gcr.io/ml-pipeline-test',
     kfp_host: 'URI' = 'http://ml-pipeline:8888',
     samples_config: list = SAMPLES_CONFIG,
 ):
-    download_src_op = download_gcs_tgz(gcs_path=context)
-    download_src_op.set_display_name('download_src')
+    # pipeline configs
+    conf = kfp.dsl.get_pipeline_conf()
+    conf.set_timeout(
+        PIPELINE_TIME_OUT
+    )  # add timeout to avoid pipelines stuck in running leak indefinetely
+
+    download_src_op = download_gcs_tgz(
+        gcs_path=context
+    ).set_cpu_limit('0.5').set_memory_limit('500Mi'
+                                           ).set_display_name('download_src')
     download_src_op.execution_options.caching_strategy.max_cache_staleness = "P0D"
-    build_kfp_launcher_op = kaniko(
-        context_artifact=download_src_op.outputs['folder'],
-        destination=launcher_destination,
+
+    def build_image(name: str, dockerfile: str) -> kfp.dsl.ContainerOp:
+        task: kfp.dsl.ContainerOp = kaniko(
+            context_artifact=download_src_op.outputs['folder'],
+            destination=f'{image_registry}/{name}',
+            dockerfile=dockerfile,
+        )
+        # CPU request/limit can be more flexible (request < limit), because being assigned to a node
+        # with insufficient CPU resource will only slow the task down, but not fail.
+        task.container.set_cpu_request('1').set_cpu_limit('2')
+        # Memory request/limit needs to be more rigid (request == limit), because in a node without
+        # enough memory, the task can hang indefinetely or OOM.
+        task.container.set_memory_request('4Gi').set_memory_limit('4Gi')
+        task.set_display_name(f'build-image-{name}')
+        task.set_retry(
+            1, policy='Always'
+        )  # Always -> retry on both system error and user code failure.
+        return task
+
+    # build v2 compatible image
+    build_kfp_launcher_op = build_image(
+        name='kfp-launcher',
         dockerfile='v2/container/launcher/Dockerfile',
     )
-    build_kfp_launcher_op.set_display_name('build_kfp_launcher')
-    build_samples_image_op = kaniko(
-        context_artifact=download_src_op.outputs['folder'],
-        destination=samples_destination,
+    # build sample test image
+    build_samples_image_op = build_image(
+        name='v2-sample-test',
         dockerfile='v2/test/Dockerfile',
     )
-    build_samples_image_op.execution_options.caching_strategy.max_cache_staleness = "P0D"
-    build_samples_image_op.set_display_name('build_samples_image')
+    # build v2 engine images
+    build_kfp_launcher_v2_op = build_image(
+        name='kfp-launcher-v2',
+        dockerfile='v2/container/launcher-v2/Dockerfile'
+    )
+    build_kfp_driver_op = build_image(
+        name='kfp-driver', dockerfile='v2/container/driver/Dockerfile'
+    )
+
+    # run test samples in parallel
     with kfp.dsl.ParallelFor(samples_config) as sample:
-        run_sample_op = run_sample(
+        run_sample_op: kfp.dsl.ContainerOp = run_sample(
             name=sample.name,
             sample_path=sample.path,
             gcs_root=gcs_root,
             external_host=kfp_host,
-            launcher_image=build_kfp_launcher_op.outputs['digest']
+            launcher_image=build_kfp_launcher_op.outputs['digest'],
+            launcher_v2_image=build_kfp_launcher_v2_op.outputs['digest'],
+            driver_image=build_kfp_driver_op.outputs['digest'],
         )
         run_sample_op.container.image = build_samples_image_op.outputs['digest']
         run_sample_op.set_display_name(f'sample_{sample.name}')
-        run_sample_op.set_retry(3)
+        run_sample_op.set_retry(1, policy='Always')
 
 
 def main(
@@ -92,16 +123,15 @@ def main(
         v2_sample_test,
         {
             'context': context,
-            'launcher_destination': f'{gcr_root}/kfp-launcher',
+            'image_registry': f'{gcr_root}/test',
             'gcs_root': gcs_root,
-            'samples_destination': f'{gcr_root}/v2-sample-test',
             'kfp_host': host,
         },
         experiment_name=experiment,
     )
     print("Run details page URL:")
     print(f"{host}/#/runs/details/{run_result.run_id}")
-    run_response = run_result.wait_for_run_completion(40 * MINUTE)
+    run_response = run_result.wait_for_run_completion(PIPELINE_TIME_OUT)
     run = run_response.run
     from pprint import pprint
     # Hide verbose content