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));
}
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()));
}
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);
}
Aggregations