From fdd9e19554236540382a2819b99804f836a6815b Mon Sep 17 00:00:00 2001 From: Christian Yarros Date: Tue, 27 Aug 2024 20:20:58 +0000 Subject: [PATCH 01/10] add supervised_fine_tuning --- .../hooks/vertex_ai/supervised_fine_tuning.py | 101 +++++++++++++++ .../operators/vertex_ai/generative_model.py | 2 +- .../vertex_ai/supervised_fine_tuning.py | 117 ++++++++++++++++++ .../operators/cloud/vertex_ai.rst | 14 +++ .../vertex_ai/test_supervised_fine_tuning.py | 80 ++++++++++++ .../vertex_ai/test_supervised_fine_tuning.py | 75 +++++++++++ ...xample_vertex_ai_supervised_fine_tuning.py | 68 ++++++++++ 7 files changed, 456 insertions(+), 1 deletion(-) create mode 100644 airflow/providers/google/cloud/hooks/vertex_ai/supervised_fine_tuning.py create mode 100644 airflow/providers/google/cloud/operators/vertex_ai/supervised_fine_tuning.py create mode 100644 tests/providers/google/cloud/hooks/vertex_ai/test_supervised_fine_tuning.py create mode 100644 tests/providers/google/cloud/operators/vertex_ai/test_supervised_fine_tuning.py create mode 100644 tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_supervised_fine_tuning.py diff --git a/airflow/providers/google/cloud/hooks/vertex_ai/supervised_fine_tuning.py b/airflow/providers/google/cloud/hooks/vertex_ai/supervised_fine_tuning.py new file mode 100644 index 0000000000000..e0d2769cfecf4 --- /dev/null +++ b/airflow/providers/google/cloud/hooks/vertex_ai/supervised_fine_tuning.py @@ -0,0 +1,101 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""This module contains a Google Cloud Vertex AI Generative Model hook.""" + +from __future__ import annotations + +import time +from typing import TYPE_CHECKING, Sequence + +import vertexai +from vertexai.preview.tuning import sft + +if TYPE_CHECKING: + from vertexai.preview.tuning.sft import SupervisedTuningJob + + +from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook + + +class SupervisedFineTuningHook(GoogleBaseHook): + """Hook for Google Cloud Vertex AI Supervised Fine Tuning APIs.""" + + def __init__( + self, + gcp_conn_id: str = "google_cloud_default", + impersonation_chain: str | Sequence[str] | None = None, + **kwargs, + ): + if kwargs.get("delegate_to") is not None: + raise RuntimeError( + "The `delegate_to` parameter has been deprecated before and finally removed in this version" + " of Google Provider. You MUST convert it to `impersonate_chain`" + ) + super().__init__(gcp_conn_id=gcp_conn_id, impersonation_chain=impersonation_chain, **kwargs) + + @GoogleBaseHook.fallback_to_default_project_id + def train( + self, + source_model: str, + train_dataset: str, + location: str, + tuned_model_display_name: str | None = None, + validation_dataset: str | None = None, + epochs: int | None = None, + adapter_size: int | None = None, + learning_rate_multiplier: float | None = None, + project_id: str = PROVIDE_PROJECT_ID, + ) -> SupervisedTuningJob: + """ + Use the Supervised Fine Tuning API to create a tuning job. + + :param source_model: Required. A pre-trained model optimized for performing natural + language tasks such as classification, summarization, extraction, content + creation, and ideation. + :param training_dataset: Required. Cloud Storage URI of your training dataset. The dataset + must be formatted as a JSONL file. For best results, provide at least 100 to 500 examples. + :param location: Required. The ID of the Google Cloud location that the service belongs to. + :param tuned_model_display_name: Optional. Display name of the TunedModel. The name can be up + to 128 characters long and can consist of any UTF-8 characters. + :param validation_dataset: Optional. Cloud Storage URI of your training dataset. The dataset must be + formatted as a JSONL file. For best results, provide at least 100 to 500 examples. + :param epochs: Optional. To optimize performance on a specific dataset, try using a higher + epoch value. Increasing the number of epochs might improve results. However, be cautious + about over-fitting, especially when dealing with small datasets. If over-fitting occurs, + consider lowering the epoch number. + :param adapter_size: Optional. Adapter size for tuning. + :param learning_rate_multiplier: Optional. Multiplier for adjusting the default learning rate. + """ + vertexai.init(project=project_id, location=location, credentials=self.get_credentials()) + + sft_tuning_job = sft.train( + source_model=source_model, + train_dataset=train_dataset, + validation_dataset=validation_dataset, + epochs=epochs, + adapter_size=adapter_size, + learning_rate_multiplier=learning_rate_multiplier, + tuned_model_display_name=tuned_model_display_name, + ) + + # Polling for job completion + while not sft_tuning_job.has_ended: + time.sleep(60) + sft_tuning_job.refresh() + + return sft_tuning_job diff --git a/airflow/providers/google/cloud/operators/vertex_ai/generative_model.py b/airflow/providers/google/cloud/operators/vertex_ai/generative_model.py index cfcd0014fb73a..5583e5e8bc664 100644 --- a/airflow/providers/google/cloud/operators/vertex_ai/generative_model.py +++ b/airflow/providers/google/cloud/operators/vertex_ai/generative_model.py @@ -525,7 +525,7 @@ class GenerativeModelGenerateContentOperator(GoogleCloudBaseOperator): account from the list granting this role to the originating account (templated). """ - template_fields = ("location", "project_id", "impersonation_chain", "contents") + template_fields = ("location", "project_id", "impersonation_chain", "contents", "pretrained_model") def __init__( self, diff --git a/airflow/providers/google/cloud/operators/vertex_ai/supervised_fine_tuning.py b/airflow/providers/google/cloud/operators/vertex_ai/supervised_fine_tuning.py new file mode 100644 index 0000000000000..a703bebc2ef69 --- /dev/null +++ b/airflow/providers/google/cloud/operators/vertex_ai/supervised_fine_tuning.py @@ -0,0 +1,117 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""This module contains Google Vertex AI Generative AI operators.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Sequence + +from airflow.providers.google.cloud.hooks.vertex_ai.supervised_fine_tuning import SupervisedFineTuningHook +from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator + +if TYPE_CHECKING: + from airflow.utils.context import Context + + +class SupervisedFineTuningTrainOperator(GoogleCloudBaseOperator): + """ + Use the Supervised Fine Tuning API to create a tuning job. + + :param source_model: Required. A pre-trained model optimized for performing natural + language tasks such as classification, summarization, extraction, content + creation, and ideation. + :param training_dataset: Required. Cloud Storage URI of your training dataset. The dataset + must be formatted as a JSONL file. For best results, provide at least 100 to 500 examples. + :param project_id: Required. The ID of the Google Cloud project that the + service belongs to. + :param location: Required. The ID of the Google Cloud location that the service belongs to. + :param tuned_model_display_name: Optional. Display name of the TunedModel. The name can be up + to 128 characters long and can consist of any UTF-8 characters. + :param validation_dataset: Optional. Cloud Storage URI of your training dataset. The dataset must be + formatted as a JSONL file. For best results, provide at least 100 to 500 examples. + :param epochs: Optional. To optimize performance on a specific dataset, try using a higher + epoch value. Increasing the number of epochs might improve results. However, be cautious + about over-fitting, especially when dealing with small datasets. If over-fitting occurs, + consider lowering the epoch number. + :param adapter_size: Optional. Adapter size for tuning. + :param learning_multiplier_rate: Optional. Multiplier for adjusting the default learning rate. + :param gcp_conn_id: The connection ID to use connecting to Google Cloud. + :param impersonation_chain: Optional service account to impersonate using short-term + credentials, or chained list of accounts required to get the access_token + of the last account in the list, which will be impersonated in the request. + If set as a string, the account must grant the originating account + the Service Account Token Creator IAM role. + If set as a sequence, the identities from the list must grant + Service Account Token Creator IAM role to the directly preceding identity, with first + account from the list granting this role to the originating account (templated). + """ + + template_fields = ("location", "project_id", "impersonation_chain", "train_dataset", "validation_dataset") + + def __init__( + self, + *, + source_model: str, + train_dataset: str, + project_id: str, + location: str, + tuned_model_display_name: str | None = None, + validation_dataset: str | None = None, + epochs: int | None = None, + adapter_size: int | None = None, + learning_rate_multiplier: float | None = None, + gcp_conn_id: str = "google_cloud_default", + impersonation_chain: str | Sequence[str] | None = None, + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.source_model = source_model + self.train_dataset = train_dataset + self.tuned_model_display_name = tuned_model_display_name + self.validation_dataset = validation_dataset + self.epochs = epochs + self.adapter_size = adapter_size + self.learning_rate_multiplier = learning_rate_multiplier + self.project_id = project_id + self.location = location + self.gcp_conn_id = gcp_conn_id + self.impersonation_chain = impersonation_chain + + def execute(self, context: Context): + self.hook = SupervisedFineTuningHook( + gcp_conn_id=self.gcp_conn_id, + impersonation_chain=self.impersonation_chain, + ) + response = self.hook.train( + source_model=self.source_model, + train_dataset=self.train_dataset, + project_id=self.project_id, + location=self.location, + validation_dataset=self.validation_dataset, + epochs=self.epochs, + adapter_size=self.adapter_size, + learning_rate_multiplier=self.learning_rate_multiplier, + tuned_model_display_name=self.tuned_model_display_name, + ) + + self.log.info("Tuned Model Name: %s", response.tuned_model_name) + self.log.info("Tuned Model Endpoint Name: %s", response.tuned_model_endpoint_name) + + self.xcom_push(context, key="tuned_model_name", value=response.tuned_model_name) + self.xcom_push(context, key="tuned_model_endpoint_name", value=response.tuned_model_endpoint_name) + return response diff --git a/docs/apache-airflow-providers-google/operators/cloud/vertex_ai.rst b/docs/apache-airflow-providers-google/operators/cloud/vertex_ai.rst index 5df8c60c8c67b..fb1ca4bf60cd7 100644 --- a/docs/apache-airflow-providers-google/operators/cloud/vertex_ai.rst +++ b/docs/apache-airflow-providers-google/operators/cloud/vertex_ai.rst @@ -615,6 +615,20 @@ The operator returns the model's response in :ref:`XCom ` under ` :start-after: [START how_to_cloud_vertex_ai_generative_model_generate_content_operator] :end-before: [END how_to_cloud_vertex_ai_generative_model_generate_content_operator] + +Performing Supervised Fine Tuning +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +To train and deploy a tuned generative model to an endpoint you can use +:class:`~airflow.providers.google.cloud.operators.vertex_ai.supervised_fine_tuning.SupervisedFineTuningTrainOperator`. +The operator returns the tuned model's endpoint name in :ref:`XCom ` under ``tuned_model_endpoint_name`` key. + +.. exampleinclude:: /../../tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_supervised_fine_tuning.py + :language: python + :dedent: 4 + :start-after: [START how_to_cloud_vertex_ai_supervised_fine_tuning_train_operator] + :end-before: [END how_to_cloud_vertex_ai_supervised_fine_tuning_train_operator] + Reference ^^^^^^^^^ diff --git a/tests/providers/google/cloud/hooks/vertex_ai/test_supervised_fine_tuning.py b/tests/providers/google/cloud/hooks/vertex_ai/test_supervised_fine_tuning.py new file mode 100644 index 0000000000000..b218bc89033d6 --- /dev/null +++ b/tests/providers/google/cloud/hooks/vertex_ai/test_supervised_fine_tuning.py @@ -0,0 +1,80 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from unittest import mock + +import pytest + +from airflow.providers.google.cloud.hooks.vertex_ai.supervised_fine_tuning import ( + SupervisedFineTuningHook, +) +from tests.providers.google.cloud.utils.base_gcp_mock import ( + mock_base_gcp_hook_default_project_id, +) + +# For no Pydantic environment, we need to skip the tests +pytest.importorskip("google.cloud.aiplatform_v1") +vertexai = pytest.importorskip("vertexai.preview.tuning.sft") + + +TEST_GCP_CONN_ID: str = "test-gcp-conn-id" +GCP_PROJECT = "test-project" +GCP_LOCATION = "us-central1" + +SOURCE_MODEL = "gemini-1.0-pro-002" +TRAIN_DATASET = "gs://cloud-samples-data/ai-platform/generative_ai/sft_train_data.jsonl" + +BASE_STRING = "airflow.providers.google.common.hooks.base_google.{}" +SUPERVISED_FINE_TUNING_STRING = "airflow.providers.google.cloud.hooks.vertex_ai.supervised_fine_tuning.{}" + + +def assert_warning(msg: str, warnings): + assert any(msg in str(w) for w in warnings) + + +class TestSupervisedFineTuningWithDefaultProjectIdHook: + def dummy_get_credentials(self): + pass + + def setup_method(self): + with mock.patch( + BASE_STRING.format("GoogleBaseHook.__init__"), new=mock_base_gcp_hook_default_project_id + ): + self.hook = SupervisedFineTuningHook(gcp_conn_id=TEST_GCP_CONN_ID) + self.hook.get_credentials = self.dummy_get_credentials + + @mock.patch("vertexai.preview.tuning.sft.train") + def test_train(self, mock_train) -> None: + self.hook.train( + project_id=GCP_PROJECT, + location=GCP_LOCATION, + source_model=SOURCE_MODEL, + train_dataset=TRAIN_DATASET, + ) + + # Assertions + mock_train.assert_called_once_with( + source_model=SOURCE_MODEL, + train_dataset=TRAIN_DATASET, + validation_dataset=None, + epochs=None, + adapter_size=None, + learning_rate_multiplier=None, + tuned_model_display_name=None, + ) diff --git a/tests/providers/google/cloud/operators/vertex_ai/test_supervised_fine_tuning.py b/tests/providers/google/cloud/operators/vertex_ai/test_supervised_fine_tuning.py new file mode 100644 index 0000000000000..24f82b164ede5 --- /dev/null +++ b/tests/providers/google/cloud/operators/vertex_ai/test_supervised_fine_tuning.py @@ -0,0 +1,75 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from unittest import mock + +import pytest + +from airflow.providers.google.cloud.operators.vertex_ai.supervised_fine_tuning import ( + SupervisedFineTuningTrainOperator, +) + +# For no Pydantic environment, we need to skip the tests +pytest.importorskip("google.cloud.aiplatform_v1") +vertexai = pytest.importorskip("vertexai.preview.tuning.sft") + + +VERTEX_AI_PATH = "airflow.providers.google.cloud.operators.vertex_ai.{}" + +TASK_ID = "test_task_id" +GCP_PROJECT = "test-project" +GCP_LOCATION = "test-location" +GCP_CONN_ID = "test-conn" +IMPERSONATION_CHAIN = ["ACCOUNT_1", "ACCOUNT_2", "ACCOUNT_3"] + + +def assert_warning(msg: str, warnings): + assert any(msg in str(w) for w in warnings) + + +class TestVertexAISupervisedFineTuningTrainOperator: + @mock.patch(VERTEX_AI_PATH.format("supervised_fine_tuning.SupervisedFineTuningHook")) + def test_execute(self, mock_hook): + source_model = "gemini-1.0-pro-002" + train_dataset = "gs://cloud-samples-data/ai-platform/generative_ai/sft_train_data.jsonl" + + op = SupervisedFineTuningTrainOperator( + task_id=TASK_ID, + project_id=GCP_PROJECT, + location=GCP_LOCATION, + source_model=source_model, + train_dataset=train_dataset, + gcp_conn_id=GCP_CONN_ID, + impersonation_chain=IMPERSONATION_CHAIN, + ) + op.execute(context={"ti": mock.MagicMock()}) + mock_hook.assert_called_once_with( + gcp_conn_id=GCP_CONN_ID, + impersonation_chain=IMPERSONATION_CHAIN, + ) + mock_hook.return_value.train.assert_called_once_with( + project_id=GCP_PROJECT, + location=GCP_LOCATION, + source_model=source_model, + train_dataset=train_dataset, + adapter_size=None, + epochs=None, + learning_rate_multiplier=None, + tuned_model_display_name=None, + validation_dataset=None, + ) diff --git a/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_supervised_fine_tuning.py b/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_supervised_fine_tuning.py new file mode 100644 index 0000000000000..4689abfddbb98 --- /dev/null +++ b/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_supervised_fine_tuning.py @@ -0,0 +1,68 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +""" +Example Airflow DAG for Google Vertex AI Supervised Fine Tuning Jobs. +""" + +from __future__ import annotations + +import os +from datetime import datetime + +from airflow.models.dag import DAG +from airflow.providers.google.cloud.operators.vertex_ai.supervised_fine_tuning import ( + SupervisedFineTuningTrainOperator, +) + +PROJECT_ID = os.environ.get("SYSTEM_TESTS_GCP_PROJECT", "default") +DAG_ID = "vertex_ai_supervised_fine_tuning_dag" +REGION = "us-central1" +SOURCE_MODEL = "gemini-1.0-pro-002" +TRAIN_DATASET = "gs://cloud-samples-data/ai-platform/generative_ai/sft_train_data.jsonl" +TUNED_MODEL_DISPLAY_NAME = "my_tuned_gemini_model" + +with DAG( + dag_id=DAG_ID, + description="Sample DAG with supervised fine tuning jobs.", + schedule="@once", + start_date=datetime(2024, 1, 1), + catchup=False, + tags=["example", "vertex_ai", "supervised_fine_tuning"], +) as dag: + # [START how_to_cloud_vertex_ai_supervised_fine_tuning_train_operator] + sft_train_task = SupervisedFineTuningTrainOperator( + task_id="sft_train_task", + project_id=PROJECT_ID, + location=REGION, + source_model=SOURCE_MODEL, + train_dataset=TRAIN_DATASET, + tuned_model_display_name=TUNED_MODEL_DISPLAY_NAME, + ) + # [END how_to_cloud_vertex_ai_supervised_fine_tuning_train_operator] + + from tests.system.utils.watcher import watcher + + # This test needs watcher in order to properly mark success/failure + # when "tearDown" task with trigger rule is part of the DAG + list(dag.tasks) >> watcher() + +from tests.system.utils import get_test_run # noqa: E402 + +# Needed to run the example DAG with pytest (see: tests/system/README.md#run_via_pytest) +test_run = get_test_run(dag) From 6b3eac6328b1069dbfdb1edfabfc5e2b227ed70c Mon Sep 17 00:00:00 2001 From: Christian Yarros Date: Tue, 27 Aug 2024 21:54:09 +0000 Subject: [PATCH 02/10] build fix --- .../google/cloud/hooks/vertex_ai/supervised_fine_tuning.py | 2 +- .../cloud/operators/vertex_ai/supervised_fine_tuning.py | 2 +- airflow/providers/google/provider.yaml | 4 +++- generated/provider_dependencies.json | 2 +- 4 files changed, 6 insertions(+), 4 deletions(-) diff --git a/airflow/providers/google/cloud/hooks/vertex_ai/supervised_fine_tuning.py b/airflow/providers/google/cloud/hooks/vertex_ai/supervised_fine_tuning.py index e0d2769cfecf4..c162317263132 100644 --- a/airflow/providers/google/cloud/hooks/vertex_ai/supervised_fine_tuning.py +++ b/airflow/providers/google/cloud/hooks/vertex_ai/supervised_fine_tuning.py @@ -15,7 +15,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -"""This module contains a Google Cloud Vertex AI Generative Model hook.""" +"""This module contains a Google Cloud Vertex AI Supervised Fine Tuning hook.""" from __future__ import annotations diff --git a/airflow/providers/google/cloud/operators/vertex_ai/supervised_fine_tuning.py b/airflow/providers/google/cloud/operators/vertex_ai/supervised_fine_tuning.py index a703bebc2ef69..cc73edd35b55b 100644 --- a/airflow/providers/google/cloud/operators/vertex_ai/supervised_fine_tuning.py +++ b/airflow/providers/google/cloud/operators/vertex_ai/supervised_fine_tuning.py @@ -15,7 +15,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -"""This module contains Google Vertex AI Generative AI operators.""" +"""This module contains Google Vertex AI Supervised Fine Tuning operators.""" from __future__ import annotations diff --git a/airflow/providers/google/provider.yaml b/airflow/providers/google/provider.yaml index e6f38ef26b114..b85158187a961 100644 --- a/airflow/providers/google/provider.yaml +++ b/airflow/providers/google/provider.yaml @@ -112,7 +112,7 @@ dependencies: - google-api-python-client>=2.0.2 - google-auth>=2.29.0 - google-auth-httplib2>=0.0.1 - - google-cloud-aiplatform>=1.57.0 + - google-cloud-aiplatform>=1.63.0 - google-cloud-automl>=2.12.0 # google-cloud-bigquery version 3.21.0 introduced a performance enhancement in QueryJob.result(), # which has led to backward compatibility issues @@ -691,6 +691,7 @@ operators: - airflow.providers.google.cloud.operators.vertex_ai.model_service - airflow.providers.google.cloud.operators.vertex_ai.pipeline_job - airflow.providers.google.cloud.operators.vertex_ai.generative_model + - airflow.providers.google.cloud.operators.vertex_ai.supervised_fine_tuning - integration-name: Google Looker python-modules: - airflow.providers.google.cloud.operators.looker @@ -951,6 +952,7 @@ hooks: - airflow.providers.google.cloud.hooks.vertex_ai.pipeline_job - airflow.providers.google.cloud.hooks.vertex_ai.generative_model - airflow.providers.google.cloud.hooks.vertex_ai.prediction_service + - airflow.providers.google.cloud.hooks.vertex_ai.supervised_fine_tuning - integration-name: Google Looker python-modules: - airflow.providers.google.cloud.hooks.looker diff --git a/generated/provider_dependencies.json b/generated/provider_dependencies.json index 3d0e7841c5286..d4439324a244e 100644 --- a/generated/provider_dependencies.json +++ b/generated/provider_dependencies.json @@ -618,7 +618,7 @@ "google-api-python-client>=2.0.2", "google-auth-httplib2>=0.0.1", "google-auth>=2.29.0", - "google-cloud-aiplatform>=1.57.0", + "google-cloud-aiplatform>=1.63.0", "google-cloud-automl>=2.12.0", "google-cloud-batch>=0.13.0", "google-cloud-bigquery-datatransfer>=3.13.0", From e55bccec5e621d30ec8607bb9ca912ff0f87d7ae Mon Sep 17 00:00:00 2001 From: Christian Yarros Date: Wed, 28 Aug 2024 18:17:41 +0000 Subject: [PATCH 03/10] build,test fix --- .../google/cloud/hooks/vertex_ai/test_supervised_fine_tuning.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/providers/google/cloud/hooks/vertex_ai/test_supervised_fine_tuning.py b/tests/providers/google/cloud/hooks/vertex_ai/test_supervised_fine_tuning.py index b218bc89033d6..8c80ce84e46b6 100644 --- a/tests/providers/google/cloud/hooks/vertex_ai/test_supervised_fine_tuning.py +++ b/tests/providers/google/cloud/hooks/vertex_ai/test_supervised_fine_tuning.py @@ -30,8 +30,6 @@ # For no Pydantic environment, we need to skip the tests pytest.importorskip("google.cloud.aiplatform_v1") -vertexai = pytest.importorskip("vertexai.preview.tuning.sft") - TEST_GCP_CONN_ID: str = "test-gcp-conn-id" GCP_PROJECT = "test-project" From d8ada16b789c5bb6c97a84c5741c3b8cb2b68629 Mon Sep 17 00:00:00 2001 From: Christian Yarros Date: Wed, 28 Aug 2024 19:34:26 +0000 Subject: [PATCH 04/10] unit test build fix --- .../cloud/hooks/vertex_ai/test_supervised_fine_tuning.py | 7 +++---- .../operators/vertex_ai/test_supervised_fine_tuning.py | 8 +++----- 2 files changed, 6 insertions(+), 9 deletions(-) diff --git a/tests/providers/google/cloud/hooks/vertex_ai/test_supervised_fine_tuning.py b/tests/providers/google/cloud/hooks/vertex_ai/test_supervised_fine_tuning.py index 8c80ce84e46b6..3a084e4df2fd5 100644 --- a/tests/providers/google/cloud/hooks/vertex_ai/test_supervised_fine_tuning.py +++ b/tests/providers/google/cloud/hooks/vertex_ai/test_supervised_fine_tuning.py @@ -21,6 +21,9 @@ import pytest +# For no Pydantic environment, we need to skip the tests +pytest.importorskip("google.cloud.aiplatform_v1") + from airflow.providers.google.cloud.hooks.vertex_ai.supervised_fine_tuning import ( SupervisedFineTuningHook, ) @@ -28,9 +31,6 @@ mock_base_gcp_hook_default_project_id, ) -# For no Pydantic environment, we need to skip the tests -pytest.importorskip("google.cloud.aiplatform_v1") - TEST_GCP_CONN_ID: str = "test-gcp-conn-id" GCP_PROJECT = "test-project" GCP_LOCATION = "us-central1" @@ -39,7 +39,6 @@ TRAIN_DATASET = "gs://cloud-samples-data/ai-platform/generative_ai/sft_train_data.jsonl" BASE_STRING = "airflow.providers.google.common.hooks.base_google.{}" -SUPERVISED_FINE_TUNING_STRING = "airflow.providers.google.cloud.hooks.vertex_ai.supervised_fine_tuning.{}" def assert_warning(msg: str, warnings): diff --git a/tests/providers/google/cloud/operators/vertex_ai/test_supervised_fine_tuning.py b/tests/providers/google/cloud/operators/vertex_ai/test_supervised_fine_tuning.py index 24f82b164ede5..5e672003df9da 100644 --- a/tests/providers/google/cloud/operators/vertex_ai/test_supervised_fine_tuning.py +++ b/tests/providers/google/cloud/operators/vertex_ai/test_supervised_fine_tuning.py @@ -20,14 +20,12 @@ import pytest -from airflow.providers.google.cloud.operators.vertex_ai.supervised_fine_tuning import ( - SupervisedFineTuningTrainOperator, -) - # For no Pydantic environment, we need to skip the tests pytest.importorskip("google.cloud.aiplatform_v1") -vertexai = pytest.importorskip("vertexai.preview.tuning.sft") +from airflow.providers.google.cloud.operators.vertex_ai.supervised_fine_tuning import ( + SupervisedFineTuningTrainOperator, +) VERTEX_AI_PATH = "airflow.providers.google.cloud.operators.vertex_ai.{}" From 619ec5eb1aff464f7e5dbdb2276f2378e77a28c8 Mon Sep 17 00:00:00 2001 From: Christian Yarros Date: Wed, 28 Aug 2024 21:55:27 +0000 Subject: [PATCH 05/10] xcom fix --- .../google/cloud/operators/vertex_ai/supervised_fine_tuning.py | 1 - 1 file changed, 1 deletion(-) diff --git a/airflow/providers/google/cloud/operators/vertex_ai/supervised_fine_tuning.py b/airflow/providers/google/cloud/operators/vertex_ai/supervised_fine_tuning.py index cc73edd35b55b..dea32e3ca2742 100644 --- a/airflow/providers/google/cloud/operators/vertex_ai/supervised_fine_tuning.py +++ b/airflow/providers/google/cloud/operators/vertex_ai/supervised_fine_tuning.py @@ -114,4 +114,3 @@ def execute(self, context: Context): self.xcom_push(context, key="tuned_model_name", value=response.tuned_model_name) self.xcom_push(context, key="tuned_model_endpoint_name", value=response.tuned_model_endpoint_name) - return response From 748d964b1090a19db01552b3c7734a2dcfeeff56 Mon Sep 17 00:00:00 2001 From: Christian Yarros Date: Thu, 29 Aug 2024 23:39:33 +0000 Subject: [PATCH 06/10] refactor supervised tuning into generative_model module, PR feedback, tests --- .../cloud/hooks/vertex_ai/generative_model.py | 59 ++++++++- .../hooks/vertex_ai/supervised_fine_tuning.py | 101 --------------- .../operators/vertex_ai/generative_model.py | 91 ++++++++++++++ .../vertex_ai/supervised_fine_tuning.py | 116 ------------------ .../operators/cloud/vertex_ai.rst | 6 +- .../hooks/vertex_ai/test_generative_model.py | 23 ++++ .../vertex_ai/test_supervised_fine_tuning.py | 77 ------------ .../vertex_ai/test_generative_model.py | 39 ++++++ .../vertex_ai/test_supervised_fine_tuning.py | 73 ----------- ...mple_vertex_ai_generative_model_tuning.py} | 10 +- 10 files changed, 217 insertions(+), 378 deletions(-) delete mode 100644 airflow/providers/google/cloud/hooks/vertex_ai/supervised_fine_tuning.py delete mode 100644 airflow/providers/google/cloud/operators/vertex_ai/supervised_fine_tuning.py delete mode 100644 tests/providers/google/cloud/hooks/vertex_ai/test_supervised_fine_tuning.py delete mode 100644 tests/providers/google/cloud/operators/vertex_ai/test_supervised_fine_tuning.py rename tests/system/providers/google/cloud/vertex_ai/{example_vertex_ai_supervised_fine_tuning.py => example_vertex_ai_generative_model_tuning.py} (86%) diff --git a/airflow/providers/google/cloud/hooks/vertex_ai/generative_model.py b/airflow/providers/google/cloud/hooks/vertex_ai/generative_model.py index f494ee2241d30..13c6ebe81fd5f 100644 --- a/airflow/providers/google/cloud/hooks/vertex_ai/generative_model.py +++ b/airflow/providers/google/cloud/hooks/vertex_ai/generative_model.py @@ -19,16 +19,21 @@ from __future__ import annotations -from typing import Sequence +import time +from typing import TYPE_CHECKING, Sequence import vertexai from deprecated import deprecated from vertexai.generative_models import GenerativeModel, Part from vertexai.language_models import TextEmbeddingModel, TextGenerationModel +from vertexai.preview.tuning import sft from airflow.exceptions import AirflowProviderDeprecationWarning from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook +if TYPE_CHECKING: + from google.cloud.aiplatform_v1 import types + class GenerativeModelHook(GoogleBaseHook): """Hook for Google Cloud Vertex AI Generative Model APIs.""" @@ -348,3 +353,55 @@ def generative_model_generate_content( ) return response.text + + @GoogleBaseHook.fallback_to_default_project_id + def supervised_fine_tuning_train( + self, + source_model: str, + train_dataset: str, + location: str, + tuned_model_display_name: str | None = None, + validation_dataset: str | None = None, + epochs: int | None = None, + adapter_size: int | None = None, + learning_rate_multiplier: float | None = None, + project_id: str = PROVIDE_PROJECT_ID, + ) -> types.TuningJob: + """ + Use the Supervised Fine Tuning API to create a tuning job. + + :param source_model: Required. A pre-trained model optimized for performing natural + language tasks such as classification, summarization, extraction, content + creation, and ideation. + :param train_dataset: Required. Cloud Storage URI of your training dataset. The dataset + must be formatted as a JSONL file. For best results, provide at least 100 to 500 examples. + :param location: Required. The ID of the Google Cloud location that the service belongs to. + :param tuned_model_display_name: Optional. Display name of the TunedModel. The name can be up + to 128 characters long and can consist of any UTF-8 characters. + :param validation_dataset: Optional. Cloud Storage URI of your training dataset. The dataset must be + formatted as a JSONL file. For best results, provide at least 100 to 500 examples. + :param epochs: Optional. To optimize performance on a specific dataset, try using a higher + epoch value. Increasing the number of epochs might improve results. However, be cautious + about over-fitting, especially when dealing with small datasets. If over-fitting occurs, + consider lowering the epoch number. + :param adapter_size: Optional. Adapter size for tuning. + :param learning_rate_multiplier: Optional. Multiplier for adjusting the default learning rate. + """ + vertexai.init(project=project_id, location=location, credentials=self.get_credentials()) + + sft_tuning_job = sft.train( + source_model=source_model, + train_dataset=train_dataset, + validation_dataset=validation_dataset, + epochs=epochs, + adapter_size=adapter_size, + learning_rate_multiplier=learning_rate_multiplier, + tuned_model_display_name=tuned_model_display_name, + ) + + # Polling for job completion + while not sft_tuning_job.has_ended: + time.sleep(60) + sft_tuning_job.refresh() + + return sft_tuning_job diff --git a/airflow/providers/google/cloud/hooks/vertex_ai/supervised_fine_tuning.py b/airflow/providers/google/cloud/hooks/vertex_ai/supervised_fine_tuning.py deleted file mode 100644 index c162317263132..0000000000000 --- a/airflow/providers/google/cloud/hooks/vertex_ai/supervised_fine_tuning.py +++ /dev/null @@ -1,101 +0,0 @@ -# -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. -"""This module contains a Google Cloud Vertex AI Supervised Fine Tuning hook.""" - -from __future__ import annotations - -import time -from typing import TYPE_CHECKING, Sequence - -import vertexai -from vertexai.preview.tuning import sft - -if TYPE_CHECKING: - from vertexai.preview.tuning.sft import SupervisedTuningJob - - -from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook - - -class SupervisedFineTuningHook(GoogleBaseHook): - """Hook for Google Cloud Vertex AI Supervised Fine Tuning APIs.""" - - def __init__( - self, - gcp_conn_id: str = "google_cloud_default", - impersonation_chain: str | Sequence[str] | None = None, - **kwargs, - ): - if kwargs.get("delegate_to") is not None: - raise RuntimeError( - "The `delegate_to` parameter has been deprecated before and finally removed in this version" - " of Google Provider. You MUST convert it to `impersonate_chain`" - ) - super().__init__(gcp_conn_id=gcp_conn_id, impersonation_chain=impersonation_chain, **kwargs) - - @GoogleBaseHook.fallback_to_default_project_id - def train( - self, - source_model: str, - train_dataset: str, - location: str, - tuned_model_display_name: str | None = None, - validation_dataset: str | None = None, - epochs: int | None = None, - adapter_size: int | None = None, - learning_rate_multiplier: float | None = None, - project_id: str = PROVIDE_PROJECT_ID, - ) -> SupervisedTuningJob: - """ - Use the Supervised Fine Tuning API to create a tuning job. - - :param source_model: Required. A pre-trained model optimized for performing natural - language tasks such as classification, summarization, extraction, content - creation, and ideation. - :param training_dataset: Required. Cloud Storage URI of your training dataset. The dataset - must be formatted as a JSONL file. For best results, provide at least 100 to 500 examples. - :param location: Required. The ID of the Google Cloud location that the service belongs to. - :param tuned_model_display_name: Optional. Display name of the TunedModel. The name can be up - to 128 characters long and can consist of any UTF-8 characters. - :param validation_dataset: Optional. Cloud Storage URI of your training dataset. The dataset must be - formatted as a JSONL file. For best results, provide at least 100 to 500 examples. - :param epochs: Optional. To optimize performance on a specific dataset, try using a higher - epoch value. Increasing the number of epochs might improve results. However, be cautious - about over-fitting, especially when dealing with small datasets. If over-fitting occurs, - consider lowering the epoch number. - :param adapter_size: Optional. Adapter size for tuning. - :param learning_rate_multiplier: Optional. Multiplier for adjusting the default learning rate. - """ - vertexai.init(project=project_id, location=location, credentials=self.get_credentials()) - - sft_tuning_job = sft.train( - source_model=source_model, - train_dataset=train_dataset, - validation_dataset=validation_dataset, - epochs=epochs, - adapter_size=adapter_size, - learning_rate_multiplier=learning_rate_multiplier, - tuned_model_display_name=tuned_model_display_name, - ) - - # Polling for job completion - while not sft_tuning_job.has_ended: - time.sleep(60) - sft_tuning_job.refresh() - - return sft_tuning_job diff --git a/airflow/providers/google/cloud/operators/vertex_ai/generative_model.py b/airflow/providers/google/cloud/operators/vertex_ai/generative_model.py index 5583e5e8bc664..0aa7cbc6d9324 100644 --- a/airflow/providers/google/cloud/operators/vertex_ai/generative_model.py +++ b/airflow/providers/google/cloud/operators/vertex_ai/generative_model.py @@ -22,6 +22,7 @@ from typing import TYPE_CHECKING, Sequence from deprecated import deprecated +from google.cloud.aiplatform_v1 import types from airflow.exceptions import AirflowProviderDeprecationWarning from airflow.providers.google.cloud.hooks.vertex_ai.generative_model import GenerativeModelHook @@ -571,3 +572,93 @@ def execute(self, context: Context): self.xcom_push(context, key="model_response", value=response) return response + + +class SupervisedFineTuningTrainOperator(GoogleCloudBaseOperator): + """ + Use the Supervised Fine Tuning API to create a tuning job. + + :param source_model: Required. A pre-trained model optimized for performing natural + language tasks such as classification, summarization, extraction, content + creation, and ideation. + :param train_dataset: Required. Cloud Storage URI of your training dataset. The dataset + must be formatted as a JSONL file. For best results, provide at least 100 to 500 examples. + :param project_id: Required. The ID of the Google Cloud project that the + service belongs to. + :param location: Required. The ID of the Google Cloud location that the service belongs to. + :param tuned_model_display_name: Optional. Display name of the TunedModel. The name can be up + to 128 characters long and can consist of any UTF-8 characters. + :param validation_dataset: Optional. Cloud Storage URI of your training dataset. The dataset must be + formatted as a JSONL file. For best results, provide at least 100 to 500 examples. + :param epochs: Optional. To optimize performance on a specific dataset, try using a higher + epoch value. Increasing the number of epochs might improve results. However, be cautious + about over-fitting, especially when dealing with small datasets. If over-fitting occurs, + consider lowering the epoch number. + :param adapter_size: Optional. Adapter size for tuning. + :param learning_multiplier_rate: Optional. Multiplier for adjusting the default learning rate. + :param gcp_conn_id: The connection ID to use connecting to Google Cloud. + :param impersonation_chain: Optional service account to impersonate using short-term + credentials, or chained list of accounts required to get the access_token + of the last account in the list, which will be impersonated in the request. + If set as a string, the account must grant the originating account + the Service Account Token Creator IAM role. + If set as a sequence, the identities from the list must grant + Service Account Token Creator IAM role to the directly preceding identity, with first + account from the list granting this role to the originating account (templated). + """ + + template_fields = ("location", "project_id", "impersonation_chain", "train_dataset", "validation_dataset") + + def __init__( + self, + *, + source_model: str, + train_dataset: str, + project_id: str, + location: str, + tuned_model_display_name: str | None = None, + validation_dataset: str | None = None, + epochs: int | None = None, + adapter_size: int | None = None, + learning_rate_multiplier: float | None = None, + gcp_conn_id: str = "google_cloud_default", + impersonation_chain: str | Sequence[str] | None = None, + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.source_model = source_model + self.train_dataset = train_dataset + self.tuned_model_display_name = tuned_model_display_name + self.validation_dataset = validation_dataset + self.epochs = epochs + self.adapter_size = adapter_size + self.learning_rate_multiplier = learning_rate_multiplier + self.project_id = project_id + self.location = location + self.gcp_conn_id = gcp_conn_id + self.impersonation_chain = impersonation_chain + + def execute(self, context: Context): + self.hook = GenerativeModelHook( + gcp_conn_id=self.gcp_conn_id, + impersonation_chain=self.impersonation_chain, + ) + response = self.hook.supervised_fine_tuning_train( + source_model=self.source_model, + train_dataset=self.train_dataset, + project_id=self.project_id, + location=self.location, + validation_dataset=self.validation_dataset, + epochs=self.epochs, + adapter_size=self.adapter_size, + learning_rate_multiplier=self.learning_rate_multiplier, + tuned_model_display_name=self.tuned_model_display_name, + ) + + self.log.info("Tuned Model Name: %s", response.tuned_model_name) + self.log.info("Tuned Model Endpoint Name: %s", response.tuned_model_endpoint_name) + + self.xcom_push(context, key="tuned_model_name", value=response.tuned_model_name) + self.xcom_push(context, key="tuned_model_endpoint_name", value=response.tuned_model_endpoint_name) + + return types.TuningJob.to_dict(response) diff --git a/airflow/providers/google/cloud/operators/vertex_ai/supervised_fine_tuning.py b/airflow/providers/google/cloud/operators/vertex_ai/supervised_fine_tuning.py deleted file mode 100644 index dea32e3ca2742..0000000000000 --- a/airflow/providers/google/cloud/operators/vertex_ai/supervised_fine_tuning.py +++ /dev/null @@ -1,116 +0,0 @@ -# -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. -"""This module contains Google Vertex AI Supervised Fine Tuning operators.""" - -from __future__ import annotations - -from typing import TYPE_CHECKING, Sequence - -from airflow.providers.google.cloud.hooks.vertex_ai.supervised_fine_tuning import SupervisedFineTuningHook -from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator - -if TYPE_CHECKING: - from airflow.utils.context import Context - - -class SupervisedFineTuningTrainOperator(GoogleCloudBaseOperator): - """ - Use the Supervised Fine Tuning API to create a tuning job. - - :param source_model: Required. A pre-trained model optimized for performing natural - language tasks such as classification, summarization, extraction, content - creation, and ideation. - :param training_dataset: Required. Cloud Storage URI of your training dataset. The dataset - must be formatted as a JSONL file. For best results, provide at least 100 to 500 examples. - :param project_id: Required. The ID of the Google Cloud project that the - service belongs to. - :param location: Required. The ID of the Google Cloud location that the service belongs to. - :param tuned_model_display_name: Optional. Display name of the TunedModel. The name can be up - to 128 characters long and can consist of any UTF-8 characters. - :param validation_dataset: Optional. Cloud Storage URI of your training dataset. The dataset must be - formatted as a JSONL file. For best results, provide at least 100 to 500 examples. - :param epochs: Optional. To optimize performance on a specific dataset, try using a higher - epoch value. Increasing the number of epochs might improve results. However, be cautious - about over-fitting, especially when dealing with small datasets. If over-fitting occurs, - consider lowering the epoch number. - :param adapter_size: Optional. Adapter size for tuning. - :param learning_multiplier_rate: Optional. Multiplier for adjusting the default learning rate. - :param gcp_conn_id: The connection ID to use connecting to Google Cloud. - :param impersonation_chain: Optional service account to impersonate using short-term - credentials, or chained list of accounts required to get the access_token - of the last account in the list, which will be impersonated in the request. - If set as a string, the account must grant the originating account - the Service Account Token Creator IAM role. - If set as a sequence, the identities from the list must grant - Service Account Token Creator IAM role to the directly preceding identity, with first - account from the list granting this role to the originating account (templated). - """ - - template_fields = ("location", "project_id", "impersonation_chain", "train_dataset", "validation_dataset") - - def __init__( - self, - *, - source_model: str, - train_dataset: str, - project_id: str, - location: str, - tuned_model_display_name: str | None = None, - validation_dataset: str | None = None, - epochs: int | None = None, - adapter_size: int | None = None, - learning_rate_multiplier: float | None = None, - gcp_conn_id: str = "google_cloud_default", - impersonation_chain: str | Sequence[str] | None = None, - **kwargs, - ) -> None: - super().__init__(**kwargs) - self.source_model = source_model - self.train_dataset = train_dataset - self.tuned_model_display_name = tuned_model_display_name - self.validation_dataset = validation_dataset - self.epochs = epochs - self.adapter_size = adapter_size - self.learning_rate_multiplier = learning_rate_multiplier - self.project_id = project_id - self.location = location - self.gcp_conn_id = gcp_conn_id - self.impersonation_chain = impersonation_chain - - def execute(self, context: Context): - self.hook = SupervisedFineTuningHook( - gcp_conn_id=self.gcp_conn_id, - impersonation_chain=self.impersonation_chain, - ) - response = self.hook.train( - source_model=self.source_model, - train_dataset=self.train_dataset, - project_id=self.project_id, - location=self.location, - validation_dataset=self.validation_dataset, - epochs=self.epochs, - adapter_size=self.adapter_size, - learning_rate_multiplier=self.learning_rate_multiplier, - tuned_model_display_name=self.tuned_model_display_name, - ) - - self.log.info("Tuned Model Name: %s", response.tuned_model_name) - self.log.info("Tuned Model Endpoint Name: %s", response.tuned_model_endpoint_name) - - self.xcom_push(context, key="tuned_model_name", value=response.tuned_model_name) - self.xcom_push(context, key="tuned_model_endpoint_name", value=response.tuned_model_endpoint_name) diff --git a/docs/apache-airflow-providers-google/operators/cloud/vertex_ai.rst b/docs/apache-airflow-providers-google/operators/cloud/vertex_ai.rst index fb1ca4bf60cd7..6cabd41fb47db 100644 --- a/docs/apache-airflow-providers-google/operators/cloud/vertex_ai.rst +++ b/docs/apache-airflow-providers-google/operators/cloud/vertex_ai.rst @@ -615,11 +615,7 @@ The operator returns the model's response in :ref:`XCom ` under ` :start-after: [START how_to_cloud_vertex_ai_generative_model_generate_content_operator] :end-before: [END how_to_cloud_vertex_ai_generative_model_generate_content_operator] - -Performing Supervised Fine Tuning -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - -To train and deploy a tuned generative model to an endpoint you can use +To run a supervised fine tuning job you can use :class:`~airflow.providers.google.cloud.operators.vertex_ai.supervised_fine_tuning.SupervisedFineTuningTrainOperator`. The operator returns the tuned model's endpoint name in :ref:`XCom ` under ``tuned_model_endpoint_name`` key. diff --git a/tests/providers/google/cloud/hooks/vertex_ai/test_generative_model.py b/tests/providers/google/cloud/hooks/vertex_ai/test_generative_model.py index 03c8ca2809335..cee720d519392 100644 --- a/tests/providers/google/cloud/hooks/vertex_ai/test_generative_model.py +++ b/tests/providers/google/cloud/hooks/vertex_ai/test_generative_model.py @@ -70,6 +70,9 @@ TEST_MEDIA_GCS_PATH = "gs://download.tensorflow.org/example_images/320px-Felis_catus-cat_on_snow.jpg" TEST_MIME_TYPE = "image/jpeg" +SOURCE_MODEL = "gemini-1.0-pro-002" +TRAIN_DATASET = "gs://cloud-samples-data/ai-platform/generative_ai/sft_train_data.jsonl" + BASE_STRING = "airflow.providers.google.common.hooks.base_google.{}" GENERATIVE_MODEL_STRING = "airflow.providers.google.cloud.hooks.vertex_ai.generative_model.{}" @@ -194,3 +197,23 @@ def test_generative_model_generate_content(self, mock_model) -> None: generation_config=TEST_GENERATION_CONFIG, safety_settings=TEST_SAFETY_SETTINGS, ) + + @mock.patch("vertexai.preview.tuning.sft.train") + def test_supervised_fine_tuning_train(self, mock_sft_train) -> None: + self.hook.supervised_fine_tuning_train( + project_id=GCP_PROJECT, + location=GCP_LOCATION, + source_model=SOURCE_MODEL, + train_dataset=TRAIN_DATASET, + ) + + # Assertions + mock_sft_train.assert_called_once_with( + source_model=SOURCE_MODEL, + train_dataset=TRAIN_DATASET, + validation_dataset=None, + epochs=None, + adapter_size=None, + learning_rate_multiplier=None, + tuned_model_display_name=None, + ) diff --git a/tests/providers/google/cloud/hooks/vertex_ai/test_supervised_fine_tuning.py b/tests/providers/google/cloud/hooks/vertex_ai/test_supervised_fine_tuning.py deleted file mode 100644 index 3a084e4df2fd5..0000000000000 --- a/tests/providers/google/cloud/hooks/vertex_ai/test_supervised_fine_tuning.py +++ /dev/null @@ -1,77 +0,0 @@ -# -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. -from __future__ import annotations - -from unittest import mock - -import pytest - -# For no Pydantic environment, we need to skip the tests -pytest.importorskip("google.cloud.aiplatform_v1") - -from airflow.providers.google.cloud.hooks.vertex_ai.supervised_fine_tuning import ( - SupervisedFineTuningHook, -) -from tests.providers.google.cloud.utils.base_gcp_mock import ( - mock_base_gcp_hook_default_project_id, -) - -TEST_GCP_CONN_ID: str = "test-gcp-conn-id" -GCP_PROJECT = "test-project" -GCP_LOCATION = "us-central1" - -SOURCE_MODEL = "gemini-1.0-pro-002" -TRAIN_DATASET = "gs://cloud-samples-data/ai-platform/generative_ai/sft_train_data.jsonl" - -BASE_STRING = "airflow.providers.google.common.hooks.base_google.{}" - - -def assert_warning(msg: str, warnings): - assert any(msg in str(w) for w in warnings) - - -class TestSupervisedFineTuningWithDefaultProjectIdHook: - def dummy_get_credentials(self): - pass - - def setup_method(self): - with mock.patch( - BASE_STRING.format("GoogleBaseHook.__init__"), new=mock_base_gcp_hook_default_project_id - ): - self.hook = SupervisedFineTuningHook(gcp_conn_id=TEST_GCP_CONN_ID) - self.hook.get_credentials = self.dummy_get_credentials - - @mock.patch("vertexai.preview.tuning.sft.train") - def test_train(self, mock_train) -> None: - self.hook.train( - project_id=GCP_PROJECT, - location=GCP_LOCATION, - source_model=SOURCE_MODEL, - train_dataset=TRAIN_DATASET, - ) - - # Assertions - mock_train.assert_called_once_with( - source_model=SOURCE_MODEL, - train_dataset=TRAIN_DATASET, - validation_dataset=None, - epochs=None, - adapter_size=None, - learning_rate_multiplier=None, - tuned_model_display_name=None, - ) diff --git a/tests/providers/google/cloud/operators/vertex_ai/test_generative_model.py b/tests/providers/google/cloud/operators/vertex_ai/test_generative_model.py index dc070b79d37ad..cdf1bc5f2a7cc 100644 --- a/tests/providers/google/cloud/operators/vertex_ai/test_generative_model.py +++ b/tests/providers/google/cloud/operators/vertex_ai/test_generative_model.py @@ -35,6 +35,7 @@ PromptLanguageModelOperator, PromptMultimodalModelOperator, PromptMultimodalModelWithMediaOperator, + SupervisedFineTuningTrainOperator, TextEmbeddingModelGetEmbeddingsOperator, TextGenerationModelPredictOperator, ) @@ -390,3 +391,41 @@ def test_execute(self, mock_hook): safety_settings=safety_settings, pretrained_model=pretrained_model, ) + + +class TestVertexAISupervisedFineTuningTrainOperator: + @mock.patch(VERTEX_AI_PATH.format("generative_model.GenerativeModelHook")) + @mock.patch("google.cloud.aiplatform_v1.types.TuningJob.to_dict") + def test_execute( + self, + to_dict_mock, + mock_hook, + ): + source_model = "gemini-1.0-pro-002" + train_dataset = "gs://cloud-samples-data/ai-platform/generative_ai/sft_train_data.jsonl" + + op = SupervisedFineTuningTrainOperator( + task_id=TASK_ID, + project_id=GCP_PROJECT, + location=GCP_LOCATION, + source_model=source_model, + train_dataset=train_dataset, + gcp_conn_id=GCP_CONN_ID, + impersonation_chain=IMPERSONATION_CHAIN, + ) + op.execute(context={"ti": mock.MagicMock()}) + mock_hook.assert_called_once_with( + gcp_conn_id=GCP_CONN_ID, + impersonation_chain=IMPERSONATION_CHAIN, + ) + mock_hook.return_value.supervised_fine_tuning_train.assert_called_once_with( + project_id=GCP_PROJECT, + location=GCP_LOCATION, + source_model=source_model, + train_dataset=train_dataset, + adapter_size=None, + epochs=None, + learning_rate_multiplier=None, + tuned_model_display_name=None, + validation_dataset=None, + ) diff --git a/tests/providers/google/cloud/operators/vertex_ai/test_supervised_fine_tuning.py b/tests/providers/google/cloud/operators/vertex_ai/test_supervised_fine_tuning.py deleted file mode 100644 index 5e672003df9da..0000000000000 --- a/tests/providers/google/cloud/operators/vertex_ai/test_supervised_fine_tuning.py +++ /dev/null @@ -1,73 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. -from __future__ import annotations - -from unittest import mock - -import pytest - -# For no Pydantic environment, we need to skip the tests -pytest.importorskip("google.cloud.aiplatform_v1") - -from airflow.providers.google.cloud.operators.vertex_ai.supervised_fine_tuning import ( - SupervisedFineTuningTrainOperator, -) - -VERTEX_AI_PATH = "airflow.providers.google.cloud.operators.vertex_ai.{}" - -TASK_ID = "test_task_id" -GCP_PROJECT = "test-project" -GCP_LOCATION = "test-location" -GCP_CONN_ID = "test-conn" -IMPERSONATION_CHAIN = ["ACCOUNT_1", "ACCOUNT_2", "ACCOUNT_3"] - - -def assert_warning(msg: str, warnings): - assert any(msg in str(w) for w in warnings) - - -class TestVertexAISupervisedFineTuningTrainOperator: - @mock.patch(VERTEX_AI_PATH.format("supervised_fine_tuning.SupervisedFineTuningHook")) - def test_execute(self, mock_hook): - source_model = "gemini-1.0-pro-002" - train_dataset = "gs://cloud-samples-data/ai-platform/generative_ai/sft_train_data.jsonl" - - op = SupervisedFineTuningTrainOperator( - task_id=TASK_ID, - project_id=GCP_PROJECT, - location=GCP_LOCATION, - source_model=source_model, - train_dataset=train_dataset, - gcp_conn_id=GCP_CONN_ID, - impersonation_chain=IMPERSONATION_CHAIN, - ) - op.execute(context={"ti": mock.MagicMock()}) - mock_hook.assert_called_once_with( - gcp_conn_id=GCP_CONN_ID, - impersonation_chain=IMPERSONATION_CHAIN, - ) - mock_hook.return_value.train.assert_called_once_with( - project_id=GCP_PROJECT, - location=GCP_LOCATION, - source_model=source_model, - train_dataset=train_dataset, - adapter_size=None, - epochs=None, - learning_rate_multiplier=None, - tuned_model_display_name=None, - validation_dataset=None, - ) diff --git a/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_supervised_fine_tuning.py b/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_generative_model_tuning.py similarity index 86% rename from tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_supervised_fine_tuning.py rename to tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_generative_model_tuning.py index 4689abfddbb98..1bc02f6e92b21 100644 --- a/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_supervised_fine_tuning.py +++ b/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_generative_model_tuning.py @@ -17,7 +17,7 @@ # under the License. """ -Example Airflow DAG for Google Vertex AI Supervised Fine Tuning Jobs. +Example Airflow DAG for Google Vertex AI Generative Model Tuning Tasks. """ from __future__ import annotations @@ -26,12 +26,12 @@ from datetime import datetime from airflow.models.dag import DAG -from airflow.providers.google.cloud.operators.vertex_ai.supervised_fine_tuning import ( +from airflow.providers.google.cloud.operators.vertex_ai.generative_model_tuning import ( SupervisedFineTuningTrainOperator, ) PROJECT_ID = os.environ.get("SYSTEM_TESTS_GCP_PROJECT", "default") -DAG_ID = "vertex_ai_supervised_fine_tuning_dag" +DAG_ID = "vertex_ai_generative_model_tuning_dag" REGION = "us-central1" SOURCE_MODEL = "gemini-1.0-pro-002" TRAIN_DATASET = "gs://cloud-samples-data/ai-platform/generative_ai/sft_train_data.jsonl" @@ -39,11 +39,11 @@ with DAG( dag_id=DAG_ID, - description="Sample DAG with supervised fine tuning jobs.", + description="Sample DAG with generative model tuning tasks.", schedule="@once", start_date=datetime(2024, 1, 1), catchup=False, - tags=["example", "vertex_ai", "supervised_fine_tuning"], + tags=["example", "vertex_ai", "generative_model_tuning"], ) as dag: # [START how_to_cloud_vertex_ai_supervised_fine_tuning_train_operator] sft_train_task = SupervisedFineTuningTrainOperator( From 3b176770b699cd0bfdfe1e29ece75778a4654617 Mon Sep 17 00:00:00 2001 From: Christian Yarros Date: Thu, 29 Aug 2024 23:44:42 +0000 Subject: [PATCH 07/10] minor system test fix --- .../vertex_ai/example_vertex_ai_generative_model_tuning.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_generative_model_tuning.py b/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_generative_model_tuning.py index 1bc02f6e92b21..18958cb409e65 100644 --- a/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_generative_model_tuning.py +++ b/tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_generative_model_tuning.py @@ -26,7 +26,7 @@ from datetime import datetime from airflow.models.dag import DAG -from airflow.providers.google.cloud.operators.vertex_ai.generative_model_tuning import ( +from airflow.providers.google.cloud.operators.vertex_ai.generative_model import ( SupervisedFineTuningTrainOperator, ) @@ -43,7 +43,7 @@ schedule="@once", start_date=datetime(2024, 1, 1), catchup=False, - tags=["example", "vertex_ai", "generative_model_tuning"], + tags=["example", "vertex_ai", "generative_model"], ) as dag: # [START how_to_cloud_vertex_ai_supervised_fine_tuning_train_operator] sft_train_task = SupervisedFineTuningTrainOperator( From 838f16f7243163f4b16964215b4ee2efc227fffd Mon Sep 17 00:00:00 2001 From: Christian Yarros Date: Thu, 29 Aug 2024 23:54:18 +0000 Subject: [PATCH 08/10] update provider.yaml --- airflow/providers/google/provider.yaml | 2 -- 1 file changed, 2 deletions(-) diff --git a/airflow/providers/google/provider.yaml b/airflow/providers/google/provider.yaml index b85158187a961..b241f7c44cb66 100644 --- a/airflow/providers/google/provider.yaml +++ b/airflow/providers/google/provider.yaml @@ -691,7 +691,6 @@ operators: - airflow.providers.google.cloud.operators.vertex_ai.model_service - airflow.providers.google.cloud.operators.vertex_ai.pipeline_job - airflow.providers.google.cloud.operators.vertex_ai.generative_model - - airflow.providers.google.cloud.operators.vertex_ai.supervised_fine_tuning - integration-name: Google Looker python-modules: - airflow.providers.google.cloud.operators.looker @@ -952,7 +951,6 @@ hooks: - airflow.providers.google.cloud.hooks.vertex_ai.pipeline_job - airflow.providers.google.cloud.hooks.vertex_ai.generative_model - airflow.providers.google.cloud.hooks.vertex_ai.prediction_service - - airflow.providers.google.cloud.hooks.vertex_ai.supervised_fine_tuning - integration-name: Google Looker python-modules: - airflow.providers.google.cloud.hooks.looker From 415e752729740128f2bafc3f5d826773dc16bafc Mon Sep 17 00:00:00 2001 From: Christian Yarros Date: Fri, 30 Aug 2024 00:39:09 +0000 Subject: [PATCH 09/10] doc fix --- .../operators/cloud/vertex_ai.rst | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/apache-airflow-providers-google/operators/cloud/vertex_ai.rst b/docs/apache-airflow-providers-google/operators/cloud/vertex_ai.rst index 6cabd41fb47db..1bdd1f2dbec8e 100644 --- a/docs/apache-airflow-providers-google/operators/cloud/vertex_ai.rst +++ b/docs/apache-airflow-providers-google/operators/cloud/vertex_ai.rst @@ -582,7 +582,7 @@ To get a pipeline job list you can use :start-after: [START how_to_cloud_vertex_ai_list_pipeline_job_operator] :end-before: [END how_to_cloud_vertex_ai_list_pipeline_job_operator] -Interacting with a Generative Model +Interacting with Generative AI ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ To generate a prediction via language model you can use @@ -616,7 +616,7 @@ The operator returns the model's response in :ref:`XCom ` under ` :end-before: [END how_to_cloud_vertex_ai_generative_model_generate_content_operator] To run a supervised fine tuning job you can use -:class:`~airflow.providers.google.cloud.operators.vertex_ai.supervised_fine_tuning.SupervisedFineTuningTrainOperator`. +:class:`~airflow.providers.google.cloud.operators.vertex_ai.generative_model.SupervisedFineTuningTrainOperator`. The operator returns the tuned model's endpoint name in :ref:`XCom ` under ``tuned_model_endpoint_name`` key. .. exampleinclude:: /../../tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_supervised_fine_tuning.py From ddf915049a6eb97140e759b9d931165ca398d66c Mon Sep 17 00:00:00 2001 From: Christian Yarros Date: Fri, 30 Aug 2024 01:32:16 +0000 Subject: [PATCH 10/10] Update Vertex AI Documentation --- .../operators/cloud/vertex_ai.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/apache-airflow-providers-google/operators/cloud/vertex_ai.rst b/docs/apache-airflow-providers-google/operators/cloud/vertex_ai.rst index 1bdd1f2dbec8e..e359cd21d25cf 100644 --- a/docs/apache-airflow-providers-google/operators/cloud/vertex_ai.rst +++ b/docs/apache-airflow-providers-google/operators/cloud/vertex_ai.rst @@ -619,7 +619,7 @@ To run a supervised fine tuning job you can use :class:`~airflow.providers.google.cloud.operators.vertex_ai.generative_model.SupervisedFineTuningTrainOperator`. The operator returns the tuned model's endpoint name in :ref:`XCom ` under ``tuned_model_endpoint_name`` key. -.. exampleinclude:: /../../tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_supervised_fine_tuning.py +.. exampleinclude:: /../../tests/system/providers/google/cloud/vertex_ai/example_vertex_ai_generative_model_tuning.py :language: python :dedent: 4 :start-after: [START how_to_cloud_vertex_ai_supervised_fine_tuning_train_operator]