From be2e860ad457749399ca26155cb84f04e3fbb200 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nelson=20Arap=C3=A9?= Date: Tue, 24 Jan 2023 08:42:31 +0100 Subject: [PATCH] Local engine classes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Nelson Arapé --- .../flyte/localengine/LocalEngineTest.java | 14 ++--- .../CollatzConjectureStepWorkflow.java | 58 ++++++++++--------- .../examples/FibonacciWorkflow.java | 34 +++++------ .../examples/InnerSubWorkflow.java | 9 +-- .../flyte/localengine/examples/ListTask.java | 10 ++-- .../localengine/examples/ListWorkflow.java | 16 ++--- .../flyte/localengine/examples/MapTask.java | 4 +- .../localengine/examples/MapWorkflow.java | 19 ++++-- .../examples/NestedSubWorkflow.java | 9 +-- .../examples/OuterSubWorkflow.java | 14 ++--- .../examples/RetryableWorkflow.java | 4 +- .../flyte/localengine/examples/SumTask.java | 14 ++--- .../examples/TestTuple3IntegerInput.java | 34 +++++++++++ .../examples/TestUnaryIntegerOutput.java | 39 ------------- 14 files changed, 141 insertions(+), 137 deletions(-) create mode 100644 flytekit-local-engine/src/test/java/org/flyte/localengine/examples/TestTuple3IntegerInput.java diff --git a/flytekit-local-engine/src/test/java/org/flyte/localengine/LocalEngineTest.java b/flytekit-local-engine/src/test/java/org/flyte/localengine/LocalEngineTest.java index 2ff089db5..4ce27d35a 100644 --- a/flytekit-local-engine/src/test/java/org/flyte/localengine/LocalEngineTest.java +++ b/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 > Map loadAll( @AutoService(SdkWorkflow.class) public static class TestCaseExhaustivenessWorkflow - extends SdkWorkflow { + 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 of(SdkBindingData x) { - return new NoOp().withInput("x", x); - } } @AutoValue diff --git a/flytekit-local-engine/src/test/java/org/flyte/localengine/examples/CollatzConjectureStepWorkflow.java b/flytekit-local-engine/src/test/java/org/flyte/localengine/examples/CollatzConjectureStepWorkflow.java index 80859c22c..cd7cff942 100644 --- a/flytekit-local-engine/src/test/java/org/flyte/localengine/examples/CollatzConjectureStepWorkflow.java +++ b/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 { +public class CollatzConjectureStepWorkflow + extends SdkWorkflow { public CollatzConjectureStepWorkflow() { - super(new TestUnaryIntegerOutput.SdkType()); + super( + JacksonSdkType.of(CollatzConjectureStepWorkflow.Input.class), + JacksonSdkType.of(TestUnaryIntegerOutput.class)); + } + + @AutoValue + abstract static class Input { + abstract SdkBindingData x(); + + public static Input create(SdkBindingData x) { + return new AutoValue_CollatzConjectureStepWorkflow_Input(x); + } } @Override public void expand(SdkWorkflowBuilder builder) { SdkBindingData x = builder.inputOfInteger("x"); SdkBindingData 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 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,7 +76,7 @@ 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 @@ -75,9 +84,8 @@ public abstract static class Input { public abstract SdkBindingData x(); - public static Input create(Long x) { - return new AutoValue_CollatzConjectureStepWorkflow_IsEvenTask_Input( - SdkBindingData.ofInteger(x)); + public static Input create(SdkBindingData x) { + return new AutoValue_CollatzConjectureStepWorkflow_IsEvenTask_Input(x); } } @@ -86,9 +94,8 @@ public abstract static class Output { public abstract SdkBindingData res(); - public static Output create(boolean res) { - return new AutoValue_CollatzConjectureStepWorkflow_IsEvenTask_Output( - SdkBindingData.ofBoolean(res)); + public static Output create(SdkBindingData res) { + return new AutoValue_CollatzConjectureStepWorkflow_IsEvenTask_Output(res); } } } @@ -98,7 +105,7 @@ public static class Divide extends SdkRunnableTask 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 num, SdkBindingData den) { + return new AutoValue_CollatzConjectureStepWorkflow_Divide_Input(num, den); } } @@ -124,9 +130,8 @@ public abstract static class Output { public abstract SdkBindingData res(); - public static Output create(long res) { - return new AutoValue_CollatzConjectureStepWorkflow_Divide_Output( - SdkBindingData.ofInteger(res)); + public static Output create(SdkBindingData 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 x(); - public static Input create(long x) { - return new AutoValue_CollatzConjectureStepWorkflow_ThreeXPlusOne_Input( - SdkBindingData.ofInteger(x)); + public static Input create(SdkBindingData x) { + return new AutoValue_CollatzConjectureStepWorkflow_ThreeXPlusOne_Input(x); } } } diff --git a/flytekit-local-engine/src/test/java/org/flyte/localengine/examples/FibonacciWorkflow.java b/flytekit-local-engine/src/test/java/org/flyte/localengine/examples/FibonacciWorkflow.java index da9ee444f..556c3d174 100644 --- a/flytekit-local-engine/src/test/java/org/flyte/localengine/examples/FibonacciWorkflow.java +++ b/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 { +public class FibonacciWorkflow + extends SdkWorkflow { 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 fib1 = builder.inputOfInteger("fib1"); SdkBindingData 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 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 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 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 fib0(); + + public abstract SdkBindingData fib1(); + } + @AutoValue public abstract static class Output { public abstract SdkBindingData fib4(); diff --git a/flytekit-local-engine/src/test/java/org/flyte/localengine/examples/InnerSubWorkflow.java b/flytekit-local-engine/src/test/java/org/flyte/localengine/examples/InnerSubWorkflow.java index e68ef2fe6..3fc00c19d 100644 --- a/flytekit-local-engine/src/test/java/org/flyte/localengine/examples/InnerSubWorkflow.java +++ b/flytekit-local-engine/src/test/java/org/flyte/localengine/examples/InnerSubWorkflow.java @@ -23,9 +23,9 @@ import org.flyte.flytekit.jackson.JacksonSdkType; @AutoService(SdkWorkflow.class) -public class InnerSubWorkflow extends SdkWorkflow { +public class InnerSubWorkflow extends SdkWorkflow { public InnerSubWorkflow() { - super(JacksonSdkType.of(TestUnaryIntegerOutput.class)); + super(JacksonSdkType.of(SumTask.Input.class), JacksonSdkType.of(TestUnaryIntegerOutput.class)); } @Override @@ -33,10 +33,7 @@ public void expand(SdkWorkflowBuilder builder) { SdkBindingData a = builder.inputOfInteger("a"); SdkBindingData b = builder.inputOfInteger("b"); SdkBindingData 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); } diff --git a/flytekit-local-engine/src/test/java/org/flyte/localengine/examples/ListTask.java b/flytekit-local-engine/src/test/java/org/flyte/localengine/examples/ListTask.java index ab43adfb0..1f1dd0c87 100644 --- a/flytekit-local-engine/src/test/java/org/flyte/localengine/examples/ListTask.java +++ b/flytekit-local-engine/src/test/java/org/flyte/localengine/examples/ListTask.java @@ -33,15 +33,15 @@ 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(); - public static Input create(List list) { - return new AutoValue_ListTask_Input(SdkBindingData.ofIntegerCollection(list)); + public static Input create(SdkBindingData> list) { + return new AutoValue_ListTask_Input(list); } } @@ -49,8 +49,8 @@ public static Input create(List list) { public abstract static class Output { public abstract SdkBindingData> list(); - public static Output create(List list) { - return new AutoValue_ListTask_Output(SdkBindingData.ofIntegerCollection(list)); + public static Output create(SdkBindingData> list) { + return new AutoValue_ListTask_Output(list); } } } diff --git a/flytekit-local-engine/src/test/java/org/flyte/localengine/examples/ListWorkflow.java b/flytekit-local-engine/src/test/java/org/flyte/localengine/examples/ListWorkflow.java index 7f29050a9..44fb97f15 100644 --- a/flytekit-local-engine/src/test/java/org/flyte/localengine/examples/ListWorkflow.java +++ b/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 { +public class ListWorkflow extends SdkWorkflow { public ListWorkflow() { - super(JacksonSdkType.of(ListTask.Output.class)); + super(SdkTypes.nulls(), JacksonSdkType.of(ListTask.Output.class)); } @Override public void expand(SdkWorkflowBuilder builder) { SdkNode 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 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 = 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 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()); } diff --git a/flytekit-local-engine/src/test/java/org/flyte/localengine/examples/MapTask.java b/flytekit-local-engine/src/test/java/org/flyte/localengine/examples/MapTask.java index 15f8d4151..47d76801f 100644 --- a/flytekit-local-engine/src/test/java/org/flyte/localengine/examples/MapTask.java +++ b/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(); - public static Input create(Map map) { - return new AutoValue_MapTask_Input(SdkBindingData.ofIntegerMap(map)); + public static Input create(SdkBindingData> map) { + return new AutoValue_MapTask_Input(map); } } diff --git a/flytekit-local-engine/src/test/java/org/flyte/localengine/examples/MapWorkflow.java b/flytekit-local-engine/src/test/java/org/flyte/localengine/examples/MapWorkflow.java index 3aa08d272..7642b80f0 100644 --- a/flytekit-local-engine/src/test/java/org/flyte/localengine/examples/MapWorkflow.java +++ b/flytekit-local-engine/src/test/java/org/flyte/localengine/examples/MapWorkflow.java @@ -16,6 +16,8 @@ */ package org.flyte.localengine.examples; +import static org.flyte.flytekit.SdkBindingData.ofInteger; + import com.google.auto.service.AutoService; import com.google.auto.value.AutoValue; import java.util.Map; @@ -23,14 +25,15 @@ 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; @AutoService(SdkWorkflow.class) -public class MapWorkflow extends SdkWorkflow { +public class MapWorkflow extends SdkWorkflow { public MapWorkflow() { - super(JacksonSdkType.of(MapWorkflow.Output.class)); + super(SdkTypes.nulls(), JacksonSdkType.of(MapWorkflow.Output.class)); } @AutoValue @@ -46,17 +49,23 @@ public static MapWorkflow.Output create(Map map) { @Override public void expand(SdkWorkflowBuilder builder) { SdkBindingData sum1 = - builder.apply("sum-1", new SumTask().withInput("a", 1).withInput("b", 2)).getOutputs().o(); + builder + .apply("sum-1", new SumTask(), SumTask.Input.create(ofInteger(1), ofInteger(2))) + .getOutputs() + .o(); SdkBindingData sum2 = - builder.apply("sum-2", new SumTask().withInput("a", 3).withInput("b", 4)).getOutputs().o(); + builder + .apply("sum-2", new SumTask(), SumTask.Input.create(ofInteger(3), ofInteger(4))) + .getOutputs() + .o(); SdkBindingData> map = SdkBindingData.ofBindingMap( LiteralType.ofMapValueType(LiteralType.ofSimpleType(SimpleType.INTEGER)), Map.of("e", sum1, "f", sum2)); - SdkNode map1 = builder.apply("map-1", new MapTask().withInput("map", map)); + SdkNode map1 = builder.apply("map-1", new MapTask(), MapTask.Input.create(map)); builder.output("map", map1.getOutputs().map()); } diff --git a/flytekit-local-engine/src/test/java/org/flyte/localengine/examples/NestedSubWorkflow.java b/flytekit-local-engine/src/test/java/org/flyte/localengine/examples/NestedSubWorkflow.java index bfc760ee3..d9eae6d0f 100644 --- a/flytekit-local-engine/src/test/java/org/flyte/localengine/examples/NestedSubWorkflow.java +++ b/flytekit-local-engine/src/test/java/org/flyte/localengine/examples/NestedSubWorkflow.java @@ -23,10 +23,12 @@ import org.flyte.flytekit.jackson.JacksonSdkType; @AutoService(SdkWorkflow.class) -public class NestedSubWorkflow extends SdkWorkflow { +public class NestedSubWorkflow extends SdkWorkflow { public NestedSubWorkflow() { - super(JacksonSdkType.of(TestUnaryIntegerOutput.class)); + super( + JacksonSdkType.of(TestTuple3IntegerInput.class), + JacksonSdkType.of(TestUnaryIntegerOutput.class)); } @Override @@ -37,8 +39,7 @@ public void expand(SdkWorkflowBuilder builder) { SdkBindingData result = builder .apply( - "nested-workflow", - new OuterSubWorkflow().withInput("a", a).withInput("b", b).withInput("c", c)) + "nested-workflow", new OuterSubWorkflow(), TestTuple3IntegerInput.create(a, b, c)) .getOutputs() .o(); builder.output("o", result); diff --git a/flytekit-local-engine/src/test/java/org/flyte/localengine/examples/OuterSubWorkflow.java b/flytekit-local-engine/src/test/java/org/flyte/localengine/examples/OuterSubWorkflow.java index 55722b00f..c5eb4d544 100644 --- a/flytekit-local-engine/src/test/java/org/flyte/localengine/examples/OuterSubWorkflow.java +++ b/flytekit-local-engine/src/test/java/org/flyte/localengine/examples/OuterSubWorkflow.java @@ -21,12 +21,15 @@ import org.flyte.flytekit.SdkWorkflow; import org.flyte.flytekit.SdkWorkflowBuilder; import org.flyte.flytekit.jackson.JacksonSdkType; +import org.flyte.localengine.examples.SumTask.Input; @AutoService(SdkWorkflow.class) -public class OuterSubWorkflow extends SdkWorkflow { +public class OuterSubWorkflow extends SdkWorkflow { public OuterSubWorkflow() { - super(JacksonSdkType.of(TestUnaryIntegerOutput.class)); + super( + JacksonSdkType.of(TestTuple3IntegerInput.class), + JacksonSdkType.of(TestUnaryIntegerOutput.class)); } @Override @@ -35,13 +38,10 @@ public void expand(SdkWorkflowBuilder builder) { SdkBindingData b = builder.inputOfInteger("b"); SdkBindingData c = builder.inputOfInteger("c"); SdkBindingData ab = - builder - .apply("outer-sum-a-b", new SumTask().withInput("a", a).withInput("b", b)) - .getOutputs() - .o(); + builder.apply("outer-sum-a-b", new SumTask(), Input.create(a, b)).getOutputs().o(); SdkBindingData res = builder - .apply("outer-sum-ab-c", new InnerSubWorkflow().withInput("a", ab).withInput("b", c)) + .apply("outer-sum-ab-c", new InnerSubWorkflow(), SumTask.Input.create(ab, c)) .getOutputs() .o(); builder.output("o", res); diff --git a/flytekit-local-engine/src/test/java/org/flyte/localengine/examples/RetryableWorkflow.java b/flytekit-local-engine/src/test/java/org/flyte/localengine/examples/RetryableWorkflow.java index 27317e5a4..9f1d4583a 100644 --- a/flytekit-local-engine/src/test/java/org/flyte/localengine/examples/RetryableWorkflow.java +++ b/flytekit-local-engine/src/test/java/org/flyte/localengine/examples/RetryableWorkflow.java @@ -22,9 +22,9 @@ import org.flyte.flytekit.SdkWorkflowBuilder; @AutoService(SdkWorkflow.class) -public class RetryableWorkflow extends SdkWorkflow { +public class RetryableWorkflow extends SdkWorkflow { public RetryableWorkflow() { - super(SdkTypes.nulls()); + super(SdkTypes.nulls(), SdkTypes.nulls()); } @Override diff --git a/flytekit-local-engine/src/test/java/org/flyte/localengine/examples/SumTask.java b/flytekit-local-engine/src/test/java/org/flyte/localengine/examples/SumTask.java index 54c614c1d..7f37f1409 100644 --- a/flytekit-local-engine/src/test/java/org/flyte/localengine/examples/SumTask.java +++ b/flytekit-local-engine/src/test/java/org/flyte/localengine/examples/SumTask.java @@ -21,29 +21,29 @@ import org.flyte.flytekit.SdkBindingData; import org.flyte.flytekit.SdkRunnableTask; import org.flyte.flytekit.jackson.JacksonSdkType; +import org.flyte.localengine.examples.SumTask.Input; @AutoService(SdkRunnableTask.class) -public class SumTask extends SdkRunnableTask { +public class SumTask extends SdkRunnableTask { private static final long serialVersionUID = -7796919693971619417L; public SumTask() { - super(JacksonSdkType.of(SumInput.class), new TestUnaryIntegerOutput.SdkType()); + super(JacksonSdkType.of(Input.class), JacksonSdkType.of(TestUnaryIntegerOutput.class)); } @AutoValue - public abstract static class SumInput { + public abstract static class Input { public abstract SdkBindingData a(); public abstract SdkBindingData b(); - public static SumInput create(long a, long b) { - return new AutoValue_SumTask_SumInput( - SdkBindingData.ofInteger(a), SdkBindingData.ofInteger(b)); + public static Input create(SdkBindingData a, SdkBindingData b) { + return new AutoValue_SumTask_Input(a, b); } } @Override - public TestUnaryIntegerOutput run(SumInput input) { + public TestUnaryIntegerOutput run(Input input) { return TestUnaryIntegerOutput.create( SdkBindingData.ofInteger(input.a().get() + input.b().get())); } diff --git a/flytekit-local-engine/src/test/java/org/flyte/localengine/examples/TestTuple3IntegerInput.java b/flytekit-local-engine/src/test/java/org/flyte/localengine/examples/TestTuple3IntegerInput.java new file mode 100644 index 000000000..f88a43601 --- /dev/null +++ b/flytekit-local-engine/src/test/java/org/flyte/localengine/examples/TestTuple3IntegerInput.java @@ -0,0 +1,34 @@ +/* + * Copyright 2021 Flyte Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.flyte.localengine.examples; + +import com.google.auto.value.AutoValue; +import org.flyte.flytekit.SdkBindingData; + +@AutoValue +public abstract class TestTuple3IntegerInput { + public abstract SdkBindingData a(); + + public abstract SdkBindingData b(); + + public abstract SdkBindingData c(); + + public static TestTuple3IntegerInput create( + SdkBindingData a, SdkBindingData b, SdkBindingData c) { + return new AutoValue_TestTuple3IntegerInput(a, b, c); + } +} diff --git a/flytekit-local-engine/src/test/java/org/flyte/localengine/examples/TestUnaryIntegerOutput.java b/flytekit-local-engine/src/test/java/org/flyte/localengine/examples/TestUnaryIntegerOutput.java index 337ff472e..cd7ed92f4 100644 --- a/flytekit-local-engine/src/test/java/org/flyte/localengine/examples/TestUnaryIntegerOutput.java +++ b/flytekit-local-engine/src/test/java/org/flyte/localengine/examples/TestUnaryIntegerOutput.java @@ -17,13 +17,6 @@ package org.flyte.localengine.examples; import com.google.auto.value.AutoValue; -import java.util.Map; -import org.flyte.api.v1.Literal; -import org.flyte.api.v1.LiteralType; -import org.flyte.api.v1.Primitive; -import org.flyte.api.v1.Scalar; -import org.flyte.api.v1.SimpleType; -import org.flyte.api.v1.Variable; import org.flyte.flytekit.SdkBindingData; @AutoValue @@ -33,36 +26,4 @@ public abstract class TestUnaryIntegerOutput { public static TestUnaryIntegerOutput create(SdkBindingData o) { return new AutoValue_TestUnaryIntegerOutput(o); } - - public static class SdkType extends org.flyte.flytekit.SdkType { - - private static final String VAR = "o"; - private static final LiteralType LITERAL_TYPE = LiteralType.ofSimpleType(SimpleType.INTEGER); - - @Override - public Map toLiteralMap(TestUnaryIntegerOutput value) { - return Map.of( - VAR, Literal.ofScalar(Scalar.ofPrimitive(Primitive.ofIntegerValue(value.o().get())))); - } - - @Override - public TestUnaryIntegerOutput fromLiteralMap(Map value) { - return create(SdkBindingData.ofInteger(value.get(VAR).scalar().primitive().integerValue())); - } - - @Override - public TestUnaryIntegerOutput promiseFor(String nodeId) { - return create(SdkBindingData.ofOutputReference(nodeId, VAR, LITERAL_TYPE)); - } - - @Override - public Map getVariableMap() { - return Map.of(VAR, Variable.builder().literalType(LITERAL_TYPE).build()); - } - - @Override - public Map> toSdkBindingMap(TestUnaryIntegerOutput value) { - return Map.of(VAR, value.o()); - } - } }