fix(cache): include artifact URIs in cache key to prevent incorrect reuse

Previously only artifact names were used in cache key generation, causing
incorrect cache hits when different data sources had same artifact names.
Now uses 'name@uri' format to ensure different URIs produce different keys.

Signed-off-by: Aman-Cool <aman017102007@gmail.com>

Author: Aman-Cool

api/v2alpha1/cache_key.proto

@@ -21,6 +21,10 @@ import "google/protobuf/struct.proto";
 import "pipeline_spec.proto";
 
 message CacheKey {
+  // Input artifact identifiers. Each string in ArtifactNameList now uses the format
+  // "name@uri" to include both the artifact name and URI, ensuring different data
+  // sources produce different cache keys. Previously only names were used, which
+  // caused incorrect cache reuse when artifacts had the same name but different URIs.
   map<string, ArtifactNameList> inputArtifactNames = 1;
   map<string, Value> inputParameters = 2 [deprecated = true];
   map<string, RuntimeArtifact> outputArtifactsSpec = 3;
@@ -36,5 +40,8 @@ message ContainerSpec {
 }
 
 message ArtifactNameList {
+  // Artifact identifiers. Format: "name@uri" when URI is present, or just "name"
+  // when URI is empty. This format ensures cache keys differ for artifacts with
+  // the same name but different data sources.
   repeated string artifactNames = 1;
 }

backend/src/v2/cacheutils/cache.go

@@ -187,7 +187,15 @@ func (c *client) GenerateCacheKey(
 	for inputArtifactName, inputArtifactList := range inputs.GetArtifacts() {
 		inputArtifactNameList := cachekey.ArtifactNameList{ArtifactNames: make([]string, 0)}
 		for _, artifact := range inputArtifactList.Artifacts {
-			inputArtifactNameList.ArtifactNames = append(inputArtifactNameList.ArtifactNames, artifact.GetName())
+			// CRITICAL FIX: Include both name AND URI in the cache key identifier.
+			// Previously only the name was used, which caused incorrect cache reuse
+			// when different data sources had artifacts with the same logical name.
+			// The format "name@uri" ensures different URIs produce different cache keys.
+			artifactIdentifier := artifact.GetName()
+			if uri := artifact.GetUri(); uri != "" {
+				artifactIdentifier = artifact.GetName() + "@" + uri
+			}
+			inputArtifactNameList.ArtifactNames = append(inputArtifactNameList.ArtifactNames, artifactIdentifier)
 		}
 		cacheKey.InputArtifactNames[inputArtifactName] = &inputArtifactNameList
 	}

backend/src/v2/cacheutils/cache_test.go

@@ -98,9 +98,11 @@ func TestGenerateCacheKey(t *testing.T) {
 			cmdArgs: []string{"sh", "ec", "test"},
 			image:   "python:3.11",
 			want: &cachekey.CacheKey{
+				// CRITICAL: Artifact identifiers now include URIs in "name@uri" format
+				// to ensure different data sources produce different cache keys.
 				InputArtifactNames: map[string]*cachekey.ArtifactNameList{
-					"dataset_one": {ArtifactNames: []string{"1"}},
-					"dataset_two": {ArtifactNames: []string{"2"}},
+					"dataset_one": {ArtifactNames: []string{"1@gs://some-bucket/dataset-one"}},
+					"dataset_two": {ArtifactNames: []string{"2@gs://some-bucket/dataset-two"}},
 				},
 				InputParameterValues: map[string]*structpb.Value{
 					"message":   {Kind: &structpb.Value_StringValue{StringValue: "Some string value"}},
@@ -187,8 +189,9 @@ func TestGenerateCacheKey(t *testing.T) {
 			image:    "python:3.11",
 			pvcNames: []string{"workspace-pvc", "data-pvc"},
 			want: &cachekey.CacheKey{
+				// CRITICAL: Artifact identifiers now include URIs in "name@uri" format
 				InputArtifactNames: map[string]*cachekey.ArtifactNameList{
-					"dataset_one": {ArtifactNames: []string{"1"}},
+					"dataset_one": {ArtifactNames: []string{"1@gs://some-bucket/dataset-one"}},
 				},
 				InputParameterValues: map[string]*structpb.Value{
 					"message":   {Kind: &structpb.Value_StringValue{StringValue: "Some string value"}},
@@ -353,6 +356,129 @@ func TestGenerateFingerPrint(t *testing.T) {
 	}
 }
 
+// TestGenerateFingerPrint_ConsidersArtifactURIs verifies that different artifact URIs
+// produce different fingerprints, even when artifact names are identical.
+// This test validates the fix for the critical cache key bug where only artifact names
+// (not URIs) were included in the cache key, causing incorrect cache reuse across
+// different data sources.
+func TestGenerateFingerPrint_ConsidersArtifactURIs(t *testing.T) {
+	cacheClient, err := NewClient("ml-pipeline.kubeflow", "8887", false, &tls.Config{})
+	require.NoError(t, err)
+
+	// Create two inputs with same artifact name but different URIs
+	inputsWithURIA := &pipelinespec.ExecutorInput_Inputs{
+		Artifacts: map[string]*pipelinespec.ArtifactList{
+			"dataset": {
+				Artifacts: []*pipelinespec.RuntimeArtifact{
+					{
+						Name: "data",
+						Uri:  "gs://bucket-a/dataset/v1",
+					},
+				},
+			},
+		},
+	}
+
+	inputsWithURIB := &pipelinespec.ExecutorInput_Inputs{
+		Artifacts: map[string]*pipelinespec.ArtifactList{
+			"dataset": {
+				Artifacts: []*pipelinespec.RuntimeArtifact{
+					{
+						Name: "data", // Same name as above
+						Uri:  "gs://bucket-b/dataset/v2", // Different URI
+					},
+				},
+			},
+		},
+	}
+
+	inputsWithSameURI := &pipelinespec.ExecutorInput_Inputs{
+		Artifacts: map[string]*pipelinespec.ArtifactList{
+			"dataset": {
+				Artifacts: []*pipelinespec.RuntimeArtifact{
+					{
+						Name: "data",
+						Uri:  "gs://bucket-a/dataset/v1", // Same URI as inputsWithURIA
+					},
+				},
+			},
+		},
+	}
+
+	outputs := &pipelinespec.ExecutorInput_Outputs{
+		Parameters: map[string]*pipelinespec.ExecutorInput_OutputParameter{
+			"output": {OutputFile: "/tmp/output"},
+		},
+	}
+	outputParamTypes := map[string]string{"output": "STRING"}
+
+	// Generate cache keys
+	cacheKeyA, err := cacheClient.GenerateCacheKey(inputsWithURIA, outputs, outputParamTypes, []string{"echo"}, "python:3.11", nil)
+	require.NoError(t, err)
+
+	cacheKeyB, err := cacheClient.GenerateCacheKey(inputsWithURIB, outputs, outputParamTypes, []string{"echo"}, "python:3.11", nil)
+	require.NoError(t, err)
+
+	cacheKeySameAsA, err := cacheClient.GenerateCacheKey(inputsWithSameURI, outputs, outputParamTypes, []string{"echo"}, "python:3.11", nil)
+	require.NoError(t, err)
+
+	// Generate fingerprints
+	fingerprintA, err := cacheClient.GenerateFingerPrint(cacheKeyA)
+	require.NoError(t, err)
+
+	fingerprintB, err := cacheClient.GenerateFingerPrint(cacheKeyB)
+	require.NoError(t, err)
+
+	fingerprintSameAsA, err := cacheClient.GenerateFingerPrint(cacheKeySameAsA)
+	require.NoError(t, err)
+
+	// CRITICAL ASSERTIONS:
+	// Different URIs MUST produce different fingerprints to prevent incorrect cache reuse
+	assert.NotEqual(t, fingerprintA, fingerprintB,
+		"CRITICAL BUG: Same artifact name with different URIs produced same fingerprint! "+
+			"This would cause incorrect cache reuse across different data sources.")
+
+	// Same URIs should produce same fingerprints (cache hit is correct)
+	assert.Equal(t, fingerprintA, fingerprintSameAsA,
+		"Same artifact name and URI should produce same fingerprint for correct cache hits.")
+
+	// Verify the cache key contains the URI in the artifact identifier
+	assert.Contains(t, cacheKeyA.InputArtifactNames["dataset"].ArtifactNames[0], "gs://bucket-a/dataset/v1",
+		"Cache key should contain the artifact URI")
+	assert.Contains(t, cacheKeyB.InputArtifactNames["dataset"].ArtifactNames[0], "gs://bucket-b/dataset/v2",
+		"Cache key should contain the artifact URI")
+}
+
+// TestGenerateCacheKey_ArtifactWithoutURI verifies that artifacts without URIs
+// still work correctly (backward compatibility).
+func TestGenerateCacheKey_ArtifactWithoutURI(t *testing.T) {
+	cacheClient, err := NewClient("ml-pipeline.kubeflow", "8887", false, &tls.Config{})
+	require.NoError(t, err)
+
+	inputsNoURI := &pipelinespec.ExecutorInput_Inputs{
+		Artifacts: map[string]*pipelinespec.ArtifactList{
+			"dataset": {
+				Artifacts: []*pipelinespec.RuntimeArtifact{
+					{
+						Name: "data",
+						Uri:  "", // No URI
+					},
+				},
+			},
+		},
+	}
+
+	outputs := &pipelinespec.ExecutorInput_Outputs{}
+	outputParamTypes := map[string]string{}
+
+	cacheKey, err := cacheClient.GenerateCacheKey(inputsNoURI, outputs, outputParamTypes, []string{"echo"}, "python:3.11", nil)
+	require.NoError(t, err)
+
+	// When URI is empty, the artifact identifier should just be the name
+	assert.Equal(t, "data", cacheKey.InputArtifactNames["dataset"].ArtifactNames[0],
+		"Artifact without URI should use just the name as identifier")
+}
+
 func TestGenerateFingerPrint_ConsidersPVCNames(t *testing.T) {
 	base := &cachekey.CacheKey{
 		InputArtifactNames: map[string]*cachekey.ArtifactNameList{