Local engine classes

Signed-off-by: Nelson Arapé <[email protected]>

flytekit-local-engine/src/test/java/org/flyte/localengine/LocalEngineTest.java

@@ -48,7 +48,6 @@
 import org.flyte.api.v1.WorkflowTemplateRegistrar;
 import org.flyte.flytekit.SdkBindingData;
 import org.flyte.flytekit.SdkRunnableTask;
-import org.flyte.flytekit.SdkTransform;
 import org.flyte.flytekit.SdkWorkflow;
 import org.flyte.flytekit.SdkWorkflowBuilder;
 import org.flyte.flytekit.jackson.JacksonSdkType;
@@ -549,10 +548,11 @@ static <K, V, T extends Registrar<K, V>> Map<K, V> loadAll(
 
   @AutoService(SdkWorkflow.class)
   public static class TestCaseExhaustivenessWorkflow
-      extends SdkWorkflow<TestCaseExhaustivenessWorkflow.NoOpType> {
+      extends SdkWorkflow<
+          TestCaseExhaustivenessWorkflow.NoOpType, TestCaseExhaustivenessWorkflow.NoOpType> {
 
     public TestCaseExhaustivenessWorkflow() {
-      super(JacksonSdkType.of(NoOpType.class));
+      super(JacksonSdkType.of(NoOpType.class), JacksonSdkType.of(NoOpType.class));
     }
 
     @Override
@@ -562,8 +562,8 @@ public void expand(SdkWorkflowBuilder builder) {
           builder
               .apply(
                   "decide",
-                  when("eq_1", eq(ofInteger(1L), x), NoOp.of(x))
-                      .when("eq_2", eq(ofInteger(2L), x), NoOp.of(x)))
+                  when("eq_1", eq(ofInteger(1L), x), new NoOp(), NoOpType.create(x))
+                      .when("eq_2", eq(ofInteger(2L), x), new NoOp(), NoOpType.create(x)))
               .getOutputs()
               .x();
 
@@ -582,10 +582,6 @@ public NoOp() {
       public NoOpType run(NoOpType input) {
         return NoOpType.create(input.x());
       }
-
-      static SdkTransform<NoOpType> of(SdkBindingData<Long> x) {
-        return new NoOp().withInput("x", x);
-      }
     }
 
     @AutoValue

flytekit-local-engine/src/test/java/org/flyte/localengine/examples/CollatzConjectureStepWorkflow.java

@@ -30,27 +30,36 @@
 
 // if x is even, then x/2 else 3x+1
 @AutoService(SdkWorkflow.class)
-public class CollatzConjectureStepWorkflow extends SdkWorkflow<TestUnaryIntegerOutput> {
+public class CollatzConjectureStepWorkflow
+    extends SdkWorkflow<CollatzConjectureStepWorkflow.Input, TestUnaryIntegerOutput> {
 
   public CollatzConjectureStepWorkflow() {
-    super(new TestUnaryIntegerOutput.SdkType());
+    super(
+        JacksonSdkType.of(CollatzConjectureStepWorkflow.Input.class),
+        JacksonSdkType.of(TestUnaryIntegerOutput.class));
+  }
+
+  @AutoValue
+  abstract static class Input {
+    abstract SdkBindingData<Long> x();
+
+    public static Input create(SdkBindingData<Long> x) {
+      return new AutoValue_CollatzConjectureStepWorkflow_Input(x);
+    }
   }
 
   @Override
   public void expand(SdkWorkflowBuilder builder) {
     SdkBindingData<Long> x = builder.inputOfInteger("x");
     SdkBindingData<Boolean> isOdd =
-        builder.apply("is_odd", new IsEvenTask().withInput("x", x)).getOutputs().res();
+        builder.apply("is_odd", new IsEvenTask(), IsEvenTask.Input.create(x)).getOutputs().res();
 
     SdkBindingData<Long> nextX =
         builder
             .apply(
                 "decide",
-                when(
-                        "was_even",
-                        isTrue(isOdd),
-                        new Divide().withInput("num", x).withInput("den", ofInteger(2L)))
-                    .otherwise("was_odd", new ThreeXPlusOne().withInput("x", x)))
+                when("was_even", isTrue(isOdd), new Divide(), Divide.Input.create(x, ofInteger(2L)))
+                    .otherwise("was_odd", new ThreeXPlusOne(), ThreeXPlusOne.Input.create(x)))
             .getOutputs()
             .o();
 
@@ -67,17 +76,16 @@ public IsEvenTask() {
 
     @Override
     public IsEvenTask.Output run(IsEvenTask.Input input) {
-      return IsEvenTask.Output.create(input.x().get() % 2 == 0);
+      return IsEvenTask.Output.create(SdkBindingData.ofBoolean(input.x().get() % 2 == 0));
     }
 
     @AutoValue
     public abstract static class Input {
 
       public abstract SdkBindingData<Long> x();
 
-      public static Input create(Long x) {
-        return new AutoValue_CollatzConjectureStepWorkflow_IsEvenTask_Input(
-            SdkBindingData.ofInteger(x));
+      public static Input create(SdkBindingData<Long> x) {
+        return new AutoValue_CollatzConjectureStepWorkflow_IsEvenTask_Input(x);
       }
     }
 
@@ -86,9 +94,8 @@ public abstract static class Output {
 
       public abstract SdkBindingData<Boolean> res();
 
-      public static Output create(boolean res) {
-        return new AutoValue_CollatzConjectureStepWorkflow_IsEvenTask_Output(
-            SdkBindingData.ofBoolean(res));
+      public static Output create(SdkBindingData<Boolean> res) {
+        return new AutoValue_CollatzConjectureStepWorkflow_IsEvenTask_Output(res);
       }
     }
   }
@@ -98,7 +105,7 @@ public static class Divide extends SdkRunnableTask<Divide.Input, TestUnaryIntege
     private static final long serialVersionUID = -526903889896397227L;
 
     public Divide() {
-      super(JacksonSdkType.of(Divide.Input.class), new TestUnaryIntegerOutput.SdkType());
+      super(JacksonSdkType.of(Divide.Input.class), JacksonSdkType.of(TestUnaryIntegerOutput.class));
     }
 
     @Override
@@ -113,9 +120,8 @@ public abstract static class Input {
 
       public abstract SdkBindingData<Long> den();
 
-      public static Input create(long num, long den) {
-        return new AutoValue_CollatzConjectureStepWorkflow_Divide_Input(
-            SdkBindingData.ofInteger(num), SdkBindingData.ofInteger(den));
+      public static Input create(SdkBindingData<Long> num, SdkBindingData<Long> den) {
+        return new AutoValue_CollatzConjectureStepWorkflow_Divide_Input(num, den);
       }
     }
 
@@ -124,9 +130,8 @@ public abstract static class Output {
 
       public abstract SdkBindingData<Long> res();
 
-      public static Output create(long res) {
-        return new AutoValue_CollatzConjectureStepWorkflow_Divide_Output(
-            SdkBindingData.ofInteger(res));
+      public static Output create(SdkBindingData<Long> res) {
+        return new AutoValue_CollatzConjectureStepWorkflow_Divide_Output(res);
       }
     }
   }
@@ -138,7 +143,9 @@ public static class ThreeXPlusOne
     private static final long serialVersionUID = 932934331328064751L;
 
     public ThreeXPlusOne() {
-      super(JacksonSdkType.of(ThreeXPlusOne.Input.class), new TestUnaryIntegerOutput.SdkType());
+      super(
+          JacksonSdkType.of(ThreeXPlusOne.Input.class),
+          JacksonSdkType.of(TestUnaryIntegerOutput.class));
     }
 
     @Override
@@ -150,9 +157,8 @@ public TestUnaryIntegerOutput run(ThreeXPlusOne.Input input) {
     public abstract static class Input {
       public abstract SdkBindingData<Long> x();
 
-      public static Input create(long x) {
-        return new AutoValue_CollatzConjectureStepWorkflow_ThreeXPlusOne_Input(
-            SdkBindingData.ofInteger(x));
+      public static Input create(SdkBindingData<Long> x) {
+        return new AutoValue_CollatzConjectureStepWorkflow_ThreeXPlusOne_Input(x);
       }
     }
   }

flytekit-local-engine/src/test/java/org/flyte/localengine/examples/FibonacciWorkflow.java

@@ -24,10 +24,13 @@
 import org.flyte.flytekit.jackson.JacksonSdkType;
 
 @AutoService(SdkWorkflow.class)
-public class FibonacciWorkflow extends SdkWorkflow<FibonacciWorkflow.Output> {
+public class FibonacciWorkflow
+    extends SdkWorkflow<FibonacciWorkflow.Input, FibonacciWorkflow.Output> {
 
   public FibonacciWorkflow() {
-    super(JacksonSdkType.of(FibonacciWorkflow.Output.class));
+    super(
+        JacksonSdkType.of(FibonacciWorkflow.Input.class),
+        JacksonSdkType.of(FibonacciWorkflow.Output.class));
   }
 
   @Override
@@ -36,33 +39,28 @@ public void expand(SdkWorkflowBuilder builder) {
     SdkBindingData<Long> fib1 = builder.inputOfInteger("fib1");
 
     SdkBindingData<Long> fib2 =
-        builder
-            .apply("fib-2", new SumTask().withInput("a", fib0).withInput("b", fib1))
-            .getOutputs()
-            .o();
+        builder.apply("fib-2", new SumTask(), SumTask.Input.create(fib0, fib1)).getOutputs().o();
 
     SdkBindingData<Long> fib3 =
-        builder
-            .apply("fib-3", new SumTask().withInput("a", fib1).withInput("b", fib2))
-            .getOutputs()
-            .o();
+        builder.apply("fib-3", new SumTask(), SumTask.Input.create(fib1, fib2)).getOutputs().o();
 
     SdkBindingData<Long> fib4 =
-        builder
-            .apply("fib-4", new SumTask().withInput("a", fib2).withInput("b", fib3))
-            .getOutputs()
-            .o();
+        builder.apply("fib-4", new SumTask(), SumTask.Input.create(fib2, fib3)).getOutputs().o();
 
     SdkBindingData<Long> fib5 =
-        builder
-            .apply("fib-5", new SumTask().withInput("a", fib3).withInput("b", fib4))
-            .getOutputs()
-            .o();
+        builder.apply("fib-5", new SumTask(), SumTask.Input.create(fib3, fib4)).getOutputs().o();
 
     builder.output("fib4", fib4);
     builder.output("fib5", fib5);
   }
 
+  @AutoValue
+  public abstract static class Input {
+    public abstract SdkBindingData<Long> fib0();
+
+    public abstract SdkBindingData<Long> fib1();
+  }
+
   @AutoValue
   public abstract static class Output {
     public abstract SdkBindingData<Long> fib4();

flytekit-local-engine/src/test/java/org/flyte/localengine/examples/InnerSubWorkflow.java

@@ -23,20 +23,17 @@
 import org.flyte.flytekit.jackson.JacksonSdkType;
 
 @AutoService(SdkWorkflow.class)
-public class InnerSubWorkflow extends SdkWorkflow<TestUnaryIntegerOutput> {
+public class InnerSubWorkflow extends SdkWorkflow<SumTask.Input, TestUnaryIntegerOutput> {
   public InnerSubWorkflow() {
-    super(JacksonSdkType.of(TestUnaryIntegerOutput.class));
+    super(JacksonSdkType.of(SumTask.Input.class), JacksonSdkType.of(TestUnaryIntegerOutput.class));
   }
 
   @Override
   public void expand(SdkWorkflowBuilder builder) {
     SdkBindingData<Long> a = builder.inputOfInteger("a");
     SdkBindingData<Long> b = builder.inputOfInteger("b");
     SdkBindingData<Long> c =
-        builder
-            .apply("inner-sum-a-b", new SumTask().withInput("a", a).withInput("b", b))
-            .getOutputs()
-            .o();
+        builder.apply("inner-sum-a-b", new SumTask(), SumTask.Input.create(a, b)).getOutputs().o();
 
     builder.output("o", c);
   }

flytekit-local-engine/src/test/java/org/flyte/localengine/examples/ListTask.java

@@ -33,24 +33,24 @@ public ListTask() {
 
   @Override
   public Output run(Input input) {
-    return Output.create(input.list().get());
+    return Output.create(SdkBindingData.ofIntegerCollection(input.list().get()));
   }
 
   @AutoValue
   public abstract static class Input {
     public abstract SdkBindingData<List<Long>> list();
 
-    public static Input create(List<Long> list) {
-      return new AutoValue_ListTask_Input(SdkBindingData.ofIntegerCollection(list));
+    public static Input create(SdkBindingData<List<Long>> list) {
+      return new AutoValue_ListTask_Input(list);
     }
   }
 
   @AutoValue
   public abstract static class Output {
     public abstract SdkBindingData<List<Long>> list();
 
-    public static Output create(List<Long> list) {
-      return new AutoValue_ListTask_Output(SdkBindingData.ofIntegerCollection(list));
+    public static Output create(SdkBindingData<List<Long>> list) {
+      return new AutoValue_ListTask_Output(list);
     }
   }
 }

flytekit-local-engine/src/test/java/org/flyte/localengine/examples/ListWorkflow.java

@@ -16,37 +16,39 @@
  */
 package org.flyte.localengine.examples;
 
+import static org.flyte.flytekit.SdkBindingData.ofInteger;
+
 import com.google.auto.service.AutoService;
 import java.util.List;
 import org.flyte.api.v1.LiteralType;
 import org.flyte.api.v1.SimpleType;
 import org.flyte.flytekit.SdkBindingData;
 import org.flyte.flytekit.SdkNode;
+import org.flyte.flytekit.SdkTypes;
 import org.flyte.flytekit.SdkWorkflow;
 import org.flyte.flytekit.SdkWorkflowBuilder;
 import org.flyte.flytekit.jackson.JacksonSdkType;
-import org.flyte.localengine.ImmutableList;
 
 @AutoService(SdkWorkflow.class)
-public class ListWorkflow extends SdkWorkflow<ListTask.Output> {
+public class ListWorkflow extends SdkWorkflow<Void, ListTask.Output> {
   public ListWorkflow() {
-    super(JacksonSdkType.of(ListTask.Output.class));
+    super(SdkTypes.nulls(), JacksonSdkType.of(ListTask.Output.class));
   }
 
   @Override
   public void expand(SdkWorkflowBuilder builder) {
     SdkNode<TestUnaryIntegerOutput> sum1 =
-        builder.apply("sum-1", new SumTask().withInput("a", 1).withInput("b", 2));
+        builder.apply("sum-1", new SumTask(), SumTask.Input.create(ofInteger(1), ofInteger(2)));
     SdkNode<TestUnaryIntegerOutput> sum2 =
-        builder.apply("sum-2", new SumTask().withInput("a", 3).withInput("b", 4));
+        builder.apply("sum-2", new SumTask(), SumTask.Input.create(ofInteger(3), ofInteger(4)));
 
     SdkBindingData<List<Long>> list =
         SdkBindingData.ofBindingCollection(
             LiteralType.ofCollectionType(LiteralType.ofSimpleType(SimpleType.INTEGER)),
-            ImmutableList.of(sum1.getOutputs().o(), sum2.getOutputs().o()));
+            List.of(sum1.getOutputs().o(), sum2.getOutputs().o()));
 
     SdkNode<ListTask.Output> list1 =
-        builder.apply("list-1", new ListTask().withInput("list", list));
+        builder.apply("list-1", new ListTask(), ListTask.Input.create(list));
 
     builder.output("list", list1.getOutputs().list());
   }

flytekit-local-engine/src/test/java/org/flyte/localengine/examples/MapTask.java

@@ -40,8 +40,8 @@ public Output run(Input input) {
   public abstract static class Input {
     public abstract SdkBindingData<Map<String, Long>> map();
 
-    public static Input create(Map<String, Long> map) {
-      return new AutoValue_MapTask_Input(SdkBindingData.ofIntegerMap(map));
+    public static Input create(SdkBindingData<Map<String, Long>> map) {
+      return new AutoValue_MapTask_Input(map);
     }
   }