diff --git a/.github/sync-repo-settings.yaml b/.github/sync-repo-settings.yaml index 8920c91903..5a1844bb70 100644 --- a/.github/sync-repo-settings.yaml +++ b/.github/sync-repo-settings.yaml @@ -14,4 +14,8 @@ branchProtectionRules: requiresStrictStatusChecks: true requiredStatusCheckContexts: - 'cla/google' - - 'Presubmit - Unit Tests' + - 'Presubmit - Lint and Coverage' + - 'Presubmit - Unit Tests Python 3.8' + - 'Presubmit - Unit Tests Python 3.9' + - 'Presubmit - Unit Tests Python 3.10' + - 'Presubmit - Unit Tests Python 3.11' \ No newline at end of file diff --git a/.release-please-manifest.json b/.release-please-manifest.json index ea40fbb6c7..b126785d7e 100644 --- a/.release-please-manifest.json +++ b/.release-please-manifest.json @@ -1,3 +1,3 @@ { - ".": "1.33.0" + ".": "1.33.1" } diff --git a/CHANGELOG.md b/CHANGELOG.md index ade59b087a..250a1b982b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,12 @@ # Changelog +## [1.33.1](https://github.com/googleapis/python-aiplatform/compare/v1.33.0...v1.33.1) (2023-09-20) + + +### Bug Fixes + +* Lightning trainer fails to be unwrapped in remote training ([8271301](https://github.com/googleapis/python-aiplatform/commit/8271301454814b233a630d1c18ebe5e4833fcec2)) + ## [1.33.0](https://github.com/googleapis/python-aiplatform/compare/v1.32.0...v1.33.0) (2023-09-18) diff --git a/google/cloud/aiplatform/gapic_version.py b/google/cloud/aiplatform/gapic_version.py index 11cb6e1c3a..c2a948a447 100644 --- a/google/cloud/aiplatform/gapic_version.py +++ b/google/cloud/aiplatform/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.33.0" # {x-release-please-version} +__version__ = "1.33.1" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1/schema/predict/instance/gapic_version.py b/google/cloud/aiplatform/v1/schema/predict/instance/gapic_version.py index 11cb6e1c3a..c2a948a447 100644 --- a/google/cloud/aiplatform/v1/schema/predict/instance/gapic_version.py +++ b/google/cloud/aiplatform/v1/schema/predict/instance/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.33.0" # {x-release-please-version} +__version__ = "1.33.1" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1/schema/predict/instance_v1/gapic_version.py b/google/cloud/aiplatform/v1/schema/predict/instance_v1/gapic_version.py index 11cb6e1c3a..c2a948a447 100644 --- a/google/cloud/aiplatform/v1/schema/predict/instance_v1/gapic_version.py +++ b/google/cloud/aiplatform/v1/schema/predict/instance_v1/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.33.0" # {x-release-please-version} +__version__ = "1.33.1" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1/schema/predict/params/gapic_version.py b/google/cloud/aiplatform/v1/schema/predict/params/gapic_version.py index 11cb6e1c3a..c2a948a447 100644 --- a/google/cloud/aiplatform/v1/schema/predict/params/gapic_version.py +++ b/google/cloud/aiplatform/v1/schema/predict/params/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.33.0" # {x-release-please-version} +__version__ = "1.33.1" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1/schema/predict/params_v1/gapic_version.py b/google/cloud/aiplatform/v1/schema/predict/params_v1/gapic_version.py index 11cb6e1c3a..c2a948a447 100644 --- a/google/cloud/aiplatform/v1/schema/predict/params_v1/gapic_version.py +++ b/google/cloud/aiplatform/v1/schema/predict/params_v1/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.33.0" # {x-release-please-version} +__version__ = "1.33.1" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1/schema/predict/prediction/gapic_version.py b/google/cloud/aiplatform/v1/schema/predict/prediction/gapic_version.py index 11cb6e1c3a..c2a948a447 100644 --- a/google/cloud/aiplatform/v1/schema/predict/prediction/gapic_version.py +++ b/google/cloud/aiplatform/v1/schema/predict/prediction/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.33.0" # {x-release-please-version} +__version__ = "1.33.1" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1/schema/predict/prediction_v1/gapic_version.py b/google/cloud/aiplatform/v1/schema/predict/prediction_v1/gapic_version.py index 11cb6e1c3a..c2a948a447 100644 --- a/google/cloud/aiplatform/v1/schema/predict/prediction_v1/gapic_version.py +++ b/google/cloud/aiplatform/v1/schema/predict/prediction_v1/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.33.0" # {x-release-please-version} +__version__ = "1.33.1" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1/schema/trainingjob/definition/gapic_version.py b/google/cloud/aiplatform/v1/schema/trainingjob/definition/gapic_version.py index 11cb6e1c3a..c2a948a447 100644 --- a/google/cloud/aiplatform/v1/schema/trainingjob/definition/gapic_version.py +++ b/google/cloud/aiplatform/v1/schema/trainingjob/definition/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.33.0" # {x-release-please-version} +__version__ = "1.33.1" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1/schema/trainingjob/definition_v1/gapic_version.py b/google/cloud/aiplatform/v1/schema/trainingjob/definition_v1/gapic_version.py index 11cb6e1c3a..c2a948a447 100644 --- a/google/cloud/aiplatform/v1/schema/trainingjob/definition_v1/gapic_version.py +++ b/google/cloud/aiplatform/v1/schema/trainingjob/definition_v1/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.33.0" # {x-release-please-version} +__version__ = "1.33.1" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1beta1/schema/predict/instance/gapic_version.py b/google/cloud/aiplatform/v1beta1/schema/predict/instance/gapic_version.py index 11cb6e1c3a..c2a948a447 100644 --- a/google/cloud/aiplatform/v1beta1/schema/predict/instance/gapic_version.py +++ b/google/cloud/aiplatform/v1beta1/schema/predict/instance/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.33.0" # {x-release-please-version} +__version__ = "1.33.1" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1beta1/schema/predict/instance_v1beta1/gapic_version.py b/google/cloud/aiplatform/v1beta1/schema/predict/instance_v1beta1/gapic_version.py index 11cb6e1c3a..c2a948a447 100644 --- a/google/cloud/aiplatform/v1beta1/schema/predict/instance_v1beta1/gapic_version.py +++ b/google/cloud/aiplatform/v1beta1/schema/predict/instance_v1beta1/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.33.0" # {x-release-please-version} +__version__ = "1.33.1" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1beta1/schema/predict/params/gapic_version.py b/google/cloud/aiplatform/v1beta1/schema/predict/params/gapic_version.py index 11cb6e1c3a..c2a948a447 100644 --- a/google/cloud/aiplatform/v1beta1/schema/predict/params/gapic_version.py +++ b/google/cloud/aiplatform/v1beta1/schema/predict/params/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.33.0" # {x-release-please-version} +__version__ = "1.33.1" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1beta1/schema/predict/params_v1beta1/gapic_version.py b/google/cloud/aiplatform/v1beta1/schema/predict/params_v1beta1/gapic_version.py index 11cb6e1c3a..c2a948a447 100644 --- a/google/cloud/aiplatform/v1beta1/schema/predict/params_v1beta1/gapic_version.py +++ b/google/cloud/aiplatform/v1beta1/schema/predict/params_v1beta1/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.33.0" # {x-release-please-version} +__version__ = "1.33.1" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1beta1/schema/predict/prediction/gapic_version.py b/google/cloud/aiplatform/v1beta1/schema/predict/prediction/gapic_version.py index 11cb6e1c3a..c2a948a447 100644 --- a/google/cloud/aiplatform/v1beta1/schema/predict/prediction/gapic_version.py +++ b/google/cloud/aiplatform/v1beta1/schema/predict/prediction/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.33.0" # {x-release-please-version} +__version__ = "1.33.1" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1beta1/schema/predict/prediction_v1beta1/gapic_version.py b/google/cloud/aiplatform/v1beta1/schema/predict/prediction_v1beta1/gapic_version.py index 11cb6e1c3a..c2a948a447 100644 --- a/google/cloud/aiplatform/v1beta1/schema/predict/prediction_v1beta1/gapic_version.py +++ b/google/cloud/aiplatform/v1beta1/schema/predict/prediction_v1beta1/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.33.0" # {x-release-please-version} +__version__ = "1.33.1" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1beta1/schema/trainingjob/definition/gapic_version.py b/google/cloud/aiplatform/v1beta1/schema/trainingjob/definition/gapic_version.py index 11cb6e1c3a..c2a948a447 100644 --- a/google/cloud/aiplatform/v1beta1/schema/trainingjob/definition/gapic_version.py +++ b/google/cloud/aiplatform/v1beta1/schema/trainingjob/definition/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.33.0" # {x-release-please-version} +__version__ = "1.33.1" # {x-release-please-version} diff --git a/google/cloud/aiplatform/v1beta1/schema/trainingjob/definition_v1beta1/gapic_version.py b/google/cloud/aiplatform/v1beta1/schema/trainingjob/definition_v1beta1/gapic_version.py index 11cb6e1c3a..c2a948a447 100644 --- a/google/cloud/aiplatform/v1beta1/schema/trainingjob/definition_v1beta1/gapic_version.py +++ b/google/cloud/aiplatform/v1beta1/schema/trainingjob/definition_v1beta1/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.33.0" # {x-release-please-version} +__version__ = "1.33.1" # {x-release-please-version} diff --git a/google/cloud/aiplatform/version.py b/google/cloud/aiplatform/version.py index fcef003494..c4fd607991 100644 --- a/google/cloud/aiplatform/version.py +++ b/google/cloud/aiplatform/version.py @@ -15,4 +15,4 @@ # limitations under the License. # -__version__ = "1.33.0" +__version__ = "1.33.1" diff --git a/google/cloud/aiplatform_v1/gapic_version.py b/google/cloud/aiplatform_v1/gapic_version.py index 11cb6e1c3a..c2a948a447 100644 --- a/google/cloud/aiplatform_v1/gapic_version.py +++ b/google/cloud/aiplatform_v1/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.33.0" # {x-release-please-version} +__version__ = "1.33.1" # {x-release-please-version} diff --git a/google/cloud/aiplatform_v1beta1/gapic_version.py b/google/cloud/aiplatform_v1beta1/gapic_version.py index 11cb6e1c3a..c2a948a447 100644 --- a/google/cloud/aiplatform_v1beta1/gapic_version.py +++ b/google/cloud/aiplatform_v1beta1/gapic_version.py @@ -13,4 +13,4 @@ # See the License for the specific language governing permissions and # limitations under the License. # -__version__ = "1.33.0" # {x-release-please-version} +__version__ = "1.33.1" # {x-release-please-version} diff --git a/samples/generated_samples/snippet_metadata_google.cloud.aiplatform.v1.json b/samples/generated_samples/snippet_metadata_google.cloud.aiplatform.v1.json index 1895c381f9..8eb6336968 100644 --- a/samples/generated_samples/snippet_metadata_google.cloud.aiplatform.v1.json +++ b/samples/generated_samples/snippet_metadata_google.cloud.aiplatform.v1.json @@ -8,7 +8,7 @@ ], "language": "PYTHON", "name": "google-cloud-aiplatform", - "version": "1.33.0" + "version": "1.33.1" }, "snippets": [ { diff --git a/samples/generated_samples/snippet_metadata_google.cloud.aiplatform.v1beta1.json b/samples/generated_samples/snippet_metadata_google.cloud.aiplatform.v1beta1.json index d20d9babfa..5e3844840f 100644 --- a/samples/generated_samples/snippet_metadata_google.cloud.aiplatform.v1beta1.json +++ b/samples/generated_samples/snippet_metadata_google.cloud.aiplatform.v1beta1.json @@ -8,7 +8,7 @@ ], "language": "PYTHON", "name": "google-cloud-aiplatform", - "version": "1.33.0" + "version": "1.33.1" }, "snippets": [ { diff --git a/tests/system/aiplatform/test_model_evaluation.py b/tests/system/aiplatform/test_model_evaluation.py index 565edaefa6..e897b47787 100644 --- a/tests/system/aiplatform/test_model_evaluation.py +++ b/tests/system/aiplatform/test_model_evaluation.py @@ -20,6 +20,7 @@ import pytest +from google import auth from google.cloud import storage from google.cloud import aiplatform @@ -83,7 +84,13 @@ def staging_bucket(self, storage_client): yield bucket def test_model_evaluate_custom_tabular_model(self, staging_bucket, shared_state): - aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) + aiplatform.init( + project=_TEST_PROJECT, + location=_TEST_LOCATION, + credentials=auth.default( + scopes=["https://www.googleapis.com/auth/cloud-platform"] + ), + ) custom_model = aiplatform.Model( model_name=_TEST_PERMANENT_CUSTOM_MODEL_CLASSIFICATION_RESOURCE_NAME diff --git a/tests/system/vertexai/test_pytorch.py b/tests/system/vertexai/test_pytorch.py index 0a4b2f3c7f..e0064b9ce0 100644 --- a/tests/system/vertexai/test_pytorch.py +++ b/tests/system/vertexai/test_pytorch.py @@ -18,12 +18,10 @@ import os from unittest import mock +from google.cloud import aiplatform import vertexai from tests.system.aiplatform import e2e_base from vertexai.preview._workflow.executor import training -from vertexai.preview._workflow.serialization_engine import ( - any_serializer, -) from vertexai.preview._workflow.serialization_engine import ( serializers, ) @@ -143,15 +141,27 @@ def predict(self, X): model.train(train_loader, num_epochs=100, lr=0.05) # Assert the right serializer is being used - serializer = any_serializer.AnySerializer() - assert ( - serializer._get_predefined_serializer(model.__class__.__mro__[2]) - is serializers.TorchModelSerializer + remote_job = aiplatform.CustomJob.list( + filter=f'display_name="{model.train.vertex.remote_config.display_name}"' + )[0] + base_path = remote_job.job_spec.base_output_directory.output_uri_prefix + + input_estimator_metadata = serializers._get_metadata( + os.path.join(base_path, "input/input_estimator") ) - assert ( - serializer._get_predefined_serializer(train_loader.__class__) - is serializers.TorchDataLoaderSerializer + assert input_estimator_metadata["serializer"] == "TorchModelSerializer" + + output_estimator_metadata = serializers._get_metadata( + os.path.join(base_path, "output/output_estimator") + ) + assert output_estimator_metadata["serializer"] == "TorchModelSerializer" + + train_loader_metadata = serializers._get_metadata( + os.path.join(base_path, "input/dataloader") ) + assert train_loader_metadata["serializer"] == "TorchDataLoaderSerializer" + + shared_state["resources"] = [remote_job] # Remote prediction on Torch custom model model.predict.vertex.remote_config.display_name = self._make_display_name( @@ -161,7 +171,7 @@ def predict(self, X): # Register trained model registered_model = vertexai.preview.register(model) - shared_state["resources"] = [registered_model] + shared_state["resources"].append(registered_model) # Load the registered model pulled_model = vertexai.preview.from_pretrained( @@ -175,12 +185,24 @@ def predict(self, X): pulled_model.train(retrain_loader, num_epochs=100, lr=0.05) # Assert the right serializer is being used - serializer = any_serializer.AnySerializer() - assert ( - serializer._get_predefined_serializer(pulled_model.__class__.__mro__[2]) - is serializers.TorchModelSerializer + remote_job = aiplatform.CustomJob.list( + filter=f'display_name="{pulled_model.train.vertex.remote_config.display_name}"' + )[0] + base_path = remote_job.job_spec.base_output_directory.output_uri_prefix + + input_estimator_metadata = serializers._get_metadata( + os.path.join(base_path, "input/input_estimator") + ) + assert input_estimator_metadata["serializer"] == "TorchModelSerializer" + + output_estimator_metadata = serializers._get_metadata( + os.path.join(base_path, "output/output_estimator") ) - assert ( - serializer._get_predefined_serializer(retrain_loader.__class__) - is serializers.TorchDataLoaderSerializer + assert output_estimator_metadata["serializer"] == "TorchModelSerializer" + + train_loader_metadata = serializers._get_metadata( + os.path.join(base_path, "input/dataloader") ) + assert train_loader_metadata["serializer"] == "TorchDataLoaderSerializer" + + shared_state["resources"].append(remote_job) diff --git a/tests/system/vertexai/test_sklearn.py b/tests/system/vertexai/test_sklearn.py index 98bd82dcd1..1f7d58bfbe 100644 --- a/tests/system/vertexai/test_sklearn.py +++ b/tests/system/vertexai/test_sklearn.py @@ -18,12 +18,10 @@ import os from unittest import mock +from google.cloud import aiplatform import vertexai from tests.system.aiplatform import e2e_base from vertexai.preview._workflow.executor import training -from vertexai.preview._workflow.serialization_engine import ( - any_serializer, -) from vertexai.preview._workflow.serialization_engine import ( serializers, ) @@ -90,11 +88,22 @@ def test_remote_execution_sklearn(self, shared_state): X_train = transformer.fit_transform(X_train) # Assert the right serializer is being used - serializer = any_serializer.AnySerializer() - assert ( - serializer._get_predefined_serializer(transformer.__class__.__mro__[-2]) - is serializers.SklearnEstimatorSerializer + remote_job = aiplatform.CustomJob.list( + filter=f'display_name="{transformer.fit_transform.vertex.remote_config.display_name}"' + )[0] + base_path = remote_job.job_spec.base_output_directory.output_uri_prefix + + input_estimator_metadata = serializers._get_metadata( + os.path.join(base_path, "input/input_estimator") + ) + assert input_estimator_metadata["serializer"] == "SklearnEstimatorSerializer" + + output_estimator_metadata = serializers._get_metadata( + os.path.join(base_path, "output/output_estimator") ) + assert output_estimator_metadata["serializer"] == "SklearnEstimatorSerializer" + + shared_state["resources"] = [remote_job] # Remote transform on test dataset transformer.transform.vertex.set_config( @@ -103,11 +112,22 @@ def test_remote_execution_sklearn(self, shared_state): X_test = transformer.transform(X_test) # Assert the right serializer is being used - serializer = any_serializer.AnySerializer() - assert ( - serializer._get_predefined_serializer(transformer.__class__.__mro__[-2]) - is serializers.SklearnEstimatorSerializer + remote_job = aiplatform.CustomJob.list( + filter=f'display_name="{transformer.transform.vertex.remote_config.display_name}"' + )[0] + base_path = remote_job.job_spec.base_output_directory.output_uri_prefix + + input_estimator_metadata = serializers._get_metadata( + os.path.join(base_path, "input/input_estimator") + ) + assert input_estimator_metadata["serializer"] == "SklearnEstimatorSerializer" + + output_estimator_metadata = serializers._get_metadata( + os.path.join(base_path, "output/output_estimator") ) + assert output_estimator_metadata["serializer"] == "SklearnEstimatorSerializer" + + shared_state["resources"].append(remote_job) # Local transform on retrain data vertexai.preview.init(remote=False) @@ -126,11 +146,22 @@ def test_remote_execution_sklearn(self, shared_state): model.fit(X_train, y_train) # Assert the right serializer is being used - serializer = any_serializer.AnySerializer() - assert ( - serializer._get_predefined_serializer(model.__class__.__mro__[-2]) - is serializers.SklearnEstimatorSerializer + remote_job = aiplatform.CustomJob.list( + filter=f'display_name="{model.fit.vertex.remote_config.display_name}"' + )[0] + base_path = remote_job.job_spec.base_output_directory.output_uri_prefix + + input_estimator_metadata = serializers._get_metadata( + os.path.join(base_path, "input/input_estimator") ) + assert input_estimator_metadata["serializer"] == "SklearnEstimatorSerializer" + + output_estimator_metadata = serializers._get_metadata( + os.path.join(base_path, "output/output_estimator") + ) + assert output_estimator_metadata["serializer"] == "SklearnEstimatorSerializer" + + shared_state["resources"].append(remote_job) # Remote prediction on sklearn model.predict.vertex.remote_config.display_name = self._make_display_name( @@ -140,7 +171,7 @@ def test_remote_execution_sklearn(self, shared_state): # Register trained model registered_model = vertexai.preview.register(model) - shared_state["resources"] = [registered_model] + shared_state["resources"].append(registered_model) # Load the registered model pulled_model = vertexai.preview.from_pretrained( @@ -151,8 +182,19 @@ def test_remote_execution_sklearn(self, shared_state): pulled_model.fit(X_retrain_df, y_retrain_df) # Assert the right serializer is being used - serializer = any_serializer.AnySerializer() - assert ( - serializer._get_predefined_serializer(pulled_model.__class__.__mro__[-2]) - is serializers.SklearnEstimatorSerializer + remote_job = aiplatform.CustomJob.list( + filter=f'display_name="{pulled_model.fit.vertex.remote_config.display_name}"' + )[0] + base_path = remote_job.job_spec.base_output_directory.output_uri_prefix + + input_estimator_metadata = serializers._get_metadata( + os.path.join(base_path, "input/input_estimator") + ) + assert input_estimator_metadata["serializer"] == "SklearnEstimatorSerializer" + + output_estimator_metadata = serializers._get_metadata( + os.path.join(base_path, "output/output_estimator") ) + assert output_estimator_metadata["serializer"] == "SklearnEstimatorSerializer" + + shared_state["resources"].append(remote_job) diff --git a/tests/system/vertexai/test_tensorflow.py b/tests/system/vertexai/test_tensorflow.py index ea7e1fa1c8..f42cce89dd 100644 --- a/tests/system/vertexai/test_tensorflow.py +++ b/tests/system/vertexai/test_tensorflow.py @@ -18,12 +18,10 @@ import os from unittest import mock +from google.cloud import aiplatform import vertexai from tests.system.aiplatform import e2e_base from vertexai.preview._workflow.executor import training -from vertexai.preview._workflow.serialization_engine import ( - any_serializer, -) from vertexai.preview._workflow.serialization_engine import ( serializers, ) @@ -108,15 +106,25 @@ def test_remote_execution_keras(self, shared_state): model.fit(tf_train_dataset, epochs=10) # Assert the right serializer is being used - serializer = any_serializer.AnySerializer() - assert ( - serializer._get_predefined_serializer(model.__class__.__mro__[3]) - is serializers.KerasModelSerializer + remote_job = aiplatform.CustomJob.list( + filter=f'display_name="{model.fit.vertex.remote_config.display_name}"' + )[0] + base_path = remote_job.job_spec.base_output_directory.output_uri_prefix + + input_estimator_metadata = serializers._get_metadata( + os.path.join(base_path, "input/input_estimator") ) - assert ( - serializer._get_predefined_serializer(tf_train_dataset.__class__.__mro__[2]) - is serializers.TFDatasetSerializer + assert input_estimator_metadata["serializer"] == "KerasModelSerializer" + + output_estimator_metadata = serializers._get_metadata( + os.path.join(base_path, "output/output_estimator") ) + assert output_estimator_metadata["serializer"] == "KerasModelSerializer" + + train_x_metadata = serializers._get_metadata(os.path.join(base_path, "input/x")) + assert train_x_metadata["serializer"] == "TFDatasetSerializer" + + shared_state["resources"] = [remote_job] # Remote prediction on keras model.predict.vertex.remote_config.display_name = self._make_display_name( @@ -126,7 +134,7 @@ def test_remote_execution_keras(self, shared_state): # Register trained model registered_model = vertexai.preview.register(model) - shared_state["resources"] = [registered_model] + shared_state["resources"].append(registered_model) # Load the registered model pulled_model = vertexai.preview.from_pretrained( @@ -141,14 +149,22 @@ def test_remote_execution_keras(self, shared_state): pulled_model.fit(tf_retrain_dataset, epochs=10) # Assert the right serializer is being used - serializer = any_serializer.AnySerializer() - assert ( - serializer._get_predefined_serializer(pulled_model.__class__.__mro__[3]) - is serializers.KerasModelSerializer + remote_job = aiplatform.CustomJob.list( + filter=f'display_name="{pulled_model.fit.vertex.remote_config.display_name}"' + )[0] + base_path = remote_job.job_spec.base_output_directory.output_uri_prefix + + input_estimator_metadata = serializers._get_metadata( + os.path.join(base_path, "input/input_estimator") ) - assert ( - serializer._get_predefined_serializer( - tf_retrain_dataset.__class__.__mro__[2] - ) - is serializers.TFDatasetSerializer + assert input_estimator_metadata["serializer"] == "KerasModelSerializer" + + output_estimator_metadata = serializers._get_metadata( + os.path.join(base_path, "output/output_estimator") ) + assert output_estimator_metadata["serializer"] == "KerasModelSerializer" + + train_x_metadata = serializers._get_metadata(os.path.join(base_path, "input/x")) + assert train_x_metadata["serializer"] == "TFDatasetSerializer" + + shared_state["resources"].append(remote_job) diff --git a/vertexai/preview/_workflow/driver/__init__.py b/vertexai/preview/_workflow/driver/__init__.py index f17ed3c97a..cdc9030ffa 100644 --- a/vertexai/preview/_workflow/driver/__init__.py +++ b/vertexai/preview/_workflow/driver/__init__.py @@ -197,18 +197,33 @@ def _unwrapper(instance: Any) -> Callable[..., Any]: config_map = dict() - for attr_name, attr_value in inspect.getmembers(instance): - if isinstance(attr_value, VertexRemoteFunctor): - config_map[attr_name] = ( - attr_value.vertex, - attr_value._remote_executor, - attr_value._remote_executor_kwargs, - ) - setattr(instance, attr_name, attr_value._method) - if not wrapped_in_place: + for ( + attr_name, + attr_value, + remote_executor, + remote_executor_kwargs, + ) in _supported_member_iter(instance): + if isinstance(attr_value, VertexRemoteFunctor): + config_map[attr_name] = ( + attr_value.vertex, + remote_executor, + remote_executor_kwargs, + ) + setattr(instance, attr_name, attr_value._method) + instance.__class__ = super_class + else: + for attr_name, attr_value in inspect.getmembers(instance): + if isinstance(attr_value, VertexRemoteFunctor): + config_map[attr_name] = ( + attr_value.vertex, + attr_value._remote_executor, + attr_value._remote_executor_kwargs, + ) + setattr(instance, attr_name, attr_value._method) + return functools.partial( _rewrapper, wrapped_class=current_class, config_map=config_map ) diff --git a/vertexai/preview/_workflow/serialization_engine/any_serializer.py b/vertexai/preview/_workflow/serialization_engine/any_serializer.py index 75797f9d23..2df2abc8c1 100644 --- a/vertexai/preview/_workflow/serialization_engine/any_serializer.py +++ b/vertexai/preview/_workflow/serialization_engine/any_serializer.py @@ -18,11 +18,9 @@ """Defines the Serializer classes.""" import json import os -import tempfile from typing import Any, Dict, Union, List, TypeVar, Type from google.cloud.aiplatform import base -from google.cloud.aiplatform.utils import gcs_utils from vertexai.preview._workflow.serialization_engine import ( serializers, serializers_base, @@ -243,19 +241,7 @@ def serialize(self, to_serialize: T, gcs_path: str, **kwargs) -> Dict[str, Any]: def deserialize(self, serialized_gcs_path: str, **kwargs) -> T: """Routes the corresponding Serializer based on the metadata.""" - metadata_path = serializers.get_metadata_path_from_file_gcs_uri( - serialized_gcs_path - ) - - if metadata_path.startswith("gs://"): - with tempfile.NamedTemporaryFile() as temp_file: - gcs_utils.download_file_from_gcs(metadata_path, temp_file.name) - with open(temp_file.name, mode="rb") as f: - metadata = json.load(f) - else: - with open(metadata_path, mode="rb") as f: - metadata = json.load(f) - + metadata = serializers._get_metadata(serialized_gcs_path) _LOGGER.debug( "deserializing from %s, metadata is %s", serialized_gcs_path, metadata ) diff --git a/vertexai/preview/_workflow/serialization_engine/serializers.py b/vertexai/preview/_workflow/serialization_engine/serializers.py index 223407f880..cb6c1b30cd 100644 --- a/vertexai/preview/_workflow/serialization_engine/serializers.py +++ b/vertexai/preview/_workflow/serialization_engine/serializers.py @@ -25,7 +25,7 @@ import pickle import shutil import tempfile -from typing import Any, Optional, Union +from typing import Any, Dict, Optional, Union import uuid from google.cloud.aiplatform.utils import gcs_utils @@ -158,6 +158,20 @@ def get_metadata_path_from_file_gcs_uri(gcs_uri: str) -> str: ) +def _get_metadata(gcs_uri: str) -> Dict[str, Any]: + metadata_file = get_metadata_path_from_file_gcs_uri(gcs_uri) + if metadata_file.startswith("gs://"): + with tempfile.NamedTemporaryFile() as temp_file: + gcs_utils.download_file_from_gcs(metadata_file, temp_file.name) + with open(temp_file.name, mode="rb") as f: + metadata = json.load(f) + else: + with open(metadata_file, "rb") as f: + metadata = json.load(f) + + return metadata + + def _is_valid_gcs_path(path: str) -> bool: """checks if a path is a valid gcs path. @@ -245,8 +259,7 @@ def serialize( save_format=save_format, ) gcs_utils.upload_to_gcs( - temp_file_or_dir.name if is_file else temp_file_or_dir, - gcs_path + temp_file_or_dir.name if is_file else temp_file_or_dir, gcs_path ) else: to_serialize.save(gcs_path, save_format=save_format) @@ -269,14 +282,8 @@ def deserialize(self, serialized_gcs_path: str, **kwargs) -> KerasModel: del kwargs if not _is_valid_gcs_path(serialized_gcs_path): raise ValueError(f"Invalid gcs path: {serialized_gcs_path}") - metadata_file = get_metadata_path_from_file_gcs_uri(serialized_gcs_path) - if metadata_file.startswith("gs://"): - with tempfile.NamedTemporaryFile(suffix=".json") as temp_file: - gcs_utils.download_file_from_gcs(metadata_file, temp_file.name) - metadata = json.load(temp_file) - else: - with open(metadata_file, "r") as f: - metadata = json.load(f) + + metadata = _get_metadata(serialized_gcs_path) # For backward compatibility, if the metadata doesn't contain # save_format, we assume the model was saved as saved_model format. save_format = metadata.get("save_format", "tf")