Set model_garden_source_model_name for some model garden deployment notebooks.

PiperOrigin-RevId: 721846318
This commit is contained in:
Dustin Luong
2025-01-31 11:39:27 -08:00
committed by Copybara-Service
parent a5944510d6
commit 3c06e4797a
7 changed files with 164 additions and 18 deletions
@@ -292,6 +292,7 @@
" serving_container_image_uri: str,\n",
" model_pth_file: str,\n",
" model_cfg_yaml_file: str,\n",
" publisher_model_id: str,\n",
" test_threshold: float = 0.5,\n",
"):\n",
"\n",
@@ -311,6 +312,9 @@
" serving_container_predict_route=\"/predictions/detectron2_serving\",\n",
" serving_container_health_route=\"/ping\",\n",
" serving_container_environment_variables=serving_env,\n",
" model_garden_source_model_name=(\n",
" f\"publishers/google/models/{publisher_model_id}\"\n",
" ),\n",
" )\n",
"\n",
" model.wait()\n",
@@ -1022,6 +1026,7 @@
"remainder_args_list += [\"SOLVER.STEPS\"] + [\"[]\"]\n",
"remainder_args_list += [\"MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE\"] + [\"128\"]\n",
"remainder_args_list += [\"MODEL.ROI_HEADS.NUM_CLASSES\"] + [\"1\"]\n",
"publisher_model_id = \"mask-r-cnn\"\n",
"#################################################\n",
"\n",
"# #################################################\n",
@@ -1038,6 +1043,7 @@
"# remainder_args_list += ['MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE'] + ['128']\n",
"# remainder_args_list += ['MODEL.ROI_HEADS.NUM_CLASSES'] + ['1']\n",
"# remainder_args_list += ['MODEL.RETINANET.NUM_CLASSES'] + ['1']\n",
"# publisher_model_id = \"retinanet\"\n",
"# #################################################\n",
"\n",
"# #################################################\n",
@@ -1053,6 +1059,7 @@
"# remainder_args_list += ['SOLVER.STEPS'] + ['[]']\n",
"# remainder_args_list += ['MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE'] + ['128']\n",
"# remainder_args_list += ['MODEL.ROI_HEADS.NUM_CLASSES'] + ['1']\n",
"# publisher_model_id = \"faster-r-cnn'\"\n",
"# #################################################\n",
"\n",
"# Create argument list for docker.\n",
@@ -1146,6 +1153,7 @@
" serving_container_image_uri=PREDICTION_CONTAINER_URI,\n",
" model_pth_file=PRETRAINED_MODEL_PTH_FILE,\n",
" model_cfg_yaml_file=PRETRAINED_MODEL_CFG_YAML_FILE,\n",
" publisher_model_id=publisher_model_id,\n",
" test_threshold=TEST_THRESHOLD,\n",
")\n",
"print(\"The uploaded model name is: \", PREDICTION_DISPLAY_NAME)"
@@ -192,6 +192,7 @@
"def deploy_model(\n",
" model_name: str,\n",
" model_id: str,\n",
" publisher_model_id: str,\n",
" service_account: str,\n",
" machine_type: str = \"g2-standard-8\",\n",
" accelerator_type: str = \"NVIDIA_L4\",\n",
@@ -224,6 +225,9 @@
" serving_container_predict_route=\"/generate\",\n",
" serving_container_health_route=\"/ping\",\n",
" serving_container_environment_variables=env_vars,\n",
" model_garden_source_model_name=(\n",
" f\"publishers/meta/models/{publisher_model_id}\",\n",
" ),\n",
" )\n",
" print(\n",
" f\"Deploying {model_id} on {machine_type} with {accelerator_count} {accelerator_type} GPU(s).\"\n",
@@ -454,12 +458,14 @@
"if quantized_model_id == \"Custom quantized model\":\n",
" # quantization_method, base_model_name, and quantization_output_dir are set in the Quantize section\n",
" model_id = quantization_output_dir\n",
" publisher_model_id = \"llama-2-quantized\"\n",
"else:\n",
" model_id = quantized_model_id\n",
" quantization_method = quantized_model_id.split(\"-\")[-1].lower()\n",
" base_model_name = \"-\".join(\n",
" quantized_model_id.split(\"/\")[-1].split(\"-\")[:-1]\n",
" ).lower()\n",
" publisher_model_id = \"llama2\"\n",
"\n",
"\n",
"# @markdown Deploying a quantized model requires much less GPU.\n",
@@ -522,6 +528,7 @@
"model_quantized_vllm, endpoint_quantized_vllm = deploy_model(\n",
" model_name=get_job_name_with_datetime(prefix=\"llama2-serve-vllm-quantized\"),\n",
" model_id=model_id,\n",
" publisher_model_id=publisher_model_id,\n",
" service_account=SERVICE_ACCOUNT,\n",
" machine_type=machine_type,\n",
" accelerator_type=accelerator_type,\n",
@@ -283,6 +283,8 @@
"def deploy_model(\n",
" model_name,\n",
" model_id,\n",
" publisher,\n",
" publisher_model_id,\n",
" finetuned_lora_model_path,\n",
" service_account,\n",
" task,\n",
@@ -304,6 +306,9 @@
" serving_container_predict_route=\"/predictions/peft_serving\",\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",
" model.deploy(\n",
" endpoint=endpoint,\n",
@@ -480,6 +485,8 @@
"model, endpoint = deploy_model(\n",
" model_name=get_job_name_with_datetime(prefix=\"openllama-peft-serve\"),\n",
" model_id=model_id,\n",
" publisher=\"openlm-research\",\n",
" publisher_model_id=\"openllama\",\n",
" finetuned_lora_model_path=output_dir,\n",
" service_account=SERVICE_ACCOUNT,\n",
" task=\"causal-language-modeling-lora\",\n",
@@ -720,6 +727,8 @@
"model, endpoint = deploy_model(\n",
" model_name=get_job_name_with_datetime(prefix=\"falcon-peft-serve\"),\n",
" model_id=model_id,\n",
" publisher=\"tiiuae\",\n",
" publisher_model_id=\"falcon-instruct-7b-peft\",\n",
" finetuned_lora_model_path=os.path.join(output_dir, \"checkpoint-\" + str(max_steps)),\n",
" service_account=SERVICE_ACCOUNT,\n",
" task=\"instruct-lora\",\n",
@@ -938,9 +947,23 @@
},
"outputs": [],
"source": [
"if model_id == \"bert-base-uncased\":\n",
" publisher = \"google\"\n",
" publisher_model_id = \"bert-base-uncased\"\n",
"elif model_id == \"roberta-large\":\n",
" publisher = \"meta\"\n",
" publisher_model_id = \"roberta-large\"\n",
"elif model_id == \"xlm-roberta-large\":\n",
" publisher = \"meta\"\n",
" publisher_model_id = \"xlm-roberta-large\"\n",
"else:\n",
" raise ValueError(f\"Unsupported model id: {model_id}\")\n",
"\n",
"model, endpoint = deploy_model(\n",
" model_name=get_job_name_with_datetime(prefix=\"sequence-classification-peft-serve\"),\n",
" model_id=model_id,\n",
" publisher=publisher,\n",
" publisher_model_id=publisher_model_id,\n",
" finetuned_lora_model_path=output_dir,\n",
" service_account=SERVICE_ACCOUNT,\n",
" task=\"sequence-classification-lora\",\n",
@@ -384,6 +384,54 @@
" return endpoint\n",
"\n",
"\n",
"def get_publisher_model_id(model_name: str) -> str | None:\n",
" \"\"\"Returns the corresponding task name for the given model_name.\"\"\"\n",
"\n",
" model_to_publisher_dict = {\n",
" \"stabilityai/stable-diffusion-2-1\": (\"stability-ai\", \"stable-diffusion-2-1\"),\n",
" \"stabilityai/stable-diffusion-xl-base-1.0\": (\n",
" \"stability-ai\",\n",
" \"stable-diffusion-xl-base\",\n",
" ),\n",
" \"stabilityai/stable-diffusion-xl-base-1.0 - refiner\": (\n",
" \"stability-ai\",\n",
" \"stable-diffusion-xl-base\",\n",
" ),\n",
" \"latent-consistency/lcm-sdxl\": (\"stability-ai\", \"stable-diffusion-xl-lcm\"),\n",
" \"latent-consistency/lcm-lora-sdxl\": (\"stability-ai\", \"stable-diffusion-xl-lcm\"),\n",
" \"stabilityai/sdxl-turbo\": None,\n",
" \"bytedance/sdxl-lightning\": (\"bytedance\", \"stable-diffusion-xl-lightning\"),\n",
" \"runDiffusion/Juggernaut-XL-Lightning\": None,\n",
" \"stablediffusionapi/juggernaut-xl-v9\": None,\n",
" \"lykon/dreamshaper-xl-v2-turbo\": None,\n",
" \"stablediffusionapi/anything-xl\": None,\n",
" \"stablediffusionapi/turbovision_xl\": None,\n",
" \"kandinsky-community/kandinsky-2-2-decoder-inpaint\": None,\n",
" \"diffusers/stable-diffusion-xl-1.0-inpainting-0.1\": (\n",
" \"runwayml\",\n",
" \"stable-diffusion-inpainting\",\n",
" ),\n",
" \"timbrooks/instruct-pix2pix\": (\"timbrooks\", \"instruct-pix2pix\"),\n",
" \"lllyasviel/sd-controlnet-canny\": (\"lllyasviel\", \"control-net\"),\n",
" \"stabilityai/stable-diffusion-x4-upscaler\": (\n",
" \"stability-ai\",\n",
" \"stable-diffusion-4x-upscaler\",\n",
" ),\n",
" }\n",
"\n",
" if model_name not in model_to_publisher_dict.keys():\n",
" print(model_name)\n",
" raise gr.Error(\"Select a valid model name for Endpoint creation.\")\n",
"\n",
" if model_to_publisher_dict[model_name]:\n",
" publisher = model_to_publisher_dict[model_name][0]\n",
" publisher_model_id = model_to_publisher_dict[model_name][1]\n",
" else:\n",
" publisher = \"hf-\" + model_name.split(\"/\")[0]\n",
" publisher_model_id = model_name.split(\"/\")[1]\n",
" return f\"publishers/{publisher}/models/{publisher_model_id}\"\n",
"\n",
"\n",
"def get_task_name(model_name: str) -> str:\n",
" \"\"\"Returns the corresponding task name for the given model_name.\"\"\"\n",
"\n",
@@ -420,13 +468,14 @@
" raise gr.Error(\"Select a valid model name for model list.\")\n",
"\n",
" model_id = model_name.split(\": \")[1]\n",
" publisher_model_id = get_publisher_model_id(model_id)\n",
" task_name = get_task_name(model_id)\n",
"\n",
" return deploy_model_vertex(model_id, task_name)\n",
" return deploy_model_vertex(model_id, publisher_model_id, task_name)\n",
"\n",
"\n",
"def deploy_model_vertex(\n",
" model_id: str, task_name: str, lora_id: str = \"\"\n",
" model_id: str, publisher_model_id: str | None, task_name: str, lora_id: str = \"\"\n",
") -> aiplatform.Endpoint:\n",
" \"\"\"\n",
" Creates a new Vertex prediction endpoint and deploys a model to it.\n",
@@ -464,9 +513,13 @@
" }\n",
"\n",
" if \"text-to-image\" in task_name:\n",
" model = upload_model_pytorch_inference_container(model_id, serving_env)\n",
" model = upload_model_pytorch_inference_container(\n",
" model_id, publisher_model_id, serving_env\n",
" )\n",
" else:\n",
" model = upload_model_pytorch_diffusers_serve_container(model_id, serving_env)\n",
" model = upload_model_pytorch_diffusers_serve_container(\n",
" model_id, publisher_model_id, serving_env\n",
" )\n",
"\n",
" machine_type = \"g2-standard-8\"\n",
" accelerator_type = \"NVIDIA_L4\"\n",
@@ -505,7 +558,7 @@
"\n",
"\n",
"def upload_model_pytorch_inference_container(\n",
" model_id: str, serving_env: dict\n",
" model_id: str, publisher_model_id: str | None, serving_env: dict\n",
") -> aiplatform.Model:\n",
" return aiplatform.Model.upload(\n",
" display_name=model_id,\n",
@@ -514,6 +567,7 @@
" serving_container_predict_route=\"/predict\",\n",
" serving_container_health_route=\"/health\",\n",
" serving_container_environment_variables=serving_env,\n",
" model_garden_source_model_name=publisher_model_id,\n",
" )\n",
"\n",
"\n",
@@ -522,7 +576,7 @@
"\n",
"\n",
"def upload_model_pytorch_diffusers_serve_container(\n",
" model_id: str, serving_env: dict\n",
" model_id: str, publisher_model_id: str | None, serving_env: dict\n",
") -> aiplatform.Model:\n",
" return aiplatform.Model.upload(\n",
" display_name=model_id,\n",
@@ -531,6 +585,7 @@
" serving_container_predict_route=\"/predictions/diffusers_serving\",\n",
" serving_container_health_route=\"/ping\",\n",
" serving_container_environment_variables=serving_env,\n",
" model_garden_source_model_name=publisher_model_id,\n",
" )\n",
"\n",
"\n",
@@ -215,7 +215,7 @@
"source": [
"# @title Load a pretrained model\n",
"\n",
"# @markdown `MODEL_NAME :` The model you want to train and serve.\n",
"# @markdown `MODEL_NAME`: The model you want to train and serve.\n",
"\n",
"# We use a ViT model as the example.\n",
"MODEL_NAME = \"vit_tiny_patch16_224\" # @param [\"vit_tiny_patch16_224\", \"beit_base_patch16_224\", \"deit3_small_patch16_224\", \"efficientnet_b2\", \"mobilenetv2_100\", \"resnet50\", \"resnest50d\", \"convnext_base\", \"cspdarknet53\", \"inception_v4\"]\n",
@@ -539,6 +539,30 @@
" \"DEPLOY_SOURCE\": \"notebook\",\n",
" }\n",
"\n",
"match MODEL_NAME:\n",
" case \"vit_tiny_patch16_224\":\n",
" publisher_model_id = \"vit-base-patch16-224\"\n",
" case \"beit_base_patch16_224\":\n",
" publisher_model_id = \"beit-base-patch16-224\"\n",
" case \"deit3_small_patch16_224\":\n",
" publihser_model_id = \"deit-base-patch16-224\"\n",
" case \"efficientnet_b2\":\n",
" publisher_model_id = \"efficientnetv2-rw-s\"\n",
" case \"mobilenetv2_100\":\n",
" publisher_model_id = \"mobilenetv2-100\"\n",
" case \"resnet50\":\n",
" publisher_model_id = \"resnet-50\"\n",
" case \"resnest50d\":\n",
" publisher_model_id = \"resnest50d\"\n",
" case \"convnext_base\":\n",
" publisher_model_id = \"convnext-base\"\n",
" case \"cspdarknet53\":\n",
" publisher_model_id = \"cspdarknet53\"\n",
" case \"inception_v4\":\n",
" publisher_model_id = \"inception-v4\"\n",
" case _:\n",
" raise ValueError(f\"Unknown model: {MODEL_NAME}\")\n",
"\n",
"models[\"timm-model\"] = aiplatform.Model.upload(\n",
" display_name=MODEL_NAME,\n",
" serving_container_image_uri=SERVE_DOCKER_URI,\n",
@@ -546,25 +570,25 @@
" serving_container_predict_route=\"/predictions/timm_serving\",\n",
" serving_container_health_route=\"/ping\",\n",
" serving_container_environment_variables=serving_env,\n",
" model_garden_source_model_name=(\n",
" f\"publishers/timm/models/{publisher_model_id}\",\n",
" ),\n",
")\n",
"# Or reuse a pre-uploaded model.\n",
"# models[\"timm-model\"] = aiplatform.Model('projects/123456789/locations/us-central1/models/123456789@1')\n",
"\n",
"# Create an endpoint.\n",
"endpoints[\"timm-endpoint\"] = aiplatform.Endpoint.create(\n",
" display_name=\"pytorch-timm-endpoint\"\n",
")\n",
"endpoints[\"timm-endpoint\"] = aiplatform.Endpoint.create(display_name=\"pytorch-timm-endpoint\")\n",
"# Or reuse a pre-created endpoint.\n",
"# endpoints[\"timm-endpoint\"] = aiplatform.Endpoint('projects/123456789/locations/us-central1/endpoints/123456789')\n",
"\n",
"# Check quota.\n",
"common_util.check_quota(\n",
" project_id=PROJECT_ID,\n",
" region=REGION,\n",
" accelerator_type=SERVING_ACCELERATOR_TYPE,\n",
" accelerator_count=1,\n",
" is_for_training=False,\n",
")\n",
" project_id=PROJECT_ID,\n",
" region=REGION,\n",
" accelerator_type=SERVING_ACCELERATOR_TYPE,\n",
" accelerator_count=1,\n",
" is_for_training=False,)\n",
"\n",
"# Deploy model to endpoint.\n",
"models[\"timm-model\"].deploy(\n",
@@ -574,7 +598,9 @@
" accelerator_count=1,\n",
" traffic_percentage=100,\n",
" service_account=SERVICE_ACCOUNT,\n",
" system_labels={\"NOTEBOOK_NAME\": \"model_garden_pytorch_timm.ipynb\"},\n",
" system_labels={\n",
" \"NOTEBOOK_NAME\": \"model_garden_pytorch_timm.ipynb\"\n",
" },\n",
")"
]
},
@@ -674,6 +674,15 @@
" \"MODEL_ID\": \"tensorflow-hub-efficientnetv2\",\n",
" \"DEPLOY_SOURCE\": \"notebook\",\n",
"}\n",
"match experiment:\n",
" case \"Efficientnetv2-m\":\n",
" publisher_model_id = \"imageclassification-efficientnet\"\n",
" case \"ViT-ti16\" | \"ViT-s16\" | \"ViT-b16\" | \"ViT-l16\":\n",
" publisher_model_id = \"imageclassification-vit\"\n",
" case \"MaxViT\":\n",
" publisher_model_id = \"imageclassification-maxvit\"\n",
" case _:\n",
" raise ValueError(f\"Unknown experiment: {experiment}\")\n",
"\n",
"models[\"model_icn\"] = aiplatform.Model.upload(\n",
" display_name=upload_job_name,\n",
@@ -681,6 +690,9 @@
" serving_container_image_uri=PREDICTION_CONTAINER_URI,\n",
" serving_container_args=SERVING_CONTAINER_ARGS,\n",
" serving_container_environment_variables=serving_env,\n",
" model_garden_source_model_name=(\n",
" f\"publishers/google/models/{publisher_model_id}\"\n",
" ),\n",
")\n",
"\n",
"models[\"model_icn\"].wait()\n",
@@ -708,7 +720,9 @@
" accelerator_count=1,\n",
" min_replica_count=1,\n",
" max_replica_count=1,\n",
" system_labels={\"NOTEBOOK_NAME\": \"model_garden_tfvision_image_classification.ipynb\"},\n",
" system_labels={\n",
" \"NOTEBOOK_NAME\": \"model_garden_tfvision_image_classification.ipynb\"\n",
" },\n",
")\n",
"\n",
"endpoint_id = endpoints[\"endpoint_icn\"].name\n",
@@ -916,12 +916,25 @@
" \"DEPLOY_SOURCE\": \"notebook\",\n",
"}\n",
"\n",
"match experiment:\n",
" case experiment.startswith(\"retinanet\"):\n",
" publisher_model_id = \"imageobjectdetection-spinenet\"\n",
" case \"scaled_yolo_v4\":\n",
" publisher_model_id = \"imageobjectdetection-yolo\"\n",
" case \"yolov7\":\n",
" publisher_model_id = \"tfvision-yolov7\"\n",
" case _:\n",
" raise ValueError(f\"Unknown experiment: {experiment}\")\n",
"\n",
"model = aiplatform.Model.upload(\n",
" display_name=upload_job_name,\n",
" artifact_uri=trained_model_dir,\n",
" serving_container_image_uri=PREDICTION_CONTAINER_URI,\n",
" serving_container_args=SERVING_CONTAINER_ARGS,\n",
" serving_container_environment_variables=serving_env,\n",
" model_garden_source_model_name=(\n",
" f\"publishers/google/models/{publisher_model_id}\",\n",
" ),\n",
")\n",
"\n",
"model.wait()\n",