Search in sources :

Example 1 with DynamicWorkflowTask

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

Example 2 with DynamicWorkflowTask

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

the class ExecuteLocal method call.

@Override
public Integer call() {
    Map<String, ClassLoader> modules = ExecuteLocalLoader.loadModules(packageDir);
    Map<String, String> env = ImmutableMap.of("FLYTE_INTERNAL_DOMAIN", "development", "FLYTE_INTERNAL_VERSION", "test", "FLYTE_INTERNAL_PROJECT", "flytetester");
    Map<String, RunnableTask> runnableTasks = ExecuteLocalLoader.loadTasks(modules, env);
    Map<String, DynamicWorkflowTask> dynamicWorkflowTasks = // TODO support dynamic tasks
    emptyMap();
    Map<String, WorkflowTemplate> workflows = ExecuteLocalLoader.loadWorkflows(modules, env);
    WorkflowTemplate workflow = Preconditions.checkNotNull(workflows.get(workflowName), "workflow not found [%s]", workflowName);
    String synopsis = getCustomSynopsis();
    List<String> inputArgsList = inputArgs == null ? Collections.emptyList() : Arrays.asList(inputArgs);
    Map<String, Literal> inputs = getArgsParser().parseInputs(synopsis, workflow.interface_().inputs(), inputArgsList);
    try {
        // TODO, use logging listener here
        ExecutionListener listener = NoopExecutionListener.create();
        Map<String, Literal> outputs = LocalEngine.compileAndExecute(workflow, runnableTasks, dynamicWorkflowTasks, inputs, listener);
        LOG.info("Outputs: " + StringUtil.serializeLiteralMap(outputs));
        return 0;
    } catch (Throwable e) {
        return handleException(e);
    }
}
Also used : WorkflowTemplate(org.flyte.api.v1.WorkflowTemplate) NoopExecutionListener(org.flyte.localengine.NoopExecutionListener) ExecutionListener(org.flyte.localengine.ExecutionListener) DynamicWorkflowTask(org.flyte.api.v1.DynamicWorkflowTask) Literal(org.flyte.api.v1.Literal) RunnableTask(org.flyte.api.v1.RunnableTask)

Example 3 with DynamicWorkflowTask

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

the class ProjectClosure method createTaskTemplates.

static Map<TaskIdentifier, TaskTemplate> createTaskTemplates(ExecutionConfig config, Map<TaskIdentifier, RunnableTask> runnableTasks, Map<TaskIdentifier, DynamicWorkflowTask> dynamicWorkflowTasks) {
    Map<TaskIdentifier, TaskTemplate> taskTemplates = new HashMap<>();
    runnableTasks.forEach((id, task) -> {
        TaskTemplate taskTemplate = createTaskTemplateForRunnableTask(task, config.image());
        taskTemplates.put(id, taskTemplate);
    });
    dynamicWorkflowTasks.forEach((id, task) -> {
        TaskTemplate taskTemplate = createTaskTemplateForDynamicWorkflow(task, config.image());
        taskTemplates.put(id, taskTemplate);
    });
    return taskTemplates;
}
Also used : TaskTemplate(org.flyte.api.v1.TaskTemplate) TaskIdentifier(org.flyte.api.v1.TaskIdentifier) PartialTaskIdentifier(org.flyte.api.v1.PartialTaskIdentifier) HashMap(java.util.HashMap)

Example 4 with DynamicWorkflowTask

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

the class SdkDynamicWorkflowTaskRegistrar method load.

@Override
@SuppressWarnings("rawtypes")
public Map<TaskIdentifier, DynamicWorkflowTask> load(Map<String, String> env, ClassLoader classLoader) {
    ServiceLoader<SdkDynamicWorkflowTask> loader = ServiceLoader.load(SdkDynamicWorkflowTask.class, classLoader);
    LOG.fine("Discovering SdkDynamicWorkflowTask");
    Map<TaskIdentifier, DynamicWorkflowTask> tasks = new HashMap<>();
    SdkConfig sdkConfig = SdkConfig.load(env);
    for (SdkDynamicWorkflowTask<?, ?> sdkTask : loader) {
        String name = sdkTask.getName();
        TaskIdentifier taskId = TaskIdentifier.builder().domain(sdkConfig.domain()).project(sdkConfig.project()).name(name).version(sdkConfig.version()).build();
        LOG.fine(String.format("Discovered [%s]", name));
        DynamicWorkflowTask task = new DynamicWorkflowTaskImpl<>(sdkTask);
        DynamicWorkflowTask previous = tasks.put(taskId, task);
        if (previous != null) {
            throw new IllegalArgumentException(String.format("Discovered a duplicate dynamic workflow task [%s] [%s] [%s]", name, task, previous));
        }
    }
    return tasks;
}
Also used : TaskIdentifier(org.flyte.api.v1.TaskIdentifier) DynamicWorkflowTask(org.flyte.api.v1.DynamicWorkflowTask) HashMap(java.util.HashMap)

