Read pipeline run node details from mlmd context

frontend/src/concepts/pipelines/apiHooks/mlmd/types.ts

@@ -0,0 +1,7 @@
+import { Context } from '~/third_party/mlmd';
+
+export type MlmdContext = Context;
+
+export enum MlmdContextTypes {
+  RUN = 'system.PipelineRun',
+}

frontend/src/concepts/pipelines/apiHooks/mlmd/useExecutionsFromMlmdContext.ts

@@ -0,0 +1,31 @@
+import React from 'react';
+import { MlmdContext } from '~/concepts/pipelines/apiHooks/mlmd/types';
+import { usePipelinesAPI } from '~/concepts/pipelines/context';
+import { Execution, GetExecutionsByContextRequest } from '~/third_party/mlmd';
+import useFetchState, {
+  FetchState,
+  FetchStateCallbackPromise,
+  NotReadyError,
+} from '~/utilities/useFetchState';
+
+export const useExecutionsFromMlmdContext = (
+  context: MlmdContext | null,
+  refreshRate?: number,
+): FetchState<Execution[] | null> => {
+  const { metadataStoreServiceClient } = usePipelinesAPI();
+
+  const call = React.useCallback<FetchStateCallbackPromise<Execution[] | null>>(async () => {
+    if (!context) {
+      return Promise.reject(new NotReadyError('No context'));
+    }
+
+    const request = new GetExecutionsByContextRequest();
+    request.setContextId(context.getId());
+    const res = await metadataStoreServiceClient.getExecutionsByContext(request);
+    return res.getExecutionsList();
+  }, [metadataStoreServiceClient, context]);
+
+  return useFetchState(call, null, {
+    refreshRate,
+  });
+};

frontend/src/concepts/pipelines/apiHooks/mlmd/useMlmdContext.ts

@@ -0,0 +1,45 @@
+import React from 'react';
+import { MlmdContext, MlmdContextTypes } from '~/concepts/pipelines/apiHooks/mlmd/types';
+import { usePipelinesAPI } from '~/concepts/pipelines/context';
+import { GetContextByTypeAndNameRequest } from '~/third_party/mlmd';
+import useFetchState, {
+  FetchState,
+  FetchStateCallbackPromise,
+  NotReadyError,
+} from '~/utilities/useFetchState';
+
+/**
+ * A hook used to use the MLMD service and fetch the MLMD context
+ * If being used without name/type, this hook will throw an error
+ * @param name The identifier to query a specific type of MLMD context. e.g. The runID for a pipeline run
+ */
+export const useMlmdContext = (
+  name?: string,
+  type?: MlmdContextTypes,
+  refreshRate?: number,
+): FetchState<MlmdContext | null> => {
+  const { metadataStoreServiceClient } = usePipelinesAPI();
+
+  const call = React.useCallback<FetchStateCallbackPromise<MlmdContext | null>>(async () => {
+    if (!type) {
+      return Promise.reject(new NotReadyError('No context type'));
+    }
+    if (!name) {
+      return Promise.reject(new NotReadyError('No context name'));
+    }
+
+    const request = new GetContextByTypeAndNameRequest();
+    request.setTypeName(type);
+    request.setContextName(name);
+    const res = await metadataStoreServiceClient.getContextByTypeAndName(request);
+    const context = res.getContext();
+    if (!context) {
+      return Promise.reject(new Error('Cannot find specified context'));
+    }
+    return context;
+  }, [metadataStoreServiceClient, type, name]);
+
+  return useFetchState(call, null, {
+    refreshRate,
+  });
+};

frontend/src/concepts/pipelines/apiHooks/mlmd/usePipelineRunMlmdContext.ts

@@ -0,0 +1,8 @@
+import { FetchState } from '~/utilities/useFetchState';
+import { MlmdContext, MlmdContextTypes } from '~/concepts/pipelines/apiHooks/mlmd/types';
+import { useMlmdContext } from '~/concepts/pipelines/apiHooks/mlmd/useMlmdContext';
+
+export const usePipelineRunMlmdContext = (
+  runID?: string,
+  refreshRate?: number,
+): FetchState<MlmdContext | null> => useMlmdContext(runID, MlmdContextTypes.RUN, refreshRate);

