Examples with BranchNode org.flyte.api.v1.BranchNode used on opensource projects

Example 1 with BranchNode

use of org.flyte.api.v1.BranchNode in project flytekit-java by flyteorg.

the class ProjectClosureTest method testCollectSubWorkflows.

@Test
public void testCollectSubWorkflows() {
    TypedInterface emptyInterface = TypedInterface.builder().inputs(ImmutableMap.of()).outputs(ImmutableMap.of()).build();
    WorkflowMetadata emptyMetadata = WorkflowMetadata.builder().build();
    PartialWorkflowIdentifier rewrittenSubWorkflowRef = PartialWorkflowIdentifier.builder().project("project").name("name").version("version").domain("domain").build();
    WorkflowIdentifier subWorkflowRef = WorkflowIdentifier.builder().project("project").name("name").version("version").domain("domain").build();
    WorkflowIdentifier otherSubWorkflowRef = WorkflowIdentifier.builder().project("project").name("other-name").version("version").domain("domain").build();
    PartialWorkflowIdentifier rewrittenNestedSubWorkflowRef = PartialWorkflowIdentifier.builder().project("project").name("nested").version("version").domain("domain").build();
    WorkflowIdentifier nestedSubWorkflowRef = WorkflowIdentifier.builder().project("project").name("nested").version("version").domain("domain").build();
    PartialWorkflowIdentifier rewrittenNestedOtherSubWorkflowRef = PartialWorkflowIdentifier.builder().project("project").name("nested-other").version("version").domain("domain").build();
    WorkflowIdentifier nestedOtherSubWorkflowRef = WorkflowIdentifier.builder().project("project").name("nested-other").version("version").domain("domain").build();
    WorkflowNode workflowNode = WorkflowNode.builder().reference(WorkflowNode.Reference.ofSubWorkflowRef(rewrittenSubWorkflowRef)).build();
    WorkflowNode nestedWorkflowNode = WorkflowNode.builder().reference(WorkflowNode.Reference.ofSubWorkflowRef(rewrittenNestedSubWorkflowRef)).build();
    WorkflowNode nestedOtherWorkflowNode = WorkflowNode.builder().reference(WorkflowNode.Reference.ofSubWorkflowRef(rewrittenNestedOtherSubWorkflowRef)).build();
    WorkflowTemplate emptyWorkflowTemplate = WorkflowTemplate.builder().interface_(emptyInterface).metadata(emptyMetadata).nodes(ImmutableList.of()).outputs(ImmutableList.of()).build();
    WorkflowTemplate nestedWorkflowTemplate = WorkflowTemplate.builder().interface_(emptyInterface).metadata(emptyMetadata).nodes(ImmutableList.of(Node.builder().id("nested-node").inputs(ImmutableList.of()).upstreamNodeIds(ImmutableList.of()).workflowNode(nestedOtherWorkflowNode).build())).outputs(ImmutableList.of()).build();
    Operand opTrue = Operand.ofPrimitive(Primitive.ofBooleanValue(true));
    BooleanExpression exprTrue = BooleanExpression.ofComparison(ComparisonExpression.builder().leftValue(opTrue).rightValue(opTrue).operator(ComparisonExpression.Operator.EQ).build());
    List<Node> nodes = ImmutableList.of(Node.builder().id("node-1").inputs(ImmutableList.of()).upstreamNodeIds(ImmutableList.of()).workflowNode(workflowNode).build(), // Same sub-workflow
    Node.builder().id("node-2").inputs(ImmutableList.of()).upstreamNodeIds(ImmutableList.of()).workflowNode(workflowNode).build(), // Sub-workflow which has a nested sub-workflow in branch (nestedOtherWorkflowNode)
    Node.builder().id("node-3").inputs(ImmutableList.of()).upstreamNodeIds(ImmutableList.of()).branchNode(BranchNode.builder().ifElse(IfElseBlock.builder().case_(IfBlock.builder().condition(exprTrue).thenNode(Node.builder().id("node-4").inputs(ImmutableList.of()).upstreamNodeIds(ImmutableList.of()).workflowNode(nestedWorkflowNode).build()).build()).other(ImmutableList.of()).build()).build()).build());
    // nestedOtherWorkflowNode is not in the previous list because
    // that node belongs to the template of a sub-workflow
    Map<WorkflowIdentifier, WorkflowTemplate> allWorkflows = ImmutableMap.of(subWorkflowRef, emptyWorkflowTemplate, otherSubWorkflowRef, emptyWorkflowTemplate, nestedSubWorkflowRef, nestedWorkflowTemplate, nestedOtherSubWorkflowRef, emptyWorkflowTemplate);
    Map<WorkflowIdentifier, WorkflowTemplate> collectedSubWorkflows = ProjectClosure.collectSubWorkflows(nodes, allWorkflows);
    assertThat(collectedSubWorkflows, equalTo(ImmutableMap.of(subWorkflowRef, emptyWorkflowTemplate, nestedSubWorkflowRef, nestedWorkflowTemplate, nestedOtherSubWorkflowRef, emptyWorkflowTemplate)));
}

