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