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

Example 46 with Environment

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

the class DefaultJobBundleFactoryTest method createsMultipleEnvironmentsWithSdkWorkerParallelism.

@Test
public void createsMultipleEnvironmentsWithSdkWorkerParallelism() throws Exception {
    ServerFactory serverFactory = ServerFactory.createDefault();
    Environment environmentA = Environment.newBuilder().setUrn("env:urn:a").setPayload(ByteString.copyFrom(new byte[1])).build();
    EnvironmentFactory envFactoryA = mock(EnvironmentFactory.class);
    when(envFactoryA.createEnvironment(eq(environmentA), any())).thenReturn(remoteEnvironment);
    EnvironmentFactory.Provider environmentProviderFactoryA = mock(EnvironmentFactory.Provider.class);
    when(environmentProviderFactoryA.createEnvironmentFactory(any(), any(), any(), any(), any(), any())).thenReturn(envFactoryA);
    when(environmentProviderFactoryA.getServerFactory()).thenReturn(serverFactory);
    Map<String, Provider> environmentFactoryProviderMap = ImmutableMap.of(environmentA.getUrn(), environmentProviderFactoryA);
    PortablePipelineOptions portableOptions = PipelineOptionsFactory.as(PortablePipelineOptions.class);
    portableOptions.setSdkWorkerParallelism(2);
    Struct pipelineOptions = PipelineOptionsTranslation.toProto(portableOptions);
    try (DefaultJobBundleFactory bundleFactory = new DefaultJobBundleFactory(JobInfo.create("testJob", "testJob", "token", pipelineOptions), environmentFactoryProviderMap, stageIdGenerator, serverInfo)) {
        bundleFactory.forStage(getExecutableStage(environmentA));
        verify(environmentProviderFactoryA, Mockito.times(1)).createEnvironmentFactory(any(), any(), any(), any(), any(), any());
        verify(envFactoryA, Mockito.times(1)).createEnvironment(eq(environmentA), any());
        bundleFactory.forStage(getExecutableStage(environmentA));
        verify(environmentProviderFactoryA, Mockito.times(2)).createEnvironmentFactory(any(), any(), any(), any(), any(), any());
        verify(envFactoryA, Mockito.times(2)).createEnvironment(eq(environmentA), any());
        // round robin, no new environment created
        bundleFactory.forStage(getExecutableStage(environmentA));
        verify(environmentProviderFactoryA, Mockito.times(2)).createEnvironmentFactory(any(), any(), any(), any(), any(), any());
        verify(envFactoryA, Mockito.times(2)).createEnvironment(eq(environmentA), any());
    }
    portableOptions.setSdkWorkerParallelism(0);
    pipelineOptions = PipelineOptionsTranslation.toProto(portableOptions);
    Mockito.reset(envFactoryA);
    when(envFactoryA.createEnvironment(eq(environmentA), any())).thenReturn(remoteEnvironment);
    int expectedParallelism = Math.max(1, Runtime.getRuntime().availableProcessors() - 1);
    try (DefaultJobBundleFactory bundleFactory = new DefaultJobBundleFactory(JobInfo.create("testJob", "testJob", "token", pipelineOptions), environmentFactoryProviderMap, stageIdGenerator, serverInfo)) {
        HashSet<StageBundleFactory> stageBundleFactorySet = new HashSet<>();
        // more factories than parallelism for round-robin
        int numStageBundleFactories = expectedParallelism + 5;
        for (int i = 0; i < numStageBundleFactories; i++) {
            stageBundleFactorySet.add(bundleFactory.forStage(getExecutableStage(environmentA)));
        }
        verify(envFactoryA, Mockito.times(expectedParallelism)).createEnvironment(eq(environmentA), any());
        Assert.assertEquals(numStageBundleFactories, stageBundleFactorySet.size());
    }
}

Example 47 with Environment

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

the class ExternalEnvironmentFactory method createEnvironment.

/**
 * Creates a new, active {@link RemoteEnvironment} backed by an unmanaged worker.
 */
