Examples with TaskTemplate org.flyte.api.v1.TaskTemplate used on opensource projects

Example 1 with TaskTemplate

use of org.flyte.api.v1.TaskTemplate 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"));
}

Example 2 with TaskTemplate

use of org.flyte.api.v1.TaskTemplate 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));
}

Example 3 with TaskTemplate

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

the class ProjectClosure method applyCustom.

private static TaskSpec applyCustom(TaskSpec taskSpec, JFlyteCustom custom) {
    Struct rewrittenCustom = merge(custom.serializeToStruct(), taskSpec.taskTemplate().custom());
    TaskTemplate rewrittenTaskTemplate = taskSpec.taskTemplate().toBuilder().custom(rewrittenCustom).build();
    return TaskSpec.create(rewrittenTaskTemplate);
}

Example 4 with TaskTemplate

use of org.flyte.api.v1.TaskTemplate 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()));
}

Example 5 with TaskTemplate

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

the class Execute method execute.

private void execute() {
    Config config = Config.load();
    Collection<ClassLoader> modules = ClassLoaders.forModuleDir(config.moduleDir()).values();
    Map<String, FileSystem> fileSystems = FileSystemLoader.loadFileSystems(modules);
    FileSystem outputFs = FileSystemLoader.getFileSystem(fileSystems, outputPrefix);
    ProtoWriter protoWriter = new ProtoWriter(outputPrefix, outputFs);
    try {
        FileSystem inputFs = FileSystemLoader.getFileSystem(fileSystems, inputs);
        ProtoReader protoReader = new ProtoReader(inputFs);
        TaskTemplate taskTemplate = protoReader.getTaskTemplate(taskTemplatePath);
        ClassLoader packageClassLoader = PackageLoader.load(fileSystems, taskTemplate);
        // before we run anything, switch class loader, otherwise,
        // ServiceLoaders and other things wouldn't work, for instance,
        // FileSystemRegister in Apache Beam
        Map<String, Literal> outputs = withClassLoader(packageClassLoader, () -> {
            Map<String, Literal> input = protoReader.getInput(inputs);
            RunnableTask runnableTask = getTask(task);
            return runnableTask.run(input);
        });
        protoWriter.writeOutputs(outputs);
    } catch (ContainerError e) {
        LOG.error("failed to run task", e);
        protoWriter.writeError(ProtoUtil.serializeContainerError(e));
    } catch (Throwable e) {
        LOG.error("failed to run task", e);
        protoWriter.writeError(ProtoUtil.serializeThrowable(e));
    }
}