port: kfp2 (#2302)

* port: kfp2 * fix import importer_node * fix: selector * fix: format * lint * set node * fix: types * fix: type * types * yaml/json * yaml/json * train image * train image * key error * fix: outputs * fix: outputs * fix: outputs * fix: outputs * fix: outputs
2026-09-26 14:42:04 +00:00 · 2023-10-06 20:40:51 +00:00
parent feee6149a1
commit 682786616c
1 changed files with 23 additions and 30 deletions
@@ -49,8 +49,8 @@
        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
        "      Open in Vertex AI Workbench\n",
        "    </a>\n",
-        "  </td>                                                                                               \n",
-        "</table>"
+        "  </td>                                                                                     \n",
+        "</table>\n"
      ]
    },
    {
@@ -138,8 +138,8 @@
        "import os\n",
        "\n",
        "! pip3 install --upgrade google-cloud-aiplatform \\\n",
-        "                         'google-cloud-pipeline-components<2' --quiet\n",
-        "! pip3 install --upgrade 'kfp<2' --quiet\n",
+        "                         google-cloud-pipeline-components --quiet\n",
+        "! pip3 install --upgrade kfp --quiet\n",
        "! pip3 install --upgrade tensorflow==2.7 --quiet"
      ]
    },
@@ -337,7 +337,7 @@
      },
      "outputs": [],
      "source": [
-        "! gsutil mb -l $REGION $BUCKET_URI"
+        "! gsutil mb -l {REGION} {BUCKET_URI}"
      ]
    },
    {
@@ -443,8 +443,8 @@
        "from google.cloud import aiplatform\n",
        "from google_cloud_pipeline_components.v1.custom_job import \\\n",
        "    create_custom_training_job_from_component\n",
-        "from kfp.v2 import compiler, dsl\n",
-        "from kfp.v2.dsl import component"
+        "from kfp import compiler, dsl\n",
+        "from kfp.dsl import component"
      ]
    },
    {
@@ -530,15 +530,8 @@
      "outputs": [],
      "source": [
        "TF = \"2.5\".replace(\".\", \"-\")\n",
-        "\n",
-        "if TRAIN_GPU:\n",
-        "    TRAIN_VERSION = \"tf-gpu.{}\".format(TF)\n",
-        "else:\n",
-        "    TRAIN_VERSION = \"tf-cpu.{}\".format(TF)\n",
-        "if DEPLOY_GPU:\n",
-        "    DEPLOY_VERSION = \"tf2-gpu.{}\".format(TF)\n",
-        "else:\n",
-        "    DEPLOY_VERSION = \"tf2-cpu.{}\".format(TF)\n",
+        "TRAIN_VERSION = \"tf-gpu.{}\".format(TF)\n",
+        "DEPLOY_VERSION = \"tf2-gpu.{}\".format(TF)\n",
        "\n",
        "\n",
        "TRAIN_IMAGE = \"{}-docker.pkg.dev/vertex-ai/training/{}:latest\".format(\n",
@@ -622,7 +615,6 @@
      "outputs": [],
      "source": [
        "@component(\n",
-        "    output_component_file=\"demo_componet.yaml\",\n",
        "    base_image=TRAIN_IMAGE,\n",
        "    packages_to_install=[\"tensorflow\"],\n",
        ")\n",
@@ -670,7 +662,10 @@
        "    train_model(x_train, y_train, model, epochs)\n",
        "\n",
        "    model.save(model_dir)\n",
-        "    return model_dir"
+        "    return model_dir\n",
+        "\n",
+        "\n",
+        "compiler.Compiler().compile(self_contained_training_component, \"demo_componet.yaml\")"
      ]
    },
    {
@@ -714,16 +709,14 @@
        "def pipeline(epochs: int, model_dir: str, project: str = PROJECT_ID):\n",
        "    from google_cloud_pipeline_components.types import artifact_types\n",
        "    from google_cloud_pipeline_components.v1.model import ModelUploadOp\n",
-        "    from kfp.v2.components import importer_node\n",
+        "    from kfp.dsl import importer_node\n",
        "\n",
        "    training_job_task = (\n",
        "        self_contained_training_component(epochs=epochs, model_dir=model_dir)\n",
        "        .set_display_name(\"self-contained-training\")\n",
        "        .set_cpu_limit(CPU_LIMIT)\n",
        "        .set_memory_limit(MEMORY_LIMIT)\n",
-        "        .add_node_selector_constraint(\n",
-        "            value=TRAIN_GPU.name, label_name=\"cloud.google.com/gke-accelerator\"\n",
-        "        )\n",
+        "        .add_node_selector_constraint(\"NVIDIA_TESLA_K80\")\n",
        "        .set_gpu_limit(TRAIN_NGPU)\n",
        "    )\n",
        "\n",
@@ -769,12 +762,12 @@
      "source": [
        "compiler.Compiler().compile(\n",
        "    pipeline_func=pipeline,\n",
-        "    package_path=\"component_level_settings.json\",\n",
+        "    package_path=\"component_level_settings.yaml\",\n",
        ")\n",
        "\n",
        "pipeline = aiplatform.PipelineJob(\n",
        "    display_name=\"component-level-settings\",\n",
-        "    template_path=\"component_level_settings.json\",\n",
+        "    template_path=\"component_level_settings.yaml\",\n",
        "    pipeline_root=PIPELINE_ROOT,\n",
        "    parameter_values={\"model_dir\": BUCKET_URI, \"epochs\": 20, \"project\": PROJECT_ID},\n",
        "    enable_caching=False,\n",
@@ -782,7 +775,7 @@
        "\n",
        "pipeline.run()\n",
        "\n",
-        "! rm -rf component_level_settings.json"
+        "! rm -rf component_level_settings.yaml"
      ]
    },
    {
@@ -997,14 +990,14 @@
        "):\n",
        "    from google_cloud_pipeline_components.types import artifact_types\n",
        "    from google_cloud_pipeline_components.v1.model import ModelUploadOp\n",
-        "    from kfp.v2.components import importer_node\n",
+        "    from kfp.dsl import importer_node\n",
        "\n",
        "    training_job_task = custom_job_op(\n",
        "        epochs=epochs, model_dir=model_dir, project=project, location=region\n",
        "    )\n",
        "\n",
        "    import_unmanaged_model_task = importer_node.importer(\n",
-        "        artifact_uri=training_job_task.outputs[\"output\"],\n",
+        "        artifact_uri=training_job_task.outputs[\"Output\"],\n",
        "        artifact_class=artifact_types.UnmanagedContainerModel,\n",
        "        metadata={\n",
        "            \"containerSpec\": {\n",
@@ -1045,12 +1038,12 @@
      "source": [
        "compiler.Compiler().compile(\n",
        "    pipeline_func=pipeline,\n",
-        "    package_path=\"customjob_level_settings.json\",\n",
+        "    package_path=\"customjob_level_settings.yaml\",\n",
        ")\n",
        "\n",
        "pipeline = aiplatform.PipelineJob(\n",
        "    display_name=\"customjob-level-settings\",\n",
-        "    template_path=\"customjob_level_settings.json\",\n",
+        "    template_path=\"customjob_level_settings.yaml\",\n",
        "    pipeline_root=PIPELINE_ROOT,\n",
        "    parameter_values={\"model_dir\": BUCKET_URI, \"epochs\": 20, \"project\": PROJECT_ID},\n",
        "    enable_caching=False,\n",
@@ -1058,7 +1051,7 @@
        "\n",
        "pipeline.run()\n",
        "\n",
-        "! rm -rf customjob_level_settings.json"
+        "! rm -rf customjob_level_settings.yaml"
      ]
    },
    {