Search in sources :

Example 1 with RetryStrategy

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

the class ProtoUtilTest method shouldSerDeTaskTemplate.

@Test
void shouldSerDeTaskTemplate() {
    Container container = Container.builder().command(ImmutableList.of("echo")).args(ImmutableList.of("hello world")).env(ImmutableList.of(KeyValuePair.of("key", "value"))).image("alpine:3.7").build();
    Variable stringVar = ApiUtils.createVar(SimpleType.STRING);
    Variable integerVar = ApiUtils.createVar(SimpleType.INTEGER);
    TypedInterface interface_ = TypedInterface.builder().inputs(ImmutableMap.of("x", stringVar)).outputs(ImmutableMap.of("y", integerVar)).build();
    RetryStrategy retries = RetryStrategy.builder().retries(4).build();
    TaskTemplate template = TaskTemplate.builder().container(container).type("custom-task").interface_(interface_).retries(retries).custom(Struct.of(ImmutableMap.of("custom-prop", Struct.Value.ofStringValue("custom-value")))).build();
    Tasks.TaskTemplate templateProto = Tasks.TaskTemplate.newBuilder().setContainer(Tasks.Container.newBuilder().setImage("alpine:3.7").addCommand("echo").addArgs("hello world").addEnv(Literals.KeyValuePair.newBuilder().setKey("key").setValue("value").build()).build()).setMetadata(Tasks.TaskMetadata.newBuilder().setRuntime(Tasks.RuntimeMetadata.newBuilder().setType(Tasks.RuntimeMetadata.RuntimeType.FLYTE_SDK).setFlavor(ProtoUtil.RUNTIME_FLAVOR).setVersion(ProtoUtil.RUNTIME_VERSION).build()).setRetries(Literals.RetryStrategy.newBuilder().setRetries(4).build()).build()).setInterface(Interface.TypedInterface.newBuilder().setInputs(Interface.VariableMap.newBuilder().putVariables("x", Interface.Variable.newBuilder().setType(Types.LiteralType.newBuilder().setSimple(Types.SimpleType.STRING).build()).build()).build()).setOutputs(Interface.VariableMap.newBuilder().putVariables("y", Interface.Variable.newBuilder().setType(Types.LiteralType.newBuilder().setSimple(Types.SimpleType.INTEGER).build()).build()).build()).build()).setType("custom-task").setCustom(com.google.protobuf.Struct.newBuilder().putFields("custom-prop", Value.newBuilder().setStringValue("custom-value").build()).build()).build();
    Tasks.TaskTemplate serializedTemplate = ProtoUtil.serialize(template);
    TaskTemplate deserializedTemplate = ProtoUtil.deserialize(templateProto);
    assertThat(serializedTemplate, equalTo(templateProto));
    assertThat(deserializedTemplate, equalTo(template));
}
Also used : TaskTemplate(org.flyte.api.v1.TaskTemplate) TypedInterface(org.flyte.api.v1.TypedInterface) Container(org.flyte.api.v1.Container) Variable(org.flyte.api.v1.Variable) Tasks(flyteidl.core.Tasks) RetryStrategy(org.flyte.api.v1.RetryStrategy) Test(org.junit.jupiter.api.Test) ParameterizedTest(org.junit.jupiter.params.ParameterizedTest)

Example 2 with RetryStrategy

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

the class FlyteAdminClientTest method shouldPropagateCreateTaskToStub.

@Test
public void shouldPropagateCreateTaskToStub() {
    TaskIdentifier identifier = TaskIdentifier.builder().domain(DOMAIN).project(PROJECT).name(TASK_NAME).version(TASK_VERSION).build();
    TypedInterface interface_ = TypedInterface.builder().inputs(ImmutableMap.of("x", createVar(SimpleType.STRING))).outputs(ImmutableMap.of("y", createVar(SimpleType.INTEGER))).build();
    Container container = Container.builder().command(ImmutableList.of(COMMAND)).args(ImmutableList.of()).image(IMAGE_NAME).env(ImmutableList.of(KeyValuePair.of("key", "value"))).build();
    RetryStrategy retries = RetryStrategy.builder().retries(4).build();
    TaskTemplate template = TaskTemplate.builder().container(container).type("custom-task").interface_(interface_).custom(Struct.of(emptyMap())).retries(retries).build();
    client.createTask(identifier, template);
    assertThat(stubService.createTaskRequest, equalTo(TaskOuterClass.TaskCreateRequest.newBuilder().setId(newIdentifier(ResourceType.TASK, TASK_NAME, TASK_VERSION)).setSpec(newTaskSpec()).build()));
}
Also used : TaskTemplate(org.flyte.api.v1.TaskTemplate) TypedInterface(org.flyte.api.v1.TypedInterface) Container(org.flyte.api.v1.Container) TaskIdentifier(org.flyte.api.v1.TaskIdentifier) PartialTaskIdentifier(org.flyte.api.v1.PartialTaskIdentifier) RetryStrategy(org.flyte.api.v1.RetryStrategy) Test(org.junit.Test)

