Search in sources :

Example 1 with ClassLoaders.withClassLoader

use of org.flyte.jflyte.ClassLoaders.withClassLoader 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));
    }
}
Also used : DynamicWorkflowTaskRegistrar(org.flyte.api.v1.DynamicWorkflowTaskRegistrar) WorkflowTemplate(org.flyte.api.v1.WorkflowTemplate) TaskIdentifier(org.flyte.api.v1.TaskIdentifier) PartialTaskIdentifier(org.flyte.api.v1.PartialTaskIdentifier) Struct(org.flyte.api.v1.Struct) WorkflowIdentifier(org.flyte.api.v1.WorkflowIdentifier) PartialWorkflowIdentifier(org.flyte.api.v1.PartialWorkflowIdentifier) RunnableTaskRegistrar(org.flyte.api.v1.RunnableTaskRegistrar) FileSystem(org.flyte.jflyte.api.FileSystem) Literal(org.flyte.api.v1.Literal) ClassLoaders.withClassLoader(org.flyte.jflyte.ClassLoaders.withClassLoader) TaskTemplate(org.flyte.api.v1.TaskTemplate) WorkflowTemplateRegistrar(org.flyte.api.v1.WorkflowTemplateRegistrar) DynamicWorkflowTask(org.flyte.api.v1.DynamicWorkflowTask) RunnableTask(org.flyte.api.v1.RunnableTask) DynamicJobSpec(org.flyte.api.v1.DynamicJobSpec) ContainerError(org.flyte.api.v1.ContainerError)

Aggregations

ContainerError (org.flyte.api.v1.ContainerError)1 DynamicJobSpec (org.flyte.api.v1.DynamicJobSpec)1 DynamicWorkflowTask (org.flyte.api.v1.DynamicWorkflowTask)1 DynamicWorkflowTaskRegistrar (org.flyte.api.v1.DynamicWorkflowTaskRegistrar)1 Literal (org.flyte.api.v1.Literal)1 PartialTaskIdentifier (org.flyte.api.v1.PartialTaskIdentifier)1 PartialWorkflowIdentifier (org.flyte.api.v1.PartialWorkflowIdentifier)1 RunnableTask (org.flyte.api.v1.RunnableTask)1 RunnableTaskRegistrar (org.flyte.api.v1.RunnableTaskRegistrar)1 Struct (org.flyte.api.v1.Struct)1 TaskIdentifier (org.flyte.api.v1.TaskIdentifier)1 TaskTemplate (org.flyte.api.v1.TaskTemplate)1 WorkflowIdentifier (org.flyte.api.v1.WorkflowIdentifier)1 WorkflowTemplate (org.flyte.api.v1.WorkflowTemplate)1 WorkflowTemplateRegistrar (org.flyte.api.v1.WorkflowTemplateRegistrar)1 ClassLoaders.withClassLoader (org.flyte.jflyte.ClassLoaders.withClassLoader)1 FileSystem (org.flyte.jflyte.api.FileSystem)1