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