feat: Add pipeline_job_name to allow PipelineJob.get(pipeline_job_name)

vertex-sdk-bot · copybara-github · commit 32e3b22993a8 · 2024-06-13T16:50:51.000-07:00
PiperOrigin-RevId: 643162661
diff --git a/tests/unit/aiplatform/test_language_models.py b/tests/unit/aiplatform/test_language_models.py
@@ -2385,7 +2385,8 @@ def test_tune_text_embedding_model(
             model = language_models_module.TextEmbeddingModel.from_pretrained(
                 base_model_version_id
             )
-            tuning_job = model.tune_model(**tune_args)
+            tuning_result = model.tune_model(**tune_args)
+
             call_kwargs = mock_pipeline_service_create.call_args[1]
             pipeline_arguments = dict(
                 call_kwargs["pipeline_job"].runtime_config.parameter_values
@@ -2402,7 +2403,8 @@ def test_tune_text_embedding_model(
             assert pipeline_arguments == expected_pipeline_args
 
             # Testing the tuned model
-            tuned_model = tuning_job.deploy_tuned_model()
+            assert tuning_result.pipeline_job_name.startswith("sample-test-pipeline")
+            tuned_model = tuning_result.deploy_tuned_model()
             assert (
                 tuned_model._endpoint_name
                 == test_constants.EndpointConstants._TEST_ENDPOINT_NAME
diff --git a/vertexai/language_models/_language_models.py b/vertexai/language_models/_language_models.py
@@ -3932,6 +3932,11 @@ def _get_tuned_model_name(self) -> str:
         self._tuned_model_name = vertex_model_name
         return vertex_model_name
 
+    @property
+    def pipeline_job_name(self) -> str:
+        """Pipeline job name."""
+        return self._job.name
+
     @property
     def _status(self) -> Optional[aiplatform_types.pipeline_state.PipelineState]:
         """Job status."""
