Merge branch 'main' into mor--reduction-server-integration

googleapis · Oct 15, 2021 · 7fe79b0 · 7fe79b0
2 parents a51daae + c1c2326
commit 7fe79b0
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 1 deletion.
diff --git a/google/cloud/aiplatform/training_jobs.py b/google/cloud/aiplatform/training_jobs.py
@@ -4486,7 +4486,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,
@@ -4498,6 +4498,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."""
@@ -4506,6 +4514,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
@@ -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