Repository: kubeflow/kfp-tekton (language: python)

Issue:
Artifact paths for cached tasks are incorrect
/kind bug

**What steps did you take and what happened:**
[A clear and concise description of what the bug is.]

Artifact paths for tasks which utilize the caching feature by adding `pipelines.kubeflow.org/cache_enabled: 'true'` label are incorrect. These artifact paths use current pipeline ID rather than the cached pipeline ID.

<img width="593" alt="Screen Shot 2021-03-24 at 2 11 12 PM" src="https://user-images.githubusercontent.com/32145094/112383794-cefdca80-8caa-11eb-86d1-98056bd75799.png">


**What did you expect to happen:**
 The artifact paths use the cached pipeline ID.

**Additional information:**
[Miscellaneous information that will assist in solving the issue.]


**Environment:**

* Python Version (use `python --version`):
* SDK Version: 
* Tekton Version (use `tkn version`):
* Kubernetes Version (use `kubectl version`):
* OS (e.g. from `/etc/os-release`):



Discussion / hints:

Fix patch:
diff --git a/backend/src/cache/server/mutation.go b/backend/src/cache/server/mutation.go
index 09f9e7a592..8854bed60d 100644
--- a/backend/src/cache/server/mutation.go
+++ b/backend/src/cache/server/mutation.go
@@ -58,6 +58,8 @@ const (
 	TaskName                   string = "tekton.dev/pipelineTask"
 	PipelineName               string = "pipelines.kubeflow.org/pipelinename"
 	Generation                 string = "pipelines.kubeflow.org/generation"
+	PipelineRun               string = "tekton.dev/pipelineRun"
+	CachedPipeline            string = "pipelines.kubeflow.org/cached_pipeline_run"
 
 	TektonGroup        string = "tekton.dev/v1beta1"
 	TektonTaskKind     string = "TaskRun"
@@ -181,6 +183,8 @@ func MutatePodIfCached(req *v1beta1.AdmissionRequest, clientMgr ClientManagerInt
 
 		result := getValueFromSerializedMap(cachedExecution.ExecutionOutput, TektonTaskrunOutputs)
 		annotations[TektonTaskrunOutputs] = result
+		cachedPipelineRun := getValueFromSerializedMap(cachedExecution.ExecutionOutput, CachedPipeline)
+		annotations[CachedPipeline] = cachedPipelineRun
 		labels[CacheIDLabelKey] = strconv.FormatInt(cachedExecution.ID, 10)
 		labels[KFPCachedLabelKey] = KFPCachedLabelValue // This label indicates the pod is taken from cache.
 
@@ -257,6 +261,7 @@ func prepareMainContainer(pod *corev1.Pod, result string, logger *zap.SugaredLog
 	}
 
 	args := []string{}
+	args = append(args, "printf 'This step output is taken from cache.\n\n'")
 	for _, result := range results {
 		arg := fmt.Sprintf("printf '%s' | tee /tekton/results/%s", result.Value, result.Name)
 		args = append(args, arg)
diff --git a/backend/src/cache/server/watcher.go b/backend/src/cache/server/watcher.go
index aa44bce3ed..f67a9dc1f0 100644
--- a/backend/src/cache/server/watcher.go
+++ b/backend/src/cache/server/watcher.go
@@ -79,6 +79,7 @@ func WatchPods(namespaceToWatch string, clientManager ClientManagerInterface) {
 			executionOutputMap := make(map[string]interface{})
 			executionOutputMap[TektonTaskrunOutputs] = executionOutput
 			executionOutputMap[MetadataExecutionIDKey] = pod.ObjectMeta.Labels[MetadataExecutionIDKey]
+			executionOutputMap[CachedPipeline] = pod.ObjectMeta.Labels[PipelineRun]
 			executionOutputJSON, _ := json.Marshal(executionOutputMap)
 
 			executionMaxCacheStaleness, exists := pod.ObjectMeta.Annotations[MaxCacheStalenessKey]
diff --git a/frontend/src/lib/WorkflowParser.ts b/frontend/src/lib/WorkflowParser.ts
index b91895c2a2..b2dfef1f53 100644
--- a/frontend/src/lib/WorkflowParser.ts
+++ b/frontend/src/lib/WorkflowParser.ts
@@ -134,6 +134,7 @@ export default class WorkflowParser {
 
         // Checks if the task is an any-sequencer and if so, adds the dependencies from the task list
         if (
+          task['taskSpec'] &&
           task['taskSpec']['steps'] &&
           task['taskSpec']['steps'][0] &&
           task['taskSpec']['steps'][0]['args'] &&
@@ -359,6 +360,7 @@ export default class WorkflowParser {
   public static getNodeInputOutputArtifacts(
     workflow?: any,
     nodeId?: string,
+    cachedPipelineRun?: string,
   ): Record<'inputArtifacts' | 'outputArtifacts', Array<KeyValue<S3Artifact>>> {
     type ParamList = Array<KeyValue<S3Artifact>>;
     let inputArtifacts: ParamList = [];
@@ -368,7 +370,7 @@ export default class WorkflowParser {
       !workflow ||
       !workflow.metadata ||
       !workflow.status ||
-      !workflow.status.taskRuns ||
+      !(workflow.status.taskRuns || workflow.status.runs) ||
       !workflow.metadata.annotations
     )
       return { inputArtifacts, outputArtifacts };
@@ -384,6 +386,15 @@ export default class WorkflowParser {
       }
     }
 
+    // Loop for custom tasks
+    for (const runId of Object.getOwnPropertyNames(workflow.status.runs || {})) {
+      const run = workflow.status.runs[runId];
+      if (run.status && run.pipelineTaskName === nodeId) {
+        taskName = run.pipelineTaskName
+        taskStatus = this.getStatus(run);
+      }
+    }
+
     if (!taskName) return { inputArtifacts, outputArtifacts };
 
     const annotations = workflow.metadata.annotations;
@@ -404,6 +415,17 @@ export default class WorkflowParser {
     };
 
     inputArtifacts = (rawInputArtifacts[taskName] || []).map((artifact: any) => {
+      const parentTask = artifact.parent_task;
+      let pipelineRunID = workflow.metadata.name || '';
+      // Use the cached pipelineRun as the pipelineRunID if this artifact is stored by a taskRun with caching enabled
+      workflow.spec.pipelineSpec.tasks.forEach((task: any) => {
+        if (
+          parentTask === task.name &&
+          task.taskSpec.metadata.labels['pipelines.kubeflow.org/cache_enabled'] === 'true' &&
+          cachedPipelineRun
+        )
+          pipelineRunID = cachedPipelineRun;
+      });
       return [
         artifact.name,
         {
@@ -412,9 +434,9 @@ export default class WorkflowParser {
           insecure: template.insecure,
           accessKeySecret: template.accessKeySecret,
           secretKeySecret: template.secretKeySecret,
-          key: `artifacts/${workflow.metadata.name}/${
-            artifact.parent_task
-          }/${artifact.name.substring(artifact.parent_task.length + 1)}.tgz`,
+          key: `artifacts/${pipelineRunID}/${parentTask}/${artifact.name.substring(
+            artifact.parent_task.length + 1,
+          )}.tgz`,
         },
       ];
     });
@@ -430,6 +452,18 @@ export default class WorkflowParser {
       return { inputArtifacts, outputArtifacts };
 
     outputArtifacts = rawOutputArtifacts[taskName].map((artifact: any) => {
+      const split = artifact.key.split('/');
+      const taskName = split[2];
+      let pipelineRunID = workflow.metadata.name || '';
+      // Use the cached pipelineRun as the pipelineRunID if this artifact is stored by a taskRun with caching enabled
+      workflow.spec.pipelineSpec.tasks.forEach((task: any) => {
+        if (
+          taskName === task.name &&
+          task.taskSpec.metadata.labels['pipelines.kubeflow.org/cache_enabled'] === 'true' &&
+          cachedPipelineRun
+        )
+          pipelineRunID = cachedPipelineRun;
+      });
       return [
         artifact.name,
         {
@@ -438,7 +472,7 @@ export default class WorkflowParser {
           insecure: template.insecure,
           accessKeySecret: template.accessKeySecret,
           secretKeySecret: template.secretKeySecret,
-          key: artifact.key.replace('$PIPELINERUN', workflow.metadata.name || ''),
+          key: artifact.key.replace('$PIPELINERUN', pipelineRunID),
         },
       ];
     });
diff --git a/frontend/src/pages/RunDetails.tsx b/frontend/src/pages/RunDetails.tsx
index f3e2821ad4..19cf4c5227 100644
--- a/frontend/src/pages/RunDetails.tsx
+++ b/frontend/src/pages/RunDetails.tsx
@@ -131,6 +131,7 @@ interface RunDetailsState {
   mlmdRunContext?: Context;
   mlmdExecutions?: Execution[];
   showReducedGraph?: boolean;
+  cachedPipelineRun?: string;
 }
 
 export const css = stylesheet({
@@ -241,6 +242,7 @@ class RunDetails extends Page<RunDetailsInternalProps, RunDetailsState> {
       workflow,
       mlmdExecutions,
       showReducedGraph,
+      cachedPipelineRun,
     } = this.state;
     const { projectId, clusterName } = this.props.gkeMetadata;
     const selectedNodeId = selectedNodeDetails?.id || '';
@@ -258,6 +260,7 @@ class RunDetails extends Page<RunDetailsInternalProps, RunDetailsState> {
     const { inputArtifacts, outputArtifacts } = WorkflowParser.getNodeInputOutputArtifacts(
       workflow,
       selectedNodeId,
+      cachedPipelineRun,
     );
     const selectedExecution = mlmdExecutions?.find(
       execution => ExecutionHelpers.getKfpPod(execution) === selectedNodeId,
@@ -694,6 +697,41 @@ class RunDetails extends Page<RunDetailsInternalProps, RunDetailsState> {
 
       const jsonWorkflow = JSON.parse(runDetail.pipeline_runtime!.workflow_manifest || '{}');
 
+      let cachedTaskRun = '';
+      // Check if any of the pods are cached and if they are then set the cached pipelineRun
+      if (jsonWorkflow.status && jsonWorkflow.status.taskRuns) {
+        jsonWorkflow.spec.pipelineSpec.tasks.forEach((task: any) => {
+          if (
+            task.taskSpec &&
+            task.taskSpec.metadata &&
+            task.taskSpec.metadata.labels &&
+            task.taskSpec.metadata.labels['pipelines.kubeflow.org/cache_enabled'] === 'true' &&
+            !cachedTaskRun
+          ) {
+            Object.keys(jsonWorkflow.status.taskRuns).forEach((taskRun: string) => {
+              if (task.name === jsonWorkflow.status.taskRuns[taskRun].pipelineTaskName) {
+                cachedTaskRun = jsonWorkflow.status.taskRuns[taskRun].status.podName;
+              }
+            });
+          }
+        });
+
+        if (cachedTaskRun) {
+          const podInfo = await Apis.getPodInfo(cachedTaskRun, jsonWorkflow?.metadata?.namespace);
+          if (
+            podInfo &&
+            podInfo.metadata &&
+            podInfo.metadata['annotations'] &&
+            podInfo.metadata['annotations']['pipelines.kubeflow.org/cached_pipeline_run']
+          ) {
+            this.setStateSafe({
+              cachedPipelineRun:
+                podInfo.metadata['annotations']['pipelines.kubeflow.org/cached_pipeline_run'],
+            });
+          }
+        }
+      }
+
       if (
         jsonWorkflow.status &&
         !jsonWorkflow.status.nodes &&


Test patch:
diff --git a/backend/src/cache/server/mutation_test.go b/backend/src/cache/server/mutation_test.go
index 18fbeefa19..1053816a6b 100644
--- a/backend/src/cache/server/mutation_test.go
+++ b/backend/src/cache/server/mutation_test.go
@@ -57,7 +57,7 @@ var (
 					Name:    "main",
 					Image:   "test_image",
 					Command: []string{"python"},
-					Args:    []string{"-entrypoint", "/bin/bash", "--", "-c", "printf '200'"},
+					Args:    []string{"-entrypoint", "/bin/bash", "--", "-c", "printf 'This step output is taken from cache.\n\n'", "printf '200'"},
 				},
 			},
 		},