Example 2 with BranchNode

use of org.flyte.api.v1.BranchNode in project flytekit-java by flyteorg.

the class ProtoUtilTest method shouldSerializeBranchNode.

@Test
void shouldSerializeBranchNode() {
    ComparisonExpression comparison = ComparisonExpression.builder().operator(ComparisonExpression.Operator.EQ).leftValue(Operand.ofVar("a")).rightValue(Operand.ofVar("b")).build();
    Condition.ComparisonExpression comparisonProto = Condition.ComparisonExpression.newBuilder().setOperator(Condition.ComparisonExpression.Operator.EQ).setLeftValue(Condition.Operand.newBuilder().setVar("a").build()).setRightValue(Condition.Operand.newBuilder().setVar("b").build()).build();
    IfBlock ifBlock = IfBlock.builder().condition(BooleanExpression.ofComparison(comparison)).thenNode(Node.builder().id("node-1").upstreamNodeIds(ImmutableList.of()).inputs(ImmutableList.of()).build()).build();
    Workflow.IfBlock ifBlockProto = Workflow.IfBlock.newBuilder().setThenNode(Workflow.Node.newBuilder().setId("node-1").build()).setCondition(Condition.BooleanExpression.newBuilder().setComparison(comparisonProto).build()).build();
    IfElseBlock ifElse = IfElseBlock.builder().case_(ifBlock).other(ImmutableList.of(ifBlock)).elseNode(ifBlock.thenNode()).build();
    Workflow.IfElseBlock ifElseProto = Workflow.IfElseBlock.newBuilder().setCase(ifBlockProto).addOther(ifBlockProto).setElseNode(ifBlockProto.getThenNode()).build();
    BranchNode branchNode = BranchNode.builder().ifElse(ifElse).build();
    Workflow.BranchNode branchNodeProto = Workflow.BranchNode.newBuilder().setIfElse(ifElseProto).build();
    assertThat(ProtoUtil.serialize(branchNode), equalTo(branchNodeProto));
}

Example 3 with BranchNode

use of org.flyte.api.v1.BranchNode in project flytekit-java by flyteorg.

the class IdentifierRewriteTest method shouldRewriteBranchNodes.

