chore: LLM - Switched to a more robust way to get the tuned model resource name from the pipeline job

PiperOrigin-RevId: 553633219

Author: Ark-kun

tests/unit/aiplatform/test_language_models.py

@@ -360,12 +360,14 @@ def make_pipeline_job(state):
             task_details=[
                 gca_pipeline_job.PipelineTaskDetail(
                     task_id=456,
-                    task_name="upload-llm-model",
+                    task_name="tune-large-model-20230724214903",
                     execution=GapicExecution(
-                        name="test-execution-name",
-                        display_name="evaluation_metrics",
+                        name="projects/123/locations/europe-west4/metadataStores/default/executions/...",
+                        display_name="tune-large-model-20230724214903",
+                        schema_title="system.Run",
                         metadata={
-                            "output:model_resource_name": "projects/123/locations/us-central1/models/456"
+                            "output:model_resource_name": "projects/123/locations/us-central1/models/456",
+                            "output:endpoint_resource_name": "projects/123/locations/us-central1/endpoints/456",
                         },
                     ),
                 ),

vertexai/language_models/_language_models.py

@@ -1037,19 +1037,19 @@ def result(self) -> "_LanguageModel":
         if self._model:
             return self._model
         self._job.wait()
-        upload_model_tasks = [
-            task_info
-            for task_info in self._job.gca_resource.job_detail.task_details
-            if task_info.task_name == "upload-llm-model"
+        root_pipeline_tasks = [
+            task_detail
+            for task_detail in self._job.gca_resource.job_detail.task_details
+            if task_detail.execution.schema_title == "system.Run"
         ]
-        if len(upload_model_tasks) != 1:
+        if len(root_pipeline_tasks) != 1:
             raise RuntimeError(
                 f"Failed to get the model name from the tuning pipeline: {self._job.name}"
             )
-        upload_model_task = upload_model_tasks[0]
+        root_pipeline_task = root_pipeline_tasks[0]
 
         # Trying to get model name from output parameter
-        vertex_model_name = upload_model_task.execution.metadata[
+        vertex_model_name = root_pipeline_task.execution.metadata[
             "output:model_resource_name"
         ].strip()
         _LOGGER.info(f"Tuning has completed. Created Vertex Model: {vertex_model_name}")