[ui] Move stepStats out of RunRootQuery (#26610)

## Summary & Motivation

Move `stepStats` out of the root Run query, since it can be expensive. Instead, query for it at the `RunMetadataProvider`.

This should allow the header, Gantt chart, and logs to load even while the step stats panel of the page continues to load.

## How I Tested These Changes

View a very long-running run. Verify that the Run header, Gantt chart, and logs appear fairly quickly, and the right panel takes a bit longer to load.

## Changelog

[ui] Improve performance of Run page for very long-running runs.

js_modules/dagster-ui/packages/ui-core/src/runs/RunFragments.tsx

@@ -46,20 +46,6 @@ export const RUN_FRAGMENT = gql`
     }
     stepKeysToExecute
     updateTime
-    stepStats {
-      stepKey
-      status
-      startTime
-      endTime
-      attempts {
-        startTime
-        endTime
-      }
-      markers {
-        startTime
-        endTime
-      }
-    }
     ...RunTimingFragment
   }
 

js_modules/dagster-ui/packages/ui-core/src/runs/RunMetadataProvider.tsx

@@ -1,11 +1,17 @@
 import * as React from 'react';
+import {useMemo} from 'react';
 
 import {LogsProviderLogs} from './LogsProvider';
 import {RunContext} from './RunContext';
-import {gql} from '../apollo-client';
+import {gql, useQuery} from '../apollo-client';
 import {flattenOneLevel} from '../util/flattenOneLevel';
 import {RunFragment} from './types/RunFragments.types';
-import {RunMetadataProviderMessageFragment} from './types/RunMetadataProvider.types';
+import {
+  RunMetadataProviderMessageFragment,
+  RunStepStatsFragment,
+  RunStepStatsQuery,
+  RunStepStatsQueryVariables,
+} from './types/RunMetadataProvider.types';
 import {StepEventStatus} from '../graphql/types';
 import {METADATA_ENTRY_FRAGMENT} from '../metadata/MetadataEntryFragment';
 
@@ -103,13 +109,17 @@ export const extractLogCaptureStepsFromLegacySteps = (stepKeys: string[]) => {
 
 const fromTimestamp = (ts: number | null) => (ts ? Math.floor(ts * 1000) : undefined);
 
-function extractMetadataFromRun(run?: RunFragment): IRunMetadataDict {
+function extractMetadataFromRun(
+  run: RunFragment | null = null,
+  stepStats: RunStepStatsFragment['stepStats'] = [],
+): IRunMetadataDict {
   const metadata: IRunMetadataDict = {
     firstLogAt: 0,
     mostRecentLogAt: 0,
     globalMarkers: [],
     steps: {},
   };
+
   if (!run) {
     return metadata;
   }
@@ -120,7 +130,7 @@ function extractMetadataFromRun(run?: RunFragment): IRunMetadataDict {
     metadata.exitedAt = fromTimestamp(run.endTime);
   }
 
-  run.stepStats.forEach((stepStat) => {
+  stepStats.forEach((stepStat) => {
     metadata.steps[stepStat.stepKey] = {
       // state:
       // current state
@@ -370,7 +380,18 @@ interface IRunMetadataProviderProps {
 
 export const RunMetadataProvider = ({logs, children}: IRunMetadataProviderProps) => {
   const run = React.useContext(RunContext);
-  const runMetadata = React.useMemo(() => extractMetadataFromRun(run), [run]);
+
+  // Step stats can be expensive to load, so we separate them from the main run query.
+  const {data} = useQuery<RunStepStatsQuery, RunStepStatsQueryVariables>(RUN_STEP_STATS_QUERY, {
+    variables: run ? {runId: run.id} : undefined,
+    skip: !run,
+  });
+
+  const stepStats = useMemo(() => {
+    return data?.pipelineRunOrError.__typename === 'Run' ? data.pipelineRunOrError.stepStats : [];
+  }, [data]);
+
+  const runMetadata = React.useMemo(() => extractMetadataFromRun(run, stepStats), [run, stepStats]);
   const metadata = React.useMemo(
     () =>
       logs.loading ? runMetadata : extractMetadataFromLogs(flattenOneLevel(logs.allNodeChunks)),
@@ -379,6 +400,35 @@ export const RunMetadataProvider = ({logs, children}: IRunMetadataProviderProps)
   return <>{children(metadata)}</>;
 };
 
+const RUN_STEP_STATS_QUERY = gql`
+  query RunStepStatsQuery($runId: ID!) {
+    pipelineRunOrError(runId: $runId) {
+      ... on Run {
+        id
+        ...RunStepStatsFragment
+      }
+    }
+  }
+
+  fragment RunStepStatsFragment on Run {
+    id
+    stepStats {
+      stepKey
+      status
+      startTime
+      endTime
+      attempts {
+        startTime
+        endTime
+      }
+      markers {
+        startTime
+        endTime
+      }
+    }
+  }
+`;
+
 export const RUN_METADATA_PROVIDER_MESSAGE_FRAGMENT = gql`
   fragment RunMetadataProviderMessageFragment on DagsterRunEvent {
     ... on MessageEvent {