Try to automatically set region

GoogleCloudPlatform · Apr 23, 2024 · d67e5b0 · d67e5b0
1 parent 43c050a
commit d67e5b0
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 2 deletions.
diff --git a/notebooks/community/model_garden/model_garden_pytorch_mixtral_deployment.ipynb b/notebooks/community/model_garden/model_garden_pytorch_mixtral_deployment.ipynb
@@ -217,10 +217,17 @@
         "    if accelerator_type in [\"NVIDIA_TESLA_T4\", \"NVIDIA_TESLA_V100\"]:\n",
         "        dtype = \"float16\"\n",
         "\n",
+        "    if \"asia\" in REGION:\n",
+        "        region_suffix = \"asia\"\n",
+        "    elif \"europe\" in REGION:\n",
+        "        region_suffix = \"eu\"\n",
+        "    else:\n",
+        "        region_suffix = \"us\"\n",
+        "\n",
         "    vllm_args = [\n",
         "        \"--host=0.0.0.0\",\n",
         "        \"--port=7080\",\n",
-        "        f\"--model=gs://vertex-model-garden-public-us/{model_id}\",\n",
+        "        f\"--model=gs://vertex-model-garden-public-{region_suffix}/{model_id}\",\n",
         "        f\"--tensor-parallel-size={accelerator_count}\",\n",
         "        \"--swap-space=16\",\n",
         "        f\"--dtype={dtype}\",\n",

diff --git a/notebooks/community/model_garden/model_garden_pytorch_mixtral_peft_tuning.ipynb b/notebooks/community/model_garden/model_garden_pytorch_mixtral_peft_tuning.ipynb
@@ -261,7 +261,13 @@
         "\n",
         "# Huggingface dataset name or gs:// URI to a custom JSONL dataset.\n",
         "base_model_id = \"mistralai/Mixtral-8x7B\"\n",
-        "gcs_model_id = f\"gs://vertex-model-garden-public-us/{base_model_id}\"\n",
+        "if \"asia\" in REGION:\n",
+        "    region_suffix = \"asia\"\n",
+        "elif \"europe\" in REGION:\n",
+        "    region_suffix = \"eu\"\n",
+        "else:\n",
+        "    region_suffix = \"us\"\n",
+        "gcs_model_id = f\"gs://vertex-model-garden-public-{region_suffix}/{base_model_id}\"\n",
         "dataset_name = \"fredmo/vertexai-qna-500\"  # @param {type:\"string\"}\n",
         "\n",
         "# Optional. Template name or gs:// URI to a custom template.\n",