Examples with RemoteGrpcPort org.apache.beam.model.fnexecution.v1.BeamFnApi.RemoteGrpcPort used on opensource projects

Example 1 with RemoteGrpcPort

use of org.apache.beam.model.fnexecution.v1.BeamFnApi.RemoteGrpcPort in project beam by apache.

the class SdkHarnessClientTest method setup.

@Before
public void setup() throws Exception {
    MockitoAnnotations.initMocks(this);
    sdkHarnessClient = SdkHarnessClient.usingFnApiClient(fnApiControlClient, dataService);
    Pipeline userPipeline = Pipeline.create();
    TupleTag<String> outputTag = new TupleTag<>();
    userPipeline.apply("create", Create.of("foo")).apply("proc", ParDo.of(new TestFn()).withOutputTags(outputTag, TupleTagList.empty()));
    RunnerApi.Pipeline userProto = PipelineTranslation.toProto(userPipeline);
    ProcessBundleDescriptor.Builder pbdBuilder = ProcessBundleDescriptor.newBuilder().setId("my_id").putAllEnvironments(userProto.getComponents().getEnvironmentsMap()).putAllWindowingStrategies(userProto.getComponents().getWindowingStrategiesMap()).putAllCoders(userProto.getComponents().getCodersMap());
    RunnerApi.Coder fullValueCoder = CoderTranslation.toProto(WindowedValue.getFullCoder(StringUtf8Coder.of(), Coder.INSTANCE)).getCoder();
    pbdBuilder.putCoders("wire_coder", fullValueCoder);
    PTransform targetProcessor = userProto.getComponents().getTransformsOrThrow("proc");
    RemoteGrpcPort port = RemoteGrpcPort.newBuilder().setApiServiceDescriptor(harness.dataEndpoint()).setCoderId("wire_coder").build();
    RemoteGrpcPortRead readNode = RemoteGrpcPortRead.readFromPort(port, getOnlyElement(targetProcessor.getInputsMap().values()));
    RemoteGrpcPortWrite writeNode = RemoteGrpcPortWrite.writeToPort(getOnlyElement(targetProcessor.getOutputsMap().values()), port);
    // TODO: Ensure cross-env (Runner <-> SDK GRPC Read/Write Node) coders are length-prefixed
    for (String pc : targetProcessor.getInputsMap().values()) {
        pbdBuilder.putPcollections(pc, userProto.getComponents().getPcollectionsOrThrow(pc));
    }
    for (String pc : targetProcessor.getOutputsMap().values()) {
        pbdBuilder.putPcollections(pc, userProto.getComponents().getPcollectionsOrThrow(pc));
    }
    pbdBuilder.putTransforms("proc", targetProcessor).putTransforms(SDK_GRPC_READ_TRANSFORM, readNode.toPTransform()).putTransforms(SDK_GRPC_WRITE_TRANSFORM, writeNode.toPTransform());
    descriptor = pbdBuilder.build();
}

Example 2 with RemoteGrpcPort

use of org.apache.beam.model.fnexecution.v1.BeamFnApi.RemoteGrpcPort in project beam by apache.

the class ProcessBundleDescriptors method addStageOutput.

private static OutputEncoding addStageOutput(ApiServiceDescriptor dataEndpoint, Components.Builder components, PCollectionNode outputPCollection, WireCoderSetting wireCoderSetting) throws IOException {
    String outputWireCoderId = WireCoders.addSdkWireCoder(outputPCollection, components, wireCoderSetting);
    @SuppressWarnings("unchecked") Coder<WindowedValue<?>> wireCoder = (Coder) WireCoders.instantiateRunnerWireCoder(outputPCollection, components.build(), wireCoderSetting);
    RemoteGrpcPort outputPort = RemoteGrpcPort.newBuilder().setApiServiceDescriptor(dataEndpoint).setCoderId(outputWireCoderId).build();
    RemoteGrpcPortWrite outputWrite = RemoteGrpcPortWrite.writeToPort(outputPCollection.getId(), outputPort);
    String outputId = uniqueId(String.format("fn/write/%s", outputPCollection.getId()), components::containsTransforms);
    PTransform outputTransform = outputWrite.toPTransform();
    components.putTransforms(outputId, outputTransform);
    return new AutoValue_ProcessBundleDescriptors_OutputEncoding(outputId, wireCoder);
}

Example 3 with RemoteGrpcPort

use of org.apache.beam.model.fnexecution.v1.BeamFnApi.RemoteGrpcPort in project beam by apache.

the class RemoteGrpcPortReadTest method getPortSucceeds.

@Test
public void getPortSucceeds() {
    RemoteGrpcPort port = RemoteGrpcPort.newBuilder().setApiServiceDescriptor(ApiServiceDescriptor.newBuilder().setUrl("foo").setAuthentication(AuthenticationSpec.getDefaultInstance()).build()).build();
    RemoteGrpcPortRead read = RemoteGrpcPortRead.readFromPort(port, "myPort");
    assertThat(read.getPort(), equalTo(port));
}

Example 4 with RemoteGrpcPort

use of org.apache.beam.model.fnexecution.v1.BeamFnApi.RemoteGrpcPort in project beam by apache.

the class RemoteGrpcPortReadTest method toFromPTransform.

@Test
public void toFromPTransform() throws InvalidProtocolBufferException {
    RemoteGrpcPort port = RemoteGrpcPort.newBuilder().setApiServiceDescriptor(ApiServiceDescriptor.newBuilder().setUrl("foo").setAuthentication(AuthenticationSpec.getDefaultInstance()).build()).build();
    RemoteGrpcPortRead read = RemoteGrpcPortRead.readFromPort(port, "myPort");
    PTransform ptransform = PTransform.parseFrom(read.toPTransform().toByteArray());
    RemoteGrpcPortRead serDeRead = RemoteGrpcPortRead.fromPTransform(ptransform);
    assertThat(serDeRead, equalTo(read));
    assertThat(serDeRead.getPort(), equalTo(read.getPort()));
    assertThat(serDeRead.toPTransform(), equalTo(ptransform));
}

Example 5 with RemoteGrpcPort

use of org.apache.beam.model.fnexecution.v1.BeamFnApi.RemoteGrpcPort in project beam by apache.

the class RemoteGrpcPortRead method fromPTransform.

public static RemoteGrpcPortRead fromPTransform(PTransform pTransform) throws InvalidProtocolBufferException {
    checkArgument(URN.equals(pTransform.getSpec().getUrn()), "Expected URN for %s, got %s", RemoteGrpcPortRead.class.getSimpleName(), pTransform.getSpec().getUrn());
    checkArgument(pTransform.getOutputsCount() == 1, "Expected exactly one output, got %s", pTransform.getOutputsCount());
    RemoteGrpcPort port = RemoteGrpcPort.parseFrom(pTransform.getSpec().getPayload());
    String outputPcollection = Iterables.getOnlyElement(pTransform.getOutputsMap().values());
    return readFromPort(port, outputPcollection);
}