fix importer_node

notebooks/official/pipelines/custom_model_training_and_batch_prediction.ipynb

@@ -151,11 +151,11 @@
       "source": [
         "! pip3 install --upgrade --quiet google-cloud-aiplatform \\\n",
         "                                 google-cloud-storage \\\n",
-        "                                 'kfp<2' \\\n",
-        "                                 'google-cloud-pipeline-components<2'\n",
+        "                                 kfp \\\n",
+        "                                 google-cloud-pipeline-components\n",
         "\n",
         "\n",
-        "! pip3 install --upgrade --force-reinstall tensorflow 'kfp<2' google-cloud-aiplatform google-cloud-storage 'google-cloud-pipeline-components<2' -q"
+        "! pip3 install --upgrade --force-reinstall tensorflow kfp google-cloud-aiplatform google-cloud-storage google-cloud-pipeline-components -q"
       ]
     },
     {
@@ -297,10 +297,8 @@
       },
       "outputs": [],
       "source": [
-        "IS_COLAB = False\n",
         "# from google.colab import auth\n",
-        "# auth.authenticate_user()\n",
-        "# IS_COLAB = True"
+        "# auth.authenticate_user()"
       ]
     },
     {
@@ -324,6 +322,17 @@
         "Create a storage bucket to store intermediate artifacts such as datasets."
       ]
     },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "MzGDU7TWdts_"
+      },
+      "outputs": [],
+      "source": [
+        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
+      ]
+    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -337,11 +346,11 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "id": "MzGDU7TWdts_"
+        "id": "NIq7R4HZCfIc"
       },
       "outputs": [],
       "source": [
-        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
+        "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
       ]
     },
     {
@@ -374,6 +383,9 @@
       },
       "outputs": [],
       "source": [
+        "import sys\n",
+        "\n",
+        "IS_COLAB = \"google.colab\" in sys.modules\n",
         "if (\n",
         "    SERVICE_ACCOUNT == \"\"\n",
         "    or SERVICE_ACCOUNT is None\n",
@@ -437,9 +449,9 @@
       "outputs": [],
       "source": [
         "import google.cloud.aiplatform as aip\n",
-        "from google_cloud_pipeline_components.experimental.custom_job import utils\n",
-        "from kfp.v2 import compiler, dsl\n",
-        "from kfp.v2.dsl import component"
+        "from google_cloud_pipeline_components.v1.custom_job import utils\n",
+        "from kfp import compiler, dsl\n",
+        "from kfp.dsl import component"
       ]
     },
     {
@@ -516,23 +528,9 @@
       },
       "outputs": [],
       "source": [
-        "import os\n",
+        "TRAIN_GPU, TRAIN_NGPU = (None, None)\n",
         "\n",
-        "if os.getenv(\"IS_TESTING_TRAIN_GPU\"):\n",
-        "    TRAIN_GPU, TRAIN_NGPU = (\n",
-        "        aip.gapic.AcceleratorType.NVIDIA_TESLA_K80,\n",
-        "        int(os.getenv(\"IS_TESTING_TRAIN_GPU\")),\n",
-        "    )\n",
-        "else:\n",
-        "    TRAIN_GPU, TRAIN_NGPU = (None, None)\n",
-        "\n",
-        "if os.getenv(\"IS_TESTING_DEPLOY_GPU\"):\n",
-        "    DEPLOY_GPU, DEPLOY_NGPU = (\n",
-        "        aip.gapic.AcceleratorType.NVIDIA_TESLA_K80,\n",
-        "        int(os.getenv(\"IS_TESTING_DEPLOY_GPU\")),\n",
-        "    )\n",
-        "else:\n",
-        "    DEPLOY_GPU, DEPLOY_NGPU = (None, None)"
+        "DEPLOY_GPU, DEPLOY_NGPU = (None, None)"
       ]
     },
     {
@@ -560,29 +558,16 @@
       },
       "outputs": [],
       "source": [
-        "if os.getenv(\"IS_TESTING_TF\"):\n",
-        "    TF = os.getenv(\"IS_TESTING_TF\")\n",
-        "else:\n",
-        "    TF = \"2-5\"\n",
+        "TF = \"2-5\"\n",
         "\n",
-        "if TF[0] == \"2\":\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",
+        "if TRAIN_GPU:\n",
+        "    TRAIN_VERSION = \"tf-gpu.{}\".format(TF)\n",
         "else:\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 = \"tf-gpu.{}\".format(TF)\n",
-        "    else:\n",
-        "        DEPLOY_VERSION = \"tf-cpu.{}\".format(TF)\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",
         "\n",
         "TRAIN_IMAGE = \"gcr.io/cloud-aiplatform/training/{}:latest\".format(TRAIN_VERSION)\n",
         "DEPLOY_IMAGE = \"gcr.io/cloud-aiplatform/prediction/{}:latest\".format(DEPLOY_VERSION)\n",
@@ -624,20 +609,12 @@
       },
       "outputs": [],
       "source": [
-        "if os.getenv(\"IS_TESTING_TRAIN_MACHINE\"):\n",
-        "    MACHINE_TYPE = os.getenv(\"IS_TESTING_TRAIN_MACHINE\")\n",
-        "else:\n",
-        "    MACHINE_TYPE = \"n1-standard\"\n",
+        "MACHINE_TYPE = \"n1-standard\"\n",
         "\n",
         "VCPU = \"4\"\n",
         "TRAIN_COMPUTE = MACHINE_TYPE + \"-\" + VCPU\n",
         "print(\"Train machine type\", TRAIN_COMPUTE)\n",
         "\n",
-        "if os.getenv(\"IS_TESTING_DEPLOY_MACHINE\"):\n",
-        "    MACHINE_TYPE = os.getenv(\"IS_TESTING_DEPLOY_MACHINE\")\n",
-        "else:\n",
-        "    MACHINE_TYPE = \"n1-standard\"\n",
-        "\n",
         "VCPU = \"4\"\n",
         "DEPLOY_COMPUTE = MACHINE_TYPE + \"-\" + VCPU\n",
         "print(\"Deploy machine type\", DEPLOY_COMPUTE)"
@@ -997,7 +974,7 @@
         "    from google_cloud_pipeline_components.v1.batch_predict_job import \\\n",
         "        ModelBatchPredictOp\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",
         "    custom_producer_task = custom_job_distributed_training_op(\n",
         "        model_dir=model_dir,\n",
@@ -1063,14 +1040,14 @@
       "outputs": [],
       "source": [
         "compiler.Compiler().compile(\n",
-        "    pipeline_func=pipeline, package_path=\"custom_model_training_spec.json\"\n",
+        "    pipeline_func=pipeline, package_path=\"custom_model_training_spec.yaml\"\n",
         ")\n",
         "\n",
         "DISPLAY_NAME = \"cifar10\"\n",
         "\n",
         "job = aip.PipelineJob(\n",
         "    display_name=DISPLAY_NAME,\n",
-        "    template_path=\"custom_model_training_spec.json\",\n",
+        "    template_path=\"custom_model_training_spec.yaml\",\n",
         "    pipeline_root=PIPELINE_ROOT,\n",
         ")\n",
         "\n",
@@ -1199,6 +1176,8 @@
       },
       "outputs": [],
       "source": [
+        "import os\n",
+        "\n",
         "job.delete()\n",
         "\n",
         "model = aip.Model(model_id)\n",