use of org.flyte.api.v1.TaskIdentifier in project flytekit-java by flyteorg.
the class IdentifierRewrite method getLatestTaskVersion.
private String getLatestTaskVersion(String project, String domain, String name) {
TaskIdentifier latestTaskId = adminClient().fetchLatestTaskId(NamedEntityIdentifier.builder().domain(domain).project(project).name(name).build());
Verify.verifyNotNull(latestTaskId, "task not found domain=[%s], project=[%s], name=[%s]", domain, project, name);
return latestTaskId.version();
}
use of org.flyte.api.v1.TaskIdentifier in project flytekit-java by flyteorg.
the class FlyteAdminClientTest method fetchLatestTaskIdShouldReturnFirstTaskFromList.
@Test
public void fetchLatestTaskIdShouldReturnFirstTaskFromList() {
stubService.taskLists = Arrays.asList(TaskOuterClass.Task.newBuilder().setId(newIdentifier(ResourceType.TASK, TASK_NAME, TASK_VERSION)).build(), TaskOuterClass.Task.newBuilder().setId(newIdentifier(ResourceType.TASK, TASK_NAME, TASK_OLD_VERSION)).build());
TaskIdentifier taskId = client.fetchLatestTaskId(NamedEntityIdentifier.builder().project(PROJECT).domain(DOMAIN).name(TASK_NAME).build());
assertThat(taskId, equalTo(TaskIdentifier.builder().project(PROJECT).domain(DOMAIN).name(TASK_NAME).version(TASK_VERSION).build()));
}
use of org.flyte.api.v1.TaskIdentifier in project flytekit-java by flyteorg.
the class FlyteAdminClientTest method fetchLatestTaskIdShouldReturnNullWhenEmptyList.
@Test
public void fetchLatestTaskIdShouldReturnNullWhenEmptyList() {
stubService.taskLists = Collections.emptyList();
TaskIdentifier taskId = client.fetchLatestTaskId(NamedEntityIdentifier.builder().project(PROJECT).domain(DOMAIN).name(TASK_NAME).build());
assertThat(taskId, nullValue());
}
use of org.flyte.api.v1.TaskIdentifier 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()));
}
use of org.flyte.api.v1.TaskIdentifier 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));
}
}
Aggregations