diff --git a/providers/src/airflow/providers/google/cloud/hooks/vertex_ai/model_service.py b/providers/src/airflow/providers/google/cloud/hooks/vertex_ai/model_service.py index 02fc7378380f2..ee3b171eb63ad 100644 --- a/providers/src/airflow/providers/google/cloud/hooks/vertex_ai/model_service.py +++ b/providers/src/airflow/providers/google/cloud/hooks/vertex_ai/model_service.py @@ -219,7 +219,7 @@ def upload_model( :param project_id: Required. The ID of the Google Cloud project that the service belongs to. :param region: Required. The ID of the Google Cloud region that the service belongs to. :param model: Required. The Model to create. - :param parent_model: The name of the parent model to create a new version under. + :param parent_model: The ID of the parent model to create a new version under. :param retry: Designation of what errors, if any, should be retried. :param timeout: The timeout for this request. :param metadata: Strings which should be sent along with the request as metadata. @@ -233,7 +233,7 @@ def upload_model( } if parent_model: - request["parent_model"] = parent_model + request["parent_model"] = client.model_path(project_id, region, parent_model) result = client.upload_model( request=request, diff --git a/providers/src/airflow/providers/google/cloud/operators/vertex_ai/model_service.py b/providers/src/airflow/providers/google/cloud/operators/vertex_ai/model_service.py index 04d45f253f568..48fb9154ce051 100644 --- a/providers/src/airflow/providers/google/cloud/operators/vertex_ai/model_service.py +++ b/providers/src/airflow/providers/google/cloud/operators/vertex_ai/model_service.py @@ -361,8 +361,9 @@ class UploadModelOperator(GoogleCloudBaseOperator): :param project_id: Required. The ID of the Google Cloud project that the service belongs to. :param region: Required. The ID of the Google Cloud region that the service belongs to. - :param model: Required. The Model to create. - :param parent_model: The name of the parent model to create a new version under. + :param model: Required. The Model to create. Creating model with the name that already + exists leads to creating new version of existing model. + :param parent_model: The ID of the parent model to create a new version under. :param retry: Designation of what errors, if any, should be retried. :param timeout: The timeout for this request. :param metadata: Strings which should be sent along with the request as metadata. @@ -377,7 +378,7 @@ class UploadModelOperator(GoogleCloudBaseOperator): account from the list granting this role to the originating account (templated). """ - template_fields = ("region", "project_id", "model", "impersonation_chain") + template_fields = ("region", "project_id", "model", "parent_model", "impersonation_chain") operator_extra_links = (VertexAIModelLink(),) def __init__( diff --git a/providers/tests/google/cloud/hooks/vertex_ai/test_model_service.py b/providers/tests/google/cloud/hooks/vertex_ai/test_model_service.py index c95d3fecd492f..ea337576e0cb8 100644 --- a/providers/tests/google/cloud/hooks/vertex_ai/test_model_service.py +++ b/providers/tests/google/cloud/hooks/vertex_ai/test_model_service.py @@ -37,7 +37,7 @@ TEST_REGION: str = "test-region" TEST_PROJECT_ID: str = "test-project-id" TEST_MODEL = None -TEST_PARENT_MODEL = "test-parent-model" +TEST_PARENT_MODEL = "projects/test-project-id/locations/test-region/models/test-parent-model" TEST_MODEL_NAME: str = "test_model_name" TEST_OUTPUT_CONFIG: dict = {} @@ -148,7 +148,7 @@ def test_upload_model_with_parent_model(self, mock_client) -> None: request=dict( parent=mock_client.return_value.common_location_path.return_value, model=TEST_MODEL, - parent_model=TEST_PARENT_MODEL, + parent_model=mock_client.return_value.model_path.return_value, ), metadata=(), retry=DEFAULT, @@ -352,7 +352,7 @@ def test_upload_model_with_parent_model(self, mock_client) -> None: request=dict( parent=mock_client.return_value.common_location_path.return_value, model=TEST_MODEL, - parent_model=TEST_PARENT_MODEL, + parent_model=mock_client.return_value.model_path.return_value, ), metadata=(), retry=DEFAULT, diff --git a/providers/tests/system/google/cloud/vertex_ai/example_vertex_ai_model_service.py b/providers/tests/system/google/cloud/vertex_ai/example_vertex_ai_model_service.py index 272069b2f87f0..1aaeccc9c069a 100644 --- a/providers/tests/system/google/cloud/vertex_ai/example_vertex_ai_model_service.py +++ b/providers/tests/system/google/cloud/vertex_ai/example_vertex_ai_model_service.py @@ -113,6 +113,19 @@ "health_route": "", }, } +MODEL_OBJ_V2 = { + "display_name": f"model-{ENV_ID}-v2", + "artifact_uri": "{{ti.xcom_pull('custom_task')['artifactUri']}}", + "container_spec": { + "image_uri": MODEL_SERVING_CONTAINER_URI, + "command": [], + "args": [], + "env": [], + "ports": [], + "predict_route": "", + "health_route": "", + }, +} with DAG( @@ -229,13 +242,14 @@ project_id=PROJECT_ID, model=MODEL_OBJ, ) + upload_model_v1 = upload_model.output["model_id"] # [END how_to_cloud_vertex_ai_upload_model_operator] upload_model_with_parent_model = UploadModelOperator( task_id="upload_model_with_parent_model", region=REGION, project_id=PROJECT_ID, - model=MODEL_OBJ, - parent_model=MODEL_DISPLAY_NAME, + model=MODEL_OBJ_V2, + parent_model=upload_model_v1, ) # [START how_to_cloud_vertex_ai_export_model_operator]