refactor: Batch prediction using GCS supports global endpoint (#2394)

Author: gericdong

gemini/batch-prediction/intro_batch_prediction.ipynb

@@ -150,44 +150,6 @@
         "%pip install --upgrade --quiet google-genai pandas google-cloud-storage google-cloud-bigquery"
       ]
     },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "R5Xep4W9lq-Z"
-      },
-      "source": [
-        "### Restart runtime\n",
-        "\n",
-        "To use the newly installed packages in this Jupyter runtime, you must restart the runtime. You can do this by running the cell below, which restarts the current kernel.\n",
-        "\n",
-        "The restart might take a minute or longer. After it's restarted, continue to the next step."
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "XRvKdaPDTznN"
-      },
-      "outputs": [],
-      "source": [
-        "import IPython\n",
-        "\n",
-        "app = IPython.Application.instance()\n",
-        "app.kernel.do_shutdown(True)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "SbmM4z7FOBpM"
-      },
-      "source": [
-        "<div class=\"alert alert-block alert-warning\">\n",
-        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
-        "</div>\n"
-      ]
-    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -269,7 +231,7 @@
         "if not PROJECT_ID or PROJECT_ID == \"[your-project-id]\":\n",
         "    PROJECT_ID = str(os.environ.get(\"GOOGLE_CLOUD_PROJECT\"))\n",
         "\n",
-        "LOCATION = os.environ.get(\"GOOGLE_CLOUD_REGION\", \"us-central1\")"
+        "LOCATION = os.environ.get(\"GOOGLE_CLOUD_REGION\", \"global\")"
       ]
     },
     {
@@ -293,7 +255,7 @@
         "\n",
         "You can find a list of the Gemini models that support batch predictions in the [Multimodal models that support batch predictions](https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/batch-prediction-gemini#multimodal_models_that_support_batch_predictions) page.\n",
         "\n",
-        "This tutorial uses Gemini 2.0 Flash (`gemini-2.0-flash-001`) model."
+        "This tutorial uses Gemini 2.5 Flash (`gemini-2.5-flash`) model."
       ]
     },
     {
@@ -304,7 +266,7 @@
       },
       "outputs": [],
       "source": [
-        "MODEL_ID = \"gemini-2.0-flash-001\"  # @param {type:\"string\", isTemplate: true}"
+        "MODEL_ID = \"gemini-2.5-flash\"  # @param {type:\"string\", isTemplate: true}"
       ]
     },
     {
@@ -378,12 +340,13 @@
       "outputs": [],
       "source": [
         "BUCKET_URI = \"[your-cloud-storage-bucket]\"  # @param {type:\"string\"}\n",
+        "GCS_LOCATION = \"us-central1\"  # @param {type:\"string\"}\n",
         "\n",
         "if BUCKET_URI == \"[your-cloud-storage-bucket]\":\n",
         "    TIMESTAMP = datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
         "    BUCKET_URI = f\"gs://{PROJECT_ID}-{TIMESTAMP}\"\n",
         "\n",
-        "    ! gsutil mb -l {LOCATION} -p {PROJECT_ID} {BUCKET_URI}"
+        "    ! gsutil mb -l {GCS_LOCATION} -p {PROJECT_ID} {BUCKET_URI}"
       ]
     },
     {
@@ -477,7 +440,11 @@
       "outputs": [],
       "source": [
         "# Refresh the job until complete\n",
-        "while gcs_batch_job.state == \"JOB_STATE_RUNNING\":\n",
+        "while gcs_batch_job.state in (\n",
+        "    \"JOB_STATE_RUNNING\",\n",
+        "    \"JOB_STATE_PENDING\",\n",
+        "    \"JOB_STATE_QUEUED\",\n",
+        "):\n",
         "    time.sleep(5)\n",
         "    gcs_batch_job = client.batches.get(name=gcs_batch_job.name)\n",
         "\n",
@@ -544,7 +511,20 @@
         "id": "bfb2a462a7c6"
       },
       "source": [
-        "## BigQuery"
+        "## BigQuery\n",
+        "\n",
+        "⚠️ Batch predictions using BigQuery currently does not support Global endpoints. "
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "9c694724e069"
+      },
+      "outputs": [],
+      "source": [
+        "client = genai.Client(vertexai=True, project=PROJECT_ID, location=\"us-central1\")"
       ]
     },
     {
@@ -755,7 +735,11 @@
       "outputs": [],
       "source": [
         "# Refresh the job until complete\n",
-        "while bq_batch_job.state == \"JOB_STATE_RUNNING\":\n",
+        "while bq_batch_job.state in (\n",
+        "    \"JOB_STATE_RUNNING\",\n",
+        "    \"JOB_STATE_PENDING\",\n",
+        "    \"JOB_STATE_QUEUED\",\n",
+        "):\n",
         "    time.sleep(5)\n",
         "    bq_batch_job = client.batches.get(name=bq_batch_job.name)\n",
         "\n",