Examples with Pipeline org.apache.beam.model.pipeline.v1.RunnerApi.Pipeline used on opensource projects

Example 41 with Pipeline

use of org.apache.beam.model.pipeline.v1.RunnerApi.Pipeline in project beam by apache.

the class InMemoryJobService method getPipeline.

@Override
public void getPipeline(GetJobPipelineRequest request, StreamObserver<GetJobPipelineResponse> responseObserver) {
    LOG.trace("{} {}", GetJobPipelineRequest.class.getSimpleName(), request);
    String invocationId = request.getJobId();
    try {
        JobInvocation invocation = getInvocation(invocationId);
        RunnerApi.Pipeline pipeline = invocation.getPipeline();
        GetJobPipelineResponse response = GetJobPipelineResponse.newBuilder().setPipeline(pipeline).build();
        responseObserver.onNext(response);
        responseObserver.onCompleted();
    } catch (StatusRuntimeException | StatusException e) {
        responseObserver.onError(e);
    } catch (Exception e) {
        String errMessage = String.format("Encountered Unexpected Exception for Invocation %s", invocationId);
        LOG.error(errMessage, e);
        responseObserver.onError(Status.INTERNAL.withCause(e).asException());
    }
}

Example 42 with Pipeline

use of org.apache.beam.model.pipeline.v1.RunnerApi.Pipeline in project beam by apache.

the class SingleEnvironmentInstanceJobBundleFactoryTest method closeShutsDownEnvironmentsWhenSomeFail.

@Test
public void closeShutsDownEnvironmentsWhenSomeFail() throws Exception {
    Pipeline p = Pipeline.create();
    ExperimentalOptions.addExperiment(p.getOptions().as(ExperimentalOptions.class), "beam_fn_api");
    p.apply("Create", Create.of(1, 2, 3));
    ExecutableStage firstEnvStage = GreedyPipelineFuser.fuse(PipelineTranslation.toProto(p)).getFusedStages().stream().findFirst().get();
    ExecutableStagePayload basePayload = ExecutableStagePayload.parseFrom(firstEnvStage.toPTransform("foo").getSpec().getPayload());
    Environment secondEnv = Environments.createDockerEnvironment("second_env");
    ExecutableStage secondEnvStage = ExecutableStage.fromPayload(basePayload.toBuilder().setEnvironment(secondEnv).build());
    Environment thirdEnv = Environments.createDockerEnvironment("third_env");
    ExecutableStage thirdEnvStage = ExecutableStage.fromPayload(basePayload.toBuilder().setEnvironment(thirdEnv).build());
    RemoteEnvironment firstRemoteEnv = mock(RemoteEnvironment.class, "First Remote Env");
    RemoteEnvironment secondRemoteEnv = mock(RemoteEnvironment.class, "Second Remote Env");
    RemoteEnvironment thirdRemoteEnv = mock(RemoteEnvironment.class, "Third Remote Env");
    when(environmentFactory.createEnvironment(firstEnvStage.getEnvironment(), GENERATED_ID)).thenReturn(firstRemoteEnv);
    when(environmentFactory.createEnvironment(secondEnvStage.getEnvironment(), GENERATED_ID)).thenReturn(secondRemoteEnv);
    when(environmentFactory.createEnvironment(thirdEnvStage.getEnvironment(), GENERATED_ID)).thenReturn(thirdRemoteEnv);
    when(firstRemoteEnv.getInstructionRequestHandler()).thenReturn(instructionRequestHandler);
    when(secondRemoteEnv.getInstructionRequestHandler()).thenReturn(instructionRequestHandler);
    when(thirdRemoteEnv.getInstructionRequestHandler()).thenReturn(instructionRequestHandler);
    factory.forStage(firstEnvStage);
    factory.forStage(secondEnvStage);
    factory.forStage(thirdEnvStage);
    IllegalStateException firstException = new IllegalStateException("first stage");
    doThrow(firstException).when(firstRemoteEnv).close();
    IllegalStateException thirdException = new IllegalStateException("third stage");
    doThrow(thirdException).when(thirdRemoteEnv).close();
    try {
        factory.close();
        fail("Factory close should have thrown");
    } catch (IllegalStateException expected) {
        if (expected.equals(firstException)) {
            assertThat(ImmutableList.copyOf(expected.getSuppressed()), contains(thirdException));
        } else if (expected.equals(thirdException)) {
            assertThat(ImmutableList.copyOf(expected.getSuppressed()), contains(firstException));
        } else {
            throw expected;
        }
        verify(firstRemoteEnv).close();
        verify(secondRemoteEnv).close();
        verify(thirdRemoteEnv).close();
    }
}

Example 43 with Pipeline

use of org.apache.beam.model.pipeline.v1.RunnerApi.Pipeline in project beam by apache.

the class DataflowRunnerTest method testSdkHarnessConfiguration.

