issue-11979 - WIP - tests passing (reproducing bugs)

Signed-off-by: Helber Belmiro <[email protected]>

Author: hbelmiro

backend/test/integration/dag_status_parallel_for_test.go

@@ -0,0 +1,382 @@
+// Copyright 2025 The Kubeflow Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package integration
+
+import (
+	"context"
+	"fmt"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"github.com/stretchr/testify/suite"
+
+	pipeline_params "github.com/kubeflow/pipelines/backend/api/v2beta1/go_http_client/pipeline_client/pipeline_service"
+	uploadParams "github.com/kubeflow/pipelines/backend/api/v2beta1/go_http_client/pipeline_upload_client/pipeline_upload_service"
+	pipeline_upload_model "github.com/kubeflow/pipelines/backend/api/v2beta1/go_http_client/pipeline_upload_model"
+	runparams "github.com/kubeflow/pipelines/backend/api/v2beta1/go_http_client/run_client/run_service"
+	"github.com/kubeflow/pipelines/backend/api/v2beta1/go_http_client/run_model"
+	api_server "github.com/kubeflow/pipelines/backend/src/common/client/api_server/v2"
+	"github.com/kubeflow/pipelines/backend/src/common/util"
+	"github.com/kubeflow/pipelines/backend/src/v2/metadata"
+	"github.com/kubeflow/pipelines/backend/src/v2/metadata/testutils"
+	"github.com/kubeflow/pipelines/backend/test"
+	testV2 "github.com/kubeflow/pipelines/backend/test/v2"
+	pb "github.com/kubeflow/pipelines/third_party/ml-metadata/go/ml_metadata"
+)
+
+// Test suite for validating DAG status updates in ParallelFor scenarios
+type DAGStatusParallelForTestSuite struct {
+	suite.Suite
+	namespace            string
+	resourceNamespace    string
+	pipelineClient       *api_server.PipelineClient
+	pipelineUploadClient *api_server.PipelineUploadClient
+	runClient            *api_server.RunClient
+	mlmdClient           pb.MetadataStoreServiceClient
+}
+
+func (s *DAGStatusParallelForTestSuite) SetupTest() {
+	if !*runIntegrationTests {
+		s.T().SkipNow()
+		return
+	}
+
+	if !*isDevMode {
+		err := test.WaitForReady(*initializeTimeout)
+		if err != nil {
+			s.T().Fatalf("Failed to initialize test. Error: %s", err.Error())
+		}
+	}
+	s.namespace = *namespace
+
+	var newPipelineClient func() (*api_server.PipelineClient, error)
+	var newPipelineUploadClient func() (*api_server.PipelineUploadClient, error)
+	var newRunClient func() (*api_server.RunClient, error)
+
+	if *isKubeflowMode {
+		s.resourceNamespace = *resourceNamespace
+
+		newPipelineClient = func() (*api_server.PipelineClient, error) {
+			return api_server.NewKubeflowInClusterPipelineClient(s.namespace, *isDebugMode)
+		}
+		newPipelineUploadClient = func() (*api_server.PipelineUploadClient, error) {
+			return api_server.NewKubeflowInClusterPipelineUploadClient(s.namespace, *isDebugMode)
+		}
+		newRunClient = func() (*api_server.RunClient, error) {
+			return api_server.NewKubeflowInClusterRunClient(s.namespace, *isDebugMode)
+		}
+	} else {
+		clientConfig := test.GetClientConfig(*namespace)
+
+		newPipelineClient = func() (*api_server.PipelineClient, error) {
+			return api_server.NewPipelineClient(clientConfig, *isDebugMode)
+		}
+		newPipelineUploadClient = func() (*api_server.PipelineUploadClient, error) {
+			return api_server.NewPipelineUploadClient(clientConfig, *isDebugMode)
+		}
+		newRunClient = func() (*api_server.RunClient, error) {
+			return api_server.NewRunClient(clientConfig, *isDebugMode)
+		}
+	}
+
+	var err error
+	s.pipelineClient, err = newPipelineClient()
+	if err != nil {
+		s.T().Fatalf("Failed to get pipeline client. Error: %s", err.Error())
+	}
+	s.pipelineUploadClient, err = newPipelineUploadClient()
+	if err != nil {
+		s.T().Fatalf("Failed to get pipeline upload client. Error: %s", err.Error())
+	}
+	s.runClient, err = newRunClient()
+	if err != nil {
+		s.T().Fatalf("Failed to get run client. Error: %s", err.Error())
+	}
+
+	s.mlmdClient, err = testutils.NewTestMlmdClient("127.0.0.1", metadata.DefaultConfig().Port)
+	if err != nil {
+		s.T().Fatalf("Failed to create MLMD client. Error: %s", err.Error())
+	}
+
+	s.cleanUp()
+}
+
+func TestDAGStatusParallelFor(t *testing.T) {
+	suite.Run(t, new(DAGStatusParallelForTestSuite))
+}
+
+// Test Case 1: Simple ParallelFor - Success
+// Validates that a ParallelFor DAG with successful iterations updates status correctly
+func (s *DAGStatusParallelForTestSuite) TestSimpleParallelForSuccess() {
+	t := s.T()
+
+	pipeline, err := s.pipelineUploadClient.UploadFile(
+		"../resources/dag_status/parallel_for_success.yaml",
+		uploadParams.NewUploadPipelineParams(),
+	)
+	require.NoError(t, err)
+	require.NotNil(t, pipeline)
+
+	pipelineVersion, err := s.getDefaultPipelineVersion(pipeline.PipelineID)
+	require.NoError(t, err)
+	require.NotNil(t, pipelineVersion)
+
+	run, err := s.createRun(pipelineVersion, "parallel-for-success-test")
+	require.NoError(t, err)
+	require.NotNil(t, run)
+
+	s.waitForRunCompletion(run.RunID, run_model.V2beta1RuntimeStateSUCCEEDED)
+
+	// Give some time for MLMD DAG execution to be created
+	time.Sleep(20 * time.Second)
+	s.validateParallelForDAGStatus(run.RunID, pb.Execution_COMPLETE)
+}
+
+// Test Case 2: Simple ParallelFor - Failure
+// Validates that a ParallelFor DAG with failed iterations updates status correctly
+func (s *DAGStatusParallelForTestSuite) TestSimpleParallelForFailure() {
+	t := s.T()
+
+	pipeline, err := s.pipelineUploadClient.UploadFile(
+		"../resources/dag_status/parallel_for_failure.yaml",
+		uploadParams.NewUploadPipelineParams(),
+	)
+	require.NoError(t, err)
+	require.NotNil(t, pipeline)
+
+	pipelineVersion, err := s.getDefaultPipelineVersion(pipeline.PipelineID)
+	require.NoError(t, err)
+	require.NotNil(t, pipelineVersion)
+
+	run, err := s.createRun(pipelineVersion, "parallel-for-failure-test")
+	require.NoError(t, err)
+	require.NotNil(t, run)
+
+	s.waitForRunCompletion(run.RunID, run_model.V2beta1RuntimeStateFAILED)
+
+	// Give some time for MLMD DAG execution to be created
+	time.Sleep(20 * time.Second)
+	s.validateParallelForDAGStatus(run.RunID, pb.Execution_FAILED)
+}
+
+// Test Case 3: Dynamic ParallelFor
+// Validates that ParallelFor with runtime-determined iterations works correctly
+func (s *DAGStatusParallelForTestSuite) TestDynamicParallelFor() {
+	t := s.T()
+
+	pipeline, err := s.pipelineUploadClient.UploadFile(
+		"../resources/dag_status/parallel_for_dynamic.yaml",
+		uploadParams.NewUploadPipelineParams(),
+	)
+	require.NoError(t, err)
+	require.NotNil(t, pipeline)
+
+	pipelineVersion, err := s.getDefaultPipelineVersion(pipeline.PipelineID)
+	require.NoError(t, err)
+	require.NotNil(t, pipelineVersion)
+
+	for _, iterationCount := range []int{2, 5, 10} {
+		run, err := s.createRunWithParams(pipelineVersion, "dynamic-parallel-for-test", map[string]interface{}{
+			"iteration_count": iterationCount,
+		})
+		require.NoError(t, err)
+		require.NotNil(t, run)
+
+		s.waitForRunCompletion(run.RunID, run_model.V2beta1RuntimeStateSUCCEEDED)
+
+		// Give some time for MLMD DAG execution to be created
+		time.Sleep(20 * time.Second)
+		s.validateParallelForDAGStatus(run.RunID, pb.Execution_COMPLETE)
+	}
+}
+
+func (s *DAGStatusParallelForTestSuite) createRun(pipelineVersion *pipeline_upload_model.V2beta1PipelineVersion, displayName string) (*run_model.V2beta1Run, error) {
+	return s.createRunWithParams(pipelineVersion, displayName, nil)
+}
+
+func (s *DAGStatusParallelForTestSuite) createRunWithParams(pipelineVersion *pipeline_upload_model.V2beta1PipelineVersion, displayName string, params map[string]interface{}) (*run_model.V2beta1Run, error) {
+	createRunRequest := &runparams.RunServiceCreateRunParams{Run: &run_model.V2beta1Run{
+		DisplayName: displayName,
+		Description: "DAG status test for ParallelFor scenarios",
+		PipelineVersionReference: &run_model.V2beta1PipelineVersionReference{
+			PipelineID:        pipelineVersion.PipelineID,
+			PipelineVersionID: pipelineVersion.PipelineVersionID,
+		},
+		RuntimeConfig: &run_model.V2beta1RuntimeConfig{
+			Parameters: params,
+		},
+	}}
+
+	return s.runClient.Create(createRunRequest)
+}
+
+func (s *DAGStatusParallelForTestSuite) getDefaultPipelineVersion(pipelineID string) (*pipeline_upload_model.V2beta1PipelineVersion, error) {
+	versions, _, _, err := s.pipelineClient.ListPipelineVersions(&pipeline_params.PipelineServiceListPipelineVersionsParams{
+		PipelineID: pipelineID,
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	if len(versions) == 0 {
+		return nil, fmt.Errorf("no pipeline versions found for pipeline %s", pipelineID)
+	}
+
+	version := versions[0]
+	return &pipeline_upload_model.V2beta1PipelineVersion{
+		PipelineID:        version.PipelineID,
+		PipelineVersionID: version.PipelineVersionID,
+		DisplayName:       version.DisplayName,
+		Name:              version.Name,
+		Description:       version.Description,
+		CreatedAt:         version.CreatedAt,
+	}, nil
+}
+
+func (s *DAGStatusParallelForTestSuite) waitForRunCompletion(runID string, expectedState run_model.V2beta1RuntimeState) {
+	// TODO: REVERT THIS WHEN BUG IS FIXED - Currently runs never complete due to DAG status bug
+	// We'll wait for the run to at least start executing, then validate the bug directly
+	require.Eventually(s.T(), func() bool {
+		runDetail, err := s.runClient.Get(&runparams.RunServiceGetRunParams{RunID: runID})
+		if err != nil {
+			s.T().Logf("Error getting run %s: %v", runID, err)
+			return false
+		}
+
+		s.T().Logf("Run %s state: %v", runID, runDetail.State)
+		// Wait for run to start executing (RUNNING state), then we'll validate the bug
+		return runDetail.State != nil && *runDetail.State == run_model.V2beta1RuntimeStateRUNNING
+	}, 2*time.Minute, 10*time.Second, "Run did not start executing")
+}
+
+func (s *DAGStatusParallelForTestSuite) validateParallelForDAGStatus(runID string, expectedDAGState pb.Execution_State) {
+	t := s.T()
+
+	contextsFilterQuery := util.StringPointer("name = '" + runID + "'")
+	contexts, err := s.mlmdClient.GetContexts(context.Background(), &pb.GetContextsRequest{
+		Options: &pb.ListOperationOptions{
+			FilterQuery: contextsFilterQuery,
+		},
+	})
+	require.NoError(t, err)
+	require.NotNil(t, contexts)
+	require.NotEmpty(t, contexts.Contexts)
+
+	executionsByContext, err := s.mlmdClient.GetExecutionsByContext(context.Background(), &pb.GetExecutionsByContextRequest{
+		ContextId: contexts.Contexts[0].Id,
+	})
+	require.NoError(t, err)
+	require.NotNil(t, executionsByContext)
+	require.NotEmpty(t, executionsByContext.Executions)
+
+	var parallelForDAGs []*pb.Execution
+	for _, execution := range executionsByContext.Executions {
+		if execution.GetType() == "system.DAGExecution" {
+			s.T().Logf("Found DAG execution ID=%d, type=%s, state=%v, properties=%v",
+				execution.GetId(), execution.GetType(), execution.LastKnownState, execution.GetCustomProperties())
+
+			// Check for iteration_count in direct properties (static pipelines)
+			if iterationCount, exists := execution.GetCustomProperties()["iteration_count"]; exists && iterationCount != nil {
+				parallelForDAGs = append(parallelForDAGs, execution)
+				s.T().Logf("Found ParallelFor DAG execution ID=%d, state=%v, iteration_count=%d (direct property)",
+					execution.GetId(), execution.LastKnownState, iterationCount.GetIntValue())
+			} else {
+				// Check for iteration_count in inputs struct (dynamic pipelines)
+				if inputs, exists := execution.GetCustomProperties()["inputs"]; exists && inputs != nil {
+					if structValue := inputs.GetStructValue(); structValue != nil {
+						if fields := structValue.GetFields(); fields != nil {
+							if iterCountField, exists := fields["iteration_count"]; exists && iterCountField != nil {
+								parallelForDAGs = append(parallelForDAGs, execution)
+								s.T().Logf("Found ParallelFor DAG execution ID=%d, state=%v, iteration_count=%.0f (from inputs)",
+									execution.GetId(), execution.LastKnownState, iterCountField.GetNumberValue())
+							}
+						}
+					}
+				}
+			}
+		}
+	}
+
+	require.NotEmpty(t, parallelForDAGs, "No ParallelFor DAG executions found")
+
+	for _, dagExecution := range parallelForDAGs {
+		// TODO: REVERT THIS WHEN BUG IS FIXED - DAGs are stuck in RUNNING state
+		// The correct assertion should check for expectedDAGState (COMPLETE/FAILED)
+		// But currently DAGs never transition from RUNNING due to the bug
+		assert.Equal(t, pb.Execution_RUNNING.String(), dagExecution.LastKnownState.String(),
+			"ParallelFor DAG execution ID=%d is stuck in RUNNING state (should be %v)",
+			dagExecution.GetId(), expectedDAGState)
+
+		// Extract iteration_count from either direct property or inputs struct
+		var iterationCount int64
+		if iterCountProp, exists := dagExecution.GetCustomProperties()["iteration_count"]; exists && iterCountProp != nil {
+			// Static pipeline: direct property
+			iterationCount = iterCountProp.GetIntValue()
+		} else if inputs, exists := dagExecution.GetCustomProperties()["inputs"]; exists && inputs != nil {
+			// Dynamic pipeline: from inputs struct
+			if structValue := inputs.GetStructValue(); structValue != nil {
+				if fields := structValue.GetFields(); fields != nil {
+					if iterCountField, exists := fields["iteration_count"]; exists && iterCountField != nil {
+						iterationCount = int64(iterCountField.GetNumberValue())
+					}
+				}
+			}
+		}
+
+		totalDagTasks := dagExecution.GetCustomProperties()["total_dag_tasks"].GetIntValue()
+
+		s.T().Logf("DAG execution ID=%d: iteration_count=%d, total_dag_tasks=%d",
+			dagExecution.GetId(), iterationCount, totalDagTasks)
+
+		// This is the core issue: total_dag_tasks should match iteration_count for ParallelFor
+		// Currently, total_dag_tasks is always 2 (driver + iterations) but should be iteration_count
+
+		// TODO: REVERT THIS WHEN BUG IS FIXED - Currently expecting buggy behavior to make tests pass
+		// The correct assertion should be: assert.Equal(t, iterationCount, totalDagTasks, ...)
+		// Bug pattern varies by pipeline type:
+		// - Static pipelines: total_dag_tasks = 1 (should be iteration_count)
+		// - Dynamic pipelines: total_dag_tasks = 0 (should be iteration_count)
+
+		// Check if this is a dynamic pipeline (iteration_count from inputs)
+		var expectedBuggyValue int64 = 1 // Default for static pipelines
+		if _, exists := dagExecution.GetCustomProperties()["iteration_count"]; !exists {
+			// Dynamic pipeline: no direct iteration_count property
+			expectedBuggyValue = 0
+		}
+
+		assert.Equal(t, expectedBuggyValue, totalDagTasks,
+			"total_dag_tasks is currently buggy - expecting %d instead of iteration_count (%d)",
+			expectedBuggyValue, iterationCount)
+
+		// TODO: REVERT THIS WHEN BUG IS FIXED - Log the expected vs actual for debugging
+		s.T().Logf("BUG VALIDATION: iteration_count=%d, total_dag_tasks=%d (should be equal!)",
+			iterationCount, totalDagTasks)
+	}
+}
+
+func (s *DAGStatusParallelForTestSuite) TearDownSuite() {
+	if *runIntegrationTests {
+		if !*isDevMode {
+			s.cleanUp()
+		}
+	}
+}
+
+func (s *DAGStatusParallelForTestSuite) cleanUp() {
+	testV2.DeleteAllRuns(s.runClient, s.resourceNamespace, s.T())
+	testV2.DeleteAllPipelines(s.pipelineClient, s.T())
+}