From c1c2326b2342ab1b6f4c4ce3852e63376eae740d Mon Sep 17 00:00:00 2001 From: mansiachuthan Date: Thu, 14 Oct 2021 14:06:13 -0700 Subject: [PATCH] feat: enabling AutoML Forecasting training response to include BigQuery location of exported evaluated examples (#657) * Added BigQuery location of exported evaluated examples from AutoML Forecasting Training Job in training output * Testing BQ uri output feature for AutoML Forecasting * Testing BQ uri output feature for AutoML Forecasting * Testing BQ uri output feature for AutoML Forecasting * Testing BQ uri output feature for AutoML Forecasting * Testing BQ uri output feature for AutoML Forecasting * Testing BQ uri output feature for AutoML Forecasting * Testing BQ uri output feature for AutoML Forecasting * Adding changes from dev to master branch * Removed extra try/except block Removed the unnecessary try/except block from the evaluated bigquery uri property method in Automl Forecasting * Fix lint errors * Address comments * Add metadata existence check Co-authored-by: Mansi Achuthan Co-authored-by: Vinny Senthil Co-authored-by: Hardik Vala Co-authored-by: sasha-gitg <44654632+sasha-gitg@users.noreply.github.com> Co-authored-by: thehardikv <78449654+thehardikv@users.noreply.github.com> --- google/cloud/aiplatform/training_jobs.py | 27 ++++++++++++++++++- .../test_automl_forecasting_training_jobs.py | 3 +++ 2 files changed, 29 insertions(+), 1 deletion(-) diff --git a/google/cloud/aiplatform/training_jobs.py b/google/cloud/aiplatform/training_jobs.py index 9436f19cfe..87b0fa6d02 100644 --- a/google/cloud/aiplatform/training_jobs.py +++ b/google/cloud/aiplatform/training_jobs.py @@ -4418,7 +4418,7 @@ def _run( encryption_spec=self._model_encryption_spec, ) - return self._run_job( + new_model = self._run_job( training_task_definition=training_task_definition, training_task_inputs=training_task_inputs_dict, dataset=dataset, @@ -4430,6 +4430,14 @@ def _run( model=model, ) + if export_evaluated_data_items: + _LOGGER.info( + "Exported examples available at:\n%s" + % self.evaluated_data_items_bigquery_uri + ) + + return new_model + @property def _model_upload_fail_string(self) -> str: """Helper property for model upload failure.""" @@ -4438,6 +4446,23 @@ def _model_upload_fail_string(self) -> str: "Model." ) + @property + def evaluated_data_items_bigquery_uri(self) -> Optional[str]: + """BigQuery location of exported evaluated examples from the Training Job + Returns: + str: BigQuery uri for the exported evaluated examples if the export + feature is enabled for training. + None: If the export feature was not enabled for training. + """ + + self._assert_gca_resource_is_available() + + metadata = self._gca_resource.training_task_metadata + if metadata and "evaluatedDataItemsBigqueryUri" in metadata: + return metadata["evaluatedDataItemsBigqueryUri"] + + return None + def _add_additional_experiments(self, additional_experiments: List[str]): """Add experiment flags to the training job. Args: diff --git a/tests/unit/aiplatform/test_automl_forecasting_training_jobs.py b/tests/unit/aiplatform/test_automl_forecasting_training_jobs.py index dc2e00b658..5aad3225c7 100644 --- a/tests/unit/aiplatform/test_automl_forecasting_training_jobs.py +++ b/tests/unit/aiplatform/test_automl_forecasting_training_jobs.py @@ -151,6 +151,9 @@ def mock_pipeline_service_get(): name=_TEST_PIPELINE_RESOURCE_NAME, state=gca_pipeline_state.PipelineState.PIPELINE_STATE_SUCCEEDED, model_to_upload=gca_model.Model(name=_TEST_MODEL_NAME), + training_task_metadata={ + "evaluatedDataItemsBigqueryUri": _TEST_TRAINING_EXPORT_EVALUATED_DATA_ITEMS_BIGQUERY_DESTINATION_URI + }, ) yield mock_get_training_pipeline