Skip to content

Commit

Permalink
Local engine classes
Browse files Browse the repository at this point in the history 
Signed-off-by: Nelson Arapé <[email protected]>
  • Loading branch information
@narape
narape committed Jan 24, 2023
1 parent 9561bcb commit 39b162e
Show file tree
Hide file tree
Showing 14 changed files with 119 additions and 100 deletions.
12 changes: 4 additions & 8 deletions flytekit-local-engine/src/test/java/org/flyte/localengine/LocalEngineTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -549,10 +548,10 @@ 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
Expand All @@ -562,8 +561,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();

Expand All @@ -583,9 +582,6 @@ public NoOpType run(NoOpType input) {
return NoOpType.create(input.x());
}

static SdkTransform<NoOpType> of(SdkBindingData<Long> x) {
return new NoOp().withInput("x", x);
}
}

@AutoValue
Expand Down
50 changes: 27 additions & 23 deletions ...al-engine/src/test/java/org/flyte/localengine/examples/CollatzConjectureStepWorkflow.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,17 +30,26 @@

// 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
static abstract 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
Expand All @@ -49,8 +58,8 @@ public void expand(SdkWorkflowBuilder builder) {
when(
"was_even",
isTrue(isOdd),
new Divide().withInput("num", x).withInput("den", ofInteger(2L)))
.otherwise("was_odd", new ThreeXPlusOne().withInput("x", x)))
new Divide(), Divide.Input.create(x, ofInteger(2L)))
.otherwise("was_odd", new ThreeXPlusOne(), ThreeXPlusOne.Input.create(x)))
.getOutputs()
.o();

Expand All @@ -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);
}
}

Expand All @@ -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);
}
}
}
Expand All @@ -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
Expand All @@ -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);
}
}

Expand All @@ -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);
}
}
}
Expand All @@ -138,7 +143,7 @@ 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
Expand All @@ -150,9 +155,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);
}
}
}
Expand Down
20 changes: 14 additions & 6 deletions flytekit-local-engine/src/test/java/org/flyte/localengine/examples/FibonacciWorkflow.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,10 +24,10 @@
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
Expand All @@ -37,32 +37,40 @@ public void expand(SdkWorkflowBuilder builder) {

SdkBindingData<Long> fib2 =
builder
.apply("fib-2", new SumTask().withInput("a", fib0).withInput("b", fib1))
.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))
.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))
.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))
.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();
Expand Down
6 changes: 3 additions & 3 deletions flytekit-local-engine/src/test/java/org/flyte/localengine/examples/InnerSubWorkflow.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,9 +23,9 @@
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
Expand All @@ -34,7 +34,7 @@ public void expand(SdkWorkflowBuilder builder) {
SdkBindingData<Long> b = builder.inputOfInteger("b");
SdkBindingData<Long> c =
builder
.apply("inner-sum-a-b", new SumTask().withInput("a", a).withInput("b", b))
.apply("inner-sum-a-b", new SumTask(), SumTask.Input.create(a, b))
.getOutputs()
.o();

Expand Down
8 changes: 4 additions & 4 deletions flytekit-local-engine/src/test/java/org/flyte/localengine/examples/ListTask.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,17 +40,17 @@ public Output run(Input input) {
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);
}
}
}
16 changes: 9 additions & 7 deletions flytekit-local-engine/src/test/java/org/flyte/localengine/examples/ListWorkflow.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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());
}
Expand Down
4 changes: 2 additions & 2 deletions flytekit-local-engine/src/test/java/org/flyte/localengine/examples/MapTask.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
}

Expand Down
13 changes: 8 additions & 5 deletions flytekit-local-engine/src/test/java/org/flyte/localengine/examples/MapWorkflow.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,21 +16,24 @@
*/
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;
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;

@AutoService(SdkWorkflow.class)
public class MapWorkflow extends SdkWorkflow<MapWorkflow.Output> {
public class MapWorkflow extends SdkWorkflow<Void, MapWorkflow.Output> {
public MapWorkflow() {
super(JacksonSdkType.of(MapWorkflow.Output.class));
super(SdkTypes.nulls(), JacksonSdkType.of(MapWorkflow.Output.class));
}

@AutoValue
Expand All @@ -46,17 +49,17 @@ public static MapWorkflow.Output create(Map<String, String> map) {
@Override
public void expand(SdkWorkflowBuilder builder) {
SdkBindingData<Long> 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<Long> 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<String, Long>> map =
SdkBindingData.ofBindingMap(
LiteralType.ofMapValueType(LiteralType.ofSimpleType(SimpleType.INTEGER)),
Map.of("e", sum1, "f", sum2));

SdkNode<MapTask.Output> map1 = builder.apply("map-1", new MapTask().withInput("map", map));
SdkNode<MapTask.Output> map1 = builder.apply("map-1", new MapTask(), MapTask.Input.create(map));

builder.output("map", map1.getOutputs().map());
}
Expand Down
6 changes: 3 additions & 3 deletions flytekit-local-engine/src/test/java/org/flyte/localengine/examples/NestedSubWorkflow.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,10 +23,10 @@
import org.flyte.flytekit.jackson.JacksonSdkType;

@AutoService(SdkWorkflow.class)
public class NestedSubWorkflow extends SdkWorkflow<TestUnaryIntegerOutput> {
public class NestedSubWorkflow extends SdkWorkflow<TestTuple3IntegerInput, TestUnaryIntegerOutput> {

public NestedSubWorkflow() {
super(JacksonSdkType.of(TestUnaryIntegerOutput.class));
super(JacksonSdkType.of(TestTuple3IntegerInput.class), JacksonSdkType.of(TestUnaryIntegerOutput.class));
}

@Override
Expand All @@ -38,7 +38,7 @@ public void expand(SdkWorkflowBuilder builder) {
builder
.apply(
"nested-workflow",
new OuterSubWorkflow().withInput("a", a).withInput("b", b).withInput("c", c))
new OuterSubWorkflow(), TestTuple3IntegerInput.create(a, b, c))
.getOutputs()
.o();
builder.output("o", result);
Expand Down
Loading

0 comments on commit 39b162e

Please sign in to comment.