diff --git a/CHANGELOG.md b/CHANGELOG.md index 04e51a15bd..340873320d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,30 @@ # Changelog +### [1.13.1](https://github.com/googleapis/python-aiplatform/compare/v1.13.0...v1.13.1) (2022-05-26) + + +### Features + +* add batch_size kwarg for batch prediction jobs ([#1194](https://github.com/googleapis/python-aiplatform/issues/1194)) ([50bdb01](https://github.com/googleapis/python-aiplatform/commit/50bdb01504740ed31de788d8a160f3e2be7f55df)) +* add update endpoint ([#1162](https://github.com/googleapis/python-aiplatform/issues/1162)) ([0ecfe1e](https://github.com/googleapis/python-aiplatform/commit/0ecfe1e7ab8687c13cb4267985e8b6ebc7bd2534)) +* support autoscaling metrics when deploying models ([#1197](https://github.com/googleapis/python-aiplatform/issues/1197)) ([095717c](https://github.com/googleapis/python-aiplatform/commit/095717c8b77dc5d66e677413a437ea6ed92e0b1a)) + + +### Bug Fixes + +* check in service proto file ([#1174](https://github.com/googleapis/python-aiplatform/issues/1174)) ([5fdf151](https://github.com/googleapis/python-aiplatform/commit/5fdf151ee0d0a630c07a75dc8f19906e7ad1aa8a)) +* regenerate pb2 files using grpcio-tools ([#1394](https://github.com/googleapis/python-aiplatform/issues/1394)) ([406c868](https://github.com/googleapis/python-aiplatform/commit/406c868344280d424f4191c98bcbbdeaf947b2d1)) + + +### Documentation + +* update aiplatform SDK arrangement for Sphinx ([#1163](https://github.com/googleapis/python-aiplatform/issues/1163)) ([e9510ea](https://github.com/googleapis/python-aiplatform/commit/e9510ea6344a296e0c93ddf32280cf4c010ee4f1)) + + +### Miscellaneous Chores + +* release 1.13.1 ([#1395](https://github.com/googleapis/python-aiplatform/issues/1395)) ([df78407](https://github.com/googleapis/python-aiplatform/commit/df78407b2f14c95c9e84b4b1375a8de5bc9c7bb5)) + ## [1.13.0](https://github.com/googleapis/python-aiplatform/compare/v1.12.1...v1.13.0) (2022-05-09) diff --git a/google/cloud/aiplatform/compat/__init__.py b/google/cloud/aiplatform/compat/__init__.py index dd141b6653..6aea51d133 100644 --- a/google/cloud/aiplatform/compat/__init__.py +++ b/google/cloud/aiplatform/compat/__init__.py @@ -40,6 +40,10 @@ ) services.metadata_service_client = services.metadata_service_client_v1beta1 services.tensorboard_service_client = services.tensorboard_service_client_v1beta1 + services.index_service_client = services.index_service_client_v1beta1 + services.index_endpoint_service_client = ( + services.index_endpoint_service_client_v1beta1 + ) types.accelerator_type = types.accelerator_type_v1beta1 types.annotation = types.annotation_v1beta1 @@ -71,6 +75,8 @@ types.featurestore_online_service = types.featurestore_online_service_v1beta1 types.featurestore_service = types.featurestore_service_v1beta1 types.hyperparameter_tuning_job = types.hyperparameter_tuning_job_v1beta1 + types.index = types.index_v1beta1 + types.index_endpoint = types.index_endpoint_v1beta1 types.io = types.io_v1beta1 types.job_service = types.job_service_v1beta1 types.job_state = types.job_state_v1beta1 @@ -79,8 +85,8 @@ types.matching_engine_deployed_index_ref = ( types.matching_engine_deployed_index_ref_v1beta1 ) - types.matching_engine_index = types.matching_engine_index_v1beta1 - types.matching_engine_index_endpoint = types.matching_engine_index_endpoint_v1beta1 + types.matching_engine_index = types.index_v1beta1 + types.matching_engine_index_endpoint = types.index_endpoint_v1beta1 types.metadata_service = types.metadata_service_v1beta1 types.metadata_store = types.metadata_store_v1beta1 types.model = types.model_v1beta1 @@ -88,6 +94,7 @@ types.model_evaluation_slice = types.model_evaluation_slice_v1beta1 types.model_service = types.model_service_v1beta1 types.operation = types.operation_v1beta1 + types.pipeline_job = types.pipeline_job_v1beta1 types.pipeline_service = types.pipeline_service_v1beta1 types.pipeline_state = types.pipeline_state_v1beta1 types.prediction_service = types.prediction_service_v1beta1 @@ -102,6 +109,7 @@ types.tensorboard_service = types.tensorboard_service_v1beta1 types.tensorboard_time_series = types.tensorboard_time_series_v1beta1 types.training_pipeline = types.training_pipeline_v1beta1 + types.types = types.types_v1beta1 if DEFAULT_VERSION == V1: @@ -117,6 +125,8 @@ services.prediction_service_client = services.prediction_service_client_v1 services.specialist_pool_service_client = services.specialist_pool_service_client_v1 services.tensorboard_service_client = services.tensorboard_service_client_v1 + services.index_service_client = services.index_service_client_v1 + services.index_endpoint_service_client = services.index_endpoint_service_client_v1 types.accelerator_type = types.accelerator_type_v1 types.annotation = types.annotation_v1 @@ -147,6 +157,8 @@ types.featurestore_online_service = types.featurestore_online_service_v1 types.featurestore_service = types.featurestore_service_v1 types.hyperparameter_tuning_job = types.hyperparameter_tuning_job_v1 + types.index = types.index_v1 + types.index_endpoint = types.index_endpoint_v1 types.io = types.io_v1 types.job_service = types.job_service_v1 types.job_state = types.job_state_v1 @@ -155,8 +167,8 @@ types.matching_engine_deployed_index_ref = ( types.matching_engine_deployed_index_ref_v1 ) - types.matching_engine_index = types.matching_engine_index_v1 - types.matching_engine_index_endpoint = types.matching_engine_index_endpoint_v1 + types.matching_engine_index = types.index_v1 + types.matching_engine_index_endpoint = types.index_endpoint_v1 types.metadata_service = types.metadata_service_v1 types.metadata_store = types.metadata_store_v1 types.model = types.model_v1 @@ -164,6 +176,7 @@ types.model_evaluation_slice = types.model_evaluation_slice_v1 types.model_service = types.model_service_v1 types.operation = types.operation_v1 + types.pipeline_job = types.pipeline_job_v1 types.pipeline_service = types.pipeline_service_v1 types.pipeline_state = types.pipeline_state_v1 types.prediction_service = types.prediction_service_v1 @@ -178,6 +191,7 @@ types.tensorboard_service = types.tensorboard_service_v1 types.tensorboard_time_series = types.tensorboard_time_series_v1 types.training_pipeline = types.training_pipeline_v1 + types.types = types.types_v1 __all__ = ( DEFAULT_VERSION, diff --git a/google/cloud/aiplatform/compat/types/__init__.py b/google/cloud/aiplatform/compat/types/__init__.py index fc8e7c0b30..14ff93f011 100644 --- a/google/cloud/aiplatform/compat/types/__init__.py +++ b/google/cloud/aiplatform/compat/types/__init__.py @@ -46,8 +46,8 @@ featurestore_monitoring as featurestore_monitoring_v1beta1, featurestore_online_service as featurestore_online_service_v1beta1, featurestore_service as featurestore_service_v1beta1, - index as matching_engine_index_v1beta1, - index_endpoint as matching_engine_index_endpoint_v1beta1, + index as index_v1beta1, + index_endpoint as index_endpoint_v1beta1, hyperparameter_tuning_job as hyperparameter_tuning_job_v1beta1, io as io_v1beta1, job_service as job_service_v1beta1, @@ -75,6 +75,7 @@ tensorboard_service as tensorboard_service_v1beta1, tensorboard_time_series as tensorboard_time_series_v1beta1, training_pipeline as training_pipeline_v1beta1, + types as types_v1beta1, ) from google.cloud.aiplatform_v1.types import ( accelerator_type as accelerator_type_v1, @@ -107,8 +108,8 @@ featurestore_online_service as featurestore_online_service_v1, featurestore_service as featurestore_service_v1, hyperparameter_tuning_job as hyperparameter_tuning_job_v1, - index as matching_engine_index_v1, - index_endpoint as matching_engine_index_endpoint_v1, + index as index_v1, + index_endpoint as index_endpoint_v1, io as io_v1, job_service as job_service_v1, job_state as job_state_v1, @@ -135,6 +136,7 @@ tensorboard_service as tensorboard_service_v1, tensorboard_time_series as tensorboard_time_series_v1, training_pipeline as training_pipeline_v1, + types as types_v1, ) __all__ = ( @@ -174,8 +176,8 @@ machine_resources_v1, manual_batch_tuning_parameters_v1, matching_engine_deployed_index_ref_v1, - matching_engine_index_v1, - matching_engine_index_endpoint_v1, + index_v1, + index_endpoint_v1, metadata_service_v1, metadata_store_v1, model_v1, @@ -196,6 +198,7 @@ tensorboard_service_v1, tensorboard_time_series_v1, training_pipeline_v1, + types_v1, # v1beta1 accelerator_type_v1beta1, annotation_v1beta1, @@ -233,8 +236,8 @@ machine_resources_v1beta1, manual_batch_tuning_parameters_v1beta1, matching_engine_deployed_index_ref_v1beta1, - matching_engine_index_v1beta1, - matching_engine_index_endpoint_v1beta1, + index_v1beta1, + index_endpoint_v1beta1, metadata_service_v1beta1, metadata_store_v1beta1, model_v1beta1, @@ -255,4 +258,5 @@ tensorboard_service_v1beta1, tensorboard_time_series_v1beta1, training_pipeline_v1beta1, + types_v1beta1, ) diff --git a/google/cloud/aiplatform/datasets/time_series_dataset.py b/google/cloud/aiplatform/datasets/time_series_dataset.py index ec5546f12a..6bde6be7a5 100644 --- a/google/cloud/aiplatform/datasets/time_series_dataset.py +++ b/google/cloud/aiplatform/datasets/time_series_dataset.py @@ -46,6 +46,7 @@ def create( labels: Optional[Dict[str, str]] = None, encryption_spec_key_name: Optional[str] = None, sync: bool = True, + create_request_timeout: Optional[float] = None, ) -> "TimeSeriesDataset": """Creates a new time series dataset. @@ -102,6 +103,8 @@ def create( Whether to execute this method synchronously. If False, this method will be executed in concurrent Future and any downstream object will be immediately returned and synced when the Future has completed. + create_request_timeout (float): + Optional. The timeout for the create request in seconds. Returns: time_series_dataset (TimeSeriesDataset): @@ -141,6 +144,7 @@ def create( encryption_spec_key_name=encryption_spec_key_name ), sync=sync, + create_request_timeout=create_request_timeout, ) def import_data(self): diff --git a/google/cloud/aiplatform/matching_engine/_protos/match_service_pb2.py b/google/cloud/aiplatform/matching_engine/_protos/match_service_pb2.py index 8d2adac367..6b3ab988b2 100644 --- a/google/cloud/aiplatform/matching_engine/_protos/match_service_pb2.py +++ b/google/cloud/aiplatform/matching_engine/_protos/match_service_pb2.py @@ -13,11 +13,12 @@ # 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. -# + # Generated by the protocol buffer compiler. DO NOT EDIT! -# source: match_service.proto +# source: google/cloud/aiplatform/matching_engine/_protos/match_service.proto """Generated protocol buffer code.""" from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pool as _descriptor_pool from google.protobuf import message as _message from google.protobuf import reflection as _reflection from google.protobuf import symbol_database as _symbol_database @@ -30,629 +31,29 @@ from google.rpc import status_pb2 as google_dot_rpc_dot_status__pb2 -DESCRIPTOR = _descriptor.FileDescriptor( - name="match_service.proto", - package="google.cloud.aiplatform.container.v1beta1", - syntax="proto3", - serialized_options=None, - create_key=_descriptor._internal_create_key, - serialized_pb=b'\n\x13match_service.proto\x12)google.cloud.aiplatform.container.v1beta1\x1a\x17google/rpc/status.proto"\x97\x02\n\x0cMatchRequest\x12\x19\n\x11\x64\x65ployed_index_id\x18\x01 \x01(\t\x12\x11\n\tfloat_val\x18\x02 \x03(\x02\x12\x15\n\rnum_neighbors\x18\x03 \x01(\x05\x12G\n\trestricts\x18\x04 \x03(\x0b\x32\x34.google.cloud.aiplatform.container.v1beta1.Namespace\x12,\n$per_crowding_attribute_num_neighbors\x18\x05 \x01(\x05\x12\x1c\n\x14\x61pprox_num_neighbors\x18\x06 \x01(\x05\x12-\n%leaf_nodes_to_search_percent_override\x18\x07 \x01(\x05"\x8e\x01\n\rMatchResponse\x12S\n\x08neighbor\x18\x01 \x03(\x0b\x32\x41.google.cloud.aiplatform.container.v1beta1.MatchResponse.Neighbor\x1a(\n\x08Neighbor\x12\n\n\x02id\x18\x01 \x01(\t\x12\x10\n\x08\x64istance\x18\x02 \x01(\x01"\x9f\x02\n\x11\x42\x61tchMatchRequest\x12h\n\x08requests\x18\x01 \x03(\x0b\x32V.google.cloud.aiplatform.container.v1beta1.BatchMatchRequest.BatchMatchRequestPerIndex\x1a\x9f\x01\n\x19\x42\x61tchMatchRequestPerIndex\x12\x19\n\x11\x64\x65ployed_index_id\x18\x01 \x01(\t\x12I\n\x08requests\x18\x02 \x03(\x0b\x32\x37.google.cloud.aiplatform.container.v1beta1.MatchRequest\x12\x1c\n\x14low_level_batch_size\x18\x03 \x01(\x05"\xac\x02\n\x12\x42\x61tchMatchResponse\x12k\n\tresponses\x18\x01 \x03(\x0b\x32X.google.cloud.aiplatform.container.v1beta1.BatchMatchResponse.BatchMatchResponsePerIndex\x1a\xa8\x01\n\x1a\x42\x61tchMatchResponsePerIndex\x12\x19\n\x11\x64\x65ployed_index_id\x18\x01 \x01(\t\x12K\n\tresponses\x18\x02 \x03(\x0b\x32\x38.google.cloud.aiplatform.container.v1beta1.MatchResponse\x12"\n\x06status\x18\x03 \x01(\x0b\x32\x12.google.rpc.Status"D\n\tNamespace\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x14\n\x0c\x61llow_tokens\x18\x02 \x03(\t\x12\x13\n\x0b\x64\x65ny_tokens\x18\x03 \x03(\t2\x9a\x02\n\x0cMatchService\x12|\n\x05Match\x12\x37.google.cloud.aiplatform.container.v1beta1.MatchRequest\x1a\x38.google.cloud.aiplatform.container.v1beta1.MatchResponse"\x00\x12\x8b\x01\n\nBatchMatch\x12<.google.cloud.aiplatform.container.v1beta1.BatchMatchRequest\x1a=.google.cloud.aiplatform.container.v1beta1.BatchMatchResponse"\x00\x62\x06proto3', - dependencies=[ - google_dot_rpc_dot_status__pb2.DESCRIPTOR, - ], +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile( + b'\nCgoogle/cloud/aiplatform/matching_engine/_protos/match_service.proto\x12)google.cloud.aiplatform.container.v1beta1\x1a\x17google/rpc/status.proto"\x97\x02\n\x0cMatchRequest\x12\x19\n\x11\x64\x65ployed_index_id\x18\x01 \x01(\t\x12\x11\n\tfloat_val\x18\x02 \x03(\x02\x12\x15\n\rnum_neighbors\x18\x03 \x01(\x05\x12G\n\trestricts\x18\x04 \x03(\x0b\x32\x34.google.cloud.aiplatform.container.v1beta1.Namespace\x12,\n$per_crowding_attribute_num_neighbors\x18\x05 \x01(\x05\x12\x1c\n\x14\x61pprox_num_neighbors\x18\x06 \x01(\x05\x12-\n%leaf_nodes_to_search_percent_override\x18\x07 \x01(\x05"\x8e\x01\n\rMatchResponse\x12S\n\x08neighbor\x18\x01 \x03(\x0b\x32\x41.google.cloud.aiplatform.container.v1beta1.MatchResponse.Neighbor\x1a(\n\x08Neighbor\x12\n\n\x02id\x18\x01 \x01(\t\x12\x10\n\x08\x64istance\x18\x02 \x01(\x01"\x9f\x02\n\x11\x42\x61tchMatchRequest\x12h\n\x08requests\x18\x01 \x03(\x0b\x32V.google.cloud.aiplatform.container.v1beta1.BatchMatchRequest.BatchMatchRequestPerIndex\x1a\x9f\x01\n\x19\x42\x61tchMatchRequestPerIndex\x12\x19\n\x11\x64\x65ployed_index_id\x18\x01 \x01(\t\x12I\n\x08requests\x18\x02 \x03(\x0b\x32\x37.google.cloud.aiplatform.container.v1beta1.MatchRequest\x12\x1c\n\x14low_level_batch_size\x18\x03 \x01(\x05"\xac\x02\n\x12\x42\x61tchMatchResponse\x12k\n\tresponses\x18\x01 \x03(\x0b\x32X.google.cloud.aiplatform.container.v1beta1.BatchMatchResponse.BatchMatchResponsePerIndex\x1a\xa8\x01\n\x1a\x42\x61tchMatchResponsePerIndex\x12\x19\n\x11\x64\x65ployed_index_id\x18\x01 \x01(\t\x12K\n\tresponses\x18\x02 \x03(\x0b\x32\x38.google.cloud.aiplatform.container.v1beta1.MatchResponse\x12"\n\x06status\x18\x03 \x01(\x0b\x32\x12.google.rpc.Status"D\n\tNamespace\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x14\n\x0c\x61llow_tokens\x18\x02 \x03(\t\x12\x13\n\x0b\x64\x65ny_tokens\x18\x03 \x03(\t2\x9a\x02\n\x0cMatchService\x12|\n\x05Match\x12\x37.google.cloud.aiplatform.container.v1beta1.MatchRequest\x1a\x38.google.cloud.aiplatform.container.v1beta1.MatchResponse"\x00\x12\x8b\x01\n\nBatchMatch\x12<.google.cloud.aiplatform.container.v1beta1.BatchMatchRequest\x1a=.google.cloud.aiplatform.container.v1beta1.BatchMatchResponse"\x00\x62\x06proto3' ) -_MATCHREQUEST = _descriptor.Descriptor( - name="MatchRequest", - full_name="google.cloud.aiplatform.container.v1beta1.MatchRequest", - filename=None, - file=DESCRIPTOR, - containing_type=None, - create_key=_descriptor._internal_create_key, - fields=[ - _descriptor.FieldDescriptor( - name="deployed_index_id", - full_name="google.cloud.aiplatform.container.v1beta1.MatchRequest.deployed_index_id", - index=0, - number=1, - type=9, - cpp_type=9, - label=1, - has_default_value=False, - default_value=b"".decode("utf-8"), - message_type=None, - enum_type=None, - containing_type=None, - is_extension=False, - extension_scope=None, - serialized_options=None, - file=DESCRIPTOR, - create_key=_descriptor._internal_create_key, - ), - _descriptor.FieldDescriptor( - name="float_val", - full_name="google.cloud.aiplatform.container.v1beta1.MatchRequest.float_val", - index=1, - number=2, - type=2, - cpp_type=6, - label=3, - has_default_value=False, - default_value=[], - message_type=None, - enum_type=None, - containing_type=None, - is_extension=False, - extension_scope=None, - serialized_options=None, - file=DESCRIPTOR, - create_key=_descriptor._internal_create_key, - ), - _descriptor.FieldDescriptor( - name="num_neighbors", - full_name="google.cloud.aiplatform.container.v1beta1.MatchRequest.num_neighbors", - index=2, - number=3, - type=5, - cpp_type=1, - label=1, - has_default_value=False, - default_value=0, - message_type=None, - enum_type=None, - containing_type=None, - is_extension=False, - extension_scope=None, - serialized_options=None, - file=DESCRIPTOR, - create_key=_descriptor._internal_create_key, - ), - _descriptor.FieldDescriptor( - name="restricts", - full_name="google.cloud.aiplatform.container.v1beta1.MatchRequest.restricts", - index=3, - number=4, - type=11, - cpp_type=10, - label=3, - has_default_value=False, - default_value=[], - message_type=None, - enum_type=None, - containing_type=None, - is_extension=False, - extension_scope=None, - serialized_options=None, - file=DESCRIPTOR, - create_key=_descriptor._internal_create_key, - ), - _descriptor.FieldDescriptor( - name="per_crowding_attribute_num_neighbors", - full_name="google.cloud.aiplatform.container.v1beta1.MatchRequest.per_crowding_attribute_num_neighbors", - index=4, - number=5, - type=5, - cpp_type=1, - label=1, - has_default_value=False, - default_value=0, - message_type=None, - enum_type=None, - containing_type=None, - is_extension=False, - extension_scope=None, - serialized_options=None, - file=DESCRIPTOR, - create_key=_descriptor._internal_create_key, - ), - _descriptor.FieldDescriptor( - name="approx_num_neighbors", - full_name="google.cloud.aiplatform.container.v1beta1.MatchRequest.approx_num_neighbors", - index=5, - number=6, - type=5, - cpp_type=1, - label=1, - has_default_value=False, - default_value=0, - message_type=None, - enum_type=None, - containing_type=None, - is_extension=False, - extension_scope=None, - serialized_options=None, - file=DESCRIPTOR, - create_key=_descriptor._internal_create_key, - ), - _descriptor.FieldDescriptor( - name="leaf_nodes_to_search_percent_override", - full_name="google.cloud.aiplatform.container.v1beta1.MatchRequest.leaf_nodes_to_search_percent_override", - index=6, - number=7, - type=5, - cpp_type=1, - label=1, - has_default_value=False, - default_value=0, - message_type=None, - enum_type=None, - containing_type=None, - is_extension=False, - extension_scope=None, - serialized_options=None, - file=DESCRIPTOR, - create_key=_descriptor._internal_create_key, - ), - ], - extensions=[], - nested_types=[], - enum_types=[], - serialized_options=None, - is_extendable=False, - syntax="proto3", - extension_ranges=[], - oneofs=[], - serialized_start=92, - serialized_end=371, +_MATCHREQUEST = DESCRIPTOR.message_types_by_name["MatchRequest"] +_MATCHRESPONSE = DESCRIPTOR.message_types_by_name["MatchResponse"] +_MATCHRESPONSE_NEIGHBOR = _MATCHRESPONSE.nested_types_by_name["Neighbor"] +_BATCHMATCHREQUEST = DESCRIPTOR.message_types_by_name["BatchMatchRequest"] +_BATCHMATCHREQUEST_BATCHMATCHREQUESTPERINDEX = _BATCHMATCHREQUEST.nested_types_by_name[ + "BatchMatchRequestPerIndex" +] +_BATCHMATCHRESPONSE = DESCRIPTOR.message_types_by_name["BatchMatchResponse"] +_BATCHMATCHRESPONSE_BATCHMATCHRESPONSEPERINDEX = ( + _BATCHMATCHRESPONSE.nested_types_by_name["BatchMatchResponsePerIndex"] ) - - -_MATCHRESPONSE_NEIGHBOR = _descriptor.Descriptor( - name="Neighbor", - full_name="google.cloud.aiplatform.container.v1beta1.MatchResponse.Neighbor", - filename=None, - file=DESCRIPTOR, - containing_type=None, - create_key=_descriptor._internal_create_key, - fields=[ - _descriptor.FieldDescriptor( - name="id", - full_name="google.cloud.aiplatform.container.v1beta1.MatchResponse.Neighbor.id", - index=0, - number=1, - type=9, - cpp_type=9, - label=1, - has_default_value=False, - default_value=b"".decode("utf-8"), - message_type=None, - enum_type=None, - containing_type=None, - is_extension=False, - extension_scope=None, - serialized_options=None, - file=DESCRIPTOR, - create_key=_descriptor._internal_create_key, - ), - _descriptor.FieldDescriptor( - name="distance", - full_name="google.cloud.aiplatform.container.v1beta1.MatchResponse.Neighbor.distance", - index=1, - number=2, - type=1, - cpp_type=5, - label=1, - has_default_value=False, - default_value=float(0), - message_type=None, - enum_type=None, - containing_type=None, - is_extension=False, - extension_scope=None, - serialized_options=None, - file=DESCRIPTOR, - create_key=_descriptor._internal_create_key, - ), - ], - extensions=[], - nested_types=[], - enum_types=[], - serialized_options=None, - is_extendable=False, - syntax="proto3", - extension_ranges=[], - oneofs=[], - serialized_start=476, - serialized_end=516, -) - -_MATCHRESPONSE = _descriptor.Descriptor( - name="MatchResponse", - full_name="google.cloud.aiplatform.container.v1beta1.MatchResponse", - filename=None, - file=DESCRIPTOR, - containing_type=None, - create_key=_descriptor._internal_create_key, - fields=[ - _descriptor.FieldDescriptor( - name="neighbor", - full_name="google.cloud.aiplatform.container.v1beta1.MatchResponse.neighbor", - index=0, - number=1, - type=11, - cpp_type=10, - label=3, - has_default_value=False, - default_value=[], - message_type=None, - enum_type=None, - containing_type=None, - is_extension=False, - extension_scope=None, - serialized_options=None, - file=DESCRIPTOR, - create_key=_descriptor._internal_create_key, - ), - ], - extensions=[], - nested_types=[ - _MATCHRESPONSE_NEIGHBOR, - ], - enum_types=[], - serialized_options=None, - is_extendable=False, - syntax="proto3", - extension_ranges=[], - oneofs=[], - serialized_start=374, - serialized_end=516, -) - - -_BATCHMATCHREQUEST_BATCHMATCHREQUESTPERINDEX = _descriptor.Descriptor( - name="BatchMatchRequestPerIndex", - full_name="google.cloud.aiplatform.container.v1beta1.BatchMatchRequest.BatchMatchRequestPerIndex", - filename=None, - file=DESCRIPTOR, - containing_type=None, - create_key=_descriptor._internal_create_key, - fields=[ - _descriptor.FieldDescriptor( - name="deployed_index_id", - full_name="google.cloud.aiplatform.container.v1beta1.BatchMatchRequest.BatchMatchRequestPerIndex.deployed_index_id", - index=0, - number=1, - type=9, - cpp_type=9, - label=1, - has_default_value=False, - default_value=b"".decode("utf-8"), - message_type=None, - enum_type=None, - containing_type=None, - is_extension=False, - extension_scope=None, - serialized_options=None, - file=DESCRIPTOR, - create_key=_descriptor._internal_create_key, - ), - _descriptor.FieldDescriptor( - name="requests", - full_name="google.cloud.aiplatform.container.v1beta1.BatchMatchRequest.BatchMatchRequestPerIndex.requests", - index=1, - number=2, - type=11, - cpp_type=10, - label=3, - has_default_value=False, - default_value=[], - message_type=None, - enum_type=None, - containing_type=None, - is_extension=False, - extension_scope=None, - serialized_options=None, - file=DESCRIPTOR, - create_key=_descriptor._internal_create_key, - ), - _descriptor.FieldDescriptor( - name="low_level_batch_size", - full_name="google.cloud.aiplatform.container.v1beta1.BatchMatchRequest.BatchMatchRequestPerIndex.low_level_batch_size", - index=2, - number=3, - type=5, - cpp_type=1, - label=1, - has_default_value=False, - default_value=0, - message_type=None, - enum_type=None, - containing_type=None, - is_extension=False, - extension_scope=None, - serialized_options=None, - file=DESCRIPTOR, - create_key=_descriptor._internal_create_key, - ), - ], - extensions=[], - nested_types=[], - enum_types=[], - serialized_options=None, - is_extendable=False, - syntax="proto3", - extension_ranges=[], - oneofs=[], - serialized_start=647, - serialized_end=806, -) - -_BATCHMATCHREQUEST = _descriptor.Descriptor( - name="BatchMatchRequest", - full_name="google.cloud.aiplatform.container.v1beta1.BatchMatchRequest", - filename=None, - file=DESCRIPTOR, - containing_type=None, - create_key=_descriptor._internal_create_key, - fields=[ - _descriptor.FieldDescriptor( - name="requests", - full_name="google.cloud.aiplatform.container.v1beta1.BatchMatchRequest.requests", - index=0, - number=1, - type=11, - cpp_type=10, - label=3, - has_default_value=False, - default_value=[], - message_type=None, - enum_type=None, - containing_type=None, - is_extension=False, - extension_scope=None, - serialized_options=None, - file=DESCRIPTOR, - create_key=_descriptor._internal_create_key, - ), - ], - extensions=[], - nested_types=[ - _BATCHMATCHREQUEST_BATCHMATCHREQUESTPERINDEX, - ], - enum_types=[], - serialized_options=None, - is_extendable=False, - syntax="proto3", - extension_ranges=[], - oneofs=[], - serialized_start=519, - serialized_end=806, -) - - -_BATCHMATCHRESPONSE_BATCHMATCHRESPONSEPERINDEX = _descriptor.Descriptor( - name="BatchMatchResponsePerIndex", - full_name="google.cloud.aiplatform.container.v1beta1.BatchMatchResponse.BatchMatchResponsePerIndex", - filename=None, - file=DESCRIPTOR, - containing_type=None, - create_key=_descriptor._internal_create_key, - fields=[ - _descriptor.FieldDescriptor( - name="deployed_index_id", - full_name="google.cloud.aiplatform.container.v1beta1.BatchMatchResponse.BatchMatchResponsePerIndex.deployed_index_id", - index=0, - number=1, - type=9, - cpp_type=9, - label=1, - has_default_value=False, - default_value=b"".decode("utf-8"), - message_type=None, - enum_type=None, - containing_type=None, - is_extension=False, - extension_scope=None, - serialized_options=None, - file=DESCRIPTOR, - create_key=_descriptor._internal_create_key, - ), - _descriptor.FieldDescriptor( - name="responses", - full_name="google.cloud.aiplatform.container.v1beta1.BatchMatchResponse.BatchMatchResponsePerIndex.responses", - index=1, - number=2, - type=11, - cpp_type=10, - label=3, - has_default_value=False, - default_value=[], - message_type=None, - enum_type=None, - containing_type=None, - is_extension=False, - extension_scope=None, - serialized_options=None, - file=DESCRIPTOR, - create_key=_descriptor._internal_create_key, - ), - _descriptor.FieldDescriptor( - name="status", - full_name="google.cloud.aiplatform.container.v1beta1.BatchMatchResponse.BatchMatchResponsePerIndex.status", - index=2, - number=3, - type=11, - cpp_type=10, - label=1, - has_default_value=False, - default_value=None, - message_type=None, - enum_type=None, - containing_type=None, - is_extension=False, - extension_scope=None, - serialized_options=None, - file=DESCRIPTOR, - create_key=_descriptor._internal_create_key, - ), - ], - extensions=[], - nested_types=[], - enum_types=[], - serialized_options=None, - is_extendable=False, - syntax="proto3", - extension_ranges=[], - oneofs=[], - serialized_start=941, - serialized_end=1109, -) - -_BATCHMATCHRESPONSE = _descriptor.Descriptor( - name="BatchMatchResponse", - full_name="google.cloud.aiplatform.container.v1beta1.BatchMatchResponse", - filename=None, - file=DESCRIPTOR, - containing_type=None, - create_key=_descriptor._internal_create_key, - fields=[ - _descriptor.FieldDescriptor( - name="responses", - full_name="google.cloud.aiplatform.container.v1beta1.BatchMatchResponse.responses", - index=0, - number=1, - type=11, - cpp_type=10, - label=3, - has_default_value=False, - default_value=[], - message_type=None, - enum_type=None, - containing_type=None, - is_extension=False, - extension_scope=None, - serialized_options=None, - file=DESCRIPTOR, - create_key=_descriptor._internal_create_key, - ), - ], - extensions=[], - nested_types=[ - _BATCHMATCHRESPONSE_BATCHMATCHRESPONSEPERINDEX, - ], - enum_types=[], - serialized_options=None, - is_extendable=False, - syntax="proto3", - extension_ranges=[], - oneofs=[], - serialized_start=809, - serialized_end=1109, -) - - -_NAMESPACE = _descriptor.Descriptor( - name="Namespace", - full_name="google.cloud.aiplatform.container.v1beta1.Namespace", - filename=None, - file=DESCRIPTOR, - containing_type=None, - create_key=_descriptor._internal_create_key, - fields=[ - _descriptor.FieldDescriptor( - name="name", - full_name="google.cloud.aiplatform.container.v1beta1.Namespace.name", - index=0, - number=1, - type=9, - cpp_type=9, - label=1, - has_default_value=False, - default_value=b"".decode("utf-8"), - message_type=None, - enum_type=None, - containing_type=None, - is_extension=False, - extension_scope=None, - serialized_options=None, - file=DESCRIPTOR, - create_key=_descriptor._internal_create_key, - ), - _descriptor.FieldDescriptor( - name="allow_tokens", - full_name="google.cloud.aiplatform.container.v1beta1.Namespace.allow_tokens", - index=1, - number=2, - type=9, - cpp_type=9, - label=3, - has_default_value=False, - default_value=[], - message_type=None, - enum_type=None, - containing_type=None, - is_extension=False, - extension_scope=None, - serialized_options=None, - file=DESCRIPTOR, - create_key=_descriptor._internal_create_key, - ), - _descriptor.FieldDescriptor( - name="deny_tokens", - full_name="google.cloud.aiplatform.container.v1beta1.Namespace.deny_tokens", - index=2, - number=3, - type=9, - cpp_type=9, - label=3, - has_default_value=False, - default_value=[], - message_type=None, - enum_type=None, - containing_type=None, - is_extension=False, - extension_scope=None, - serialized_options=None, - file=DESCRIPTOR, - create_key=_descriptor._internal_create_key, - ), - ], - extensions=[], - nested_types=[], - enum_types=[], - serialized_options=None, - is_extendable=False, - syntax="proto3", - extension_ranges=[], - oneofs=[], - serialized_start=1111, - serialized_end=1179, -) - -_MATCHREQUEST.fields_by_name["restricts"].message_type = _NAMESPACE -_MATCHRESPONSE_NEIGHBOR.containing_type = _MATCHRESPONSE -_MATCHRESPONSE.fields_by_name["neighbor"].message_type = _MATCHRESPONSE_NEIGHBOR -_BATCHMATCHREQUEST_BATCHMATCHREQUESTPERINDEX.fields_by_name[ - "requests" -].message_type = _MATCHREQUEST -_BATCHMATCHREQUEST_BATCHMATCHREQUESTPERINDEX.containing_type = _BATCHMATCHREQUEST -_BATCHMATCHREQUEST.fields_by_name[ - "requests" -].message_type = _BATCHMATCHREQUEST_BATCHMATCHREQUESTPERINDEX -_BATCHMATCHRESPONSE_BATCHMATCHRESPONSEPERINDEX.fields_by_name[ - "responses" -].message_type = _MATCHRESPONSE -_BATCHMATCHRESPONSE_BATCHMATCHRESPONSEPERINDEX.fields_by_name[ - "status" -].message_type = google_dot_rpc_dot_status__pb2._STATUS -_BATCHMATCHRESPONSE_BATCHMATCHRESPONSEPERINDEX.containing_type = _BATCHMATCHRESPONSE -_BATCHMATCHRESPONSE.fields_by_name[ - "responses" -].message_type = _BATCHMATCHRESPONSE_BATCHMATCHRESPONSEPERINDEX -DESCRIPTOR.message_types_by_name["MatchRequest"] = _MATCHREQUEST -DESCRIPTOR.message_types_by_name["MatchResponse"] = _MATCHRESPONSE -DESCRIPTOR.message_types_by_name["BatchMatchRequest"] = _BATCHMATCHREQUEST -DESCRIPTOR.message_types_by_name["BatchMatchResponse"] = _BATCHMATCHRESPONSE -DESCRIPTOR.message_types_by_name["Namespace"] = _NAMESPACE -_sym_db.RegisterFileDescriptor(DESCRIPTOR) - +_NAMESPACE = DESCRIPTOR.message_types_by_name["Namespace"] MatchRequest = _reflection.GeneratedProtocolMessageType( "MatchRequest", (_message.Message,), { "DESCRIPTOR": _MATCHREQUEST, - "__module__": "match_service_pb2" + "__module__": "google.cloud.aiplatform.matching_engine._protos.match_service_pb2" # @@protoc_insertion_point(class_scope:google.cloud.aiplatform.container.v1beta1.MatchRequest) }, ) @@ -667,12 +68,12 @@ (_message.Message,), { "DESCRIPTOR": _MATCHRESPONSE_NEIGHBOR, - "__module__": "match_service_pb2" + "__module__": "google.cloud.aiplatform.matching_engine._protos.match_service_pb2" # @@protoc_insertion_point(class_scope:google.cloud.aiplatform.container.v1beta1.MatchResponse.Neighbor) }, ), "DESCRIPTOR": _MATCHRESPONSE, - "__module__": "match_service_pb2" + "__module__": "google.cloud.aiplatform.matching_engine._protos.match_service_pb2" # @@protoc_insertion_point(class_scope:google.cloud.aiplatform.container.v1beta1.MatchResponse) }, ) @@ -688,12 +89,12 @@ (_message.Message,), { "DESCRIPTOR": _BATCHMATCHREQUEST_BATCHMATCHREQUESTPERINDEX, - "__module__": "match_service_pb2" + "__module__": "google.cloud.aiplatform.matching_engine._protos.match_service_pb2" # @@protoc_insertion_point(class_scope:google.cloud.aiplatform.container.v1beta1.BatchMatchRequest.BatchMatchRequestPerIndex) }, ), "DESCRIPTOR": _BATCHMATCHREQUEST, - "__module__": "match_service_pb2" + "__module__": "google.cloud.aiplatform.matching_engine._protos.match_service_pb2" # @@protoc_insertion_point(class_scope:google.cloud.aiplatform.container.v1beta1.BatchMatchRequest) }, ) @@ -709,12 +110,12 @@ (_message.Message,), { "DESCRIPTOR": _BATCHMATCHRESPONSE_BATCHMATCHRESPONSEPERINDEX, - "__module__": "match_service_pb2" + "__module__": "google.cloud.aiplatform.matching_engine._protos.match_service_pb2" # @@protoc_insertion_point(class_scope:google.cloud.aiplatform.container.v1beta1.BatchMatchResponse.BatchMatchResponsePerIndex) }, ), "DESCRIPTOR": _BATCHMATCHRESPONSE, - "__module__": "match_service_pb2" + "__module__": "google.cloud.aiplatform.matching_engine._protos.match_service_pb2" # @@protoc_insertion_point(class_scope:google.cloud.aiplatform.container.v1beta1.BatchMatchResponse) }, ) @@ -726,47 +127,32 @@ (_message.Message,), { "DESCRIPTOR": _NAMESPACE, - "__module__": "match_service_pb2" + "__module__": "google.cloud.aiplatform.matching_engine._protos.match_service_pb2" # @@protoc_insertion_point(class_scope:google.cloud.aiplatform.container.v1beta1.Namespace) }, ) _sym_db.RegisterMessage(Namespace) - -_MATCHSERVICE = _descriptor.ServiceDescriptor( - name="MatchService", - full_name="google.cloud.aiplatform.container.v1beta1.MatchService", - file=DESCRIPTOR, - index=0, - serialized_options=None, - create_key=_descriptor._internal_create_key, - serialized_start=1182, - serialized_end=1464, - methods=[ - _descriptor.MethodDescriptor( - name="Match", - full_name="google.cloud.aiplatform.container.v1beta1.MatchService.Match", - index=0, - containing_service=None, - input_type=_MATCHREQUEST, - output_type=_MATCHRESPONSE, - serialized_options=None, - create_key=_descriptor._internal_create_key, - ), - _descriptor.MethodDescriptor( - name="BatchMatch", - full_name="google.cloud.aiplatform.container.v1beta1.MatchService.BatchMatch", - index=1, - containing_service=None, - input_type=_BATCHMATCHREQUEST, - output_type=_BATCHMATCHRESPONSE, - serialized_options=None, - create_key=_descriptor._internal_create_key, - ), - ], -) -_sym_db.RegisterServiceDescriptor(_MATCHSERVICE) - -DESCRIPTOR.services_by_name["MatchService"] = _MATCHSERVICE - +_MATCHSERVICE = DESCRIPTOR.services_by_name["MatchService"] +if _descriptor._USE_C_DESCRIPTORS == False: + + DESCRIPTOR._options = None + _MATCHREQUEST._serialized_start = 140 + _MATCHREQUEST._serialized_end = 419 + _MATCHRESPONSE._serialized_start = 422 + _MATCHRESPONSE._serialized_end = 564 + _MATCHRESPONSE_NEIGHBOR._serialized_start = 524 + _MATCHRESPONSE_NEIGHBOR._serialized_end = 564 + _BATCHMATCHREQUEST._serialized_start = 567 + _BATCHMATCHREQUEST._serialized_end = 854 + _BATCHMATCHREQUEST_BATCHMATCHREQUESTPERINDEX._serialized_start = 695 + _BATCHMATCHREQUEST_BATCHMATCHREQUESTPERINDEX._serialized_end = 854 + _BATCHMATCHRESPONSE._serialized_start = 857 + _BATCHMATCHRESPONSE._serialized_end = 1157 + _BATCHMATCHRESPONSE_BATCHMATCHRESPONSEPERINDEX._serialized_start = 989 + _BATCHMATCHRESPONSE_BATCHMATCHRESPONSEPERINDEX._serialized_end = 1157 + _NAMESPACE._serialized_start = 1159 + _NAMESPACE._serialized_end = 1227 + _MATCHSERVICE._serialized_start = 1230 + _MATCHSERVICE._serialized_end = 1512 # @@protoc_insertion_point(module_scope) diff --git a/google/cloud/aiplatform/matching_engine/_protos/match_service_pb2_grpc.py b/google/cloud/aiplatform/matching_engine/_protos/match_service_pb2_grpc.py index 9c99081a16..2c0c14f8ed 100644 --- a/google/cloud/aiplatform/matching_engine/_protos/match_service_pb2_grpc.py +++ b/google/cloud/aiplatform/matching_engine/_protos/match_service_pb2_grpc.py @@ -14,7 +14,23 @@ # See the License for the specific language governing permissions and # limitations under the License. # + +# Copyright 2022 Google LLC +# +# Licensed 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. + # Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT! + """Client and server classes corresponding to protobuf-defined services.""" from google.cloud.aiplatform.matching_engine._protos import match_service_pb2 diff --git a/google/cloud/aiplatform/models.py b/google/cloud/aiplatform/models.py index 95f3044cbe..7d104b3112 100644 --- a/google/cloud/aiplatform/models.py +++ b/google/cloud/aiplatform/models.py @@ -40,7 +40,6 @@ from google.cloud.aiplatform.compat.types import ( encryption_spec as gca_encryption_spec, endpoint as gca_endpoint_compat, - endpoint_v1 as gca_endpoint_v1, explanation as gca_explanation_compat, io as gca_io_compat, machine_resources as gca_machine_resources_compat, @@ -52,6 +51,7 @@ from google.protobuf import field_mask_pb2, json_format _DEFAULT_MACHINE_TYPE = "n1-standard-2" +_DEPLOYING_MODEL_TRAFFIC_SPLIT_KEY = "0" _LOGGER = base.Logger(__name__) @@ -486,7 +486,7 @@ def _allocate_traffic( new_traffic_split[deployed_model] += 1 unallocated_traffic -= 1 - new_traffic_split["0"] = traffic_percentage + new_traffic_split[_DEPLOYING_MODEL_TRAFFIC_SPLIT_KEY] = traffic_percentage return new_traffic_split @@ -612,7 +612,6 @@ def _validate_deploy_args( raise ValueError("Traffic percentage cannot be negative.") elif traffic_split: - # TODO(b/172678233) verify every referenced deployed model exists if sum(traffic_split.values()) != 100: raise ValueError( "Sum of all traffic within traffic split needs to be 100." @@ -644,6 +643,8 @@ def deploy( metadata: Optional[Sequence[Tuple[str, str]]] = (), sync=True, deploy_request_timeout: Optional[float] = None, + autoscaling_target_cpu_utilization: Optional[int] = None, + autoscaling_target_accelerator_duty_cycle: Optional[int] = None, ) -> None: """Deploys a Model to the Endpoint. @@ -717,6 +718,13 @@ def deploy( be immediately returned and synced when the Future has completed. deploy_request_timeout (float): Optional. The timeout for the deploy request in seconds. + autoscaling_target_cpu_utilization (int): + Target CPU Utilization to use for Autoscaling Replicas. + A default value of 60 will be used if not specified. + autoscaling_target_accelerator_duty_cycle (int): + Target Accelerator Duty Cycle. + Must also set accelerator_type and accelerator_count if specified. + A default value of 60 will be used if not specified. """ self._sync_gca_resource_if_skipped() @@ -747,6 +755,8 @@ def deploy( metadata=metadata, sync=sync, deploy_request_timeout=deploy_request_timeout, + autoscaling_target_cpu_utilization=autoscaling_target_cpu_utilization, + autoscaling_target_accelerator_duty_cycle=autoscaling_target_accelerator_duty_cycle, ) @base.optional_sync() @@ -767,6 +777,8 @@ def _deploy( metadata: Optional[Sequence[Tuple[str, str]]] = (), sync=True, deploy_request_timeout: Optional[float] = None, + autoscaling_target_cpu_utilization: Optional[int] = None, + autoscaling_target_accelerator_duty_cycle: Optional[int] = None, ) -> None: """Deploys a Model to the Endpoint. @@ -840,6 +852,13 @@ def _deploy( be immediately returned and synced when the Future has completed. deploy_request_timeout (float): Optional. The timeout for the deploy request in seconds. + autoscaling_target_cpu_utilization (int): + Target CPU Utilization to use for Autoscaling Replicas. + A default value of 60 will be used if not specified. + autoscaling_target_accelerator_duty_cycle (int): + Target Accelerator Duty Cycle. + Must also set accelerator_type and accelerator_count if specified. + A default value of 60 will be used if not specified. Raises: ValueError: If there is not current traffic split and traffic percentage is not 0 or 100. @@ -866,6 +885,8 @@ def _deploy( explanation_parameters=explanation_parameters, metadata=metadata, deploy_request_timeout=deploy_request_timeout, + autoscaling_target_cpu_utilization=autoscaling_target_cpu_utilization, + autoscaling_target_accelerator_duty_cycle=autoscaling_target_accelerator_duty_cycle, ) _LOGGER.log_action_completed_against_resource("model", "deployed", self) @@ -892,6 +913,8 @@ def _deploy_call( explanation_parameters: Optional[explain.ExplanationParameters] = None, metadata: Optional[Sequence[Tuple[str, str]]] = (), deploy_request_timeout: Optional[float] = None, + autoscaling_target_cpu_utilization: Optional[int] = None, + autoscaling_target_accelerator_duty_cycle: Optional[int] = None, ): """Helper method to deploy model to endpoint. @@ -965,6 +988,13 @@ def _deploy_call( be immediately returned and synced when the Future has completed. deploy_request_timeout (float): Optional. The timeout for the deploy request in seconds. + autoscaling_target_cpu_utilization (int): + Optional. Target CPU Utilization to use for Autoscaling Replicas. + A default value of 60 will be used if not specified. + autoscaling_target_accelerator_duty_cycle (int): + Optional. Target Accelerator Duty Cycle. + Must also set accelerator_type and accelerator_count if specified. + A default value of 60 will be used if not specified. Raises: ValueError: If there is not current traffic split and traffic percentage is not 0 or 100. @@ -980,6 +1010,14 @@ def _deploy_call( "Both `accelerator_type` and `accelerator_count` should be specified or None." ) + if autoscaling_target_accelerator_duty_cycle is not None and ( + not accelerator_type or not accelerator_count + ): + raise ValueError( + "Both `accelerator_type` and `accelerator_count` should be set " + "when specifying autoscaling_target_accelerator_duty_cycle`" + ) + deployed_model = gca_endpoint_compat.DeployedModel( model=model.resource_name, display_name=deployed_model_display_name, @@ -995,7 +1033,11 @@ def _deploy_call( in model.supported_deployment_resources_types ) provided_custom_machine_spec = ( - machine_type or accelerator_type or accelerator_count + machine_type + or accelerator_type + or accelerator_count + or autoscaling_target_accelerator_duty_cycle + or autoscaling_target_cpu_utilization ) # If the model supports both automatic and dedicated deployment resources, @@ -1007,7 +1049,9 @@ def _deploy_call( if provided_custom_machine_spec and not use_dedicated_resources: _LOGGER.info( "Model does not support dedicated deployment resources. " - "The machine_type, accelerator_type and accelerator_count parameters are ignored." + "The machine_type, accelerator_type and accelerator_count," + "autoscaling_target_accelerator_duty_cycle," + "autoscaling_target_cpu_utilization parameters are ignored." ) if use_dedicated_resources and not machine_type: @@ -1015,22 +1059,41 @@ def _deploy_call( _LOGGER.info(f"Using default machine_type: {machine_type}") if use_dedicated_resources: + + dedicated_resources = gca_machine_resources_compat.DedicatedResources( + min_replica_count=min_replica_count, + max_replica_count=max_replica_count, + ) + machine_spec = gca_machine_resources_compat.MachineSpec( machine_type=machine_type ) + if autoscaling_target_cpu_utilization: + autoscaling_metric_spec = gca_machine_resources_compat.AutoscalingMetricSpec( + metric_name="aiplatform.googleapis.com/prediction/online/cpu/utilization", + target=autoscaling_target_cpu_utilization, + ) + dedicated_resources.autoscaling_metric_specs.extend( + [autoscaling_metric_spec] + ) + if accelerator_type and accelerator_count: utils.validate_accelerator_type(accelerator_type) machine_spec.accelerator_type = accelerator_type machine_spec.accelerator_count = accelerator_count - deployed_model.dedicated_resources = ( - gca_machine_resources_compat.DedicatedResources( - machine_spec=machine_spec, - min_replica_count=min_replica_count, - max_replica_count=max_replica_count, - ) - ) + if autoscaling_target_accelerator_duty_cycle: + autoscaling_metric_spec = gca_machine_resources_compat.AutoscalingMetricSpec( + metric_name="aiplatform.googleapis.com/prediction/online/accelerator/duty_cycle", + target=autoscaling_target_accelerator_duty_cycle, + ) + dedicated_resources.autoscaling_metric_specs.extend( + [autoscaling_metric_spec] + ) + + dedicated_resources.machine_spec = machine_spec + deployed_model.dedicated_resources = dedicated_resources elif supports_automatic_resources: deployed_model.automatic_resources = ( @@ -1227,6 +1290,110 @@ def _instantiate_prediction_client( prediction_client=True, ) + def update( + self, + display_name: Optional[str] = None, + description: Optional[str] = None, + labels: Optional[Dict[str, str]] = None, + traffic_split: Optional[Dict[str, int]] = None, + request_metadata: Optional[Sequence[Tuple[str, str]]] = (), + update_request_timeout: Optional[float] = None, + ) -> "Endpoint": + """Updates an endpoint. + + Example usage: + + my_endpoint = my_endpoint.update( + display_name='my-updated-endpoint', + description='my updated description', + labels={'key': 'value'}, + traffic_split={ + '123456': 20, + '234567': 80, + }, + ) + + Args: + display_name (str): + Optional. The display name of the Endpoint. + The name can be up to 128 characters long and can be consist of any UTF-8 + characters. + description (str): + Optional. The description of the Endpoint. + labels (Dict[str, str]): + Optional. The labels with user-defined metadata to organize your Endpoints. + Label keys and values can be no longer than 64 characters + (Unicode codepoints), can only contain lowercase letters, numeric + characters, underscores and dashes. International characters are allowed. + See https://goo.gl/xmQnxf for more information and examples of labels. + traffic_split (Dict[str, int]): + Optional. A map from a DeployedModel's ID to the percentage of this Endpoint's + traffic that should be forwarded to that DeployedModel. + If a DeployedModel's ID is not listed in this map, then it receives no traffic. + The traffic percentage values must add up to 100, or map must be empty if + the Endpoint is to not accept any traffic at a moment. + request_metadata (Sequence[Tuple[str, str]]): + Optional. Strings which should be sent along with the request as metadata. + update_request_timeout (float): + Optional. The timeout for the update request in seconds. + + Returns: + Endpoint - Updated endpoint resource. + + Raises: + ValueError: If `labels` is not the correct format. + """ + + self.wait() + + current_endpoint_proto = self.gca_resource + copied_endpoint_proto = current_endpoint_proto.__class__(current_endpoint_proto) + + update_mask: List[str] = [] + + if display_name: + utils.validate_display_name(display_name) + copied_endpoint_proto.display_name = display_name + update_mask.append("display_name") + + if description: + copied_endpoint_proto.description = description + update_mask.append("description") + + if labels: + utils.validate_labels(labels) + copied_endpoint_proto.labels = labels + update_mask.append("labels") + + if traffic_split: + update_mask.append("traffic_split") + copied_endpoint_proto.traffic_split = traffic_split + + update_mask = field_mask_pb2.FieldMask(paths=update_mask) + + _LOGGER.log_action_start_against_resource( + "Updating", + "endpoint", + self, + ) + + update_endpoint_lro = self.api_client.update_endpoint( + endpoint=copied_endpoint_proto, + update_mask=update_mask, + metadata=request_metadata, + timeout=update_request_timeout, + ) + + _LOGGER.log_action_started_against_resource_with_lro( + "Update", "endpoint", self.__class__, update_endpoint_lro + ) + + update_endpoint_lro.result() + + _LOGGER.log_action_completed_against_resource("endpoint", "updated", self) + + return self + def predict( self, instances: List, @@ -1382,15 +1549,15 @@ def list( credentials=credentials, ) - def list_models(self) -> Sequence[gca_endpoint_v1.DeployedModel]: + def list_models(self) -> List[gca_endpoint_compat.DeployedModel]: """Returns a list of the models deployed to this Endpoint. Returns: - deployed_models (Sequence[aiplatform.gapic.DeployedModel]): + deployed_models (List[aiplatform.gapic.DeployedModel]): A list of the models deployed in this Endpoint. """ self._sync_gca_resource() - return self._gca_resource.deployed_models + return list(self._gca_resource.deployed_models) def undeploy_all(self, sync: bool = True) -> "Endpoint": """Undeploys every model deployed to this Endpoint. @@ -1995,6 +2162,8 @@ def deploy( encryption_spec_key_name: Optional[str] = None, sync=True, deploy_request_timeout: Optional[float] = None, + autoscaling_target_cpu_utilization: Optional[int] = None, + autoscaling_target_accelerator_duty_cycle: Optional[int] = None, ) -> Endpoint: """Deploys model to endpoint. Endpoint will be created if unspecified. @@ -2079,6 +2248,13 @@ def deploy( be immediately returned and synced when the Future has completed. deploy_request_timeout (float): Optional. The timeout for the deploy request in seconds. + autoscaling_target_cpu_utilization (int): + Optional. Target CPU Utilization to use for Autoscaling Replicas. + A default value of 60 will be used if not specified. + autoscaling_target_accelerator_duty_cycle (int): + Optional. Target Accelerator Duty Cycle. + Must also set accelerator_type and accelerator_count if specified. + A default value of 60 will be used if not specified. Returns: endpoint ("Endpoint"): Endpoint with the deployed model. @@ -2113,6 +2289,8 @@ def deploy( or initializer.global_config.encryption_spec_key_name, sync=sync, deploy_request_timeout=deploy_request_timeout, + autoscaling_target_cpu_utilization=autoscaling_target_cpu_utilization, + autoscaling_target_accelerator_duty_cycle=autoscaling_target_accelerator_duty_cycle, ) @base.optional_sync(return_input_arg="endpoint", bind_future_to_self=False) @@ -2134,6 +2312,8 @@ def _deploy( encryption_spec_key_name: Optional[str] = None, sync: bool = True, deploy_request_timeout: Optional[float] = None, + autoscaling_target_cpu_utilization: Optional[int] = None, + autoscaling_target_accelerator_duty_cycle: Optional[int] = None, ) -> Endpoint: """Deploys model to endpoint. Endpoint will be created if unspecified. @@ -2218,6 +2398,13 @@ def _deploy( be immediately returned and synced when the Future has completed. deploy_request_timeout (float): Optional. The timeout for the deploy request in seconds. + autoscaling_target_cpu_utilization (int): + Optional. Target CPU Utilization to use for Autoscaling Replicas. + A default value of 60 will be used if not specified. + autoscaling_target_accelerator_duty_cycle (int): + Optional. Target Accelerator Duty Cycle. + Must also set accelerator_type and accelerator_count if specified. + A default value of 60 will be used if not specified. Returns: endpoint ("Endpoint"): Endpoint with the deployed model. @@ -2253,6 +2440,8 @@ def _deploy( explanation_parameters=explanation_parameters, metadata=metadata, deploy_request_timeout=deploy_request_timeout, + autoscaling_target_cpu_utilization=autoscaling_target_cpu_utilization, + autoscaling_target_accelerator_duty_cycle=autoscaling_target_accelerator_duty_cycle, ) _LOGGER.log_action_completed_against_resource("model", "deployed", endpoint) diff --git a/google/cloud/aiplatform/pipeline_jobs.py b/google/cloud/aiplatform/pipeline_jobs.py index 4c8a3ad806..90d7e0f86d 100644 --- a/google/cloud/aiplatform/pipeline_jobs.py +++ b/google/cloud/aiplatform/pipeline_jobs.py @@ -30,24 +30,22 @@ from google.protobuf import json_format from google.cloud.aiplatform.compat.types import ( - pipeline_job_v1 as gca_pipeline_job_v1, - pipeline_state_v1 as gca_pipeline_state_v1, + pipeline_job as gca_pipeline_job, + pipeline_state as gca_pipeline_state, ) _LOGGER = base.Logger(__name__) _PIPELINE_COMPLETE_STATES = set( [ - gca_pipeline_state_v1.PipelineState.PIPELINE_STATE_SUCCEEDED, - gca_pipeline_state_v1.PipelineState.PIPELINE_STATE_FAILED, - gca_pipeline_state_v1.PipelineState.PIPELINE_STATE_CANCELLED, - gca_pipeline_state_v1.PipelineState.PIPELINE_STATE_PAUSED, + gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED, + gca_pipeline_state.PipelineState.PIPELINE_STATE_FAILED, + gca_pipeline_state.PipelineState.PIPELINE_STATE_CANCELLED, + gca_pipeline_state.PipelineState.PIPELINE_STATE_PAUSED, ] ) -_PIPELINE_ERROR_STATES = set( - [gca_pipeline_state_v1.PipelineState.PIPELINE_STATE_FAILED] -) +_PIPELINE_ERROR_STATES = set([gca_pipeline_state.PipelineState.PIPELINE_STATE_FAILED]) # Pattern for valid names used as a Vertex resource name. _VALID_NAME_PATTERN = re.compile("^[a-z][-a-z0-9]{0,127}$") @@ -205,7 +203,7 @@ def __init__( builder.update_runtime_parameters(parameter_values) runtime_config_dict = builder.build() - runtime_config = gca_pipeline_job_v1.PipelineJob.RuntimeConfig()._pb + runtime_config = gca_pipeline_job.PipelineJob.RuntimeConfig()._pb json_format.ParseDict(runtime_config_dict, runtime_config) pipeline_name = pipeline_job["pipelineSpec"]["pipelineInfo"]["name"] @@ -225,7 +223,7 @@ def __init__( if enable_caching is not None: _set_enable_caching_value(pipeline_job["pipelineSpec"], enable_caching) - self._gca_resource = gca_pipeline_job_v1.PipelineJob( + self._gca_resource = gca_pipeline_job.PipelineJob( display_name=display_name, pipeline_spec=pipeline_job["pipelineSpec"], labels=labels, @@ -326,7 +324,7 @@ def pipeline_spec(self): return self._gca_resource.pipeline_spec @property - def state(self) -> Optional[gca_pipeline_state_v1.PipelineState]: + def state(self) -> Optional[gca_pipeline_state.PipelineState]: """Current pipeline state.""" self._sync_gca_resource() return self._gca_resource.state @@ -337,7 +335,7 @@ def has_failed(self) -> bool: False otherwise. """ - return self.state == gca_pipeline_state_v1.PipelineState.PIPELINE_STATE_FAILED + return self.state == gca_pipeline_state.PipelineState.PIPELINE_STATE_FAILED def _dashboard_uri(self) -> str: """Helper method to compose the dashboard uri where pipeline can be diff --git a/google/cloud/aiplatform/training_jobs.py b/google/cloud/aiplatform/training_jobs.py index 2dbd130555..a6244e08ca 100644 --- a/google/cloud/aiplatform/training_jobs.py +++ b/google/cloud/aiplatform/training_jobs.py @@ -4037,6 +4037,14 @@ def run( model_display_name: Optional[str] = None, model_labels: Optional[Dict[str, str]] = None, additional_experiments: Optional[List[str]] = None, + hierarchy_group_columns: Optional[List[str]] = None, + hierarchy_group_total_weight: Optional[float] = None, + hierarchy_temporal_total_weight: Optional[float] = None, + hierarchy_group_temporal_total_weight: Optional[float] = None, + window_column: Optional[str] = None, + window_stride_length: Optional[int] = None, + window_max_count: Optional[int] = None, + holiday_regions: Optional[List[str]] = None, sync: bool = True, create_request_timeout: Optional[float] = None, ) -> models.Model: @@ -4157,7 +4165,7 @@ def run( Applies only if [export_evaluated_data_items] is True and [export_evaluated_data_items_bigquery_destination_uri] is specified. quantiles (List[float]): - Quantiles to use for the `minimize-quantile-loss` + Quantiles to use for the ``minimize-quantile-loss`` [AutoMLForecastingTrainingJob.optimization_objective]. This argument is required in this case. @@ -4200,10 +4208,54 @@ def run( Optional. Additional experiment flags for the time series forcasting training. create_request_timeout (float): Optional. The timeout for the create request in seconds. + hierarchy_group_columns (List[str]): + Optional. A list of time series attribute column names that + define the time series hierarchy. Only one level of hierarchy is + supported, ex. ``region`` for a hierarchy of stores or + ``department`` for a hierarchy of products. If multiple columns + are specified, time series will be grouped by their combined + values, ex. (``blue``, ``large``) for ``color`` and ``size``, up + to 5 columns are accepted. If no group columns are specified, + all time series are considered to be part of the same group. + hierarchy_group_total_weight (float): + Optional. The weight of the loss for predictions aggregated over + time series in the same hierarchy group. + hierarchy_temporal_total_weight (float): + Optional. The weight of the loss for predictions aggregated over + the horizon for a single time series. + hierarchy_group_temporal_total_weight (float): + Optional. The weight of the loss for predictions aggregated over + both the horizon and time series in the same hierarchy group. + window_column (str): + Optional. Name of the column that should be used to filter input + rows. The column should contain either booleans or string + booleans; if the value of the row is True, generate a sliding + window from that row. + window_stride_length (int): + Optional. Step length used to generate input examples. Every + ``window_stride_length`` rows will be used to generate a sliding + window. + window_max_count (int): + Optional. Number of rows that should be used to generate input + examples. If the total row count is larger than this number, the + input data will be randomly sampled to hit the count. + holiday_regions (List[str]): + Optional. The geographical regions to use when creating holiday + features. This option is only allowed when data_granularity_unit + is ``day``. Acceptable values can come from any of the following + levels: + Top level: GLOBAL + Second level: continental regions + NA: North America + JAPAC: Japan and Asia Pacific + EMEA: Europe, the Middle East and Africa + LAC: Latin America and the Caribbean + Third level: countries from ISO 3166-1 Country codes. sync (bool): - Whether to execute this method synchronously. If False, this method + Optional. Whether to execute this method synchronously. If False, this method will be executed in concurrent Future and any downstream object will be immediately returned and synced when the Future has completed. + Returns: model: The trained Vertex AI Model resource or None if training did not produce a Vertex AI Model. @@ -4254,6 +4306,14 @@ def run( validation_options=validation_options, model_display_name=model_display_name, model_labels=model_labels, + hierarchy_group_columns=hierarchy_group_columns, + hierarchy_group_total_weight=hierarchy_group_total_weight, + hierarchy_temporal_total_weight=hierarchy_temporal_total_weight, + hierarchy_group_temporal_total_weight=hierarchy_group_temporal_total_weight, + window_column=window_column, + window_stride_length=window_stride_length, + window_max_count=window_max_count, + holiday_regions=holiday_regions, sync=sync, create_request_timeout=create_request_timeout, ) @@ -4286,6 +4346,14 @@ def _run( budget_milli_node_hours: int = 1000, model_display_name: Optional[str] = None, model_labels: Optional[Dict[str, str]] = None, + hierarchy_group_columns: Optional[List[str]] = None, + hierarchy_group_total_weight: Optional[float] = None, + hierarchy_temporal_total_weight: Optional[float] = None, + hierarchy_group_temporal_total_weight: Optional[float] = None, + window_column: Optional[str] = None, + window_stride_length: Optional[int] = None, + window_max_count: Optional[int] = None, + holiday_regions: Optional[List[str]] = None, sync: bool = True, create_request_timeout: Optional[float] = None, ) -> models.Model: @@ -4453,12 +4521,56 @@ def _run( are allowed. See https://goo.gl/xmQnxf for more information and examples of labels. + hierarchy_group_columns (List[str]): + Optional. A list of time series attribute column names that + define the time series hierarchy. Only one level of hierarchy is + supported, ex. ``region`` for a hierarchy of stores or + ``department`` for a hierarchy of products. If multiple columns + are specified, time series will be grouped by their combined + values, ex. (``blue``, ``large``) for ``color`` and ``size``, up + to 5 columns are accepted. If no group columns are specified, + all time series are considered to be part of the same group. + hierarchy_group_total_weight (float): + Optional. The weight of the loss for predictions aggregated over + time series in the same hierarchy group. + hierarchy_temporal_total_weight (float): + Optional. The weight of the loss for predictions aggregated over + the horizon for a single time series. + hierarchy_group_temporal_total_weight (float): + Optional. The weight of the loss for predictions aggregated over + both the horizon and time series in the same hierarchy group. + window_column (str): + Optional. Name of the column that should be used to filter input + rows. The column should contain either booleans or string + booleans; if the value of the row is True, generate a sliding + window from that row. + window_stride_length (int): + Optional. Step length used to generate input examples. Every + ``window_stride_length`` rows will be used to generate a sliding + window. + window_max_count (int): + Optional. Number of rows that should be used to generate input + examples. If the total row count is larger than this number, the + input data will be randomly sampled to hit the count. + holiday_regions (List[str]): + Optional. The geographical regions to use when creating holiday + features. This option is only allowed when data_granularity_unit + is ``day``. Acceptable values can come from any of the following + levels: + Top level: GLOBAL + Second level: continental regions + NA: North America + JAPAC: Japan and Asia Pacific + EMEA: Europe, the Middle East and Africa + LAC: Latin America and the Caribbean + Third level: countries from ISO 3166-1 Country codes. sync (bool): Whether to execute this method synchronously. If False, this method will be executed in concurrent Future and any downstream object will be immediately returned and synced when the Future has completed. create_request_timeout (float): Optional. The timeout for the create request in seconds. + Returns: model: The trained Vertex AI Model resource or None if training did not produce a Vertex AI Model. @@ -4482,6 +4594,12 @@ def _run( % column_names ) + window_config = self._create_window_config( + column=window_column, + stride_length=window_stride_length, + max_count=window_max_count, + ) + training_task_inputs_dict = { # required inputs "targetColumn": target_column, @@ -4503,8 +4621,27 @@ def _run( "quantiles": quantiles, "validationOptions": validation_options, "optimizationObjective": self._optimization_objective, + "holidayRegions": holiday_regions, } + # TODO(TheMichaelHu): Remove the ifs once the API supports these inputs. + if any( + [ + hierarchy_group_columns, + hierarchy_group_total_weight, + hierarchy_temporal_total_weight, + hierarchy_group_temporal_total_weight, + ] + ): + training_task_inputs_dict["hierarchyConfig"] = { + "groupColumns": hierarchy_group_columns, + "groupTotalWeight": hierarchy_group_total_weight, + "temporalTotalWeight": hierarchy_temporal_total_weight, + "groupTemporalTotalWeight": hierarchy_group_temporal_total_weight, + } + if window_config: + training_task_inputs_dict["windowConfig"] = window_config + final_export_eval_bq_uri = export_evaluated_data_items_bigquery_destination_uri if final_export_eval_bq_uri and not final_export_eval_bq_uri.startswith( "bq://" @@ -4582,6 +4719,29 @@ def _add_additional_experiments(self, additional_experiments: List[str]): """ self._additional_experiments.extend(additional_experiments) + @staticmethod + def _create_window_config( + column: Optional[str] = None, + stride_length: Optional[int] = None, + max_count: Optional[int] = None, + ) -> Optional[Dict[str, Union[int, str]]]: + """Creates a window config from training job arguments.""" + configs = { + "column": column, + "strideLength": stride_length, + "maxCount": max_count, + } + present_configs = {k: v for k, v in configs.items() if v is not None} + if not present_configs: + return None + if len(present_configs) > 1: + raise ValueError( + "More than one windowing strategy provided. Make sure only one " + "of window_column, window_stride_length, or window_max_count " + "is specified." + ) + return present_configs + class AutoMLImageTrainingJob(_TrainingJob): _supported_training_schemas = ( diff --git a/google/cloud/aiplatform/version.py b/google/cloud/aiplatform/version.py index 2313afcd27..ec04f5d40b 100644 --- a/google/cloud/aiplatform/version.py +++ b/google/cloud/aiplatform/version.py @@ -15,4 +15,4 @@ # limitations under the License. # -__version__ = "1.13.0" +__version__ = "1.13.1" diff --git a/tests/system/aiplatform/test_dataset.py b/tests/system/aiplatform/test_dataset.py index f152ed0e32..7cd3c0416c 100644 --- a/tests/system/aiplatform/test_dataset.py +++ b/tests/system/aiplatform/test_dataset.py @@ -51,6 +51,9 @@ ) _TEST_DATASET_DISPLAY_NAME = "permanent_50_flowers_dataset" _TEST_TABULAR_CLASSIFICATION_GCS_SOURCE = "gs://ucaip-sample-resources/iris_1000.csv" +_TEST_FORECASTING_BQ_SOURCE = ( + "bq://ucaip-sample-tests:ucaip_test_us_central1.2020_sales_train" +) _TEST_TEXT_ENTITY_EXTRACTION_GCS_SOURCE = f"gs://{TEST_BUCKET}/ai-platform-unified/sdk/datasets/text_entity_extraction_dataset.jsonl" _TEST_IMAGE_OBJECT_DETECTION_GCS_SOURCE = ( "gs://ucaip-test-us-central1/dataset/salads_oid_ml_use_public_unassigned.jsonl" @@ -189,6 +192,7 @@ def test_get_new_dataset_and_import(self, dataset_gapic_client): my_dataset.import_data( gcs_source=_TEST_TEXT_ENTITY_EXTRACTION_GCS_SOURCE, import_schema_uri=_TEST_TEXT_ENTITY_IMPORT_SCHEMA, + import_request_timeout=500, ) data_items_post_import = dataset_gapic_client.list_data_items( @@ -305,6 +309,30 @@ def test_create_tabular_dataset_from_dataframe_with_provided_schema( finally: tabular_dataset.delete() + def test_create_time_series_dataset(self): + """Use the Dataset.create() method to create a new time series dataset. + Then confirm the dataset was successfully created and references GCS source.""" + + try: + time_series_dataset = aiplatform.TimeSeriesDataset.create( + display_name=self._make_display_name(key="create_time_series_dataset"), + bq_source=[_TEST_FORECASTING_BQ_SOURCE], + create_request_timeout=None, + ) + + gapic_metadata = time_series_dataset.to_dict()["metadata"] + bq_source_uri = gapic_metadata["inputConfig"]["bigquerySource"]["uri"] + + assert _TEST_FORECASTING_BQ_SOURCE == bq_source_uri + assert ( + time_series_dataset.metadata_schema_uri + == aiplatform.schema.dataset.metadata.time_series + ) + + finally: + if time_series_dataset is not None: + time_series_dataset.delete() + def test_export_data(self, storage_client, staging_bucket): """Get an existing dataset, export data to a newly created folder in Google Cloud Storage, then verify data was successfully exported.""" diff --git a/tests/system/aiplatform/test_e2e_forecasting.py b/tests/system/aiplatform/test_e2e_forecasting.py new file mode 100644 index 0000000000..b0f3e19711 --- /dev/null +++ b/tests/system/aiplatform/test_e2e_forecasting.py @@ -0,0 +1,127 @@ +# -*- coding: utf-8 -*- + +# Copyright 2022 Google LLC +# +# Licensed 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 google.cloud import aiplatform +from google.cloud.aiplatform.compat.types import job_state +from google.cloud.aiplatform.compat.types import pipeline_state +import pytest +from tests.system.aiplatform import e2e_base + +_TRAINING_DATASET_BQ_PATH = ( + "bq://ucaip-sample-tests:ucaip_test_us_central1.2020_sales_train" +) +_PREDICTION_DATASET_BQ_PATH = ( + "bq://ucaip-sample-tests:ucaip_test_us_central1.2021_sales_predict" +) + + +@pytest.mark.usefixtures("prepare_staging_bucket", "delete_staging_bucket") +class TestEndToEndForecasting(e2e_base.TestEndToEnd): + """End to end system test of the Vertex SDK with forecasting data.""" + + _temp_prefix = "temp-vertex-sdk-e2e-forecasting" + + def test_end_to_end_forecasting(self, shared_state): + """Builds a dataset, trains models, and gets batch predictions.""" + ds = None + automl_job = None + automl_model = None + automl_batch_prediction_job = None + + aiplatform.init( + project=e2e_base._PROJECT, + location=e2e_base._LOCATION, + staging_bucket=shared_state["staging_bucket_name"], + ) + try: + # Create and import to single managed dataset for both training + # jobs. + ds = aiplatform.TimeSeriesDataset.create( + display_name=self._make_display_name("dataset"), + bq_source=[_TRAINING_DATASET_BQ_PATH], + sync=False, + create_request_timeout=180.0, + ) + + time_column = "date" + time_series_identifier_column = "store_name" + target_column = "sale_dollars" + column_specs = { + time_column: "timestamp", + target_column: "numeric", + "city": "categorical", + "zip_code": "categorical", + "county": "categorical", + } + + # Define both training jobs + # TODO(humichael): Add seq2seq job. + automl_job = aiplatform.AutoMLForecastingTrainingJob( + display_name=self._make_display_name("train-housing-automl"), + optimization_objective="minimize-rmse", + column_specs=column_specs, + ) + + # Kick off both training jobs, AutoML job will take approx one hour + # to run. + automl_model = automl_job.run( + dataset=ds, + target_column=target_column, + time_column=time_column, + time_series_identifier_column=time_series_identifier_column, + available_at_forecast_columns=[time_column], + unavailable_at_forecast_columns=[target_column], + time_series_attribute_columns=["city", "zip_code", "county"], + forecast_horizon=30, + context_window=30, + data_granularity_unit="day", + data_granularity_count=1, + budget_milli_node_hours=1000, + model_display_name=self._make_display_name("automl-liquor-model"), + sync=False, + ) + + automl_batch_prediction_job = automl_model.batch_predict( + job_display_name=self._make_display_name("automl-liquor-model"), + instances_format="bigquery", + machine_type="n1-standard-4", + bigquery_source=_PREDICTION_DATASET_BQ_PATH, + gcs_destination_prefix=( + f'gs://{shared_state["staging_bucket_name"]}/bp_results/' + ), + sync=False, + ) + + automl_batch_prediction_job.wait() + + assert ( + automl_job.state + == pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED + ) + assert ( + automl_batch_prediction_job.state + == job_state.JobState.JOB_STATE_SUCCEEDED + ) + finally: + if ds is not None: + ds.delete() + if automl_job is not None: + automl_job.delete() + if automl_model is not None: + automl_model.delete() + if automl_batch_prediction_job is not None: + automl_batch_prediction_job.delete() diff --git a/tests/system/aiplatform/test_e2e_tabular.py b/tests/system/aiplatform/test_e2e_tabular.py index 31b9bb9769..6746e816a9 100644 --- a/tests/system/aiplatform/test_e2e_tabular.py +++ b/tests/system/aiplatform/test_e2e_tabular.py @@ -46,7 +46,9 @@ } -@pytest.mark.usefixtures("prepare_staging_bucket", "delete_staging_bucket") +@pytest.mark.usefixtures( + "prepare_staging_bucket", "delete_staging_bucket", "tear_down_resources" +) class TestEndToEndTabular(e2e_base.TestEndToEnd): """End to end system test of the Vertex SDK with tabular data adapted from reference notebook http://shortn/_eyoNx3SN0X""" diff --git a/tests/system/aiplatform/test_featurestore.py b/tests/system/aiplatform/test_featurestore.py index a7abf3ce84..9573ea6556 100644 --- a/tests/system/aiplatform/test_featurestore.py +++ b/tests/system/aiplatform/test_featurestore.py @@ -51,6 +51,7 @@ "delete_staging_bucket", "prepare_bigquery_dataset", "delete_bigquery_dataset", + "tear_down_resources", ) class TestFeaturestore(e2e_base.TestEndToEnd): diff --git a/tests/system/aiplatform/test_metadata.py b/tests/system/aiplatform/test_metadata.py index 3d2116fa0a..d5c076e75b 100644 --- a/tests/system/aiplatform/test_metadata.py +++ b/tests/system/aiplatform/test_metadata.py @@ -15,6 +15,8 @@ # limitations under the License. # +import pytest + from google.cloud import aiplatform from tests.system.aiplatform import e2e_base @@ -24,6 +26,7 @@ METRICS = {"sdk-metric-test-1": 0.8, "sdk-metric-test-2": 100} +@pytest.mark.usefixtures("tear_down_resources") class TestMetadata(e2e_base.TestEndToEnd): _temp_prefix = "temp-vertex-sdk-e2e-test" diff --git a/tests/system/aiplatform/test_model_upload.py b/tests/system/aiplatform/test_model_upload.py index 3acc20c103..48e6169af4 100644 --- a/tests/system/aiplatform/test_model_upload.py +++ b/tests/system/aiplatform/test_model_upload.py @@ -28,7 +28,7 @@ _XGBOOST_MODEL_URI = "gs://cloud-samples-data-us-central1/vertex-ai/google-cloud-aiplatform-ci-artifacts/models/iris_xgboost/model.bst" -@pytest.mark.usefixtures("delete_staging_bucket") +@pytest.mark.usefixtures("delete_staging_bucket", "tear_down_resources") class TestModel(e2e_base.TestEndToEnd): _temp_prefix = "temp_vertex_sdk_e2e_model_upload_test" @@ -76,3 +76,12 @@ def test_upload_and_deploy_xgboost_model(self, shared_state): assert model.display_name == "new_name" assert model.description == "new_description" assert model.labels == {"my_label": "updated"} + + assert len(endpoint.list_models()) == 1 + endpoint.deploy(model, traffic_percentage=100) + assert len(endpoint.list_models()) == 2 + traffic_split = { + deployed_model.id: 50 for deployed_model in endpoint.list_models() + } + endpoint.update(traffic_split=traffic_split) + assert endpoint.traffic_split == traffic_split diff --git a/tests/system/aiplatform/test_tensorboard.py b/tests/system/aiplatform/test_tensorboard.py index 501205122f..a1f4634bd9 100644 --- a/tests/system/aiplatform/test_tensorboard.py +++ b/tests/system/aiplatform/test_tensorboard.py @@ -15,10 +15,13 @@ # limitations under the License. # +import pytest + from google.cloud import aiplatform from tests.system.aiplatform import e2e_base +@pytest.mark.usefixtures("tear_down_resources") class TestTensorboard(e2e_base.TestEndToEnd): _temp_prefix = "temp-vertex-sdk-e2e-test" diff --git a/tests/unit/aiplatform/conftest.py b/tests/unit/aiplatform/conftest.py new file mode 100644 index 0000000000..1a2e9c54f1 --- /dev/null +++ b/tests/unit/aiplatform/conftest.py @@ -0,0 +1,32 @@ +# -*- coding: utf-8 -*- + +# Copyright 2022 Google LLC +# +# Licensed 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. +# + +import pytest + +from google import auth +from google.auth import credentials as auth_credentials +from unittest.mock import patch + + +@pytest.fixture(scope="module") +def google_auth_mock(): + with patch.object(auth, "default") as google_auth_mock: + google_auth_mock.return_value = ( + auth_credentials.AnonymousCredentials(), + "test-project", + ) + yield google_auth_mock diff --git a/tests/unit/aiplatform/test_automl_forecasting_training_jobs.py b/tests/unit/aiplatform/test_automl_forecasting_training_jobs.py index 6a96d656e8..21ca78da2e 100644 --- a/tests/unit/aiplatform/test_automl_forecasting_training_jobs.py +++ b/tests/unit/aiplatform/test_automl_forecasting_training_jobs.py @@ -26,13 +26,12 @@ from google.cloud.aiplatform import schema from google.cloud.aiplatform.training_jobs import AutoMLForecastingTrainingJob -from google.cloud.aiplatform_v1.services.model_service import ( - client as model_service_client, +from google.cloud.aiplatform.compat.services import ( + model_service_client, + pipeline_service_client, ) -from google.cloud.aiplatform_v1.services.pipeline_service import ( - client as pipeline_service_client, -) -from google.cloud.aiplatform_v1.types import ( + +from google.cloud.aiplatform.compat.types import ( dataset as gca_dataset, model as gca_model, pipeline_state as gca_pipeline_state, @@ -82,6 +81,14 @@ _TEST_TRAINING_WEIGHT_COLUMN = "weight" _TEST_TRAINING_OPTIMIZATION_OBJECTIVE_NAME = "minimize-rmse" _TEST_ADDITIONAL_EXPERIMENTS = ["exp1", "exp2"] +_TEST_HIERARCHY_GROUP_COLUMNS = [] +_TEST_HIERARCHY_GROUP_TOTAL_WEIGHT = 1 +_TEST_HIERARCHY_TEMPORAL_TOTAL_WEIGHT = None +_TEST_HIERARCHY_GROUP_TEMPORAL_TOTAL_WEIGHT = None +_TEST_WINDOW_COLUMN = None +_TEST_WINDOW_STRIDE_LENGTH = 1 +_TEST_WINDOW_MAX_COUNT = None +_TEST_TRAINING_HOLIDAY_REGIONS = ["GLOBAL"] _TEST_TRAINING_TASK_INPUTS_DICT = { # required inputs "targetColumn": _TEST_TRAINING_TARGET_COLUMN, @@ -107,6 +114,16 @@ "quantiles": _TEST_TRAINING_QUANTILES, "validationOptions": _TEST_TRAINING_VALIDATION_OPTIONS, "optimizationObjective": _TEST_TRAINING_OPTIMIZATION_OBJECTIVE_NAME, + "hierarchyConfig": { + "groupColumns": _TEST_HIERARCHY_GROUP_COLUMNS, + "groupTotalWeight": _TEST_HIERARCHY_GROUP_TOTAL_WEIGHT, + "temporalTotalWeight": _TEST_HIERARCHY_TEMPORAL_TOTAL_WEIGHT, + "groupTemporalTotalWeight": _TEST_HIERARCHY_GROUP_TEMPORAL_TOTAL_WEIGHT, + }, + "windowConfig": { + "strideLength": _TEST_WINDOW_STRIDE_LENGTH, + }, + "holidayRegions": _TEST_TRAINING_HOLIDAY_REGIONS, } _TEST_TRAINING_TASK_INPUTS_WITH_ADDITIONAL_EXPERIMENTS = json_format.ParseDict( @@ -248,6 +265,7 @@ def mock_dataset_nontimeseries(): return ds +@pytest.mark.usefixtures("google_auth_mock") class TestAutoMLForecastingTrainingJob: def setup_method(self): importlib.reload(initializer) @@ -297,8 +315,16 @@ def test_run_call_pipeline_service_create( quantiles=_TEST_TRAINING_QUANTILES, validation_options=_TEST_TRAINING_VALIDATION_OPTIONS, additional_experiments=_TEST_ADDITIONAL_EXPERIMENTS, + hierarchy_group_columns=_TEST_HIERARCHY_GROUP_COLUMNS, + hierarchy_group_total_weight=_TEST_HIERARCHY_GROUP_TOTAL_WEIGHT, + hierarchy_temporal_total_weight=_TEST_HIERARCHY_TEMPORAL_TOTAL_WEIGHT, + hierarchy_group_temporal_total_weight=_TEST_HIERARCHY_GROUP_TEMPORAL_TOTAL_WEIGHT, + window_column=_TEST_WINDOW_COLUMN, + window_stride_length=_TEST_WINDOW_STRIDE_LENGTH, + window_max_count=_TEST_WINDOW_MAX_COUNT, sync=sync, create_request_timeout=None, + holiday_regions=_TEST_TRAINING_HOLIDAY_REGIONS, ) if not sync: @@ -385,8 +411,16 @@ def test_run_call_pipeline_service_create_with_timeout( quantiles=_TEST_TRAINING_QUANTILES, validation_options=_TEST_TRAINING_VALIDATION_OPTIONS, additional_experiments=_TEST_ADDITIONAL_EXPERIMENTS, + hierarchy_group_columns=_TEST_HIERARCHY_GROUP_COLUMNS, + hierarchy_group_total_weight=_TEST_HIERARCHY_GROUP_TOTAL_WEIGHT, + hierarchy_temporal_total_weight=_TEST_HIERARCHY_TEMPORAL_TOTAL_WEIGHT, + hierarchy_group_temporal_total_weight=_TEST_HIERARCHY_GROUP_TEMPORAL_TOTAL_WEIGHT, + window_column=_TEST_WINDOW_COLUMN, + window_stride_length=_TEST_WINDOW_STRIDE_LENGTH, + window_max_count=_TEST_WINDOW_MAX_COUNT, sync=sync, create_request_timeout=180.0, + holiday_regions=_TEST_TRAINING_HOLIDAY_REGIONS, ) if not sync: @@ -455,8 +489,16 @@ def test_run_call_pipeline_if_no_model_display_name_nor_model_labels( export_evaluated_data_items_override_destination=_TEST_TRAINING_EXPORT_EVALUATED_DATA_ITEMS_OVERRIDE_DESTINATION, quantiles=_TEST_TRAINING_QUANTILES, validation_options=_TEST_TRAINING_VALIDATION_OPTIONS, + hierarchy_group_columns=_TEST_HIERARCHY_GROUP_COLUMNS, + hierarchy_group_total_weight=_TEST_HIERARCHY_GROUP_TOTAL_WEIGHT, + hierarchy_temporal_total_weight=_TEST_HIERARCHY_TEMPORAL_TOTAL_WEIGHT, + hierarchy_group_temporal_total_weight=_TEST_HIERARCHY_GROUP_TEMPORAL_TOTAL_WEIGHT, + window_column=_TEST_WINDOW_COLUMN, + window_stride_length=_TEST_WINDOW_STRIDE_LENGTH, + window_max_count=_TEST_WINDOW_MAX_COUNT, sync=sync, create_request_timeout=None, + holiday_regions=_TEST_TRAINING_HOLIDAY_REGIONS, ) if not sync: @@ -525,8 +567,16 @@ def test_run_call_pipeline_if_set_additional_experiments( export_evaluated_data_items_override_destination=_TEST_TRAINING_EXPORT_EVALUATED_DATA_ITEMS_OVERRIDE_DESTINATION, quantiles=_TEST_TRAINING_QUANTILES, validation_options=_TEST_TRAINING_VALIDATION_OPTIONS, + hierarchy_group_columns=_TEST_HIERARCHY_GROUP_COLUMNS, + hierarchy_group_total_weight=_TEST_HIERARCHY_GROUP_TOTAL_WEIGHT, + hierarchy_temporal_total_weight=_TEST_HIERARCHY_TEMPORAL_TOTAL_WEIGHT, + hierarchy_group_temporal_total_weight=_TEST_HIERARCHY_GROUP_TEMPORAL_TOTAL_WEIGHT, + window_column=_TEST_WINDOW_COLUMN, + window_stride_length=_TEST_WINDOW_STRIDE_LENGTH, + window_max_count=_TEST_WINDOW_MAX_COUNT, sync=sync, create_request_timeout=None, + holiday_regions=_TEST_TRAINING_HOLIDAY_REGIONS, ) if not sync: @@ -592,7 +642,15 @@ def test_run_called_twice_raises( export_evaluated_data_items_override_destination=_TEST_TRAINING_EXPORT_EVALUATED_DATA_ITEMS_OVERRIDE_DESTINATION, quantiles=_TEST_TRAINING_QUANTILES, validation_options=_TEST_TRAINING_VALIDATION_OPTIONS, + hierarchy_group_columns=_TEST_HIERARCHY_GROUP_COLUMNS, + hierarchy_group_total_weight=_TEST_HIERARCHY_GROUP_TOTAL_WEIGHT, + hierarchy_temporal_total_weight=_TEST_HIERARCHY_TEMPORAL_TOTAL_WEIGHT, + hierarchy_group_temporal_total_weight=_TEST_HIERARCHY_GROUP_TEMPORAL_TOTAL_WEIGHT, + window_column=_TEST_WINDOW_COLUMN, + window_stride_length=_TEST_WINDOW_STRIDE_LENGTH, + window_max_count=_TEST_WINDOW_MAX_COUNT, sync=sync, + holiday_regions=_TEST_TRAINING_HOLIDAY_REGIONS, ) with pytest.raises(RuntimeError): @@ -616,7 +674,15 @@ def test_run_called_twice_raises( export_evaluated_data_items_override_destination=_TEST_TRAINING_EXPORT_EVALUATED_DATA_ITEMS_OVERRIDE_DESTINATION, quantiles=_TEST_TRAINING_QUANTILES, validation_options=_TEST_TRAINING_VALIDATION_OPTIONS, + hierarchy_group_columns=_TEST_HIERARCHY_GROUP_COLUMNS, + hierarchy_group_total_weight=_TEST_HIERARCHY_GROUP_TOTAL_WEIGHT, + hierarchy_temporal_total_weight=_TEST_HIERARCHY_TEMPORAL_TOTAL_WEIGHT, + hierarchy_group_temporal_total_weight=_TEST_HIERARCHY_GROUP_TEMPORAL_TOTAL_WEIGHT, + window_column=_TEST_WINDOW_COLUMN, + window_stride_length=_TEST_WINDOW_STRIDE_LENGTH, + window_max_count=_TEST_WINDOW_MAX_COUNT, sync=sync, + holiday_regions=_TEST_TRAINING_HOLIDAY_REGIONS, ) @pytest.mark.parametrize("sync", [True, False]) @@ -656,7 +722,15 @@ def test_run_raises_if_pipeline_fails( export_evaluated_data_items_override_destination=_TEST_TRAINING_EXPORT_EVALUATED_DATA_ITEMS_OVERRIDE_DESTINATION, quantiles=_TEST_TRAINING_QUANTILES, validation_options=_TEST_TRAINING_VALIDATION_OPTIONS, + hierarchy_group_columns=_TEST_HIERARCHY_GROUP_COLUMNS, + hierarchy_group_total_weight=_TEST_HIERARCHY_GROUP_TOTAL_WEIGHT, + hierarchy_temporal_total_weight=_TEST_HIERARCHY_TEMPORAL_TOTAL_WEIGHT, + hierarchy_group_temporal_total_weight=_TEST_HIERARCHY_GROUP_TEMPORAL_TOTAL_WEIGHT, + window_column=_TEST_WINDOW_COLUMN, + window_stride_length=_TEST_WINDOW_STRIDE_LENGTH, + window_max_count=_TEST_WINDOW_MAX_COUNT, sync=sync, + holiday_regions=_TEST_TRAINING_HOLIDAY_REGIONS, ) if not sync: @@ -731,8 +805,16 @@ def test_splits_fraction( export_evaluated_data_items_override_destination=_TEST_TRAINING_EXPORT_EVALUATED_DATA_ITEMS_OVERRIDE_DESTINATION, quantiles=_TEST_TRAINING_QUANTILES, validation_options=_TEST_TRAINING_VALIDATION_OPTIONS, + hierarchy_group_columns=_TEST_HIERARCHY_GROUP_COLUMNS, + hierarchy_group_total_weight=_TEST_HIERARCHY_GROUP_TOTAL_WEIGHT, + hierarchy_temporal_total_weight=_TEST_HIERARCHY_TEMPORAL_TOTAL_WEIGHT, + hierarchy_group_temporal_total_weight=_TEST_HIERARCHY_GROUP_TEMPORAL_TOTAL_WEIGHT, + window_column=_TEST_WINDOW_COLUMN, + window_stride_length=_TEST_WINDOW_STRIDE_LENGTH, + window_max_count=_TEST_WINDOW_MAX_COUNT, sync=sync, create_request_timeout=None, + holiday_regions=_TEST_TRAINING_HOLIDAY_REGIONS, ) if not sync: @@ -819,8 +901,16 @@ def test_splits_timestamp( export_evaluated_data_items_override_destination=_TEST_TRAINING_EXPORT_EVALUATED_DATA_ITEMS_OVERRIDE_DESTINATION, quantiles=_TEST_TRAINING_QUANTILES, validation_options=_TEST_TRAINING_VALIDATION_OPTIONS, + hierarchy_group_columns=_TEST_HIERARCHY_GROUP_COLUMNS, + hierarchy_group_total_weight=_TEST_HIERARCHY_GROUP_TOTAL_WEIGHT, + hierarchy_temporal_total_weight=_TEST_HIERARCHY_TEMPORAL_TOTAL_WEIGHT, + hierarchy_group_temporal_total_weight=_TEST_HIERARCHY_GROUP_TEMPORAL_TOTAL_WEIGHT, + window_column=_TEST_WINDOW_COLUMN, + window_stride_length=_TEST_WINDOW_STRIDE_LENGTH, + window_max_count=_TEST_WINDOW_MAX_COUNT, sync=sync, create_request_timeout=None, + holiday_regions=_TEST_TRAINING_HOLIDAY_REGIONS, ) if not sync: @@ -905,8 +995,16 @@ def test_splits_predefined( export_evaluated_data_items_override_destination=_TEST_TRAINING_EXPORT_EVALUATED_DATA_ITEMS_OVERRIDE_DESTINATION, quantiles=_TEST_TRAINING_QUANTILES, validation_options=_TEST_TRAINING_VALIDATION_OPTIONS, + hierarchy_group_columns=_TEST_HIERARCHY_GROUP_COLUMNS, + hierarchy_group_total_weight=_TEST_HIERARCHY_GROUP_TOTAL_WEIGHT, + hierarchy_temporal_total_weight=_TEST_HIERARCHY_TEMPORAL_TOTAL_WEIGHT, + hierarchy_group_temporal_total_weight=_TEST_HIERARCHY_GROUP_TEMPORAL_TOTAL_WEIGHT, + window_column=_TEST_WINDOW_COLUMN, + window_stride_length=_TEST_WINDOW_STRIDE_LENGTH, + window_max_count=_TEST_WINDOW_MAX_COUNT, sync=sync, create_request_timeout=None, + holiday_regions=_TEST_TRAINING_HOLIDAY_REGIONS, ) if not sync: @@ -986,8 +1084,16 @@ def test_splits_default( export_evaluated_data_items_override_destination=_TEST_TRAINING_EXPORT_EVALUATED_DATA_ITEMS_OVERRIDE_DESTINATION, quantiles=_TEST_TRAINING_QUANTILES, validation_options=_TEST_TRAINING_VALIDATION_OPTIONS, + hierarchy_group_columns=_TEST_HIERARCHY_GROUP_COLUMNS, + hierarchy_group_total_weight=_TEST_HIERARCHY_GROUP_TOTAL_WEIGHT, + hierarchy_temporal_total_weight=_TEST_HIERARCHY_TEMPORAL_TOTAL_WEIGHT, + hierarchy_group_temporal_total_weight=_TEST_HIERARCHY_GROUP_TEMPORAL_TOTAL_WEIGHT, + window_column=_TEST_WINDOW_COLUMN, + window_stride_length=_TEST_WINDOW_STRIDE_LENGTH, + window_max_count=_TEST_WINDOW_MAX_COUNT, sync=sync, create_request_timeout=None, + holiday_regions=_TEST_TRAINING_HOLIDAY_REGIONS, ) if not sync: diff --git a/tests/unit/aiplatform/test_automl_image_training_jobs.py b/tests/unit/aiplatform/test_automl_image_training_jobs.py index 95e3c3f641..30c51f2e62 100644 --- a/tests/unit/aiplatform/test_automl_image_training_jobs.py +++ b/tests/unit/aiplatform/test_automl_image_training_jobs.py @@ -30,13 +30,11 @@ from google.cloud.aiplatform import schema from google.cloud.aiplatform import training_jobs -from google.cloud.aiplatform_v1.services.model_service import ( - client as model_service_client, +from google.cloud.aiplatform.compat.services import ( + model_service_client, + pipeline_service_client, ) -from google.cloud.aiplatform_v1.services.pipeline_service import ( - client as pipeline_service_client, -) -from google.cloud.aiplatform_v1.types import ( +from google.cloud.aiplatform.compat.types import ( dataset as gca_dataset, encryption_spec as gca_encryption_spec, model as gca_model, @@ -215,6 +213,7 @@ def mock_model(): yield model +@pytest.mark.usefixtures("google_auth_mock") class TestAutoMLImageTrainingJob: def setup_method(self): importlib.reload(initializer) diff --git a/tests/unit/aiplatform/test_automl_tabular_training_jobs.py b/tests/unit/aiplatform/test_automl_tabular_training_jobs.py index 3119793c14..c72f664e1d 100644 --- a/tests/unit/aiplatform/test_automl_tabular_training_jobs.py +++ b/tests/unit/aiplatform/test_automl_tabular_training_jobs.py @@ -26,13 +26,12 @@ from google.cloud.aiplatform import schema from google.cloud.aiplatform import training_jobs -from google.cloud.aiplatform_v1.services.model_service import ( - client as model_service_client, +from google.cloud.aiplatform.compat.services import ( + model_service_client, + pipeline_service_client, ) -from google.cloud.aiplatform_v1.services.pipeline_service import ( - client as pipeline_service_client, -) -from google.cloud.aiplatform_v1.types import ( + +from google.cloud.aiplatform.compat.types import ( dataset as gca_dataset, encryption_spec as gca_encryption_spec, model as gca_model, @@ -315,6 +314,7 @@ def mock_dataset_nontabular(): return ds +@pytest.mark.usefixtures("google_auth_mock") class TestAutoMLTabularTrainingJob: def setup_method(self): importlib.reload(initializer) diff --git a/tests/unit/aiplatform/test_automl_text_training_jobs.py b/tests/unit/aiplatform/test_automl_text_training_jobs.py index 7712c758fa..e9c0f8cc07 100644 --- a/tests/unit/aiplatform/test_automl_text_training_jobs.py +++ b/tests/unit/aiplatform/test_automl_text_training_jobs.py @@ -27,13 +27,11 @@ from google.cloud.aiplatform import schema from google.cloud.aiplatform import training_jobs -from google.cloud.aiplatform_v1.services.model_service import ( - client as model_service_client, +from google.cloud.aiplatform.compat.services import ( + model_service_client, + pipeline_service_client, ) -from google.cloud.aiplatform_v1.services.pipeline_service import ( - client as pipeline_service_client, -) -from google.cloud.aiplatform_v1.types import ( +from google.cloud.aiplatform.compat.types import ( dataset as gca_dataset, encryption_spec as gca_encryption_spec, model as gca_model, @@ -201,6 +199,7 @@ def mock_model(): yield model +@pytest.mark.usefixtures("google_auth_mock") class TestAutoMLTextTrainingJob: def setup_method(self): importlib.reload(initializer) diff --git a/tests/unit/aiplatform/test_automl_video_training_jobs.py b/tests/unit/aiplatform/test_automl_video_training_jobs.py index b8abf02058..3468704204 100644 --- a/tests/unit/aiplatform/test_automl_video_training_jobs.py +++ b/tests/unit/aiplatform/test_automl_video_training_jobs.py @@ -30,13 +30,11 @@ from google.cloud.aiplatform import schema from google.cloud.aiplatform import training_jobs -from google.cloud.aiplatform_v1.services.model_service import ( - client as model_service_client, +from google.cloud.aiplatform.compat.services import ( + model_service_client, + pipeline_service_client, ) -from google.cloud.aiplatform_v1.services.pipeline_service import ( - client as pipeline_service_client, -) -from google.cloud.aiplatform_v1.types import ( +from google.cloud.aiplatform.compat.types import ( dataset as gca_dataset, encryption_spec as gca_encryption_spec, model as gca_model, @@ -199,6 +197,7 @@ def mock_model(): yield model +@pytest.mark.usefixtures("google_auth_mock") class TestAutoMLVideoTrainingJob: def setup_method(self): importlib.reload(initializer) diff --git a/tests/unit/aiplatform/test_custom_job.py b/tests/unit/aiplatform/test_custom_job.py index c0062a65e9..c09df26ad2 100644 --- a/tests/unit/aiplatform/test_custom_job.py +++ b/tests/unit/aiplatform/test_custom_job.py @@ -36,7 +36,7 @@ from google.cloud.aiplatform.compat.types import ( encryption_spec as gca_encryption_spec_compat, ) -from google.cloud.aiplatform_v1.services.job_service import client as job_service_client +from google.cloud.aiplatform.compat.services import job_service_client _TEST_PROJECT = "test-project" _TEST_LOCATION = "us-central1" @@ -265,6 +265,7 @@ def create_custom_job_mock_fail(): yield create_custom_job_mock +@pytest.mark.usefixtures("google_auth_mock") class TestCustomJob: def setup_method(self): reload(aiplatform.initializer) diff --git a/tests/unit/aiplatform/test_datasets.py b/tests/unit/aiplatform/test_datasets.py index 13ef13aebd..0624264e4c 100644 --- a/tests/unit/aiplatform/test_datasets.py +++ b/tests/unit/aiplatform/test_datasets.py @@ -15,7 +15,6 @@ # limitations under the License. # -import os import pandas as pd @@ -26,7 +25,6 @@ from unittest.mock import patch from google.api_core import operation -from google.auth.exceptions import GoogleAuthError from google.auth import credentials as auth_credentials from google.cloud import aiplatform @@ -39,11 +37,9 @@ from google.cloud import bigquery from google.cloud import storage -from google.cloud.aiplatform_v1.services.dataset_service import ( - client as dataset_service_client, -) +from google.cloud.aiplatform.compat.services import dataset_service_client -from google.cloud.aiplatform_v1.types import ( +from google.cloud.aiplatform.compat.types import ( dataset as gca_dataset, dataset_service as gca_dataset_service, encryption_spec as gca_encryption_spec, @@ -532,6 +528,7 @@ def bigquery_table_schema_mock(): # TODO(b/171333554): Move reusable test fixtures to conftest.py file +@pytest.mark.usefixtures("google_auth_mock") class TestDataset: def setup_method(self): reload(initializer) @@ -608,17 +605,6 @@ def test_init_dataset_with_id_only(self, get_dataset_mock): name=_TEST_NAME, retry=base._DEFAULT_RETRY ) - @pytest.mark.usefixtures("get_dataset_without_name_mock") - @patch.dict( - os.environ, {"GOOGLE_CLOUD_PROJECT": "", "GOOGLE_APPLICATION_CREDENTIALS": ""} - ) - def test_init_dataset_with_id_only_without_project_or_location(self): - with pytest.raises(GoogleAuthError): - datasets._Dataset( - dataset_name=_TEST_ID, - credentials=auth_credentials.AnonymousCredentials(), - ) - def test_init_dataset_with_location_override(self, get_dataset_mock): aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) datasets._Dataset(dataset_name=_TEST_ID, location=_TEST_ALT_LOCATION) @@ -1011,6 +997,7 @@ def test_delete_dataset(self, delete_dataset_mock, sync): delete_dataset_mock.assert_called_once_with(name=my_dataset.resource_name) +@pytest.mark.usefixtures("google_auth_mock") class TestImageDataset: def setup_method(self): reload(initializer) @@ -1226,6 +1213,7 @@ def test_create_dataset_with_labels(self, create_dataset_mock, sync): ) +@pytest.mark.usefixtures("google_auth_mock") class TestTabularDataset: def setup_method(self): reload(initializer) diff --git a/tests/unit/aiplatform/test_end_to_end.py b/tests/unit/aiplatform/test_end_to_end.py index c31b17ab1c..5f42b7628d 100644 --- a/tests/unit/aiplatform/test_end_to_end.py +++ b/tests/unit/aiplatform/test_end_to_end.py @@ -25,7 +25,7 @@ from google.cloud.aiplatform import models from google.cloud.aiplatform import schema -from google.cloud.aiplatform_v1.types import ( +from google.cloud.aiplatform.compat.types import ( dataset as gca_dataset, encryption_spec as gca_encryption_spec, io as gca_io, @@ -66,6 +66,7 @@ ) +@pytest.mark.usefixtures("google_auth_mock") class TestEndToEnd: def setup_method(self): reload(initializer) diff --git a/tests/unit/aiplatform/test_endpoints.py b/tests/unit/aiplatform/test_endpoints.py index 23c4781406..3e79328b34 100644 --- a/tests/unit/aiplatform/test_endpoints.py +++ b/tests/unit/aiplatform/test_endpoints.py @@ -25,6 +25,8 @@ from google.api_core import operation as ga_operation from google.auth import credentials as auth_credentials +from google.protobuf import field_mask_pb2 + from google.cloud import aiplatform from google.cloud.aiplatform import base from google.cloud.aiplatform import initializer @@ -32,14 +34,10 @@ from google.cloud.aiplatform import models from google.cloud.aiplatform import utils -from google.cloud.aiplatform_v1.services.model_service import ( - client as model_service_client, -) -from google.cloud.aiplatform_v1.services.endpoint_service import ( - client as endpoint_service_client, -) -from google.cloud.aiplatform_v1.services.prediction_service import ( - client as prediction_service_client, +from google.cloud.aiplatform.compat.services import ( + model_service_client, + endpoint_service_client, + prediction_service_client, ) from google.cloud.aiplatform.compat.types import ( endpoint as gca_endpoint, @@ -62,6 +60,8 @@ _TEST_ID_2 = "4366591682456584192" _TEST_ID_3 = "5820582938582924817" _TEST_DESCRIPTION = "test-description" +_TEST_REQUEST_METADATA = () +_TEST_TIMEOUT = None _TEST_ENDPOINT_NAME = ( f"projects/{_TEST_PROJECT}/locations/{_TEST_LOCATION}/endpoints/{_TEST_ID}" @@ -107,6 +107,13 @@ _TEST_ACCELERATOR_TYPE = "NVIDIA_TESLA_P100" _TEST_ACCELERATOR_COUNT = 2 +_TEST_METRIC_NAME_CPU_UTILIZATION = ( + "aiplatform.googleapis.com/prediction/online/cpu/utilization" +) +_TEST_METRIC_NAME_GPU_UTILIZATION = ( + "aiplatform.googleapis.com/prediction/online/accelerator/duty_cycle" +) + _TEST_EXPLANATIONS = [gca_prediction_service.explanation.Explanation(attributions=[])] _TEST_ATTRIBUTIONS = [ @@ -267,6 +274,16 @@ def create_endpoint_mock(): yield create_endpoint_mock +@pytest.fixture +def update_endpoint_mock(): + with mock.patch.object( + endpoint_service_client.EndpointServiceClient, "update_endpoint" + ) as update_endpoint_mock: + update_endpoint_lro_mock = mock.Mock(ga_operation.Operation) + update_endpoint_mock.return_value = update_endpoint_lro_mock + yield update_endpoint_mock + + @pytest.fixture def deploy_model_mock(): with mock.patch.object( @@ -401,6 +418,7 @@ def predict_client_explain_mock(): yield predict_mock +@pytest.mark.usefixtures("google_auth_mock") class TestEndpoint: def setup_method(self): reload(initializer) @@ -722,6 +740,54 @@ def test_create_with_labels(self, create_endpoint_mock, sync): timeout=None, ) + @pytest.mark.usefixtures("get_endpoint_mock") + def test_update_endpoint(self, update_endpoint_mock): + endpoint = models.Endpoint(_TEST_ENDPOINT_NAME) + endpoint.update( + display_name=_TEST_DISPLAY_NAME, + description=_TEST_DESCRIPTION, + labels=_TEST_LABELS, + ) + + expected_endpoint = gca_endpoint.Endpoint( + name=_TEST_ENDPOINT_NAME, + display_name=_TEST_DISPLAY_NAME, + description=_TEST_DESCRIPTION, + labels=_TEST_LABELS, + encryption_spec=_TEST_ENCRYPTION_SPEC, + ) + + expected_update_mask = field_mask_pb2.FieldMask( + paths=["display_name", "description", "labels"] + ) + + update_endpoint_mock.assert_called_once_with( + endpoint=expected_endpoint, + update_mask=expected_update_mask, + metadata=_TEST_REQUEST_METADATA, + timeout=_TEST_TIMEOUT, + ) + + @pytest.mark.usefixtures("get_endpoint_with_models_mock") + def test_update_traffic_split(self, update_endpoint_mock): + endpoint = models.Endpoint(_TEST_ENDPOINT_NAME) + endpoint.update(traffic_split={_TEST_ID: 10, _TEST_ID_2: 80, _TEST_ID_3: 10}) + + expected_endpoint = gca_endpoint.Endpoint( + name=_TEST_ENDPOINT_NAME, + display_name=_TEST_DISPLAY_NAME, + deployed_models=_TEST_DEPLOYED_MODELS, + traffic_split={_TEST_ID: 10, _TEST_ID_2: 80, _TEST_ID_3: 10}, + ) + expected_update_mask = field_mask_pb2.FieldMask(paths=["traffic_split"]) + + update_endpoint_mock.assert_called_once_with( + endpoint=expected_endpoint, + update_mask=expected_update_mask, + metadata=_TEST_REQUEST_METADATA, + timeout=_TEST_TIMEOUT, + ) + @pytest.mark.usefixtures("get_endpoint_mock", "get_model_mock") @pytest.mark.parametrize("sync", [True, False]) def test_deploy(self, deploy_model_mock, sync): @@ -916,7 +982,7 @@ def test_deploy_raise_error_max_replica(self, sync): ) test_endpoint.deploy(model=test_model, max_replica_count=-2, sync=sync) - @pytest.mark.usefixtures("get_endpoint_mock", "get_model_mock") + @pytest.mark.usefixtures("get_endpoint_with_models_mock", "get_model_mock") @pytest.mark.parametrize("sync", [True, False]) def test_deploy_raise_error_traffic_split(self, sync): with pytest.raises(ValueError): @@ -969,48 +1035,39 @@ def test_deploy_with_traffic_percent(self, deploy_model_mock, sync): timeout=None, ) - @pytest.mark.usefixtures("get_model_mock") + @pytest.mark.usefixtures("get_endpoint_with_models_mock", "get_model_mock") @pytest.mark.parametrize("sync", [True, False]) def test_deploy_with_traffic_split(self, deploy_model_mock, sync): - with mock.patch.object( - endpoint_service_client.EndpointServiceClient, "get_endpoint" - ) as get_endpoint_mock: - get_endpoint_mock.return_value = gca_endpoint.Endpoint( - display_name=_TEST_DISPLAY_NAME, - name=_TEST_ENDPOINT_NAME, - traffic_split={"model1": 100}, - ) - - test_endpoint = models.Endpoint(_TEST_ENDPOINT_NAME) - test_model = models.Model(_TEST_ID) - test_model._gca_resource.supported_deployment_resources_types.append( - aiplatform.gapic.Model.DeploymentResourcesType.AUTOMATIC_RESOURCES - ) - test_endpoint.deploy( - model=test_model, - traffic_split={"model1": 30, "0": 70}, - sync=sync, - deploy_request_timeout=None, - ) + test_endpoint = models.Endpoint(_TEST_ENDPOINT_NAME) + test_model = models.Model(_TEST_ID) + test_model._gca_resource.supported_deployment_resources_types.append( + aiplatform.gapic.Model.DeploymentResourcesType.AUTOMATIC_RESOURCES + ) + test_endpoint.deploy( + model=test_model, + traffic_split={_TEST_ID: 10, _TEST_ID_2: 40, _TEST_ID_3: 10, "0": 40}, + sync=sync, + deploy_request_timeout=None, + ) - if not sync: - test_endpoint.wait() - automatic_resources = gca_machine_resources.AutomaticResources( - min_replica_count=1, - max_replica_count=1, - ) - deployed_model = gca_endpoint.DeployedModel( - automatic_resources=automatic_resources, - model=test_model.resource_name, - display_name=None, - ) - deploy_model_mock.assert_called_once_with( - endpoint=test_endpoint.resource_name, - deployed_model=deployed_model, - traffic_split={"model1": 30, "0": 70}, - metadata=(), - timeout=None, - ) + if not sync: + test_endpoint.wait() + automatic_resources = gca_machine_resources.AutomaticResources( + min_replica_count=1, + max_replica_count=1, + ) + deployed_model = gca_endpoint.DeployedModel( + automatic_resources=automatic_resources, + model=test_model.resource_name, + display_name=None, + ) + deploy_model_mock.assert_called_once_with( + endpoint=test_endpoint.resource_name, + deployed_model=deployed_model, + traffic_split={_TEST_ID: 10, _TEST_ID_2: 40, _TEST_ID_3: 10, "0": 40}, + metadata=(), + timeout=None, + ) @pytest.mark.usefixtures("get_endpoint_mock", "get_model_mock") @pytest.mark.parametrize("sync", [True, False]) @@ -1057,6 +1114,138 @@ def test_deploy_with_dedicated_resources(self, deploy_model_mock, sync): timeout=None, ) + @pytest.mark.usefixtures("get_endpoint_mock", "get_model_mock") + @pytest.mark.parametrize("sync", [True, False]) + def test_deploy_with_autoscaling_target_cpu_utilization( + self, deploy_model_mock, sync + ): + test_endpoint = models.Endpoint(_TEST_ENDPOINT_NAME) + test_model = models.Model(_TEST_ID) + test_model._gca_resource.supported_deployment_resources_types.append( + aiplatform.gapic.Model.DeploymentResourcesType.DEDICATED_RESOURCES + ) + test_endpoint.deploy( + model=test_model, + machine_type=_TEST_MACHINE_TYPE, + service_account=_TEST_SERVICE_ACCOUNT, + sync=sync, + deploy_request_timeout=None, + autoscaling_target_cpu_utilization=70, + ) + + if not sync: + test_endpoint.wait() + + expected_machine_spec = gca_machine_resources.MachineSpec( + machine_type=_TEST_MACHINE_TYPE, + ) + + expected_autoscaling_metric_spec = gca_machine_resources.AutoscalingMetricSpec( + metric_name=_TEST_METRIC_NAME_CPU_UTILIZATION, + target=70, + ) + + expected_dedicated_resources = gca_machine_resources.DedicatedResources( + machine_spec=expected_machine_spec, + min_replica_count=1, + max_replica_count=1, + ) + expected_dedicated_resources.autoscaling_metric_specs.extend( + [expected_autoscaling_metric_spec] + ) + + expected_deployed_model = gca_endpoint.DeployedModel( + dedicated_resources=expected_dedicated_resources, + model=test_model.resource_name, + display_name=None, + service_account=_TEST_SERVICE_ACCOUNT, + ) + deploy_model_mock.assert_called_once_with( + endpoint=test_endpoint.resource_name, + deployed_model=expected_deployed_model, + traffic_split={"0": 100}, + metadata=(), + timeout=None, + ) + + @pytest.mark.usefixtures("get_endpoint_mock", "get_model_mock") + @pytest.mark.parametrize("sync", [True, False]) + def test_deploy_with_autoscaling_target_accelerator_duty_cycle( + self, deploy_model_mock, sync + ): + test_endpoint = models.Endpoint(_TEST_ENDPOINT_NAME) + test_model = models.Model(_TEST_ID) + test_model._gca_resource.supported_deployment_resources_types.append( + aiplatform.gapic.Model.DeploymentResourcesType.DEDICATED_RESOURCES + ) + test_endpoint.deploy( + model=test_model, + machine_type=_TEST_MACHINE_TYPE, + accelerator_type=_TEST_ACCELERATOR_TYPE, + accelerator_count=_TEST_ACCELERATOR_COUNT, + service_account=_TEST_SERVICE_ACCOUNT, + sync=sync, + deploy_request_timeout=None, + autoscaling_target_accelerator_duty_cycle=70, + ) + + if not sync: + test_endpoint.wait() + + expected_machine_spec = gca_machine_resources.MachineSpec( + machine_type=_TEST_MACHINE_TYPE, + accelerator_type=_TEST_ACCELERATOR_TYPE, + accelerator_count=_TEST_ACCELERATOR_COUNT, + ) + + expected_autoscaling_metric_spec = gca_machine_resources.AutoscalingMetricSpec( + metric_name=_TEST_METRIC_NAME_GPU_UTILIZATION, + target=70, + ) + + expected_dedicated_resources = gca_machine_resources.DedicatedResources( + machine_spec=expected_machine_spec, + min_replica_count=1, + max_replica_count=1, + ) + expected_dedicated_resources.autoscaling_metric_specs.extend( + [expected_autoscaling_metric_spec] + ) + + expected_deployed_model = gca_endpoint.DeployedModel( + dedicated_resources=expected_dedicated_resources, + model=test_model.resource_name, + display_name=None, + service_account=_TEST_SERVICE_ACCOUNT, + ) + deploy_model_mock.assert_called_once_with( + endpoint=test_endpoint.resource_name, + deployed_model=expected_deployed_model, + traffic_split={"0": 100}, + metadata=(), + timeout=None, + ) + + @pytest.mark.usefixtures("get_endpoint_mock", "get_model_mock") + @pytest.mark.parametrize("sync", [True, False]) + def test_deploy_with_autoscaling_target_accelerator_duty_cycle_and_no_accelerator_type_or_count_raises( + self, sync + ): + with pytest.raises(ValueError): + test_endpoint = models.Endpoint(_TEST_ENDPOINT_NAME) + test_model = models.Model(_TEST_ID) + test_model._gca_resource.supported_deployment_resources_types.append( + aiplatform.gapic.Model.DeploymentResourcesType.DEDICATED_RESOURCES + ) + test_endpoint.deploy( + model=test_model, + sync=sync, + autoscaling_target_accelerator_duty_cycle=70, + ) + + if not sync: + test_endpoint.wait() + @pytest.mark.usefixtures("get_endpoint_mock", "get_model_mock") @pytest.mark.parametrize("sync", [True, False]) def test_deploy_with_explanations(self, deploy_model_with_explanations_mock, sync): diff --git a/tests/unit/aiplatform/test_explain_lit.py b/tests/unit/aiplatform/test_explain_lit.py index c4cc538868..e5a03c2a2e 100644 --- a/tests/unit/aiplatform/test_explain_lit.py +++ b/tests/unit/aiplatform/test_explain_lit.py @@ -36,11 +36,9 @@ open_lit, set_up_and_open_lit, ) -from google.cloud.aiplatform_v1.services.endpoint_service import ( - client as endpoint_service_client, -) -from google.cloud.aiplatform_v1.services.prediction_service import ( - client as prediction_service_client, +from google.cloud.aiplatform.compat.services import ( + endpoint_service_client, + prediction_service_client, ) from importlib import reload from lit_nlp.api import types as lit_types diff --git a/tests/unit/aiplatform/test_featurestores.py b/tests/unit/aiplatform/test_featurestores.py index 3691189325..66f7e9706c 100644 --- a/tests/unit/aiplatform/test_featurestores.py +++ b/tests/unit/aiplatform/test_featurestores.py @@ -35,13 +35,11 @@ from google.cloud.aiplatform.utils import resource_manager_utils from google.cloud.aiplatform.utils import featurestore_utils -from google.cloud.aiplatform_v1.services.featurestore_service import ( - client as featurestore_service_client, +from google.cloud.aiplatform.compat.services import ( + featurestore_service_client, + featurestore_online_serving_service_client, ) -from google.cloud.aiplatform_v1.services.featurestore_online_serving_service import ( - client as featurestore_online_serving_service_client, -) -from google.cloud.aiplatform_v1.types import ( +from google.cloud.aiplatform.compat.types import ( encryption_spec as gca_encryption_spec, entity_type as gca_entity_type, feature as gca_feature, @@ -704,6 +702,7 @@ def batch_create_features_mock(): yield batch_create_features_mock +@pytest.mark.usefixtures("google_auth_mock") class TestFeaturestoreUtils: @pytest.mark.parametrize( "resource_id", diff --git a/tests/unit/aiplatform/test_hyperparameter_tuning_job.py b/tests/unit/aiplatform/test_hyperparameter_tuning_job.py index 727f106fb5..30a2ea40be 100644 --- a/tests/unit/aiplatform/test_hyperparameter_tuning_job.py +++ b/tests/unit/aiplatform/test_hyperparameter_tuning_job.py @@ -33,7 +33,7 @@ job_state as gca_job_state_compat, study as gca_study_compat, ) -from google.cloud.aiplatform_v1.services.job_service import client as job_service_client +from google.cloud.aiplatform.compat.services import job_service_client import test_custom_job @@ -347,6 +347,7 @@ def create_hyperparameter_tuning_job_mock_with_tensorboard(): yield create_hyperparameter_tuning_job_mock +@pytest.mark.usefixtures("google_auth_mock") class TestHyperparameterTuningJob: def setup_method(self): reload(aiplatform.initializer) diff --git a/tests/unit/aiplatform/test_initializer.py b/tests/unit/aiplatform/test_initializer.py index d7e3e3ad8c..6a31a316e1 100644 --- a/tests/unit/aiplatform/test_initializer.py +++ b/tests/unit/aiplatform/test_initializer.py @@ -30,8 +30,8 @@ from google.cloud.aiplatform import utils from google.cloud.aiplatform.utils import resource_manager_utils -from google.cloud.aiplatform_v1.services.model_service import ( - client as model_service_client, +from google.cloud.aiplatform.compat.services import ( + model_service_client, ) _TEST_PROJECT = "test-project" @@ -44,6 +44,7 @@ _TEST_STAGING_BUCKET = "test-bucket" +@pytest.mark.usefixtures("google_auth_mock") class TestInit: def setup_method(self): importlib.reload(initializer) diff --git a/tests/unit/aiplatform/test_jobs.py b/tests/unit/aiplatform/test_jobs.py index 73a4f8da0c..364d0ef17b 100644 --- a/tests/unit/aiplatform/test_jobs.py +++ b/tests/unit/aiplatform/test_jobs.py @@ -40,7 +40,9 @@ manual_batch_tuning_parameters as gca_manual_batch_tuning_parameters_compat, ) -from google.cloud.aiplatform_v1.services.job_service import client as job_service_client +from google.cloud.aiplatform.compat.services import ( + job_service_client, +) _TEST_API_CLIENT = job_service_client.JobServiceClient @@ -178,6 +180,7 @@ def fake_job_cancel_mock(): yield fake_job_cancel_mock +@pytest.mark.usefixtures("google_auth_mock") class TestJob: class FakeJob(jobs._Job): _job_type = "custom-job" @@ -394,6 +397,7 @@ def bq_list_rows_mock(): yield list_rows_mock +@pytest.mark.usefixtures("google_auth_mock") class TestBatchPredictionJob: def setup_method(self): reload(initializer) diff --git a/tests/unit/aiplatform/test_matching_engine_index.py b/tests/unit/aiplatform/test_matching_engine_index.py index bf4c3d1232..78e3a608a9 100644 --- a/tests/unit/aiplatform/test_matching_engine_index.py +++ b/tests/unit/aiplatform/test_matching_engine_index.py @@ -28,11 +28,11 @@ from google.cloud import aiplatform from google.cloud.aiplatform import base from google.cloud.aiplatform import initializer -from google.cloud.aiplatform_v1.services.index_service import ( - client as index_service_client, +from google.cloud.aiplatform.compat.services import ( + index_service_client, ) -from google.cloud.aiplatform_v1.types import index as gca_index +from google.cloud.aiplatform.compat.types import index as gca_index # project _TEST_PROJECT = "test-project" @@ -167,6 +167,7 @@ def create_index_mock(): yield create_index_mock +@pytest.mark.usefixtures("google_auth_mock") class TestMatchingEngineIndex: def setup_method(self): reload(initializer) diff --git a/tests/unit/aiplatform/test_matching_engine_index_endpoint.py b/tests/unit/aiplatform/test_matching_engine_index_endpoint.py index cd3dd78ed2..58ff16ed56 100644 --- a/tests/unit/aiplatform/test_matching_engine_index_endpoint.py +++ b/tests/unit/aiplatform/test_matching_engine_index_endpoint.py @@ -26,18 +26,15 @@ from google.cloud.aiplatform import initializer from google.cloud.aiplatform.compat.types import ( matching_engine_deployed_index_ref as gca_matching_engine_deployed_index_ref, - matching_engine_index_endpoint as gca_matching_engine_index_endpoint, -) -from google.cloud.aiplatform_v1.services.index_endpoint_service import ( - client as index_endpoint_service_client, -) -from google.cloud.aiplatform_v1.services.index_service import ( - client as index_service_client, -) -from google.cloud.aiplatform_v1.types import ( - index as gca_index, index_endpoint as gca_index_endpoint, + index as gca_index, +) + +from google.cloud.aiplatform.compat.services import ( + index_endpoint_service_client, + index_service_client, ) + from google.protobuf import field_mask_pb2 import pytest @@ -254,7 +251,7 @@ def get_index_endpoint_mock(): description=_TEST_INDEX_ENDPOINT_DESCRIPTION, ) index_endpoint.deployed_indexes = [ - gca_matching_engine_index_endpoint.DeployedIndex( + gca_index_endpoint.DeployedIndex( id=_TEST_DEPLOYED_INDEX_ID, index=_TEST_INDEX_NAME, display_name=_TEST_DEPLOYED_INDEX_DISPLAY_NAME, @@ -265,14 +262,14 @@ def get_index_endpoint_mock(): "min_replica_count": _TEST_MIN_REPLICA_COUNT, "max_replica_count": _TEST_MAX_REPLICA_COUNT, }, - deployed_index_auth_config=gca_matching_engine_index_endpoint.DeployedIndexAuthConfig( - auth_provider=gca_matching_engine_index_endpoint.DeployedIndexAuthConfig.AuthProvider( + deployed_index_auth_config=gca_index_endpoint.DeployedIndexAuthConfig( + auth_provider=gca_index_endpoint.DeployedIndexAuthConfig.AuthProvider( audiences=_TEST_AUTH_CONFIG_AUDIENCES, allowed_issuers=_TEST_AUTH_CONFIG_ALLOWED_ISSUERS, ) ), ), - gca_matching_engine_index_endpoint.DeployedIndex( + gca_index_endpoint.DeployedIndex( id=f"{_TEST_DEPLOYED_INDEX_ID}_2", index=f"{_TEST_INDEX_NAME}_2", display_name=_TEST_DEPLOYED_INDEX_DISPLAY_NAME, @@ -283,8 +280,8 @@ def get_index_endpoint_mock(): "min_replica_count": _TEST_MIN_REPLICA_COUNT, "max_replica_count": _TEST_MAX_REPLICA_COUNT, }, - deployed_index_auth_config=gca_matching_engine_index_endpoint.DeployedIndexAuthConfig( - auth_provider=gca_matching_engine_index_endpoint.DeployedIndexAuthConfig.AuthProvider( + deployed_index_auth_config=gca_index_endpoint.DeployedIndexAuthConfig( + auth_provider=gca_index_endpoint.DeployedIndexAuthConfig.AuthProvider( audiences=_TEST_AUTH_CONFIG_AUDIENCES, allowed_issuers=_TEST_AUTH_CONFIG_ALLOWED_ISSUERS, ) @@ -383,6 +380,7 @@ def create_index_endpoint_mock(): yield create_index_endpoint_mock +@pytest.mark.usefixtures("google_auth_mock") class TestMatchingEngineIndexEndpoint: def setup_method(self): reload(initializer) @@ -519,7 +517,7 @@ def test_deploy_index(self, deploy_index_mock, undeploy_index_mock): deploy_index_mock.assert_called_once_with( index_endpoint=my_index_endpoint.resource_name, - deployed_index=gca_matching_engine_index_endpoint.DeployedIndex( + deployed_index=gca_index_endpoint.DeployedIndex( id=_TEST_DEPLOYED_INDEX_ID, index=my_index.resource_name, display_name=_TEST_DEPLOYED_INDEX_DISPLAY_NAME, @@ -530,8 +528,8 @@ def test_deploy_index(self, deploy_index_mock, undeploy_index_mock): "min_replica_count": _TEST_MIN_REPLICA_COUNT, "max_replica_count": _TEST_MAX_REPLICA_COUNT, }, - deployed_index_auth_config=gca_matching_engine_index_endpoint.DeployedIndexAuthConfig( - auth_provider=gca_matching_engine_index_endpoint.DeployedIndexAuthConfig.AuthProvider( + deployed_index_auth_config=gca_index_endpoint.DeployedIndexAuthConfig( + auth_provider=gca_index_endpoint.DeployedIndexAuthConfig.AuthProvider( audiences=_TEST_AUTH_CONFIG_AUDIENCES, allowed_issuers=_TEST_AUTH_CONFIG_ALLOWED_ISSUERS, ) @@ -567,7 +565,7 @@ def test_mutate_deployed_index(self, mutate_deployed_index_mock): mutate_deployed_index_mock.assert_called_once_with( index_endpoint=_TEST_INDEX_ENDPOINT_NAME, - deployed_index=gca_matching_engine_index_endpoint.DeployedIndex( + deployed_index=gca_index_endpoint.DeployedIndex( id=_TEST_DEPLOYED_INDEX_ID, automatic_resources={ "min_replica_count": _TEST_MIN_REPLICA_COUNT_UPDATED, diff --git a/tests/unit/aiplatform/test_metadata.py b/tests/unit/aiplatform/test_metadata.py index 84ad3949ed..acbfa9098b 100644 --- a/tests/unit/aiplatform/test_metadata.py +++ b/tests/unit/aiplatform/test_metadata.py @@ -377,6 +377,7 @@ def _assert_frame_equal_with_sorted_columns(dataframe_1, dataframe_2): ) +@pytest.mark.usefixtures("google_auth_mock") class TestMetadata: def setup_method(self): reload(initializer) diff --git a/tests/unit/aiplatform/test_metadata_resources.py b/tests/unit/aiplatform/test_metadata_resources.py index c4ca4b74d9..2a7180adfd 100644 --- a/tests/unit/aiplatform/test_metadata_resources.py +++ b/tests/unit/aiplatform/test_metadata_resources.py @@ -356,6 +356,7 @@ def update_artifact_mock(): yield update_artifact_mock +@pytest.mark.usefixtures("google_auth_mock") class TestContext: def setup_method(self): reload(initializer) diff --git a/tests/unit/aiplatform/test_metadata_store.py b/tests/unit/aiplatform/test_metadata_store.py index 658972f689..b6dfe8e032 100644 --- a/tests/unit/aiplatform/test_metadata_store.py +++ b/tests/unit/aiplatform/test_metadata_store.py @@ -15,15 +15,12 @@ # limitations under the License. # -import os from importlib import reload from unittest import mock from unittest.mock import patch import pytest from google.api_core import operation -from google.auth import credentials as auth_credentials -from google.auth.exceptions import GoogleAuthError from google.cloud import aiplatform from google.cloud.aiplatform import base @@ -31,8 +28,8 @@ from google.cloud.aiplatform.metadata import metadata_store from google.cloud.aiplatform_v1 import MetadataServiceClient from google.cloud.aiplatform_v1 import MetadataStore as GapicMetadataStore -from google.cloud.aiplatform_v1.types import encryption_spec as gca_encryption_spec -from google.cloud.aiplatform_v1.types import metadata_service +from google.cloud.aiplatform.compat.types import encryption_spec as gca_encryption_spec +from google.cloud.aiplatform.compat.types import metadata_service # project _TEST_PROJECT = "test-project" @@ -145,6 +142,7 @@ def setup_method(self): def teardown_method(self): initializer.global_pool.shutdown(wait=True) + @pytest.mark.usefixtures("google_auth_mock") def test_init_metadata_store(self, get_metadata_store_mock): aiplatform.init(project=_TEST_PROJECT) metadata_store._MetadataStore(metadata_store_name=_TEST_NAME) @@ -166,17 +164,6 @@ def test_init_metadata_store_with_default_id(self, get_metadata_store_mock): name=_TEST_DEFAULT_NAME, retry=base._DEFAULT_RETRY ) - @pytest.mark.usefixtures("get_metadata_store_without_name_mock") - @patch.dict( - os.environ, {"GOOGLE_CLOUD_PROJECT": "", "GOOGLE_APPLICATION_CREDENTIALS": ""} - ) - def test_init_metadata_store_with_id_without_project_or_location(self): - with pytest.raises(GoogleAuthError): - metadata_store._MetadataStore( - metadata_store_name=_TEST_ID, - credentials=auth_credentials.AnonymousCredentials(), - ) - def test_init_metadata_store_with_location_override(self, get_metadata_store_mock): aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) metadata_store._MetadataStore( diff --git a/tests/unit/aiplatform/test_model_evaluation.py b/tests/unit/aiplatform/test_model_evaluation.py index c5c5cd9ac3..d45e9a1506 100644 --- a/tests/unit/aiplatform/test_model_evaluation.py +++ b/tests/unit/aiplatform/test_model_evaluation.py @@ -23,14 +23,16 @@ from google.cloud.aiplatform import base from google.cloud.aiplatform import models -from google.cloud.aiplatform_v1.services.model_service import ( - client as model_service_client, +from google.cloud.aiplatform.compat.services import ( + model_service_client, ) from google.cloud.aiplatform.compat.types import model as gca_model -from google.cloud.aiplatform_v1.types import model_evaluation as gca_model_evaluation +from google.cloud.aiplatform.compat.types import ( + model_evaluation as gca_model_evaluation, +) _TEST_PROJECT = "test-project" _TEST_LOCATION = "us-central1" @@ -125,6 +127,7 @@ def mock_model_eval_get(): yield mock_get_model_eval +@pytest.mark.usefixtures("google_auth_mock") class TestModelEvaluation: def test_init_model_evaluation_with_only_resource_name(self, mock_model_eval_get): aiplatform.init(project=_TEST_PROJECT) diff --git a/tests/unit/aiplatform/test_models.py b/tests/unit/aiplatform/test_models.py index eaf63d9fdd..23f933128a 100644 --- a/tests/unit/aiplatform/test_models.py +++ b/tests/unit/aiplatform/test_models.py @@ -32,12 +32,10 @@ from google.cloud.aiplatform import models from google.cloud.aiplatform import utils -from google.cloud.aiplatform_v1.services.endpoint_service import ( - client as endpoint_service_client, -) -from google.cloud.aiplatform_v1.services.job_service import client as job_service_client -from google.cloud.aiplatform_v1.services.model_service import ( - client as model_service_client, +from google.cloud.aiplatform.compat.services import ( + endpoint_service_client, + model_service_client, + job_service_client, ) from google.cloud.aiplatform.compat.services import pipeline_service_client from google.cloud.aiplatform.compat.types import ( @@ -543,6 +541,7 @@ def list_model_evaluations_mock(): yield list_model_evaluations_mock +@pytest.mark.usefixtures("google_auth_mock") class TestModel: def setup_method(self): importlib.reload(initializer) diff --git a/tests/unit/aiplatform/test_pipeline_jobs.py b/tests/unit/aiplatform/test_pipeline_jobs.py index df5e294b03..159400f8ce 100644 --- a/tests/unit/aiplatform/test_pipeline_jobs.py +++ b/tests/unit/aiplatform/test_pipeline_jobs.py @@ -32,12 +32,12 @@ from google.cloud import storage from google.protobuf import json_format -from google.cloud.aiplatform_v1.services.pipeline_service import ( - client as pipeline_service_client_v1, +from google.cloud.aiplatform.compat.services import ( + pipeline_service_client, ) -from google.cloud.aiplatform_v1.types import ( - pipeline_job as gca_pipeline_job_v1, - pipeline_state as gca_pipeline_state_v1, +from google.cloud.aiplatform.compat.types import ( + pipeline_job as gca_pipeline_job, + pipeline_state as gca_pipeline_state, ) _TEST_PROJECT = "test-project" @@ -188,11 +188,11 @@ @pytest.fixture def mock_pipeline_service_create(): with mock.patch.object( - pipeline_service_client_v1.PipelineServiceClient, "create_pipeline_job" + pipeline_service_client.PipelineServiceClient, "create_pipeline_job" ) as mock_create_pipeline_job: - mock_create_pipeline_job.return_value = gca_pipeline_job_v1.PipelineJob( + mock_create_pipeline_job.return_value = gca_pipeline_job.PipelineJob( name=_TEST_PIPELINE_JOB_NAME, - state=gca_pipeline_state_v1.PipelineState.PIPELINE_STATE_SUCCEEDED, + state=gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED, create_time=_TEST_PIPELINE_CREATE_TIME, service_account=_TEST_SERVICE_ACCOUNT, network=_TEST_NETWORK, @@ -201,7 +201,7 @@ def mock_pipeline_service_create(): def make_pipeline_job(state): - return gca_pipeline_job_v1.PipelineJob( + return gca_pipeline_job.PipelineJob( name=_TEST_PIPELINE_JOB_NAME, state=state, create_time=_TEST_PIPELINE_CREATE_TIME, @@ -213,35 +213,33 @@ def make_pipeline_job(state): @pytest.fixture def mock_pipeline_service_get(): with mock.patch.object( - pipeline_service_client_v1.PipelineServiceClient, "get_pipeline_job" + pipeline_service_client.PipelineServiceClient, "get_pipeline_job" ) as mock_get_pipeline_job: mock_get_pipeline_job.side_effect = [ + make_pipeline_job(gca_pipeline_state.PipelineState.PIPELINE_STATE_RUNNING), make_pipeline_job( - gca_pipeline_state_v1.PipelineState.PIPELINE_STATE_RUNNING + gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED ), make_pipeline_job( - gca_pipeline_state_v1.PipelineState.PIPELINE_STATE_SUCCEEDED + gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED ), make_pipeline_job( - gca_pipeline_state_v1.PipelineState.PIPELINE_STATE_SUCCEEDED + gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED ), make_pipeline_job( - gca_pipeline_state_v1.PipelineState.PIPELINE_STATE_SUCCEEDED + gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED ), make_pipeline_job( - gca_pipeline_state_v1.PipelineState.PIPELINE_STATE_SUCCEEDED + gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED ), make_pipeline_job( - gca_pipeline_state_v1.PipelineState.PIPELINE_STATE_SUCCEEDED + gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED ), make_pipeline_job( - gca_pipeline_state_v1.PipelineState.PIPELINE_STATE_SUCCEEDED + gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED ), make_pipeline_job( - gca_pipeline_state_v1.PipelineState.PIPELINE_STATE_SUCCEEDED - ), - make_pipeline_job( - gca_pipeline_state_v1.PipelineState.PIPELINE_STATE_SUCCEEDED + gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED ), ] @@ -251,18 +249,12 @@ def mock_pipeline_service_get(): @pytest.fixture def mock_pipeline_service_get_with_fail(): with mock.patch.object( - pipeline_service_client_v1.PipelineServiceClient, "get_pipeline_job" + pipeline_service_client.PipelineServiceClient, "get_pipeline_job" ) as mock_get_pipeline_job: mock_get_pipeline_job.side_effect = [ - make_pipeline_job( - gca_pipeline_state_v1.PipelineState.PIPELINE_STATE_RUNNING - ), - make_pipeline_job( - gca_pipeline_state_v1.PipelineState.PIPELINE_STATE_RUNNING - ), - make_pipeline_job( - gca_pipeline_state_v1.PipelineState.PIPELINE_STATE_FAILED - ), + make_pipeline_job(gca_pipeline_state.PipelineState.PIPELINE_STATE_RUNNING), + make_pipeline_job(gca_pipeline_state.PipelineState.PIPELINE_STATE_RUNNING), + make_pipeline_job(gca_pipeline_state.PipelineState.PIPELINE_STATE_FAILED), ] yield mock_get_pipeline_job @@ -271,7 +263,7 @@ def mock_pipeline_service_get_with_fail(): @pytest.fixture def mock_pipeline_service_cancel(): with mock.patch.object( - pipeline_service_client_v1.PipelineServiceClient, "cancel_pipeline_job" + pipeline_service_client.PipelineServiceClient, "cancel_pipeline_job" ) as mock_cancel_pipeline_job: yield mock_cancel_pipeline_job @@ -279,7 +271,7 @@ def mock_pipeline_service_cancel(): @pytest.fixture def mock_pipeline_service_list(): with mock.patch.object( - pipeline_service_client_v1.PipelineServiceClient, "list_pipeline_jobs" + pipeline_service_client.PipelineServiceClient, "list_pipeline_jobs" ) as mock_list_pipeline_jobs: yield mock_list_pipeline_jobs @@ -291,6 +283,7 @@ def mock_load_yaml_and_json(job_spec): yield mock_load_yaml_and_json +@pytest.mark.usefixtures("google_auth_mock") class TestPipelineJob: class FakePipelineJob(pipeline_jobs.PipelineJob): @@ -351,14 +344,14 @@ def test_run_call_pipeline_service_create( "gcsOutputDirectory": _TEST_GCS_BUCKET_NAME, "parameterValues": _TEST_PIPELINE_PARAMETER_VALUES, } - runtime_config = gca_pipeline_job_v1.PipelineJob.RuntimeConfig()._pb + runtime_config = gca_pipeline_job.PipelineJob.RuntimeConfig()._pb json_format.ParseDict(expected_runtime_config_dict, runtime_config) job_spec = yaml.safe_load(job_spec) pipeline_spec = job_spec.get("pipelineSpec") or job_spec # Construct expected request - expected_gapic_pipeline_job = gca_pipeline_job_v1.PipelineJob( + expected_gapic_pipeline_job = gca_pipeline_job.PipelineJob( display_name=_TEST_PIPELINE_JOB_DISPLAY_NAME, pipeline_spec={ "components": {}, @@ -383,7 +376,7 @@ def test_run_call_pipeline_service_create( ) assert job._gca_resource == make_pipeline_job( - gca_pipeline_state_v1.PipelineState.PIPELINE_STATE_SUCCEEDED + gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED ) @pytest.mark.parametrize( @@ -432,14 +425,14 @@ def test_run_call_pipeline_service_create_with_timeout( "gcsOutputDirectory": _TEST_GCS_BUCKET_NAME, "parameterValues": _TEST_PIPELINE_PARAMETER_VALUES, } - runtime_config = gca_pipeline_job_v1.PipelineJob.RuntimeConfig()._pb + runtime_config = gca_pipeline_job.PipelineJob.RuntimeConfig()._pb json_format.ParseDict(expected_runtime_config_dict, runtime_config) job_spec = yaml.safe_load(job_spec) pipeline_spec = job_spec.get("pipelineSpec") or job_spec # Construct expected request - expected_gapic_pipeline_job = gca_pipeline_job_v1.PipelineJob( + expected_gapic_pipeline_job = gca_pipeline_job.PipelineJob( display_name=_TEST_PIPELINE_JOB_DISPLAY_NAME, pipeline_spec={ "components": {}, @@ -464,7 +457,7 @@ def test_run_call_pipeline_service_create_with_timeout( # ) # assert job._gca_resource == make_pipeline_job( - # gca_pipeline_state_v1.PipelineState.PIPELINE_STATE_SUCCEEDED + # gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED # ) @pytest.mark.parametrize( @@ -512,14 +505,14 @@ def test_run_call_pipeline_service_create_with_timeout_not_explicitly_set( "gcsOutputDirectory": _TEST_GCS_BUCKET_NAME, "parameterValues": _TEST_PIPELINE_PARAMETER_VALUES, } - runtime_config = gca_pipeline_job_v1.PipelineJob.RuntimeConfig()._pb + runtime_config = gca_pipeline_job.PipelineJob.RuntimeConfig()._pb json_format.ParseDict(expected_runtime_config_dict, runtime_config) job_spec = yaml.safe_load(job_spec) pipeline_spec = job_spec.get("pipelineSpec") or job_spec # Construct expected request - expected_gapic_pipeline_job = gca_pipeline_job_v1.PipelineJob( + expected_gapic_pipeline_job = gca_pipeline_job.PipelineJob( display_name=_TEST_PIPELINE_JOB_DISPLAY_NAME, pipeline_spec={ "components": {}, @@ -585,14 +578,14 @@ def test_run_call_pipeline_service_create_legacy( "gcsOutputDirectory": _TEST_GCS_BUCKET_NAME, "parameters": {"string_param": {"stringValue": "hello"}}, } - runtime_config = gca_pipeline_job_v1.PipelineJob.RuntimeConfig()._pb + runtime_config = gca_pipeline_job.PipelineJob.RuntimeConfig()._pb json_format.ParseDict(expected_runtime_config_dict, runtime_config) job_spec = yaml.safe_load(job_spec) pipeline_spec = job_spec.get("pipelineSpec") or job_spec # Construct expected request - expected_gapic_pipeline_job = gca_pipeline_job_v1.PipelineJob( + expected_gapic_pipeline_job = gca_pipeline_job.PipelineJob( display_name=_TEST_PIPELINE_JOB_DISPLAY_NAME, pipeline_spec={ "components": {}, @@ -617,7 +610,7 @@ def test_run_call_pipeline_service_create_legacy( ) assert job._gca_resource == make_pipeline_job( - gca_pipeline_state_v1.PipelineState.PIPELINE_STATE_SUCCEEDED + gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED ) @pytest.mark.parametrize( @@ -666,14 +659,14 @@ def test_run_call_pipeline_service_create_tfx( "gcsOutputDirectory": _TEST_GCS_BUCKET_NAME, "parameters": {"string_param": {"stringValue": "hello"}}, } - runtime_config = gca_pipeline_job_v1.PipelineJob.RuntimeConfig()._pb + runtime_config = gca_pipeline_job.PipelineJob.RuntimeConfig()._pb json_format.ParseDict(expected_runtime_config_dict, runtime_config) job_spec = yaml.safe_load(job_spec) pipeline_spec = job_spec.get("pipelineSpec") or job_spec # Construct expected request - expected_gapic_pipeline_job = gca_pipeline_job_v1.PipelineJob( + expected_gapic_pipeline_job = gca_pipeline_job.PipelineJob( display_name=_TEST_PIPELINE_JOB_DISPLAY_NAME, pipeline_spec={ "components": {}, @@ -699,7 +692,7 @@ def test_run_call_pipeline_service_create_tfx( ) assert job._gca_resource == make_pipeline_job( - gca_pipeline_state_v1.PipelineState.PIPELINE_STATE_SUCCEEDED + gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED ) @pytest.mark.parametrize( @@ -738,14 +731,14 @@ def test_submit_call_pipeline_service_pipeline_job_create( "gcsOutputDirectory": _TEST_GCS_BUCKET_NAME, "parameterValues": _TEST_PIPELINE_PARAMETER_VALUES, } - runtime_config = gca_pipeline_job_v1.PipelineJob.RuntimeConfig()._pb + runtime_config = gca_pipeline_job.PipelineJob.RuntimeConfig()._pb json_format.ParseDict(expected_runtime_config_dict, runtime_config) job_spec = yaml.safe_load(job_spec) pipeline_spec = job_spec.get("pipelineSpec") or job_spec # Construct expected request - expected_gapic_pipeline_job = gca_pipeline_job_v1.PipelineJob( + expected_gapic_pipeline_job = gca_pipeline_job.PipelineJob( display_name=_TEST_PIPELINE_JOB_DISPLAY_NAME, pipeline_spec={ "components": {}, @@ -774,7 +767,7 @@ def test_submit_call_pipeline_service_pipeline_job_create( ) assert job._gca_resource == make_pipeline_job( - gca_pipeline_state_v1.PipelineState.PIPELINE_STATE_SUCCEEDED + gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED ) @pytest.mark.parametrize( @@ -851,14 +844,14 @@ def test_submit_call_pipeline_service_pipeline_job_create_legacy( "parameters": {"string_param": {"stringValue": "hello"}}, "gcsOutputDirectory": _TEST_GCS_BUCKET_NAME, } - runtime_config = gca_pipeline_job_v1.PipelineJob.RuntimeConfig()._pb + runtime_config = gca_pipeline_job.PipelineJob.RuntimeConfig()._pb json_format.ParseDict(expected_runtime_config_dict, runtime_config) job_spec = yaml.safe_load(job_spec) pipeline_spec = job_spec.get("pipelineSpec") or job_spec # Construct expected request - expected_gapic_pipeline_job = gca_pipeline_job_v1.PipelineJob( + expected_gapic_pipeline_job = gca_pipeline_job.PipelineJob( display_name=_TEST_PIPELINE_JOB_DISPLAY_NAME, pipeline_spec={ "components": {}, @@ -887,7 +880,7 @@ def test_submit_call_pipeline_service_pipeline_job_create_legacy( ) assert job._gca_resource == make_pipeline_job( - gca_pipeline_state_v1.PipelineState.PIPELINE_STATE_SUCCEEDED + gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED ) @pytest.mark.usefixtures("mock_pipeline_service_get") diff --git a/tests/unit/aiplatform/test_tensorboard.py b/tests/unit/aiplatform/test_tensorboard.py index 40b20c7b70..37e3376875 100644 --- a/tests/unit/aiplatform/test_tensorboard.py +++ b/tests/unit/aiplatform/test_tensorboard.py @@ -15,7 +15,6 @@ # limitations under the License. # -import os import pytest @@ -24,19 +23,17 @@ from importlib import reload from google.api_core import operation -from google.auth.exceptions import GoogleAuthError -from google.auth import credentials as auth_credentials from google.cloud import aiplatform from google.cloud.aiplatform import base from google.cloud.aiplatform import initializer from google.cloud.aiplatform import tensorboard -from google.cloud.aiplatform_v1.services.tensorboard_service import ( - client as tensorboard_service_client, +from google.cloud.aiplatform.compat.services import ( + tensorboard_service_client, ) -from google.cloud.aiplatform_v1.types import ( +from google.cloud.aiplatform.compat.types import ( encryption_spec as gca_encryption_spec, tensorboard as gca_tensorboard, tensorboard_experiment as gca_tensorboard_experiment, @@ -266,6 +263,7 @@ def list_tensorboard_run_mock(): yield list_tensorboard_run_mock +@pytest.mark.usefixtures("google_auth_mock") class TestTensorboard: def setup_method(self): reload(initializer) @@ -330,16 +328,6 @@ def test_init_tensorboard_with_project_and_alt_location(self): location=_TEST_ALT_LOCATION, ) - @patch.dict( - os.environ, {"GOOGLE_CLOUD_PROJECT": "", "GOOGLE_APPLICATION_CREDENTIALS": ""} - ) - def test_init_tensorboard_with_id_only_without_project_or_location(self): - with pytest.raises(GoogleAuthError): - tensorboard.Tensorboard( - tensorboard_name=_TEST_ID, - credentials=auth_credentials.AnonymousCredentials(), - ) - def test_init_tensorboard_with_location_override(self, get_tensorboard_mock): aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION) tensorboard.Tensorboard(tensorboard_name=_TEST_ID, location=_TEST_ALT_LOCATION) diff --git a/tests/unit/aiplatform/test_training_jobs.py b/tests/unit/aiplatform/test_training_jobs.py index 21d9baac3e..a956584663 100644 --- a/tests/unit/aiplatform/test_training_jobs.py +++ b/tests/unit/aiplatform/test_training_jobs.py @@ -17,6 +17,7 @@ from distutils import core import copy +import os import functools import importlib import logging @@ -45,16 +46,13 @@ from google.cloud.aiplatform.utils import source_utils from google.cloud.aiplatform.utils import worker_spec_utils - -from google.cloud.aiplatform_v1.services.job_service import client as job_service_client -from google.cloud.aiplatform_v1.services.model_service import ( - client as model_service_client, -) -from google.cloud.aiplatform_v1.services.pipeline_service import ( - client as pipeline_service_client, +from google.cloud.aiplatform.compat.services import ( + model_service_client, + pipeline_service_client, + job_service_client, ) -from google.cloud.aiplatform_v1.types import ( +from google.cloud.aiplatform.compat.types import ( custom_job as gca_custom_job, dataset as gca_dataset, encryption_spec as gca_encryption_spec, @@ -76,7 +74,8 @@ _TEST_GCS_PATH = f"{_TEST_BUCKET_NAME}/{_TEST_GCS_PATH_WITHOUT_BUCKET}" _TEST_GCS_PATH_WITH_TRAILING_SLASH = f"{_TEST_GCS_PATH}/" _TEST_LOCAL_SCRIPT_FILE_NAME = "____test____script.py" -_TEST_LOCAL_SCRIPT_FILE_PATH = f"path/to/{_TEST_LOCAL_SCRIPT_FILE_NAME}" +_TEST_TEMPDIR = tempfile.mkdtemp() +_TEST_LOCAL_SCRIPT_FILE_PATH = os.path.join(_TEST_TEMPDIR, _TEST_LOCAL_SCRIPT_FILE_NAME) _TEST_PYTHON_SOURCE = """ print('hello world') """ @@ -331,6 +330,7 @@ def mock_get_backing_custom_job_with_enable_web_access(): yield get_custom_job_mock +@pytest.mark.usefixtures("google_auth_mock") class TestTrainingScriptPythonPackagerHelpers: def setup_method(self): importlib.reload(initializer) @@ -446,15 +446,16 @@ def test_get_python_executable_returns_python_executable(self): assert "python" in source_utils._get_python_executable().lower() +@pytest.mark.usefixtures("google_auth_mock") class TestTrainingScriptPythonPackager: def setup_method(self): importlib.reload(initializer) importlib.reload(aiplatform) - with open(_TEST_LOCAL_SCRIPT_FILE_NAME, "w") as fp: + with open(_TEST_LOCAL_SCRIPT_FILE_PATH, "w") as fp: fp.write(_TEST_PYTHON_SOURCE) def teardown_method(self): - pathlib.Path(_TEST_LOCAL_SCRIPT_FILE_NAME).unlink() + pathlib.Path(_TEST_LOCAL_SCRIPT_FILE_PATH).unlink() python_package_file = f"{source_utils._TrainingScriptPythonPackager._ROOT_MODULE}-{source_utils._TrainingScriptPythonPackager._SETUP_PY_VERSION}.tar.gz" if pathlib.Path(python_package_file).is_file(): pathlib.Path(python_package_file).unlink() @@ -468,14 +469,14 @@ def teardown_method(self): ) def test_packager_creates_and_copies_python_package(self): - tsp = source_utils._TrainingScriptPythonPackager(_TEST_LOCAL_SCRIPT_FILE_NAME) + tsp = source_utils._TrainingScriptPythonPackager(_TEST_LOCAL_SCRIPT_FILE_PATH) tsp.package_and_copy(copy_method=local_copy_method) assert pathlib.Path( f"{tsp._ROOT_MODULE}-{tsp._SETUP_PY_VERSION}.tar.gz" ).is_file() def test_created_package_module_is_installable_and_can_be_run(self): - tsp = source_utils._TrainingScriptPythonPackager(_TEST_LOCAL_SCRIPT_FILE_NAME) + tsp = source_utils._TrainingScriptPythonPackager(_TEST_LOCAL_SCRIPT_FILE_PATH) source_dist_path = tsp.package_and_copy(copy_method=local_copy_method) subprocess.check_output(["pip3", "install", source_dist_path]) module_output = subprocess.check_output( @@ -485,7 +486,7 @@ def test_created_package_module_is_installable_and_can_be_run(self): def test_requirements_are_in_package(self): tsp = source_utils._TrainingScriptPythonPackager( - _TEST_LOCAL_SCRIPT_FILE_NAME, requirements=_TEST_REQUIREMENTS + _TEST_LOCAL_SCRIPT_FILE_PATH, requirements=_TEST_REQUIREMENTS ) source_dist_path = tsp.package_and_copy(copy_method=local_copy_method) with tarfile.open(source_dist_path) as tf: @@ -504,7 +505,7 @@ def test_packaging_fails_whith_RuntimeError(self): mock_subprocess.returncode = 1 mock_popen.return_value = mock_subprocess tsp = source_utils._TrainingScriptPythonPackager( - _TEST_LOCAL_SCRIPT_FILE_NAME + _TEST_LOCAL_SCRIPT_FILE_PATH ) with pytest.raises(RuntimeError): tsp.package_and_copy(copy_method=local_copy_method) @@ -512,7 +513,7 @@ def test_packaging_fails_whith_RuntimeError(self): def test_package_and_copy_to_gcs_copies_to_gcs(self, mock_client_bucket): mock_client_bucket, mock_blob = mock_client_bucket - tsp = source_utils._TrainingScriptPythonPackager(_TEST_LOCAL_SCRIPT_FILE_NAME) + tsp = source_utils._TrainingScriptPythonPackager(_TEST_LOCAL_SCRIPT_FILE_PATH) gcs_path = tsp.package_and_copy_to_gcs( gcs_staging_dir=_TEST_BUCKET_NAME, project=_TEST_PROJECT @@ -834,11 +835,14 @@ def mock_nontabular_dataset(): return ds +@pytest.mark.usefixtures("google_auth_mock") class TestCustomTrainingJob: def setup_method(self): importlib.reload(initializer) importlib.reload(aiplatform) - self._local_script_file_name = f"{uuid.uuid4()}-{_TEST_LOCAL_SCRIPT_FILE_NAME}" + self._local_script_file_name = os.path.join( + _TEST_TEMPDIR, f"{uuid.uuid4()}-{_TEST_LOCAL_SCRIPT_FILE_NAME}" + ) with open(self._local_script_file_name, "w") as fp: fp.write(_TEST_PYTHON_SOURCE) @@ -2749,6 +2753,7 @@ def test_cancel_training_job_without_running(self, mock_pipeline_service_cancel) assert e.match(regexp=r"TrainingJob has not been launched") +@pytest.mark.usefixtures("google_auth_mock") class TestCustomContainerTrainingJob: def setup_method(self): importlib.reload(initializer) @@ -4708,6 +4713,7 @@ def test_machine_spec_handles_missing_pools(self): assert spec.pool_specs == true_pool_spec +@pytest.mark.usefixtures("google_auth_mock") class TestCustomPythonPackageTrainingJob: def setup_method(self): importlib.reload(initializer) diff --git a/tests/unit/aiplatform/test_uploader.py b/tests/unit/aiplatform/test_uploader.py index 29c435d40f..44cb1bb11a 100644 --- a/tests/unit/aiplatform/test_uploader.py +++ b/tests/unit/aiplatform/test_uploader.py @@ -46,19 +46,20 @@ from google.cloud.aiplatform.tensorboard.plugins.tf_profiler import profile_uploader import google.cloud.aiplatform.tensorboard.uploader as uploader_lib from google.cloud import storage -from google.cloud.aiplatform_v1.services.tensorboard_service import ( - client as tensorboard_service_client, +from google.cloud.aiplatform.compat.services import ( + tensorboard_service_client, ) from google.cloud.aiplatform_v1.services.tensorboard_service.transports import ( grpc as transports_grpc, ) -from google.cloud.aiplatform_v1.types import tensorboard_data -from google.cloud.aiplatform_v1.types import tensorboard_service -from google.cloud.aiplatform_v1.types import ( + +from google.cloud.aiplatform.compat.types import tensorboard_data +from google.cloud.aiplatform.compat.types import tensorboard_service +from google.cloud.aiplatform.compat.types import ( tensorboard_experiment as tensorboard_experiment_type, ) -from google.cloud.aiplatform_v1.types import tensorboard_run as tensorboard_run_type -from google.cloud.aiplatform_v1.types import ( +from google.cloud.aiplatform.compat.types import tensorboard_run as tensorboard_run_type +from google.cloud.aiplatform.compat.types import ( tensorboard_time_series as tensorboard_time_series_type, ) from google.protobuf import timestamp_pb2 diff --git a/tests/unit/aiplatform/test_uploader_main.py b/tests/unit/aiplatform/test_uploader_main.py index 79c86b22fc..417c865c27 100644 --- a/tests/unit/aiplatform/test_uploader_main.py +++ b/tests/unit/aiplatform/test_uploader_main.py @@ -25,7 +25,9 @@ from google.cloud.aiplatform.tensorboard import uploader_main from google.cloud.aiplatform.compat.types import job_state as gca_job_state_compat from google.cloud.aiplatform.compat.types import custom_job as gca_custom_job_compat -from google.cloud.aiplatform_v1.services.job_service import client as job_service_client +from google.cloud.aiplatform.compat.services import ( + job_service_client, +) _TEST_PROJECT = "test-project" _TEST_LOCATION = "us-central1" @@ -67,6 +69,7 @@ def get_custom_job_mock(): yield get_custom_job_mock +@pytest.mark.usefixtures("google_auth_mock") class TestUploaderMain: def setup_method(self): reload(initializer) diff --git a/tests/unit/aiplatform/test_utils.py b/tests/unit/aiplatform/test_utils.py index ecbb325e69..c700271590 100644 --- a/tests/unit/aiplatform/test_utils.py +++ b/tests/unit/aiplatform/test_utils.py @@ -272,6 +272,7 @@ def test_extract_bucket_and_prefix_from_gcs_path(gcs_path: str, expected: tuple) assert expected == utils.extract_bucket_and_prefix_from_gcs_path(gcs_path) +@pytest.mark.usefixtures("google_auth_mock") def test_wrapped_client(): test_client_info = gapic_v1.client_info.ClientInfo() test_client_options = client_options.ClientOptions()