Set model_garden_source_model_name for pytorch inference deployments.

PiperOrigin-RevId: 720311822
This commit is contained in:
Dustin Luong
2025-01-27 14:38:52 -08:00
committed by Copybara-Service
parent b7b14ba7d8
commit 52925286ec
2 changed files with 14 additions and 0 deletions
@@ -183,6 +183,8 @@
"def deploy_model(\n",
" model_name: str,\n",
" model_id: str,\n",
" publisher: str,\n",
" publisher_model_id: str,\n",
" task: str,\n",
" machine_type: str = \"g2-standard-12\",\n",
" accelerator_type: str = \"NVIDIA_L4\",\n",
@@ -224,6 +226,9 @@
" serving_container_predict_route=\"/pred\",\n",
" serving_container_health_route=\"/ping\",\n",
" serving_container_environment_variables=serving_env,\n",
" model_garden_source_model_name=(\n",
" f\"publishers/{publisher}/models/{publisher_model_id}\"\n",
" ),\n",
" )\n",
"\n",
" model.deploy(\n",
@@ -243,6 +248,8 @@
"models[\"pytorch_inference_gpu\"], endpoints[\"pytorch_inference_gpu\"] = deploy_model(\n",
" model_name=common_util.get_job_name_with_datetime(prefix=MODEL_ID),\n",
" model_id=MODEL_ID,\n",
" publisher=\"hf-distilbert\",\n",
" publisher_model_id=\"distilbert-base-uncased-finetuned-sst-2-english\",\n",
" task=TASK,\n",
" machine_type=machine_type,\n",
" accelerator_type=accelerator_type,\n",
@@ -319,6 +319,8 @@
"def deploy_model(\n",
" model_name: str,\n",
" model_id: str,\n",
" publisher: str,\n",
" publisher_model_id: str,\n",
" task: str,\n",
" machine_type: str = \"g2-standard-12\",\n",
" accelerator_type: str = \"NVIDIA_L4\",\n",
@@ -360,6 +362,9 @@
" serving_container_predict_route=\"/pred\",\n",
" serving_container_health_route=\"/ping\",\n",
" serving_container_environment_variables=serving_env,\n",
" model_garden_source_model_name=(\n",
" f\"publishers/{publisher}/models/{publisher_model_id}\"\n",
" ),\n",
" )\n",
"\n",
" model.deploy(\n",
@@ -379,6 +384,8 @@
"models[\"pytorch_inference_gpu\"], endpoints[\"pytorch_inference_gpu\"] = deploy_model(\n",
" model_name=common_util.get_job_name_with_datetime(prefix=\"prompt-guard-serve\"),\n",
" model_id=model_id,\n",
" publisher=\"meta\",\n",
" publisher_model_id=\"prompt-guard\",\n",
" task=task,\n",
" machine_type=machine_type,\n",
" accelerator_type=accelerator_type,\n",