feat: Populate task_unique_name from initial pipeline run in Pipeline Task Rerun Configs for pipeline job rerun

PiperOrigin-RevId: 781659183

Author: vertex-sdk-bot

google/cloud/aiplatform/preview/pipelinejob/pipeline_jobs.py

@@ -578,12 +578,34 @@ def rerun(
             pipeline_job.original_pipeline_job_id = int(
                 original_pipeline_job.labels["vertex-ai-pipelines-run-billing-id"]
             )
+            original_pipeline_task_details = (
+                original_pipeline_job.job_detail.task_details
+            )
         except Exception as e:
             raise ValueError(
                 f"Failed to get original pipeline job: {original_pipelinejob_name}"
             ) from e
 
-        pipeline_job.pipeline_task_rerun_configs = pipeline_task_rerun_configs
+        task_id_to_task_rerun_config = {}
+        for task_rerun_config in pipeline_task_rerun_configs:
+            task_id_to_task_rerun_config[task_rerun_config.task_id] = task_rerun_config
+
+        pipeline_job.pipeline_task_rerun_configs = []
+        for task_detail in original_pipeline_task_details:
+            if task_detail.task_id in task_id_to_task_rerun_config:
+                task_rerun_config = task_id_to_task_rerun_config[task_detail.task_id]
+                if task_detail.task_unique_name:
+                    task_rerun_config.task_name = task_detail.task_unique_name
+                pipeline_job.pipeline_task_rerun_configs.append(task_rerun_config)
+            else:
+                pipeline_job.pipeline_task_rerun_configs.append(
+                    aiplatform_v1beta1.PipelineTaskRerunConfig(
+                        task_id=task_detail.task_id,
+                        task_name=task_detail.task_unique_name,
+                        skip_task=task_detail.state
+                        == aiplatform_v1beta1.PipelineTaskDetail.State.SUCCEEDED,
+                    )
+                )
 
         if parameter_values:
             runtime_config = self._v1_beta1_pipeline_job.runtime_config

tests/unit/aiplatform/test_pipeline_jobs.py

@@ -74,6 +74,7 @@
 _TEST_PIPELINE_JOB_DISPLAY_NAME_2 = "sample-pipeline-job-display-name-2"
 _TEST_PIPELINE_JOB_ID = "sample-test-pipeline-202111111"
 _TEST_PIPELINE_JOB_ID_2 = "sample-test-pipeline-202111112"
+_TEST_PIPELINE_RERUN_JOB_ID = "sample-test-pipeline-rerun"
 _TEST_GCS_BUCKET_NAME = "my-bucket"
 _TEST_GCS_OUTPUT_DIRECTORY = f"gs://{_TEST_GCS_BUCKET_NAME}/output_artifacts/"
 _TEST_CREDENTIALS = auth_credentials.AnonymousCredentials()
@@ -339,6 +340,41 @@ def mock_pipeline_v1beta1_service_get():
         yield mock_get_pipeline_job
 
 
+@pytest.fixture
+def mock_pipeline_v1beta1_service_get_with_task_details():
+    pipeline_job = make_v1beta1_pipeline_job(
+        _TEST_PIPELINE_JOB_NAME,
+        gca_pipeline_state.PipelineState.PIPELINE_STATE_FAILED,
+    )
+    pipeline_job.job_detail.task_details = [
+        v1beta1_pipeline_job.PipelineTaskDetail(
+            task_id=1,
+            task_name="task-1",
+            state=aiplatform_v1beta1.PipelineTaskDetail.State.SUCCEEDED,
+            task_unique_name="task-1",
+        ),
+        v1beta1_pipeline_job.PipelineTaskDetail(
+            task_id=2,
+            task_name="task-2",
+            state=aiplatform_v1beta1.PipelineTaskDetail.State.FAILED,
+            task_unique_name="task-2",
+        ),
+        v1beta1_pipeline_job.PipelineTaskDetail(
+            task_id=3,
+            task_name="task-3",
+            state=aiplatform_v1beta1.PipelineTaskDetail.State.SUCCEEDED,
+            task_unique_name="task-3",
+        ),
+    ]
+
+    with mock.patch.object(
+        v1beta1_pipeline_service.PipelineServiceClient, "get_pipeline_job"
+    ) as mock_get_pipeline_job:
+        mock_get_pipeline_job.side_effect = [pipeline_job]
+
+        yield mock_get_pipeline_job
+
+
 @pytest.fixture
 def mock_pipeline_v1_service_batch_cancel():
     with patch.object(
@@ -2417,6 +2453,63 @@ def test_rerun_v1beta1_pipeline_job_returns_response(
         assert mock_pipeline_v1beta1_service_get.call_count == 1
         assert mock_pipeline_v1beta1_service_create.call_count == 2
 
+    @pytest.mark.usefixtures(
+        "mock_pipeline_v1beta1_service_create",
+        "mock_pipeline_v1beta1_service_get_with_task_details",
+    )
+    @pytest.mark.parametrize(
+        "job_spec",
+        [_TEST_PIPELINE_SPEC_JSON, _TEST_PIPELINE_SPEC_YAML, _TEST_PIPELINE_JOB],
+    )
+    def test_rerun_v1beta1_pipeline_job_returns_response_with_task_details(
+        self,
+        mock_load_yaml_and_json,
+        job_spec,
+        mock_pipeline_v1beta1_service_create,
+        mock_pipeline_v1beta1_service_get_with_task_details,
+    ):
+        aiplatform.init(
+            project=_TEST_PROJECT,
+            staging_bucket=_TEST_GCS_BUCKET_NAME,
+            credentials=_TEST_CREDENTIALS,
+        )
+
+        job = preview_pipeline_jobs._PipelineJob(
+            display_name=_TEST_PIPELINE_JOB_DISPLAY_NAME,
+            template_path=_TEST_TEMPLATE_PATH,
+            job_id=_TEST_PIPELINE_JOB_ID,
+        )
+
+        job.submit()
+
+        job.rerun(
+            job_id=_TEST_PIPELINE_RERUN_JOB_ID,
+            original_pipelinejob_name=_TEST_PIPELINE_JOB_NAME,
+            pipeline_task_rerun_configs=[
+                aiplatform_v1beta1.PipelineTaskRerunConfig(
+                    task_id=1,
+                    skip_task=False,
+                )
+            ],
+            parameter_values={"param-1": "value-1"},
+        )
+
+        assert mock_pipeline_v1beta1_service_get_with_task_details.call_count == 1
+        assert mock_pipeline_v1beta1_service_create.call_count == 2
+
+        rerun_config = mock_pipeline_v1beta1_service_create.call_args[1][
+            "request"
+        ].pipeline_job.pipeline_task_rerun_configs
+        assert rerun_config[0].task_name == "task-1"
+        assert not rerun_config[0].skip_task
+        assert rerun_config[0].task_id == 1
+        assert rerun_config[1].task_name == "task-2"
+        assert not rerun_config[1].skip_task
+        assert rerun_config[1].task_id == 2
+        assert rerun_config[2].task_name == "task-3"
+        assert rerun_config[2].skip_task
+        assert rerun_config[2].task_id == 3
+
     @pytest.mark.parametrize(
         "job_spec",
         [_TEST_PIPELINE_SPEC_JSON, _TEST_PIPELINE_SPEC_YAML, _TEST_PIPELINE_JOB],