Example 5 with DynamicWorkflowTask

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

the class ExecutionNodeCompiler method compile.

static ExecutionNode compile(Node node, Map<String, RunnableTask> runnableTasks, Map<String, DynamicWorkflowTask> dynamicWorkflowTasks) {
    List<String> upstreamNodeIds = new ArrayList<>();
    node.inputs().stream().map(Binding::binding).flatMap(ExecutionNodeCompiler::unpackBindingData).filter(x -> x.kind() == BindingData.Kind.PROMISE).map(x -> x.promise().nodeId()).forEach(upstreamNodeIds::add);
    upstreamNodeIds.addAll(node.upstreamNodeIds());
    if (upstreamNodeIds.isEmpty()) {
        upstreamNodeIds.add(START_NODE_ID);
    }
    if (node.branchNode() != null) {
        throw new IllegalArgumentException("BranchNode isn't yet supported for local execution");
    }
    if (node.workflowNode() != null) {
        throw new IllegalArgumentException("WorkflowNode isn't yet supported for local execution");
    }
    String taskName = node.taskNode().referenceId().name();
    DynamicWorkflowTask dynamicWorkflowTask = dynamicWorkflowTasks.get(taskName);
    RunnableTask runnableTask = runnableTasks.get(taskName);
    if (dynamicWorkflowTask != null) {
        throw new IllegalArgumentException("DynamicWorkflowTask isn't yet supported for local execution");
    }
    Objects.requireNonNull(runnableTask, () -> String.format("Couldn't find task [%s]", taskName));
    int attempts = runnableTask.getRetries().retries() + 1;
    return ExecutionNode.builder().nodeId(node.id()).bindings(node.inputs()).runnableTask(runnableTask).upstreamNodeIds(upstreamNodeIds).attempts(attempts).build();
}
Also used : Binding(org.flyte.api.v1.Binding) Node(org.flyte.api.v1.Node) Collections.emptyList(java.util.Collections.emptyList) START_NODE_ID(org.flyte.api.v1.Node.START_NODE_ID) DynamicWorkflowTask(org.flyte.api.v1.DynamicWorkflowTask) Set(java.util.Set) HashMap(java.util.HashMap) Binding(org.flyte.api.v1.Binding) Deque(java.util.Deque) RunnableTask(org.flyte.api.v1.RunnableTask) Collections.singletonList(java.util.Collections.singletonList) ArrayList(java.util.ArrayList) HashSet(java.util.HashSet) Objects(java.util.Objects) Collectors.toList(java.util.stream.Collectors.toList) List(java.util.List) BindingData(org.flyte.api.v1.BindingData) Stream(java.util.stream.Stream) Map(java.util.Map) ArrayDeque(java.util.ArrayDeque) Comparator(java.util.Comparator) DynamicWorkflowTask(org.flyte.api.v1.DynamicWorkflowTask) ArrayList(java.util.ArrayList) RunnableTask(org.flyte.api.v1.RunnableTask)

Aggregations

DynamicWorkflowTask (org.flyte.api.v1.DynamicWorkflowTask)4 HashMap (java.util.HashMap)3 RunnableTask (org.flyte.api.v1.RunnableTask)3 TaskIdentifier (org.flyte.api.v1.TaskIdentifier)3 Literal (org.flyte.api.v1.Literal)2 PartialTaskIdentifier (org.flyte.api.v1.PartialTaskIdentifier)2 TaskTemplate (org.flyte.api.v1.TaskTemplate)2 WorkflowTemplate (org.flyte.api.v1.WorkflowTemplate)2 ArrayDeque (java.util.ArrayDeque)1 ArrayList (java.util.ArrayList)1 Collections.emptyList (java.util.Collections.emptyList)1 Collections.singletonList (java.util.Collections.singletonList)1 Comparator (java.util.Comparator)1 Deque (java.util.Deque)1 HashSet (java.util.HashSet)1 List (java.util.List)1 Map (java.util.Map)1 Objects (java.util.Objects)1 Set (java.util.Set)1 Collectors.toList (java.util.stream.Collectors.toList)1