@Test
void shouldRewriteBranchNodes() {
    ComparisonExpression comparison = ComparisonExpression.builder().operator(ComparisonExpression.Operator.EQ).leftValue(Operand.ofVar("a")).rightValue(Operand.ofVar("b")).build();
    BooleanExpression condition = BooleanExpression.ofComparison(comparison);
    PartialTaskIdentifier partialReference = PartialTaskIdentifier.builder().name("task-name").build();
    PartialTaskIdentifier rewrittenReference = PartialTaskIdentifier.builder().name("task-name").domain("rewritten-domain").version("rewritten-version").project("rewritten-project").build();
    TaskNode partialTaskNode = TaskNode.builder().referenceId(partialReference).build();
    TaskNode rewrittenTaskNode = TaskNode.builder().referenceId(rewrittenReference).build();
    Node partialNode = Node.builder().id("node-1").inputs(ImmutableList.of()).upstreamNodeIds(ImmutableList.of()).taskNode(partialTaskNode).build();
    Node rewrittenNode = Node.builder().id("node-1").inputs(ImmutableList.of()).upstreamNodeIds(ImmutableList.of()).taskNode(rewrittenTaskNode).build();
    IfBlock partialIfBlock = IfBlock.builder().condition(condition).thenNode(partialNode).build();
    IfBlock rewrittenIfBlock = IfBlock.builder().condition(condition).thenNode(rewrittenNode).build();
    BranchNode partialBranchNode = BranchNode.builder().ifElse(IfElseBlock.builder().case_(partialIfBlock).other(ImmutableList.of(partialIfBlock)).elseNode(partialNode).build()).build();
    BranchNode rewrittenBranchNode = BranchNode.builder().ifElse(IfElseBlock.builder().case_(rewrittenIfBlock).other(ImmutableList.of(rewrittenIfBlock)).elseNode(rewrittenNode).build()).build();
    assertThat(rewriter.visitor().visitBranchNode(partialBranchNode), equalTo(rewrittenBranchNode));
}

Example 4 with BranchNode

use of org.flyte.api.v1.BranchNode in project flytekit-java by flyteorg.

the class SdkWorkflowBuilderTest method testConditionalWorkflowIdl.

@Test
void testConditionalWorkflowIdl() {
    SdkWorkflowBuilder builder = new SdkWorkflowBuilder();
    new ConditionalWorkflow().expand(builder);
    Node caseNode = Node.builder().id("neq").taskNode(TaskNode.builder().referenceId(PartialTaskIdentifier.builder().name("org.flyte.flytekit.SdkWorkflowBuilderTest$MultiplicationTask").build()).build()).inputs(asList(Binding.builder().var_("a").binding(BindingData.ofOutputReference(OutputReference.builder().var("in").nodeId(Node.START_NODE_ID).build())).build(), Binding.builder().var_("b").binding(BindingData.ofScalar(Scalar.ofPrimitive(Primitive.ofIntegerValue(2L)))).build())).upstreamNodeIds(emptyList()).build();
    IfElseBlock ifElse = IfElseBlock.builder().case_(IfBlock.builder().condition(BooleanExpression.ofComparison(ComparisonExpression.builder().leftValue(Operand.ofVar("$0")).rightValue(Operand.ofPrimitive(Primitive.ofIntegerValue(2L))).operator(ComparisonExpression.Operator.NEQ).build())).thenNode(caseNode).build()).error(NodeError.builder().message("No cases matched").failedNodeId("square").build()).other(emptyList()).build();
    Node expectedNode = Node.builder().id("square").branchNode(BranchNode.builder().ifElse(ifElse).build()).inputs(singletonList(Binding.builder().var_("$0").binding(BindingData.ofOutputReference(OutputReference.builder().var("in").nodeId(Node.START_NODE_ID).build())).build())).upstreamNodeIds(emptyList()).build();
    WorkflowTemplate expected = WorkflowTemplate.builder().metadata(WorkflowMetadata.builder().build()).interface_(expectedInterface()).outputs(expectedOutputs()).nodes(singletonList(expectedNode)).build();
    assertEquals(expected, builder.toIdlTemplate());
}

Example 5 with BranchNode

use of org.flyte.api.v1.BranchNode in project flytekit-java by flyteorg.

the class SdkBranchNode method toIdl.

@Override
public Node toIdl() {
    NodeError nodeError = NodeError.builder().failedNodeId(nodeId).message("No cases matched").build();
    Map<String, Binding> extraInputs = new HashMap<>();
    @Var IfElseBlock ifElseBlock = IfBlockIdl.toIdl(ifElse, extraInputs);
    if (ifElseBlock.elseNode() == null) {
        ifElseBlock = ifElseBlock.toBuilder().error(nodeError).build();
    }
    return Node.builder().id(nodeId).branchNode(BranchNode.builder().ifElse(ifElseBlock).build()).inputs(unmodifiableList(new ArrayList<>(extraInputs.values()))).upstreamNodeIds(upstreamNodeIds).build();
}