Use trtllm_region for trtllm deployment

PiperOrigin-RevId: 764029854

notebooks/community/model_garden/model_garden_pytorch_deepseek_deployment.ipynb

@@ -1248,6 +1248,7 @@
         "# @markdown Find Vertex AI prediction supported accelerators and regions at https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
         "trtllm_accelerator_type = \"NVIDIA_H200_141GB\"  # @param [\"NVIDIA_H200_141GB\"] {isTemplate:true}\n",
         "accelerator_count = 8\n",
+        "trtllm_region = \"us-east4\"  # @param [\"us-east4\"] {isTemplate:true}\n",
         "if trtllm_accelerator_type == \"NVIDIA_H200_141GB\":\n",
         "    machine_type = \"a3-ultragpu-8g\"\n",
         "    multihost_gpu_node_count = 1\n",
@@ -1257,7 +1258,7 @@
         "\n",
         "check_quota(\n",
         "    project_id=PROJECT_ID,\n",
-        "    region=REGION,\n",
+        "    region=trtllm_region,\n",
         "    resource_id=resource_id,\n",
         "    accelerator_count=int(accelerator_count * multihost_gpu_node_count),\n",
         ")\n",
@@ -1269,7 +1270,7 @@
         "GPU_MEMORY_UTILIZATION = 0.55\n",
         "\n",
         "\n",
-        "def poll_operation(op_name: str) -> bool:  # noqa: F811\n",
+        "def poll_operation(op_name: str, trtllm_region: str) -> bool:  # noqa: F811\n",
         "    creds, _ = auth.default()\n",
         "    auth_req = auth.transport.requests.Request()\n",
         "    creds.refresh(auth_req)\n",
@@ -1277,7 +1278,7 @@
         "        \"Authorization\": f\"Bearer {creds.token}\",\n",
         "    }\n",
         "    get_resp = requests.get(\n",
-        "        f\"https://{REGION}-aiplatform.googleapis.com/ui/{op_name}\",\n",
+        "        f\"https://{trtllm_region}-aiplatform.googleapis.com/ui/{op_name}\",\n",
         "        headers=headers,\n",
         "    )\n",
         "    opjs = get_resp.json()\n",
@@ -1286,9 +1287,11 @@
         "    return opjs.get(\"done\", False)\n",
         "\n",
         "\n",
-        "def poll_and_wait(op_name: str, total_wait: int, interval: int = 60):  # noqa: F811\n",
+        "def poll_and_wait_trtllm(\n",
+        "    op_name: str, total_wait: int, trtllm_region: str, interval: int = 60\n",
+        "):  # noqa: F811\n",
         "    waited = 0\n",
-        "    while not poll_operation(op_name):\n",
+        "    while not poll_operation(op_name, trtllm_region):\n",
         "        if waited > total_wait:\n",
         "            raise TimeoutError(\"Operation timed out\")\n",
         "        print(\n",
@@ -1319,6 +1322,7 @@
         "    enable_chunked_prefill: bool = False,\n",
         "    use_dedicated_endpoint: bool = False,\n",
         "    is_spot: bool = True,\n",
+        "    trtllm_region: str = REGION,\n",
         ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
         "    \"\"\"Deploys trained models with TensorRT-LLM on Vertex AI.\"\"\"\n",
         "    endpoint = aiplatform.Endpoint.create(\n",
@@ -1390,7 +1394,7 @@
         "    auth_req = auth.transport.requests.Request()\n",
         "    creds.refresh(auth_req)\n",
         "\n",
-        "    url = f\"https://{REGION}-aiplatform.googleapis.com/ui/projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint.name}:deployModel\"\n",
+        "    url = f\"https://{trtllm_region}-aiplatform.googleapis.com/ui/projects/{PROJECT_ID}/locations/{trtllm_region}/endpoints/{endpoint.name}:deployModel\"\n",
         "    headers = {\n",
         "        \"Content-Type\": \"application/json\",\n",
         "        \"Authorization\": f\"Bearer {creds.token}\",\n",
@@ -1423,7 +1427,7 @@
         "    print(f\"Deploy Model response: {response.json()}\")\n",
         "    if response.status_code != 200 or \"name\" not in response.json():\n",
         "        raise ValueError(f\"Failed to deploy model: {response.text}\")\n",
-        "    poll_and_wait(response.json()[\"name\"], 7200)\n",
+        "    poll_and_wait_trtllm(response.json()[\"name\"], 7200, trtllm_region)\n",
         "    print(\"endpoint_name:\", endpoint.name)\n",
         "\n",
         "    return model, endpoint\n",
