port: kfp2

notebooks/official/pipelines/lightweight_functions_component_io_kfp.ipynb

@@ -146,8 +146,8 @@
       "source": [
         "! pip3 install --upgrade --quiet google-cloud-aiplatform \\\n",
         "                                 google-cloud-storage \\\n",
-        "                                 'kfp<2' \\\n",
-        "                                 'google-cloud-pipeline-components<2' "
+        "                                 kfp \\\n",
+        "                                 google-cloud-pipeline-components"
       ]
     },
     {
@@ -498,9 +498,9 @@
         "from typing import NamedTuple\n",
         "\n",
         "import kfp\n",
-        "from kfp.v2 import dsl\n",
-        "from kfp.v2.dsl import (Artifact, Dataset, Input, InputPath, Model, Output,\n",
-        "                        OutputPath, component)"
+        "from kfp import compiler, dsl\n",
+        "from kfp.dsl import (Artifact, Dataset, Input, InputPath, Model, Output,\n",
+        "                     OutputPath, component)"
       ]
     },
     {
@@ -748,7 +748,7 @@
         "        num_steps=5,\n",
         "    )\n",
         "    read_task = read_artifact_input(  # noqa: F841\n",
-        "        train_task.outputs[\"generic_artifact\"]\n",
+        "        generic=train_task.outputs[\"generic_artifact\"]\n",
         "    )"
       ]
     },
@@ -771,10 +771,8 @@
       },
       "outputs": [],
       "source": [
-        "from kfp.v2 import compiler  # noqa: F811\n",
-        "\n",
         "compiler.Compiler().compile(\n",
-        "    pipeline_func=pipeline, package_path=\"lightweight_pipeline.json\"\n",
+        "    pipeline_func=pipeline, package_path=\"lightweight_pipeline.yaml\"\n",
         ")"
       ]
     },
@@ -801,7 +799,7 @@
         "\n",
         "job = aip.PipelineJob(\n",
         "    display_name=DISPLAY_NAME,\n",
-        "    template_path=\"lightweight_pipeline.json\",\n",
+        "    template_path=\"lightweight_pipeline.yaml\",\n",
         "    pipeline_root=PIPELINE_ROOT,\n",
         "    parameter_values={\"message\": \"Hello, World\"},\n",
         ")\n",
@@ -855,7 +853,9 @@
         "delete_bucket = False\n",
         "\n",
         "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
-        "    ! gsutil rm -r $BUCKET_URI"
+        "    ! gsutil rm -r $BUCKET_URI\n",
+        "\n",
+        "! rm lightweight_pipeline.yaml"
       ]
     }
   ],