Search in sources :

Example 1 with Resources

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

the class ProjectClosureTest method testCreateTaskTemplateForRunnableTaskWithResources.

@Test
public void testCreateTaskTemplateForRunnableTaskWithResources() {
    // given
    Map<Resources.ResourceName, String> resourceValues = new HashMap<>();
    resourceValues.put(Resources.ResourceName.MEMORY, "0.5Gi");
    Resources expectedResources = Resources.builder().limits(resourceValues).requests(resourceValues).build();
    RunnableTask task = createRunnableTask(expectedResources);
    String image = "my-image";
    // when
    TaskTemplate result = ProjectClosure.createTaskTemplateForRunnableTask(task, image);
    // then
    Container container = result.container();
    assertNotNull(container);
    assertThat(container.image(), equalTo(image));
    assertThat(container.resources(), equalTo(expectedResources));
    assertThat(result.interface_(), equalTo(TypedInterface.builder().inputs(SdkTypes.nulls().getVariableMap()).outputs(SdkTypes.nulls().getVariableMap()).build()));
    assertThat(result.custom(), equalTo(Struct.of(emptyMap())));
    assertThat(result.retries(), equalTo(RetryStrategy.builder().retries(0).build()));
    assertThat(result.type(), equalTo("java-task"));
}
Also used : TaskTemplate(org.flyte.api.v1.TaskTemplate) Container(org.flyte.api.v1.Container) HashMap(java.util.HashMap) ByteString(com.google.protobuf.ByteString) Resources(org.flyte.api.v1.Resources) RunnableTask(org.flyte.api.v1.RunnableTask) Test(org.junit.jupiter.api.Test)

Example 2 with Resources

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

the class ProtoUtil method serialize.

@VisibleForTesting
static Tasks.Container serialize(Container container) {
    Tasks.Container.Builder builder = Tasks.Container.newBuilder().setImage(container.image()).addAllCommand(container.command()).addAllArgs(container.args());
    container.env().forEach(pair -> builder.addEnv(serialize(pair)));
    Resources resources = container.resources();
    if (resources != null) {
        builder.setResources(serialize(resources));
    }
    return builder.build();
}
Also used : Container(org.flyte.api.v1.Container) Resources(org.flyte.api.v1.Resources) VisibleForTesting(com.google.common.annotations.VisibleForTesting)

Example 3 with Resources

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

the class ProjectClosureTest method testCreateTaskTemplateForRunnableTask.

@Test
public void testCreateTaskTemplateForRunnableTask() {
    // given
    RunnableTask task = createRunnableTask(null);
    String image = "my-image";
    Resources expectedResources = Resources.builder().build();
    // when
    TaskTemplate result = ProjectClosure.createTaskTemplateForRunnableTask(task, image);
    // then
    Container container = result.container();
    assertNotNull(container);
    assertThat(container.image(), equalTo(image));
    assertThat(container.resources(), equalTo(expectedResources));
    assertThat(result.interface_(), equalTo(TypedInterface.builder().inputs(SdkTypes.nulls().getVariableMap()).outputs(SdkTypes.nulls().getVariableMap()).build()));
    assertThat(result.custom(), equalTo(Struct.of(emptyMap())));
    assertThat(result.retries(), equalTo(RetryStrategy.builder().retries(0).build()));
    assertThat(result.type(), equalTo("java-task"));
}
Also used : TaskTemplate(org.flyte.api.v1.TaskTemplate) Container(org.flyte.api.v1.Container) RunnableTask(org.flyte.api.v1.RunnableTask) ByteString(com.google.protobuf.ByteString) Resources(org.flyte.api.v1.Resources) Test(org.junit.jupiter.api.Test)

Example 4 with Resources

use of org.flyte.api.v1.Resources 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

Resources (org.flyte.api.v1.Resources)4 Container (org.flyte.api.v1.Container)3 RunnableTask (org.flyte.api.v1.RunnableTask)3 Test (org.junit.jupiter.api.Test)3 ByteString (com.google.protobuf.ByteString)2 HashMap (java.util.HashMap)2 TaskTemplate (org.flyte.api.v1.TaskTemplate)2 VisibleForTesting (com.google.common.annotations.VisibleForTesting)1 RetryStrategy (org.flyte.api.v1.RetryStrategy)1 TaskIdentifier (org.flyte.api.v1.TaskIdentifier)1 TypedInterface (org.flyte.api.v1.TypedInterface)1