frontend/src/concepts/pipelines/apiHooks/usePipelineRunById.ts

@@ -9,6 +9,15 @@ import { PipelineRunKFv2, RuntimeStateKF, runtimeStateLabels } from '~/concepts/
 import { FAST_POLL_INTERVAL } from '~/utilities/const';
 import { computeRunStatus } from '~/concepts/pipelines/content/utils';
 
+export const isPipelineRunFinished = (run?: PipelineRunKFv2 | null): boolean => {
+  const { label } = computeRunStatus(run);
+  return [
+    runtimeStateLabels[RuntimeStateKF.SUCCEEDED],
+    runtimeStateLabels[RuntimeStateKF.FAILED],
+    runtimeStateLabels[RuntimeStateKF.CANCELED],
+  ].includes(label);
+};
+
 const usePipelineRunById = (
   pipelineRunId?: string,
   refreshForDetails?: boolean,
@@ -32,18 +41,13 @@ const usePipelineRunById = (
   });
 
   const [run] = runData;
-  const { label } = computeRunStatus(run);
-  const isComplete = [
-    runtimeStateLabels[RuntimeStateKF.SUCCEEDED],
-    runtimeStateLabels[RuntimeStateKF.FAILED],
-    runtimeStateLabels[RuntimeStateKF.CANCELED],
-  ].includes(label);
+  const isFinished = isPipelineRunFinished(run);
 
   React.useEffect(() => {
-    if (isComplete) {
+    if (isFinished) {
       setPipelineFinished(true);
     }
-  }, [isComplete]);
+  }, [isFinished]);
 
   return runData;
 };

frontend/src/concepts/pipelines/content/pipelinesDetails/pipelineRun/PipelineRunDetails.tsx

@@ -36,6 +36,7 @@ import { usePipelineTaskTopology } from '~/concepts/pipelines/topology';
 import { PipelineRunType } from '~/pages/pipelines/global/runs/types';
 import { routePipelineRunsNamespace } from '~/routes';
 import PipelineJobReferenceName from '~/concepts/pipelines/content/PipelineJobReferenceName';