@Override
public RemoteEnvironment createEnvironment(Environment environment, String workerId) throws Exception {
    Preconditions.checkState(environment.getUrn().equals(BeamUrns.getUrn(RunnerApi.StandardEnvironments.Environments.EXTERNAL)), "The passed environment does not contain an ExternalPayload.");
    final RunnerApi.ExternalPayload externalPayload = RunnerApi.ExternalPayload.parseFrom(environment.getPayload());
    BeamFnApi.StartWorkerRequest startWorkerRequest = BeamFnApi.StartWorkerRequest.newBuilder().setWorkerId(workerId).setControlEndpoint(controlServiceServer.getApiServiceDescriptor()).setLoggingEndpoint(loggingServiceServer.getApiServiceDescriptor()).setArtifactEndpoint(retrievalServiceServer.getApiServiceDescriptor()).setProvisionEndpoint(provisioningServiceServer.getApiServiceDescriptor()).putAllParams(externalPayload.getParamsMap()).build();
    LOG.debug("Requesting worker ID {}", workerId);
    final ManagedChannel managedChannel = ManagedChannelFactory.createDefault().forDescriptor(externalPayload.getEndpoint());
    BeamFnApi.StartWorkerResponse startWorkerResponse = BeamFnExternalWorkerPoolGrpc.newBlockingStub(managedChannel).startWorker(startWorkerRequest);
    if (!startWorkerResponse.getError().isEmpty()) {
        throw new RuntimeException(startWorkerResponse.getError());
    }
    // Wait on a client from the gRPC server.
    InstructionRequestHandler instructionHandler = null;
    while (instructionHandler == null) {
        try {
            instructionHandler = clientSource.take(workerId, Duration.ofMinutes(2));
        } catch (TimeoutException timeoutEx) {
            LOG.info("Still waiting for startup of environment from {} for worker id {}", externalPayload.getEndpoint().getUrl(), workerId);
        } catch (InterruptedException interruptEx) {
            Thread.currentThread().interrupt();
            managedChannel.shutdownNow();
            throw new RuntimeException(interruptEx);
        }
    }
    final InstructionRequestHandler finalInstructionHandler = instructionHandler;
    return new RemoteEnvironment() {

        @Override
        public Environment getEnvironment() {
            return environment;
        }

        @Override
        public InstructionRequestHandler getInstructionRequestHandler() {
            return finalInstructionHandler;
        }

        @Override
        public void close() throws Exception {
            try {
                finalInstructionHandler.close();
                BeamFnApi.StopWorkerRequest stopWorkerRequest = BeamFnApi.StopWorkerRequest.newBuilder().setWorkerId(workerId).build();
                LOG.debug("Closing worker ID {}", workerId);
                BeamFnApi.StopWorkerResponse stopWorkerResponse = BeamFnExternalWorkerPoolGrpc.newBlockingStub(managedChannel).stopWorker(stopWorkerRequest);
                if (!stopWorkerResponse.getError().isEmpty()) {
                    throw new RuntimeException(stopWorkerResponse.getError());
                }
            } finally {
                managedChannel.shutdown();
                managedChannel.awaitTermination(10, TimeUnit.SECONDS);
                if (!managedChannel.isTerminated()) {
                    managedChannel.shutdownNow();
                }
            }
        }
    };
}

Example 48 with Environment

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

the class BatchSideInputHandlerFactoryTest method createExecutableStage.

private static ExecutableStage createExecutableStage(Collection<SideInputReference> sideInputs) {
    Components components = Components.getDefaultInstance();
    Environment environment = Environment.getDefaultInstance();
    PCollectionNode inputCollection = PipelineNode.pCollection("collection-id", RunnerApi.PCollection.getDefaultInstance());
    return ImmutableExecutableStage.of(components, environment, inputCollection, sideInputs, Collections.emptyList(), Collections.emptyList(), Collections.emptyList(), Collections.emptyList(), DEFAULT_WIRE_CODER_SETTINGS);
}

Example 49 with Environment

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

the class ProcessEnvironmentFactoryTest method createsMultipleEnvironments.

@Test
public void createsMultipleEnvironments() throws Exception {
    Environment fooEnv = Environments.createProcessEnvironment("", "", "foo", Collections.emptyMap());
    RemoteEnvironment fooHandle = factory.createEnvironment(fooEnv, "workerId");
    assertThat(fooHandle.getEnvironment(), is(equalTo(fooEnv)));
    Environment barEnv = Environments.createProcessEnvironment("", "", "bar", Collections.emptyMap());
    RemoteEnvironment barHandle = factory.createEnvironment(barEnv, "workerId");
    assertThat(barHandle.getEnvironment(), is(equalTo(barEnv)));
}