fix: Don't throw exception when getting representation of unrun GCA objects (#1071)

samgoodman · web-flow · commit c9ba0603e6a8 · 2022-03-14T10:10:24.000-04:00
diff --git a/google/cloud/aiplatform/base.py b/google/cloud/aiplatform/base.py
@@ -655,6 +655,15 @@ def gca_resource(self) -> proto.Message:
         self._assert_gca_resource_is_available()
         return self._gca_resource
 
+    @property
+    def _resource_is_available(self) -> bool:
+        """Returns True if GCA resource has been created and is available, otherwise False"""
+        try:
+            self._assert_gca_resource_is_available()
+            return True
+        except RuntimeError:
+            return False
+
     def _assert_gca_resource_is_available(self) -> None:
         """Helper method to raise when property is not accessible.
 
@@ -1153,7 +1162,7 @@ def delete(self, sync: bool = True) -> None:
         _LOGGER.log_action_completed_against_resource("deleted.", "", self)
 
     def __repr__(self) -> str:
-        if self._gca_resource:
+        if self._gca_resource and self._resource_is_available:
             return VertexAiResourceNoun.__repr__(self)
 
         return FutureManager.__repr__(self)
diff --git a/tests/unit/aiplatform/test_custom_job.py b/tests/unit/aiplatform/test_custom_job.py
@@ -661,3 +661,38 @@ def test_create_from_local_script_with_all_args(
         assert (
             job._gca_resource.state == gca_job_state_compat.JobState.JOB_STATE_SUCCEEDED
         )
+
+    @pytest.mark.usefixtures("get_custom_job_mock", "create_custom_job_mock")
+    def test_check_custom_job_availability(self):
+        aiplatform.init(
+            project=_TEST_PROJECT,
+            location=_TEST_LOCATION,
+            staging_bucket=_TEST_STAGING_BUCKET,
+            encryption_spec_key_name=_TEST_DEFAULT_ENCRYPTION_KEY_NAME,
+        )
+
+        job = aiplatform.CustomJob(
+            display_name=_TEST_DISPLAY_NAME,
+            worker_pool_specs=_TEST_WORKER_POOL_SPEC,
+            base_output_dir=_TEST_BASE_OUTPUT_DIR,
+            labels=_TEST_LABELS,
+        )
+
+        assert not job._resource_is_available
+        assert job.__repr__().startswith(
+            "<google.cloud.aiplatform.jobs.CustomJob object"
+        )
+
+        job.run(
+            service_account=_TEST_SERVICE_ACCOUNT,
+            network=_TEST_NETWORK,
+            timeout=_TEST_TIMEOUT,
+            restart_job_on_worker_restart=_TEST_RESTART_JOB_ON_WORKER_RESTART,
+        )
+
+        job.wait_for_resource_creation()
+
+        assert job._resource_is_available
+        assert "resource name" in job.__repr__()
+
+        job.wait()
