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