Example 3 with RetryStrategy

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

the class SdkRunnableTaskRegistrarTest method shouldLoadRunnableTasksFromDiscoveredRegistries.

@Test
void shouldLoadRunnableTasksFromDiscoveredRegistries() {
    // given
    String testTaskName = "org.flyte.flytekit.SdkRunnableTaskRegistrarTest$TestTask";
    String otherTestTaskName = "org.flyte.flytekit.SdkRunnableTaskRegistrarTest$OtherTestTask";
    TaskIdentifier expectedTestTaskId = TaskIdentifier.builder().project("project").domain("domain").name(testTaskName).version("version").build();
    TypedInterface typedInterface = TypedInterface.builder().inputs(SdkTypes.nulls().getVariableMap()).outputs(SdkTypes.nulls().getVariableMap()).build();
    RetryStrategy retries = RetryStrategy.builder().retries(0).build();
    RetryStrategy otherRetries = RetryStrategy.builder().retries(1).build();
    Map<Resources.ResourceName, String> limits = new HashMap<>();
    limits.put(Resources.ResourceName.CPU, "0.5");
    limits.put(Resources.ResourceName.MEMORY, "2Gi");
    Map<Resources.ResourceName, String> requests = new HashMap<>();
    requests.put(Resources.ResourceName.CPU, "2");
    requests.put(Resources.ResourceName.MEMORY, "5Gi");
    Resources resources = Resources.builder().limits(limits).requests(requests).build();
    RunnableTask expectedTask = createRunnableTask(testTaskName, typedInterface, retries, null);
    TaskIdentifier expectedOtherTestTaskId = TaskIdentifier.builder().project("project").domain("domain").name(otherTestTaskName).version("version").build();
    RunnableTask expectedOtherTask = createRunnableTask(otherTestTaskName, typedInterface, otherRetries, resources);
    // when
    Map<TaskIdentifier, RunnableTask> tasks = registrar.load(ENV);
    // then
    assertAll(() -> assertThat(tasks, hasKey(is(expectedTestTaskId))), () -> assertThat(tasks, hasKey(is(expectedOtherTestTaskId))));
    assertTaskEquals(tasks.get(expectedTestTaskId), expectedTask);
    assertTaskEquals(tasks.get(expectedOtherTestTaskId), expectedOtherTask);
}
Also used : TypedInterface(org.flyte.api.v1.TypedInterface) TaskIdentifier(org.flyte.api.v1.TaskIdentifier) HashMap(java.util.HashMap) Resources(org.flyte.api.v1.Resources) RunnableTask(org.flyte.api.v1.RunnableTask) RetryStrategy(org.flyte.api.v1.RetryStrategy) Test(org.junit.jupiter.api.Test)

Aggregations

RetryStrategy (org.flyte.api.v1.RetryStrategy)3 TypedInterface (org.flyte.api.v1.TypedInterface)3 Container (org.flyte.api.v1.Container)2 TaskIdentifier (org.flyte.api.v1.TaskIdentifier)2 TaskTemplate (org.flyte.api.v1.TaskTemplate)2 Test (org.junit.jupiter.api.Test)2 Tasks (flyteidl.core.Tasks)1 HashMap (java.util.HashMap)1 PartialTaskIdentifier (org.flyte.api.v1.PartialTaskIdentifier)1 Resources (org.flyte.api.v1.Resources)1 RunnableTask (org.flyte.api.v1.RunnableTask)1 Variable (org.flyte.api.v1.Variable)1 Test (org.junit.Test)1 ParameterizedTest (org.junit.jupiter.params.ParameterizedTest)1