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

Example 1 with DynamicJobSpec

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

the class ProtoUtilTest method shouldSerializeDynamicJobSpec.

@Test
void shouldSerializeDynamicJobSpec() {
    DynamicJobSpec dynamicJobSpec = DynamicJobSpec.builder().nodes(emptyList()).subWorkflows(emptyMap()).tasks(emptyMap()).outputs(emptyList()).build();
    DynamicJob.DynamicJobSpec proto = ProtoUtil.serialize(dynamicJobSpec);
    assertThat(proto, equalTo(DynamicJob.DynamicJobSpec.newBuilder().build()));
}

Example 2 with DynamicJobSpec

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

the class ProtoWriterTest method shouldWriteFutures.

@Test
void shouldWriteFutures() throws IOException {
    ProtoWriter protoWriter = new ProtoWriter("/home/test/prefix", new InMemoryFileSystem(extension.getFileSystem()));
    DynamicJobSpec spec = DynamicJobSpec.builder().nodes(singletonList(Node.builder().id("echo").inputs(singletonList(Binding.builder().var_("in").binding(BindingData.ofScalar(Scalar.ofPrimitive(Primitive.ofStringValue("ping")))).build())).upstreamNodeIds(emptyList()).build())).outputs(singletonList(Binding.builder().var_("out").binding(BindingData.ofOutputReference(OutputReference.builder().nodeId("echo").var("out").build())).build())).subWorkflows(emptyMap()).tasks(emptyMap()).build();
    protoWriter.writeFutures(spec);
    Path futurePath = extension.getFileSystem().getPath("/home/test/prefix/futures.pb");
    DynamicJob.DynamicJobSpec actualError = DynamicJob.DynamicJobSpec.parseFrom(Files.newInputStream(futurePath));
    assertThat(actualError, equalTo(DynamicJob.DynamicJobSpec.newBuilder().addNodes(Workflow.Node.newBuilder().setId("echo").addInputs(Literals.Binding.newBuilder().setVar("in").setBinding(Literals.BindingData.newBuilder().setScalar(Literals.Scalar.newBuilder().setPrimitive(Literals.Primitive.newBuilder().setStringValue("ping").build()).build()).build()).build()).build()).addOutputs(Literals.Binding.newBuilder().setVar("out").setBinding(Literals.BindingData.newBuilder().setPromise(Types.OutputReference.newBuilder().setNodeId("echo").setVar("out").build()).build()).build()).build()));
}

Example 3 with DynamicJobSpec

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

the class ExecuteDynamicWorkflow method execute.

private void execute() {
    Config config = Config.load();
    ExecutionConfig executionConfig = ExecutionConfig.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);
        Map<String, String> env = getEnv();
        Map<WorkflowIdentifier, WorkflowTemplate> workflowTemplates = ClassLoaders.withClassLoader(packageClassLoader, () -> Registrars.loadAll(WorkflowTemplateRegistrar.class, env));
        Map<TaskIdentifier, RunnableTask> runnableTasks = ClassLoaders.withClassLoader(packageClassLoader, () -> Registrars.loadAll(RunnableTaskRegistrar.class, env));
        Map<TaskIdentifier, DynamicWorkflowTask> dynamicWorkflowTasks = ClassLoaders.withClassLoader(packageClassLoader, () -> Registrars.loadAll(DynamicWorkflowTaskRegistrar.class, env));
        // before we run anything, switch class loader, otherwise,
        // ServiceLoaders and other things wouldn't work, for instance,
        // FileSystemRegister in Apache Beam
        // we don't take the whole "custom" field, but only jflyte part, for that we ser-de it
        Struct custom = JFlyteCustom.deserializeFromStruct(taskTemplate.custom()).serializeToStruct();
        // all tasks already have staged jars, we can reuse 'jflyte' custom from current task to get
        // it
        Map<TaskIdentifier, TaskTemplate> taskTemplates = mapValues(ProjectClosure.createTaskTemplates(executionConfig, runnableTasks, dynamicWorkflowTasks), template -> template.toBuilder().custom(ProjectClosure.merge(template.custom(), custom)).build());
        DynamicJobSpec futures = withClassLoader(packageClassLoader, () -> {
            Map<String, Literal> input = protoReader.getInput(inputs);
            DynamicWorkflowTask task = getDynamicWorkflowTask(this.task);
            return task.run(input);
        });
        DynamicJobSpec rewrittenFutures = rewrite(executionConfig, futures, taskTemplates, workflowTemplates);
        if (rewrittenFutures.nodes().isEmpty()) {
            Map<String, Literal> outputs = getLiteralMap(rewrittenFutures.outputs());
            protoWriter.writeOutputs(outputs);
        } else {
            protoWriter.writeFutures(rewrittenFutures);
        }
    } catch (ContainerError e) {
        LOG.error("failed to run dynamic workflow", e);
        protoWriter.writeError(ProtoUtil.serializeContainerError(e));
    } catch (Throwable e) {
        LOG.error("failed to run dynamic workflow", e);
        protoWriter.writeError(ProtoUtil.serializeThrowable(e));
    }
}