fix: train job not using TRAIN_IMAGE

notebooks/official/pipelines/get_started_with_machine_management.ipynb

@@ -284,10 +284,8 @@
       },
       "outputs": [],
       "source": [
-        "IS_COLAB = False\n",
         "# from google.colab import auth\n",
-        "# auth.authenticate_user()\n",
-        "# IS_COLAB=True"
+        "# auth.authenticate_user()"
       ]
     },
     {
@@ -372,6 +370,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",
@@ -622,7 +623,7 @@
       "source": [
         "@component(\n",
         "    output_component_file=\"demo_componet.yaml\",\n",
-        "    base_image=\"python:3.9\",\n",
+        "    base_image=TRAIN_IMAGE,\n",
         "    packages_to_install=[\"tensorflow\"],\n",
         ")\n",
         "def self_contained_training_component(\n",
@@ -1162,6 +1163,8 @@
       },
       "outputs": [],
       "source": [
+        "import os\n",
+        "\n",
         "# Set this to true only if you'd like to delete your bucket\n",
         "delete_bucket = False\n",
         "\n",