Handle conditional in dynamic (#301)

Signed-off-by: Andres Gomez Ferrer <andresgomezfrr@users.noreply.github.com> Co-authored-by: Andres Gomez Ferrer <andresgomezfrr@users.noreply.github.com>
flyteorg · Jul 1, 2024 · 2beaff7 · 2beaff7
1 parent 4b98408
commit 2beaff7
Showing 1 changed file with 24 additions and 0 deletions.
diff --git a/jflyte/src/main/java/org/flyte/jflyte/ExecuteDynamicWorkflow.java b/jflyte/src/main/java/org/flyte/jflyte/ExecuteDynamicWorkflow.java
@@ -24,6 +24,7 @@
 
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Maps;
+import java.util.ArrayList;
 import java.util.Collection;
 import java.util.HashMap;
 import java.util.List;
@@ -41,6 +42,7 @@
 import org.flyte.api.v1.DynamicJobSpec;
 import org.flyte.api.v1.DynamicWorkflowTask;
 import org.flyte.api.v1.DynamicWorkflowTaskRegistrar;
+import org.flyte.api.v1.IfBlock;
 import org.flyte.api.v1.Literal;
 import org.flyte.api.v1.NamedEntityIdentifier;
 import org.flyte.api.v1.Node;
@@ -280,6 +282,28 @@ private static List<Node> collectAllUsedTaskTemplates(
             flyteAdminClient,
             cache);
 
+    // collect task templates used by conditionals
+    spec.nodes().stream()
+        .filter(node -> node.branchNode() != null)
+        .forEach(
+            node -> {
+              List<Node> nodes = new ArrayList<>();
+              nodes.add(node.branchNode().ifElse().case_().thenNode());
+              nodes.add(node.branchNode().ifElse().elseNode());
+              nodes.addAll(
+                  node.branchNode().ifElse().other().stream()
+                      .map(IfBlock::thenNode)
+                      .collect(toList()));
+
+              collectTaskTemplates(
+                  nodes,
+                  nodesRewriter,
+                  allUsedTaskTemplates,
+                  allTaskTemplates,
+                  flyteAdminClient,
+                  cache);
+            });
+
     // collect task templates used by subworkflows
     allUsedSubWorkflows
         .values()