@@ -1446,6 +1450,7 @@
         "    enable_trust_remote_code=True,\n",
         "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
         "    is_spot=is_spot,\n",
+        "    trtllm_region=trtllm_region,\n",
         ")\n",
         "# @markdown Click \"Show Code\" to see more details."
       ]

notebooks/community/model_garden/model_garden_pytorch_llama3_3_deployment.ipynb

@@ -701,10 +701,11 @@
         "accelerator_count = 8\n",
         "machine_type = \"a3-highgpu-8g\"\n",
         "multihost_gpu_node_count = 1\n",
+        "trtllm_region = REGION\n",
         "\n",
         "common_util.check_quota(\n",
         "    project_id=PROJECT_ID,\n",
-        "    region=REGION,\n",
+        "    region=trtllm_region,\n",
         "    accelerator_type=accelerator_type,\n",
         "    accelerator_count=accelerator_count,\n",
         "    is_for_training=False,\n",
@@ -714,7 +715,7 @@
         "MAX_MODEL_LEN = 131072\n",
         "\n",
         "\n",
-        "def poll_operation(op_name: str) -> bool:  # noqa: F811\n",
+        "def poll_operation(op_name: str, trtllm_region: str) -> bool:  # noqa: F811\n",
         "    creds, _ = auth.default()\n",
         "    auth_req = auth.transport.requests.Request()\n",
         "    creds.refresh(auth_req)\n",
@@ -722,7 +723,7 @@
         "        \"Authorization\": f\"Bearer {creds.token}\",\n",
         "    }\n",
         "    get_resp = requests.get(\n",
-        "        f\"https://{REGION}-aiplatform.googleapis.com/ui/{op_name}\",\n",
+        "        f\"https://{trtllm_region}-aiplatform.googleapis.com/ui/{op_name}\",\n",
         "        headers=headers,\n",
         "    )\n",
         "    opjs = get_resp.json()\n",
@@ -731,9 +732,11 @@
         "    return opjs.get(\"done\", False)\n",
         "\n",
         "\n",
-        "def poll_and_wait(op_name: str, total_wait: int, interval: int = 60):  # noqa: F811\n",
+        "def poll_and_wait_trtllm(\n",
+        "    op_name: str, total_wait: int, trtllm_region: str, interval: int = 60\n",
+        "):  # noqa: F811\n",
         "    waited = 0\n",
-        "    while not poll_operation(op_name):\n",
+        "    while not poll_operation(op_name, trtllm_region):\n",
         "        if waited > total_wait:\n",
         "            raise TimeoutError(\"Operation timed out\")\n",
         "        print(\n",
@@ -764,6 +767,7 @@
         "    enable_chunked_prefill: bool = False,\n",
         "    use_dedicated_endpoint: bool = False,\n",
         "    is_spot: bool = True,\n",
+        "    trtllm_region: str = REGION,\n",
         ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
         "    \"\"\"Deploys trained models with TensorRT-LLM on Vertex AI.\"\"\"\n",
         "    endpoint = aiplatform.Endpoint.create(\n",
@@ -835,7 +839,7 @@
         "    auth_req = auth.transport.requests.Request()\n",
         "    creds.refresh(auth_req)\n",
         "\n",
-        "    url = f\"https://{REGION}-aiplatform.googleapis.com/ui/projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint.name}:deployModel\"\n",
+        "    url = f\"https://{trtllm_region}-aiplatform.googleapis.com/ui/projects/{PROJECT_ID}/locations/{trtllm_region}/endpoints/{endpoint.name}:deployModel\"\n",
         "    headers = {\n",
         "        \"Content-Type\": \"application/json\",\n",
         "        \"Authorization\": f\"Bearer {creds.token}\",\n",
@@ -868,7 +872,7 @@
         "    print(f\"Deploy Model response: {response.json()}\")\n",
         "    if response.status_code != 200 or \"name\" not in response.json():\n",
         "        raise ValueError(f\"Failed to deploy model: {response.text}\")\n",
-        "    poll_and_wait(response.json()[\"name\"], 7200)\n",
+        "    poll_and_wait_trtllm(response.json()[\"name\"], 7200, trtllm_region)\n",
         "    print(\"endpoint_name:\", endpoint.name)\n",
         "\n",
         "    return model, endpoint\n",
@@ -888,6 +892,7 @@
         "    enable_chunked_prefill=True,\n",
         "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
         "    is_spot=is_spot,\n",
+        "    trtllm_region=trtllm_region,\n",
         ")\n",
         "# @markdown Click \"Show Code\" to see more details."
       ]