Skip to content

Commit

Permalink
fix: Don't throw exception when getting representation of unrun GCA o…
Browse files Browse the repository at this point in the history
…bjects (#1071)
  • Loading branch information
samgoodman authored Mar 14, 2022
1 parent be0ccc4 commit c9ba060
Show file tree
Hide file tree
Showing 2 changed files with 45 additions and 1 deletion.
11 changes: 10 additions & 1 deletion google/cloud/aiplatform/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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)
Expand Down
35 changes: 35 additions & 0 deletions tests/unit/aiplatform/test_custom_job.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()

0 comments on commit c9ba060

Please sign in to comment.