File tree

2 files changed

+45
-0
lines changed

2 files changed

+45
-0
lines changed
Original file line numberDiff line numberDiff line change
@@ -331,6 +331,43 @@ def _block_until_complete(self):
331331
else:
332332
_LOGGER.log_action_completed_against_resource("run", "completed", self)
333333

334+
@classmethod
335+
def get(
336+
cls,
337+
resource_name: str,
338+
project: Optional[str] = None,
339+
location: Optional[str] = None,
340+
credentials: Optional[auth_credentials.Credentials] = None,
341+
) -> "PipelineJob":
342+
"""Get a Vertex AI Pipeline Job for the given resource_name.
343+
344+
Args:
345+
resource_name (str):
346+
Required. A fully-qualified resource name or ID.
347+
project (str):
348+
Optional. Project to retrieve dataset from. If not set, project
349+
set in aiplatform.init will be used.
350+
location (str):
351+
Optional. Location to retrieve dataset from. If not set,
352+
location set in aiplatform.init will be used.
353+
credentials (auth_credentials.Credentials):
354+
Optional. Custom credentials to use to upload this model.
355+
Overrides credentials set in aiplatform.init.
356+
357+
Returns:
358+
A Vertex AI PipelineJob.
359+
"""
360+
self = cls._empty_constructor(
361+
project=project,
362+
location=location,
363+
credentials=credentials,
364+
resource_name=resource_name,
365+
)
366+
367+
self._gca_resource = self._get_gca_resource(resource_name=resource_name)
368+
369+
return self
370+
334371
def cancel(self) -> None:
335372
"""Starts asynchronous cancellation on the PipelineJob. The server
336373
makes a best effort to cancel the job, but success is not guaranteed.
Original file line numberDiff line numberDiff line change
@@ -267,6 +267,14 @@ def test_cancel_pipeline_job(
267267
name=_TEST_PIPELINE_JOB_NAME
268268
)
269269

270+
@pytest.mark.usefixtures("mock_pipeline_service_get")
271+
def test_get_training_job(self, mock_pipeline_service_get):
272+
aiplatform.init(project=_TEST_PROJECT)
273+
job = pipeline_jobs.PipelineJob.get(resource_name=_TEST_PIPELINE_JOB_ID)
274+
275+
mock_pipeline_service_get.assert_called_once_with(name=_TEST_PIPELINE_JOB_NAME)
276+
assert isinstance(job, pipeline_jobs.PipelineJob)
277+
270278
@pytest.mark.usefixtures(
271279
"mock_pipeline_service_create", "mock_pipeline_service_get", "mock_load_json",
272280
)

0 commit comments

Comments
 (0)