@Test
public void testSdkHarnessConfiguration() throws IOException {
    DataflowPipelineOptions options = buildPipelineOptions();
    ExperimentalOptions.addExperiment(options, "use_runner_v2");
    Pipeline p = Pipeline.create(options);
    p.apply(Create.of(Arrays.asList(1, 2, 3)));
    String defaultSdkContainerImage = DataflowRunner.getContainerImageForJob(options);
    SdkComponents sdkComponents = SdkComponents.create();
    RunnerApi.Environment defaultEnvironmentForDataflow = Environments.createDockerEnvironment(defaultSdkContainerImage);
    sdkComponents.registerEnvironment(defaultEnvironmentForDataflow.toBuilder().build());
    RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(p, sdkComponents, true);
    Job job = DataflowPipelineTranslator.fromOptions(options).translate(p, pipelineProto, sdkComponents, DataflowRunner.fromOptions(options), Collections.emptyList()).getJob();
    DataflowRunner.configureSdkHarnessContainerImages(options, pipelineProto, job);
    List<SdkHarnessContainerImage> sdks = job.getEnvironment().getWorkerPools().get(0).getSdkHarnessContainerImages();
    Map<String, String> expectedEnvIdsAndContainerImages = pipelineProto.getComponents().getEnvironmentsMap().entrySet().stream().filter(x -> BeamUrns.getUrn(RunnerApi.StandardEnvironments.Environments.DOCKER).equals(x.getValue().getUrn())).collect(Collectors.toMap(x -> x.getKey(), x -> {
        RunnerApi.DockerPayload payload;
        try {
            payload = RunnerApi.DockerPayload.parseFrom(x.getValue().getPayload());
        } catch (InvalidProtocolBufferException e) {
            throw new RuntimeException(e);
        }
        return payload.getContainerImage();
    }));
    assertEquals(1, expectedEnvIdsAndContainerImages.size());
    assertEquals(1, sdks.size());
    assertEquals(expectedEnvIdsAndContainerImages, sdks.stream().collect(Collectors.toMap(SdkHarnessContainerImage::getEnvironmentId, SdkHarnessContainerImage::getContainerImage)));
}

Example 44 with Pipeline

use of org.apache.beam.model.pipeline.v1.RunnerApi.Pipeline in project beam by apache.

the class DataflowPipelineTranslatorTest method testMultiGraphPipelineSerialization.

@Test
public void testMultiGraphPipelineSerialization() throws Exception {
    DataflowPipelineOptions options = buildPipelineOptions();
    Pipeline p = Pipeline.create(options);
    PCollection<Integer> input = p.begin().apply(Create.of(1, 2, 3));
    input.apply(new UnrelatedOutputCreator());
    input.apply(new UnboundOutputCreator());
    DataflowPipelineTranslator t = DataflowPipelineTranslator.fromOptions(PipelineOptionsFactory.as(DataflowPipelineOptions.class));
    // Check that translation doesn't fail.
    SdkComponents sdkComponents = createSdkComponents(options);
    RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(p, sdkComponents, true);
    JobSpecification jobSpecification = t.translate(p, pipelineProto, sdkComponents, DataflowRunner.fromOptions(options), Collections.emptyList());
    assertAllStepOutputsHaveUniqueIds(jobSpecification.getJob());
}

Example 45 with Pipeline

use of org.apache.beam.model.pipeline.v1.RunnerApi.Pipeline in project beam by apache.

the class DataflowPipelineTranslatorTest method testStreamingSplittableParDoTranslation.

/**
 * Smoke test to fail fast if translation of a splittable ParDo in streaming breaks.
 */
@Test
public void testStreamingSplittableParDoTranslation() throws Exception {
    DataflowPipelineOptions options = buildPipelineOptions();
    DataflowRunner runner = DataflowRunner.fromOptions(options);
    options.setStreaming(true);
    DataflowPipelineTranslator translator = DataflowPipelineTranslator.fromOptions(options);
    Pipeline pipeline = Pipeline.create(options);
    PCollection<String> windowedInput = pipeline.apply(Create.of("a")).apply(Window.into(FixedWindows.of(Duration.standardMinutes(1))));
    windowedInput.apply(ParDo.of(new TestSplittableFn()));
    runner.replaceV1Transforms(pipeline);
    SdkComponents sdkComponents = createSdkComponents(options);
    RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(pipeline, sdkComponents, true);
    Job job = translator.translate(pipeline, pipelineProto, sdkComponents, runner, Collections.emptyList()).getJob();
    // The job should contain a SplittableParDo.ProcessKeyedElements step, translated as
    // "SplittableProcessKeyed".
    List<Step> steps = job.getSteps();
    Step processKeyedStep = null;
    for (Step step : steps) {
        if ("SplittableProcessKeyed".equals(step.getKind())) {
            assertNull(processKeyedStep);
            processKeyedStep = step;
        }
    }
    assertNotNull(processKeyedStep);
    @SuppressWarnings({ "unchecked", "rawtypes" }) DoFnInfo<String, Integer> fnInfo = (DoFnInfo<String, Integer>) SerializableUtils.deserializeFromByteArray(jsonStringToByteArray(getString(processKeyedStep.getProperties(), PropertyNames.SERIALIZED_FN)), "DoFnInfo");
    assertThat(fnInfo.getDoFn(), instanceOf(TestSplittableFn.class));
    assertThat(fnInfo.getWindowingStrategy().getWindowFn(), Matchers.<WindowFn>equalTo(FixedWindows.of(Duration.standardMinutes(1))));
    assertThat(fnInfo.getInputCoder(), instanceOf(StringUtf8Coder.class));
    Coder<?> restrictionCoder = CloudObjects.coderFromCloudObject((CloudObject) Structs.getObject(processKeyedStep.getProperties(), PropertyNames.RESTRICTION_CODER));
    assertEquals(KvCoder.of(SerializableCoder.of(OffsetRange.class), VoidCoder.of()), restrictionCoder);
}