+import useExecutionsForPipelineRun from '~/concepts/pipelines/content/pipelinesDetails/pipelineRun/useExecutionsForPipelineRun';
 
 const PipelineRunDetails: PipelineCoreDetailsPageComponent = ({ breadcrumbPath, contextPath }) => {
   const { runId } = useParams();
@@ -52,17 +53,16 @@ const PipelineRunDetails: PipelineCoreDetailsPageComponent = ({ breadcrumbPath,
     RunDetailsTabs.DETAILS,
   );
   const [selectedId, setSelectedId] = React.useState<string | null>(null);
-  const { taskMap, nodes } = usePipelineTaskTopology(pipelineSpec, runResource ?? undefined);
+
+  const [executions, executionsLoaded, executionsError] = useExecutionsForPipelineRun(runResource);
+  const { taskMap, nodes } = usePipelineTaskTopology(
+    pipelineSpec,
+    runResource?.run_details,
+    executions,
+  );
 
   const loaded = runLoaded && (versionLoaded || !!runResource?.pipeline_spec);
   const error = versionError || runError;
-  if (!loaded && !error) {
-    return (
-      <Bullseye>
-        <Spinner />
-      </Bullseye>
-    );
-  }
 
   if (error) {
     return (
@@ -77,6 +77,14 @@ const PipelineRunDetails: PipelineCoreDetailsPageComponent = ({ breadcrumbPath,
     );
   }
 
+  if (!loaded || (!executionsLoaded && !executionsError)) {
+    return (
+      <Bullseye>
+        <Spinner />
+      </Bullseye>
+    );
+  }
+
   return (
     <>
       <Drawer isExpanded={!!selectedId}>
@@ -136,12 +144,10 @@ const PipelineRunDetails: PipelineCoreDetailsPageComponent = ({ breadcrumbPath,
                     </Breadcrumb>
                   }
                   headerAction={
-                    loaded && (
-                      <PipelineRunDetailsActions
-                        run={runResource}
-                        onDelete={() => setDeleting(true)}
-                      />
-                    )
+                    <PipelineRunDetailsActions
+                      run={runResource}
+                      onDelete={() => setDeleting(true)}
+                    />
                   }
                   empty={false}
                 >

frontend/src/concepts/pipelines/content/pipelinesDetails/pipelineRun/runLogs/LogsTab.tsx

@@ -40,12 +40,12 @@ import DownloadDropdown from '~/concepts/pipelines/content/pipelinesDetails/pipe
 import { PodStepStateType } from '~/types';
 import useDebounceCallback from '~/utilities/useDebounceCallback';
 import { PipelineTask } from '~/concepts/pipelines/topology';
-import { RuntimeStateKF } from '~/concepts/pipelines/kfTypes';
+import { ExecutionStateKF } from '~/concepts/pipelines/kfTypes';
 
 // TODO: If this gets large enough we should look to make this its own component file
 const LogsTab: React.FC<{ task: PipelineTask }> = ({ task }) => {
   const podName = task.status?.podName;
-  const isFailedPod = task.status?.state === RuntimeStateKF.FAILED;
+  const isFailedPod = task.status?.state === ExecutionStateKF.FAILED;
 
   if (!podName) {
     return <>No content</>;

frontend/src/concepts/pipelines/content/pipelinesDetails/pipelineRun/useExecutionsForPipelineRun.ts

@@ -0,0 +1,21 @@
+import { useExecutionsFromMlmdContext } from '~/concepts/pipelines/apiHooks/mlmd/useExecutionsFromMlmdContext';
+import { usePipelineRunMlmdContext } from '~/concepts/pipelines/apiHooks/mlmd/usePipelineRunMlmdContext';
+import { isPipelineRunFinished } from '~/concepts/pipelines/apiHooks/usePipelineRunById';
+import { PipelineRunKFv2 } from '~/concepts/pipelines/kfTypes';
+import { Execution } from '~/third_party/mlmd';
+import { FAST_POLL_INTERVAL } from '~/utilities/const';
+
+const useExecutionsForPipelineRun = (
+  run: PipelineRunKFv2 | null,
+): [executions: Execution[] | null, loaded: boolean, error?: Error] => {
+  const isFinished = isPipelineRunFinished(run);
+  const refreshRate = isFinished ? 0 : FAST_POLL_INTERVAL;
+  // contextError means mlmd service is not available, no need to check executions
+  const [context, , contextError] = usePipelineRunMlmdContext(run?.run_id, refreshRate);
+  // executionsLoaded is the flag to show the spinner or not
+  const [executions, executionsLoaded] = useExecutionsFromMlmdContext(context, refreshRate);
+
+  return [executions, executionsLoaded, contextError];
+};
+
+export default useExecutionsForPipelineRun;

frontend/src/concepts/pipelines/kfTypes.ts

@@ -466,6 +466,22 @@ export enum RuntimeStateKF {
   PAUSED = 'PAUSED',
 }
 
+export enum ExecutionStateKF {
+  NEW = 'New',
+  RUNNING = 'Running',
+  COMPLETE = 'Complete',
+  CANCELED = 'Canceled',
+  FAILED = 'Failed',
+  CACHED = 'Cached',
+}
+
+export enum ArtifactStateKF {
+  PENDING = 'Pending',
+  LIVE = 'Live',
+  MARKED_FOR_DELETION = 'Marked for deletion',
+  DELETED = 'Deleted',
+}
+
 export const runtimeStateLabels = {
   [RuntimeStateKF.RUNTIME_STATE_UNSPECIFIED]: 'Unspecified',
   [RuntimeStateKF.PENDING]: 'Pending',