Skip to content

Commit a154859

Browse files
connor-mccarthycopybara-github
authored andcommitted
fix: fix bug when checking PipelineJob failure status
PiperOrigin-RevId: 538299878
1 parent 76465e2 commit a154859

File tree

2 files changed

+56
-8
lines changed

2 files changed

+56
-8
lines changed

google/cloud/aiplatform/pipeline_jobs.py

+1-8
Original file line numberDiff line numberDiff line change
@@ -628,13 +628,6 @@ def done(self) -> bool:
628628

629629
return self.state in _PIPELINE_COMPLETE_STATES
630630

631-
def _has_failed(self) -> bool:
632-
"""Return True if PipelineJob has Failed."""
633-
if not self._gca_resource:
634-
return False
635-
636-
return self.state in _PIPELINE_ERROR_STATES
637-
638631
def _get_context(self) -> context.Context:
639632
"""Returns the PipelineRun Context for this PipelineJob in the MetadataStore.
640633
@@ -655,7 +648,7 @@ def _get_context(self) -> context.Context:
655648
time.sleep(1)
656649

657650
if not pipeline_run_context:
658-
if self._has_failed:
651+
if self.has_failed:
659652
raise RuntimeError(
660653
f"Cannot associate PipelineJob to Experiment: {self.gca_resource.error}"
661654
)

tests/unit/aiplatform/test_pipeline_jobs.py

+55
Original file line numberDiff line numberDiff line change
@@ -1567,6 +1567,61 @@ def test_pipeline_failure_raises(self, mock_load_yaml_and_json, sync):
15671567
if not sync:
15681568
job.wait()
15691569

1570+
@pytest.mark.usefixtures(
1571+
"mock_pipeline_service_create",
1572+
"mock_pipeline_service_get_with_fail",
1573+
"mock_pipeline_bucket_exists",
1574+
)
1575+
@pytest.mark.parametrize(
1576+
"job_spec",
1577+
[_TEST_PIPELINE_SPEC_JSON, _TEST_PIPELINE_SPEC_YAML, _TEST_PIPELINE_JOB],
1578+
)
1579+
def test_pipeline_job_has_failed_property(self, mock_load_yaml_and_json):
1580+
aiplatform.init(
1581+
project=_TEST_PROJECT,
1582+
staging_bucket=_TEST_GCS_BUCKET_NAME,
1583+
location=_TEST_LOCATION,
1584+
credentials=_TEST_CREDENTIALS,
1585+
)
1586+
1587+
job = pipeline_jobs.PipelineJob(
1588+
display_name=_TEST_PIPELINE_JOB_DISPLAY_NAME,
1589+
template_path=_TEST_TEMPLATE_PATH,
1590+
job_id=_TEST_PIPELINE_JOB_ID,
1591+
parameter_values=_TEST_PIPELINE_PARAMETER_VALUES,
1592+
enable_caching=True,
1593+
)
1594+
1595+
job.submit(
1596+
service_account=_TEST_SERVICE_ACCOUNT,
1597+
network=_TEST_NETWORK,
1598+
)
1599+
1600+
assert job.state == gca_pipeline_state.PipelineState.PIPELINE_STATE_RUNNING
1601+
assert job.state == gca_pipeline_state.PipelineState.PIPELINE_STATE_RUNNING
1602+
assert job.has_failed
1603+
1604+
@pytest.mark.parametrize(
1605+
"job_spec",
1606+
[_TEST_PIPELINE_SPEC_JSON, _TEST_PIPELINE_SPEC_YAML, _TEST_PIPELINE_JOB],
1607+
)
1608+
def test_pipeline_job_has_failed_property_with_no_submit(
1609+
self, mock_load_yaml_and_json
1610+
):
1611+
job = pipeline_jobs.PipelineJob(
1612+
display_name=_TEST_PIPELINE_JOB_DISPLAY_NAME,
1613+
template_path=_TEST_TEMPLATE_PATH,
1614+
job_id=_TEST_PIPELINE_JOB_ID,
1615+
parameter_values=_TEST_PIPELINE_PARAMETER_VALUES,
1616+
enable_caching=True,
1617+
)
1618+
1619+
with pytest.raises(
1620+
RuntimeError,
1621+
match=r"PipelineJob resource has not been created\.",
1622+
):
1623+
assert job.has_failed
1624+
15701625
@pytest.mark.parametrize(
15711626
"job_spec",
15721627
[_TEST_PIPELINE_SPEC_JSON, _TEST_PIPELINE_SPEC_YAML, _TEST_PIPELINE_JOB],

0 commit comments

Comments
 (0)