gate node in dynamic task (#729)

* fix: gate node in dynamic Signed-off-by: James <james@union.ai> * fix: upgraded version Signed-off-by: James <james@union.ai> * fix: compiledNode in PausedTasksComponent and ExecutionDetailsActions Signed-off-by: James <james@union.ai> * fix: upgrade version Signed-off-by: James <james@union.ai> * chore: lockfile Signed-off-by: Carina Ursu <carina@union.ai> * fix: remove caching for workflow closure Signed-off-by: James <james@union.ai> * fix: break link between cache and context state Signed-off-by: James <james@union.ai> --------- Signed-off-by: James <james@union.ai> Signed-off-by: Carina Ursu <carina@union.ai> Co-authored-by: Carina Ursu <carina@union.ai> Signed-off-by: 4nalog <4nalog@protonmail.com>
flyteorg · Mar 30, 2023 · 4107f72 · 4107f72
1 parent 77932de
commit 4107f72
Show file tree

Hide file tree

Showing 12 changed files with 75 additions and 52 deletions.
diff --git a/packages/console/src/components/Executions/ExecutionDetails/ExecutionDetailsActions.tsx b/packages/console/src/components/Executions/ExecutionDetails/ExecutionDetailsActions.tsx
@@ -12,13 +12,15 @@ import {
 import { literalsToLiteralValueMap } from 'components/Launch/LaunchForm/utils';
 import { TaskInitialLaunchParameters } from 'components/Launch/LaunchForm/types';
 import { NodeExecutionPhase } from 'models/Execution/enums';
+import { extractCompiledNodes } from 'components/hooks/utils';
 import Close from '@material-ui/icons/Close';
 import { useEffect, useState } from 'react';
 import classnames from 'classnames';
 import { NodeExecutionDetails } from '../types';
 import t from './strings';
 import { ExecutionNodeDeck } from './ExecutionNodeDeck';
 import { useNodeExecutionContext } from '../contextProvider/NodeExecutionDetails';
+import { NodeExecutionsByIdContext } from '../contexts';
 
 const useStyles = makeStyles((theme: Theme) => {
   return {
@@ -94,9 +96,14 @@ export const ExecutionDetailsActions = ({
   const execution = useNodeExecution(nodeExecutionId);
   const { compiledWorkflowClosure } = useNodeExecutionContext();
   const id = details?.taskTemplate?.id;
-  const compiledNode = (
-    compiledWorkflowClosure?.primary.template.nodes ?? []
-  ).find(node => node.id === nodeExecutionId.nodeId);
+  const { nodeExecutionsById } = React.useContext(NodeExecutionsByIdContext);
+
+  const compiledNode = extractCompiledNodes(compiledWorkflowClosure).find(
+    node =>
+      node.id ===
+        nodeExecutionsById[nodeExecutionId.nodeId]?.metadata?.specNodeId ||
+      node.id === nodeExecutionId.nodeId,
+  );
 
   useEffect(() => {
     if (!id) {

diff --git a/...s/console/src/components/Executions/ExecutionDetails/NodeExecutionDetailsPanelContent.tsx b/...s/console/src/components/Executions/ExecutionDetails/NodeExecutionDetailsPanelContent.tsx
@@ -38,6 +38,7 @@ import { TaskVersionDetailsLink } from 'components/Entities/VersionDetails/Versi
 import { Identifier } from 'models/Common/types';
 import { isMapTaskV1 } from 'models/Task/utils';
 import { merge } from 'lodash';
+import { extractCompiledNodes } from 'components/hooks/utils';
 import { NodeExecutionCacheStatus } from '../NodeExecutionCacheStatus';
 import {
   makeListTaskExecutionsQuery,
@@ -264,9 +265,11 @@ export const NodeExecutionDetailsPanelContent: React.FC<
     NodeExecutionsByIdContext,
   );
   const isGateNode = isNodeGateNode(
-    compiledWorkflowClosure?.primary.template.nodes ?? [],
-    nodeExecutionId,
+    extractCompiledNodes(compiledWorkflowClosure),
+    nodeExecutionsById[nodeExecutionId.nodeId]?.metadata?.specNodeId ||
+      nodeExecutionId.nodeId,
   );
+
   const [nodeExecutionLoading, setNodeExecutionLoading] =
     useState<boolean>(false);
 
@@ -440,7 +443,7 @@ export const NodeExecutionDetailsPanelContent: React.FC<
 
   const frontendPhase = useMemo(
     () => getNodeFrontendPhase(nodePhase, isGateNode),
-    [nodePhase],
+    [nodePhase, isGateNode],
   );
 
   const isRunningPhase = useMemo(

diff --git a/packages/console/src/components/Executions/Tables/NodeExecutionActions.tsx b/packages/console/src/components/Executions/Tables/NodeExecutionActions.tsx
@@ -12,6 +12,7 @@ import { TaskInitialLaunchParameters } from 'components/Launch/LaunchForm/types'
 import { literalsToLiteralValueMap } from 'components/Launch/LaunchForm/utils';
 import { useContext, useEffect, useState } from 'react';
 import { NodeExecutionPhase } from 'models/Execution/enums';
+import { extractCompiledNodes } from 'components/hooks/utils';
 import { useNodeExecutionContext } from '../contextProvider/NodeExecutionDetails';
 import { NodeExecutionDetails } from '../types';
 import t from './strings';
@@ -42,13 +43,16 @@ export const NodeExecutionActions = ({
   const id = nodeExecutionDetails?.taskTemplate?.id;
 
   const isGateNode = isNodeGateNode(
-    compiledWorkflowClosure?.primary.template.nodes ?? [],
-    execution.id,
+    extractCompiledNodes(compiledWorkflowClosure),
+    execution.metadata?.specNodeId || execution.id.nodeId,
   );
+
   const phase = getNodeFrontendPhase(execution.closure.phase, isGateNode);
-  const compiledNode = (
-    compiledWorkflowClosure?.primary.template.nodes ?? []
-  ).find(node => node.id === execution.id.nodeId);
+  const compiledNode = extractCompiledNodes(compiledWorkflowClosure).find(
+    node =>
+      node.id === execution.metadata?.specNodeId ||
+      node.id === execution.id.nodeId,
+  );
 
   useEffect(() => {
     let isCurrent = true;

diff --git a/packages/console/src/components/Executions/Tables/NodeExecutionsTable.tsx b/packages/console/src/components/Executions/Tables/NodeExecutionsTable.tsx
@@ -9,6 +9,7 @@ import { dateToTimestamp } from 'common/utils';
 import React, { useMemo, useEffect, useState, useContext } from 'react';
 import { useQueryClient } from 'react-query';
 import { merge, eq } from 'lodash';
+import { extractCompiledNodes } from 'components/hooks/utils';
 import { ExecutionsTableHeader } from './ExecutionsTableHeader';
 import { generateColumns } from './nodeExecutionColumns';
 import { NoExecutionsContent } from './NoExecutionsContent';
@@ -58,13 +59,10 @@ export const NodeExecutionsTable: React.FC<NodeExecutionsTableProps> = ({
 
   const columnStyles = useColumnStyles();
   // Memoizing columns so they won't be re-generated unless the styles change
+  const compiledNodes = extractCompiledNodes(compiledWorkflowClosure);
   const columns = useMemo(
-    () =>
-      generateColumns(
-        columnStyles,
-        compiledWorkflowClosure?.primary.template.nodes ?? [],
-      ),
-    [columnStyles],
+    () => generateColumns(columnStyles, compiledNodes),
+    [columnStyles, compiledNodes],
   );
 
   useEffect(() => {

diff --git a/packages/console/src/components/Executions/Tables/nodeExecutionColumns.tsx b/packages/console/src/components/Executions/Tables/nodeExecutionColumns.tsx
@@ -105,7 +105,11 @@ export function generateColumns(
     },
     {
       cellRenderer: ({ execution }) => {
-        const isGateNode = isNodeGateNode(nodes, execution.id);
+        const isGateNode = isNodeGateNode(
+          nodes,
+          execution.metadata?.specNodeId || execution.id.nodeId,
+        );
+
         const phase = getNodeFrontendPhase(
           execution.closure?.phase ?? NodeExecutionPhase.UNDEFINED,
           isGateNode,

diff --git a/packages/console/src/components/Executions/contextProvider/NodeExecutionDetails/index.tsx b/packages/console/src/components/Executions/contextProvider/NodeExecutionDetails/index.tsx
@@ -96,12 +96,12 @@ export const NodeExecutionDetailsContextProvider = (props: ProviderProps) => {
         name,
         version,
       };
-      const workflow = await fetchWorkflow(queryClient, workflowId);
-      if (!workflow) {
+      const result = await fetchWorkflow(queryClient, workflowId);
+      if (!result) {
         resetState();
         return;
       }
-
+      const workflow = JSON.parse(JSON.stringify(result));
       const tree = createExecutionDetails(workflow);
       if (isCurrent) {
         setClosure(workflow.closure?.compiledWorkflow ?? null);

diff --git a/packages/console/src/components/Executions/test/utils.test.ts b/packages/console/src/components/Executions/test/utils.test.ts
@@ -180,34 +180,23 @@ describe('isNodeGateNode', () => {
   const executionId = { project: 'project', domain: 'domain', name: 'name' };
 
   it('should return true if nodeId is in the list and has a gateNode field', () => {
-    expect(
-      isNodeGateNode(mockNodesWithGateNode, {
-        nodeId: 'GateNode',
-        executionId,
-      }),
-    ).toBeTruthy();
+    expect(isNodeGateNode(mockNodesWithGateNode, 'GateNode')).toBeTruthy();
   });
 
   it('should return false if nodeId is in the list, but a gateNode field is missing', () => {
-    expect(
-      isNodeGateNode(mockNodes, { nodeId: 'BasicNode', executionId }),
-    ).toBeFalsy();
+    expect(isNodeGateNode(mockNodes, 'BasicNode')).toBeFalsy();
   });
 
   it('should return false if nodeId is not in the list, but has a gateNode field', () => {
-    expect(
-      isNodeGateNode(mockNodes, { nodeId: 'GateNode', executionId }),
-    ).toBeFalsy();
+    expect(isNodeGateNode(mockNodes, 'GateNode')).toBeFalsy();
   });
 
   it('should return false if nodeId is a gateNode, but the list is empty', () => {
-    expect(isNodeGateNode([], { nodeId: 'GateNode', executionId })).toBeFalsy();
+    expect(isNodeGateNode([], 'GateNode')).toBeFalsy();
   });
 
   it('should return false if nodeId is not a gateNode and the list is empty', () => {
-    expect(
-      isNodeGateNode([], { nodeId: 'BasicNode', executionId }),
-    ).toBeFalsy();
+    expect(isNodeGateNode([], 'BasicNode')).toBeFalsy();
   });
 });
 

diff --git a/packages/console/src/components/Executions/utils.ts b/packages/console/src/components/Executions/utils.ts
@@ -203,11 +203,8 @@ export function isExecutionArchived(execution: Execution): boolean {
 }
 
 /** Returns true if current node (by nodeId) has 'gateNode' field in the list of nodes on compiledWorkflowClosure */
-export function isNodeGateNode(
-  nodes: CompiledNode[],
-  executionId: NodeExecutionIdentifier,
-): boolean {
-  const node = nodes.find(n => n.id === executionId.nodeId);
+export function isNodeGateNode(nodes: CompiledNode[], id: string): boolean {
+  const node = nodes.find(n => n.id === id);
   return !!node?.gateNode;
 }
 

diff --git a/packages/console/src/components/flytegraph/ReactFlow/PausedTasksComponent.tsx b/packages/console/src/components/flytegraph/ReactFlow/PausedTasksComponent.tsx
@@ -10,6 +10,7 @@ import { nodeExecutionPhaseConstants } from 'components/Executions/constants';
 import { LaunchFormDialog } from 'components/Launch/LaunchForm/LaunchFormDialog';
 import { useNodeExecutionContext } from 'components/Executions/contextProvider/NodeExecutionDetails';
 import { NodeExecutionsByIdContext } from 'components/Executions/contexts';
+import { extractCompiledNodes } from 'components/hooks/utils';
 import {
   graphButtonContainer,
   graphButtonStyle,
@@ -73,9 +74,12 @@ export const PausedTasksComponent: React.FC<PausedTasksComponentProps> = ({
     setShowResumeForm(true);
   };
 
-  const compiledNode = (
-    compiledWorkflowClosure?.primary.template.nodes ?? []
-  ).find(node => node.id === selectedNodeId);
+  const compiledNode = extractCompiledNodes(compiledWorkflowClosure).find(
+    node =>
+      (selectedNodeId &&
+        node.id === nodeExecutionsById[selectedNodeId]?.metadata?.specNodeId) ||
+      node.id === selectedNodeId,
+  );
 
   const selectedNode = (pausedNodes ?? []).find(
     node => node.id === selectedNodeId,

diff --git a/packages/console/src/components/flytegraph/ReactFlow/ReactFlowGraphComponent.tsx b/packages/console/src/components/flytegraph/ReactFlow/ReactFlowGraphComponent.tsx
@@ -11,6 +11,7 @@ import { dNode } from 'models/Graph/types';
 import { useQueryClient } from 'react-query';
 import { fetchTaskExecutionList } from 'components/Executions/taskExecutionQueries';
 import { isMapTaskV1 } from 'models/Task/utils';
+import { extractCompiledNodes } from 'components/hooks/utils';
 import { ExternalResource, LogsByPhase } from 'models/Execution/types';
 import { getGroupedLogs } from 'components/Executions/TaskExecutionsList/utils';
 import { LargeLoadingSpinner } from 'components/common/LoadingSpinner';
@@ -177,8 +178,8 @@ export const ReactFlowGraphComponent = ({
       if (nodeExecution) {
         const phase = nodeExecution?.closure.phase;
         const isGateNode = isNodeGateNode(
-          compiledWorkflowClosure?.primary.template.nodes ?? [],
-          nodeExecution.id,
+          extractCompiledNodes(compiledWorkflowClosure),
+          nodeExecution.metadata?.specNodeId || nodeExecution.id.nodeId,
         );
         return isGateNode && phase === NodeExecutionPhase.RUNNING;
       }

diff --git a/packages/console/src/components/flytegraph/ReactFlow/customNodeComponents.tsx b/packages/console/src/components/flytegraph/ReactFlow/customNodeComponents.tsx
@@ -12,6 +12,7 @@ import { CacheStatus } from 'components/Executions/CacheStatus';
 import { LaunchFormDialog } from 'components/Launch/LaunchForm/LaunchFormDialog';
 import { useNodeExecutionContext } from 'components/Executions/contextProvider/NodeExecutionDetails';
 import { NodeExecutionsByIdContext } from 'components/Executions/contexts';
+import { extractCompiledNodes } from 'components/hooks/utils';
 import {
   COLOR_GRAPH_BACKGROUND,
   getGraphHandleStyle,
@@ -234,9 +235,11 @@ export const ReactFlowGateNode = ({ data }: RFNode) => {
   const styles = getGraphNodeStyle(nodeType, phase);
   const [showResumeForm, setShowResumeForm] = useState<boolean>(false);
 
-  const compiledNode = (
-    compiledWorkflowClosure?.primary.template.nodes ?? []
-  ).find(node => node.id === nodeExecutionsById[scopedId]?.id?.nodeId);
+  const compiledNode = extractCompiledNodes(compiledWorkflowClosure).find(
+    node =>
+      node.id === nodeExecutionsById[scopedId]?.metadata?.specNodeId ||
+      node.id === nodeExecutionsById[scopedId]?.id?.nodeId,
+  );
 
   const iconStyles: React.CSSProperties = {
     width: '10px',

diff --git a/packages/console/src/components/hooks/utils.ts b/packages/console/src/components/hooks/utils.ts
@@ -1,6 +1,19 @@
-import { GloballyUniqueNode } from 'models/Node/types';
+import { CompiledNode, GloballyUniqueNode } from 'models/Node/types';
 import { TaskTemplate } from 'models/Task/types';
-import { Workflow } from 'models/Workflow/types';
+import { CompiledWorkflowClosure, Workflow } from 'models/Workflow/types';
+
+export function extractCompiledNodes(
+  compiledWorkflowClosure: CompiledWorkflowClosure | null,
+): CompiledNode[] {
+  if (!compiledWorkflowClosure) return [];
+
+  const { primary, subWorkflows = [] } = compiledWorkflowClosure;
+
+  return subWorkflows.reduce(
+    (out, subWorkflow) => [...out, ...subWorkflow.template.nodes],
+    primary.template.nodes,
+  );
+}
 
 export function extractTaskTemplates(workflow: Workflow): TaskTemplate[] {
   if (!workflow.closure || !workflow.closure.compiledWorkflow) {