feat: add get method for PipelineJob (#561)

* feat: get method for pipeline job * Update google/cloud/aiplatform/pipeline_jobs.py Co-authored-by: sasha-gitg <44654632+sasha-gitg@users.noreply.github.com> Co-authored-by: sasha-gitg <44654632+sasha-gitg@users.noreply.github.com>
googleapis · Jul 21, 2021 · fe5e6e4 · fe5e6e4
1 parent c9da662
commit fe5e6e4
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 0 deletions.
diff --git a/google/cloud/aiplatform/pipeline_jobs.py b/google/cloud/aiplatform/pipeline_jobs.py
@@ -331,6 +331,43 @@ def _block_until_complete(self):
         else:
             _LOGGER.log_action_completed_against_resource("run", "completed", self)
 
+    @classmethod
+    def get(
+        cls,
+        resource_name: str,
+        project: Optional[str] = None,
+        location: Optional[str] = None,
+        credentials: Optional[auth_credentials.Credentials] = None,
+    ) -> "PipelineJob":
+        """Get a Vertex AI Pipeline Job for the given resource_name.
+
+        Args:
+            resource_name (str):
+                Required. A fully-qualified resource name or ID.
+            project (str):
+                Optional. Project to retrieve dataset from. If not set, project
+                set in aiplatform.init will be used.
+            location (str):
+                Optional. Location to retrieve dataset from. If not set,
+                location set in aiplatform.init will be used.
+            credentials (auth_credentials.Credentials):
+                Optional. Custom credentials to use to upload this model.
+                Overrides credentials set in aiplatform.init.
+
+        Returns:
+            A Vertex AI PipelineJob.
+        """
+        self = cls._empty_constructor(
+            project=project,
+            location=location,
+            credentials=credentials,
+            resource_name=resource_name,
+        )
+
+        self._gca_resource = self._get_gca_resource(resource_name=resource_name)
+
+        return self
+
     def cancel(self) -> None:
         """Starts asynchronous cancellation on the PipelineJob. The server
         makes a best effort to cancel the job, but success is not guaranteed.

diff --git a/tests/unit/aiplatform/test_pipeline_jobs.py b/tests/unit/aiplatform/test_pipeline_jobs.py
@@ -267,6 +267,14 @@ def test_cancel_pipeline_job(
             name=_TEST_PIPELINE_JOB_NAME
         )
 
+    @pytest.mark.usefixtures("mock_pipeline_service_get")
+    def test_get_training_job(self, mock_pipeline_service_get):
+        aiplatform.init(project=_TEST_PROJECT)
+        job = pipeline_jobs.PipelineJob.get(resource_name=_TEST_PIPELINE_JOB_ID)
+
+        mock_pipeline_service_get.assert_called_once_with(name=_TEST_PIPELINE_JOB_NAME)
+        assert isinstance(job, pipeline_jobs.PipelineJob)
+
     @pytest.mark.usefixtures(
         "mock_pipeline_service_create", "mock_pipeline_service_get", "mock_load_json",
     )