chore: merge dev into mb-experimental-release (#185)

* fix: unblock builds (#132)

* chore: Update README with Experimental verbiage (#131)

* fix: Fixed comments (#116)

Co-authored-by: Ivan Cheung <[email protected]>

* feat: Implements a wrapped client that instantiates the client at every API invocation (#139)

* feat: Added optional model args for custom training (#129)

* Added optional model args

* fix: Removed etag

* fix: Added predict schemata and fixed type error

* fix: Added description and fixed predict_schemata

* Added _model_serving_container_command, _model_serving_container_args,                 env=self._model_serving_container_environment_variables and _model_serving_container_ports

* fix: Ran linter

* fix: Added tests for model_instance_schema_uri, model_parameters_schema_uri and model_prediction_schema_uri

* fix: Fixed env and ports and added tests

* fix: Removed model_labels

* fix: Moved container spec creation into init function

* fix: Fixed docstrings

* fix: Moved import to be alphabetical

* fix: Moved model creation to init function

* fix: Fixed predict_schemata

* fix: simplified predict schemata

* fix: added linter

* fix: Fixed trailing comma

* fix: Removed CustomTrainingJob private fields

* fix: Fixed model tests

* fix: Set managed_model to None

Co-authored-by: Ivan Cheung <[email protected]>

* Fix: refactor class constructor for retrieving resource (#125)

* Added property and abstract method _getter_method and _resource_noun, implemented method _get_gca_resource to class AiPlatformResourceNoun; Added _resource_noun, _getter_method, to Dataset, Model, Endpoint, subclasses of _Job, _TrainingJob, refactored (_)get_* and utils.full_resource_name in class constructor to self._get_gca_resource to Dataset, Model, Endpoint, _Job

* Added return value in _get_gca_resource, added method _sync_gca_resource in AiPlatformResourceNoun class; removed job_type, updated status method with _sync_gca_resource in _Job class

* fix: added return type and lint issues

* fix: merge conflict issue with models.py

* fix: F401 'abc' imported but unused

* chore: merge main into dev (#154)

* test: Dataset integration tests (#126)

* Add dataset.metadata.text to schemas

* Add first integation tests, Dataset class

* Make teardown work if test fails, update asserts

* Change test folder name, enable system tests

* Hide test_base, test_end_to_end for Kokoro CI bug

* Add GCP Project env var to Kokoro presubmit cfg

* Restore presubmit cfg, drop --quiet in unit tests

* Restore test_base, test_end_to_end to find timeout

* Skip tests depending on persistent resources

* Use auth default creds for system tests

* Drop unused import os

* feat: specialized dataset classes, fix: datasets refactor (#153)

* feat: Refactored Dataset by removing intermediate layers

* Added image_dataset and tabular_dataset subclass

* Moved metadata_schema_uri responsibility to subclass to enable forecasting

* Moved validation logic for tabular into Dataset._create_tabular

* Added validation in image_dataset and fixed bounding_box schema error

* Removed import_config

* Fixed metadata_schema_uri

* Fixed import and subclasses

* Added EmptyNontabularDatasource

* change import_metadata to ioformat

* added datasources.py

* added support of multiple gcs_sources

* fix: default (empty) dataset_metadata need to be set to {}, not None

* 1) imported datasources 2) added _support_metadata_schema_uris and _support_import_schema_classes 3) added getter and setter/validation for resource_metadata_schema_uri, metadata_schema_uri, and import_schema_uri 4) fixed request_metadata, data_item_labels 5) encapsulated dataset_metadata, and import_data_configs 6) added datasource configuration logic

* added image_dataset.py and tabular_dataset.py

* fix: refactor - create datasets modeule

* fix: cleanup __init__.py

* fix: data_item_labels

* fix: docstring

* fix:
- changed NonTabularDatasource.dataset_metadata default to None
- updated NonTabularDatasource docstring
- changed gcs_source type hint with Union
- changed _create_and_import to _create_encapsulated with datasource
- removed subclass.__init__ and irrelevant parameters in create

* fix: import the module instead of the classes for datasources

* fix: removed all validation for import_schema_uri

* fix: set parameter default to immutable

* fix: replaced Datasource / DatasourceImportable abstract class instead of a concrete type

* fix: added examples for gcs_source

* fix:
- remove Sequence from utils.py
- refactor datasources.py to _datasources.py
- change docstring format to arg_name (arg_type): convention
- change and include the type signature _supported_metadata_schema_uris
- change _validate_metadata_schema_uri
- refactor _create_encapsulated to _create_and_import
- refactor to module level imports
- add tests for ImageDataset and TabularDataset

* fix: remove all labels

* fix: remove Optional in docstring, add example for bq_source

* test: add import_data raise for tabular dataset test

* fix: refactor datasource creation with create_datasource

* fix: lint

Co-authored-by: Ivan Cheung <[email protected]>

* feat: Add AutoML Image Training Job class (#152)

* Add AutoMLImageTrainingJob, tests, constants

* Address reviewer comments

* feat: Add custom container support (#164)

* chore: merge main into dev (#162)

* fix: suppress no project id warning (#160)

* fix: suppress no project id warning

* fix: temporary suppress logging.WARNING and set credentials as google.auth.default credentials

* fix: move default credentials config to credentials property

* fix: add property setter for credentials to avoid everytime reset

* fix: Fixed wrong key value for multilabel (#168)

Co-authored-by: Ivan Cheung <[email protected]>

* feat: Add delete methods, add list_models and undeploy_all for Endpoint class (#165)

* Endpoint list_models, delete, undeploy_all WIP

* Finish delete + undeploy methods, tests

* Add global pool teardowns for test timeout issue

* Address reviewer comments, add async support

* fix: Fixed bug causing training failure for object detection (#171)

Co-authored-by: Ivan Cheung <[email protected]>

* fix: Support intermediary BQ Table for Custom Training (#166)

* chore: add AutoMLImageTrainingJob to aiplatform namespace (#173)

* fix: Unblock build (#174)

* fix: default credentials config related test failures (#167)

* fix: suppress no project id warning

* fix: temporary suppress logging.WARNING and set credentials as google.auth.default credentials

* fix: move default credentials config to credentials property

* fix: add property setter for credentials to avoid everytime reset

* fix: tests for set credentials to default when default not provided

* fix: change credentials with initializer default when not provided in AiPlatformResourceNoun

* fix: use credential mock in tests

* fix: lint

Co-authored-by: sasha-gitg <[email protected]>

* Fix: pass bq_destination to input data config when using training script (#181)

* fix: pass bigquery destination

* fix: add tests and formatting

Co-authored-by: Ivan Cheung <[email protected]>
Co-authored-by: Ivan Cheung <[email protected]>
Co-authored-by: Morgan Du <[email protected]>
Co-authored-by: Vinny Senthil <[email protected]>

Author: 5 people

google/cloud/aiplatform/training_jobs.py

@@ -1580,6 +1580,7 @@ def run(
             managed_model=managed_model,
             args=args,
             base_output_dir=base_output_dir,
+            bigquery_destination=bigquery_destination,
             training_fraction_split=training_fraction_split,
             validation_fraction_split=validation_fraction_split,
             test_fraction_split=test_fraction_split,
@@ -1596,6 +1597,7 @@ def _run(
         managed_model: Optional[gca_model.Model] = None,
         args: Optional[List[Union[str, float, int]]] = None,
         base_output_dir: Optional[str] = None,
+        bigquery_destination: Optional[str] = None,
         training_fraction_split: float = 0.8,
         validation_fraction_split: float = 0.1,
         test_fraction_split: float = 0.1,
@@ -1618,6 +1620,21 @@ def _run(
             base_output_dir (str):
                 GCS output directory of job. If not provided a
                 timestamped directory in the staging directory will be used.
+            bigquery_destination (str):
+                Provide this field if `dataset` is a BiqQuery dataset.
+                The BigQuery project location where the training data is to
+                be written to. In the given project a new dataset is created
+                with name
+                ``dataset_<dataset-id>_<annotation-type>_<timestamp-of-training-call>``
+                where timestamp is in YYYY_MM_DDThh_mm_ss_sssZ format. All
+                training input data will be written into that dataset. In
+                the dataset three tables will be created, ``training``,
+                ``validation`` and ``test``.
+
+                -  AIP_DATA_FORMAT = "bigquery".
+                -  AIP_TRAINING_DATA_URI ="bigquery_destination.dataset_*.training"
+                -  AIP_VALIDATION_DATA_URI = "bigquery_destination.dataset_*.validation"
+                -  AIP_TEST_DATA_URI = "bigquery_destination.dataset_*.test"
             training_fraction_split (float):
                 The fraction of the input data that is to be
                 used to train the Model.
@@ -1679,6 +1696,7 @@ def _run(
             predefined_split_column_name=predefined_split_column_name,
             model=managed_model,
             gcs_destination_uri_prefix=base_output_dir,
+            bigquery_destination=bigquery_destination,
         )
 
         return model

tests/unit/aiplatform/test_training_jobs.py

@@ -581,6 +581,150 @@ def test_run_call_pipeline_service_create(
 
         assert job.state == gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED
 
+    @pytest.mark.parametrize("sync", [True, False])
+    def test_run_call_pipeline_service_create_with_bigquery_destination(
+        self,
+        mock_pipeline_service_create,
+        mock_python_package_to_gcs,
+        mock_dataset,
+        mock_model_service_get,
+        sync,
+    ):
+        aiplatform.init(project=_TEST_PROJECT, staging_bucket=_TEST_BUCKET_NAME)
+
+        job = training_jobs.CustomTrainingJob(
+            display_name=_TEST_DISPLAY_NAME,
+            script_path=_TEST_LOCAL_SCRIPT_FILE_NAME,
+            container_uri=_TEST_TRAINING_CONTAINER_IMAGE,
+            model_serving_container_image_uri=_TEST_SERVING_CONTAINER_IMAGE,
+            model_serving_container_predict_route=_TEST_SERVING_CONTAINER_PREDICTION_ROUTE,
+            model_serving_container_health_route=_TEST_SERVING_CONTAINER_HEALTH_ROUTE,
+            model_instance_schema_uri=_TEST_MODEL_INSTANCE_SCHEMA_URI,
+            model_parameters_schema_uri=_TEST_MODEL_PARAMETERS_SCHEMA_URI,
+            model_prediction_schema_uri=_TEST_MODEL_PREDICTION_SCHEMA_URI,
+            model_serving_container_command=_TEST_MODEL_SERVING_CONTAINER_COMMAND,
+            model_serving_container_args=_TEST_MODEL_SERVING_CONTAINER_ARGS,
+            model_serving_container_environment_variables=_TEST_MODEL_SERVING_CONTAINER_ENVIRONMENT_VARIABLES,
+            model_serving_container_ports=_TEST_MODEL_SERVING_CONTAINER_PORTS,
+            model_description=_TEST_MODEL_DESCRIPTION,
+        )
+
+        model_from_job = job.run(
+            dataset=mock_dataset,
+            base_output_dir=_TEST_BASE_OUTPUT_DIR,
+            bigquery_destination=_TEST_BIGQUERY_DESTINATION,
+            args=_TEST_RUN_ARGS,
+            replica_count=1,
+            machine_type=_TEST_MACHINE_TYPE,
+            accelerator_type=_TEST_ACCELERATOR_TYPE,
+            accelerator_count=_TEST_ACCELERATOR_COUNT,
+            model_display_name=_TEST_MODEL_DISPLAY_NAME,
+            training_fraction_split=_TEST_TRAINING_FRACTION_SPLIT,
+            validation_fraction_split=_TEST_VALIDATION_FRACTION_SPLIT,
+            test_fraction_split=_TEST_TEST_FRACTION_SPLIT,
+            predefined_split_column_name=_TEST_PREDEFINED_SPLIT_COLUMN_NAME,
+            sync=sync,
+        )
+
+        if not sync:
+            model_from_job.wait()
+
+        true_args = _TEST_RUN_ARGS
+
+        true_worker_pool_spec = {
+            "replicaCount": _TEST_REPLICA_COUNT,
+            "machineSpec": {
+                "machineType": _TEST_MACHINE_TYPE,
+                "acceleratorType": _TEST_ACCELERATOR_TYPE,
+                "acceleratorCount": _TEST_ACCELERATOR_COUNT,
+            },
+            "pythonPackageSpec": {
+                "executorImageUri": _TEST_TRAINING_CONTAINER_IMAGE,
+                "pythonModule": training_jobs._TrainingScriptPythonPackager.module_name,
+                "packageUris": [_TEST_OUTPUT_PYTHON_PACKAGE_PATH],
+                "args": true_args,
+            },
+        }
+
+        true_fraction_split = gca_training_pipeline.FractionSplit(
+            training_fraction=_TEST_TRAINING_FRACTION_SPLIT,
+            validation_fraction=_TEST_VALIDATION_FRACTION_SPLIT,
+            test_fraction=_TEST_TEST_FRACTION_SPLIT,
+        )
+
+        env = [
+            env_var.EnvVar(name=str(key), value=str(value))
+            for key, value in _TEST_MODEL_SERVING_CONTAINER_ENVIRONMENT_VARIABLES.items()
+        ]
+
+        ports = [
+            gca_model.Port(container_port=port)
+            for port in _TEST_MODEL_SERVING_CONTAINER_PORTS
+        ]
+
+        true_container_spec = gca_model.ModelContainerSpec(
+            image_uri=_TEST_SERVING_CONTAINER_IMAGE,
+            predict_route=_TEST_SERVING_CONTAINER_PREDICTION_ROUTE,
+            health_route=_TEST_SERVING_CONTAINER_HEALTH_ROUTE,
+            command=_TEST_MODEL_SERVING_CONTAINER_COMMAND,
+            args=_TEST_MODEL_SERVING_CONTAINER_ARGS,
+            env=env,
+            ports=ports,
+        )
+
+        true_managed_model = gca_model.Model(
+            display_name=_TEST_MODEL_DISPLAY_NAME,
+            description=_TEST_MODEL_DESCRIPTION,
+            container_spec=true_container_spec,
+            predict_schemata=gca_model.PredictSchemata(
+                instance_schema_uri=_TEST_MODEL_INSTANCE_SCHEMA_URI,
+                parameters_schema_uri=_TEST_MODEL_PARAMETERS_SCHEMA_URI,
+                prediction_schema_uri=_TEST_MODEL_PREDICTION_SCHEMA_URI,
+            ),
+        )
+
+        true_input_data_config = gca_training_pipeline.InputDataConfig(
+            fraction_split=true_fraction_split,
+            predefined_split=gca_training_pipeline.PredefinedSplit(
+                key=_TEST_PREDEFINED_SPLIT_COLUMN_NAME
+            ),
+            dataset_id=mock_dataset.name,
+            bigquery_destination=gca_io.BigQueryDestination(
+                output_uri=_TEST_BIGQUERY_DESTINATION
+            ),
+        )
+
+        true_training_pipeline = gca_training_pipeline.TrainingPipeline(
+            display_name=_TEST_DISPLAY_NAME,
+            training_task_definition=schema.training_job.definition.custom_task,
+            training_task_inputs=json_format.ParseDict(
+                {
+                    "workerPoolSpecs": [true_worker_pool_spec],
+                    "baseOutputDirectory": {"output_uri_prefix": _TEST_BASE_OUTPUT_DIR},
+                },
+                struct_pb2.Value(),
+            ),
+            model_to_upload=true_managed_model,
+            input_data_config=true_input_data_config,
+        )
+
+        mock_pipeline_service_create.assert_called_once_with(
+            parent=initializer.global_config.common_location_path(),
+            training_pipeline=true_training_pipeline,
+        )
+
+        assert job._gca_resource is mock_pipeline_service_create.return_value
+
+        mock_model_service_get.assert_called_once_with(name=_TEST_MODEL_NAME)
+
+        assert model_from_job._gca_resource is mock_model_service_get.return_value
+
+        assert job.get_model()._gca_resource is mock_model_service_get.return_value
+
+        assert not job.has_failed
+
+        assert job.state == gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED
+
     @pytest.mark.parametrize("sync", [True, False])
     def test_run_called_twice_raises(
         self,